[Model] Aura support - Non async chunk path (#4257)

Signed-off-by: Rein Yang <ruiruyang2@gmail.com>
Signed-off-by: R2-Y <ruiruyang2@gmail.com>

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`).

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())

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 \
+  "$@"

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/).

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()

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())

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 <<EOF
+{
+  "model": "$MODEL",
+  "modalities": ["text", "audio"],
+  "sampling_params_list": [
+    {"temperature": 0.0, "top_p": 1.0, "top_k": -1, "max_tokens": 256, "seed": 42},
+    {"temperature": 0.5, "top_p": 1.0, "top_k": -1, "max_tokens": 256, "seed": 42, "repetition_penalty": 1.0},
+    {"temperature": 0.9, "top_k": 50, "max_tokens": 4096, "seed": 42, "detokenize": false, "repetition_penalty": 1.05, "stop_token_ids": [2150]},
+    {"temperature": 0.0, "top_p": 1.0, "top_k": -1, "max_tokens": 65536, "seed": 42, "repetition_penalty": 1.0}
+  ],
+  "additional_information": {
+    "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. Otherwise output '<|silent|>'. 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

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