From afbcc51672cbcdb10cd5886a5351f034ed70ea1d Mon Sep 17 00:00:00 2001 From: rein yang <73573651+R2-Y@users.noreply.github.com> Date: Sun, 21 Jun 2026 22:08:26 +0800 Subject: [PATCH] [Model] Aura support - Non async chunk path (#4257) Signed-off-by: Rein Yang Signed-off-by: R2-Y --- .../examples/online_serving/aura_omni.md | 65 +++ .../offline_inference/aura_omni/README.md | 61 +++ .../offline_inference/aura_omni/end2end.py | 236 +++++++++++ .../aura_omni/run_single_prompt.sh | 8 + examples/online_serving/aura_omni/README.md | 146 +++++++ .../online_serving/aura_omni/gradio_demo.py | 305 ++++++++++++++ .../openai_chat_completion_client.py | 199 +++++++++ .../run_curl_multimodal_generation.sh | 60 +++ .../aura_omni/run_gradio_demo.sh | 79 ++++ tests/config/test_pipeline_registry.py | 16 + .../test_async_omni_engine_stage_init.py | 18 + tests/engine/test_orchestrator.py | 21 +- .../test_orchestrator_stage_input_bridge.py | 156 +++++++ ...test_serving_chat_multistage_generation.py | 76 +++- .../test_serving_chat_sampling_params.py | 31 ++ .../models/test_aura_qwen3_vl.py | 69 ++++ .../stage_input_processors/test_aura_omni.py | 197 +++++++++ tests/test_config_factory.py | 28 ++ tests/worker/test_omni_gpu_model_runner.py | 34 +- vllm_omni/config/pipeline_registry.py | 4 + vllm_omni/deploy/aura_omni.yaml | 110 +++++ vllm_omni/engine/orchestrator.py | 187 ++++++++- vllm_omni/engine/stage_init_utils.py | 1 + vllm_omni/entrypoints/openai/serving_chat.py | 286 ++++++++++++- .../models/aura_omni/pipeline.py | 84 ++++ .../models/aura_omni/qwen3_vl.py | 59 +++ .../models/qwen3_tts/qwen3_tts_talker.py | 26 +- vllm_omni/model_executor/models/registry.py | 5 + .../stage_input_processors/aura_omni.py | 390 ++++++++++++++++++ vllm_omni/worker/gpu_model_runner.py | 27 +- 30 files changed, 2958 insertions(+), 26 deletions(-) create mode 100644 docs/user_guide/examples/online_serving/aura_omni.md create mode 100644 examples/offline_inference/aura_omni/README.md create mode 100644 examples/offline_inference/aura_omni/end2end.py create mode 100755 examples/offline_inference/aura_omni/run_single_prompt.sh create mode 100644 examples/online_serving/aura_omni/README.md create mode 100644 examples/online_serving/aura_omni/gradio_demo.py create mode 100644 examples/online_serving/aura_omni/openai_chat_completion_client.py create mode 100755 examples/online_serving/aura_omni/run_curl_multimodal_generation.sh create mode 100755 examples/online_serving/aura_omni/run_gradio_demo.sh create mode 100644 tests/engine/test_orchestrator_stage_input_bridge.py create mode 100644 tests/model_executor/models/test_aura_qwen3_vl.py create mode 100644 tests/model_executor/stage_input_processors/test_aura_omni.py create mode 100644 vllm_omni/deploy/aura_omni.yaml create mode 100644 vllm_omni/model_executor/models/aura_omni/pipeline.py create mode 100644 vllm_omni/model_executor/models/aura_omni/qwen3_vl.py create mode 100644 vllm_omni/model_executor/stage_input_processors/aura_omni.py diff --git a/docs/user_guide/examples/online_serving/aura_omni.md b/docs/user_guide/examples/online_serving/aura_omni.md new file mode 100644 index 000000000..9c66d0df9 --- /dev/null +++ b/docs/user_guide/examples/online_serving/aura_omni.md @@ -0,0 +1,65 @@ +# AURA Omni Pipeline + +`aura_omni` wires ASR, AURA, and Qwen3-TTS into one vLLM-Omni pipeline: + +```text +ASR -> AURA -> Qwen3-TTS Talker -> Code2Wav +``` + +Qwen3-TTS remains two engine stages so the pipeline reuses the existing native +Talker and Code2Wav implementation. + +```bash +vllm serve aurateam/AURA \ + --omni \ + --deploy-config vllm_omni/deploy/aura_omni.yaml \ + --served-model-name aurateam/AURA \ + --trust-remote-code +``` + +Configure local checkpoints by editing per-stage `model` values in +`vllm_omni/deploy/aura_omni.yaml`. The deploy file sets +`pipeline: aura_omni`, so the four-stage topology is used even if the +command-line `--model` points at one of the component checkpoints. + +Send requests with `"model": "aurateam/AURA"`. The ASR, AURA, and Qwen3-TTS +checkpoint paths are internal stage models from the deploy YAML, not the +OpenAI-facing served model name. + +The AURA stage can emit `<|silent|>`. Silent outputs are treated as a gate: +they produce no Qwen3-TTS Talker input, so no audio is synthesized for that +turn. + +## GPU Utilization Recommendation + +`gpu_memory_utilization` in `vllm_omni/deploy/aura_omni.yaml` controls how much +VRAM each stage can reserve. Start with this split for a single GPU: + +- Stage 0 (ASR): `0.10` +- Stage 1 (AURA): `0.40` +- Stage 2 (Qwen3-TTS Talker): `0.20` +- Stage 3 (Qwen3-TTS Code2Wav): `0.20` + +## TTS Modes + +`aura_omni` can pass AURA text to Qwen3-TTS in two task modes: + +- `Base`: voice clone from `tts_ref_audio` with ICL enabled in the AURA + pipeline. Provide both `tts_ref_audio` and `tts_ref_text`. Set + `tts_x_vector_only_mode=true` to disable ICL and use speaker embedding only. +- `CustomVoice`: predefined speaker mode. Use a Qwen3-TTS CustomVoice + checkpoint for stages 2 and 3 in `aura_omni.yaml`, then pass + `tts_task_type=CustomVoice` and `tts_speaker`. + +By default, AURA responses are passed to Qwen3-TTS as text. Set +`additional_information.tts_pass_token_ids=true` to pass AURA-generated +assistant token ids directly instead. Even when token passthrough is disabled, +the stage processor uses AURA token ids when available to estimate the Talker +prefill length, so scheduling does not rely on raw character length. + +The example client exposes this as: + +```bash +python examples/online_serving/aura_omni/openai_chat_completion_client.py \ + --tts-pass-token-ids +``` diff --git a/examples/offline_inference/aura_omni/README.md b/examples/offline_inference/aura_omni/README.md new file mode 100644 index 000000000..37fabf1ba --- /dev/null +++ b/examples/offline_inference/aura_omni/README.md @@ -0,0 +1,61 @@ +# AURA Omni + +This example runs the native AURA Omni pipeline offline: + +```text +Qwen3-ASR -> AURA/Qwen3-VL -> Qwen3-TTS Talker -> Qwen3-TTS Code2Wav +``` + +The first stage consumes speech and produces a transcript. The AURA stage then combines the transcript with the video frames from the original request and returns text or `<|silent|>`. Non-silent responses are passed to Qwen3-TTS as AURA-generated token ids. + +## Run + +```bash +cd examples/offline_inference/aura_omni +bash run_single_prompt.sh +``` + +Use local media: + +```bash +python end2end.py \ + --audio-path /path/to/input.wav \ + --video-path /path/to/video.mp4 \ + --modalities text,audio +``` + +Base voice clone mode is the default. It uses the AURA reference audio unless +you override it: + +```bash +python end2end.py \ + --tts-task-type Base \ + --tts-ref-audio /data/yrr/rein_test/shuhan.mp3 \ + --tts-x-vector-only-mode +``` + +CustomVoice mode requires stages 2 and 3 in `vllm_omni/deploy/aura_omni.yaml` +to point at a Qwen3-TTS CustomVoice checkpoint: + +```bash +python end2end.py \ + --tts-task-type CustomVoice \ + --tts-speaker Vivian +``` + +For local checkpoints, edit the stage `model` entries in +`vllm_omni/deploy/aura_omni.yaml` or pass a copied deploy config with +`--deploy-config`. + +## GPU Utilization Recommendation + +Set `gpu_memory_utilization` in `vllm_omni/deploy/aura_omni.yaml` by stage. +Suggested starting point: + +- Stage 0 (ASR): `0.10` +- Stage 1 (AURA): `0.4` +- Stage 2 (Qwen3-TTS Talker): `0.20` +- Stage 3 (Qwen3-TTS Code2Wav): `0.20` + +Generated text and audio are written to `--output-dir` +(default: `output_aura_omni`). diff --git a/examples/offline_inference/aura_omni/end2end.py b/examples/offline_inference/aura_omni/end2end.py new file mode 100644 index 000000000..c08287515 --- /dev/null +++ b/examples/offline_inference/aura_omni/end2end.py @@ -0,0 +1,236 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Offline AURA Omni example. + +The native AURA pipeline is: + + Qwen3-ASR -> AURA/Qwen3-VL -> Qwen3-TTS Talker -> Code2Wav + +Stage 0 consumes audio and emits a transcript. The stage input processor then +combines that transcript with the original visual payload for AURA. +""" + +from __future__ import annotations + +import os +from typing import NamedTuple + +import numpy as np +import soundfile as sf +import vllm +from vllm import SamplingParams +from vllm.assets.audio import AudioAsset +from vllm.assets.video import VideoAsset, video_to_ndarrays +from vllm.multimodal.media.audio import load_audio +from vllm.utils.argparse_utils import FlexibleArgumentParser + +from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults +from vllm_omni.entrypoints.omni import Omni +from vllm_omni.model_executor.stage_input_processors.aura_omni import ( + DEFAULT_QWEN3_TTS_REF_TEXT, + default_qwen3_tts_ref_audio_path, +) + +SEED = 42 +DEFAULT_MODEL = "aurateam/AURA" +DEFAULT_DEPLOY_CONFIG = "vllm_omni/deploy/aura_omni.yaml" + + +class QueryResult(NamedTuple): + inputs: dict + limit_mm_per_prompt: dict[str, int] + + +def _load_audio(audio_path: str | None, sampling_rate: int) -> tuple[np.ndarray, int]: + if audio_path: + if not os.path.exists(audio_path): + raise FileNotFoundError(f"Audio file not found: {audio_path}") + audio_signal, sr = load_audio(audio_path, sr=sampling_rate) + return audio_signal.astype(np.float32), int(sr) + audio_signal, sr = AudioAsset("mary_had_lamb").audio_and_sample_rate + return audio_signal.astype(np.float32), int(sr) + + +def _load_video(video_path: str | None, num_frames: int) -> list[np.ndarray]: + if video_path: + if not os.path.exists(video_path): + raise FileNotFoundError(f"Video file not found: {video_path}") + return video_to_ndarrays(video_path, num_frames=num_frames) + return VideoAsset(name="baby_reading", num_frames=num_frames).np_ndarrays + + +def build_query(args) -> QueryResult: + prompt = ( + "<|im_start|>user\n" + "<|audio_start|><|audio_pad|><|audio_end|>" + "Please transcribe this speech and use the video context to decide whether a reply is needed.<|im_end|>\n" + "<|im_start|>assistant\n" + ) + audio_data = _load_audio(args.audio_path, args.sampling_rate) + video_frames = _load_video(args.video_path, args.num_frames) + return QueryResult( + inputs={ + "prompt": prompt, + "multi_modal_data": { + "audio": audio_data, + "video": video_frames, + }, + "additional_information": { + "aura_system_prompt": args.aura_system_prompt, + "tts_task_type": args.tts_task_type, + "tts_language": args.tts_language, + "tts_speaker": args.tts_speaker, + "tts_instruct": args.tts_instruct, + "tts_ref_audio": args.tts_ref_audio, + "tts_ref_text": args.tts_ref_text, + "tts_x_vector_only_mode": args.tts_x_vector_only_mode, + "tts_pass_token_ids": args.tts_pass_token_ids, + }, + }, + limit_mm_per_prompt={"audio": 1, "video": 1}, + ) + + +def build_sampling_params() -> list[SamplingParams]: + asr_sampling = SamplingParams( + temperature=0.0, + top_p=1.0, + top_k=-1, + max_tokens=256, + detokenize=True, + seed=SEED, + ) + aura_sampling = SamplingParams( + temperature=0.5, + top_p=1.0, + top_k=-1, + max_tokens=256, + detokenize=True, + repetition_penalty=1.0, + seed=SEED, + ) + talker_sampling = SamplingParams( + temperature=0.9, + top_k=50, + max_tokens=4096, + detokenize=False, + repetition_penalty=1.05, + stop_token_ids=[2150], + seed=SEED, + ) + code2wav_sampling = SamplingParams( + temperature=0.0, + top_p=1.0, + top_k=-1, + max_tokens=65536, + detokenize=True, + repetition_penalty=1.0, + seed=SEED, + ) + return [asr_sampling, aura_sampling, talker_sampling, code2wav_sampling] + + +def save_stage_output(stage_outputs, output_dir: str) -> None: + output = stage_outputs.request_output + request_id = output.request_id + if stage_outputs.final_output_type == "text": + text_output = output.outputs[0].text + out_txt = os.path.join(output_dir, f"{request_id}.txt") + with open(out_txt, "w", encoding="utf-8") as f: + f.write(str(text_output).strip() + "\n") + print(f"Request ID: {request_id}, text saved to {out_txt}") + print(text_output) + return + + if stage_outputs.final_output_type == "audio": + mm = output.outputs[0].multimodal_output + if not mm or "audio" not in mm: + print(f"Request ID: {request_id}, no audio output") + return + audio = mm["audio"] + if isinstance(audio, list): + import torch + + audio = torch.cat([(item if hasattr(item, "detach") else torch.tensor(item)).flatten() for item in audio]) + sample_rate = int(mm.get("sr", 24000)) + audio_numpy = audio.float().detach().cpu().numpy().flatten() + out_wav = os.path.join(output_dir, f"{request_id}.wav") + sf.write(out_wav, audio_numpy, samplerate=sample_rate, format="WAV") + print(f"Request ID: {request_id}, audio saved to {out_wav}") + + +def main(args) -> None: + print("=" * 20, "\n", f"vllm version: {vllm.__version__}", "\n", "=" * 20) + os.makedirs(args.output_dir, exist_ok=True) + + query = build_query(args) + prompts = [query.inputs for _ in range(args.num_prompts)] + if args.modalities: + modalities = [item.strip() for item in args.modalities.split(",") if item.strip()] + for prompt in prompts: + prompt["modalities"] = modalities + + omni_kwargs = vars(args).copy() + omni_kwargs["model"] = args.model + omni = Omni(**omni_kwargs) + sampling_params_list = build_sampling_params()[: omni.num_stages] + + try: + for stage_outputs in omni.generate(prompts, sampling_params_list): + save_stage_output(stage_outputs, args.output_dir) + finally: + omni.close() + + +def parse_args(): + parser = FlexibleArgumentParser(description="AURA Omni offline inference example") + parser.add_argument("--model", default=DEFAULT_MODEL, help="Top-level model name/path.") + parser.add_argument("--deploy-config", default=DEFAULT_DEPLOY_CONFIG, help="AURA Omni deploy config.") + parser.add_argument("--audio-path", "-a", default=None, help="Local audio file. Uses a built-in asset by default.") + parser.add_argument("--video-path", "-v", default=None, help="Local video file. Uses a built-in asset by default.") + parser.add_argument("--num-frames", type=int, default=16, help="Number of video frames to sample.") + parser.add_argument("--sampling-rate", type=int, default=16000, help="Audio sampling rate for ASR input.") + parser.add_argument("--num-prompts", type=int, default=1, help="Number of repeated prompts.") + parser.add_argument("--modalities", default="text,audio", help="Comma-separated final modalities.") + parser.add_argument("--output-dir", default="output_aura_omni", help="Directory for generated files.") + parser.add_argument( + "--aura-system-prompt", + default=( + "You are receiving a live video stream where the final frame is the present moment. " + "Respond only when a response is needed. Otherwise output '<|silent|>'. Respond in Chinese." + ), + help="System prompt for the AURA stage.", + ) + parser.add_argument("--tts-task-type", default="Base", choices=["Base", "CustomVoice"], help="Qwen3-TTS task.") + parser.add_argument("--tts-language", default="Chinese", help="Qwen3-TTS language.") + parser.add_argument("--tts-speaker", default="Vivian", help="Qwen3-TTS speaker.") + parser.add_argument("--tts-instruct", default="", help="Optional Qwen3-TTS style instruction.") + parser.add_argument( + "--tts-ref-audio", + default=default_qwen3_tts_ref_audio_path(), + help="Base-mode reference audio path/URL.", + ) + parser.add_argument( + "--tts-ref-text", + default=DEFAULT_QWEN3_TTS_REF_TEXT, + help="Base-mode reference audio transcript.", + ) + parser.add_argument( + "--tts-x-vector-only-mode", + action="store_true", + help="Use speaker embedding only for Base mode (disable ICL ref_text conditioning).", + ) + parser.add_argument( + "--tts-pass-token-ids", + action="store_true", + help="Pass AURA-generated assistant token ids directly to Qwen3-TTS. Defaults to sending text.", + ) + parser.add_argument("--init-timeout", type=int, default=2000, help="Overall initialization timeout.") + parser.add_argument("--stage-init-timeout", type=int, default=2000, help="Per-stage initialization timeout.") + parser.add_argument("--log-stats", action="store_true", default=False, help="Enable stats logging.") + nullify_stage_engine_defaults(parser) + return parser.parse_args() + + +if __name__ == "__main__": + main(parse_args()) diff --git a/examples/offline_inference/aura_omni/run_single_prompt.sh b/examples/offline_inference/aura_omni/run_single_prompt.sh new file mode 100755 index 000000000..f4db0ac6a --- /dev/null +++ b/examples/offline_inference/aura_omni/run_single_prompt.sh @@ -0,0 +1,8 @@ +#!/usr/bin/env bash +set -euo pipefail + +python end2end.py \ + --model aurateam/AURA \ + --deploy-config vllm_omni/deploy/aura_omni.yaml \ + --modalities text,audio \ + "$@" diff --git a/examples/online_serving/aura_omni/README.md b/examples/online_serving/aura_omni/README.md new file mode 100644 index 000000000..6abd0ef47 --- /dev/null +++ b/examples/online_serving/aura_omni/README.md @@ -0,0 +1,146 @@ +# AURA Omni Native Pipeline + +`aura_omni` serves AURA as a native multi-stage vLLM-Omni pipeline: + +```text +Qwen3-ASR -> AURA/Qwen3-VL -> Qwen3-TTS Talker -> Qwen3-TTS Code2Wav +``` + +The pipeline has three semantic modules, but four engine stages because the +existing Qwen3-TTS implementation is natively split into Talker and Code2Wav. + +Start the server with the deploy profile: + +```bash +vllm serve aurateam/AURA \ + --omni \ + --port 8091 \ + --deploy-config vllm_omni/deploy/aura_omni.yaml \ + --served-model-name aurateam/AURA \ + --trust-remote-code +``` + +The deploy file sets per-stage model repos: + +- Stage 0 ASR: `Qwen/Qwen3-ASR-1.7B` +- Stage 1 AURA: `aurateam/AURA` +- Stage 2/3 TTS: `Qwen/Qwen3-TTS-12Hz-1.7B-Base` + +For local weights, edit the `model` value on each stage in +`vllm_omni/deploy/aura_omni.yaml`. The deploy profile includes +`pipeline: aura_omni`, so the server uses this four-stage topology even when +the command-line model path points at one component checkpoint. + +Expected request shape: + +- Send microphone audio as the Stage 0 multimodal audio input. +- Include video frames in the original request `multi_modal_data`; the + `asr2aura` processor carries them forward to AURA. +- Optional `additional_information` keys: + - `aura_system_prompt` + - `tts_task_type` + - `tts_language` + - `tts_speaker` + - `tts_instruct` + - `tts_ref_audio` + - `tts_ref_text` + - `tts_x_vector_only_mode` + - `tts_pass_token_ids` + +If AURA emits `<|silent|>`, the `aura2tts` processor returns no TTS request, so +the TTS stages are skipped for that turn. + +## GPU Utilization Recommendation + +Tune `gpu_memory_utilization` per stage in `vllm_omni/deploy/aura_omni.yaml`. +Recommended baseline on one GPU for H200 + +- Stage 0 (ASR): `0.10` +- Stage 1 (AURA): `0.4` +- Stage 2 (Qwen3-TTS Talker): `0.20` +- Stage 3 (Qwen3-TTS Code2Wav): `0.20` + +## Python Client + +```bash +python examples/online_serving/aura_omni/openai_chat_completion_client.py \ + --host localhost \ + --port 8091 \ + --model aurateam/AURA \ + --modalities text,audio +``` + +Use local media: + +```bash +python examples/online_serving/aura_omni/openai_chat_completion_client.py \ + --audio-path /path/to/input.wav \ + --video-path /path/to/video.mp4 \ + --output-dir output_aura_omni_online +``` + +Base voice clone mode (default, recommended as x-vector while debugging ICL): + +```bash +python examples/online_serving/aura_omni/openai_chat_completion_client.py \ + --tts-task-type Base \ + --tts-ref-audio vllm-omni/tests/assets/qwen3_tts/clone_2.wav \ + --tts-ref-text "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you." +``` + +Enable AURA token-id passthrough explicitly: + +```bash +python examples/online_serving/aura_omni/openai_chat_completion_client.py \ + --tts-pass-token-ids +``` + +CustomVoice mode requires stages 2 and 3 in `aura_omni.yaml` to point at a +Qwen3-TTS CustomVoice checkpoint: + +```bash +python examples/online_serving/aura_omni/openai_chat_completion_client.py \ + --tts-task-type CustomVoice \ + --tts-speaker Vivian +``` + +By default, AURA responses are passed to Qwen3-TTS as text. Set +`tts_pass_token_ids=true` to pass AURA-generated assistant token ids directly +to Qwen3-TTS instead. The processor still uses AURA token ids, when available, +to estimate the Talker prompt length in the default text path. + +## Curl + +```bash +cd examples/online_serving/aura_omni +bash run_curl_multimodal_generation.sh +``` + +Set `PORT`, `MODEL`, or `OUTPUT_DIR` to override defaults: + +```bash +PORT=8666 MODEL=aurateam/AURA bash run_curl_multimodal_generation.sh +TTS_PASS_TOKEN_IDS=true PORT=8666 MODEL=aurateam/AURA bash run_curl_multimodal_generation.sh +``` + +## Gradio + +Launch the server and Gradio UI together: + +```bash +cd examples/online_serving/aura_omni +bash run_gradio_demo.sh +``` + +If the server is already running: + +```bash +python examples/online_serving/aura_omni/gradio_demo.py \ + --model aurateam/AURA \ + --api-base http://localhost:8091/v1 +``` + +## Offline + +For offline inference, see +[`examples/offline_inference/aura_omni`](../../offline_inference/aura_omni/). diff --git a/examples/online_serving/aura_omni/gradio_demo.py b/examples/online_serving/aura_omni/gradio_demo.py new file mode 100644 index 000000000..c6badee78 --- /dev/null +++ b/examples/online_serving/aura_omni/gradio_demo.py @@ -0,0 +1,305 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Gradio demo for AURA Omni online serving.""" + +from __future__ import annotations + +import argparse +import base64 +import io +from pathlib import Path +from typing import Any + +try: + import gradio as gr +except ImportError: + raise ImportError("gradio is required to run this demo. Install it with: pip install 'vllm-omni[demo]'") from None + +import numpy as np +import soundfile as sf +from openai import OpenAI + +from vllm_omni.model_executor.stage_input_processors.aura_omni import ( + DEFAULT_QWEN3_TTS_REF_TEXT, + default_qwen3_tts_ref_audio_path, +) + +SEED = 42 +DEFAULT_TTS_REF_AUDIO = default_qwen3_tts_ref_audio_path() + + +def parse_args(): + parser = argparse.ArgumentParser(description="Gradio demo for AURA Omni online inference.") + parser.add_argument("--model", default="aurateam/AURA", help="Model name served by vLLM.") + parser.add_argument("--api-base", default="http://localhost:8091/v1", help="OpenAI-compatible API base.") + parser.add_argument("--ip", default="127.0.0.1", help="Gradio host.") + parser.add_argument("--port", type=int, default=7862, help="Gradio port.") + parser.add_argument("--share", action="store_true", help="Share the Gradio demo publicly.") + return parser.parse_args() + + +def _audio_to_data_url(audio_file: Any | None) -> str | None: + if audio_file is None: + return None + + def _path_to_data_url(path_str: str) -> str | None: + if not path_str: + return None + path = Path(path_str) + if not path.exists(): + return None + with open(path, "rb") as f: + return f"data:audio/wav;base64,{base64.b64encode(f.read()).decode('utf-8')}" + + if isinstance(audio_file, str): + return _path_to_data_url(audio_file) + + # Gradio >=5 may return dict payloads with local temp path. + if isinstance(audio_file, dict): + path_val = audio_file.get("path") or audio_file.get("name") + if isinstance(path_val, str): + return _path_to_data_url(path_val) + return None + + if isinstance(audio_file, tuple): + # Common shape: (sample_rate, np.ndarray) + if len(audio_file) == 2 and isinstance(audio_file[0], (int, float)): + sample_rate, audio_np = audio_file + audio_np = np.asarray(audio_np) + if audio_np.ndim > 1: + audio_np = audio_np[:, 0] + if audio_np.dtype != np.int16: + audio_np = np.clip(audio_np.astype(np.float32), -1.0, 1.0) + buf = io.BytesIO() + sf.write(buf, audio_np, int(sample_rate), format="WAV") + return f"data:audio/wav;base64,{base64.b64encode(buf.getvalue()).decode('utf-8')}" + # Compatibility shape from some gradio builds: (filepath, (sr, np.ndarray)) + if len(audio_file) >= 1 and isinstance(audio_file[0], str): + path_url = _path_to_data_url(audio_file[0]) + if path_url is not None: + return path_url + return None + + +def _video_to_data_url(video_file: str | None) -> str | None: + if not video_file: + return None + path = Path(video_file) + mime = { + ".mp4": "video/mp4", + ".webm": "video/webm", + ".mov": "video/quicktime", + ".avi": "video/x-msvideo", + ".mkv": "video/x-matroska", + }.get(path.suffix.lower(), "video/mp4") + with open(path, "rb") as f: + return f"data:{mime};base64,{base64.b64encode(f.read()).decode('utf-8')}" + + +def _sampling_params_list() -> list[dict[str, Any]]: + return [ + {"temperature": 0.0, "top_p": 1.0, "top_k": -1, "max_tokens": 256, "seed": SEED}, + { + "temperature": 0.5, + "top_p": 1.0, + "top_k": -1, + "max_tokens": 256, + "seed": SEED, + "repetition_penalty": 1.0, + }, + { + "temperature": 0.9, + "top_k": 50, + "max_tokens": 4096, + "seed": SEED, + "detokenize": False, + "repetition_penalty": 1.05, + "stop_token_ids": [2150], + }, + { + "temperature": 0.0, + "top_p": 1.0, + "top_k": -1, + "max_tokens": 65536, + "seed": SEED, + "repetition_penalty": 1.0, + }, + ] + + +def _decode_audio_choice(audio_data: str | None) -> tuple[int, np.ndarray] | None: + if not audio_data: + return None + wav_bytes = base64.b64decode(audio_data) + audio_np, sample_rate = sf.read(io.BytesIO(wav_bytes)) + if audio_np.ndim > 1: + audio_np = audio_np[:, 0] + return int(sample_rate), audio_np.astype(np.float32) + + +def build_interface(client: OpenAI, model: str): + def run( + audio_file, + video_file, + prompt: str, + aura_system_prompt: str, + tts_task_type: str, + tts_language: str, + tts_speaker: str, + tts_instruct: str, + tts_ref_audio: str, + tts_ref_text: str, + tts_x_vector_only_mode: bool, + tts_pass_token_ids: bool, + ): + audio_url = _audio_to_data_url(audio_file) + video_url = _video_to_data_url(video_file) + if not audio_url: + return "Please provide audio input for the ASR stage.", None + if tts_task_type == "Base" and (not (tts_ref_audio or "").strip() or not (tts_ref_text or "").strip()): + return "Base TTS requires both reference audio and reference transcript.", None + content: list[dict[str, Any]] = [{"type": "audio_url", "audio_url": {"url": audio_url}}] + if video_url: + content.append({"type": "video_url", "video_url": {"url": video_url}}) + content.append( + {"type": "text", "text": prompt or "Use the audio and video together to decide whether a reply is needed."} + ) + + additional_information: dict[str, Any] = { + "aura_system_prompt": aura_system_prompt, + "tts_task_type": tts_task_type, + "tts_instruct": tts_instruct, + "tts_pass_token_ids": bool(tts_pass_token_ids), + } + if tts_task_type == "CustomVoice": + additional_information.update( + { + "tts_language": tts_language or "English", + "tts_speaker": tts_speaker or "Vivian", + } + ) + else: + additional_information.update( + { + "tts_ref_audio": tts_ref_audio.strip() if tts_ref_audio else None, + "tts_ref_text": tts_ref_text.strip() if tts_ref_text else None, + "tts_x_vector_only_mode": bool(tts_x_vector_only_mode), + } + ) + + response = client.chat.completions.create( + model=model, + messages=[{"role": "user", "content": content}], + modalities=["text", "audio"], + extra_body={ + "sampling_params_list": _sampling_params_list(), + "additional_information": additional_information, + }, + timeout=600.0, + ) + + text_parts: list[str] = [] + audio_output = None + for choice in response.choices: + message = choice.message + if message.content: + text_parts.append(str(message.content)) + if getattr(message, "audio", None): + audio_output = _decode_audio_choice(message.audio.data) + return "\n\n".join(text_parts) or "No text output.", audio_output + + default_system = ( + "You are receiving a live video stream where the final frame is the present moment. " + "Respond only when a response is needed. Otherwise output '<|silent|>'. Respond in English." + ) + + def _toggle_tts_advanced(task_type: str): + is_base = task_type == "Base" + return ( + gr.update(visible=not is_base), + gr.update(visible=is_base), + ) + + with gr.Blocks() as demo: + gr.Markdown("# AURA Omni") + gr.Markdown("ASR -> AURA/Qwen3-VL -> Qwen3-TTS Talker -> Code2Wav") + with gr.Row(): + audio_input = gr.Audio(label="Audio input", sources=["upload", "microphone"], type="numpy") + video_input = gr.Video(label="Video input (optional)", sources=["upload"]) + prompt = gr.Textbox( + label="Instruction", + value="Use the audio and video together to decide whether a reply is needed. If needed, respond briefly in English.", + lines=2, + ) + with gr.Accordion("Advanced", open=False): + aura_system_prompt = gr.Textbox(label="AURA system prompt", value=default_system, lines=4) + tts_task_type = gr.Radio(choices=["Base", "CustomVoice"], value="Base", label="TTS task type") + tts_instruct = gr.Textbox(label="TTS instruction", value="") + tts_pass_token_ids = gr.Checkbox( + label="Pass AURA token ids directly to TTS", + value=False, + ) + + with gr.Group(visible=False) as customvoice_config_group: + tts_language = gr.Dropdown( + choices=["Chinese", "English", "Japanese", "Korean", "Cantonese"], + value="Chinese", + label="CustomVoice language", + ) + tts_speaker = gr.Textbox(label="CustomVoice speaker", value="Vivian") + + with gr.Group(visible=True) as base_config_group: + tts_ref_audio = gr.Textbox( + label="Base reference audio path/URL", + value=DEFAULT_TTS_REF_AUDIO, + ) + tts_ref_text = gr.Textbox( + label="Base reference transcript", + value=DEFAULT_QWEN3_TTS_REF_TEXT, + lines=3, + ) + tts_x_vector_only_mode = gr.Checkbox( + label="Base x-vector only mode (disable ICL)", + value=False, + ) + + tts_task_type.change( + _toggle_tts_advanced, + inputs=[tts_task_type], + outputs=[customvoice_config_group, base_config_group], + ) + button = gr.Button("Generate", variant="primary") + with gr.Row(): + text_output = gr.Textbox(label="Text output", lines=8) + audio_output = gr.Audio(label="Audio output", interactive=False) + button.click( + run, + inputs=[ + audio_input, + video_input, + prompt, + aura_system_prompt, + tts_task_type, + tts_language, + tts_speaker, + tts_instruct, + tts_ref_audio, + tts_ref_text, + tts_x_vector_only_mode, + tts_pass_token_ids, + ], + outputs=[text_output, audio_output], + ) + demo.queue() + return demo + + +def main() -> None: + args = parse_args() + client = OpenAI(api_key="EMPTY", base_url=args.api_base) + demo = build_interface(client, args.model) + demo.launch(server_name=args.ip, server_port=args.port, share=args.share) + + +if __name__ == "__main__": + main() diff --git a/examples/online_serving/aura_omni/openai_chat_completion_client.py b/examples/online_serving/aura_omni/openai_chat_completion_client.py new file mode 100644 index 000000000..c8225bf0d --- /dev/null +++ b/examples/online_serving/aura_omni/openai_chat_completion_client.py @@ -0,0 +1,199 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""OpenAI-compatible client for the AURA Omni pipeline.""" + +from __future__ import annotations + +import base64 +import io +import os + +import soundfile as sf +from openai import OpenAI +from vllm.assets.audio import AudioAsset +from vllm.utils.argparse_utils import FlexibleArgumentParser + +from vllm_omni.model_executor.stage_input_processors.aura_omni import ( + DEFAULT_QWEN3_TTS_REF_TEXT, + default_qwen3_tts_ref_audio_path, +) + +SEED = 42 +DEFAULT_MODEL = "aurateam/AURA" +DEFAULT_VIDEO_URL = "https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4" + + +def _encode_file(path: str) -> str: + with open(path, "rb") as f: + return base64.b64encode(f.read()).decode("utf-8") + + +def _data_url(path: str, default_mime: str) -> str: + suffix = os.path.splitext(path)[1].lower() + mime_by_suffix = { + ".wav": "audio/wav", + ".mp3": "audio/mpeg", + ".ogg": "audio/ogg", + ".flac": "audio/flac", + ".m4a": "audio/mp4", + ".mp4": "video/mp4", + ".webm": "video/webm", + ".mov": "video/quicktime", + ".avi": "video/x-msvideo", + ".mkv": "video/x-matroska", + } + return f"data:{mime_by_suffix.get(suffix, default_mime)};base64,{_encode_file(path)}" + + +def media_url(path_or_url: str | None, *, kind: str) -> str: + if path_or_url: + if path_or_url.startswith(("http://", "https://", "data:")): + return path_or_url + if not os.path.exists(path_or_url): + raise FileNotFoundError(f"{kind} file not found: {path_or_url}") + return _data_url(path_or_url, "audio/wav" if kind == "audio" else "video/mp4") + if kind == "audio": + return AudioAsset("mary_had_lamb").url + return DEFAULT_VIDEO_URL + + +def sampling_params_list() -> list[dict]: + return [ + {"temperature": 0.0, "top_p": 1.0, "top_k": -1, "max_tokens": 256, "seed": SEED}, + { + "temperature": 0.5, + "top_p": 1.0, + "top_k": -1, + "max_tokens": 256, + "seed": SEED, + "repetition_penalty": 1.0, + }, + { + "temperature": 0.9, + "top_k": 50, + "max_tokens": 4096, + "seed": SEED, + "detokenize": False, + "repetition_penalty": 1.05, + "stop_token_ids": [2150], + }, + { + "temperature": 0.0, + "top_p": 1.0, + "top_k": -1, + "max_tokens": 65536, + "seed": SEED, + "repetition_penalty": 1.0, + }, + ] + + +def parse_modalities(value: str | None) -> list[str] | None: + if not value: + return None + return [item.strip() for item in value.split(",") if item.strip()] + + +def save_response(response, output_dir: str) -> None: + os.makedirs(output_dir, exist_ok=True) + for idx, choice in enumerate(response.choices): + message = choice.message + if message.content: + out_txt = os.path.join(output_dir, f"choice_{idx}.txt") + with open(out_txt, "w", encoding="utf-8") as f: + f.write(str(message.content).strip() + "\n") + print(f"Text saved to {out_txt}") + print(message.content) + if getattr(message, "audio", None): + audio_bytes = base64.b64decode(message.audio.data) + audio_np, sample_rate = sf.read(io.BytesIO(audio_bytes)) + out_wav = os.path.join(output_dir, f"choice_{idx}.wav") + sf.write(out_wav, audio_np, int(sample_rate), format="WAV") + print(f"Audio saved to {out_wav}") + + +def main(args) -> None: + client = OpenAI(base_url=f"http://{args.host}:{args.port}/v1", api_key="EMPTY") + messages = [ + { + "role": "user", + "content": [ + {"type": "audio_url", "audio_url": {"url": media_url(args.audio_path, kind="audio")}}, + {"type": "video_url", "video_url": {"url": media_url(args.video_path, kind="video")}}, + {"type": "text", "text": args.prompt}, + ], + } + ] + response = client.chat.completions.create( + model=args.model, + messages=messages, + modalities=parse_modalities(args.modalities), + extra_body={ + "sampling_params_list": sampling_params_list(), + "additional_information": { + "aura_system_prompt": args.aura_system_prompt, + "tts_task_type": args.tts_task_type, + "tts_language": args.tts_language, + "tts_speaker": args.tts_speaker, + "tts_instruct": args.tts_instruct, + "tts_ref_audio": args.tts_ref_audio, + "tts_ref_text": args.tts_ref_text, + "tts_x_vector_only_mode": args.tts_x_vector_only_mode, + "tts_pass_token_ids": args.tts_pass_token_ids, + }, + }, + timeout=args.timeout, + ) + save_response(response, args.output_dir) + + +def parse_args(): + parser = FlexibleArgumentParser(description="AURA Omni online serving client") + parser.add_argument("--host", default="localhost") + parser.add_argument("--port", type=int, default=8091) + parser.add_argument("--model", default=DEFAULT_MODEL) + parser.add_argument("--audio-path", default=None, help="Audio file, URL, or data URL.") + parser.add_argument("--video-path", default=None, help="Video file, URL, or data URL.") + parser.add_argument( + "--prompt", + default="Use the audio and video together to decide whether a reply is needed. If needed, respond briefly in English.", + ) + parser.add_argument("--modalities", default="text,audio") + parser.add_argument("--output-dir", default="output_aura_omni_online") + parser.add_argument( + "--aura-system-prompt", + default=( + "You are receiving a live video stream where the final frame is the present moment. " + "Respond only when a response is needed. Otherwise output '<|silent|>'. Respond in English." + ), + ) + parser.add_argument("--tts-task-type", default="Base", choices=["Base", "CustomVoice"]) + parser.add_argument("--tts-language", default="English") + parser.add_argument("--tts-speaker", default="Vivian") + parser.add_argument("--tts-instruct", default="") + parser.add_argument( + "--tts-ref-audio", + default=default_qwen3_tts_ref_audio_path(), + help="Base-mode reference audio path/URL visible to server.", + ) + parser.add_argument( + "--tts-ref-text", + default=DEFAULT_QWEN3_TTS_REF_TEXT, + help="Base-mode reference audio transcript.", + ) + parser.add_argument( + "--tts-x-vector-only-mode", + action="store_true", + help="Use speaker embedding only for Base mode (disable ICL ref_text conditioning).", + ) + parser.add_argument( + "--tts-pass-token-ids", + action="store_true", + help="Pass AURA-generated assistant token ids directly to Qwen3-TTS. Defaults to sending text.", + ) + parser.add_argument("--timeout", type=float, default=600.0) + return parser.parse_args() + + +if __name__ == "__main__": + main(parse_args()) diff --git a/examples/online_serving/aura_omni/run_curl_multimodal_generation.sh b/examples/online_serving/aura_omni/run_curl_multimodal_generation.sh new file mode 100755 index 000000000..751078109 --- /dev/null +++ b/examples/online_serving/aura_omni/run_curl_multimodal_generation.sh @@ -0,0 +1,60 @@ +#!/usr/bin/env bash +set -euo pipefail + +PORT="${PORT:-8091}" +MODEL="${MODEL:-aurateam/AURA}" +OUTPUT_DIR="${OUTPUT_DIR:-output_aura_omni_online}" +TTS_PASS_TOKEN_IDS="${TTS_PASS_TOKEN_IDS:-false}" +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +VLLM_OMNI_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)" +CLONE_REF_AUDIO="${VLLM_OMNI_ROOT}/tests/assets/qwen3_tts/clone_2.wav" +CLONE_REF_TEXT="Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you." +mkdir -p "$OUTPUT_DIR" + +MARY_HAD_LAMB_AUDIO_URL="https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/mary_had_lamb.ogg" +SAMPLE_VIDEO_URL="https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4" + +request_body=$(cat <'. Respond in English.", + "tts_task_type": "Base", + "tts_ref_audio": "file://${CLONE_REF_AUDIO}", + "tts_ref_text": "${CLONE_REF_TEXT}", + "tts_language": "English", + "tts_speaker": "Vivian", + "tts_instruct": "", + "tts_pass_token_ids": ${TTS_PASS_TOKEN_IDS} + }, + "messages": [{ + "role": "user", + "content": [ + {"type": "audio_url", "audio_url": {"url": "$MARY_HAD_LAMB_AUDIO_URL"}}, + {"type": "video_url", "video_url": {"url": "$SAMPLE_VIDEO_URL"}}, + {"type": "text", "text": "Use the audio and video together to decide whether a reply is needed. If needed, respond briefly in English."} + ] + }] +} +EOF +) + +response=$(curl -sS --retry 3 --retry-delay 3 --retry-connrefused \ + -X POST "http://localhost:${PORT}/v1/chat/completions" \ + -H "Content-Type: application/json" \ + -d "$request_body") + +echo "$response" | jq '.choices[].message.content' + +audio_b64=$(echo "$response" | jq -r '.choices[]?.message.audio.data // empty' | head -n 1) +if [[ -n "$audio_b64" ]]; then + echo "$audio_b64" | base64 -d > "${OUTPUT_DIR}/aura_omni_output.wav" + echo "Audio saved to ${OUTPUT_DIR}/aura_omni_output.wav" +fi diff --git a/examples/online_serving/aura_omni/run_gradio_demo.sh b/examples/online_serving/aura_omni/run_gradio_demo.sh new file mode 100755 index 000000000..bf591c843 --- /dev/null +++ b/examples/online_serving/aura_omni/run_gradio_demo.sh @@ -0,0 +1,79 @@ +#!/usr/bin/env bash +set -euo pipefail + +MODEL="aurateam/AURA" +SERVER_MODEL="aurateam/AURA" +DEPLOY_CONFIG="/data/yrr/vllm-omni/vllm_omni/deploy/aura_omni.yaml" +SERVER_PORT=8091 +GRADIO_PORT=7862 +SERVER_HOST="0.0.0.0" +GRADIO_IP="127.0.0.1" +GRADIO_SHARE=false + + +while [[ $# -gt 0 ]]; do + case "$1" in + --model) MODEL="$2"; shift 2 ;; + --server-model) SERVER_MODEL="$2"; shift 2 ;; + --deploy-config) DEPLOY_CONFIG="$2"; shift 2 ;; + --server-port) SERVER_PORT="$2"; shift 2 ;; + --gradio-port) GRADIO_PORT="$2"; shift 2 ;; + --server-host) SERVER_HOST="$2"; shift 2 ;; + --gradio-ip) GRADIO_IP="$2"; shift 2 ;; + --share) GRADIO_SHARE=true; shift ;; + --help) + echo "Usage: $0 [--model SERVED_MODEL_NAME] [--server-model MODEL_PATH] [--deploy-config YAML] [--server-port PORT] [--gradio-port PORT] [--share]" + exit 0 + ;; + *) echo "Unknown option: $1"; exit 1 ;; + esac +done + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +API_BASE="http://localhost:${SERVER_PORT}/v1" +LOG_FILE="/tmp/aura_omni_vllm_${SERVER_PORT}.log" + +cleanup() { + echo "Shutting down..." + [[ -n "${SERVER_PID:-}" ]] && kill "$SERVER_PID" 2>/dev/null || true + [[ -n "${GRADIO_PID:-}" ]] && kill "$GRADIO_PID" 2>/dev/null || true +} +trap cleanup SIGINT SIGTERM EXIT + +vllm serve "$SERVER_MODEL" \ + --omni \ + --host "$SERVER_HOST" \ + --port "$SERVER_PORT" \ + --deploy-config "$DEPLOY_CONFIG" \ + --served-model-name "$MODEL" \ + --trust-remote-code 2>&1 | tee "$LOG_FILE" & +SERVER_PID=$! + +echo "Waiting for server startup..." +for _ in $(seq 1 600); do + if grep -q "Application startup complete" "$LOG_FILE" 2>/dev/null; then + break + fi + if ! kill -0 "$SERVER_PID" 2>/dev/null; then + echo "vLLM server exited before startup completed" + wait "$SERVER_PID" || true + exit 1 + fi + sleep 1 +done + +# cd "$SCRIPT_DIR" +GRADIO_CMD=(python gradio_demo.py --model "$MODEL" --api-base "$API_BASE" --ip "$GRADIO_IP" --port "$GRADIO_PORT") +if [[ "$GRADIO_SHARE" == "true" ]]; then + GRADIO_CMD+=(--share) +fi +"${GRADIO_CMD[@]}" & +GRADIO_PID=$! + +echo "vLLM server: http://${SERVER_HOST}:${SERVER_PORT}" +echo "Gradio demo: http://${GRADIO_IP}:${GRADIO_PORT}" +if [[ -n "${SERVER_PID:-}" ]]; then + wait "$SERVER_PID" "$GRADIO_PID" +else + wait "$GRADIO_PID" +fi diff --git a/tests/config/test_pipeline_registry.py b/tests/config/test_pipeline_registry.py index 336f38054..ecd739857 100644 --- a/tests/config/test_pipeline_registry.py +++ b/tests/config/test_pipeline_registry.py @@ -30,6 +30,7 @@ class TestCentralRegistryDeclarations: assert "qwen2_5_omni" in _OMNI_PIPELINES assert "qwen2_5_omni_thinker_only" in _OMNI_PIPELINES assert "qwen3_omni_moe" in _OMNI_PIPELINES + assert "aura_omni" in _OMNI_PIPELINES assert "qwen3_tts" in _OMNI_PIPELINES @@ -45,6 +46,20 @@ class TestLazyLoading: assert pipeline.model_type == "qwen3_omni_moe" assert pipeline.model_arch == "Qwen3OmniMoeForConditionalGeneration" + def test_getitem_loads_aura_omni_pipeline(self): + pipeline = _PIPELINE_REGISTRY["aura_omni"] + assert pipeline.model_type == "aura_omni" + assert [stage.model_stage for stage in pipeline.stages] == [ + "asr", + "aura", + "qwen3_tts", + "code2wav", + ] + assert pipeline.stages[0].model_arch == "Qwen3ASRForConditionalGeneration" + assert pipeline.stages[1].model_arch == "AuraQwen3VLForConditionalGeneration" + assert pipeline.stages[2].model_arch == "Qwen3TTSTalkerForConditionalGeneration" + assert pipeline.stages[3].model_arch == "Qwen3TTSCode2Wav" + def test_unknown_model_type_returns_none_via_get(self): assert _PIPELINE_REGISTRY.get("not_a_real_pipeline") is None @@ -73,6 +88,7 @@ class TestLazyLoading: "attr_fail": ("vllm_omni.config.stage_config", "NON_EXISTENT_VAR"), }, ) + monkeypatch.setattr(_PIPELINE_REGISTRY, "_loaded", {}) # Assert that list(_PIPELINE_REGISTRY.values()) and dict(_PIPELINE_REGISTRY.items()) # both yield only healthy entries. diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index 69440ac3b..31850401b 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/tests/engine/test_async_omni_engine_stage_init.py @@ -591,6 +591,24 @@ def test_build_engine_args_cli_tokenizer_overrides_inferred_base_tokenizer(tmp_p assert engine_args["tokenizer"] == "/external/tokenizer" +def test_build_engine_args_stage_model_overrides_parent_model(): + from vllm_omni.engine.stage_init_utils import build_engine_args_dict + + stage_cfg = types.SimpleNamespace( + stage_id=0, + stage_type="llm", + engine_args={"model": "/stage/model"}, + default_sampling_params={}, + ) + + engine_args = build_engine_args_dict( + stage_cfg, + "/parent/model", + ) + + assert engine_args["model"] == "/stage/model" + + def test_build_engine_args_keeps_stage_owned_tokenizer_subdir(tmp_path): from vllm_omni.engine.stage_init_utils import build_engine_args_dict diff --git a/tests/engine/test_orchestrator.py b/tests/engine/test_orchestrator.py index 029c55fd1..f73504074 100644 --- a/tests/engine/test_orchestrator.py +++ b/tests/engine/test_orchestrator.py @@ -24,7 +24,12 @@ from vllm_omni.engine.messages import ( ShutdownRequestMessage, StageSubmissionMessage, ) -from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState +from vllm_omni.engine.orchestrator import ( + Orchestrator, + OrchestratorRequestState, + _build_terminal_empty_output, + _infer_stage_audio_sample_rate, +) from vllm_omni.engine.stage_pool import StagePool from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput @@ -135,6 +140,20 @@ class FakeStageClient: self._diffusion_outputs.put_nowait(output) +def test_terminal_empty_audio_output_uses_stage_sample_rate() -> None: + final_stage = FakeStageClient(final_output=True, final_output_type="audio") + final_stage.sample_rate = 44100 + final_pool = SimpleNamespace(stage_client=final_stage, _stage_vllm_config=None) + + terminal_output = _build_terminal_empty_output( + "req-1", + final_output_type="audio", + audio_sample_rate=_infer_stage_audio_sample_rate(final_pool), + ) + + assert terminal_output.outputs[0].multimodal_output["sr"] == 44100 + + class FakeCollectiveRpcStageClient(FakeStageClient): def __init__(self, *args, rpc_result: Any = None, **kwargs) -> None: super().__init__(*args, **kwargs) diff --git a/tests/engine/test_orchestrator_stage_input_bridge.py b/tests/engine/test_orchestrator_stage_input_bridge.py new file mode 100644 index 000000000..27ee11df4 --- /dev/null +++ b/tests/engine/test_orchestrator_stage_input_bridge.py @@ -0,0 +1,156 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import queue +from types import SimpleNamespace +from typing import Any + +import janus +import pytest +from vllm.outputs import CompletionOutput, RequestOutput +from vllm.sampling_params import SamplingParams + +from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState +from vllm_omni.engine.stage_pool import StagePool + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class FakeStageClient: + def __init__( + self, + *, + next_inputs: list[dict[str, Any]] | None = None, + final_output: bool = False, + ) -> None: + self.stage_id = 0 + self.replica_id = 0 + self.stage_type = "llm" + self.final_output = final_output + self.final_output_type = "text" + self.default_sampling_params = SamplingParams(max_tokens=1) + self.requires_multimodal_data = False + self.engine_input_source = [0] + self.is_comprehension = False + self.model_stage = None + self.custom_process_input_func = None + self.next_inputs = list(next_inputs or []) + self.add_request_calls: list[tuple[Any, ...]] = [] + self._engine_core_outputs = queue.Queue() + + async def add_request_async(self, *args, **_kwargs) -> None: + self.add_request_calls.append(args) + + async def get_output_async(self): + try: + return self._engine_core_outputs.get_nowait() + except queue.Empty: + return SimpleNamespace(outputs=[]) + + def process_engine_inputs(self, _source_outputs, prompt=None, streaming_context=None): + return list(self.next_inputs) + + async def abort_requests_async(self, _request_ids: list[str]) -> None: + return None + + def set_engine_outputs(self, _outputs) -> None: + return None + + def check_health(self) -> None: + return None + + def shutdown(self) -> None: + return None + + +class FakeOutputProcessor: + def add_request(self, *args, **kwargs) -> None: + return None + + +class FakeInputProcessor: + def __init__(self) -> None: + self.calls: list[dict[str, Any]] = [] + + def process_inputs(self, **kwargs): + self.calls.append(kwargs) + return SimpleNamespace( + request_id=kwargs["request_id"], + prompt_token_ids=[101, 102], + prompt_embeds=None, + external_req_id=None, + ) + + +def _request_output(request_id: str) -> RequestOutput: + completion = CompletionOutput( + index=0, + text="transcript", + token_ids=[11, 12], + cumulative_logprob=None, + logprobs=None, + finish_reason="stop", + stop_reason=None, + ) + return RequestOutput( + request_id=request_id, + prompt="prompt", + prompt_token_ids=[1, 2], + prompt_logprobs=None, + outputs=[completion], + finished=True, + metrics=None, + lora_request=None, + ) + + +@pytest.mark.asyncio +async def test_forward_text_prompt_uses_target_stage_input_processor() -> None: + stage0 = FakeStageClient(final_output=True) + stage1 = FakeStageClient( + final_output=True, + next_inputs=[{"prompt": "hello", "multi_modal_data": {"video": ["frame"]}}], + ) + stage_pools = [ + StagePool( + 0, + [stage0], + output_processor=FakeOutputProcessor(), + stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)), + ), + StagePool( + 1, + [stage1], + output_processor=FakeOutputProcessor(), + stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)), + ), + ] + request_q = janus.Queue() + output_q = janus.Queue() + rpc_q = janus.Queue() + orchestrator = Orchestrator( + request_async_queue=request_q.async_q, + output_async_queue=output_q.async_q, + rpc_async_queue=rpc_q.async_q, + stage_pools=stage_pools, + async_chunk=False, + ) + input_processor = FakeInputProcessor() + orchestrator._stage_input_processors[1] = input_processor + req_state = OrchestratorRequestState( + request_id="req-text", + prompt={"prompt": "original"}, + sampling_params_list=[SamplingParams(max_tokens=1), SamplingParams(max_tokens=1)], + final_stage_id=1, + ) + + await orchestrator._forward_to_next_stage("req-text", 0, _request_output("req-text"), req_state) + + assert input_processor.calls + assert input_processor.calls[0]["prompt"] == {"prompt": "hello", "multi_modal_data": {"video": ["frame"]}} + assert stage1.add_request_calls + submitted_request = stage1.add_request_calls[0][0] + assert submitted_request.prompt_token_ids == [101, 102] + assert submitted_request.external_req_id == "req-text" diff --git a/tests/entrypoints/openai_api/test_serving_chat_multistage_generation.py b/tests/entrypoints/openai_api/test_serving_chat_multistage_generation.py index 507018ef7..0dea97179 100644 --- a/tests/entrypoints/openai_api/test_serving_chat_multistage_generation.py +++ b/tests/entrypoints/openai_api/test_serving_chat_multistage_generation.py @@ -1,8 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 -"""Regression tests for multistage diffusion generation input construction.""" +"""Regression tests for multistage generation input construction.""" from __future__ import annotations +import asyncio from types import SimpleNamespace import pytest @@ -93,6 +94,79 @@ def test_build_multistage_generation_inputs_applies_stage_specific_overrides(ser assert engine.default_sampling_params_list[2].lora_request is None +def test_prepare_multistage_multimodal_inputs_defers_downstream_modalities(serving_chat): + from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat + + image_payload = object() + serving_chat.engine_client = SimpleNamespace( + stage_configs=[ + SimpleNamespace(model_stage="asr", requires_multimodal_data=True), + SimpleNamespace(model_stage="aura", requires_multimodal_data=True), + ] + ) + serving_chat.model_config = SimpleNamespace( + allowed_local_media_path="", + allowed_media_domains=None, + ) + request = SimpleNamespace(media_io_kwargs=None) + messages = [ + { + "role": "user", + "content": [ + {"type": "audio_url", "audio_url": {"url": "data:audio/wav;base64,abc"}}, + {"type": "image_pil", "image_pil": image_payload}, + {"type": "text", "text": "What changed?"}, + ], + } + ] + + stripped_messages, deferred_data = asyncio.run( + OmniOpenAIServingChat._prepare_multistage_multimodal_inputs( + serving_chat, + messages, + request, + ) + ) + + assert deferred_data == {"image": [image_payload]} + stripped_content = stripped_messages[0]["content"] + assert {"type": "image_pil", "image_pil": image_payload} not in stripped_content + assert {"type": "text", "text": "What changed?"} in stripped_content + + +def test_needs_multistage_multimodal_split_noops_for_single_stage(serving_chat): + serving_chat.engine_client = SimpleNamespace( + stage_configs=[ + SimpleNamespace(model_stage="thinker", requires_multimodal_data=True), + ] + ) + + assert serving_chat._needs_multistage_multimodal_split() is False + + +def test_needs_multistage_multimodal_split_noops_for_qwen3_omni(serving_chat): + serving_chat.engine_client = SimpleNamespace( + stage_configs=[ + SimpleNamespace(model_stage="thinker", requires_multimodal_data=True), + SimpleNamespace(model_stage="talker", requires_multimodal_data=False), + SimpleNamespace(model_stage="code2wav", requires_multimodal_data=False), + ] + ) + + assert serving_chat._needs_multistage_multimodal_split() is False + + +def test_needs_multistage_multimodal_split_detects_aura_handoff(serving_chat): + serving_chat.engine_client = SimpleNamespace( + stage_configs=[ + SimpleNamespace(model_stage="asr", requires_multimodal_data=True), + SimpleNamespace(model_stage="aura", requires_multimodal_data=True), + ] + ) + + assert serving_chat._needs_multistage_multimodal_split() is True + + def test_build_multistage_generation_inputs_multi_image_emits_n_img_placeholders(serving_chat): """N reference images with bot_task set must emit N placeholders. diff --git a/tests/entrypoints/openai_api/test_serving_chat_sampling_params.py b/tests/entrypoints/openai_api/test_serving_chat_sampling_params.py index c885d907c..a6c7baffe 100644 --- a/tests/entrypoints/openai_api/test_serving_chat_sampling_params.py +++ b/tests/entrypoints/openai_api/test_serving_chat_sampling_params.py @@ -6,6 +6,8 @@ Tests that standard OpenAI API parameters (max_tokens, temperature, etc.) are correctly applied to the comprehension stage while preserving YAML defaults. """ +from types import SimpleNamespace + import pytest from pytest_mock import MockerFixture from vllm.sampling_params import SamplingParams @@ -487,6 +489,35 @@ def test_build_sampling_params_list_empty_stop_preserves_yaml(serving_chat, mock assert comprehension_params.stop_token_ids == [] +def test_to_sampling_params_list_pads_missing_tail_stage_with_defaults(): + """AURA callers may pass 3 semantic model params for a 4-stage engine pipeline.""" + from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat + + instance = object.__new__(OmniOpenAIServingChat) + default_params = [ + SamplingParams(max_tokens=10), + SamplingParams(max_tokens=20), + SamplingParams(max_tokens=30), + SamplingParams(max_tokens=40), + ] + instance.engine_client = SimpleNamespace( + stage_configs=[SimpleNamespace(stage_type="llm") for _ in range(4)], + default_sampling_params_list=default_params, + ) + + result = instance._to_sampling_params_list( + [ + {"max_tokens": 1}, + {"max_tokens": 2}, + {"max_tokens": 3}, + ] + ) + + assert len(result) == 4 + assert [params.max_tokens for params in result] == [1, 2, 3, 40] + assert result[3] is not default_params[3] + + # ============================================================================= # Tests for _get_comprehension_stage_index # ============================================================================= diff --git a/tests/model_executor/models/test_aura_qwen3_vl.py b/tests/model_executor/models/test_aura_qwen3_vl.py new file mode 100644 index 000000000..9035fcbf0 --- /dev/null +++ b/tests/model_executor/models/test_aura_qwen3_vl.py @@ -0,0 +1,69 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest + +from vllm_omni.model_executor.models.aura_omni.qwen3_vl import ( + AuraQwen3VLForConditionalGeneration, + AuraQwen3VLProcessingInfo, + Qwen3VLProcessor, +) +from vllm_omni.model_executor.models.registry import OmniModelRegistry + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _FakeCtx: + def __init__(self, config): + self.config = config + self.processor_type = None + + def get_hf_config(self, typ=None): + if typ is not None: + raise TypeError("wrong config class") + return self.config + + def get_hf_processor(self, typ=None, **kwargs): + self.processor_type = typ + return SimpleNamespace(kwargs=kwargs) + + +def test_aura_qwen3_vl_processing_accepts_remote_qwen3_vl_config(): + info = object.__new__(AuraQwen3VLProcessingInfo) + info.ctx = _FakeCtx( + SimpleNamespace( + model_type="qwen3_vl", + vision_config=SimpleNamespace(spatial_merge_size=2), + text_config=SimpleNamespace(hidden_size=4096), + ) + ) + + assert info.get_hf_config().model_type == "qwen3_vl" + + +def test_aura_qwen3_vl_processing_rejects_unrelated_remote_config(): + info = object.__new__(AuraQwen3VLProcessingInfo) + info.ctx = _FakeCtx(SimpleNamespace(model_type="not_qwen3_vl")) + + with pytest.raises(TypeError, match="wrong config class"): + info.get_hf_config() + + +def test_aura_qwen3_vl_model_arch_registered(): + model_cls = OmniModelRegistry._try_load_model_cls("AuraQwen3VLForConditionalGeneration") + + assert model_cls is AuraQwen3VLForConditionalGeneration + + +def test_aura_qwen3_vl_processing_forces_upstream_processor_class(): + info = object.__new__(AuraQwen3VLProcessingInfo) + ctx = _FakeCtx(SimpleNamespace(model_type="qwen3_vl")) + info.ctx = ctx + + processor = info.get_hf_processor(extra="value") + + assert ctx.processor_type is Qwen3VLProcessor + assert processor.kwargs["use_fast"] is True + assert processor.kwargs["extra"] == "value" diff --git a/tests/model_executor/stage_input_processors/test_aura_omni.py b/tests/model_executor/stage_input_processors/test_aura_omni.py new file mode 100644 index 000000000..21f798b3a --- /dev/null +++ b/tests/model_executor/stage_input_processors/test_aura_omni.py @@ -0,0 +1,197 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest + +from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import ( + PRECOMPUTED_TEXT_IDS_KEY, +) +from vllm_omni.model_executor.stage_input_processors.aura_omni import ( + SILENT_TEXT, + asr2aura, + aura2tts, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _source_output(text: str, request_id: str = "req-1", token_ids: list[int] | None = None): + output = SimpleNamespace(text=text, cumulative_token_ids=token_ids or [1, 2, 3], multimodal_output={}) + return SimpleNamespace(request_id=request_id, outputs=[output]) + + +def _source_delta_final_output(cumulative_text: str, request_id: str = "req-1"): + output = SimpleNamespace( + text="", + cumulative_text=cumulative_text, + cumulative_token_ids=[1, 2, 3], + multimodal_output={}, + ) + return SimpleNamespace(request_id=request_id, outputs=[output]) + + +def test_asr2aura_carries_video_payload_and_transcript(): + prompt = { + "multi_modal_data": {"video": ["frame-0", "frame-1"]}, + "additional_information": {"aura_system_prompt": ["system"]}, + } + + [next_input] = asr2aura([_source_output("What is happening now?")], prompt=[prompt]) + + assert next_input["multi_modal_data"] == {"video": ["frame-0", "frame-1"]} + assert "<|video_pad|>" in next_input["prompt"] + assert "What is happening now?" in next_input["prompt"] + assert next_input["prompt"].startswith("<|im_start|>system\nsystem") + + +def test_asr2aura_drops_audio_before_qwen3_vl_stage(): + prompt = { + "multi_modal_data": { + "audio": ("wave", 16000), + "video": ["frame-0", "frame-1"], + }, + } + + [next_input] = asr2aura([_source_output("Check the video")], prompt=[prompt]) + + assert next_input["multi_modal_data"] == {"video": ["frame-0", "frame-1"]} + assert "<|video_pad|>" in next_input["prompt"] + + +def test_asr2aura_reads_video_stashed_for_downstream_stage(): + prompt = { + "multi_modal_data": {"audio": ("wave", 16000)}, + "additional_information": { + "deferred_multi_modal_data": {"video": ["frame-0", "frame-1"]}, + }, + } + + [next_input] = asr2aura([_source_output("Check the video")], prompt=[prompt]) + + assert next_input["multi_modal_data"] == {"video": ["frame-0", "frame-1"]} + assert "<|video_pad|>" in next_input["prompt"] + + +def test_asr2aura_supports_video_only_observation(): + prompt = {"multi_modal_data": {"video": ["frame-0", "frame-1"]}} + + [next_input] = asr2aura([_source_output("")], prompt=[prompt]) + + assert "<|video_pad|>" in next_input["prompt"] + assert "<|im_start|>assistant" in next_input["prompt"] + + +def test_aura2tts_builds_qwen3_tts_prompt_information(): + prompt = { + "additional_information": { + "tts_language": ["Chinese"], + "tts_instruct": ["Calm voice."], + "tts_ref_audio": ["ref.wav"], + "tts_ref_text": ["Reference transcript sample."], + } + } + + [tts_input] = aura2tts([_source_output("Hello.")], prompt=[prompt]) + + assert len(tts_input["prompt_token_ids"]) >= 32 + assert tts_input["additional_information"]["text"] == ["Hello."] + assert PRECOMPUTED_TEXT_IDS_KEY not in tts_input["additional_information"] + assert tts_input["additional_information"]["task_type"] == ["Base"] + assert tts_input["additional_information"]["language"] == ["Chinese"] + assert tts_input["additional_information"]["ref_audio"] == ["ref.wav"] + assert tts_input["additional_information"]["ref_text"] == ["Reference transcript sample."] + assert tts_input["additional_information"]["x_vector_only_mode"] == [False] + assert tts_input["additional_information"]["instruct"] == ["Calm voice."] + + +def test_aura2tts_prefers_streaming_cumulative_text(): + prompt = { + "additional_information": { + "tts_ref_audio": ["ref.wav"], + "tts_ref_text": ["Reference transcript sample."], + } + } + + [tts_input] = aura2tts( + [_source_delta_final_output("The complete AURA reply.")], + prompt=[prompt], + ) + + assert tts_input["additional_information"]["text"] == ["The complete AURA reply."] + + +def test_aura2tts_supports_base_ref_audio_override(): + prompt = { + "additional_information": { + "tts_ref_audio": ["custom.wav"], + "tts_ref_text": ["custom transcript"], + } + } + + [tts_input] = aura2tts([_source_output("Hello.")], prompt=[prompt]) + + assert tts_input["additional_information"]["task_type"] == ["Base"] + assert tts_input["additional_information"]["ref_audio"] == ["custom.wav"] + assert tts_input["additional_information"]["ref_text"] == ["custom transcript"] + assert tts_input["additional_information"]["x_vector_only_mode"] == [False] + + +def test_aura2tts_supports_x_vector_only_mode_for_base(): + prompt = { + "additional_information": { + "tts_task_type": ["Base"], + "tts_x_vector_only_mode": [True], + "tts_ref_audio": ["ref.wav"], + "tts_ref_text": ["Reference transcript sample."], + } + } + + [tts_input] = aura2tts([_source_output("Hello.")], prompt=[prompt]) + + assert tts_input["additional_information"]["x_vector_only_mode"] == [True] + + +def test_aura2tts_supports_custom_voice_mode(): + prompt = { + "additional_information": { + "tts_task_type": ["CustomVoice"], + "tts_speaker": ["vivian"], + } + } + + [tts_input] = aura2tts([_source_output("Hello.")], prompt=[prompt]) + + assert tts_input["additional_information"]["task_type"] == ["CustomVoice"] + assert tts_input["additional_information"]["speaker"] == ["Vivian"] + assert "ref_audio" not in tts_input["additional_information"] + + +def test_aura2tts_passes_token_ids_to_qwen3_tts_when_enabled(): + prompt = { + "additional_information": { + "tts_ref_audio": ["ref.wav"], + "tts_ref_text": ["Reference transcript sample."], + "tts_pass_token_ids": [True], + } + } + + [tts_input] = aura2tts( + [ + _source_output( + "Hello.", + token_ids=[151644, 77091, 198, 108386, 1773, 151645, 198], + ) + ], + prompt=[prompt], + ) + + assert tts_input["additional_information"][PRECOMPUTED_TEXT_IDS_KEY] == [ + [151644, 77091, 198, 108386, 1773, 151645, 198, 151644, 77091, 198] + ] + assert "text" not in tts_input["additional_information"] + + +def test_aura2tts_drops_silent_response(): + assert aura2tts([_source_output(SILENT_TEXT)]) == [] diff --git a/tests/test_config_factory.py b/tests/test_config_factory.py index e3ebea88f..45a4ccfab 100644 --- a/tests/test_config_factory.py +++ b/tests/test_config_factory.py @@ -1384,6 +1384,34 @@ class TestCLIOverrideFlow: assert s.runtime_overrides["enforce_eager"] is True +class TestAuraOmniDeploy: + def test_aura_omni_deploy_forces_pipeline_override(self): + deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "aura_omni.yaml" + deploy = load_deploy_config(deploy_path) + + assert deploy.pipeline == "aura_omni" + + def test_aura_omni_deploy_resolves_four_native_stages(self): + stages = StageConfigFactory._create_from_registry( + "qwen3_tts", + cli_overrides={}, + deploy_config_path=str(Path(__file__).parent.parent / "vllm_omni" / "deploy" / "aura_omni.yaml"), + ) + + assert [stage.model_stage for stage in stages] == [ + "asr", + "aura", + "qwen3_tts", + "code2wav", + ] + assert [stage.final_output for stage in stages] == [False, True, False, True] + assert [stage.final_output_type for stage in stages] == [None, "text", None, "audio"] + assert stages[0].yaml_engine_args["model_arch"] == "Qwen3ASRForConditionalGeneration" + assert stages[1].yaml_engine_args["model_arch"] == "AuraQwen3VLForConditionalGeneration" + assert stages[2].yaml_engine_args["model_arch"] == "Qwen3TTSTalkerForConditionalGeneration" + assert stages[3].yaml_engine_args["model_arch"] == "Qwen3TTSCode2Wav" + + class TestSentinelDefaultPrecedence: """Caller-typed (non-None) values win over YAML; None values fall through to YAML / dataclass defaults (#3035).""" diff --git a/tests/worker/test_omni_gpu_model_runner.py b/tests/worker/test_omni_gpu_model_runner.py index 6f4994b2c..f562d18ce 100644 --- a/tests/worker/test_omni_gpu_model_runner.py +++ b/tests/worker/test_omni_gpu_model_runner.py @@ -4,7 +4,7 @@ from types import SimpleNamespace import pytest import torch -from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner +from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner, _filter_mrope_kwargs_for_model from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -96,12 +96,44 @@ class CaptureTalkerMTP(torch.nn.Module): return req_embeds, codes +class StrictMRoPEModel: + def get_mrope_input_positions(self, input_tokens, mm_features): + raise NotImplementedError + + +class FlexibleMRoPEModel: + def get_mrope_input_positions(self, input_tokens, mm_features=None, **kwargs): + raise NotImplementedError + + @contextmanager def _noop_forward_context(*args, **kwargs): """A no-op context manager to replace vLLM forward context in CPU tests.""" yield +def test_filter_mrope_kwargs_for_strict_model_signature(): + kwargs = { + "mm_features": ["audio"], + "hf_config": object(), + "image_grid_thw": [], + } + + assert _filter_mrope_kwargs_for_model(StrictMRoPEModel(), kwargs) == { + "mm_features": ["audio"], + } + + +def test_filter_mrope_kwargs_preserves_flexible_model_kwargs(): + kwargs = { + "mm_features": ["video"], + "hf_config": object(), + "video_grid_thw": [[1, 2, 3]], + } + + assert _filter_mrope_kwargs_for_model(FlexibleMRoPEModel(), kwargs) is kwargs + + def _make_runner(req_ids=("r1", "r2"), hidden_size=4): # Create an instance without calling OmniGPUModelRunner.__init__ runner = object.__new__(OmniGPUModelRunner) diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index 032bcbc10..f7109f395 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/vllm_omni/config/pipeline_registry.py @@ -44,6 +44,10 @@ _OMNI_PIPELINES: dict[str, tuple[str, str]] = { "vllm_omni.model_executor.models.qwen3_omni.pipeline", "QWEN3_OMNI_PIPELINE", ), + "aura_omni": ( + "vllm_omni.model_executor.models.aura_omni.pipeline", + "AURA_OMNI_PIPELINE", + ), "qwen3_tts": ( "vllm_omni.model_executor.models.qwen3_tts.pipeline", "QWEN3_TTS_PIPELINE", diff --git a/vllm_omni/deploy/aura_omni.yaml b/vllm_omni/deploy/aura_omni.yaml new file mode 100644 index 000000000..55beb53a7 --- /dev/null +++ b/vllm_omni/deploy/aura_omni.yaml @@ -0,0 +1,110 @@ +# AURA Omni deploy profile: Qwen3-ASR -> AURA/Qwen3-VL -> Qwen3-TTS Talker -> Code2Wav. +# +# Pass `--deploy-config vllm_omni/deploy/aura_omni.yaml` and point the parent +# `--model` at any local placeholder repo if needed; each stage overrides its +# own model below. +pipeline: aura_omni +# AURA has a semantic ASR -> AURA -> TTS dependency. Do not prewarm downstream +# stages before AURA text is converted into Qwen3-TTS metadata by aura2tts. +async_chunk: false + +connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + shm_threshold_bytes: 65536 + codec_streaming: true + connector_get_sleep_s: 0.01 + connector_get_max_wait_first_chunk: 3000 + connector_get_max_wait: 300 + codec_chunk_frames: 25 + codec_left_context_frames: 72 + initial_codec_chunk_frames: 1 + decode_cudagraph_capture_sizes: [25, 73, 97, 169, 325] + +stages: + - stage_id: 0 + max_num_seqs: 8 + gpu_memory_utilization: 0.1 + trust_remote_code: true + enable_prefix_caching: true + async_scheduling: true + max_model_len: 8192 + max_num_batched_tokens: 4096 + devices: "0" + model: Qwen/Qwen3-ASR-1.7B + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 256 + + - stage_id: 1 + max_num_seqs: 4 + gpu_memory_utilization: 0.4 + trust_remote_code: true + enable_prefix_caching: true + async_scheduling: true + max_model_len: 262144 + max_num_batched_tokens: 15360 + devices: "0" + model: aurateam/AURA + default_sampling_params: + temperature: 0.5 + top_p: 1.0 + top_k: -1 + max_tokens: 256 + repetition_penalty: 1.0 + + - stage_id: 2 + max_num_seqs: 8 + gpu_memory_utilization: 0.2 + trust_remote_code: true + enable_prefix_caching: false + async_scheduling: true + max_num_batched_tokens: 512 + max_model_len: 4096 + devices: "0" + output_connectors: + to_stage_3: connector_of_shared_memory + # model: Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice + model: Qwen/Qwen3-TTS-12Hz-1.7B-Base + default_sampling_params: + temperature: 0.9 + top_k: 50 + max_tokens: 4096 + seed: 42 + repetition_penalty: 1.05 + subtalker_sampling_params: + do_sample: true + temperature: 0.9 + top_k: 50 + top_p: 1.0 + + - stage_id: 3 + max_num_seqs: 8 + gpu_memory_utilization: 0.2 + enforce_eager: false + trust_remote_code: true + enable_prefix_caching: false + async_scheduling: true + max_num_batched_tokens: 65536 + max_model_len: 65536 + devices: "0" + input_connectors: + from_stage_2: connector_of_shared_memory + # model: Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice + model: Qwen/Qwen3-TTS-12Hz-1.7B-Base + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 65536 + seed: 42 + repetition_penalty: 1.0 + +platforms: + npu: + stages: + - stage_id: 2 + enforce_eager: true diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 3876414fe..41207102b 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -20,6 +20,7 @@ import janus import torch from vllm.config import ModelConfig from vllm.logger import init_logger +from vllm.outputs import CompletionOutput, RequestOutput from vllm.pooling_params import PoolingParams from vllm.sampling_params import SamplingParams from vllm.v1.engine import EngineCoreOutputs @@ -52,6 +53,68 @@ from vllm_omni.outputs import OmniRequestOutput logger = init_logger(__name__) +def _build_terminal_empty_output( + request_id: str, + *, + final_output_type: str | None, + audio_sample_rate: int = 24000, +) -> RequestOutput: + """Build a terminal empty output when no downstream stage input exists.""" + completion = CompletionOutput( + index=0, + text="", + token_ids=[], + cumulative_logprob=None, + logprobs=None, + finish_reason="stop", + stop_reason=None, + ) + if final_output_type == "audio": + completion.multimodal_output = { + "audio": torch.zeros((0,), dtype=torch.float32), + "sr": audio_sample_rate, + } + return RequestOutput( + request_id=request_id, + prompt=None, + prompt_token_ids=[], + prompt_logprobs=None, + outputs=[completion], + finished=True, + ) + + +def _coerce_int_scalar(value: Any) -> int: + if value is None: + return 0 + if isinstance(value, (list, tuple)): + for item in value: + coerced = _coerce_int_scalar(item) + if coerced > 0: + return coerced + return 0 + if hasattr(value, "item"): + value = value.item() + try: + coerced = int(value) + except (TypeError, ValueError): + return 0 + return coerced if coerced > 0 else 0 + + +def _infer_stage_audio_sample_rate(stage_pool: StagePool, default: int = 24000) -> int: + """Infer the final audio stage sample rate from stage metadata when possible.""" + sample_rate_attrs = ("audio_sample_rate", "sample_rate", "sampling_rate", "output_sample_rate", "sr") + stage_client = getattr(stage_pool, "stage_client", None) + stage_config = getattr(stage_pool, "_stage_vllm_config", None) + for source in (stage_client, stage_config): + for attr in sample_rate_attrs: + sample_rate = _coerce_int_scalar(getattr(source, attr, None)) + if sample_rate > 0: + return sample_rate + return default + + def build_engine_core_request_from_tokens( request_id: str, prompt: dict[str, Any], @@ -180,6 +243,7 @@ class Orchestrator: self._init_metrics_state(stage_pools, running_counter, transfer_emitter, log_stats=log_stats) self._cfg_tracker = CfgCompanionTracker() + self._stage_input_processors: dict[int, Any] = {} self._shutdown_event = asyncio.Event() self._stages_shutdown = False @@ -884,6 +948,77 @@ class Orchestrator: def _next_stage_already_submitted(self, stage_id: int, req_state: OrchestratorRequestState) -> bool: return (stage_id + 1) in req_state.stage_submit_ts + def _get_stage_input_processor(self, stage_id: int) -> Any: + processor = self._stage_input_processors.get(stage_id) + if processor is None: + from vllm_omni.engine.stage_init_utils import build_stage0_input_processor + + processor = build_stage0_input_processor(self.stage_pools[stage_id].stage_vllm_config) + self._stage_input_processors[stage_id] = processor + return processor + + def _upgrade_processed_stage_request(self, request: Any, raw_prompt: Any) -> Any: + prompt_embeds = getattr(request, "prompt_embeds", None) + additional_information = None + + if isinstance(raw_prompt, dict): + if prompt_embeds is None: + raw_prompt_embeds = raw_prompt.get("prompt_embeds") + if isinstance(raw_prompt_embeds, torch.Tensor): + prompt_embeds = raw_prompt_embeds + additional_information = serialize_additional_information( + raw_prompt.get("additional_information"), + log_prefix="Orchestrator stage input", + ) + + if prompt_embeds is None and additional_information is None: + return request + + return OmniEngineCoreRequest.from_request( + request, + prompt_embeds=prompt_embeds, + additional_information=additional_information, + ) + + def _next_stage_input_is_tokens(self, next_input: Any) -> bool: + return isinstance(next_input, dict) and "prompt_token_ids" in next_input + + def _build_next_stage_request( + self, + req_id: str, + next_stage_id: int, + next_input: Any, + params: SamplingParams | PoolingParams, + *, + mm_features: list | None = None, + resumable: bool = False, + ) -> Any: + next_pool = self.stage_pools[next_stage_id] + if self._next_stage_input_is_tokens(next_input): + request = build_engine_core_request_from_tokens( + request_id=req_id, + prompt=next_input, + params=params, + model_config=next_pool.stage_vllm_config.model_config, + mm_features=mm_features, + resumable=resumable, + ) + request.external_req_id = request.request_id + return request + + processor = self._get_stage_input_processor(next_stage_id) + request = processor.process_inputs( + request_id=req_id, + prompt=next_input, + params=params, + supported_tasks=("generate",), + arrival_time=_time.time(), + resumable=resumable, + ) + request = self._upgrade_processed_stage_request(request, next_input) + request.external_req_id = req_id + return request + async def _handle_cfg_companion_ready(self, req_id: str) -> None: """Mark a CFG companion as done; if all companions are done, flush deferred parent.""" parent_id = self._cfg_tracker.on_companion_completed(req_id) @@ -1207,22 +1342,64 @@ class Orchestrator: ) raise + if not next_inputs: + if not getattr(output, "finished", False): + logger.debug( + "[Orchestrator] req=%s stage-%s produced no inputs for stage-%s; waiting for more outputs", + req_id, + src_stage_id, + next_logical, + ) + return + + final_stage_id = req_state.final_stage_id + final_pool = self.stage_pools[final_stage_id] + final_output_type = getattr(final_pool.stage_client, "final_output_type", None) + terminal_output = _build_terminal_empty_output( + req_id, + final_output_type=final_output_type, + audio_sample_rate=_infer_stage_audio_sample_rate(final_pool), + ) + submit_ts = _time.time() + req_state.stage_submit_ts[final_stage_id] = submit_ts + logger.info( + "[Orchestrator] req=%s stage-%s produced no terminal inputs for stage-%s; " + "returning empty %s output from final stage-%s", + req_id, + src_stage_id, + next_logical, + final_output_type or "text", + final_stage_id, + ) + await self.output_async_queue.put( + OutputMessage( + request_id=req_id, + stage_id=final_stage_id, + replica_id=0, + engine_outputs=terminal_output, + metrics=None, + finished=True, + stage_submit_ts=submit_ts, + ) + ) + await self._cleanup_request_ids([req_id, *self._cfg_tracker.cleanup_parent(req_id)]) + return + # Build and submit requests for each input for next_input in next_inputs: # Only AR thinker stages consume encoder mm_features; downstream # (talker/code2wav/…) must not see them (avoids encoder-cache misses). model_stage = getattr(getattr(next_pool.stage_vllm_config, "model_config", None), "model_stage", None) mm_features = req_state.mm_features if model_stage == "thinker" else None - request = build_engine_core_request_from_tokens( - request_id=req_id, - prompt=next_input, + request = self._build_next_stage_request( + req_id, + next_logical, + next_input, params=params, - model_config=next_pool.stage_vllm_config.model_config, mm_features=mm_features, resumable=next_stage_resumable, ) - request.external_req_id = request.request_id if already_submitted: await next_pool.submit_update(req_id, req_state, request) else: diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 5c8f91169..7960bae61 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -698,6 +698,7 @@ def build_engine_args_dict( stage_id = stage_config.stage_id engine_args_dict = _to_dict(engine_args) + model = engine_args_dict.pop("model", None) or model stage_defines_tokenizer = ( engine_args_dict.get("tokenizer") is not None or engine_args_dict.get("tokenizer_subdir") is not None ) diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 9a0a9ed1c..67f82b27e 100755 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -74,6 +74,7 @@ from vllm.entrypoints.serve.utils.api_utils import should_include_usage from vllm.entrypoints.serve.utils.tool_calls_utils import maybe_filter_parallel_tool_calls from vllm.inputs import PromptType from vllm.logger import init_logger +from vllm.multimodal.media.connector import MediaConnector from vllm.outputs import RequestOutput from vllm.reasoning import ReasoningParser from vllm.renderers import BaseRenderer, merge_kwargs @@ -104,6 +105,7 @@ from vllm_omni.entrypoints.openai.protocol.images import ( ) from vllm_omni.entrypoints.openai.stage_params import ( build_stage_sampling_params_list, + clone_sampling_params, get_default_sampling_params_list, ) from vllm_omni.entrypoints.openai.utils import ( @@ -122,6 +124,10 @@ from vllm_omni.utils.audio import audio_chunk_pcm_bytes, audio_chunk_sample_rate logger = init_logger(__name__) +async def _identity_async(value: Any) -> Any: + return value + + class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): """OpenAI-compatible chat serving for both LLM and Diffusion models. @@ -654,6 +660,13 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): default_mm_processor_kwargs=getattr(request, "mm_processor_kwargs", None), ) + deferred_multi_modal_data: dict[str, Any] | None = None + if self._needs_multistage_multimodal_split(): + messages, deferred_multi_modal_data = await self._prepare_multistage_multimodal_inputs( + messages, + request, + ) + (conversation,), (engine_prompt,) = await renderer.render_chat_async( [messages], chat_params, @@ -698,18 +711,25 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): if hasattr(request, "cache_salt") and request.cache_salt is not None: engine_prompt["cache_salt"] = request.cache_salt + additional_information = getattr(request, "additional_information", None) + if isinstance(additional_information, dict): + prompt_additional_information = self._ensure_prompt_additional_information(engine_prompt) + prompt_additional_information.update(additional_information) + + if deferred_multi_modal_data: + prompt_additional_information = self._ensure_prompt_additional_information(engine_prompt) + prompt_additional_information["deferred_multi_modal_data"] = deferred_multi_modal_data + speaker = getattr(request, "voice", None) or getattr(request, "speaker", None) normalized = validate_requested_speaker(speaker, self._get_supported_speakers()) if normalized is not None: - if "additional_information" not in engine_prompt or engine_prompt["additional_information"] is None: - engine_prompt["additional_information"] = {} - engine_prompt["additional_information"]["speaker"] = [normalized] + prompt_additional_information = self._ensure_prompt_additional_information(engine_prompt) + prompt_additional_information["speaker"] = [normalized] language = getattr(request, "language", None) if language is not None and isinstance(language, str) and language.strip(): - if "additional_information" not in engine_prompt or engine_prompt["additional_information"] is None: - engine_prompt["additional_information"] = {} - engine_prompt["additional_information"]["language"] = [language.strip()] + prompt_additional_information = self._ensure_prompt_additional_information(engine_prompt) + prompt_additional_information["language"] = [language.strip()] # Style instruction — used by Ming-flash-omni instruct TTS path # (ming_task="instruct"). For the omni speech path the thinker2talker @@ -717,20 +737,263 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): # which hardcodes instruction=None. instructions = getattr(request, "instructions", None) if instructions is not None and isinstance(instructions, str) and instructions.strip(): - if "additional_information" not in engine_prompt or engine_prompt["additional_information"] is None: - engine_prompt["additional_information"] = {} - engine_prompt["additional_information"]["instruction"] = instructions.strip() + prompt_additional_information = self._ensure_prompt_additional_information(engine_prompt) + prompt_additional_information["instruction"] = instructions.strip() return conversation, [engine_prompt] + @staticmethod + def _ensure_prompt_additional_information(engine_prompt: dict[str, Any]) -> dict[str, Any]: + additional_information = engine_prompt.get("additional_information") + if not isinstance(additional_information, dict): + additional_information = {} + engine_prompt["additional_information"] = additional_information + return additional_information + + def _needs_multistage_multimodal_split(self) -> bool: + return bool(self._deferred_multimodal_modalities()) + + def _deferred_multimodal_modalities(self) -> set[str]: + stage_configs = list(getattr(self.engine_client, "stage_configs", []) or []) + if len(stage_configs) < 2: + return set() + + first_stage_modalities = self._stage_input_modalities(stage_configs[0]) + if not first_stage_modalities: + return set() + + downstream_modalities: set[str] = set() + for stage in stage_configs[1:]: + downstream_modalities.update(self._stage_input_modalities(stage)) + return downstream_modalities - first_stage_modalities + + @staticmethod + def _stage_input_modalities(stage: Any) -> set[str]: + engine_args = getattr(stage, "engine_args", None) + explicit = ( + getattr(stage, "input_modalities", None) + or getattr(stage, "modalities", None) + or getattr(engine_args, "input_modalities", None) + or getattr(engine_args, "modalities", None) + ) + if explicit: + return {str(modality) for modality in as_list(explicit)} + + model_stage = str(getattr(engine_args, "model_stage", None) or getattr(stage, "model_stage", "")).lower() + if model_stage in {"asr", "stt"} or model_stage.endswith("_asr"): + return {"audio"} + if any(name in model_stage for name in ("vision", "vl", "aura")): + return {"image", "video"} + if any(name in model_stage for name in ("tts", "talker", "code2wav")): + return set() + + if getattr(stage, "requires_multimodal_data", False): + return {"audio", "image", "video"} + return set() + + async def _prepare_multistage_multimodal_inputs( + self, + messages: list[ChatCompletionMessageParam], + request: ChatLikeRequest | ResponsesRequest, + ) -> tuple[list[ChatCompletionMessageParam], dict[str, Any] | None]: + """Hide modalities unsupported by stage 0 and stash them for downstream stages.""" + deferred_modalities = self._deferred_multimodal_modalities() + deferred_parts: dict[str, list[Any]] = {modality: [] for modality in deferred_modalities} + stripped_messages: list[ChatCompletionMessageParam] = [] + + for message in messages: + content = message.get("content") if isinstance(message, dict) else None + if not isinstance(content, list): + stripped_messages.append(message) + continue + + stripped_content: list[Any] = [] + changed = False + for part in content: + modality, payload = self._deferred_multimodal_part(part, deferred_modalities) + if modality is not None: + if payload is not None: + deferred_parts.setdefault(modality, []).append(payload) + changed = True + continue + stripped_content.append(part) + + if changed: + stripped_message = dict(message) + stripped_message["content"] = stripped_content + stripped_messages.append(cast(ChatCompletionMessageParam, stripped_message)) + else: + stripped_messages.append(message) + + deferred_parts = {modality: parts for modality, parts in deferred_parts.items() if parts} + if not deferred_parts: + return messages, None + + media_connector = MediaConnector( + media_io_kwargs=getattr(request, "media_io_kwargs", None), + allowed_local_media_path=getattr(self.model_config, "allowed_local_media_path", "") or "", + allowed_media_domains=getattr(self.model_config, "allowed_media_domains", None), + ) + multi_modal_data: dict[str, Any] = {} + for modality, parts in deferred_parts.items(): + multi_modal_data[modality] = await self._materialize_deferred_multimodal_parts( + media_connector, + modality, + parts, + ) + return stripped_messages, multi_modal_data + + @staticmethod + def _deferred_multimodal_part(part: Any, deferred_modalities: set[str]) -> tuple[str | None, Any | None]: + if not isinstance(part, dict): + return None, None + part_type = part.get("type") + if part_type in {"video_url", "video"} and "video" in deferred_modalities: + video = part.get("video_url", part.get("video")) + if isinstance(video, dict): + video = video.get("url") + return "video", video + if part_type in {"image_url", "image_pil", "image"} and "image" in deferred_modalities: + image = part.get("image_pil", part.get("image_url", part.get("image"))) + if isinstance(image, dict): + image = image.get("url") + return "image", image + if part_type in {"audio_url", "input_audio", "audio"} and "audio" in deferred_modalities: + audio = part.get("audio_url", part.get("input_audio", part.get("audio"))) + if isinstance(audio, dict): + audio = audio.get("url") or audio.get("data") + return "audio", audio + return None, None + + @staticmethod + async def _materialize_deferred_multimodal_parts( + media_connector: MediaConnector, + modality: str, + parts: list[Any], + ) -> list[Any]: + if modality == "video": + return list(await asyncio.gather(*(media_connector.fetch_video_async(part) for part in parts))) + if modality == "image": + fetch_image = getattr(media_connector, "fetch_image_async", None) + if fetch_image is None: + return parts + return list( + await asyncio.gather( + *(fetch_image(part) if isinstance(part, str) else _identity_async(part) for part in parts) + ) + ) + if modality == "audio": + fetch_audio = getattr(media_connector, "fetch_audio_async", None) + if fetch_audio is None: + return parts + return list( + await asyncio.gather( + *(fetch_audio(part) if isinstance(part, str) else _identity_async(part) for part in parts) + ) + ) + return parts + + async def _inject_audio_from_video_urls( + self, + messages: list[ChatCompletionMessageParam], + ) -> list[ChatCompletionMessageParam]: + """Pre-extract audio from video URLs and inject as audio_url content items. + + When use_audio_in_video=True, the qwen2_5_omni_thinker multimodal + processor requires that the number of audio items equals the number of + video items (it subtracts mm_counts["video"] from mm_counts["audio"]). + The client only sends video_url items; this method adds the matching + audio_url items on the server side before the renderer processes them. + """ + import io + + from vllm_omni.entrypoints.chat_utils import extract_audio_from_video_async + + new_messages: list[ChatCompletionMessageParam] = [] + for msg in messages: + content = msg.get("content") if isinstance(msg, dict) else getattr(msg, "content", None) + if not isinstance(content, list): + new_messages.append(msg) + continue + + video_urls = [ + part.get("video_url", {}).get("url") + for part in content + if isinstance(part, dict) and part.get("type") == "video_url" and part.get("video_url", {}).get("url") + ] + + if not video_urls: + new_messages.append(msg) + continue + + audios = await asyncio.gather(*(extract_audio_from_video_async(u) for u in video_urls)) + + audio_items: list[dict] = [] + for audio_array, sample_rate in audios: + buf = io.BytesIO() + if soundfile is not None: + soundfile.write(buf, audio_array, samplerate=int(sample_rate), format="WAV") + else: + import struct + + import numpy as np + + audio_np = np.asarray(audio_array, dtype=np.float32) + sr = int(sample_rate) + num_channels = 1 + bits_per_sample = 32 + num_frames = len(audio_np) + data_size = num_frames * num_channels * (bits_per_sample // 8) + # Write minimal RIFF/WAV header + buf.write(b"RIFF") + buf.write(struct.pack(" list[Any]: """Convert request dicts to stage-typed sampling params objects. For diffusion stages, build ``OmniDiffusionSamplingParams`` so downstream ``StageDiffusionClient._sampling_params_to_dict`` (which requires a dataclass) works. For LLM stages build ``SamplingParams``. + If callers provide params for fewer stages than the native pipeline has + (for example AURA has three semantic models but four engine stages), + append cloned deploy defaults for the omitted tail stages. """ stage_configs = list(getattr(self.engine_client, "stage_configs", []) or []) + default_params_list = list(getattr(self.engine_client, "default_sampling_params_list", []) or []) final_sampling_params_list: list[Any] = [] for idx, sampling_params in enumerate(sampling_params_list): stage_type = get_stage_type(stage_configs[idx]) if idx < len(stage_configs) else "llm" @@ -750,6 +1013,11 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): final_sampling_params_list.append(target_cls(**as_dict)) else: raise ValueError(f"Invalid sampling params: {sampling_params}") + for idx in range(len(final_sampling_params_list), len(stage_configs)): + if idx < len(default_params_list): + final_sampling_params_list.append(clone_sampling_params(default_params_list[idx])) + else: + final_sampling_params_list.append(SamplingParams()) return final_sampling_params_list def _get_comprehension_stage_index(self) -> int: diff --git a/vllm_omni/model_executor/models/aura_omni/pipeline.py b/vllm_omni/model_executor/models/aura_omni/pipeline.py new file mode 100644 index 000000000..273589e1b --- /dev/null +++ b/vllm_omni/model_executor/models/aura_omni/pipeline.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""AURA Omni pipeline topology. + +Semantic modules: + 1. Qwen3-ASR: microphone audio -> transcript text + 2. AURA/Qwen3-VL: transcript + video -> text or <|silent|> + 3. Qwen3-TTS: response text -> audio + +The Qwen3-TTS module is represented as two native engine stages, Talker and +Code2Wav, to reuse the existing vLLM-Omni implementation without wrapping or +duplicating codec streaming behavior. +""" + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +_AURA_PROC = "vllm_omni.model_executor.stage_input_processors.aura_omni" +_QWEN3_TTS_PROC = "vllm_omni.model_executor.stage_input_processors.qwen3_tts" + + +AURA_OMNI_PIPELINE = PipelineConfig( + model_type="aura_omni", + model_arch="Qwen3ASRForConditionalGeneration", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="asr", + execution_type=StageExecutionType.LLM_AR, + input_sources=(), + owns_tokenizer=True, + requires_multimodal_data=True, + engine_output_type="text", + model_arch="Qwen3ASRForConditionalGeneration", + sampling_constraints={"detokenize": True}, + ), + StagePipelineConfig( + stage_id=1, + model_stage="aura", + execution_type=StageExecutionType.LLM_AR, + input_sources=(0,), + final_output=True, + final_output_type="text", + owns_tokenizer=True, + requires_multimodal_data=True, + engine_output_type="text", + model_arch="AuraQwen3VLForConditionalGeneration", + custom_process_input_func=f"{_AURA_PROC}.asr2aura", + sampling_constraints={"detokenize": True}, + ), + StagePipelineConfig( + stage_id=2, + model_stage="qwen3_tts", + execution_type=StageExecutionType.LLM_AR, + input_sources=(1,), + owns_tokenizer=True, + engine_output_type="latent", + model_arch="Qwen3TTSTalkerForConditionalGeneration", + custom_process_input_func=f"{_AURA_PROC}.aura2tts", + custom_process_next_stage_input_func=f"{_QWEN3_TTS_PROC}.talker2code2wav_full_payload", + async_chunk_process_next_stage_input_func=f"{_QWEN3_TTS_PROC}.talker2code2wav_async_chunk", + sampling_constraints={ + "detokenize": False, + "stop_token_ids": [2150], + }, + ), + StagePipelineConfig( + stage_id=3, + model_stage="code2wav", + execution_type=StageExecutionType.LLM_GENERATION, + input_sources=(2,), + final_output=True, + final_output_type="audio", + engine_output_type="audio", + model_arch="Qwen3TTSCode2Wav", + sync_process_input_func=f"{_QWEN3_TTS_PROC}.talker2code2wav_token_only", + sampling_constraints={"detokenize": True}, + extras={"tts_args": {"max_instructions_length": 500}}, + ), + ), +) diff --git a/vllm_omni/model_executor/models/aura_omni/qwen3_vl.py b/vllm_omni/model_executor/models/aura_omni/qwen3_vl.py new file mode 100644 index 000000000..65243893b --- /dev/null +++ b/vllm_omni/model_executor/models/aura_omni/qwen3_vl.py @@ -0,0 +1,59 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""AURA-compatible Qwen3-VL wrapper. + +AURA ships Qwen3-VL weights with a repository-local ``configuration_qwen3_vl`` +module in ``auto_map``. With ``trust_remote_code=True``, Transformers therefore +returns a remote ``Qwen3VLConfig`` class. The model structure is still Qwen3-VL, +but vLLM's stock multimodal processor checks for the upstream Transformers +``Qwen3VLConfig`` type exactly. This wrapper keeps the stock model and processor +logic, while accepting structurally-compatible remote Qwen3-VL configs. +""" + +from __future__ import annotations + +from typing import Any + +from vllm.model_executor.models.qwen3_vl import ( + Qwen3VLConfig, + Qwen3VLDummyInputsBuilder, + Qwen3VLForConditionalGeneration, + Qwen3VLMultiModalProcessor, + Qwen3VLProcessingInfo, + Qwen3VLProcessor, +) +from vllm.multimodal import MULTIMODAL_REGISTRY + + +class AuraQwen3VLProcessingInfo(Qwen3VLProcessingInfo): + def get_hf_config(self) -> Any: + try: + return self.ctx.get_hf_config(Qwen3VLConfig) + except TypeError: + config = self.ctx.get_hf_config() + if ( + getattr(config, "model_type", None) == "qwen3_vl" + and hasattr(config, "vision_config") + and hasattr(config, "text_config") + ): + return config + raise + + def get_hf_processor(self, **kwargs: object) -> Any: + # AURA's config advertises `processing_qwen3_vl.py` in auto_map, but + # the checkpoint snapshot does not include that file. Force the + # upstream processor class while still allowing AURA's remote config. + return self.ctx.get_hf_processor( + Qwen3VLProcessor, + use_fast=kwargs.pop("use_fast", True), + **kwargs, + ) + + +@MULTIMODAL_REGISTRY.register_processor( + Qwen3VLMultiModalProcessor, + info=AuraQwen3VLProcessingInfo, + dummy_inputs=Qwen3VLDummyInputsBuilder, +) +class AuraQwen3VLForConditionalGeneration(Qwen3VLForConditionalGeneration): + pass diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py index c1bb7c049..105672399 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py @@ -30,7 +30,7 @@ from vllm_omni.utils.speaker_cache import ( ) from .configuration_qwen3_tts import Qwen3TTSConfig, Qwen3TTSSpeakerEncoderConfig, Qwen3TTSTalkerConfig -from .prompt_embeds_builder import Qwen3TTSPromptEmbedsBuilder +from .prompt_embeds_builder import PRECOMPUTED_TEXT_IDS_KEY, Qwen3TTSPromptEmbedsBuilder from .qwen3_tts_code_predictor_vllm import Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM from .tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2 import Qwen3TTSTokenizerV2Config from .tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import Qwen3TTSTokenizerV2Encoder @@ -40,6 +40,19 @@ logger = init_logger(__name__) _TRAILING_TEXT_COMPACT_MIN_FRAMES = 64 +def _has_tts_text_conditioning(info_dict: dict[str, Any], hidden_states: Any | None = None) -> bool: + text_list = info_dict.get("text") + if isinstance(text_list, list) and bool(text_list) and bool(text_list[0]): + return True + if PRECOMPUTED_TEXT_IDS_KEY in info_dict: + return True + if isinstance(hidden_states, dict): + tail = hidden_states.get("trailing_text") + if isinstance(tail, torch.Tensor): + return True + return False + + # --------------------------------------------------------------------------- # Components ported from the HuggingFace Qwen3-TTS reference implementation. # Only the classes actually needed by the vLLM AR Talker are kept here. @@ -625,9 +638,8 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): except Exception: is_prefill = span_len > 1 - text_list = info_dict.get("text") - if not isinstance(text_list, list) or not text_list or not text_list[0]: - raise ValueError("Missing additional_information.text for Qwen3-TTS AR talker.") + if not _has_tts_text_conditioning(info_dict, hs): + raise ValueError("Missing Qwen3-TTS text conditioning: provide `text` or precomputed text token ids.") task_type = (info_dict.get("task_type") or ["CustomVoice"])[0] codec_streaming_val = meta.get("codec_streaming") @@ -650,7 +662,6 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): if is_prefill: # Prefill (prompt embeddings) prompt_embeds_cpu = embed.get("prefill") - # First prefill round: prompt_embeds_cpu is not yet populated. # Subsequent prefill rounds (multi-chunk): prompt_embeds_cpu is a Tensor stored by the first round. is_first_prefill = not isinstance(prompt_embeds_cpu, torch.Tensor) or prompt_embeds_cpu.ndim != 2 @@ -808,9 +819,8 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): hs = payload.get("hidden_states", {}) meta = payload.get("meta", {}) - text_list = info_dict.get("text") - if not isinstance(text_list, list) or not text_list or not text_list[0]: - raise ValueError("Missing additional_information.text for Qwen3-TTS AR talker.") + if not _has_tts_text_conditioning(info_dict, hs): + raise ValueError("Missing Qwen3-TTS text conditioning: provide `text` or precomputed text token ids.") task_type = (info_dict.get("task_type") or ["CustomVoice"])[0] codec_streaming_val = meta.get("codec_streaming") diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index 9906fa4e0..d057502c2 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -343,6 +343,11 @@ _OMNI_MODELS = { "minicpmo_4_5_omni_tts", "MiniCPMO45OmniTTSForConditionalGeneration", ), + "AuraQwen3VLForConditionalGeneration": ( + "aura_omni", + "qwen3_vl", + "AuraQwen3VLForConditionalGeneration", + ), } diff --git a/vllm_omni/model_executor/stage_input_processors/aura_omni.py b/vllm_omni/model_executor/stage_input_processors/aura_omni.py new file mode 100644 index 000000000..c112e1e8e --- /dev/null +++ b/vllm_omni/model_executor/stage_input_processors/aura_omni.py @@ -0,0 +1,390 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Stage processors for the AURA Omni pipeline.""" + +from __future__ import annotations + +import math +import os +from pathlib import Path +from typing import Any + +import soundfile as sf + +from vllm_omni.inputs.data import OmniTokensPrompt +from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import ( + PRECOMPUTED_TEXT_IDS_KEY, +) + +DEFAULT_AURA_SYSTEM_PROMPT = ( + "You are receiving a live video stream where the final frame is the present moment. " + "Respond only when a response is needed based on the user's message or the visual context. " + "Otherwise, output '<|silent|>' to signify silence. Respond in Chinese." +) + +SILENT_TEXT = "<|silent|>" +QWEN_IM_START_ID = 151644 +QWEN_IM_END_ID = 151645 +QWEN_ASSISTANT_ID = 77091 +QWEN_NEWLINE_ID = 198 +QWEN_ASSISTANT_PREFIX_IDS = [QWEN_IM_START_ID, QWEN_ASSISTANT_ID, QWEN_NEWLINE_ID] +QWEN_ASSISTANT_SUFFIX_IDS = [ + QWEN_IM_END_ID, + QWEN_NEWLINE_ID, + QWEN_IM_START_ID, + QWEN_ASSISTANT_ID, + QWEN_NEWLINE_ID, +] +DEFAULT_QWEN3_TTS_REF_AUDIO = "vllm-omni/tests/assets/qwen3_tts/clone_2.wav" +DEFAULT_QWEN3_TTS_REF_TEXT = ( + "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you." +) + + +def default_qwen3_tts_ref_audio_path() -> str: + """Return absolute path to the bundled ``clone_2.wav`` reference asset.""" + bundled = Path(__file__).resolve().parents[3] / "tests" / "assets" / "qwen3_tts" / "clone_2.wav" + if bundled.is_file(): + return str(bundled) + return DEFAULT_QWEN3_TTS_REF_AUDIO + + +def _as_list(value: Any) -> list[Any]: + if value is None: + return [] + return value if isinstance(value, list) else [value] + + +def _as_prompt_dict(prompt_item: Any) -> dict[str, Any]: + return prompt_item if isinstance(prompt_item, dict) else {} + + +def _first_value(value: Any, default: Any = None) -> Any: + if isinstance(value, list): + return value[0] if value else default + return default if value is None else value + + +def _first_bool(value: Any, default: bool = False) -> bool: + value = _first_value(value, default) + if isinstance(value, str): + return value.strip().lower() in {"1", "true", "yes", "on"} + return bool(value) + + +def _normalize_qwen3_tts_speaker(speaker: Any) -> Any: + if not isinstance(speaker, str): + return speaker + speaker = speaker.strip() + if not speaker: + return speaker + if "_" in speaker: + return speaker + return speaker[0].upper() + speaker[1:].lower() + + +def _extract_output(source_output: Any) -> Any: + outputs = getattr(source_output, "outputs", None) + if isinstance(outputs, list) and outputs: + return outputs[0] + return source_output + + +def _extract_text(source_output: Any) -> str: + output = _extract_output(source_output) + cumulative_text = getattr(output, "cumulative_text", None) + if isinstance(cumulative_text, str) and cumulative_text: + return cumulative_text + text = getattr(output, "text", None) + if isinstance(text, str): + return text + mm = getattr(output, "multimodal_output", None) + if isinstance(mm, dict): + for key in ("text", "transcript", "asr_text"): + value = mm.get(key) + if isinstance(value, str): + return value + if isinstance(value, list) and value and isinstance(value[0], str): + return value[0] + return "" + + +def _extract_token_ids(source_output: Any) -> list[int]: + output = _extract_output(source_output) + token_ids = getattr(output, "cumulative_token_ids", None) + if isinstance(token_ids, list): + return [int(token_id) for token_id in token_ids if isinstance(token_id, int)] + return [] + + +def _trim_aura_response_token_ids(token_ids: list[int]) -> list[int]: + ids = list(token_ids) + if ids[: len(QWEN_ASSISTANT_PREFIX_IDS)] == QWEN_ASSISTANT_PREFIX_IDS: + ids = ids[len(QWEN_ASSISTANT_PREFIX_IDS) :] + if QWEN_IM_END_ID in ids: + ids = ids[: ids.index(QWEN_IM_END_ID)] + while ids and ids[-1] in {QWEN_IM_START_ID, QWEN_IM_END_ID, QWEN_NEWLINE_ID}: + ids.pop() + return ids + + +def _qwen3_tts_assistant_token_ids_from_aura(source_output: Any) -> list[int]: + content_ids = _trim_aura_response_token_ids(_extract_token_ids(source_output)) + if not content_ids: + return [] + return QWEN_ASSISTANT_PREFIX_IDS + content_ids + QWEN_ASSISTANT_SUFFIX_IDS + + +def _source_prompt_by_request_id(source_outputs: list[Any], prompt: Any) -> dict[str, dict[str, Any]]: + prompts = _as_list(prompt) + return { + str(getattr(source_output, "request_id", idx)): _as_prompt_dict(prompt_item) + for idx, (source_output, prompt_item) in enumerate(zip(source_outputs, prompts)) + } + + +def _vision_placeholder(multi_modal_data: dict[str, Any]) -> str: + if "video" in multi_modal_data: + return "<|vision_start|><|video_pad|><|vision_end|>" + if "image" in multi_modal_data: + return "<|vision_start|><|image_pad|><|vision_end|>" + return "" + + +def _vision_multimodal_data(multi_modal_data: dict[str, Any]) -> dict[str, Any]: + return {key: value for key, value in multi_modal_data.items() if key in {"image", "video"}} + + +def _aura_prompt(system_prompt: str, transcript: str, multi_modal_data: dict[str, Any]) -> str: + vision = _vision_placeholder(multi_modal_data) + query = transcript.strip() + user_body = f"{vision}{query}" if query else vision + return ( + f"<|im_start|>system\n{system_prompt}<|im_end|>\n" + f"<|im_start|>user\n{user_body}<|im_end|>\n" + "<|im_start|>assistant\n" + ) + + +def asr2aura( + source_outputs: list[Any], + prompt: Any = None, + requires_multimodal_data: bool = True, +) -> list[dict[str, Any]]: + """Build AURA Qwen3-VL prompts from ASR transcripts and original video payloads.""" + prompt_by_request_id = _source_prompt_by_request_id(source_outputs, prompt) + next_inputs: list[dict[str, Any]] = [] + for idx, source_output in enumerate(source_outputs): + src_prompt = prompt_by_request_id.get(str(getattr(source_output, "request_id", idx)), {}) + additional_info = src_prompt.get("additional_information") or {} + system_prompt = _first_value(additional_info.get("aura_system_prompt"), DEFAULT_AURA_SYSTEM_PROMPT) + transcript = _extract_text(source_output) + multi_modal_data = {} + source_multi_modal_data = src_prompt.get("multi_modal_data") or {} + if isinstance(source_multi_modal_data, dict): + multi_modal_data.update(source_multi_modal_data) + deferred_multi_modal_data = additional_info.get("deferred_multi_modal_data") or {} + if isinstance(deferred_multi_modal_data, dict): + multi_modal_data.update(deferred_multi_modal_data) + multi_modal_data = _vision_multimodal_data(multi_modal_data) + + next_input: dict[str, Any] = { + "prompt": _aura_prompt(str(system_prompt), transcript, multi_modal_data), + } + if requires_multimodal_data: + next_input["multi_modal_data"] = multi_modal_data + if src_prompt.get("mm_processor_kwargs") is not None: + next_input["mm_processor_kwargs"] = src_prompt.get("mm_processor_kwargs") + next_inputs.append(next_input) + return next_inputs + + +def _estimate_ref_code_len_from_ref_audio(ref_audio: Any) -> int | None: + """Estimate Qwen3-TTS ref_code length from a ref-audio payload. + + For Qwen3-TTS 12Hz models, code length is approximately: + ceil(duration_seconds * 12.5) + i.e. one codec frame per 1920 samples at 24kHz. + """ + + codec_frame_rate = 24000.0 / 1920.0 + + # Unwrap common list wrappers. + item = ref_audio + while isinstance(item, list) and item: + item = item[0] + + # Accept tuple/list like (wav, sr). + if isinstance(item, (tuple, list)) and len(item) == 2 and isinstance(item[1], (int, float)): + wav, sr = item + sr_i = int(sr) + if sr_i <= 0: + return None + if hasattr(wav, "__len__"): + n_samples = len(wav) + elif hasattr(wav, "shape"): + shape = getattr(wav, "shape", None) + if not shape: + return None + n_samples = shape[-1] if len(shape) > 1 else shape[0] + else: + return None + if n_samples <= 0: + return None + return max(1, int(math.ceil((float(n_samples) / float(sr_i)) * codec_frame_rate))) + + # Accept file path (wav only). + if isinstance(item, str) and item: + audio_path = item + if not os.path.isfile(audio_path) or not audio_path.lower().endswith(".wav"): + return None + try: + info = sf.info(audio_path) + n_frames = int(info.frames) + sr = int(info.samplerate) + if n_frames <= 0 or sr <= 0: + return None + return max(1, int(math.ceil((float(n_frames) / float(sr)) * codec_frame_rate))) + except Exception: + return None + + return None + + +def _estimate_tts_prompt_len_from_token_ids( + token_ids: list[int], + *, + task_type: str = "Base", + language: str = "Chinese", + instruct: str = "", + x_vector_only_mode: bool = False, + non_streaming_mode: bool | None = None, + ref_code_len: int | None = None, +) -> int: + """Estimate Talker prefill length from prompt structure. + + This mirrors Qwen3-TTS prompt assembly at length level: + prompt_len = instruct_len + role_len + codec_prefix_len + text/icl term + """ + + # Official defaults: Base -> streaming, others -> non-streaming. + if non_streaming_mode is None: + non_streaming_mode = task_type in ("CustomVoice", "VoiceDesign") + + # We do not have tokenizer here; use char length as a monotonic proxy. + instruct_len = len(instruct.strip()) if isinstance(instruct, str) else 0 + assistant_len = max(0, len(token_ids)) + + # role_len = 3; codec_prefix_len = (prefill_len + speaker_len + 2) - 1 + # prefill_len = 4 when language_id exists else 3. Use non-auto language as + # the language-id-present proxy. + has_language_id = isinstance(language, str) and language.strip().lower() != "auto" + prefill_len = 4 if has_language_id else 3 + speaker_len = 1 if task_type in ("CustomVoice", "Base") else 0 + base_len = instruct_len + 3 + (prefill_len + speaker_len + 2 - 1) + if task_type in ("CustomVoice", "VoiceDesign"): + if non_streaming_mode: + prompt_len = base_len + max(0, assistant_len - 6) + else: + prompt_len = base_len + 1 + return max(32, min(4096, int(prompt_len))) + + if task_type == "Base": + in_context_mode = not bool(x_vector_only_mode) + if in_context_mode and ref_code_len is not None: + codec_lens = 1 + int(ref_code_len) + if non_streaming_mode: + # Exact non-streaming ICL needs ref_ids token length; unavailable + # in this processor. Keep a conservative upper estimate. + prompt_len = base_len + codec_lens + max(0, assistant_len - 8) + 1 + else: + # Streaming ICL exact length term: 1 + ref_code_len + prompt_len = base_len + codec_lens + else: + # Base x-vector-only (or missing ref_code length) follows CV shape. + if non_streaming_mode: + prompt_len = base_len + max(0, assistant_len - 6) + else: + prompt_len = base_len + 1 + return max(32, min(4096, int(prompt_len))) + + # Defensive fallback for unknown task types. + return max(32, min(4096, int(base_len + max(assistant_len, 1)))) + + +def aura2tts( + source_outputs: list[Any], + prompt: Any = None, + requires_multimodal_data: bool = False, +) -> list[OmniTokensPrompt]: + """Convert AURA text output into Qwen3-TTS Talker requests.""" + del requires_multimodal_data + prompt_by_request_id = _source_prompt_by_request_id(source_outputs, prompt) + next_inputs: list[OmniTokensPrompt] = [] + for idx, source_output in enumerate(source_outputs): + text = _extract_text(source_output).strip() + if not text or text == SILENT_TEXT: + continue + + src_prompt = prompt_by_request_id.get(str(getattr(source_output, "request_id", idx)), {}) + additional_info = src_prompt.get("additional_information") or {} + task_type = _first_value(additional_info.get("tts_task_type"), "Base") + language = _first_value(additional_info.get("tts_language"), "English") + instruct = _first_value(additional_info.get("tts_instruct"), "") + x_vector_only_mode = _first_bool(additional_info.get("tts_x_vector_only_mode"), False) + non_streaming_mode_raw = _first_value(additional_info.get("tts_non_streaming_mode"), None) + non_streaming_mode = non_streaming_mode_raw if isinstance(non_streaming_mode_raw, bool) else None + ref_code_len_raw = _first_value(additional_info.get("tts_ref_code_length"), None) + ref_code_len = int(ref_code_len_raw) if isinstance(ref_code_len_raw, int) else None + ref_audio = None + ref_text = None + if task_type == "Base" and not x_vector_only_mode and ref_code_len is None: + ref_audio = _first_value(additional_info.get("tts_ref_audio"), None) + ref_code_len = _estimate_ref_code_len_from_ref_audio(ref_audio) + + assistant_token_ids_for_len = _qwen3_tts_assistant_token_ids_from_aura(source_output) + pass_token_ids = _first_bool(additional_info.get("tts_pass_token_ids"), False) + tts_info = { + "task_type": [task_type], + "language": [language], + "instruct": [instruct], + "max_new_tokens": [int(_first_value(additional_info.get("tts_max_new_tokens"), 2048))], + } + if pass_token_ids and assistant_token_ids_for_len: + tts_info[PRECOMPUTED_TEXT_IDS_KEY] = [assistant_token_ids_for_len] + else: + tts_info["text"] = [text] + if ref_code_len is not None: + tts_info["ref_code_length"] = [int(ref_code_len)] + prompt_len = _estimate_tts_prompt_len_from_token_ids( + assistant_token_ids_for_len if assistant_token_ids_for_len else [0] * max(0, len(text)), + task_type=str(task_type), + language=str(language), + instruct=str(instruct), + x_vector_only_mode=x_vector_only_mode, + non_streaming_mode=non_streaming_mode, + ref_code_len=ref_code_len, + ) + + if task_type == "Base": + ref_audio = ref_audio or _first_value(additional_info.get("tts_ref_audio"), None) + ref_text = _first_value(additional_info.get("tts_ref_text"), None) + if not ref_audio or not ref_text: + raise ValueError("AURA Base TTS requires tts_ref_audio and tts_ref_text.") + x_vector_only_mode = _first_bool(additional_info.get("tts_x_vector_only_mode"), False) + tts_info["ref_audio"] = [ref_audio] + tts_info["ref_text"] = [ref_text] + tts_info["x_vector_only_mode"] = [x_vector_only_mode] + elif task_type == "CustomVoice": + tts_info["speaker"] = [ + _normalize_qwen3_tts_speaker(_first_value(additional_info.get("tts_speaker"), "Vivian")) + ] + next_inputs.append( + OmniTokensPrompt( + prompt_token_ids=[0] * prompt_len, + additional_information=tts_info, + multi_modal_data=None, + mm_processor_kwargs=None, + ) + ) + return next_inputs diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index cdf2286bf..184ed8251 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -1,4 +1,5 @@ import contextlib +import inspect from collections.abc import Callable from dataclasses import replace from typing import TYPE_CHECKING, Any, cast @@ -48,6 +49,30 @@ else: logger = init_logger(__name__) +def _filter_mrope_kwargs_for_model(model: object, kwargs: dict[str, Any]) -> dict[str, Any]: + """Return only M-RoPE kwargs accepted by the model implementation.""" + method = getattr(model, "get_mrope_input_positions") + try: + signature = inspect.signature(method) + except (TypeError, ValueError): + return kwargs + + if any(param.kind == inspect.Parameter.VAR_KEYWORD for param in signature.parameters.values()): + return kwargs + + accepted = { + name + for name, param in signature.parameters.items() + if name not in {"self", "input_tokens"} + and param.kind + in { + inspect.Parameter.POSITIONAL_OR_KEYWORD, + inspect.Parameter.KEYWORD_ONLY, + } + } + return {key: value for key, value in kwargs.items() if key in accepted} + + class OmniGPUModelRunner(GPUModelRunner): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) @@ -338,7 +363,7 @@ class OmniGPUModelRunner(GPUModelRunner): kwargs["target_w"] = target_w req_state.mrope_positions, req_state.mrope_position_delta = self.model.get_mrope_input_positions( req_state.prompt_token_ids, - **kwargs, + **_filter_mrope_kwargs_for_model(self.model, kwargs), ) else: req_state.mrope_positions, req_state.mrope_position_delta = MRotaryEmbedding.get_input_positions_tensor(