Merge branch 'main' into fix/rebase-pipeline-upload-all-e2e

Signed-off-by: Weiming Liao <liaowm5@gmail.com>

.buildkite/rebase-pipeline.yaml

@@ -14,7 +14,7 @@ steps:
     depends_on: image-build
     key: upload-ready-pipeline
     commands:
-      - buildkite-agent pipeline upload .buildkite/test-ready.yml
+      - python3 .buildkite/scripts/upload_pipeline.py --upload --all .buildkite/test-ready.yml
     agents:
       queue: "cpu_queue_premerge"
 
@@ -23,7 +23,7 @@ steps:
     depends_on: image-build
     key: upload-merge-pipeline
     commands:
-      - buildkite-agent pipeline upload .buildkite/test-merge.yml
+      - python3 .buildkite/scripts/upload_pipeline.py --upload --all .buildkite/test-merge.yml
     agents:
       queue: "cpu_queue_premerge"
 
@@ -33,6 +33,6 @@ steps:
     key: upload-nightly-pipeline
     if: build.env("NIGHTLY") == "1"
     commands:
-      - buildkite-agent pipeline upload .buildkite/test-nightly.yml
+      - python3 .buildkite/scripts/upload_pipeline.py --upload --all .buildkite/test-nightly.yml
     agents:
       queue: "cpu_queue_premerge"

.buildkite/scripts/upload_pipeline.py

@@ -13,13 +13,16 @@ Test pipeline mode (e.g. test-merge.yml):
     recognize this uploader-only key).
 
 Usage:
-  python3 upload_pipeline.py [--upload] <pipeline.yml>
+  python3 upload_pipeline.py [--upload] [--all] <pipeline.yml>
 
   # Bootstrap (replaces upload_pipeline_with_skip_ci.sh):
   python3 upload_pipeline.py --upload .buildkite/pipeline.yml
 
   # Test pipeline (replaces upload_test_pipeline_with_diff_skip.py):
   python3 upload_pipeline.py --upload .buildkite/test-merge.yml
+
+  # Full suite (rebase pipeline): keep all steps, still strip the uploader-only key:
+  python3 upload_pipeline.py --upload --all .buildkite/test-merge.yml
 """
 
 from __future__ import annotations
@@ -437,11 +440,17 @@ def resolve_pipeline_path(arg: str) -> Path:
     return ROOT / path
 
 
-def render_pipeline(path: Path) -> str:
+def render_pipeline(path: Path, *, force_all: bool = False) -> str:
     text = path.read_text(encoding="utf-8")
     diff_range = resolve_diff_range()
     changed_files = resolve_changed_files()
 
+    # ``--all`` forces the keep-all-steps path (no diff-aware skipping) while still
+    # stripping ``source_file_dependencies``. Used by the rebase pipeline so main builds
+    # run the full e2e suite (see .buildkite/rebase-pipeline.yaml).
+    if force_all:
+        changed_files = None
+
     if BOOTSTRAP_MARKER in text:
         rendered = render_bootstrap_pipeline(
             text,
@@ -480,6 +489,11 @@ def main() -> int:
         action="store_true",
         help="Pipe rendered YAML to buildkite-agent pipeline upload",
     )
+    parser.add_argument(
+        "--all",
+        action="store_true",
+        help="Keep all steps (disable diff-aware skipping); still strips source_file_dependencies",
+    )
     args = parser.parse_args()
 
     path = resolve_pipeline_path(args.pipeline)
@@ -487,7 +501,7 @@ def main() -> int:
         _log(f"missing pipeline file: {path}")
         return 1
 
-    rendered = render_pipeline(path)
+    rendered = render_pipeline(path, force_all=args.all)
     if args.upload:
         upload_to_buildkite(rendered)
     else:

.claude/skills/add-diffusion-model/SKILL.md

@@ -24,7 +24,7 @@ Check the model's HF repo for `model_index.json`. This determines your path:
 
 | Scenario | How to identify | Migration path |
 |----------|----------------|----------------|
-| **Already supported** | `_class_name` in `model_index.json` matches a key in `_DIFFUSION_MODELS` in `registry.py` | Skip to Step 5 (test) and Step 7 (docs) |
+| **Already supported** | `_class_name` in `model_index.json` matches a key in `_DIFFUSION_MODELS` in `registry.py` | Skip implementation, then validate model-specific examples, tests, and docs as needed |
 | **Diffusers-based** | Has standard `model_index.json` with `_diffusers_version`, subfolders for `transformer/`, `vae/`, etc. | Follow **Path A** below |
 | **Custom/private repo** | No diffusers `model_index.json`, weights in non-standard format, custom model code in a separate git repo | Follow **Path B** below |
 | **Hybrid** | Has some diffusers components (VAE) but custom transformer/fusion | Mix of Path A and Path B |
@@ -52,7 +52,7 @@ vllm_omni/diffusion/models/your_model_name/
 2. Replace attention with `vllm_omni.diffusion.attention.layer.Attention` (QKV shape: `[B, seq, heads, head_dim]`).
 3. Add `od_config: OmniDiffusionConfig | None = None` to `__init__`.
 4. Add `load_weights()` method mapping diffusers weight names to vllm-omni names.
-5. Add class attributes: `_repeated_blocks`, `_layerwise_offload_blocks_attr`.
+5. Add class attributes for acceleration features such as `_repeated_blocks` and `_layerwise_offload_blocks_attrs` (see [references/transformer-adaptation.md](references/transformer-adaptation.md) for examples).
 
 ### A4. Adapt pipeline
 
@@ -78,13 +78,34 @@ class YourPipeline(nn.Module):
 
 Add post/pre-process functions in the same pipeline file. Register them in `registry.py`.
 
+### A4.1 Add progress bar support (recommended)
+
+For pipelines with a standard denoising loop, prefer the existing progress bar pattern instead of hand-rolled logging.
+
+```python
+from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
+
+class YourPipeline(nn.Module, ProgressBarMixin):
+    def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput:
+        # ... prepare timesteps / latents ...
+        with self.progress_bar(total=len(timesteps)) as progress_bar:
+            for i, t in enumerate(timesteps):
+                # predict noise / scheduler step
+                latents = ...
+                progress_bar.update()
+
+        return DiffusionOutput(output=output)
+```
+
+For custom loop structures, follow `vllm_omni/diffusion/models/progress_bar.py` and existing pipelines using `ProgressBarMixin`.
+
 ### A5. Register, test, docs → continue at Step 4 below.
 
 ---
 
 ## Path B: Custom/Private Repo Model
 
-For models without a diffusers pipeline — weights in custom format, model code in a private repo. Real examples: DreamID-Omni, BAGEL, HunyuanImage3.
+For models without a diffusers pipeline — weights in custom format, model code in a private repo. Real examples: DreamID-Omni, BAGEL, HunyuanImage3. See [references/custom-model-patterns.md](references/custom-model-patterns.md) for concrete custom-model integration patterns.
 
 ### B1. Understand the reference repo
 
@@ -134,7 +155,7 @@ Copy the essential model files into `vllm_omni/diffusion/models/<name>/` and ada
 
 ### B4. Handle custom weight loading
 
-Custom models have two patterns for weight loading:
+Custom models have two common patterns for weight loading:
 
 **Pattern 1: Bypass standard loader** (DreamID-Omni style)
 
@@ -149,8 +170,10 @@ class CustomPipeline(nn.Module):
         self.vae = custom_init_vae(model, device=self.device)
         self.text_encoder = custom_init_text_encoder(model, device=self.device)
         self.transformer = CustomFusionModel(CONFIG)
-        load_custom_checkpoint(self.transformer,
-            checkpoint_path=os.path.join(model, "model.safetensors"))
+        load_custom_checkpoint(
+            self.transformer,
+            checkpoint_path=os.path.join(model, "model.safetensors"),
+        )
         # NO weights_sources defined — bypasses standard loader
 
     def load_weights(self, weights):
@@ -162,6 +185,8 @@ class CustomPipeline(nn.Module):
 When weights are in safetensors format but need name remapping:
 
 ```python
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+
 class CustomPipeline(nn.Module):
     def __init__(self, *, od_config, prefix=""):
         super().__init__()
@@ -187,7 +212,9 @@ class CustomPipeline(nn.Module):
             # Remap original weight names to vllm-omni module names
             name = self._remap_weight_name(name)
             if name in params:
-                default_weight_loader(params[name], tensor)
+                param = params[name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, tensor)
                 loaded.add(name)
         return loaded
 ```
@@ -281,12 +308,12 @@ For Omni or custom models, create:
 ### Step 7: Update Documentation
 
 Required updates:
-1. `docs/user_guide/diffusion/parallelism_acceleration.md` — parallelism support table
+1. `docs/user_guide/diffusion/parallelism/overview.md` — parallelism support overview/table
 2. `docs/user_guide/diffusion/cpu_offload_diffusion.md` — if CPU offload supported (add to supported models table)
-3. `docs/user_guide/diffusion/teacache.md` — if TeaCache supported
-4. `docs/user_guide/diffusion/cache_dit_acceleration.md` — if Cache-DiT supported
-5. `examples/offline_inference/xxx/README.md` — offline example docs
-6. `examples/online_serve/xxx/README.md` — online serve docs
+3. `docs/user_guide/diffusion/cache_acceleration/teacache.md` — if TeaCache supported
+4. `docs/user_guide/diffusion/cache_acceleration/cache_dit.md` — if Cache-DiT supported
+5. Offline example docs under `examples/offline_inference/<name>/` (`README.md` or category-specific `.md`)
+6. `examples/online_serving/<name>/README.md` — online serving docs
 
 ### Step 8: Add E2E Tests
 
@@ -537,7 +564,7 @@ class YourTransformer(nn.Module):
 #### 10e. Update parallelism documentation
 
 After adding parallelism support, update:
-1. `docs/user_guide/diffusion/parallelism_acceleration.md` — add your model to the support table
+1. `docs/user_guide/diffusion/parallelism/overview.md` — add your model to the support overview/table
 2. Record which parallelism methods are supported (USP, Ring, CFG, TP, HSDP, VAE-Patch)
 
 ### Step 11: Add CPU Offload Support

.claude/skills/add-diffusion-model/references/parallelism-patterns.md

@@ -569,4 +569,4 @@ Check which parallelism methods a model supports:
 | **TP** | Uses `ColumnParallelLinear` / `QKVParallelLinear`. Search: `grep -r 'ParallelLinear\|QKVParallel' vllm_omni/diffusion/models/<model>/` |
 | **HSDP** | Transformer defines `_hsdp_shard_conditions`. Search: `grep -r '_hsdp_shard_conditions' vllm_omni/diffusion/models/` |
 
-The canonical per-model support table is in `docs/user_guide/diffusion/parallelism_acceleration.md`.
+The canonical per-model support overview/table is in `docs/user_guide/diffusion/parallelism/overview.md`.

.claude/skills/add-diffusion-model/references/transformer-adaptation.md

@@ -219,13 +219,14 @@ def load_weights(self, weights):
 
 ### Class-Level Attributes for Features
 
+
 ```python
 class MyTransformer(nn.Module):
     # torch.compile: list block class names that repeat and can be compiled
     _repeated_blocks = ["MyTransformerBlock"]
 
     # CPU offload: attribute name of the nn.ModuleList containing blocks
-    _layerwise_offload_blocks_attr = "blocks"
+    _layerwise_offload_blocks_attrs = ["blocks"]
 
     # LoRA: mapping of fused param names to original param names
     packed_modules_mapping = {"to_qkv": ["to_q", "to_k", "to_v"]}

benchmarks/diffusion/diffusion_benchmark_serving.py

@@ -1031,6 +1031,14 @@ def _save_generated_outputs(
         if not out.success or not out.response_body:
             continue
 
+        if isinstance(out.response_body, bytes):
+            fname = f"req_{idx:04d}.mp4"
+            fpath = os.path.join(save_dir, fname)
+            with open(fpath, "wb") as f:
+                f.write(out.response_body)
+            saved += 1
+            continue
+
         media_urls: list[str] = []
 
         # Chat-completions style: choices[*].message.content[*].image_url.url

docker/Dockerfile.npu

@@ -1,5 +1,5 @@
-ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.18.0rc1
+ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
+ARG VLLM_ASCEND_TAG=v0.23.0-4885d8e
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
 # WORKDIR /vllm-workspace/vllm

docker/Dockerfile.npu.a3

@@ -1,5 +1,5 @@
-ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.18.0rc1-a3
+ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
+ARG VLLM_ASCEND_TAG=v0.23.0-4885d8e-a3
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
 # WORKDIR /vllm-workspace/vllm

docker/Dockerfile.npu.ci

@@ -1,5 +1,5 @@
-ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
-ARG VLLM_ASCEND_TAG=nightly-main
+ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
+ARG VLLM_ASCEND_TAG=v0.23.0-4885d8e
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
 ARG APP_DIR=/vllm-workspace/vllm-omni

docker/Dockerfile.npu.ci.a3

@@ -1,5 +1,5 @@
-ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
-ARG VLLM_ASCEND_TAG=nightly-main-a3
+ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
+ARG VLLM_ASCEND_TAG=v0.23.0-4885d8e-a3
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
 ARG APP_DIR=/vllm-workspace/vllm-omni

docs/.nav.yml

@@ -18,7 +18,6 @@ nav:
     - Offline Inference:
       - BAGEL-7B-MoT: user_guide/examples/offline_inference/bagel.md
       - GLM-Image Multistage End-to-End Inference: user_guide/examples/offline_inference/glm_image.md
-      - Helios Video Generation: user_guide/examples/offline_inference/helios.md
       - HunyuanImage-3.0 Image-to-Text Inference: user_guide/examples/offline_inference/hunyuan_image3.md
       - Image-To-Image: user_guide/examples/offline_inference/image_to_image.md
       - Image-To-Video: user_guide/examples/offline_inference/image_to_video.md

docs/mkdocs/hooks/generate_api_readme.py

@@ -144,6 +144,7 @@ def scan_package(package_name: str = "vllm_omni") -> dict[str, list[str]]:
                 "vllm_omni.entrypoints.async_diffusion",
                 "vllm_omni.entrypoints.openai",
                 "vllm_omni.model_executor.models.voxtral_tts.configuration_voxtral_tts",
+                "vllm_omni.experimental",  # optional serving deps not installed in docs build
             ]
             if any(module_path.startswith(prefix) for prefix in excluded_prefixes):
                 continue

docs/models/supported_models.md

@@ -36,7 +36,7 @@ th {
 | `ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `WanSpeechToVideoPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `LTX2Pipeline` | LTX-2-T2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |

docs/serving/videos_api.md

@@ -160,7 +160,12 @@ JSON references currently support `image_url`/`video_url`; `file_id` references
 are not implemented yet. Models may expose additional V2V controls through
 `extra_params`. For example, Cosmos3 supports
 `condition_frame_indexes_vision` and `condition_video_keep` to select which
-decoded reference frames are used as clean conditioning.
+decoded reference frames are used as clean conditioning. Cosmos3 transfer mode
+also accepts `edge`, `blur`, `depth`, `seg`, or `wsm` control hints plus
+transfer options such as `control_path`, `control_guidance`,
+`control_guidance_interval`, `num_video_frames_per_chunk`,
+`num_conditional_frames`, `show_control_condition`, and `show_input`; see the
+Cosmos3 recipe for complete examples.
 
 ### Speech-to-Video
 

docs/user_guide/examples/offline_inference/helios.md

@@ -1,203 +0,0 @@
-# Helios Video Generation
-
-Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/helios>.
-
-
-Helios is a text-to-video (T2V), image-to-video (I2V), and video-to-video (V2V) diffusion model. This example demonstrates end-to-end video generation using vLLM-Omni with three model variants:
-
-| Variant | Description | Key Features |
-|---------|-------------|--------------|
-| **Helios-Base** | Base model, Stage 1 only | Single-stage denoising, `guidance_scale=5.0` |
-| **Helios-Mid** | Mid model, Stage 2 pyramid | Multi-stage pyramid denoising, CFG-Zero* support |
-| **Helios-Distilled** | Distilled model, Stage 2+3 | Few-step inference with DMD, `guidance_scale=1.0` |
-
-## Setup
-
-Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
-
-## Run Examples
-
-Get into the example folder:
-```bash
-cd examples/offline_inference/helios
-```
-
-### Text-to-Video (T2V)
-
-**Helios-Base** (Stage 1 only):
-```bash
-python end2end.py \
-    --model BestWishYsh/Helios-Base \
-    --sample-type t2v \
-    --prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \
-    --guidance-scale 5.0 \
-    --output helios_t2v_base.mp4
-```
-
-**Helios-Mid** (Stage 2 + CFG-Zero*):
-```bash
-python end2end.py \
-    --model BestWishYsh/Helios-Mid \
-    --sample-type t2v \
-    --prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \
-    --guidance-scale 5.0 \
-    --is-enable-stage2 \
-    --pyramid-num-inference-steps-list 20 20 20 \
-    --use-cfg-zero-star --use-zero-init --zero-steps 1 \
-    --output helios_t2v_mid.mp4
-```
-
-**Helios-Distilled** (Stage 2 pyramid + DMD):
-```bash
-python end2end.py \
-    --model BestWishYsh/Helios-Distilled \
-    --sample-type t2v \
-    --prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \
-    --num-frames 240 \
-    --guidance-scale 1.0 \
-    --is-enable-stage2 \
-    --pyramid-num-inference-steps-list 2 2 2 \
-    --is-amplify-first-chunk \
-    --output helios_t2v_distilled.mp4
-```
-
-### Image-to-Video (I2V)
-
-**Helios-Base**:
-```bash
-python end2end.py \
-    --model BestWishYsh/Helios-Base \
-    --sample-type i2v \
-    --image-path /path/to/image.jpg \
-    --prompt "A towering emerald wave surges forward, its crest curling with raw power and energy." \
-    --guidance-scale 5.0 \
-    --output helios_i2v_base.mp4
-```
-
-**Helios-Mid** (Stage 2 + CFG-Zero*):
-```bash
-python end2end.py \
-    --model BestWishYsh/Helios-Mid \
-    --sample-type i2v \
-    --image-path /path/to/image.jpg \
-    --prompt "A towering emerald wave surges forward, its crest curling with raw power and energy." \
-    --guidance-scale 5.0 \
-    --is-enable-stage2 \
-    --pyramid-num-inference-steps-list 20 20 20 \
-    --use-cfg-zero-star --use-zero-init --zero-steps 1 \
-    --output helios_i2v_mid.mp4
-```
-
-**Helios-Distilled**:
-```bash
-python end2end.py \
-    --model BestWishYsh/Helios-Distilled \
-    --sample-type i2v \
-    --image-path /path/to/image.jpg \
-    --prompt "A towering emerald wave surges forward, its crest curling with raw power and energy." \
-    --num-frames 240 \
-    --guidance-scale 1.0 \
-    --is-enable-stage2 \
-    --pyramid-num-inference-steps-list 2 2 2 \
-    --is-amplify-first-chunk \
-    --output helios_i2v_distilled.mp4
-```
-
-### Video-to-Video (V2V)
-
-**Helios-Base**:
-```bash
-python end2end.py \
-    --model BestWishYsh/Helios-Base \
-    --sample-type v2v \
-    --video-path /path/to/video.mp4 \
-    --prompt "A bright yellow Lamborghini speeds along a curving mountain road." \
-    --guidance-scale 5.0 \
-    --output helios_v2v_base.mp4
-```
-
-**Helios-Mid** (Stage 2 + CFG-Zero*):
-```bash
-python end2end.py \
-    --model BestWishYsh/Helios-Mid \
-    --sample-type v2v \
-    --video-path /path/to/video.mp4 \
-    --prompt "A bright yellow Lamborghini speeds along a curving mountain road." \
-    --guidance-scale 5.0 \
-    --is-enable-stage2 \
-    --pyramid-num-inference-steps-list 20 20 20 \
-    --use-cfg-zero-star --use-zero-init --zero-steps 1 \
-    --output helios_v2v_mid.mp4
-```
-
-**Helios-Distilled**:
-```bash
-python end2end.py \
-    --model BestWishYsh/Helios-Distilled \
-    --sample-type v2v \
-    --video-path /path/to/video.mp4 \
-    --prompt "A bright yellow Lamborghini speeds along a curving mountain road." \
-    --num-frames 240 \
-    --guidance-scale 1.0 \
-    --is-enable-stage2 \
-    --pyramid-num-inference-steps-list 2 2 2 \
-    --is-amplify-first-chunk \
-    --output helios_v2v_distilled.mp4
-```
-
-## Common Parameters
-
-| Parameter | Default | Description |
-|-----------|---------|-------------|
-| `--model` | `BestWishYsh/Helios-Base` | Model ID or local path |
-| `--sample-type` | `t2v` | Generation mode: `t2v`, `i2v`, or `v2v` |
-| `--prompt` | — | Text prompt describing the video |
-| `--negative-prompt` | *(see source)* | Negative prompt for CFG (includes anti-static terms) |
-| `--image-path` | — | Input image (required for `i2v`) |
-| `--video-path` | — | Input video (required for `v2v`) |
-| `--height` | `384` | Video height in pixels |
-| `--width` | `640` | Video width in pixels |
-| `--num-frames` | `99` | Number of output frames |
-| `--num-inference-steps` | `50` | Denoising steps (Stage 1 only) |
-| `--guidance-scale` | `5.0` | CFG scale (`1.0` for Distilled) |
-| `--seed` | `42` | Random seed |
-| `--fps` | `16` | Output video frame rate |
-| `--output` | `helios_output.mp4` | Output file path |
-
-### Stage 2 / Pyramid Parameters
-
-| Parameter | Default | Description |
-|-----------|---------|-------------|
-| `--is-enable-stage2` | off | Enable pyramid multi-stage denoising |
-| `--pyramid-num-stages` | `3` | Number of pyramid stages |
-| `--pyramid-num-inference-steps-list` | `10 10 10` | Steps per pyramid stage |
-| `--is-amplify-first-chunk` | off | DMD amplification (Distilled) |
-
-### CFG-Zero* Parameters
-
-| Parameter | Default | Description |
-|-----------|---------|-------------|
-| `--use-cfg-zero-star` | off | Enable CFG-Zero* guidance (Mid) |
-| `--use-zero-init` | off | Zero init for first steps |
-| `--zero-steps` | `1` | Number of zero-init steps |
-
-### Memory & Parallelism
-
-| Parameter | Default | Description |
-|-----------|---------|-------------|
-| `--vae-use-slicing` | off | Enable VAE slicing |
-| `--vae-use-tiling` | off | Enable VAE tiling |
-| `--enforce-eager` | off | Disable torch.compile |
-| `--enable-cpu-offload` | off | CPU offloading |
-| `--enable-layerwise-offload` | off | Layerwise offloading |
-| `--ulysses-degree` | `1` | Ulysses SP degree |
-| `--ring-degree` | `1` | Ring SP degree |
-| `--cfg-parallel-size` | `1` | CFG parallel size (1 or 2) |
-| `--tensor-parallel-size` | `1` | Tensor parallelism size |
-
-## Example materials
-
-??? abstract "end2end.py"
-    ``````py
-    --8<-- "examples/offline_inference/helios/end2end.py"
-    ``````

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
+```

docs/user_guide/quantization/fp8.md

@@ -85,7 +85,7 @@ warmup_quack_fp8([(14040, 2048, 6144), (14040, 2048, 2048)])
 | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B` | Yes | Yes | All layers | None | |
 | HunyuanImage-3.0 | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | Yes | Yes | All layers; use the Hunyuan stage config for multi-stage runs | None | |
 | HunyuanVideo-1.5 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `720p_t2v`, `480p_i2v` | Yes | Yes | All layers | None | |
-| Cosmos3 | `nvidia/Cosmos3-Nano` | Yes | Not validated | All layers | None | |
+| Cosmos3 | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | Yes | Not validated | All layers | None | |
 
 ### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
 

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/offline_inference/helios/README.md

@@ -1,193 +0,0 @@
-# Helios Video Generation
-
-Helios is a text-to-video (T2V), image-to-video (I2V), and video-to-video (V2V) diffusion model. This example demonstrates end-to-end video generation using vLLM-Omni with three model variants:
-
-| Variant | Description | Key Features |
-|---------|-------------|--------------|
-| **Helios-Base** | Base model, Stage 1 only | Single-stage denoising, `guidance_scale=5.0` |
-| **Helios-Mid** | Mid model, Stage 2 pyramid | Multi-stage pyramid denoising, CFG-Zero* support |
-| **Helios-Distilled** | Distilled model, Stage 2+3 | Few-step inference with DMD, `guidance_scale=1.0` |
-
-## Setup
-
-Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
-
-## Run Examples
-
-Get into the example folder:
-```bash
-cd examples/offline_inference/helios
-```
-
-### Text-to-Video (T2V)
-
-**Helios-Base** (Stage 1 only):
-```bash
-python end2end.py \
-    --model BestWishYsh/Helios-Base \
-    --sample-type t2v \
-    --prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \
-    --guidance-scale 5.0 \
-    --output helios_t2v_base.mp4
-```
-
-**Helios-Mid** (Stage 2 + CFG-Zero*):
-```bash
-python end2end.py \
-    --model BestWishYsh/Helios-Mid \
-    --sample-type t2v \
-    --prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \
-    --guidance-scale 5.0 \
-    --is-enable-stage2 \
-    --pyramid-num-inference-steps-list 20 20 20 \
-    --use-cfg-zero-star --use-zero-init --zero-steps 1 \
-    --output helios_t2v_mid.mp4
-```
-
-**Helios-Distilled** (Stage 2 pyramid + DMD):
-```bash
-python end2end.py \
-    --model BestWishYsh/Helios-Distilled \
-    --sample-type t2v \
-    --prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \
-    --num-frames 240 \
-    --guidance-scale 1.0 \
-    --is-enable-stage2 \
-    --pyramid-num-inference-steps-list 2 2 2 \
-    --is-amplify-first-chunk \
-    --output helios_t2v_distilled.mp4
-```
-
-### Image-to-Video (I2V)
-
-**Helios-Base**:
-```bash
-python end2end.py \
-    --model BestWishYsh/Helios-Base \
-    --sample-type i2v \
-    --image-path /path/to/image.jpg \
-    --prompt "A towering emerald wave surges forward, its crest curling with raw power and energy." \
-    --guidance-scale 5.0 \
-    --output helios_i2v_base.mp4
-```
-
-**Helios-Mid** (Stage 2 + CFG-Zero*):
-```bash
-python end2end.py \
-    --model BestWishYsh/Helios-Mid \
-    --sample-type i2v \
-    --image-path /path/to/image.jpg \
-    --prompt "A towering emerald wave surges forward, its crest curling with raw power and energy." \
-    --guidance-scale 5.0 \
-    --is-enable-stage2 \
-    --pyramid-num-inference-steps-list 20 20 20 \
-    --use-cfg-zero-star --use-zero-init --zero-steps 1 \
-    --output helios_i2v_mid.mp4
-```
-
-**Helios-Distilled**:
-```bash
-python end2end.py \
-    --model BestWishYsh/Helios-Distilled \
-    --sample-type i2v \
-    --image-path /path/to/image.jpg \
-    --prompt "A towering emerald wave surges forward, its crest curling with raw power and energy." \
-    --num-frames 240 \
-    --guidance-scale 1.0 \
-    --is-enable-stage2 \
-    --pyramid-num-inference-steps-list 2 2 2 \
-    --is-amplify-first-chunk \
-    --output helios_i2v_distilled.mp4
-```
-
-### Video-to-Video (V2V)
-
-**Helios-Base**:
-```bash
-python end2end.py \
-    --model BestWishYsh/Helios-Base \
-    --sample-type v2v \
-    --video-path /path/to/video.mp4 \
-    --prompt "A bright yellow Lamborghini speeds along a curving mountain road." \
-    --guidance-scale 5.0 \
-    --output helios_v2v_base.mp4
-```
-
-**Helios-Mid** (Stage 2 + CFG-Zero*):
-```bash
-python end2end.py \
-    --model BestWishYsh/Helios-Mid \
-    --sample-type v2v \
-    --video-path /path/to/video.mp4 \
-    --prompt "A bright yellow Lamborghini speeds along a curving mountain road." \
-    --guidance-scale 5.0 \
-    --is-enable-stage2 \
-    --pyramid-num-inference-steps-list 20 20 20 \
-    --use-cfg-zero-star --use-zero-init --zero-steps 1 \
-    --output helios_v2v_mid.mp4
-```
-
-**Helios-Distilled**:
-```bash
-python end2end.py \
-    --model BestWishYsh/Helios-Distilled \
-    --sample-type v2v \
-    --video-path /path/to/video.mp4 \
-    --prompt "A bright yellow Lamborghini speeds along a curving mountain road." \
-    --num-frames 240 \
-    --guidance-scale 1.0 \
-    --is-enable-stage2 \
-    --pyramid-num-inference-steps-list 2 2 2 \
-    --is-amplify-first-chunk \
-    --output helios_v2v_distilled.mp4
-```
-
-## Common Parameters
-
-| Parameter | Default | Description |
-|-----------|---------|-------------|
-| `--model` | `BestWishYsh/Helios-Base` | Model ID or local path |
-| `--sample-type` | `t2v` | Generation mode: `t2v`, `i2v`, or `v2v` |
-| `--prompt` | — | Text prompt describing the video |
-| `--negative-prompt` | *(see source)* | Negative prompt for CFG (includes anti-static terms) |
-| `--image-path` | — | Input image (required for `i2v`) |
-| `--video-path` | — | Input video (required for `v2v`) |
-| `--height` | `384` | Video height in pixels |
-| `--width` | `640` | Video width in pixels |
-| `--num-frames` | `99` | Number of output frames |
-| `--num-inference-steps` | `50` | Denoising steps (Stage 1 only) |
-| `--guidance-scale` | `5.0` | CFG scale (`1.0` for Distilled) |
-| `--seed` | `42` | Random seed |
-| `--fps` | `16` | Output video frame rate |
-| `--output` | `helios_output.mp4` | Output file path |
-
-### Stage 2 / Pyramid Parameters
-
-| Parameter | Default | Description |
-|-----------|---------|-------------|
-| `--is-enable-stage2` | off | Enable pyramid multi-stage denoising |
-| `--pyramid-num-stages` | `3` | Number of pyramid stages |
-| `--pyramid-num-inference-steps-list` | `10 10 10` | Steps per pyramid stage |
-| `--is-amplify-first-chunk` | off | DMD amplification (Distilled) |
-
-### CFG-Zero* Parameters
-
-| Parameter | Default | Description |
-|-----------|---------|-------------|
-| `--use-cfg-zero-star` | off | Enable CFG-Zero* guidance (Mid) |
-| `--use-zero-init` | off | Zero init for first steps |
-| `--zero-steps` | `1` | Number of zero-init steps |
-
-### Memory & Parallelism
-
-| Parameter | Default | Description |
-|-----------|---------|-------------|
-| `--vae-use-slicing` | off | Enable VAE slicing |
-| `--vae-use-tiling` | off | Enable VAE tiling |
-| `--enforce-eager` | off | Disable torch.compile |
-| `--enable-cpu-offload` | off | CPU offloading |
-| `--enable-layerwise-offload` | off | Layerwise offloading |
-| `--ulysses-degree` | `1` | Ulysses SP degree |
-| `--ring-degree` | `1` | Ring SP degree |
-| `--cfg-parallel-size` | `1` | CFG parallel size (1 or 2) |
-| `--tensor-parallel-size` | `1` | Tensor parallelism size |

examples/offline_inference/helios/end2end.py

@@ -1,355 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""
-Helios video generation example supporting T2V, I2V, and V2V.
-
-Usage (T2V, Helios-Base, Stage 1 only):
-    python end2end.py \
-        --model /path/to/Helios-Base --sample-type t2v \
-        --prompt "A serene lakeside sunrise with mist over the water." \
-        --height 384 --width 640 --num-frames 99 \
-        --num-inference-steps 50 --guidance-scale 5.0
-
-Usage (I2V, Helios-Base):
-    python end2end.py \
-        --model /path/to/Helios-Base --sample-type i2v \
-        --image-path /path/to/image.jpg \
-        --prompt "Description of desired animation." \
-        --guidance-scale 5.0
-
-Usage (V2V, Helios-Base):
-    python end2end.py \
-        --model /path/to/Helios-Base --sample-type v2v \
-        --video-path /path/to/video.mp4 \
-        --prompt "Description of desired transformation." \
-        --guidance-scale 5.0
-
-Usage (Helios-Mid, Stage 2 + CFG-Zero*):
-    python end2end.py \
-        --model /path/to/Helios-Mid --sample-type t2v \
-        --prompt "A serene lakeside sunrise with mist over the water." \
-        --guidance-scale 5.0 \
-        --is-enable-stage2 \
-        --pyramid-num-inference-steps-list 20 20 20 \
-        --use-cfg-zero-star --use-zero-init --zero-steps 1
-
-Usage (Helios-Distilled, Stage 2 pyramid + DMD):
-    python end2end.py \
-        --model /path/to/Helios-Distilled --sample-type t2v \
-        --prompt "A serene lakeside sunrise with mist over the water." \
-        --guidance-scale 1.0 \
-        --is-enable-stage2 \
-        --pyramid-num-inference-steps-list 2 2 2 \
-        --is-amplify-first-chunk
-"""
-
-import argparse
-import time
-from pathlib import Path
-
-import numpy as np
-import torch
-
-from vllm_omni.diffusion.data import DiffusionParallelConfig
-from vllm_omni.entrypoints.omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-from vllm_omni.outputs import OmniRequestOutput
-from vllm_omni.platforms import current_omni_platform
-
-
-def load_image_as_tensor(image_path: str, height: int, width: int) -> torch.Tensor:
-    """Load an image and return a (1, C, H, W) tensor normalized to [-1, 1]."""
-    from PIL import Image
-    from torchvision import transforms
-
-    image = Image.open(image_path).convert("RGB").resize((width, height))
-    transform = transforms.Compose(
-        [
-            transforms.ToTensor(),
-            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
-        ]
-    )
-    return transform(image).unsqueeze(0)
-
-
-def load_video_as_tensor(video_path: str, height: int, width: int) -> torch.Tensor:
-    """Load a video and return a (1, C, T, H, W) tensor normalized to [-1, 1]."""
-    from torchvision import transforms
-    from torchvision.io import read_video
-
-    video_frames, _, _ = read_video(video_path, output_format="TCHW")
-    video_frames = video_frames.float() / 255.0
-
-    transform = transforms.Compose(
-        [
-            transforms.Resize((height, width), antialias=True),
-            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
-        ]
-    )
-    video_frames = torch.stack([transform(f) for f in video_frames])
-    return video_frames.permute(1, 0, 2, 3).unsqueeze(0)
-
-
-def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description="Generate a video with Helios (T2V / I2V / V2V).")
-    parser.add_argument(
-        "--model",
-        default="BestWishYsh/Helios-Base",
-        help="Helios model ID or local path (e.g. Helios-Base, Helios-Mid, Helios-Distilled).",
-    )
-    parser.add_argument(
-        "--sample-type",
-        choices=["t2v", "i2v", "v2v"],
-        default="t2v",
-        help="Generation mode: t2v (text-to-video), i2v (image-to-video), v2v (video-to-video).",
-    )
-    parser.add_argument("--prompt", default="A serene lakeside sunrise with mist over the water.", help="Text prompt.")
-    parser.add_argument(
-        "--negative-prompt",
-        default=(
-            "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, "
-            "images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, "
-            "incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, "
-            "misshapen limbs, fused fingers, still picture, messy background, three legs, many people "
-            "in the background, walking backwards"
-        ),
-        help="Negative prompt.",
-    )
-    parser.add_argument("--seed", type=int, default=42, help="Random seed.")
-    parser.add_argument("--guidance-scale", type=float, default=5.0, help="CFG scale.")
-    parser.add_argument("--height", type=int, default=384, help="Video height.")
-    parser.add_argument("--width", type=int, default=640, help="Video width.")
-    parser.add_argument("--num-frames", type=int, default=99, help="Number of video frames.")
-    parser.add_argument("--num-inference-steps", type=int, default=50, help="Sampling steps (Stage 1 only).")
-    parser.add_argument("--output", type=str, default="helios_output.mp4", help="Output video path.")
-    parser.add_argument("--fps", type=int, default=16, help="Frames per second for the output video.")
-
-    # I2V / V2V
-    parser.add_argument("--image-path", type=str, default=None, help="Input image path for I2V mode.")
-    parser.add_argument("--video-path", type=str, default=None, help="Input video path for V2V mode.")
-
-    # Stage 2 (pyramid multi-stage denoising)
-    parser.add_argument(
-        "--is-enable-stage2",
-        action="store_true",
-        help="Enable pyramid multi-stage denoising (Stage 2). Required for Helios-Distilled.",
-    )
-    parser.add_argument(
-        "--pyramid-num-stages",
-        type=int,
-        default=3,
-        help="Number of pyramid stages for Stage 2.",
-    )
-    parser.add_argument(
-        "--pyramid-num-inference-steps-list",
-        type=int,
-        nargs="+",
-        default=[10, 10, 10],
-        help="Inference steps per pyramid stage.",
-    )
-
-    # DMD
-    parser.add_argument(
-        "--is-amplify-first-chunk",
-        action="store_true",
-        help="Enable DMD amplification for the first chunk (Helios-Distilled).",
-    )
-
-    # CFG Zero Star
-    parser.add_argument(
-        "--use-cfg-zero-star",
-        action="store_true",
-        help="Enable CFG Zero Star guidance (recommended for Helios-Mid).",
-    )
-    parser.add_argument(
-        "--use-zero-init",
-        action="store_true",
-        help="Use zero initialization for the first denoising steps with CFG-Zero*.",
-    )
-    parser.add_argument(
-        "--zero-steps",
-        type=int,
-        default=1,
-        help="Number of initial denoising steps using zero prediction (default: 1).",
-    )
-    parser.add_argument(
-        "--enable-diffusion-pipeline-profiler",
-        action="store_true",
-        help="Enable diffusion pipeline profiler to display stage durations.",
-    )
-    parser.add_argument(
-        "--quantization",
-        type=str,
-        default=None,
-        help="Quantization method (e.g. 'fp8').",
-    )
-
-    # Memory & parallelism
-    parser.add_argument("--vae-use-slicing", action="store_true", help="Enable VAE slicing.")
-    parser.add_argument("--vae-use-tiling", action="store_true", help="Enable VAE tiling.")
-    parser.add_argument("--enforce-eager", action="store_true", help="Disable torch.compile.")
-    parser.add_argument("--enable-cpu-offload", action="store_true", help="Enable CPU offloading.")
-    parser.add_argument("--enable-layerwise-offload", action="store_true", help="Enable layerwise offloading.")
-    parser.add_argument("--ulysses-degree", type=int, default=1, help="Ulysses SP degree.")
-    parser.add_argument("--ring-degree", type=int, default=1, help="Ring SP degree.")
-    parser.add_argument("--cfg-parallel-size", type=int, default=1, choices=[1, 2], help="CFG parallel size.")
-    parser.add_argument("--tensor-parallel-size", type=int, default=1, help="Tensor parallelism size.")
-
-    return parser.parse_args()
-
-
-def main():
-    args = parse_args()
-    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
-
-    parallel_config = DiffusionParallelConfig(
-        ulysses_degree=args.ulysses_degree,
-        ring_degree=args.ring_degree,
-        cfg_parallel_size=args.cfg_parallel_size,
-        tensor_parallel_size=args.tensor_parallel_size,
-    )
-
-    omni = Omni(
-        model=args.model,
-        enable_layerwise_offload=args.enable_layerwise_offload,
-        vae_use_slicing=args.vae_use_slicing,
-        vae_use_tiling=args.vae_use_tiling,
-        enable_cpu_offload=args.enable_cpu_offload,
-        parallel_config=parallel_config,
-        enforce_eager=args.enforce_eager,
-        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
-        quantization=args.quantization,
-    )
-
-    # Validate I2V / V2V arguments
-    if args.sample_type == "i2v" and not args.image_path:
-        raise ValueError("--image-path is required for I2V mode.")
-    if args.sample_type == "v2v" and not args.video_path:
-        raise ValueError("--video-path is required for V2V mode.")
-
-    # Build extra_args for Helios-specific parameters
-    extra_args = {}
-    if args.is_enable_stage2:
-        extra_args["is_enable_stage2"] = True
-        extra_args["pyramid_num_stages"] = args.pyramid_num_stages
-        extra_args["pyramid_num_inference_steps_list"] = args.pyramid_num_inference_steps_list
-    if args.is_amplify_first_chunk:
-        extra_args["is_amplify_first_chunk"] = True
-    if args.use_cfg_zero_star:
-        extra_args["use_cfg_zero_star"] = True
-    if args.use_zero_init:
-        extra_args["use_zero_init"] = True
-        extra_args["zero_steps"] = args.zero_steps
-
-    if args.sample_type == "i2v":
-        image_tensor = load_image_as_tensor(args.image_path, args.height, args.width)
-        extra_args["image"] = image_tensor
-    elif args.sample_type == "v2v":
-        video_tensor = load_video_as_tensor(args.video_path, args.height, args.width)
-        extra_args["video"] = video_tensor
-
-    # Print generation configuration
-    print(f"\n{'=' * 60}")
-    print("Helios Generation Configuration:")
-    print(f"  Model: {args.model}")
-    print(f"  Sample type: {args.sample_type.upper()}")
-    if args.sample_type == "i2v":
-        print(f"  Image: {args.image_path}")
-    elif args.sample_type == "v2v":
-        print(f"  Video: {args.video_path}")
-    print(f"  Prompt: {args.prompt[:80]}{'...' if len(args.prompt) > 80 else ''}")
-    print(f"  Video size: {args.width}x{args.height}, {args.num_frames} frames")
-    print(f"  Inference steps: {args.num_inference_steps}")
-    print(f"  Guidance scale: {args.guidance_scale}")
-    if args.is_enable_stage2:
-        print(f"  Stage 2: enabled (stages={args.pyramid_num_stages}, steps={args.pyramid_num_inference_steps_list})")
-        if args.is_amplify_first_chunk:
-            print("  DMD amplify first chunk: enabled")
-        if args.use_cfg_zero_star:
-            print(f"  CFG Zero Star: enabled (zero_init={args.use_zero_init}, zero_steps={args.zero_steps})")
-    else:
-        if args.use_cfg_zero_star:
-            print(f"  CFG Zero Star: enabled (zero_init={args.use_zero_init}, zero_steps={args.zero_steps})")
-        print("  Stage 2: disabled (Stage 1 only)")
-    print(f"{'=' * 60}\n")
-
-    generation_start = time.perf_counter()
-    frames = omni.generate(
-        {
-            "prompt": args.prompt,
-            "negative_prompt": args.negative_prompt,
-        },
-        OmniDiffusionSamplingParams(
-            height=args.height,
-            width=args.width,
-            generator=generator,
-            guidance_scale=args.guidance_scale,
-            num_inference_steps=args.num_inference_steps,
-            num_frames=args.num_frames,
-            extra_args=extra_args,
-        ),
-    )
-    generation_end = time.perf_counter()
-    generation_time = generation_end - generation_start
-
-    print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)")
-
-    # Extract video frames from OmniRequestOutput
-    if isinstance(frames, list) and len(frames) > 0:
-        first_item = frames[0]
-
-        if hasattr(first_item, "final_output_type"):
-            if first_item.final_output_type != "image":
-                raise ValueError(
-                    f"Unexpected output type '{first_item.final_output_type}', expected 'image' for video generation."
-                )
-
-            if hasattr(first_item, "is_pipeline_output") and first_item.is_pipeline_output:
-                inner_output = first_item.request_output
-                if isinstance(inner_output, OmniRequestOutput) and hasattr(inner_output, "images"):
-                    frames = inner_output.images[0] if inner_output.images else None
-                    if frames is None:
-                        raise ValueError("No video frames found in output.")
-            elif hasattr(first_item, "images") and first_item.images:
-                frames = first_item.images
-            else:
-                raise ValueError("No video frames found in OmniRequestOutput.")
-
-    # Unwrap batch list from postprocess_video: [numpy(T,H,W,C)] -> numpy(T,H,W,C)
-    if isinstance(frames, list) and len(frames) > 0 and isinstance(frames[0], np.ndarray):
-        frames = frames[0]
-
-    output_path = Path(args.output)
-    output_path.parent.mkdir(parents=True, exist_ok=True)
-    try:
-        from diffusers.utils import export_to_video
-    except ImportError:
-        raise ImportError("diffusers is required for export_to_video.")
-
-    if isinstance(frames, torch.Tensor):
-        video_tensor = frames.detach().cpu()
-        if video_tensor.dim() == 5:
-            if video_tensor.shape[1] in (3, 4):
-                video_tensor = video_tensor[0].permute(1, 2, 3, 0)
-            else:
-                video_tensor = video_tensor[0]
-        elif video_tensor.dim() == 4 and video_tensor.shape[0] in (3, 4):
-            video_tensor = video_tensor.permute(1, 2, 3, 0)
-        if video_tensor.is_floating_point():
-            video_tensor = video_tensor.clamp(-1, 1) * 0.5 + 0.5
-        video_array = video_tensor.float().numpy()
-    elif isinstance(frames, np.ndarray):
-        video_array = frames
-    else:
-        video_array = frames
-
-    if isinstance(video_array, np.ndarray) and video_array.ndim == 5:
-        video_array = video_array[0]
-
-    export_to_video(video_array, str(output_path), fps=args.fps)
-    print(f"Saved generated video to {output_path}")
-
-
-if __name__ == "__main__":
-    main()

examples/offline_inference/magi_human/README.md

@@ -1,72 +0,0 @@
-# MagiHuman Generation
-
-MagiHuman is an advanced, omni-modality model that generates both high-quality video and lip-synced audio from a text prompt.
-
-Because MagiHuman is a very large model featuring a powerful DiT MoE backbone and a ~9B parameter T5Gemma text encoder, it natively supports **Tensor Parallelism (TP)** in vLLM-Omni to run efficiently across multi-GPU setups, reducing device memory bottlenecks.
-
-## Setup
-
-### Install MagiCompiler (recommended)
-
-MagiHuman relies on [MagiCompiler](https://github.com/SandAI-org/MagiCompiler) for custom-op registration used by the DiT attention kernels. While the pipeline can fall back to stub implementations, installing MagiCompiler is **strongly recommended** for correct behaviour.
-
-```bash
-# Clone the repo
-git clone https://github.com/SandAI-org/MagiCompiler.git
-cd MagiCompiler
-
-# System dependencies (optional, for FX graph visualization; Debian/Ubuntu)
-sudo apt update && sudo apt install -y graphviz
-
-# Python dependencies
-pip install -r requirements.txt
-
-# Install MagiCompiler
-pip install .          # end users (recommended)
-# pip install -e .     # developers (editable install)
-```
-
-### Hardware requirements
-
-Ensure your hardware has enough VRAM. For a standard node with 80GB GPUs, running with `--tensor-parallel-size 4` is recommended to shard both the MoE weights and the T5Gemma text encoder across 4 GPUs, reducing the per-GPU peak VRAM overhead significantly (by roughly ~13.5GB per GPU compared to single-device inference).
-
-Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) for further details on allocating memory.
-
-## Run Examples
-
-Get into the example folder:
-```bash
-cd examples/offline_inference/magi_human
-```
-
-### End-to-End Generation (Text to Video+Audio)
-
-Generate a video with synchronized speech natively generated by the model.
-
-```bash
-python end2end.py \
-    --model /proj-tango-pvc/users/zhipeng.wang/workspace/models/daVinci-MagiHuman \
-    --prompt "A young woman with long, wavy golden blonde hair..." \
-    --tensor-parallel-size 4 \
-    --output output_magihuman.mp4
-```
-
-## Common Parameters
-
-| Parameter | Default | Description |
-|-----------|---------|-------------|
-| `--model` | *(Required)* | Local model path or HuggingFace ID |
-| `--prompt` | *(built-in demo prompt)* | Highly detailed text prompt dictating visual look and dialogue text |
-| `--tensor-parallel-size` | `4` | Tensor parallelism size (Number of GPUs) |
-| `--height` | `256` | Initial resolution height |
-| `--width` | `448` | Initial resolution width |
-| `--num-inference-steps` | `8` | Denoising steps |
-| `--seed` | `52` | Random seed |
-| `--output` | `output_magihuman.mp4` | Output video with audio path |
-
-## Example materials
-
-??? abstract "end2end.py"
-    ``````py
-    --8<-- "examples/offline_inference/magi_human/end2end.py"
-    ``````

examples/offline_inference/magi_human/end2end.py

@@ -1,122 +0,0 @@
-import argparse
-
-from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
-from vllm_omni.entrypoints.omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-
-def parse_args():
-    parser = argparse.ArgumentParser(description="End-to-end inference script for MagiHuman.")
-    parser.add_argument("--model", type=str, required=True, help="Path or ID of the MagiHuman model.")
-    parser.add_argument(
-        "--prompt",
-        type=str,
-        default="",
-        help="Text prompt containing visual description, dialogue, and background sound.",
-    )
-    parser.add_argument(
-        "--tensor-parallel-size", "-tp", type=int, default=4, help="Tensor parallel size (number of GPUs)."
-    )
-    parser.add_argument(
-        "--output", type=str, default="output_magihuman.mp4", help="Path to save the generated mp4 file."
-    )
-    parser.add_argument("--height", type=int, default=256, help="Video height.")
-    parser.add_argument("--width", type=int, default=448, help="Video width.")
-    parser.add_argument("--num-inference-steps", type=int, default=8, help="Number of denoising steps.")
-    parser.add_argument("--seed", type=int, default=52, help="Random seed for generation.")
-    return parser.parse_args()
-
-
-def main():
-    args = parse_args()
-
-    print(f"Initializing MagiHuman pipeline with TP={args.tensor_parallel_size}...")
-    omni = Omni(
-        model=args.model,
-        init_timeout=1200,
-        tensor_parallel_size=args.tensor_parallel_size,
-        devices=list(range(args.tensor_parallel_size)),
-    )
-
-    prompt = args.prompt
-    if not prompt:
-        prompt = (
-            "A young woman with long, wavy golden blonde hair and bright blue eyes, "
-            "wearing a fitted ivory silk blouse with a delicate lace collar, sits "
-            "stationary in front of a softly lit, blurred warm-toned interior. Her "
-            "overall disposition is warm, composed, and gently confident. The camera "
-            "holds a static medium close-up, framing her from the shoulders up, "
-            "with shallow depth of field keeping her face in sharp focus. Soft "
-            "directional key light falls from the upper left, casting a gentle "
-            "highlight along her cheekbone and nose bridge. She draws a quiet breath, "
-            "the levator labii superiors relaxing as her lips part. She speaks in "
-            "clear, warm, unhurried American English: "
-            "\"The most beautiful things in life aren't things at all — "
-            "they're moments, feelings, and the people who make you feel truly alive.\" "
-            "Her jaw descends smoothly on each stressed syllable; the orbicularis oris "
-            "shapes each vowel with precision. A faint, genuine smile engages the "
-            "zygomaticus major, lifting her lip corners fractionally. Her brows rest "
-            "in a soft, neutral arch throughout. She maintains steady, forward-facing "
-            "eye contact. Head position remains level; no torso displacement occurs.\n\n"
-            "Dialogue:\n"
-            "<Young blonde woman, American English>: "
-            "\"The most beautiful things in life aren't things at all — "
-            "they're moments, feelings, and the people who make you feel truly alive.\"\n\n"
-            "Background Sound:\n"
-            "<Soft, warm indoor ambience with a faint distant piano melody>"
-        )
-
-    sampling_params = OmniDiffusionSamplingParams(
-        height=args.height,
-        width=args.width,
-        num_inference_steps=args.num_inference_steps,
-        seed=args.seed,
-        extra_args={
-            "seconds": 5,
-            "sr_height": 1080,
-            "sr_width": 1920,
-            "sr_num_inference_steps": 5,
-        },
-    )
-
-    print(f"Generating with prompt: {prompt[:80]}...")
-    outputs = omni.generate(
-        prompts=[prompt],
-        sampling_params_list=[sampling_params],
-    )
-
-    print(f"Generation complete. Output type: {type(outputs)}")
-    if outputs:
-        first = outputs[0]
-
-        if hasattr(first, "images") and first.images:
-            video_frames = first.images[0]
-            print(f"Video frames: shape={video_frames.shape}, dtype={video_frames.dtype}")
-
-            audio_waveform = None
-            mm = first.multimodal_output or {}
-            if mm:
-                audio_waveform = mm.get("audio")
-                if audio_waveform is not None:
-                    print(f"Audio waveform: shape={audio_waveform.shape}, dtype={audio_waveform.dtype}")
-
-            output_fps = float(mm.get("fps", 25))
-            output_sr = int(mm.get("audio_sample_rate", 24000))
-            print(f"Using fps={output_fps}, audio_sample_rate={output_sr} from model output")
-
-            video_bytes = mux_video_audio_bytes(
-                video_frames,
-                audio_waveform,
-                fps=output_fps,
-                audio_sample_rate=output_sr,
-            )
-            with open(args.output, "wb") as f:
-                f.write(video_bytes)
-            print(f"Saved MP4 ({len(video_bytes)} bytes) to {args.output}")
-        print("SUCCESS: MagiHuman pipeline generation completed.")
-    else:
-        print("WARNING: No outputs returned.")
-
-
-if __name__ == "__main__":
-    main()

examples/offline_inference/text_to_speech/README.md

@@ -369,6 +369,20 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \
 ```
 Streaming requires `async_chunk: true` in the stage config.
 
+### Word Timestamps
+Generate a WAV offline and a JSON sidecar with word-level timestamps from
+`Qwen/Qwen3-ForcedAligner-0.6B`:
+```bash
+python examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py \
+    --model Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
+    --forced-aligner Qwen/Qwen3-ForcedAligner-0.6B \
+    --text "Hello world." \
+    --output-dir /tmp/qwen3_tts_timestamps
+```
+The script writes `qwen3_tts_word_timestamps.wav` and
+`qwen3_tts_word_timestamps.json`. On machines without a local CUDA toolkit,
+set `VLLM_USE_FLASHINFER_SAMPLER=0` to avoid FlashInfer sampler JIT.
+
 ### Batched decoding
 The Code2Wav stage supports batched decoding through the SpeechTokenizer. Pass multiple prompts via `--txt-prompts` and set `--batch-size` accordingly. To raise `max_num_seqs` on either stage, point `--stage-configs-path` at a stage configs YAML with the desired values (see `vllm_omni/model_executor/stage_configs/` for templates):
 ```bash

examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py

@@ -0,0 +1,172 @@
+"""Offline Qwen3-TTS example with word-level timestamps.
+
+This demo runs Qwen3-TTS offline, aligns the synthesized audio in memory with
+the same shared forced aligner used by the streaming server path, then writes
+the WAV and JSON sidecar:
+
+    python examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py \
+        --model Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
+        --forced-aligner Qwen/Qwen3-ForcedAligner-0.6B
+
+On machines without a local CUDA toolkit, set ``VLLM_USE_FLASHINFER_SAMPLER=0``
+to avoid FlashInfer sampler JIT during warmup.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import os
+from pathlib import Path
+from typing import Any
+
+os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+# The in-process offline aligner consumes token_classify outputs. Keeping V1
+# multiprocessing off avoids msgspec serialization issues for token-wise logits.
+os.environ.setdefault("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
+
+import numpy as np
+import soundfile as sf
+import torch
+from end2end import _estimate_prompt_len
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+from vllm_omni import Omni
+from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
+from vllm_omni.utils.forced_aligner import align, build_forced_aligner_config
+
+
+def _default_stage_config() -> str:
+    repo_root = Path(__file__).resolve().parents[4]
+    return str(repo_root / "vllm_omni" / "deploy" / "qwen3_tts.yaml")
+
+
+def _build_custom_voice_input(args: Any) -> dict[str, Any]:
+    additional_information = {
+        "task_type": ["CustomVoice"],
+        "text": [args.text],
+        "language": [args.language],
+        "speaker": [args.speaker],
+        "instruct": [args.instructions],
+        "max_new_tokens": [args.max_new_tokens],
+    }
+    return {
+        "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, args.model),
+        "additional_information": additional_information,
+    }
+
+
+def _audio_tensor_and_sample_rate(mm: dict[str, Any]) -> tuple[torch.Tensor, int]:
+    audio_data = mm["audio"]
+    sr_raw = mm["sr"]
+    sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw
+    sample_rate = int(sr_val.item()) if hasattr(sr_val, "item") else int(sr_val)
+    audio_tensor = torch.cat(audio_data, dim=-1) if isinstance(audio_data, list) else audio_data
+    return audio_tensor.float().detach().cpu().flatten(), sample_rate
+
+
+def _float_audio_to_pcm16_bytes(audio: torch.Tensor) -> bytes:
+    samples = audio.numpy()
+    pcm = (np.clip(samples, -1.0, 1.0) * 32767.0).astype("<i2")
+    return pcm.tobytes()
+
+
+async def _run_alignment(args: Any, audio: torch.Tensor, sample_rate: int) -> list[dict[str, Any]] | None:
+    aligner_config = build_forced_aligner_config(args)
+    if aligner_config is None:
+        raise ValueError("--forced-aligner or --forced-aligner-config is required")
+    timestamps = await align(
+        audio=_float_audio_to_pcm16_bytes(audio),
+        text=args.text,
+        sample_rate=sample_rate,
+        config=aligner_config,
+        language=args.language,
+    )
+    if timestamps is None:
+        return None
+    return [
+        {
+            "word": item.word,
+            "start_ms": item.start_ms,
+            "end_ms": item.end_ms,
+        }
+        for item in timestamps
+    ]
+
+
+def main(args: Any) -> None:
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    omni_kwargs = vars(args).copy()
+    for key in (
+        "forced_aligner",
+        "text",
+        "language",
+        "speaker",
+        "instructions",
+        "max_new_tokens",
+        "output_dir",
+    ):
+        omni_kwargs.pop(key, None)
+    omni_kwargs["stage_configs_path"] = args.stage_configs_path or _default_stage_config()
+    omni_kwargs["log_stats"] = args.log_stats
+    omni = Omni(**omni_kwargs)
+
+    prompt = _build_custom_voice_input(args)
+    final_output = None
+    for stage_outputs in omni.generate([prompt]):
+        final_output = stage_outputs.request_output
+    if final_output is None:
+        raise RuntimeError("Qwen3-TTS did not produce an output.")
+
+    mm = final_output.outputs[0].multimodal_output
+    audio, sample_rate = _audio_tensor_and_sample_rate(mm)
+
+    wav_path = output_dir / "qwen3_tts_word_timestamps.wav"
+    # Align the in-memory tensor, not the saved WAV. The WAV is only a demo
+    # artifact so users can listen to the same audio referenced by the sidecar.
+    timestamps = asyncio.run(_run_alignment(args, audio, sample_rate))
+    sf.write(wav_path, audio.numpy(), samplerate=sample_rate, format="WAV")
+
+    sidecar = {
+        "text": args.text,
+        "sample_rate": sample_rate,
+        "audio_path": str(wav_path),
+        "timestamps": timestamps,
+    }
+    json_path = output_dir / "qwen3_tts_word_timestamps.json"
+    json_path.write_text(json.dumps(sidecar, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
+
+    print(f"Saved audio: {wav_path}")
+    print(f"Saved timestamps: {json_path}")
+    print(json.dumps(sidecar["timestamps"], ensure_ascii=False, indent=2))
+
+
+def parse_args() -> Any:
+    parser = FlexibleArgumentParser(description="Offline Qwen3-TTS word timestamp example")
+    parser.add_argument("--model", default="Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice", help="Qwen3-TTS model path/name")
+    parser.add_argument(
+        "--forced-aligner",
+        default=None,
+        help="Qwen3 forced aligner model path/name",
+    )
+    parser.add_argument(
+        "--forced-aligner-config",
+        default=None,
+        help="Optional YAML file for forced aligner settings (incl. gpu_memory_utilization)",
+    )
+    parser.add_argument("--stage-configs-path", default=None, help="Qwen3-TTS deploy YAML")
+    parser.add_argument("--text", default="Hello world.", help="Text to synthesize and align")
+    parser.add_argument("--language", default="English", help="Qwen3-TTS language field")
+    parser.add_argument("--speaker", default="Vivian", help="CustomVoice speaker name")
+    parser.add_argument("--instructions", default="", help="Optional speaking style instruction")
+    parser.add_argument("--max-new-tokens", type=int, default=2048, help="TTS max_new_tokens")
+    parser.add_argument("--output-dir", default="output_audio", help="Directory for WAV and JSON sidecar")
+    parser.add_argument("--log-stats", action="store_true", default=False, help="Enable vLLM-Omni stats logging")
+    nullify_stage_engine_defaults(parser)
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    main(parse_args())

examples/offline_inference/text_to_video/text_to_video.md

@@ -9,6 +9,7 @@ A unified script for text-to-video generation. Supports multiple models with mod
 | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 720x1280 | 81 | 40 | 4.0 | ~60 GiB |
 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v` | 480x832 | 121 | 50 | 6.0 | 1×A100 80GB |
 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | 720x1280 | 121 | 50 | 6.0 | FP8 + VAE tiling required |
+| `BestWishYsh/Helios-Base` / `Helios-Mid` / `Helios-Distilled` | 384x640 | 99 | 50 | 5.0 / 5.0 / 1.0 | — |
 
 ## Local CLI Usage
 
@@ -100,6 +101,49 @@ python text_to_video.py \
   --output quick_test.mp4
 ```
 
+### Helios (T2V)
+
+Helios ships three variants. Model-specific knobs (declared in
+`vllm_omni/model_extras/helios.py`) are passed via the generic `--extra-body`
+JSON flag rather than bespoke per-model flags.
+
+**Helios-Base** (Stage 1 only):
+
+```bash
+python text_to_video.py \
+  --model BestWishYsh/Helios-Base \
+  --prompt "A dynamic time-lapse of scenery rushing past the window of a speeding train." \
+  --guidance-scale 5.0 \
+  --output helios_t2v_base.mp4
+```
+
+**Helios-Mid** (Stage 2 pyramid + CFG-Zero*):
+
+```bash
+python text_to_video.py \
+  --model BestWishYsh/Helios-Mid \
+  --prompt "A dynamic time-lapse of scenery rushing past the window of a speeding train." \
+  --guidance-scale 5.0 \
+  --extra-body '{"is_enable_stage2": true, "pyramid_num_inference_steps_list": [20, 20, 20], "use_cfg_zero_star": true, "use_zero_init": true, "zero_steps": 1}' \
+  --output helios_t2v_mid.mp4
+```
+
+**Helios-Distilled** (Stage 2 pyramid + DMD, few-step):
+
+```bash
+python text_to_video.py \
+  --model BestWishYsh/Helios-Distilled \
+  --prompt "A dynamic time-lapse of scenery rushing past the window of a speeding train." \
+  --num-frames 240 \
+  --guidance-scale 1.0 \
+  --extra-body '{"is_enable_stage2": true, "pyramid_num_inference_steps_list": [2, 2, 2], "is_amplify_first_chunk": true}' \
+  --output helios_t2v_distilled.mp4
+```
+
+> Helios image-to-video (I2V) and video-to-video (V2V) require image/video
+> conditioning tensors that cannot be passed through the JSON `--extra-body`
+> flag; they are out of scope for this text-to-video example.
+
 ## Key Arguments
 
 ### Common
@@ -112,6 +156,7 @@ python text_to_video.py \
 - `--num-inference-steps`: sampling steps. Default depends on model.
 - `--fps`: frames per second for the saved MP4.
 - `--output`: path to save the generated video.
+- `--extra-body`: JSON dict of model-specific knobs (declared in `vllm_omni/model_extras/`), merged into sampling `extra_args`. See the Helios recipes above.
 - `--vae-use-slicing`: enable VAE slicing for memory optimization.
 - `--vae-use-tiling`: enable VAE tiling for memory optimization.
 - `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](../../../docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).

examples/offline_inference/text_to_video/text_to_video.py

@@ -35,6 +35,15 @@ _MODEL_PRESETS = {
         "fps": 24,
         "output": "hunyuan_video_15_output.mp4",
     },
+    "helios": {
+        "height": 384,
+        "width": 640,
+        "num_frames": 99,
+        "num_inference_steps": 50,
+        "guidance_scale": 5.0,
+        "fps": 16,
+        "output": "helios_output.mp4",
+    },
 }
 
 
@@ -42,6 +51,8 @@ def _detect_preset(model: str) -> dict:
     model_lower = model.lower()
     if "hunyuan" in model_lower:
         return _MODEL_PRESETS["hunyuan"]
+    if "helios" in model_lower:
+        return _MODEL_PRESETS["helios"]
     return _MODEL_PRESETS["wan"]
 
 
@@ -55,10 +66,26 @@ def parse_profiler_config(value: str) -> dict[str, Any]:
     return config
 
 
+def parse_extra_body(value: str) -> dict[str, Any]:
+    """Parse a JSON object of model-specific extra_body params.
+
+    Pipeline-declared knobs (see vllm_omni/model_extras/) are merged into
+    OmniDiffusionSamplingParams.extra_args, so a single generic example can
+    drive model-specific behaviour without bespoke per-model flags.
+    """
+    try:
+        body = json.loads(value)
+    except json.JSONDecodeError as e:
+        raise argparse.ArgumentTypeError(f"--extra-body must be valid JSON: {e}") from e
+    if not isinstance(body, dict):
+        raise argparse.ArgumentTypeError("--extra-body must be a JSON object")
+    return body
+
+
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(
         description="Generate a video from a text prompt. "
-        "Supports Wan2.2, HunyuanVideo-1.5, and other text-to-video models."
+        "Supports Wan2.2, HunyuanVideo-1.5, Helios, and other text-to-video models."
     )
     parser.add_argument(
         "--model",
@@ -74,6 +101,14 @@ def parse_args() -> argparse.Namespace:
     )
     parser.add_argument("--prompt", default="A serene lakeside sunrise with mist over the water.", help="Text prompt.")
     parser.add_argument("--negative-prompt", default="", help="Negative prompt.")
+    parser.add_argument(
+        "--extra-body",
+        type=parse_extra_body,
+        default=None,
+        help="JSON dict of model-specific extra_body params (declared in vllm_omni/model_extras/), "
+        "merged into sampling extra_args. Example (Helios-Distilled): "
+        '\'{"is_enable_stage2": true, "pyramid_num_inference_steps_list": [2, 2, 2], "is_amplify_first_chunk": true}\'.',
+    )
     parser.add_argument("--seed", type=int, default=42, help="Random seed.")
     parser.add_argument("--guidance-scale", type=float, default=None, help="CFG scale. Default: model-specific.")
     parser.add_argument(
@@ -357,6 +392,9 @@ def main():
     )
     if args.guidance_scale_high is not None:
         sampling_kwargs["guidance_scale_2"] = args.guidance_scale_high
+    if args.extra_body:
+        # Model-specific knobs (declared in vllm_omni/model_extras/) routed via extra_args.
+        sampling_kwargs["extra_args"] = dict(args.extra_body)
 
     generation_start = time.perf_counter()
     frames = omni.generate(

examples/offline_inference/x_to_video_audio/x_to_video_audio.md

@@ -87,6 +87,63 @@ python x_to_video_audio.py \
 ```
 
 
+## MagiHuman (`--model-type magi-human`)
+
+MagiHuman is a text → video+audio model with a DiT MoE backbone and a ~9B-param
+T5Gemma text encoder. A detailed text prompt is the only required input; an optional
+image and/or audio file may be supplied for conditioning. Natively supports Tensor
+Parallelism. For an 80GB node, `--tensor-parallel-size 4` is recommended to shard
+the MoE weights and the text encoder.
+
+> Install [MagiCompiler](https://github.com/SandAI-org/MagiCompiler) for correct
+> attention-kernel behaviour (the pipeline otherwise falls back to stubs).
+
+Text-only generation:
+
+```bash
+python x_to_video_audio.py \
+  --model-type magi-human \
+  --model /path/to/daVinci-MagiHuman \
+  --prompt "A young woman with long, wavy golden blonde hair... <dialogue and background sound>" \
+  --tensor-parallel-size 4 \
+  --height 256 --width 448 \
+  --num-inference-steps 8 \
+  --seed 52 \
+  --extra-body '{"seconds": 5, "sr_height": 1080, "sr_width": 1920, "sr_num_inference_steps": 5}' \
+  --output output_magihuman.mp4
+```
+
+With optional image and audio conditioning:
+
+```bash
+python x_to_video_audio.py \
+  --model-type magi-human \
+  --model /path/to/daVinci-MagiHuman \
+  --prompt "A young woman..." \
+  --tensor-parallel-size 4 \
+  --height 256 --width 448 \
+  --num-inference-steps 8 \
+  --seed 52 \
+  --extra-body '{"image_path": "/path/to/ref.jpg", "audio_path": "/path/to/ref.wav", "sr_height": 1080, "sr_width": 1920, "sr_num_inference_steps": 5}' \
+  --output output_magihuman.mp4
+```
+
+MagiHuman-specific arguments are passed as a JSON dict via `--extra-body` (declared
+in `vllm_omni/model_extras/magi_human.py`, routed via `extra_args`):
+
+- `seconds`: output duration in seconds (default 10; ignored when `audio_path` is set,
+  because the audio length then determines the number of frames).
+- `audio_path`: path to an audio file for audio-to-video conditioning. When provided,
+  the audio drives the frame count and is encoded as a latent condition into the DiT
+  (`is_a2v` mode). Omit for pure text-to-video+audio generation.
+- `image_path`: path to an image file for visual conditioning. Applied at both the
+  base-resolution (BR) and super-resolution (SR) stages when SR is enabled.
+- `sr_height` / `sr_width`: super-resolution output resolution. SR stage is skipped
+  when these are omitted.
+- `sr_num_inference_steps`: denoising steps for the SR stage.
+
+## DreamID-Omni arguments
+
 Key arguments:
 - `--prompt`: text description (string).
 - `--model`: path to the model local directory.

examples/offline_inference/x_to_video_audio/x_to_video_audio.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import argparse
+import json
 import re
 import time
 from typing import Any
@@ -20,14 +21,26 @@ AUDIO_EXTENSIONS = {".wav", ".mp3", ".flac", ".m4a", ".aac", ".ogg"}
 
 
 def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description="Offline inference for DreamID-Omni (video + audio).")
-    parser.add_argument("--model", required=True, help="DreamID ckpt root directory.")
-    parser.add_argument("--model-type", default="dreamid-omni", help="Model type.")
+    parser = argparse.ArgumentParser(description="Offline inference for X -> video+audio models.")
+    parser.add_argument("--model", required=True, help="Model ckpt root directory.")
+    parser.add_argument(
+        "--model-type",
+        default="dreamid-omni",
+        choices=["dreamid-omni", "magi-human"],
+        help="Model type.",
+    )
     parser.add_argument("--prompt", default=None, help="Text prompt.")
 
     parser.add_argument("--image-path", type=str, nargs="+", help="list of image-path")
     parser.add_argument("--audio-path", type=str, nargs="+", help="list of audio-path")
     parser.add_argument("--prompt-file", type=str, default=None, help="Text prompt in json format.")
+    parser.add_argument(
+        "--extra-body",
+        type=str,
+        default=None,
+        help="[magi-human] JSON dict of model-specific extra params (declared in vllm_omni/model_extras/), "
+        'merged into sampling extra_args. Example: \'{"image_path": "/path/to/img.jpg", "seconds": 5}\'.',
+    )
 
     parser.add_argument("--height", type=int, default=704, help="Video height.")
     parser.add_argument("--width", type=int, default=1280, help="Video width.")
@@ -142,8 +155,6 @@ def main() -> None:
 
     text_prompt = args.prompt
     if args.prompt_file:
-        import json
-
         with open(args.prompt_file) as f:
             text_prompt = json.load(f)
             text_prompt = re.sub(
@@ -155,25 +166,36 @@ def main() -> None:
             text_prompt = re.sub(r"\[[A-Z_]+\]", "", text_prompt)
             text_prompt = re.sub(r"\n\s*\n", "\n", text_prompt).strip()
 
-    image, audio = load_image_and_audio(args.image_path, args.audio_path)
+    if args.model_type == "magi-human":
+        prompt = text_prompt
+        extra_args = json.loads(args.extra_body) if args.extra_body else {}
+        sampling_params = OmniDiffusionSamplingParams(
+            height=args.height,
+            width=args.width,
+            num_inference_steps=args.num_inference_steps,
+            seed=args.seed,
+            extra_args=extra_args,
+        )
+    else:
+        image, audio = load_image_and_audio(args.image_path, args.audio_path)
 
-    prompt = {
-        "prompt": text_prompt,
-        "video_negative_prompt": args.video_negative_prompt,
-        "audio_negative_prompt": args.audio_negative_prompt,
-        "multi_modal_data": {"image": image, "audio": audio},
-    }
+        prompt = {
+            "prompt": text_prompt,
+            "video_negative_prompt": args.video_negative_prompt,
+            "audio_negative_prompt": args.audio_negative_prompt,
+            "multi_modal_data": {"image": image, "audio": audio},
+        }
 
-    sampling_params = OmniDiffusionSamplingParams(
-        height=args.height,
-        width=args.width,
-        num_inference_steps=args.num_inference_steps,
-        seed=args.seed,
-        extra_args={
-            "solver_name": args.solver_name,
-            "shift": args.shift,
-        },
-    )
+        sampling_params = OmniDiffusionSamplingParams(
+            height=args.height,
+            width=args.width,
+            num_inference_steps=args.num_inference_steps,
+            seed=args.seed,
+            extra_args={
+                "solver_name": args.solver_name,
+                "shift": args.shift,
+            },
+        )
 
     parallel_config = DiffusionParallelConfig(
         cfg_parallel_size=args.cfg_parallel_size,
@@ -215,22 +237,28 @@ def main() -> None:
     elapsed = time.perf_counter() - start
 
     if not outputs:
-        raise RuntimeError("No output returned from DreamID-Omni.")
+        raise RuntimeError("No output returned from the model.")
     result = outputs[0]
     if not result.images:
-        raise RuntimeError("No video frames found in DreamID-Omni output.")
+        raise RuntimeError("No video frames found in model output.")
     generated_video = result.images[0]
     mm = result.multimodal_output or {}
     generated_audio = mm.get("audio")
-    fps = int(mm.get("fps", 24))
-    sample_rate = int(mm.get("audio_sample_rate", 16000))
 
-    # DreamID-Omni returns video as (C, F, H, W) float32 in [-1, 1].
-    # mux_video_audio_bytes expects (F, H, W, C) uint8.
-    if not isinstance(generated_video, np.ndarray) or generated_video.ndim != 4:
-        raise RuntimeError(f"Unexpected video shape: {getattr(generated_video, 'shape', None)}")
-    frames = generated_video.transpose(1, 2, 3, 0)
-    frames = (np.clip((frames + 1.0) / 2.0, 0.0, 1.0) * 255.0).round().astype(np.uint8)
+    if args.model_type == "magi-human":
+        # MagiHuman returns frames already as (F, H, W, C) uint8, ready to mux.
+        fps = float(mm.get("fps", 25))
+        sample_rate = int(mm.get("audio_sample_rate", 24000))
+        frames = generated_video
+    else:
+        fps = float(mm.get("fps", 24))
+        sample_rate = int(mm.get("audio_sample_rate", 16000))
+        # DreamID-Omni returns video as (C, F, H, W) float32 in [-1, 1].
+        # mux_video_audio_bytes expects (F, H, W, C) uint8.
+        if not isinstance(generated_video, np.ndarray) or generated_video.ndim != 4:
+            raise RuntimeError(f"Unexpected video shape: {getattr(generated_video, 'shape', None)}")
+        frames = generated_video.transpose(1, 2, 3, 0)
+        frames = (np.clip((frames + 1.0) / 2.0, 0.0, 1.0) * 255.0).round().astype(np.uint8)
 
     audio_np = None
     if generated_audio is not None:

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

examples/online_serving/helios/README.md

@@ -1,175 +0,0 @@
-# Helios Video Generation via Online API
-
-This example demonstrates how to use the `/v1/videos` API with Helios models using the generic `extra_params` field.
-
-## Overview
-
-The `/v1/videos` API now supports model-specific parameters through the `extra_params` field, which accepts a JSON object containing any model-specific configuration. This allows supporting new models like Helios without modifying the API.
-
-## Helios Model Variants
-
-- **Helios-Base**: Basic T2V/I2V/V2V generation (Stage 1 only)
-- **Helios-Mid**: Advanced generation with CFG-Zero* (Stage 2)
-- **Helios-Distilled**: Fast generation with DMD (Stage 2)
-
-## API Usage Examples
-
-### 1. T2V (Text-to-Video) - Helios-Base
-
-Basic text-to-video generation:
-
-```bash
-curl -X POST http://localhost:8000/v1/videos \
-  -F "prompt=A serene lakeside sunrise with mist over the water." \
-  -F "model=BestWishYsh/Helios-Base" \
-  -F "width=640" \
-  -F "height=384" \
-  -F "num_frames=99" \
-  -F "num_inference_steps=50" \
-  -F "guidance_scale=5.0" \
-  -F "seed=42"
-```
-
-### 2. I2V (Image-to-Video) - Helios-Base
-
-Generate video from an input image:
-
-```bash
-curl -X POST http://localhost:8000/v1/videos \
-  -F "prompt=The lake water gently ripples as morning mist rises." \
-  -F "model=BestWishYsh/Helios-Base" \
-  -F "input_reference=@/path/to/image.jpg" \
-  -F "width=640" \
-  -F "height=384" \
-  -F "num_frames=99" \
-  -F "guidance_scale=5.0"
-```
-
-### 3. Helios-Mid with Stage 2 + CFG-Zero*
-
-Advanced generation with pyramid multi-stage denoising:
-
-```bash
-curl -X POST http://localhost:8000/v1/videos \
-  -F "prompt=A serene lakeside sunrise with mist over the water." \
-  -F "model=BestWishYsh/Helios-Mid" \
-  -F "width=640" \
-  -F "height=384" \
-  -F "guidance_scale=5.0" \
-  -F 'extra_params={
-    "is_enable_stage2": true,
-    "pyramid_num_stages": 3,
-    "pyramid_num_inference_steps_list": [20, 20, 20],
-    "use_cfg_zero_star": true,
-    "use_zero_init": true,
-    "zero_steps": 1
-  }'
-```
-
-### 4. Helios-Distilled with DMD
-
-Fast generation with Distribution Matching Distillation:
-
-```bash
-curl -X POST http://localhost:8000/v1/videos \
-  -F "prompt=A serene lakeside sunrise with mist over the water." \
-  -F "model=BestWishYsh/Helios-Distilled" \
-  -F "width=640" \
-  -F "height=384" \
-  -F "guidance_scale=1.0" \
-  -F 'extra_params={
-    "is_enable_stage2": true,
-    "pyramid_num_stages": 3,
-    "pyramid_num_inference_steps_list": [2, 2, 2],
-    "is_amplify_first_chunk": true
-  }'
-```
-
-## Model-Specific Parameters
-
-The `extra_params` field accepts a JSON object with model-specific parameters. For Helios models, supported parameters include:
-
-### Stage 2 Parameters
-- `is_enable_stage2` (bool): Enable pyramid multi-stage denoising
-- `pyramid_num_stages` (int): Number of pyramid stages (default: 3)
-- `pyramid_num_inference_steps_list` (array): Steps per stage, e.g., `[20, 20, 20]`
-
-### CFG Zero Star Parameters (Helios-Mid)
-- `use_cfg_zero_star` (bool): Enable CFG Zero Star guidance
-- `use_zero_init` (bool): Use zero initialization for first steps
-- `zero_steps` (int): Number of initial zero prediction steps (default: 1)
-
-### DMD Parameters (Helios-Distilled)
-- `is_amplify_first_chunk` (bool): Enable DMD amplification for first chunk
-
-### Video Input (V2V mode)
-For video-to-video generation, upload a video file via `input_reference`:
-
-```bash
-curl -X POST http://localhost:8000/v1/videos \
-  -F "prompt=Transform the video into a watercolor painting style." \
-  -F "model=BestWishYsh/Helios-Base" \
-  -F "input_reference=@/path/to/video.mp4" \
-  -F "guidance_scale=5.0"
-```
-
-## Python Example
-
-```python
-import requests
-import json
-
-url = "http://localhost:8000/v1/videos"
-
-# Helios-Mid with Stage 2
-data = {
-    "prompt": "A serene lakeside sunrise with mist over the water.",
-    "model": "BestWishYsh/Helios-Mid",
-    "width": 640,
-    "height": 384,
-    "guidance_scale": 5.0,
-    "extra_params": json.dumps({
-        "is_enable_stage2": True,
-        "pyramid_num_stages": 3,
-        "pyramid_num_inference_steps_list": [20, 20, 20],
-        "use_cfg_zero_star": True,
-        "use_zero_init": True,
-        "zero_steps": 1
-    })
-}
-
-response = requests.post(url, data=data)
-video_job = response.json()
-print(f"Video job created: {video_job['id']}")
-
-# Poll for completion
-import time
-while True:
-    status_response = requests.get(f"{url}/{video_job['id']}")
-    status = status_response.json()
-
-    if status['status'] == 'completed':
-        print(f"Video generated: {status['file_name']}")
-        break
-    elif status['status'] == 'failed':
-        print(f"Generation failed: {status.get('error')}")
-        break
-
-    print(f"Progress: {status['progress']}%")
-    time.sleep(2)
-```
-
-## Benefits of This Approach
-
-1. **No API Changes**: New models can be supported without modifying the API
-2. **Backward Compatible**: Existing clients continue to work
-3. **Flexible**: Any model-specific parameter can be passed
-4. **Type-Safe**: Parameters are validated at the model level
-5. **Future-Proof**: Supports models that don't exist yet
-
-## Notes
-
-- The `extra_params` field must be a valid JSON object
-- Parameters are passed directly to the model's `extra_args`
-- Invalid parameters will be caught by the model implementation
-- Tensor inputs (images/videos) should use `input_reference`, not `extra_params`

examples/online_serving/helios/helios_client.py

@@ -1,212 +0,0 @@
-#!/usr/bin/env python3
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""
-Helios video generation client example using the /v1/videos API.
-
-This example demonstrates how to use the generic extra_params field
-to pass Helios-specific parameters without modifying the API.
-"""
-
-import argparse
-import json
-import time
-from pathlib import Path
-
-import requests
-
-
-def create_video_job(
-    api_url: str,
-    prompt: str,
-    model: str,
-    width: int = 640,
-    height: int = 384,
-    num_frames: int = 99,
-    guidance_scale: float = 5.0,
-    seed: int = 42,
-    extra_params: dict | None = None,
-    input_image: Path | None = None,
-) -> dict:
-    """Create a video generation job."""
-    data = {
-        "prompt": prompt,
-        "model": model,
-        "width": width,
-        "height": height,
-        "num_frames": num_frames,
-        "guidance_scale": guidance_scale,
-        "seed": seed,
-    }
-
-    files = {}
-    if input_image:
-        files["input_reference"] = open(input_image, "rb")
-
-    if extra_params:
-        data["extra_params"] = json.dumps(extra_params)
-
-    response = requests.post(f"{api_url}/v1/videos", data=data, files=files)
-    response.raise_for_status()
-
-    if files:
-        files["input_reference"].close()
-
-    return response.json()
-
-
-def poll_video_status(api_url: str, video_id: str, poll_interval: int = 2) -> dict:
-    """Poll video generation status until completion."""
-    print(f"Polling video job: {video_id}")
-
-    while True:
-        response = requests.get(f"{api_url}/v1/videos/{video_id}")
-        response.raise_for_status()
-        status_data = response.json()
-
-        status = status_data["status"]
-        progress = status_data.get("progress", 0)
-
-        print(f"Status: {status}, Progress: {progress}%")
-
-        if status == "completed":
-            print("Video generation completed!")
-            return status_data
-        elif status == "failed":
-            error = status_data.get("error", {})
-            raise RuntimeError(f"Video generation failed: {error}")
-
-        time.sleep(poll_interval)
-
-
-def main():
-    parser = argparse.ArgumentParser(description="Helios video generation client")
-    parser.add_argument(
-        "--api-url",
-        default="http://localhost:8000",
-        help="API server URL",
-    )
-    parser.add_argument(
-        "--model",
-        default="BestWishYsh/Helios-Base",
-        help="Model name (Helios-Base, Helios-Mid, Helios-Distilled)",
-    )
-    parser.add_argument(
-        "--prompt",
-        default="A serene lakeside sunrise with mist over the water.",
-        help="Text prompt",
-    )
-    parser.add_argument(
-        "--width",
-        type=int,
-        default=640,
-        help="Video width",
-    )
-    parser.add_argument(
-        "--height",
-        type=int,
-        default=384,
-        help="Video height",
-    )
-    parser.add_argument(
-        "--num-frames",
-        type=int,
-        default=99,
-        help="Number of frames",
-    )
-    parser.add_argument(
-        "--guidance-scale",
-        type=float,
-        default=5.0,
-        help="Guidance scale",
-    )
-    parser.add_argument(
-        "--seed",
-        type=int,
-        default=42,
-        help="Random seed",
-    )
-    parser.add_argument(
-        "--input-image",
-        type=Path,
-        help="Input image for I2V mode",
-    )
-    parser.add_argument(
-        "--preset",
-        choices=["base", "mid-stage2", "distilled"],
-        default="base",
-        help="Helios preset configuration",
-    )
-
-    args = parser.parse_args()
-
-    # Define preset configurations
-    presets = {
-        "base": None,  # No extra params for base model
-        "mid-stage2": {
-            "is_enable_stage2": True,
-            "pyramid_num_stages": 3,
-            "pyramid_num_inference_steps_list": [20, 20, 20],
-            "use_cfg_zero_star": True,
-            "use_zero_init": True,
-            "zero_steps": 1,
-        },
-        "distilled": {
-            "is_enable_stage2": True,
-            "pyramid_num_stages": 3,
-            "pyramid_num_inference_steps_list": [2, 2, 2],
-            "is_amplify_first_chunk": True,
-        },
-    }
-
-    extra_params = presets[args.preset]
-
-    print("=" * 50)
-    print("Helios Video Generation")
-    print("=" * 50)
-    print(f"API URL: {args.api_url}")
-    print(f"Model: {args.model}")
-    print(f"Preset: {args.preset}")
-    print(f"Prompt: {args.prompt}")
-    print(f"Size: {args.width}x{args.height}")
-    print(f"Frames: {args.num_frames}")
-    print(f"Guidance Scale: {args.guidance_scale}")
-    if extra_params:
-        print(f"Extra Params: {json.dumps(extra_params, indent=2)}")
-    print()
-
-    # Create video job
-    print("Creating video generation job...")
-    job = create_video_job(
-        api_url=args.api_url,
-        prompt=args.prompt,
-        model=args.model,
-        width=args.width,
-        height=args.height,
-        num_frames=args.num_frames,
-        guidance_scale=args.guidance_scale,
-        seed=args.seed,
-        extra_params=extra_params,
-        input_image=args.input_image,
-    )
-
-    video_id = job["id"]
-    print(f"Video job created: {video_id}")
-    print()
-
-    # Poll for completion
-    result = poll_video_status(args.api_url, video_id)
-
-    print()
-    print("=" * 50)
-    print("Generation Complete")
-    print("=" * 50)
-    print(f"Video ID: {result['id']}")
-    print(f"File: {result.get('file_name', 'N/A')}")
-    print(f"Inference Time: {result.get('inference_time_s', 'N/A')}s")
-    print()
-
-
-if __name__ == "__main__":
-    main()

examples/online_serving/helios/run_helios_distilled.sh

@@ -1,79 +0,0 @@
-#!/bin/bash
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-# Helios-Distilled with DMD Example
-# Fast generation with Distribution Matching Distillation
-
-API_URL="${API_URL:-http://localhost:8000}"
-MODEL="${MODEL:-BestWishYsh/Helios-Distilled}"
-PROMPT="${PROMPT:-A serene lakeside sunrise with mist over the water.}"
-
-echo "==================================="
-echo "Helios-Distilled with DMD"
-echo "==================================="
-echo "API URL: $API_URL"
-echo "Model: $MODEL"
-echo "Prompt: $PROMPT"
-echo ""
-
-# Model-specific parameters for Helios-Distilled
-extra_params='{
-  "is_enable_stage2": true,
-  "pyramid_num_stages": 3,
-  "pyramid_num_inference_steps_list": [2, 2, 2],
-  "is_amplify_first_chunk": true
-}'
-
-echo "Model extra params: $extra_params"
-echo ""
-
-# Create video generation job
-echo "Creating video generation job..."
-RESPONSE=$(curl -s -X POST "$API_URL/v1/videos" \
-  -F "prompt=$PROMPT" \
-  -F "model=$MODEL" \
-  -F "width=640" \
-  -F "height=384" \
-  -F "guidance_scale=1.0" \
-  -F "seed=42" \
-  -F "extra_params=$extra_params")
-
-echo "Response: $RESPONSE"
-echo ""
-
-# Extract video ID
-VIDEO_ID=$(echo "$RESPONSE" | grep -o '"id":"[^"]*"' | cut -d'"' -f4)
-
-if [ -z "$VIDEO_ID" ]; then
-  echo "Error: Failed to create video job"
-  exit 1
-fi
-
-echo "Video job created: $VIDEO_ID"
-echo ""
-
-# Poll for completion
-echo "Polling for completion..."
-while true; do
-  STATUS_RESPONSE=$(curl -s "$API_URL/v1/videos/$VIDEO_ID")
-  STATUS=$(echo "$STATUS_RESPONSE" | grep -o '"status":"[^"]*"' | cut -d'"' -f4)
-  PROGRESS=$(echo "$STATUS_RESPONSE" | grep -o '"progress":[0-9]*' | cut -d':' -f2)
-
-  echo "Status: $STATUS, Progress: $PROGRESS%"
-
-  if [ "$STATUS" = "completed" ]; then
-    echo ""
-    echo "Video generation completed!"
-    echo "Full response:"
-    echo "$STATUS_RESPONSE" | jq '.'
-    break
-  elif [ "$STATUS" = "failed" ]; then
-    echo ""
-    echo "Video generation failed!"
-    echo "$STATUS_RESPONSE" | jq '.'
-    exit 1
-  fi
-
-  sleep 2
-done

examples/online_serving/helios/run_helios_mid_stage2.sh

@@ -1,81 +0,0 @@
-#!/bin/bash
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-# Helios-Mid with Stage 2 + CFG-Zero* Example
-# This demonstrates advanced generation with pyramid multi-stage denoising
-
-API_URL="${API_URL:-http://localhost:8000}"
-MODEL="${MODEL:-BestWishYsh/Helios-Mid}"
-PROMPT="${PROMPT:-A serene lakeside sunrise with mist over the water.}"
-
-echo "==================================="
-echo "Helios-Mid Stage 2 + CFG-Zero*"
-echo "==================================="
-echo "API URL: $API_URL"
-echo "Model: $MODEL"
-echo "Prompt: $PROMPT"
-echo ""
-
-# Model-specific parameters for Helios-Mid
-extra_params='{
-  "is_enable_stage2": true,
-  "pyramid_num_stages": 3,
-  "pyramid_num_inference_steps_list": [20, 20, 20],
-  "use_cfg_zero_star": true,
-  "use_zero_init": true,
-  "zero_steps": 1
-}'
-
-echo "Model extra params: $extra_params"
-echo ""
-
-# Create video generation job
-echo "Creating video generation job..."
-RESPONSE=$(curl -s -X POST "$API_URL/v1/videos" \
-  -F "prompt=$PROMPT" \
-  -F "model=$MODEL" \
-  -F "width=640" \
-  -F "height=384" \
-  -F "guidance_scale=5.0" \
-  -F "seed=42" \
-  -F "extra_params=$extra_params")
-
-echo "Response: $RESPONSE"
-echo ""
-
-# Extract video ID
-VIDEO_ID=$(echo "$RESPONSE" | grep -o '"id":"[^"]*"' | cut -d'"' -f4)
-
-if [ -z "$VIDEO_ID" ]; then
-  echo "Error: Failed to create video job"
-  exit 1
-fi
-
-echo "Video job created: $VIDEO_ID"
-echo ""
-
-# Poll for completion
-echo "Polling for completion..."
-while true; do
-  STATUS_RESPONSE=$(curl -s "$API_URL/v1/videos/$VIDEO_ID")
-  STATUS=$(echo "$STATUS_RESPONSE" | grep -o '"status":"[^"]*"' | cut -d'"' -f4)
-  PROGRESS=$(echo "$STATUS_RESPONSE" | grep -o '"progress":[0-9]*' | cut -d':' -f2)
-
-  echo "Status: $STATUS, Progress: $PROGRESS%"
-
-  if [ "$STATUS" = "completed" ]; then
-    echo ""
-    echo "Video generation completed!"
-    echo "Full response:"
-    echo "$STATUS_RESPONSE" | jq '.'
-    break
-  elif [ "$STATUS" = "failed" ]; then
-    echo ""
-    echo "Video generation failed!"
-    echo "$STATUS_RESPONSE" | jq '.'
-    exit 1
-  fi
-
-  sleep 2
-done

examples/online_serving/helios/run_helios_t2v.sh

@@ -1,69 +0,0 @@
-#!/bin/bash
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-# Helios Text-to-Video Generation Example
-# This script demonstrates how to use the /v1/videos API with Helios models
-
-API_URL="${API_URL:-http://localhost:8000}"
-MODEL="${MODEL:-BestWishYsh/Helios-Base}"
-PROMPT="${PROMPT:-A serene lakeside sunrise with mist over the water.}"
-
-echo "==================================="
-echo "Helios T2V Generation"
-echo "==================================="
-echo "API URL: $API_URL"
-echo "Model: $MODEL"
-echo "Prompt: $PROMPT"
-echo ""
-
-# Create video generation job
-echo "Creating video generation job..."
-RESPONSE=$(curl -s -X POST "$API_URL/v1/videos" \
-  -F "prompt=$PROMPT" \
-  -F "model=$MODEL" \
-  -F "width=640" \
-  -F "height=384" \
-  -F "num_frames=99" \
-  -F "num_inference_steps=50" \
-  -F "guidance_scale=5.0" \
-  -F "seed=42")
-
-echo "Response: $RESPONSE"
-echo ""
-
-# Extract video ID
-VIDEO_ID=$(echo "$RESPONSE" | grep -o '"id":"[^"]*"' | cut -d'"' -f4)
-
-if [ -z "$VIDEO_ID" ]; then
-  echo "Error: Failed to create video job"
-  exit 1
-fi
-
-echo "Video job created: $VIDEO_ID"
-echo ""
-
-# Poll for completion
-echo "Polling for completion..."
-while true; do
-  STATUS_RESPONSE=$(curl -s "$API_URL/v1/videos/$VIDEO_ID")
-  STATUS=$(echo "$STATUS_RESPONSE" | grep -o '"status":"[^"]*"' | cut -d'"' -f4)
-  PROGRESS=$(echo "$STATUS_RESPONSE" | grep -o '"progress":[0-9]*' | cut -d':' -f2)
-
-  echo "Status: $STATUS, Progress: $PROGRESS%"
-
-  if [ "$STATUS" = "completed" ]; then
-    echo ""
-    echo "Video generation completed!"
-    echo "Full response:"
-    echo "$STATUS_RESPONSE" | jq '.'
-    break
-  elif [ "$STATUS" = "failed" ]; then
-    echo ""
-    echo "Video generation failed!"
-    echo "$STATUS_RESPONSE" | jq '.'
-    exit 1
-  fi
-
-  sleep 2
-done

examples/online_serving/joyvl_interaction/bridges/asr_bridge.py

@@ -0,0 +1,131 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import argparse
+import base64
+import io
+import json
+import struct
+import wave
+
+import aiohttp
+from aiohttp import web
+
+_HEADER = struct.Struct(">iii")
+_SAMPLE_RATE = 16000
+_DEFAULT_PROMPT = "Transcribe the audio."
+_ASR_TEXT_TAG = "<asr_text>"
+
+
+def _clean_transcript(raw: str) -> str:
+    """Qwen3-ASR emits ``language <Lang><asr_text><transcript>``; keep the transcript."""
+    text = raw or ""
+    marker = text.rfind(_ASR_TEXT_TAG)
+    if marker != -1:
+        text = text[marker + len(_ASR_TEXT_TAG) :]
+    return text.replace("</asr_text>", "").strip()
+
+
+def _pcm_to_wav_b64(pcm: bytes, sample_rate: int) -> str:
+    buf = io.BytesIO()
+    with wave.open(buf, "wb") as wav:
+        wav.setnchannels(1)
+        wav.setsampwidth(2)
+        wav.setframerate(sample_rate)
+        wav.writeframes(pcm)
+    return base64.b64encode(buf.getvalue()).decode()
+
+
+def _result(text: str, *, final: bool) -> dict:
+    return {
+        "asr_response": {
+            "event_type": "IS_FINAL" if final else "IS_PARTIAL",
+            "recognition_result": {"hypothesis": [{"text": text, "confidence": None}]},
+        },
+        "mid": "",
+        "code": 0,
+        "msg": "ok",
+    }
+
+
+async def _transcribe(cfg: dict, pcm: bytes) -> str:
+    if not pcm:
+        return ""
+    wav_b64 = _pcm_to_wav_b64(pcm, cfg["sample_rate"])
+    body = {
+        "model": cfg["model"],
+        "messages": [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "input_audio", "input_audio": {"data": wav_b64, "format": "wav"}},
+                    {"type": "text", "text": cfg["prompt"]},
+                ],
+            }
+        ],
+        "temperature": 0.0,
+        "max_tokens": cfg["max_tokens"],
+    }
+    url = cfg["backend_url"].rstrip("/") + "/v1/chat/completions"
+    async with aiohttp.ClientSession() as session:
+        async with session.post(url, json=body) as resp:
+            payload = await resp.json()
+    choices = payload.get("choices") or []
+    if not choices:
+        return ""
+    return _clean_transcript(choices[0].get("message", {}).get("content") or "")
+
+
+async def _handle(request: web.Request) -> web.WebSocketResponse:
+    ws = web.WebSocketResponse(max_msg_size=0)
+    await ws.prepare(request)
+    cfg = request.app["cfg"]
+    pcm = bytearray()
+    try:
+        async for msg in ws:
+            if msg.type == aiohttp.WSMsgType.BINARY and len(msg.data) >= _HEADER.size:
+                seqid, _, _ = _HEADER.unpack(msg.data[: _HEADER.size])
+                pcm.extend(msg.data[_HEADER.size :])
+                if seqid < 0:
+                    text = await _transcribe(cfg, bytes(pcm))
+                    await ws.send_str(json.dumps(_result(text, final=True), ensure_ascii=False))
+                    pcm.clear()
+            elif msg.type in (aiohttp.WSMsgType.CLOSE, aiohttp.WSMsgType.CLOSED, aiohttp.WSMsgType.ERROR):
+                break
+    finally:
+        if not ws.closed:
+            await ws.close()
+    return ws
+
+
+def create_app(backend_url: str, model: str, sample_rate: int, prompt: str, max_tokens: int) -> web.Application:
+    app = web.Application()
+    app["cfg"] = {
+        "backend_url": backend_url,
+        "model": model,
+        "sample_rate": sample_rate,
+        "prompt": prompt,
+        "max_tokens": max_tokens,
+    }
+    app.router.add_get("/v1/asr", _handle)
+    return app
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="JoyVL webui <-> Qwen3-ASR (chat input_audio) bridge")
+    parser.add_argument("--host", default="0.0.0.0")
+    parser.add_argument("--port", type=int, default=8093)
+    parser.add_argument("--backend-url", default="http://127.0.0.1:8094")
+    parser.add_argument("--model", default="Qwen/Qwen3-ASR-1.7B")
+    parser.add_argument("--sample-rate", type=int, default=_SAMPLE_RATE)
+    parser.add_argument("--prompt", default=_DEFAULT_PROMPT)
+    parser.add_argument("--max-tokens", type=int, default=256)
+    args = parser.parse_args()
+    app = create_app(args.backend_url, args.model, args.sample_rate, args.prompt, args.max_tokens)
+    web.run_app(app, host=args.host, port=args.port)
+
+
+if __name__ == "__main__":
+    main()

examples/online_serving/joyvl_interaction/bridges/tts_bridge.py

@@ -0,0 +1,117 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import json
+import logging
+
+import aiohttp
+from aiohttp import web
+
+try:
+    import cn2an  # optional: Arabic-digit -> Chinese number normalization
+except ImportError:
+    cn2an = None
+
+logger = logging.getLogger("tts_bridge")
+
+
+def _normalize(text: str) -> str:
+    """Spell out Arabic digits in Chinese so short counting outputs synthesize
+    intelligibly: bare "1次" garbles on Qwen3-TTS (and others), while "一次"
+    is clean. ``cn2an.transform`` rewrites digits in place ("第3次"->"第三次")
+    and is pure-Python (~microseconds). No-op if cn2an is unavailable."""
+    if not cn2an or not text:
+        return text
+    try:
+        return cn2an.transform(text, "an2cn")
+    except Exception as err:  # never let TN break synthesis
+        logger.debug("cn2an normalize skipped: %s", err)
+        return text
+
+
+async def _pump_backend_to_front(back: aiohttp.ClientWebSocketResponse, front: web.WebSocketResponse) -> None:
+    async for msg in back:
+        if msg.type == aiohttp.WSMsgType.BINARY:
+            await front.send_bytes(msg.data)
+        elif msg.type == aiohttp.WSMsgType.TEXT:
+            event = json.loads(msg.data).get("type")
+            if event == "session.done":
+                await front.send_str(json.dumps({"type": "response.done"}))
+                return
+            if event == "error":
+                await front.send_str(json.dumps({"type": "error", "error": "tts backend error"}))
+                return
+        elif msg.type in (aiohttp.WSMsgType.CLOSE, aiohttp.WSMsgType.CLOSED, aiohttp.WSMsgType.ERROR):
+            return
+
+
+async def _handle(request: web.Request) -> web.WebSocketResponse:
+    front = web.WebSocketResponse(heartbeat=20, max_msg_size=0)
+    await front.prepare(request)
+    cfg = request.app["cfg"]
+
+    session = aiohttp.ClientSession()
+    back: aiohttp.ClientWebSocketResponse | None = None
+    pump: asyncio.Task | None = None
+    parts: list[str] = []
+    try:
+        async for msg in front:
+            if msg.type != web.WSMsgType.TEXT:
+                continue
+            data = json.loads(msg.data)
+
+            if "config" in data and back is None:
+                voice = (data["config"] or {}).get("voice") or cfg["voice"]
+                back = await session.ws_connect(cfg["backend_url"], max_msg_size=0)
+                await back.send_str(
+                    json.dumps({"type": "session.config", "response_format": cfg["response_format"], "voice": voice})
+                )
+                pump = asyncio.create_task(_pump_backend_to_front(back, front))
+            elif data.get("type") == "input_text.append" and back is not None:
+                parts.append(data.get("text", ""))
+            elif data.get("type") == "input_text.commit" and back is not None:
+                # Normalize the full utterance once (digits may span chunks), then send.
+                await back.send_str(json.dumps({"type": "input.text", "text": _normalize("".join(parts))}))
+                parts = []
+                await back.send_str(json.dumps({"type": "input.done"}))
+                if pump is not None:
+                    await pump
+                break
+    except Exception as err:
+        logger.warning("tts bridge error: %s", err)
+    finally:
+        if pump is not None and not pump.done():
+            pump.cancel()
+        if back is not None and not back.closed:
+            await back.close()
+        await session.close()
+        if not front.closed:
+            await front.close()
+    return front
+
+
+def create_app(backend_url: str, voice: str, response_format: str) -> web.Application:
+    app = web.Application()
+    app["cfg"] = {"backend_url": backend_url, "voice": voice, "response_format": response_format}
+    app.router.add_get("/v1/tts", _handle)
+    return app
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="JoyVL webui <-> Qwen3-TTS websocket bridge")
+    parser.add_argument("--host", default="0.0.0.0")
+    parser.add_argument("--port", type=int, default=8092)
+    parser.add_argument("--backend-url", default="ws://127.0.0.1:8091/v1/audio/speech/stream")
+    parser.add_argument("--voice", default="vivian")
+    parser.add_argument("--response-format", default="pcm")
+    args = parser.parse_args()
+    logging.basicConfig(level=logging.INFO)
+    web.run_app(create_app(args.backend_url, args.voice, args.response_format), host=args.host, port=args.port)
+
+
+if __name__ == "__main__":
+    main()

examples/online_serving/joyvl_interaction/cli/run_cli_demo.py

@@ -0,0 +1,32 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import argparse
+
+from stream_client import stream
+
+ICON = {"response": "SPEAK   ", "silence": "silence ", "delegate": "DELEGATE"}
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("video")
+    parser.add_argument("--server", default="http://127.0.0.1:8070")
+    parser.add_argument("--query", default=None)
+    parser.add_argument("--fps", type=float, default=1.0)
+    parser.add_argument("--session-id", default="cli")
+    args = parser.parse_args()
+
+    for tick in stream(args.video, args.server, args.session_id, args.query, args.fps):
+        if tick.action == "silence":
+            continue
+        line = f"[{tick.t:6.1f}s] {ICON.get(tick.action)} {tick.text}"
+        if tick.delegation:
+            line += f"  (-> {tick.delegation.get('question', '')})"
+        print(line, flush=True)
+
+
+if __name__ == "__main__":
+    main()

examples/online_serving/joyvl_interaction/cli/stream_client.py

@@ -0,0 +1,79 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import base64
+import time
+from collections.abc import Iterator
+from dataclasses import dataclass
+
+import cv2
+import requests
+
+
+@dataclass
+class Tick:
+    index: int
+    t: float
+    action: str
+    text: str
+    skipped: bool
+    latency_ms: float
+    delegation: dict | None
+
+
+def iter_frames(video_path: str, fps: float) -> Iterator[tuple[float, bytes]]:
+    cap = cv2.VideoCapture(video_path)
+    src_fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
+    step = max(1, round(src_fps / fps))
+    idx = 0
+    while True:
+        ok, frame = cap.read()
+        if not ok:
+            break
+        if idx % step == 0:
+            ok, buf = cv2.imencode(".jpg", frame)
+            if ok:
+                yield idx / src_fps, buf.tobytes()
+        idx += 1
+    cap.release()
+
+
+def stream(
+    video_path: str,
+    server: str,
+    session_id: str,
+    query: str | None = None,
+    fps: float = 1.0,
+    realtime: bool = False,
+) -> Iterator[Tick]:
+    """Send each sampled frame as one chat completion; yield the parsed action.
+
+    ``query`` is attached to the first tick to arm the session; inject later
+    queries by posting another frame turn with a new text part.
+    """
+    url = server.rstrip("/") + "/v1/chat/completions"
+    requests.post(server.rstrip("/") + "/reset", json={"session_id": session_id}, timeout=10)
+    for i, (t, jpeg) in enumerate(iter_frames(video_path, fps)):
+        data_url = "data:image/jpeg;base64," + base64.b64encode(jpeg).decode()
+        content = [{"type": "image_url", "image_url": {"url": data_url}}]
+        if i == 0 and query:
+            content.insert(0, {"type": "text", "text": query})
+        resp = requests.post(
+            url,
+            json={"session_id": session_id, "messages": [{"role": "user", "content": content}]},
+            timeout=120,
+        ).json()
+        info = resp.get("interaction", {})
+        yield Tick(
+            index=info.get("frame_index", i + 1),
+            t=t,
+            action=info.get("action", "?"),
+            text=resp["choices"][0]["message"]["content"],
+            skipped=info.get("inference_skipped", False),
+            latency_ms=info.get("latency_ms", 0.0),
+            delegation=info.get("delegation"),
+        )
+        if realtime:
+            time.sleep(1.0 / fps)

examples/online_serving/joyvl_interaction/rtsp/README.md

@@ -0,0 +1,46 @@
+# RTSP local streaming
+
+Simulate an RTSP camera from a local video file so the WebUI (or any RTSP client) can
+test the RTSP input path without a physical IP camera. RTSP is a **webui-side input** —
+the browser pulls the stream and feeds frames to the orchestrator over the normal API;
+no serving-layer code is involved.
+
+Needs [`ffmpeg`](https://ffmpeg.org/) and [MediaMTX](https://github.com/bluenviron/mediamtx/releases)
+(download the binary + `mediamtx.yml` and keep them next to these scripts, or point
+`MEDIAMTX_BIN` / `MEDIAMTX_CONFIG` at them).
+
+## 1. Start the RTSP server
+
+```bash
+bash ./mediamtx.sh        # MediaMTX, listens on :8554
+```
+
+## 2. Push a local video as an RTSP stream
+
+```bash
+bash ./rtsp.sh ./videos/example.mp4 rtsp://127.0.0.1:8554/fire1
+```
+
+`rtsp.sh` takes `[video-path] [rtsp-url]` (or set `VIDEO_PATH` / `RTSP_URL`). To stream a
+whole directory of clips, one stream per file:
+
+```bash
+bash ./rtsp_all.sh ./videos rtsp://127.0.0.1:8554
+```
+
+## 3. Use it in the WebUI
+
+Enter the URL in the WebUI RTSP box:
+
+```
+rtsp://127.0.0.1:8554/fire1
+```
+
+If the WebUI runs on another machine, replace `127.0.0.1` with the MediaMTX host's IP.
+
+## Notes
+
+- `rtsp.sh` transcodes audio to AAC; replace `-c:a aac -b:a 128k -ar 44100` with `-an`
+  if the source clip has no audio track and your `ffmpeg` errors on it.
+- If a client cannot connect, check that port `8554` is reachable and not blocked by a
+  firewall, container network, or security group.

examples/online_serving/joyvl_interaction/rtsp/mediamtx.sh

@@ -0,0 +1,11 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+# Start a local MediaMTX RTSP server (listens on :8554 by default).
+# Download the binary + mediamtx.yml from https://github.com/bluenviron/mediamtx/releases
+# and either keep them next to this script or point the vars below at them.
+
+MEDIAMTX_BIN="${MEDIAMTX_BIN:-./mediamtx}"
+MEDIAMTX_CONFIG="${MEDIAMTX_CONFIG:-./mediamtx.yml}"
+
+exec "$MEDIAMTX_BIN" "$MEDIAMTX_CONFIG"

examples/online_serving/joyvl_interaction/rtsp/rtsp.sh

@@ -0,0 +1,38 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+# Loop a local video file and push it to an RTSP URL, simulating an RTSP camera.
+#   bash ./rtsp.sh [video-path] [rtsp-output-url]
+# Defaults stream ./videos/example.mp4 to rtsp://127.0.0.1:8554/fire1
+
+VIDEO_PATH="${1:-${VIDEO_PATH:-./videos/example.mp4}}"
+RTSP_URL="${2:-${RTSP_URL:-rtsp://127.0.0.1:8554/fire1}}"
+
+if [[ ! -f "$VIDEO_PATH" ]]; then
+  echo "Video file not found: $VIDEO_PATH" >&2
+  exit 1
+fi
+
+echo "Input video: $VIDEO_PATH"
+echo "RTSP output: $RTSP_URL"
+
+# Drop "-c:a aac -b:a 128k -ar 44100" for "-an" if the source has no audio track.
+exec ffmpeg \
+  -hide_banner \
+  -loglevel info \
+  -re \
+  -stream_loop -1 \
+  -i "$VIDEO_PATH" \
+  -vf "scale='min(1280,iw)':-2" \
+  -c:v libx264 \
+  -preset veryfast \
+  -tune zerolatency \
+  -b:v 2500k \
+  -maxrate 2500k \
+  -bufsize 5000k \
+  -c:a aac \
+  -b:a 128k \
+  -ar 44100 \
+  -f rtsp \
+  -rtsp_transport udp \
+  "$RTSP_URL"

examples/online_serving/joyvl_interaction/rtsp/rtsp_all.sh

@@ -0,0 +1,64 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+# Push every .mp4 in a directory as its own RTSP stream (one ffmpeg per file),
+# named rtsp://<base>/<filename-without-extension>.
+#   bash ./rtsp_all.sh [video-dir] [rtsp-base-url]
+
+VIDEO_DIR="${1:-./videos}"
+RTSP_BASE="${2:-rtsp://127.0.0.1:8554}"
+
+if [[ ! -d "$VIDEO_DIR" ]]; then
+  echo "Video directory not found: $VIDEO_DIR" >&2
+  exit 1
+fi
+
+shopt -s nullglob
+VIDEO_FILES=("$VIDEO_DIR"/*.mp4)
+shopt -u nullglob
+
+if [[ ${#VIDEO_FILES[@]} -eq 0 ]]; then
+  echo "No video files found in: $VIDEO_DIR" >&2
+  exit 1
+fi
+
+echo "Found ${#VIDEO_FILES[@]} video(s) in $VIDEO_DIR"
+
+PIDS=()
+cleanup() {
+  echo "Stopping all streams..."
+  for pid in "${PIDS[@]}"; do
+    kill "$pid" 2>/dev/null || true
+  done
+  wait 2>/dev/null
+}
+trap cleanup EXIT INT TERM
+
+for video in "${VIDEO_FILES[@]}"; do
+  stream_name="$(basename "${video%.*}")"
+  rtsp_url="${RTSP_BASE}/${stream_name}"
+  echo "Starting stream: $rtsp_url  <-  $(basename "$video")"
+  ffmpeg \
+    -hide_banner \
+    -loglevel error \
+    -re \
+    -stream_loop -1 \
+    -i "$video" \
+    -vf "format=yuv420p,scale='min(1280,iw)':'min(720,ih)':force_original_aspect_ratio=decrease:force_divisible_by=2" \
+    -c:v libx264 \
+    -preset veryfast \
+    -tune zerolatency \
+    -b:v 4000k \
+    -maxrate 4500k \
+    -bufsize 9000k \
+    -c:a aac \
+    -b:a 128k \
+    -ar 44100 \
+    -f rtsp \
+    -rtsp_transport tcp \
+    "$rtsp_url" &
+  PIDS+=($!)
+done
+
+echo "All ${#VIDEO_FILES[@]} streams started. Press Ctrl+C to stop."
+wait

examples/online_serving/joyvl_interaction/scripts/start_all.sh

@@ -0,0 +1,108 @@
+#!/bin/bash
+# One-shot JoyVL demo: model -> interaction orchestrator -> JD webui.
+#
+# The model is the only fixed component; every external module is pluggable and
+# configured here via env (unset = disabled), matching JoyVL's design:
+#
+#   TTS_URL                       voice out  (point at our tts_bridge, or any OpenAI-Realtime-style TTS WS)
+#   ASR_URL                       voice in   (point at our asr_bridge, or any compatible ASR WS)
+#   DELEGATION_BACKEND_URL        orchestrator background brain for </delegation> (OpenAI-compatible; DELEGATION_MODEL/DELEGATION_API_KEY pick model+key)
+#                                 DELEGATION_KIND=chat (default) text/VL brain answers; =image a text-to-image model generates a picture
+#                                 chat+Claude: DELEGATION_BACKEND_URL=https://api.anthropic.com/v1/ DELEGATION_MODEL=claude-... DELEGATION_API_KEY=sk-ant-...
+#                                 image+Qwen-Image: DELEGATION_KIND=image DELEGATION_BACKEND_URL=http://127.0.0.1:8091/v1
+#   BACKGROUND_CODEX_API_URL      webui-side background agent
+#   ENABLE_MEMORY (default 1)     3-tier memory; SUMMARIZER_BACKEND_URL/MODEL pick the summarizer (default: reuse main)
+#
+# Usage:  bash start_all.sh
+#         TTS_URL=ws://host/v1/tts ASR_URL=ws://host/v1/asr GPU=0 bash start_all.sh
+set -o pipefail
+
+CODE="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../../.." && pwd)"
+VENV="${VENV:-$CODE/.venv}"
+WEBUI_DIR="${WEBUI_DIR:-$CODE/joyvl-interaction/joyvl-interaction-webui}"
+MODEL="${MODEL:-jdopensource/JoyAI-VL-Interaction-Preview}"
+SERVED_NAME="${SERVED_NAME:-JoyAI-VL-Interaction-Preview}"
+GPU="${GPU:-0}"
+MODEL_PORT="${MODEL_PORT:-8061}"
+ORCH_PORT="${ORCH_PORT:-8070}"
+WEBUI_PORT="${WEBUI_PORT:-8999}"
+WEBUI_USERNAME="${WEBUI_USERNAME:-admin}"
+WEBUI_PASSWORD="${WEBUI_PASSWORD:-changeme}"   # override via env for any exposed deployment
+ENABLE_MEMORY="${ENABLE_MEMORY:-1}"
+RESPONSE_DEDUP="${RESPONSE_DEDUP:-0.85}"   # demo drops near-duplicate narration; set 1.0 for reference (exact-only)
+TTS_URL="${TTS_URL:-}"
+ASR_URL="${ASR_URL:-}"
+BACKGROUND_CODEX_API_URL="${BACKGROUND_CODEX_API_URL:-}"
+
+LOG=/tmp/joyvl_demo
+mkdir -p "$LOG"
+
+[ -d "$WEBUI_DIR" ] || { echo "! WEBUI_DIR not found: $WEBUI_DIR (set WEBUI_DIR=...)"; exit 1; }
+
+kill_port() {
+  local p
+  p=$(ss -tlnp 2>/dev/null | grep ":$1 " | grep -oE 'pid=[0-9]+' | head -1 | cut -d= -f2)
+  [ -n "${p:-}" ] && kill "$p" 2>/dev/null
+  return 0
+}
+wait_http() {
+  for _ in $(seq 1 150); do
+    curl -sk -m2 "$1" -o /dev/null 2>/dev/null && return 0
+    sleep 2
+  done
+  echo "  ! timeout waiting for $1"
+  return 1
+}
+
+echo "Resetting ports ${MODEL_PORT} ${ORCH_PORT} ${WEBUI_PORT}…"
+for p in "$MODEL_PORT" "$ORCH_PORT" "$WEBUI_PORT"; do kill_port "$p"; done
+sleep 3
+
+[ -f "$VENV/bin/activate" ] && source "$VENV/bin/activate"
+
+# detach into a new session so servers outlive this launcher
+spawn() { local log="$1"; shift; setsid "$@" > "$log" 2>&1 < /dev/null & }
+
+echo "[1/3] model on GPU ${GPU}…"
+spawn "$LOG/model.log" env HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 CUDA_VISIBLE_DEVICES="$GPU" \
+  vllm serve "$MODEL" --served-model-name "$SERVED_NAME" --port "$MODEL_PORT" \
+  --gpu-memory-utilization 0.85 --max-model-len 131072 --enable-prefix-caching \
+  --limit-mm-per-prompt '{"image":256,"video":1}'
+wait_http "http://127.0.0.1:${MODEL_PORT}/health" && echo "  up :${MODEL_PORT}"
+
+echo "[2/3] orchestrator (memory=${ENABLE_MEMORY})…"
+ORCH_ARGS=(--port "$ORCH_PORT" --main-backend-url "http://127.0.0.1:${MODEL_PORT}/v1"
+           --main-model "$SERVED_NAME"
+           --response-dedup-threshold "$RESPONSE_DEDUP")
+# delegation: point at a stronger background brain to enable it; otherwise keep it off
+if [ -n "${DELEGATION_BACKEND_URL:-}" ]; then
+  ORCH_ARGS+=(--delegation-backend-url "$DELEGATION_BACKEND_URL")
+  ORCH_ARGS+=(--delegation-kind "${DELEGATION_KIND:-chat}")
+  [ -n "${DELEGATION_MODEL:-}" ] && ORCH_ARGS+=(--delegation-model "$DELEGATION_MODEL")
+  [ -n "${DELEGATION_API_KEY:-}" ] && ORCH_ARGS+=(--delegation-api-key "$DELEGATION_API_KEY")
+else
+  ORCH_ARGS+=(--no-delegation)
+fi
+[ "$ENABLE_MEMORY" = "0" ] && ORCH_ARGS+=(--no-memory)
+[ -n "${SUMMARIZER_BACKEND_URL:-}" ] && ORCH_ARGS+=(--summarizer-backend-url "$SUMMARIZER_BACKEND_URL")
+[ -n "${SUMMARIZER_MODEL:-}" ] && ORCH_ARGS+=(--summarizer-model "$SUMMARIZER_MODEL")
+spawn "$LOG/orch.log" env PYTHONPATH="$CODE" python -m vllm_omni.experimental.fullduplex.joyvl.serving.server "${ORCH_ARGS[@]}"
+wait_http "http://127.0.0.1:${ORCH_PORT}/health" && echo "  up :${ORCH_PORT}"
+
+echo "[3/3] JD webui…"
+spawn "$LOG/webui.log" env WEBUI_USERNAME="$WEBUI_USERNAME" WEBUI_PASSWORD="$WEBUI_PASSWORD" \
+  TTS_URL="$TTS_URL" ASR_URL="$ASR_URL" BACKGROUND_CODEX_API_URL="$BACKGROUND_CODEX_API_URL" \
+  bash -c "cd '$WEBUI_DIR' && exec bash scripts/start_server.sh --api-base 'http://127.0.0.1:${ORCH_PORT}/v1'"
+webui_up=0
+for _ in $(seq 1 60); do grep -q "Server startup complete" "$LOG/webui.log" 2>/dev/null && { webui_up=1; break; }; sleep 1; done
+[ "$webui_up" = 1 ] || { echo "! webui did not start — see $LOG/webui.log"; exit 1; }
+
+IP=$(hostname -I 2>/dev/null | awk '{print $1}')
+echo
+echo "=== JoyVL demo ready ==="
+echo "  webui  : https://${IP}:${WEBUI_PORT}   (user ${WEBUI_USERNAME} / pass ${WEBUI_PASSWORD})"
+echo "  memory : $([ "$ENABLE_MEMORY" = 0 ] && echo off || echo on)"
+echo "  TTS    : ${TTS_URL:-<unset>}"
+echo "  ASR    : ${ASR_URL:-<unset>}"
+echo "  agent  : ${BACKGROUND_CODEX_API_URL:-<unset>}"
+echo "  logs   : ${LOG}/{model,orch,webui}.log"

examples/online_serving/joyvl_interaction/scripts/start_asr.sh

@@ -0,0 +1,24 @@
+#!/bin/bash
+# Serve Qwen3-ASR (the same audio-LLM AURA uses as its ASR stage) and run the
+# webui<->ASR bridge. Then start the webui with:
+#   ASR_URL=ws://127.0.0.1:8093/v1/asr bash scripts/start_server.sh
+set -euo pipefail
+
+ASR_MODEL="${ASR_MODEL:-Qwen/Qwen3-ASR-1.7B}"
+ASR_PORT="${ASR_PORT:-8094}"
+BRIDGE_PORT="${BRIDGE_PORT:-8093}"
+GPU="${GPU:-2}"
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
+
+# Qwen3-ASR is audio-in / text-out; served as a normal multimodal chat model.
+CUDA_VISIBLE_DEVICES="${GPU}" vllm serve "${ASR_MODEL}" \
+    --host 0.0.0.0 --port "${ASR_PORT}" \
+    --gpu-memory-utilization 0.3 --trust-remote-code &
+ASR_PID=$!
+trap 'kill ${ASR_PID} 2>/dev/null || true' EXIT
+
+until curl -s "http://127.0.0.1:${ASR_PORT}/health" >/dev/null 2>&1; do sleep 2; done
+python "${SCRIPT_DIR}/../bridges/asr_bridge.py" \
+    --port "${BRIDGE_PORT}" \
+    --backend-url "http://127.0.0.1:${ASR_PORT}" \
+    --model "${ASR_MODEL}"

examples/online_serving/joyvl_interaction/scripts/start_interaction.sh

@@ -0,0 +1,23 @@
+#!/bin/bash
+# Launch the interaction orchestrator in front of a running model server.
+set -euo pipefail
+
+PORT="${PORT:-8070}"
+MAIN_BACKEND_URL="${MAIN_BACKEND_URL:-http://127.0.0.1:8061/v1}"
+MAIN_MODEL="${MAIN_MODEL:-JoyAI-VL-Interaction-Preview}"
+PERSONA="${PERSONA:-default}"
+
+# By default memory reuses the main model as its own summarizer. Set
+# SUMMARIZER_BACKEND_URL/SUMMARIZER_MODEL to a dedicated Qwen3-VL-4B server,
+# or pass --no-memory for the lightest setup.
+EXTRA=()
+[[ -n "${SUMMARIZER_BACKEND_URL:-}" ]] && EXTRA+=(--summarizer-backend-url "${SUMMARIZER_BACKEND_URL}")
+[[ -n "${SUMMARIZER_MODEL:-}" ]] && EXTRA+=(--summarizer-model "${SUMMARIZER_MODEL}")
+[[ "${NO_MEMORY:-0}" == "1" ]] && EXTRA+=(--no-memory)
+
+python -m vllm_omni.experimental.fullduplex.joyvl.serving.server \
+    --port "${PORT}" \
+    --main-backend-url "${MAIN_BACKEND_URL}" \
+    --main-model "${MAIN_MODEL}" \
+    --persona "${PERSONA}" \
+    "${EXTRA[@]}"

examples/online_serving/joyvl_interaction/scripts/start_model.sh

@@ -0,0 +1,22 @@
+#!/bin/bash
+# Serve the JoyVL interaction model as a plain VLM (NOT --omni: it is a
+# standard Qwen3-VL autoregressive model, not an omni/diffusion model).
+set -euo pipefail
+
+MODEL="${MODEL:-jdopensource/JoyAI-VL-Interaction-Preview}"
+SERVED_NAME="${SERVED_NAME:-JoyAI-VL-Interaction-Preview}"
+PORT="${PORT:-8061}"
+GPU="${GPU:-0}"
+# Covers the default short-term window (chunk_frames=100) + key_frames=0 (all frames
+# summarized). Lower MAX_MODEL_LEN / IMAGE_LIMIT (and --chunk-frames) together for ~24GB.
+MAX_MODEL_LEN="${MAX_MODEL_LEN:-131072}"
+IMAGE_LIMIT="${IMAGE_LIMIT:-256}"
+
+CUDA_VISIBLE_DEVICES="${GPU}" vllm serve "${MODEL}" \
+    --served-model-name "${SERVED_NAME}" \
+    --port "${PORT}" \
+    --max-model-len "${MAX_MODEL_LEN}" \
+    --gpu-memory-utilization 0.85 \
+    --enable-prefix-caching \
+    --enable-chunked-prefill \
+    --limit-mm-per-prompt "{\"image\":${IMAGE_LIMIT},\"video\":1}"

examples/online_serving/joyvl_interaction/scripts/start_tts.sh

@@ -0,0 +1,25 @@
+#!/bin/bash
+# Serve Qwen3-TTS (vLLM-Omni) and run the webui<->TTS bridge.
+# Then start the webui with:  TTS_URL=ws://127.0.0.1:8092/v1/tts bash scripts/start_server.sh
+set -euo pipefail
+
+TTS_MODEL="${TTS_MODEL:-Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice}"
+TTS_PORT="${TTS_PORT:-8091}"
+BRIDGE_PORT="${BRIDGE_PORT:-8092}"
+GPU="${GPU:-1}"
+VOICE="${VOICE:-vivian}"
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
+
+CUDA_VISIBLE_DEVICES="${GPU}" vllm-omni serve "${TTS_MODEL}" \
+    --deploy-config vllm_omni/deploy/qwen3_tts.yaml \
+    --host 0.0.0.0 --port "${TTS_PORT}" \
+    --gpu-memory-utilization 0.9 --trust-remote-code --omni &
+TTS_PID=$!
+trap 'kill ${TTS_PID} 2>/dev/null || true' EXIT
+
+# wait for the TTS server, then run the bridge in the foreground
+until curl -s "http://127.0.0.1:${TTS_PORT}/health" >/dev/null 2>&1; do sleep 2; done
+python "${SCRIPT_DIR}/../bridges/tts_bridge.py" \
+    --port "${BRIDGE_PORT}" \
+    --backend-url "ws://127.0.0.1:${TTS_PORT}/v1/audio/speech/stream" \
+    --voice "${VOICE}"

examples/online_serving/text_to_speech/README.md

@@ -504,6 +504,35 @@ python qwen3_tts/streaming_speech_client.py --text "Hello world. How are you? I
 python qwen3_tts/streaming_speech_client.py --text "..." --simulate-stt --stt-delay 0.1
 ```
 
+To receive word-level timestamps, launch the server with a forced aligner:
+```bash
+vllm-omni serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
+    --omni \
+    --deploy-config vllm_omni/deploy/qwen3_tts.yaml \
+    --trust-remote-code \
+    --forced-aligner Qwen/Qwen3-ForcedAligner-0.6B
+```
+Then request PCM JSON sidecar chunks:
+```bash
+python qwen3_tts/streaming_speech_client.py \
+    --text "Hello world. How are you?" \
+    --stream-audio \
+    --response-format pcm \
+    --word-timestamps
+```
+The client writes one PCM file per sentence and a matching
+`sentence_XXX_timestamps.json` sidecar.
+
+To *see* the alignment instead of reading a JSON sidecar, run the
+word-timestamp Gradio demo (server must be launched with `--forced-aligner`):
+```bash
+python qwen3_tts/word_timestamps_demo.py --api-base http://localhost:8091
+```
+Each sentence's audio plays in an `<audio>` element while its text is rendered
+as inline word spans; the current word highlights as `audio.currentTime`
+crosses each `start_ms`. The **Stop (barge-in)** button cuts playback and
+reports the last-spoken word, useful for the voice-agent barge-in case.
+
 ### Gradio demos
 ```bash
 ./qwen3_tts/run_gradio_demo.sh                              # CustomVoice (default)

examples/online_serving/text_to_speech/qwen3_tts/streaming_speech_client.py

@@ -18,6 +18,21 @@ Usage:
         --text "Hello world. How are you? I am fine." \
         --simulate-stt --stt-delay 0.1
 
+    # Receive JSON sidecar chunks with word-level timestamps
+    python streaming_speech_client.py \
+        --text "Hello world. How are you?" \
+        --stream-audio \
+        --response-format pcm \
+        --word-timestamps
+
+    # Also save every received audio.chunk frame as individual PCM/WAV files
+    python streaming_speech_client.py \
+        --text "Hello world. How are you?" \
+        --stream-audio \
+        --response-format pcm \
+        --word-timestamps \
+        --save-chunks
+
     # VoiceDesign task
     python streaming_speech_client.py \
         --text "Today is a great day. The weather is nice." \
@@ -37,8 +52,10 @@ Requirements:
 
 import argparse
 import asyncio
+import base64
 import json
 import os
+import wave
 
 try:
     import websockets
@@ -47,6 +64,37 @@ except ImportError:
     raise SystemExit(1)
 
 
+def save_audio_file(output_dir: str, sentence_index: int, response_format: str, chunks: list[bytes]) -> str:
+    audio_bytes = b"".join(chunks)
+    if response_format == "pcm":
+        raw_path = os.path.join(output_dir, f"sentence_{sentence_index:03d}.pcm")
+        with open(raw_path, "wb") as f:
+            f.write(audio_bytes)
+
+        wav_path = os.path.join(output_dir, f"sentence_{sentence_index:03d}.wav")
+        save_pcm_wav(wav_path, audio_bytes)
+        return wav_path
+
+    filename = os.path.join(output_dir, f"sentence_{sentence_index:03d}.{response_format}")
+    with open(filename, "wb") as f:
+        f.write(audio_bytes)
+    return filename
+
+
+def save_pcm_wav(path: str, audio_bytes: bytes, sample_rate: int = 24000) -> None:
+    with wave.open(path, "wb") as wf:
+        wf.setnchannels(1)
+        wf.setsampwidth(2)
+        wf.setframerate(sample_rate)
+        wf.writeframes(audio_bytes)
+
+
+def timestamp_words_text(timestamps: list[dict] | None) -> str:
+    if not timestamps:
+        return ""
+    return " ".join(str(item.get("word", "")) for item in timestamps if item.get("word"))
+
+
 async def stream_tts(
     url: str,
     text: str,
@@ -54,6 +102,7 @@ async def stream_tts(
     output_dir: str,
     simulate_stt: bool = False,
     stt_delay: float = 0.1,
+    save_chunks: bool = False,
 ) -> None:
     """Connect to the streaming TTS endpoint and process audio responses."""
     os.makedirs(output_dir, exist_ok=True)
@@ -99,8 +148,14 @@ async def stream_tts(
         sender_task = asyncio.create_task(send_text())
 
         response_format = config.get("response_format", "wav")
+        word_timestamps = bool(config.get("word_timestamps", False))
         current_sentence_index = 0
+        current_sentence_text = ""
         current_chunks: list[bytes] = []
+        current_timestamps: list[dict] = []
+        # Distinguish `null` (aligner failed) from `[]` (silence / no tokens):
+        # the sentence is a failure only if no frame ever carried timestamps.
+        alignment_frame_seen = False
 
         try:
             while True:
@@ -116,22 +171,94 @@ async def stream_tts(
 
                     if msg_type == "audio.start":
                         current_sentence_index = msg["sentence_index"]
+                        current_sentence_text = msg.get("sentence_text", "")
                         current_chunks = []
-                        print(f"  [sentence {msg['sentence_index']}] Generating: {msg['sentence_text']!r}")
+                        current_timestamps = []
+                        alignment_frame_seen = False
+                        print(f"  [sentence {msg['sentence_index']}] Generating: {current_sentence_text!r}")
+                    elif msg_type == "audio.chunk":
+                        audio = base64.b64decode(msg["audio_b64"])
+                        current_chunks.append(audio)
+                        chunk_timestamps = msg.get("timestamps")
+                        # Sentence-level alignment sends a trailing timestamp-only
+                        # frame with empty audio; don't write an empty chunk file.
+                        if save_chunks and audio:
+                            chunk_base = os.path.join(
+                                output_dir,
+                                f"sentence_{msg['sentence_index']:03d}_chunk_{msg['chunk_id']:03d}",
+                            )
+                            chunk_pcm = f"{chunk_base}.pcm"
+                            chunk_wav = f"{chunk_base}.wav"
+                            chunk_json = f"{chunk_base}_timestamps.json"
+                            with open(chunk_pcm, "wb") as f:
+                                f.write(audio)
+                            save_pcm_wav(chunk_wav, audio, int(msg.get("sample_rate", 24000)))
+                            with open(chunk_json, "w", encoding="utf-8") as f:
+                                json.dump(
+                                    {
+                                        "sentence_index": msg["sentence_index"],
+                                        "chunk_id": msg["chunk_id"],
+                                        "chunk_start_ms": msg.get("chunk_start_ms"),
+                                        "chunk_end_ms": msg.get("chunk_end_ms"),
+                                        "sentence_text": current_sentence_text,
+                                        "covered_text": timestamp_words_text(chunk_timestamps),
+                                        "timestamps": chunk_timestamps,
+                                    },
+                                    f,
+                                    ensure_ascii=False,
+                                    indent=2,
+                                )
+                        if chunk_timestamps is None:
+                            print(
+                                f"  [sentence {msg['sentence_index']}] chunk {msg['chunk_id']}: "
+                                f"{len(audio)} bytes, timestamps unavailable"
+                            )
+                        else:
+                            alignment_frame_seen = True
+                            current_timestamps.extend(chunk_timestamps)
+                            print(
+                                f"  [sentence {msg['sentence_index']}] chunk {msg['chunk_id']}: "
+                                f"{len(audio)} bytes, {len(chunk_timestamps)} timestamp(s)"
+                            )
+                        if save_chunks and audio:
+                            print(f"    saved chunk -> {chunk_wav}")
                     elif msg_type == "audio.done":
-                        filename = os.path.join(
+                        filename = save_audio_file(
                             output_dir,
-                            f"sentence_{msg['sentence_index']:03d}.{response_format}",
+                            msg["sentence_index"],
+                            response_format,
+                            current_chunks,
                         )
-                        with open(filename, "wb") as f:
-                            f.write(b"".join(current_chunks))
+                        if word_timestamps:
+                            # `null` on failure, else the accumulated list (`[]` for silence).
+                            sentence_timestamps = current_timestamps if alignment_frame_seen else None
+                            ts_filename = os.path.join(
+                                output_dir,
+                                f"sentence_{msg['sentence_index']:03d}_timestamps.json",
+                            )
+                            with open(ts_filename, "w", encoding="utf-8") as f:
+                                json.dump(
+                                    {
+                                        "sentence_index": msg["sentence_index"],
+                                        "sentence_text": current_sentence_text,
+                                        "covered_text": timestamp_words_text(sentence_timestamps),
+                                        "timestamps": sentence_timestamps,
+                                    },
+                                    f,
+                                    ensure_ascii=False,
+                                    indent=2,
+                                )
                         print(
                             f"  [sentence {msg['sentence_index']}] Done"
                             f" bytes={msg.get('total_bytes', len(b''.join(current_chunks)))}"
                             f" error={msg.get('error', False)}"
                             f" -> {filename}"
                         )
+                        if word_timestamps:
+                            print(f"  [sentence {msg['sentence_index']}] Timestamps -> {ts_filename}")
                         current_chunks = []
+                        current_timestamps = []
+                        alignment_frame_seen = False
                     elif msg_type == "session.done":
                         print(f"\nSession complete: {msg['total_sentences']} sentence(s) generated")
                         break
@@ -189,6 +316,16 @@ def main():
         action="store_true",
         help="Receive one or more PCM chunks per sentence (requires --response-format pcm)",
     )
+    parser.add_argument(
+        "--word-timestamps",
+        action="store_true",
+        help="Receive JSON sidecar chunks with word-level timestamps (requires --stream-audio --response-format pcm)",
+    )
+    parser.add_argument(
+        "--save-chunks",
+        action="store_true",
+        help="Save each audio.chunk frame as sentence_XXX_chunk_YYY.{pcm,wav,json}",
+    )
     parser.add_argument("--speed", type=float, default=1.0, help="Playback speed (0.25-4.0)")
     parser.add_argument("--max-new-tokens", type=int, default=None, help="Max tokens")
 
@@ -238,6 +375,10 @@ def main():
             config[key] = val
     if args.stream_audio:
         config["stream_audio"] = True
+    if args.word_timestamps:
+        config["word_timestamps"] = True
+        config["stream_audio"] = True
+        config["response_format"] = "pcm"
     if args.x_vector_only_mode:
         config["x_vector_only_mode"] = True
 
@@ -249,6 +390,7 @@ def main():
             output_dir=args.output_dir,
             simulate_stt=args.simulate_stt,
             stt_delay=args.stt_delay,
+            save_chunks=args.save_chunks,
         )
     )
 

examples/online_serving/text_to_speech/qwen3_tts/word_timestamps_demo.py

@@ -0,0 +1,457 @@
+"""Gradio demo for streaming TTS word-level timestamps.
+
+Connects to the WebSocket endpoint ``/v1/audio/speech/stream`` with
+``word_timestamps: true`` and visualizes the alignment: each sentence's
+audio plays in a native ``<audio>`` element while its text is rendered as
+inline word spans, the current word highlighting as ``audio.currentTime``
+crosses each ``start_ms``.
+
+A "Stop (barge-in)" button cuts playback and reports the last-spoken word,
+demonstrating the voice-agent barge-in use case.
+
+Timestamps are sentence-level: audio streams in real time, then the forced
+aligner runs once per sentence and the complete word timestamps arrive in a
+trailing frame. The demo therefore plays one sentence block at a time.
+
+Usage:
+    # Launch the server with a forced aligner first, e.g.:
+    #   vllm-omni serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni \
+    #       --deploy-config vllm_omni/deploy/qwen3_tts.yaml --trust-remote-code \
+    #       --forced-aligner Qwen/Qwen3-ForcedAligner-0.6B
+    # Then:
+    python word_timestamps_demo.py --api-base http://localhost:8000
+"""
+
+import argparse
+import asyncio
+import json
+import logging
+
+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
+try:
+    import websockets
+except ImportError:
+    raise ImportError("websockets is required to run this demo. Install it with: pip install websockets") from None
+from fastapi import FastAPI, WebSocket, WebSocketDisconnect
+from tts_common import (
+    PCM_SAMPLE_RATE,
+    SUPPORTED_LANGUAGES,
+    TASK_TYPES,
+    add_common_args,
+    fetch_voices,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def _upstream_ws_url(api_base: str) -> str:
+    """Map an http(s) API base to its WebSocket streaming endpoint."""
+    base = api_base.rstrip("/")
+    if base.startswith("https://"):
+        base = "wss://" + base[len("https://") :]
+    elif base.startswith("http://"):
+        base = "ws://" + base[len("http://") :]
+    else:
+        base = "ws://" + base
+    return base + "/v1/audio/speech/stream"
+
+
+# ── Front-end ────────────────────────────────────────────────────────
+PLAYER_HTML = """
+<div id="wt-root">
+  <div style="display:flex; align-items:center; gap:10px;">
+    <div id="wt-dot" style="width:10px;height:10px;border-radius:50%;background:#ccc;flex-shrink:0;"></div>
+    <span id="wt-status" style="font-weight:600;font-size:1.05em;">Ready</span>
+    <button id="wt-stop" onclick="window.wtStop()"
+      style="display:none; margin-left:auto; padding:5px 16px; border-radius:6px; border:1px solid #EF5552;
+             background:#fff; color:#EF5552; cursor:pointer; font-size:0.85em;">Stop (barge-in)</button>
+  </div>
+  <div id="wt-barge" style="display:none; margin-top:10px; padding:8px 12px; border-radius:6px;
+       background:#fff5f5; border:1px solid #f3c2c2; color:#b02a2a; font-size:0.9em;"></div>
+  <div id="wt-sentences" style="margin-top:12px; display:flex; flex-direction:column; gap:14px;"></div>
+</div>
+"""
+
+
+def _build_head_js(sample_rate: int) -> str:
+    """Browser-side player: WebSocket -> per-sentence <audio> + word highlight."""
+    return f"""
+<script>
+const WT_SR = {sample_rate};
+let wtWS = null, wtQueue = [], wtPlaying = false, wtCur = null, wtLastWord = "";
+
+function wtSet(text, color) {{
+    const s = document.getElementById('wt-status');
+    const d = document.getElementById('wt-dot');
+    if (s) s.textContent = text;
+    if (d) d.style.background = color || '#ccc';
+}}
+
+function wtStopBtn(show) {{
+    const b = document.getElementById('wt-stop');
+    if (b) b.style.display = show ? 'inline-block' : 'none';
+}}
+
+// Build a WAV Blob (16-bit mono PCM) from concatenated Int16 samples.
+function wtPcmToWav(int16, sr) {{
+    const n = int16.length;
+    const buf = new ArrayBuffer(44 + n * 2);
+    const v = new DataView(buf);
+    const w = (off, s) => {{ for (let i = 0; i < s.length; i++) v.setUint8(off + i, s.charCodeAt(i)); }};
+    w(0, 'RIFF'); v.setUint32(4, 36 + n * 2, true); w(8, 'WAVE');
+    w(12, 'fmt '); v.setUint32(16, 16, true); v.setUint16(20, 1, true); v.setUint16(22, 1, true);
+    v.setUint32(24, sr, true); v.setUint32(28, sr * 2, true); v.setUint16(32, 2, true); v.setUint16(34, 16, true);
+    w(36, 'data'); v.setUint32(40, n * 2, true);
+    let off = 44;
+    for (let i = 0; i < n; i++) {{ v.setInt16(off, int16[i], true); off += 2; }}
+    return new Blob([buf], {{ type: 'audio/wav' }});
+}}
+
+function wtConcat(chunks) {{
+    let total = 0;
+    for (const c of chunks) total += c.length;
+    const out = new Int16Array(total);
+    let off = 0;
+    for (const c of chunks) {{ out.set(c, off); off += c.length; }}
+    return out;
+}}
+
+// Render one finished sentence: <audio> + inline word spans.
+function wtRenderSentence(sentence) {{
+    const wrap = document.getElementById('wt-sentences');
+    if (!wrap) return;
+    const block = document.createElement('div');
+    block.style.cssText = 'border:1px solid #e3e8ef;border-radius:8px;padding:10px 12px;';
+
+    const head = document.createElement('div');
+    head.style.cssText = 'font-size:0.75em;color:#888;margin-bottom:6px;';
+    head.textContent = 'Sentence ' + sentence.index +
+        (sentence.timestamps ? ' (' + sentence.timestamps.length + ' words)' : ' (no timestamps)');
+    block.appendChild(head);
+
+    const textEl = document.createElement('div');
+    textEl.style.cssText = 'font-size:1.15em;line-height:1.8;margin-bottom:8px;';
+    if (sentence.timestamps && sentence.timestamps.length) {{
+        sentence.spans = [];
+        for (const ts of sentence.timestamps) {{
+            const span = document.createElement('span');
+            span.textContent = ts.word;
+            span.dataset.start = ts.start_ms;
+            span.dataset.end = ts.end_ms;
+            span.style.cssText = 'padding:1px 2px;border-radius:3px;transition:background 0.1s;';
+            textEl.appendChild(span);
+            // Keep spacing for spaced languages; CJK chars sit flush.
+            textEl.appendChild(document.createTextNode(' '));
+            sentence.spans.push(span);
+        }}
+    }} else {{
+        textEl.textContent = sentence.text || '(timestamps unavailable)';
+    }}
+    block.appendChild(textEl);
+
+    const audio = document.createElement('audio');
+    audio.controls = true;
+    audio.style.cssText = 'width:100%;';
+    audio.src = URL.createObjectURL(wtPcmToWav(wtConcat(sentence.chunks), WT_SR));
+    block.appendChild(audio);
+    sentence.audio = audio;
+
+    // Highlight the current word. Driven by requestAnimationFrame (~60fps)
+    // rather than the <audio> 'timeupdate' event, which only fires ~4x/sec
+    // and is too coarse for short words. We light the latest word whose
+    // start_ms has passed, so there is always exactly one active word (no
+    // inter-word flicker) and wtLastWord stays well-defined for barge-in.
+    let raf = null;
+    function highlight() {{
+        if (!sentence.spans) return;
+        const t = audio.currentTime * 1000;
+        let idx = -1;
+        for (let i = 0; i < sentence.spans.length; i++) {{
+            if (t >= +sentence.spans[i].dataset.start) idx = i;
+            else break;
+        }}
+        sentence.spans.forEach((sp, i) => {{
+            const on = i === idx;
+            sp.style.background = on ? '#ffe680' : 'transparent';
+            sp.style.fontWeight = on ? '700' : '400';
+        }});
+        if (idx >= 0) wtLastWord = sentence.spans[idx].textContent;
+    }}
+    function loop() {{
+        highlight();
+        if (!audio.paused && !audio.ended) raf = requestAnimationFrame(loop);
+    }}
+    audio.addEventListener('play', () => {{
+        wtCur = audio;
+        wtSet('Playing sentence ' + sentence.index + '...', '#64dd17');
+        cancelAnimationFrame(raf);
+        raf = requestAnimationFrame(loop);
+    }});
+    audio.addEventListener('pause', () => cancelAnimationFrame(raf));
+    audio.addEventListener('seeked', highlight);
+    audio.addEventListener('ended', () => {{
+        cancelAnimationFrame(raf);
+        highlight();
+        wtPlaying = false;
+        wtPlayNext();
+    }});
+
+    wrap.appendChild(block);
+    wtQueue.push(audio);
+    wtPlayNext();
+}}
+
+function wtPlayNext() {{
+    if (wtPlaying) return;
+    const next = wtQueue.shift();
+    if (!next) {{
+        if (wtWS === null) {{ wtSet('Done', '#64dd17'); wtStopBtn(false); }}
+        return;
+    }}
+    wtPlaying = true;
+    next.play().catch(() => {{ wtPlaying = false; }});
+}}
+
+window.wtStop = function() {{
+    // Barge-in: cut playback, stop generation, report the last-spoken word.
+    if (wtCur) {{ try {{ wtCur.pause(); }} catch (e) {{}} }}
+    wtQueue = [];
+    wtPlaying = false;
+    if (wtWS) {{ try {{ wtWS.close(); }} catch (e) {{}} wtWS = null; }}
+    wtSet('Stopped (barge-in)', '#999');
+    wtStopBtn(false);
+    const b = document.getElementById('wt-barge');
+    if (b) {{
+        b.style.display = 'block';
+        b.textContent = wtLastWord
+            ? 'Barge-in at word: "' + wtLastWord + '"  (a voice agent would resume the conversation from here)'
+            : 'Barge-in before any word was spoken.';
+    }}
+}};
+
+window.wtGenerate = function(payload) {{
+    if (wtWS) {{ try {{ wtWS.close(); }} catch (e) {{}} }}
+    wtQueue = []; wtPlaying = false; wtCur = null; wtLastWord = "";
+    document.getElementById('wt-sentences').innerHTML = '';
+    const barge = document.getElementById('wt-barge');
+    if (barge) {{ barge.style.display = 'none'; barge.textContent = ''; }}
+    wtStopBtn(true);
+    wtSet('Connecting...', '#4A90D9');
+
+    const proto = location.protocol === 'https:' ? 'wss:' : 'ws:';
+    const ws = new WebSocket(proto + '//' + location.host + '/proxy/stream');
+    wtWS = ws;
+    const sentences = {{}};
+
+    ws.onopen = () => {{
+        wtSet('Streaming...', '#4A90D9');
+        ws.send(JSON.stringify(Object.assign({{ type: 'session.config' }}, payload.config)));
+        ws.send(JSON.stringify({{ type: 'input.text', text: payload.text }}));
+        ws.send(JSON.stringify({{ type: 'input.done' }}));
+    }};
+
+    ws.onmessage = (ev) => {{
+        let msg;
+        try {{ msg = JSON.parse(ev.data); }} catch (e) {{ return; }}
+        if (msg.type === 'audio.start') {{
+            sentences[msg.sentence_index] = {{
+                index: msg.sentence_index, text: msg.sentence_text || '',
+                chunks: [], timestamps: null
+            }};
+        }} else if (msg.type === 'audio.chunk') {{
+            const s = sentences[msg.sentence_index];
+            if (!s) return;
+            if (msg.audio_b64) {{
+                const bin = atob(msg.audio_b64);
+                const bytes = new Uint8Array(bin.length);
+                for (let i = 0; i < bin.length; i++) bytes[i] = bin.charCodeAt(i);
+                const usable = bytes.length - (bytes.length % 2);
+                if (usable > 0) s.chunks.push(new Int16Array(bytes.buffer, 0, usable / 2));
+            }}
+            if (msg.timestamps != null) s.timestamps = msg.timestamps;
+        }} else if (msg.type === 'audio.done') {{
+            const s = sentences[msg.sentence_index];
+            if (s) wtRenderSentence(s);
+        }} else if (msg.type === 'session.done') {{
+            wtWS = null;
+            try {{ ws.close(); }} catch (e) {{}}
+            wtPlayNext();
+        }} else if (msg.type === 'error') {{
+            wtSet('Error: ' + (msg.message || 'unknown'), '#EF5552');
+        }}
+    }};
+
+    ws.onerror = () => wtSet('WebSocket error', '#EF5552');
+    ws.onclose = () => {{ if (wtWS === ws) wtWS = null; }};
+}};
+</script>
+"""
+
+
+def create_app(api_base: str) -> FastAPI:
+    """FastAPI app with a same-origin WebSocket proxy + Gradio UI."""
+    fastapi_app = FastAPI()
+    upstream_url = _upstream_ws_url(api_base)
+
+    @fastapi_app.websocket("/proxy/stream")
+    async def proxy_stream(client_ws: WebSocket):
+        await client_ws.accept()
+        try:
+            async with websockets.connect(upstream_url, max_size=None) as upstream:
+
+                async def client_to_upstream():
+                    try:
+                        while True:
+                            await upstream.send(await client_ws.receive_text())
+                    except (WebSocketDisconnect, Exception):
+                        await upstream.close()
+
+                async def upstream_to_client():
+                    try:
+                        async for msg in upstream:
+                            if isinstance(msg, bytes):
+                                await client_ws.send_bytes(msg)
+                            else:
+                                await client_ws.send_text(msg)
+                    except Exception:
+                        pass
+                    finally:
+                        try:
+                            await client_ws.close()
+                        except Exception:
+                            pass
+
+                await asyncio.gather(client_to_upstream(), upstream_to_client())
+        except Exception as exc:
+            logger.exception("WS proxy error")
+            try:
+                await client_ws.send_text(json.dumps({"type": "error", "message": f"proxy: {exc}"}))
+                await client_ws.close()
+            except Exception:
+                pass
+
+    voices = fetch_voices(api_base)
+
+    with gr.Blocks(title="Qwen3-TTS Word Timestamps") as demo:
+        gr.HTML(f"""
+        <div style="margin-bottom:8px;">
+          <h1 style="margin:0; font-size:1.5em;">Qwen3-TTS Streaming Word Timestamps</h1>
+          <span style="font-size:0.85em; color:#666;">
+            Served by <a href="https://github.com/vllm-project/vllm-omni" target="_blank">vLLM-Omni</a>
+            &nbsp;&middot;&nbsp; <code>{api_base}</code> &nbsp;&middot;&nbsp;
+            launch the server with <code>--forced-aligner</code>.
+          </span>
+        </div>
+        """)
+
+        with gr.Row():
+            with gr.Column(scale=3):
+                text_input = gr.Textbox(
+                    label="Text to Synthesize",
+                    value="Hello world. How are you today? I am doing great.",
+                    lines=3,
+                )
+                with gr.Row():
+                    task_type = gr.Radio(choices=TASK_TYPES, value="CustomVoice", label="Task Type", scale=2)
+                    language = gr.Dropdown(choices=SUPPORTED_LANGUAGES, value="Auto", label="Language", scale=1)
+                voice = gr.Dropdown(
+                    choices=voices,
+                    value=voices[0] if voices else None,
+                    label="Speaker",
+                    allow_custom_value=True,
+                )
+                instructions = gr.Textbox(
+                    label="Instructions",
+                    placeholder="Optional style/emotion (required for VoiceDesign)",
+                    lines=1,
+                )
+                with gr.Column(visible=False) as ref_group:
+                    ref_audio_url = gr.Textbox(
+                        label="Reference Audio URL (Base task)",
+                        placeholder="https://example.com/reference.wav",
+                        lines=1,
+                    )
+                    ref_text = gr.Textbox(label="Reference Transcript", lines=1)
+                generate_btn = gr.Button("Generate + Align", variant="primary", size="lg")
+
+            with gr.Column(scale=2):
+                gr.HTML(value=PLAYER_HTML)
+
+        hidden_payload = gr.Textbox(visible=False, elem_id="wt-payload")
+
+        def on_task_change(tt):
+            return gr.update(visible=(tt == "Base"))
+
+        task_type.change(fn=on_task_change, inputs=[task_type], outputs=[ref_group])
+
+        def on_generate(text, tt, voice_v, lang_v, instr_v, ref_url, ref_t):
+            if not text or not text.strip():
+                raise gr.Error("Please enter text to synthesize.")
+            config: dict = {
+                "task_type": tt,
+                "word_timestamps": True,
+                "stream_audio": True,
+                "response_format": "pcm",
+            }
+            if lang_v:
+                config["language"] = lang_v
+            if tt == "CustomVoice":
+                if voice_v:
+                    config["voice"] = voice_v
+                if instr_v and instr_v.strip():
+                    config["instructions"] = instr_v.strip()
+            elif tt == "VoiceDesign":
+                if not instr_v or not instr_v.strip():
+                    raise gr.Error("VoiceDesign requires voice style instructions.")
+                config["instructions"] = instr_v.strip()
+            elif tt == "Base":
+                if not ref_url or not ref_url.strip():
+                    raise gr.Error("Base (voice clone) needs a reference audio URL.")
+                config["ref_audio"] = ref_url.strip()
+                if ref_t and ref_t.strip():
+                    config["ref_text"] = ref_t.strip()
+            return json.dumps({"config": config, "text": text.strip()})
+
+        generate_btn.click(
+            fn=on_generate,
+            inputs=[text_input, task_type, voice, language, instructions, ref_audio_url, ref_text],
+            outputs=[hidden_payload],
+        ).then(
+            fn=lambda p: p,
+            inputs=[hidden_payload],
+            outputs=[hidden_payload],
+            js="(p) => { if (p && p.trim()) window.wtGenerate(JSON.parse(p)); return p; }",
+        )
+
+        demo.queue()
+
+    return gr.mount_gradio_app(
+        fastapi_app,
+        demo,
+        path="/",
+        head=_build_head_js(PCM_SAMPLE_RATE),
+    )
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Gradio demo for streaming TTS word-level timestamps.")
+    add_common_args(parser)
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    logging.basicConfig(level=logging.INFO)
+    print(f"Connecting to vLLM server at: {args.api_base}")
+    app = create_app(args.api_base)
+
+    import uvicorn
+
+    uvicorn.run(app, host=args.host, port=args.port)
+
+
+if __name__ == "__main__":
+    main()

examples/online_serving/text_to_video/README.md

@@ -10,6 +10,7 @@ This example demonstrates how to deploy text-to-video models for online video ge
 | Wan2.1 T2V (14B) | `Wan-AI/Wan2.1-T2V-14B-Diffusers` |
 | Wan2.2 T2V | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` |
 | LTX-2 | `Lightricks/LTX-2` |
+| Helios (Base / Mid / Distilled) | `BestWishYsh/Helios-Base`, `Helios-Mid`, `Helios-Distilled` |
 
 ## Wan2.2 T2V
 
@@ -323,3 +324,45 @@ curl -sS -X POST http://localhost:8098/v1/videos \
   -F "guidance_scale=3.0" \
   -F "seed=42"
 ```
+
+## Helios
+
+Helios ships three variants (`Helios-Base`, `Helios-Mid`, `Helios-Distilled`) that
+share the same server launch. Variant-specific knobs (declared in
+`vllm_omni/model_extras/helios.py`) are sent per request through the generic
+`extra_params` JSON form field — no per-model server flags required.
+
+### Start Server
+
+```bash
+vllm serve BestWishYsh/Helios-Base --omni --port 8098
+# or: MODEL=BestWishYsh/Helios-Mid bash run_server_helios.sh
+```
+
+### Send Requests (curl)
+
+```bash
+# Helios-Base (Stage 1 only)
+bash run_curl_helios.sh
+
+# Helios-Mid (Stage 2 pyramid + CFG-Zero*)
+PRESET=mid-stage2 MODEL=BestWishYsh/Helios-Mid bash run_curl_helios.sh
+
+# Helios-Distilled (Stage 2 pyramid + DMD, few-step)
+PRESET=distilled MODEL=BestWishYsh/Helios-Distilled bash run_curl_helios.sh
+```
+
+The `mid-stage2` and `distilled` presets attach an `extra_params` field, e.g. for Helios-Distilled:
+
+```bash
+curl -sS -X POST http://localhost:8098/v1/videos \
+  -H "Accept: application/json" \
+  -F "prompt=A dynamic time-lapse of scenery rushing past the window of a speeding train." \
+  -F "model=BestWishYsh/Helios-Distilled" \
+  -F "size=640x384" \
+  -F "num_frames=99" \
+  -F "fps=16" \
+  -F "guidance_scale=1.0" \
+  -F "seed=42" \
+  -F 'extra_params={"is_enable_stage2": true, "pyramid_num_inference_steps_list": [2, 2, 2], "is_amplify_first_chunk": true}'
+```

examples/online_serving/text_to_video/run_curl_helios.sh

@@ -0,0 +1,93 @@
+#!/bin/bash
+# Helios text-to-video curl example using the async video job API.
+#
+# Helios-specific knobs (declared in vllm_omni/model_extras/helios.py) are passed
+# through the generic `extra_params` JSON form field. Select a variant via PRESET:
+#   PRESET=base        -> Helios-Base, Stage 1 only (default)
+#   PRESET=mid-stage2  -> Helios-Mid, Stage 2 pyramid + CFG-Zero*
+#   PRESET=distilled   -> Helios-Distilled, Stage 2 pyramid + DMD (few-step)
+
+set -euo pipefail
+
+BASE_URL="${BASE_URL:-http://localhost:8098}"
+MODEL="${MODEL:-BestWishYsh/Helios-Base}"
+PROMPT="${PROMPT:-A dynamic time-lapse of scenery rushing past the window of a speeding train.}"
+PRESET="${PRESET:-base}"
+OUTPUT_PATH="${OUTPUT_PATH:-helios_t2v_${PRESET}.mp4}"
+POLL_INTERVAL="${POLL_INTERVAL:-2}"
+
+case "${PRESET}" in
+  base)
+    EXTRA_PARAMS=""
+    GUIDANCE_SCALE="5.0"
+    ;;
+  mid-stage2)
+    EXTRA_PARAMS='{"is_enable_stage2": true, "pyramid_num_inference_steps_list": [20, 20, 20], "use_cfg_zero_star": true, "use_zero_init": true, "zero_steps": 1}'
+    GUIDANCE_SCALE="5.0"
+    ;;
+  distilled)
+    EXTRA_PARAMS='{"is_enable_stage2": true, "pyramid_num_inference_steps_list": [2, 2, 2], "is_amplify_first_chunk": true}'
+    GUIDANCE_SCALE="1.0"
+    ;;
+  *)
+    echo "Unknown PRESET '${PRESET}' (expected base|mid-stage2|distilled)"
+    exit 1
+    ;;
+esac
+
+create_args=(
+  -sS -X POST "${BASE_URL}/v1/videos"
+  -H "Accept: application/json"
+  -F "prompt=${PROMPT}"
+  -F "model=${MODEL}"
+  -F "size=640x384"
+  -F "num_frames=99"
+  -F "fps=16"
+  -F "num_inference_steps=50"
+  -F "guidance_scale=${GUIDANCE_SCALE}"
+  -F "seed=42"
+)
+if [ -n "${EXTRA_PARAMS}" ]; then
+  create_args+=(-F "extra_params=${EXTRA_PARAMS}")
+fi
+
+create_response=$(curl "${create_args[@]}")
+
+video_id="$(echo "${create_response}" | jq -r '.id')"
+if [ -z "${video_id}" ] || [ "${video_id}" = "null" ]; then
+  echo "Failed to create video job:"
+  echo "${create_response}" | jq .
+  exit 1
+fi
+
+echo "Created video job ${video_id} (preset=${PRESET})"
+echo "${create_response}" | jq .
+
+while true; do
+  status_response="$(curl -sS "${BASE_URL}/v1/videos/${video_id}")"
+  status="$(echo "${status_response}" | jq -r '.status')"
+
+  case "${status}" in
+    queued|in_progress)
+      echo "Video job ${video_id} status: ${status}"
+      sleep "${POLL_INTERVAL}"
+      ;;
+    completed)
+      echo "${status_response}" | jq .
+      break
+      ;;
+    failed)
+      echo "Video generation failed:"
+      echo "${status_response}" | jq .
+      exit 1
+      ;;
+    *)
+      echo "Unexpected status response:"
+      echo "${status_response}" | jq .
+      exit 1
+      ;;
+  esac
+done
+
+curl -sS -L "${BASE_URL}/v1/videos/${video_id}/content" -o "${OUTPUT_PATH}"
+echo "Saved video to ${OUTPUT_PATH}"

examples/online_serving/text_to_video/run_server_helios.sh

@@ -0,0 +1,15 @@
+#!/bin/bash
+# Helios online serving startup script.
+# All three variants (Helios-Base / Helios-Mid / Helios-Distilled) share the same
+# server launch; variant-specific knobs are sent per-request via `extra_params`
+# (see run_curl_helios.sh).
+
+MODEL="${MODEL:-BestWishYsh/Helios-Base}"
+PORT="${PORT:-8098}"
+
+echo "Starting Helios server..."
+echo "Model: $MODEL"
+echo "Port: $PORT"
+
+vllm serve "$MODEL" --omni \
+    --port "$PORT"

mkdocs.yml

@@ -102,6 +102,7 @@ plugins:
         - "vllm_omni.entrypoints.async_diffusion"  # avoid importing vllm in mkdocs building
         - "vllm_omni.entrypoints.openai"  # avoid importing vllm in mkdocs building
         - "vllm_omni.entrypoints.openai.protocol"  # avoid importing vllm in mkdocs building
+        - "vllm_omni.experimental"  # avoid importing optional serving deps (openai, fastapi, uvicorn, httpx) in mkdocs building
   - mkdocstrings:
       handlers:
         python:

pyproject.toml

@@ -71,7 +71,13 @@ dev = [
 ]
 
 demo = [
-    "gradio>=6.7.0",
+    "opencv-python>=4.12.0.88",  # cli/stream_client frame reading
+    "requests>=2.28.0",  # cli/stream_client HTTP
+    "cn2an>=0.5.22",  # Chinese number normalization for TTS (e.g. "1次" -> "一次")
+]
+
+forced-aligner = [
+    "qwen-asr>=0.0.6",  # faithful Japanese/Korean word segmentation for the TTS forced aligner
 ]
 
 # Local-test-only extras
@@ -181,6 +187,7 @@ ignore = [
 # These files keep upstream conventions (star imports, long lines, math-symbol
 # argument names) so they stay easy to diff against upstream releases.
 "vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py" = ["E402", "E501", "F403", "F405", "N803"]
+"vllm_omni/experimental/fullduplex/joyvl/decision/prompts.py" = ["E501"]  # System/summarizer prompts are long prose lines
 
 [tool.ruff.lint.flake8-tidy-imports.banned-api]
 "librosa".msg = "The librosa module is banned, use vllm.multimodal helpers instead"
@@ -301,3 +308,6 @@ CANN = "CANN"
 thi = "thi"
 ND = "ND"
 byt5 = "byt5"
+# "seeked" is the standard HTML media element DOM event name (fires when a
+# seek completes), used by the word-timestamps Gradio demo. Not a typo.
+seeked = "seeked"

recipes/Helios/Helios.md

@@ -4,27 +4,27 @@
 
 - Vendor: Helios
 - Model: `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled`
-- Task: Text-to-video, image-to-video, and video-to-video generation
-- Mode: Offline inference with the bundled Helios example
+- Task: Text-to-video generation
+- Mode: Offline inference with the shared text_to_video example
 - Maintainer: Community
 
 ## When to use this recipe
 
 Use this recipe when you want a known-good starting point for running Helios
 video generation with vLLM-Omni. The concrete command below focuses on
-`BestWishYsh/Helios-Base` text-to-video generation on one NVIDIA H20 GPU. The
-same example script also supports Helios-Mid, Helios-Distilled, image-to-video,
-and video-to-video; see the linked example docs for those model-specific flags.
+`BestWishYsh/Helios-Base` text-to-video generation on one NVIDIA H20 GPU via the
+shared `text_to_video.py` example. The same example also supports Helios-Mid and
+Helios-Distilled through the generic `--extra-body` flag (see below).
+Image-to-video and video-to-video require image/video conditioning inputs and
+are not covered by the text-to-video example.
 
 ## References
 
 - Upstream repository: <https://github.com/PKU-YuanGroup/Helios>
 - Model weights:
   <https://huggingface.co/BestWishYsh/Helios-Base>
-- Upstream or canonical docs:
-  [`docs/user_guide/examples/offline_inference/helios.md`](../../docs/user_guide/examples/offline_inference/helios.md)
 - Related example under `examples/`:
-  [`examples/offline_inference/helios/README.md`](../../examples/offline_inference/helios/README.md)
+  [`examples/offline_inference/text_to_video/text_to_video.md`](../../examples/offline_inference/text_to_video/text_to_video.md)
 
 ## Hardware Support
 
@@ -46,11 +46,10 @@ and video-to-video; see the linked example docs for those model-specific flags.
 Run the baseline Helios-Base text-to-video example from the repository root:
 
 ```bash
-cd examples/offline_inference/helios
+cd examples/offline_inference/text_to_video
 
-python end2end.py \
+python text_to_video.py \
   --model BestWishYsh/Helios-Base \
-  --sample-type t2v \
   --prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \
   --guidance-scale 5.0 \
   --output helios_t2v_base.mp4
@@ -59,13 +58,12 @@ python end2end.py \
 To use cache-dit acceleration, run on a vLLM-Omni checkout that includes Helios cache-dit support and enable the cache backend:
 
 ```bash
-cd examples/offline_inference/helios
+cd examples/offline_inference/text_to_video
 
-python end2end.py \
+python text_to_video.py \
   --cache-backend cache_dit \
   --enable-cache-dit-summary \
   --model BestWishYsh/Helios-Base \
-  --sample-type t2v \
   --prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \
   --guidance-scale 5.0 \
   --output helios_t2v_base.mp4
@@ -85,22 +83,17 @@ Saved generated video to helios_t2v_base.mp4
 
 - `--model BestWishYsh/Helios-Base` selects the base Helios checkpoint used by
   this recipe.
-- `--sample-type t2v` selects text-to-video generation. Use `i2v` with
-  `--image-path` for image-to-video, or `v2v` with `--video-path` for
-  video-to-video.
 - `--guidance-scale 5.0` matches the Helios-Base recommendation. Use
   `--guidance-scale 1.0` for Helios-Distilled.
 - `--cache-backend cache_dit` enables the cache-dit acceleration path.
 - `--enable-cache-dit-summary` prints cache-dit summary information after
   diffusion forward passes.
-- No separate stage config is required for this offline recipe; the bundled
-  Helios example configures the pipeline through `end2end.py` arguments.
-- For Helios-Mid, add `--is-enable-stage2`,
-  `--pyramid-num-inference-steps-list 20 20 20`, `--use-cfg-zero-star`,
-  `--use-zero-init`, and `--zero-steps 1`.
-- For Helios-Distilled, add `--is-enable-stage2`,
-  `--pyramid-num-inference-steps-list 2 2 2`, and
-  `--is-amplify-first-chunk`.
+- No separate stage config is required for this offline recipe; the shared
+  `text_to_video.py` example configures the pipeline through its arguments.
+- Helios-specific knobs (declared in `vllm_omni/model_extras/helios.py`) are
+  passed via the generic `--extra-body` JSON flag:
+  - Helios-Mid: `--extra-body '{"is_enable_stage2": true, "pyramid_num_inference_steps_list": [20, 20, 20], "use_cfg_zero_star": true, "use_zero_init": true, "zero_steps": 1}'`
+  - Helios-Distilled: `--extra-body '{"is_enable_stage2": true, "pyramid_num_inference_steps_list": [2, 2, 2], "is_amplify_first_chunk": true}'`
 
 #### Known limitations
 

recipes/JD/JoyAI-VL-Interaction.md

@@ -0,0 +1,192 @@
+# JoyAI-VL-Interaction
+
+> Real-time streaming video-language interaction (proactive speak / silence / delegate)
+
+## Summary
+
+- Vendor: JD (Joy Future Academy)
+- Model: [`jdopensource/JoyAI-VL-Interaction-Preview`](https://huggingface.co/jdopensource/JoyAI-VL-Interaction-Preview) (8B, vanilla Qwen3-VL weights)
+- Task: Per-tick proactive interaction over a live video stream — the model decides
+  on its own each second to speak, stay silent, or delegate a hard question
+- Mode: Online serving — an OpenAI-compatible interaction orchestrator in front of
+  a plain `vllm serve` backend
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this to stand up the streaming-interaction serving layer (`vllm_omni/experimental/fullduplex/`).
+The model is served unchanged by `vllm serve`; this layer adds session state, 3-tier
+summary memory, the per-tick decision, and pluggable ASR / TTS / delegation. For the
+framework internals and how to add another full-duplex model, see
+[`vllm_omni/experimental/fullduplex/README.md`](../../vllm_omni/experimental/fullduplex/README.md).
+
+## Environment
+
+- OS: Linux
+- Python: 3.10+
+- Hardware: 1x GPU. The default `T_s=100` frame window + `--max-model-len 131072` wants
+  ≈48GB+; for ~24GB, lower `--chunk-frames`, `--max-model-len`, and the image limit together
+- vLLM / vLLM-Omni: versions from your current checkout
+
+## Start server
+
+From repository root:
+
+```bash
+# 1. Serve the model (plain `vllm serve`, NOT --omni; it is vanilla Qwen3-VL).
+#    The image limit must cover the short-term frame window (chunk_frames, default 100 = T_s);
+#    prefix caching keeps the accumulating window cheap. Lower both for smaller GPUs.
+vllm serve jdopensource/JoyAI-VL-Interaction-Preview \
+  --served-model-name JoyAI-VL-Interaction-Preview --port 8061 \
+  --max-model-len 131072 --enable-prefix-caching \
+  --limit-mm-per-prompt '{"image":256,"video":1}'
+
+# 2. Interaction orchestrator (OpenAI-compatible, :8070)
+python -m vllm_omni.experimental.fullduplex.joyvl.serving.server --port 8070 \
+  --main-backend-url http://127.0.0.1:8061/v1 --main-model JoyAI-VL-Interaction-Preview
+```
+
+Optional one-shot launch (model + orchestrator + JD webui + ASR/TTS/background,
+all env-configurable). The JD webui frontend is external — set `WEBUI_DIR` to point at it:
+
+```bash
+bash examples/online_serving/joyvl_interaction/scripts/start_all.sh
+```
+
+## Using the model
+
+The orchestrator is OpenAI-compatible: send **one user turn per video frame** (~1 fps) to
+`/v1/chat/completions` with an `x-session-id` header, and attach an optional **standing
+instruction** as a text part. Each reply carries an `interaction` block — `action` is
+`silence` / `response` / `delegate` and `text` is what to say:
+
+```bash
+curl -s http://127.0.0.1:8070/v1/chat/completions \
+  -H 'x-session-id: s1' -H 'content-type: application/json' -d '{
+    "messages": [{"role": "user", "content": [
+      {"type": "text", "text": "Alert me if a fire breaks out"},
+      {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,..."}}
+    ]}]}' | jq .interaction
+```
+
+Send the standing instruction once (it persists for the session); subsequent turns send
+just the frame. Ready-made headless client: `cli/run_cli_demo.py`. Reset a session with
+`POST /reset {"session_id": "..."}`.
+
+**What to ask it** — give a standing task and let it act on its own each second:
+
+- **Proactive alerting** — "tell me when someone enters", "alert me if a fire breaks out"
+- **Streaming Q&A** — ask about what's on screen; it answers once it has the evidence
+- **Time & memory** — "how many people have walked past?", "what did you see earlier?"
+  (the 3-tier summary memory lets it recall beyond the live frame window)
+- **Read & translate** — "translate the on-screen text as it appears"
+- **Delegate hard questions** — handed to a background brain (see below)
+- **Live commentary** — `--persona talkative` for continuous, danmaku-style narration
+
+**Personas** (`--persona`): `default` (speak on meaningful events or to answer), `silent`
+(answer only when asked, never delegate), `talkative` (proactively narrate).
+
+**Tuning:** `--chunk-frames` (short-term window `T_s`), `--response-dedup-threshold`
+(lower drops more near-duplicate narration), and `--no-memory` to disable the 3-tier
+summaries. The model stays silent until the first instruction arrives
+(`force_silence_before_query`, on by default).
+
+## Delegation (background brain)
+
+When the model judges a question too hard, it emits `</delegation> <question>` and the
+orchestrator hands it to a **background brain** — any OpenAI-compatible endpoint you
+self-host. Enable it by pointing the orchestrator at one:
+
+```bash
+python -m vllm_omni.experimental.fullduplex.joyvl.serving.server --port 8070 \
+  --main-backend-url http://127.0.0.1:8061/v1 --main-model JoyAI-VL-Interaction-Preview \
+  --delegation-backend-url <brain-endpoint>/v1 \
+  --delegation-model <brain-model> --delegation-kind chat
+```
+
+`--delegation-kind` picks the bridge:
+
+- `chat` — a stronger text/VL model answers (`/chat/completions`)
+- `image` — a text-to-image model generates a picture (`/images/generations`, e.g. Qwen-Image)
+- `edit` — an image-edit model restyles the current frame (e.g. Qwen-Image-Edit)
+- `router` — dispatch each request to chat / image / edit by inspecting it (set
+  `--delegation-image-url` / `--delegation-edit-url` for the latter two)
+- `stub` — canned answers for tests/demos only (no backend needed)
+
+`chat`/`image`/`edit`/`router` each need a backend URL — **without one, delegation stays
+off** (the model's delegate note is still spoken, but nothing is folded back). The brain is
+**bring-your-own**: a larger vLLM you serve, or any OpenAI-compatible API (e.g.
+`--delegation-backend-url https://api.anthropic.com/v1/ --delegation-model claude-...
+--delegation-api-key …`). The reference deployment instead drives the `codex` CLI as the
+brain via a separate background-agent service; that agent runs with its own credentials
+and bypasses its sandbox, so it is **not bundled here** — self-host a plain
+OpenAI-compatible endpoint instead.
+
+## Host the WebUI demo
+
+For the full browser experience — live webcam / RTSP input, voice (ASR/TTS), and the
+per-tick decision stream — run JD's official WebUI in front of the orchestrator. Clone the
+model repo and start its WebUI pointed at the orchestrator (`:8070`):
+
+```bash
+git clone https://github.com/jd-opensource/JoyAI-VL-Interaction.git
+cd JoyAI-VL-Interaction/services/webui
+uv venv && uv pip install -e .
+bash scripts/start_server.sh --api-base http://127.0.0.1:8070/v1
+```
+
+Open the printed HTTPS URL, allow the camera (or enter an RTSP URL), and give a standing
+instruction. `examples/online_serving/joyvl_interaction/scripts/start_all.sh` can launch
+the model + orchestrator + WebUI together — point `WEBUI_DIR` at your clone.
+
+## Verification
+
+```bash
+# headless: stream a clip and print the per-tick decision timeline
+# (the CLI reads video frames via OpenCV: `uv pip install opencv-python`)
+python examples/online_serving/joyvl_interaction/cli/run_cli_demo.py \
+  path/to/video.mp4 --query "Alert me if a fire breaks out"
+
+pytest tests/fullduplex   # framework + JoyVL unit tests
+```
+
+## Testing with an RTSP stream (optional)
+
+RTSP is a **webui-side input** — the browser pulls the stream and feeds frames to the
+orchestrator over the normal API; no serving-layer code is involved. To simulate an RTSP
+camera from a local video file (no physical IP camera needed), use the helper scripts in
+[`examples/online_serving/joyvl_interaction/rtsp/`](../../examples/online_serving/joyvl_interaction/rtsp/),
+which wrap [MediaMTX](https://github.com/bluenviron/mediamtx/releases) + `ffmpeg`:
+
+```bash
+cd examples/online_serving/joyvl_interaction/rtsp
+
+# 1. Local RTSP server (MediaMTX, listens on :8554)
+bash ./mediamtx.sh
+
+# 2. Push a local video file as an RTSP stream (another terminal)
+bash ./rtsp.sh ./videos/example.mp4 rtsp://127.0.0.1:8554/fire1
+
+# 3. In the WebUI RTSP box, enter:  rtsp://127.0.0.1:8554/fire1
+#    (replace 127.0.0.1 with the MediaMTX host IP if the webui runs on another machine)
+```
+
+See the directory's `README.md` for streaming a whole video folder (`rtsp_all.sh`) and
+the audio-track caveat.
+
+## Notes
+
+- `--omni` is **not** used: the model is standard Qwen3-VL, so stock `vllm serve`
+  runs the forward pass; this recipe only adds the interaction/serving layer.
+- On a host without `nvcc` / `ninja`, `vllm serve` of the 8B can crash engine-core in the
+  FlashInfer sampler JIT (`FileNotFoundError: 'ninja'`) during `profile_run`. Set
+  `VLLM_USE_FLASHINFER_SAMPLER=0` (or install `ninja`) to work around it.
+- `force_silence_before_query` is on by default — the model stays silent until an
+  instruction arrives; give a standing task (e.g. "translate the on-screen text")
+  to arm proactive output.
+- Speech is external and pluggable: point `ASR_URL` / `TTS_URL` at the bridges in
+  `examples/online_serving/joyvl_interaction/bridges/` or any compatible service.
+- The decision prompts, sampling, and 3-tier summary memory (`T_s=100`, mid→long every
+  5 chunks, `key_frames=0` = summarize all chunk frames) are aligned to the JoyVL
+  reference adapter so per-tick behavior matches the released model; the framework only
+  supplies the serving structure.

recipes/README.md

@@ -64,6 +64,7 @@ recipes/
 | [`StabilityAI/Stable-Diffusion-3.5.md`](./StabilityAI/Stable-Diffusion-3.5.md) | Text-to-image serving (SD 3.5-medium and SD 3.5-large) | 1x RTX A6000 48GB |
 | [`zai-org/GLM-TTS.md`](./zai-org/GLM-TTS.md) | Online serving for Chinese/English zero-shot voice-cloned TTS | 1x A40 48GB |
 | [`GLM/GLM-Image.md`](./GLM/GLM-Image.md) | Online serving for image generation | 1x A800 80GB / 2x A800 80GB |
+| [`JD/JoyAI-VL-Interaction.md`](./JD/JoyAI-VL-Interaction.md) | Real-time streaming video-language interaction (proactive speak/silence/delegate) | 1x GPU 24GB+ |
 
 Within a single recipe file, include different hardware support sections such
 as `GPU`, `ROCm`, and `NPU`, and add concrete tested configurations like

recipes/cosmos3/Cosmos3-Nano.md

@@ -6,7 +6,7 @@
 
 - Vendor: NVIDIA
 - Model: `nvidia/Cosmos3-Nano`
-- Task: Text-to-image (T2I), text-to-video (T2V), image-to-video (I2V), and video-to-video (V2V) generation, with optional synchronized audio (video + sound), action policy
+- Task: Text-to-image (T2I), text-to-video (T2V), image-to-video (I2V), and video-to-video (V2V) generation, with optional transfer controls, synchronized audio (video + sound), action policy
 - Mode: Online serving with the OpenAI-compatible image/video APIs, plus offline generation via the `Omni` API
 - Maintainer: Community
 
@@ -25,6 +25,9 @@ mode is selected per request:
   reference-video latent frames; use `extra_params.condition_frame_indexes_vision`
   and `extra_params.condition_video_keep` to choose which prefix/tail frames guide
   generation.
+- **Transfer V2V** — pass one or more transfer hints in `extra_params`
+  (`edge`, `blur`, `depth`, `seg`, `wsm`) to guide generation with control
+  frames. Transfer mode is video-only and cannot be combined with sound or action.
 - **T2VS / I2VS** — add `generate_sound=true` (and optional `sound_duration`) to a
   T2V/I2V `/v1/videos/sync` request to also generate synchronized audio, muxed into
   the mp4 as AAC 48 kHz stereo. See the official model card's "Video + Audio" examples.
@@ -190,6 +193,18 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync \
   -F 'video_reference={"video_url":"https://example.com/reference.mp4"}' \
   -o cosmos3_v2v_from_url.mp4
 
+# Transfer V2V with a precomputed depth control video. `control_path` can point
+# to a local image/video; edge and blur can also be computed from `input_reference`
+# by passing `"edge":true` or `"blur":true`.
+curl -sS -X POST http://localhost:8000/v1/videos/sync \
+  -H "Accept: video/mp4" \
+  -F "model=nvidia/Cosmos3-Nano" \
+  -F "prompt=Generate a realistic scene following the provided control video." \
+  -F "size=1280x720" -F "num_frames=121" \
+  -F "num_inference_steps=50" -F "seed=125" \
+  -F 'extra_params={"depth":{"control_path":"/path/to/depth_control.mp4"},"max_frames":121,"resolution":"720","num_video_frames_per_chunk":121}' \
+  -o cosmos3_transfer_depth.mp4
+
 # Text-to-video-with-sound
 curl -sS -X POST http://localhost:8000/v1/videos/sync \
   -H "Accept: video/mp4" \
@@ -312,6 +327,17 @@ vllm serve nvidia/Cosmos3-Nano-Policy-DROID \
   For V2V, `condition_frame_indexes_vision` selects the clean conditioned latent
   frame indexes (default `[0, 1]`), and `condition_video_keep` selects whether the
   API decodes the first or last needed reference frames (`"first"` by default).
+- **Transfer controls:** `extra_params` may include `edge`, `blur`, `depth`,
+  `seg`, or `wsm`. Each hint accepts `true`, a path string, or an object such as
+  `{"control_path": "/path/to/control.mp4"}`; `edge` also accepts
+  `preset_edge_threshold` and `blur` accepts `preset_blur_strength`.
+  Transfer-level options include `control_guidance`,
+  `control_guidance_interval`, `num_video_frames_per_chunk` (default `93`,
+  `101` for WSM), `num_conditional_frames` (default `1`),
+  `num_first_chunk_conditional_frames`, `max_frames`,
+  `show_control_condition`, `show_input`, and
+  `share_vision_temporal_positions`. Non-WSM transfer preserves the input video
+  fps when available; WSM defaults to 10 fps unless `fps` is supplied.
 - **DROID OpenPI observations:** include a string `prompt`, either
   `observation/image` or the three-view DROID camera keys
   (`observation/wrist_image_left`, `observation/exterior_image_1_left`,

recipes/cosmos3/Cosmos3-Super.md

@@ -57,7 +57,10 @@ vllm serve nvidia/Cosmos3-Super \
 
 Guardrails are on by default (gated `nvidia/Cosmos-1.0-Guardrail` — `pip install
 cosmos-guardrail`, accept the license, set `HF_TOKEN`); add `--no-guardrails` to
-disable. `--enable-layerwise-offload` reduces VRAM on smaller GPUs.
+disable. `--enable-layerwise-offload` reduces VRAM on smaller GPUs;
+`--quantization fp8` (online, no calibration) cuts peak VRAM for 720p video
+generation from ~83 GB to ~55 GB per GPU (2-GPU) with BF16-level quality (T2V
+composition can shift at the same seed).
 
 #### Verification
 

tests/attention/test_fish_kvcache_attn.py

@@ -92,6 +92,18 @@ def _metadata(
     )()
 
 
+def _decode_kv_cache(
+    *,
+    num_blocks: int = 8,
+    block_size: int = 16,
+    num_kv_heads: int = 8,
+    head_size: int = 128,
+    dtype: torch.dtype = torch.float16,
+) -> torch.Tensor:
+    # vLLM >=0.23.0 KV cache layout: key/value live on dim 1.
+    return torch.zeros((num_blocks, 2, block_size, num_kv_heads, head_size), dtype=dtype)
+
+
 def test_fish_kvcache_enabled_by_default(monkeypatch):
     monkeypatch.delenv("VLLM_OMNI_FISH_KVCACHE_ATTN", raising=False)
 
@@ -182,7 +194,7 @@ def test_fish_kvcache_backend_wraps_only_model_instance(monkeypatch):
     metadata = _metadata()
     query = torch.zeros((2, 32, 128), dtype=torch.float16)
     output = torch.zeros_like(query)
-    kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16)
+    kv_cache = _decode_kv_cache()
 
     result = model.layers[0].self_attn.attn.impl.forward(
         None,
@@ -201,6 +213,110 @@ def test_fish_kvcache_backend_wraps_only_model_instance(monkeypatch):
     assert fish_kvcache_backend.get_fish_kvcache_attn_stats()["small_hit_count"] == 1
 
 
+def _vllm_kv_cache_shape(num_blocks, block_size, num_kv_heads, head_size):
+    # Source of truth for the paged KV cache layout the backend must unpack.
+    from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
+
+    return FlashAttentionBackend.get_kv_cache_shape(num_blocks, block_size, num_kv_heads, head_size)
+
+
+def test_fish_kvcache_backend_unbinds_kv_on_vllm_cache_layout(monkeypatch):
+    # vLLM >=0.23.0 lays the KV cache out as (num_blocks, 2, block_size,
+    # num_kv_heads, head_size). The backend must unbind on dim 1 so each of
+    # key/value comes back as the 4-D (num_blocks, block_size, num_kv_heads,
+    # head_size) tensor the decode kernel expects.
+    monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", "1")
+    monkeypatch.setattr(fish_kvcache_backend, "is_available", lambda: True)
+    monkeypatch.setattr(fish_kvcache_backend, "load_error", lambda: None)
+    monkeypatch.setattr(fish_kvcache_backend, "can_use_fish_kvcache_attn", lambda **_: True)
+
+    num_blocks, block_size, num_kv_heads, head_size = 8, 16, 8, 128
+    cache_shape = _vllm_kv_cache_shape(num_blocks, block_size, num_kv_heads, head_size)
+    assert cache_shape == (num_blocks, 2, block_size, num_kv_heads, head_size)
+
+    captured = {}
+
+    def fake_decode(query, key_cache, value_cache, block_table, seq_lens, out, *, scale, max_seq_len):
+        del query, block_table, seq_lens, scale, max_seq_len
+        captured["key"] = tuple(key_cache.shape)
+        captured["value"] = tuple(value_cache.shape)
+        out.fill_(1)
+        return out
+
+    monkeypatch.setattr(fish_kvcache_backend, "fish_decode_kvcache_attn", fake_decode)
+
+    model = _FakeModel()
+    assert fish_kvcache_backend.install_fish_kvcache_attn_backend(model) == 1
+
+    query = torch.zeros((2, 32, head_size), dtype=torch.float16)
+    output = torch.zeros_like(query)
+    kv_cache = torch.zeros(cache_shape, dtype=torch.float16)
+
+    model.layers[0].self_attn.attn.impl.forward(None, query, None, None, kv_cache, _metadata(), output)
+
+    expected = (num_blocks, block_size, num_kv_heads, head_size)
+    assert captured["key"] == expected
+    assert captured["value"] == expected
+
+
+def test_fish_kvcache_vllm_cache_unbinds_then_falls_back_until_contiguous(monkeypatch):
+    # End-to-end with the real vLLM cache layout and the real (unmocked) guard.
+    #
+    # The point of the unbind(1) fix: unbind(0) raised "too many values to
+    # unpack" on the 5-D vLLM 0.23.0 layout, so the request crashed outright.
+    # unbind(1) makes the request *work* again -- it unpacks cleanly and the
+    # backend gracefully falls back to the original attention with correct
+    # output. This test guards that fix and must NOT depend on the fast path
+    # firing.
+    #
+    # It also documents a known limitation (tracked separately): unbind(1) on
+    # the interleaved (num_blocks, 2, ...) layout yields *non-contiguous*
+    # key/value views, which the guard rejects on the is_contiguous() check.
+    # So the dimensions are compatible (4-D, block_size=16, head_size=128) --
+    # the only thing keeping the fast path from engaging is contiguity, not
+    # shape. Making the fast path actually fire would require a .contiguous()
+    # copy and is out of scope here.
+    fish_kvcache_backend.reset_fish_kvcache_attn_stats()
+    monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", "1")
+    monkeypatch.setattr(fish_kvcache_backend, "is_available", lambda: True)
+    monkeypatch.setattr(fish_kvcache_backend, "load_error", lambda: None)
+    monkeypatch.setattr(fish_kvcache_attn, "is_available", lambda: True)
+
+    num_blocks, block_size, num_kv_heads, head_size = 8, 16, 8, 128
+    cache_shape = _vllm_kv_cache_shape(num_blocks, block_size, num_kv_heads, head_size)
+
+    decode_calls = []
+
+    def fake_decode(*args, **kwargs):
+        del args, kwargs
+        decode_calls.append(True)
+
+    monkeypatch.setattr(fish_kvcache_backend, "fish_decode_kvcache_attn", fake_decode)
+
+    model = _FakeModel()
+    assert fish_kvcache_backend.install_fish_kvcache_attn_backend(model) == 1
+
+    query = torch.zeros((2, 32, head_size), dtype=torch.float16)
+    output = torch.zeros_like(query)
+    kv_cache = torch.zeros(cache_shape, dtype=torch.float16)
+
+    # Root cause record: unbind(1) gives the right 4-D shape but a non-contiguous view.
+    key_cache, value_cache = kv_cache.unbind(1)
+    assert tuple(key_cache.shape) == (num_blocks, block_size, num_kv_heads, head_size)
+    assert not key_cache.is_contiguous()
+
+    # unbind(1) no longer crashes (unbind(0) would have raised here) and the
+    # request completes via the correct fallback path.
+    result = model.layers[0].self_attn.attn.impl.forward(None, query, None, None, kv_cache, _metadata(), output)
+
+    assert result is output
+    assert decode_calls == []
+    assert model.layers[0].self_attn.attn.impl.original_calls == 1
+    assert torch.equal(output, torch.full_like(output, 2))
+    stats = fish_kvcache_backend.get_fish_kvcache_attn_stats()
+    assert stats["fallback_count_by_reason"] == {"guard_miss": 1}
+
+
 def test_fish_kvcache_backend_install_is_idempotent(monkeypatch):
     monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", "1")
     monkeypatch.setattr(fish_kvcache_backend, "is_available", lambda: True)
@@ -215,7 +331,7 @@ def test_fish_kvcache_backend_install_is_idempotent(monkeypatch):
 
     query = torch.zeros((2, 32, 128), dtype=torch.float16)
     output = torch.zeros_like(query)
-    kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16)
+    kv_cache = _decode_kv_cache()
 
     model.layers[0].self_attn.attn.impl.forward(
         None,
@@ -250,7 +366,7 @@ def test_fish_kvcache_backend_falls_back_to_original_forward_on_guard_miss(monke
 
     query = torch.zeros((2, 32, 128), dtype=torch.float16)
     output = torch.zeros_like(query)
-    kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16)
+    kv_cache = _decode_kv_cache()
 
     result = model.layers[0].self_attn.attn.impl.forward(
         None,
@@ -297,7 +413,7 @@ def test_fish_kvcache_backend_uses_decode_seq_lens_upper_bound(monkeypatch):
 
     query = torch.zeros((2, 32, 128), dtype=torch.float16)
     output = torch.zeros_like(query)
-    kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16)
+    kv_cache = _decode_kv_cache()
 
     model.layers[0].self_attn.attn.impl.forward(
         None,
@@ -340,7 +456,7 @@ def test_fish_kvcache_backend_uses_cpu_upper_bound_above_metadata_max(monkeypatc
 
     query = torch.zeros((2, 32, 128), dtype=torch.float16)
     output = torch.zeros_like(query)
-    kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16)
+    kv_cache = _decode_kv_cache()
 
     model.layers[0].self_attn.attn.impl.forward(
         None,
@@ -383,7 +499,7 @@ def test_fish_kvcache_backend_slices_upper_bound_by_active_batch(monkeypatch):
 
     query = torch.zeros((2, 32, 128), dtype=torch.float16)
     output = torch.zeros_like(query)
-    kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16)
+    kv_cache = _decode_kv_cache()
 
     model.layers[0].self_attn.attn.impl.forward(
         None,
@@ -428,7 +544,7 @@ def test_fish_kvcache_backend_falls_back_on_short_upper_bound_shape(monkeypatch)
 
     query = torch.zeros((2, 32, 128), dtype=torch.float16)
     output = torch.zeros_like(query)
-    kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16)
+    kv_cache = _decode_kv_cache()
 
     result = model.layers[0].self_attn.attn.impl.forward(
         None,
@@ -475,7 +591,7 @@ def test_fish_kvcache_backend_falls_back_on_non_cpu_upper_bound(monkeypatch):
 
     query = torch.zeros((2, 32, 128), dtype=torch.float16)
     output = torch.zeros_like(query)
-    kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16)
+    kv_cache = _decode_kv_cache()
 
     result = model.layers[0].self_attn.attn.impl.forward(
         None,
@@ -509,7 +625,7 @@ def test_fish_kvcache_backend_required_rejects_non_cpu_upper_bound(monkeypatch):
 
     query = torch.zeros((2, 32, 128), dtype=torch.float16)
     output = torch.zeros_like(query)
-    kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16)
+    kv_cache = _decode_kv_cache()
 
     with pytest.raises(RuntimeError, match="must be a CPU tensor"):
         model.layers[0].self_attn.attn.impl.forward(
@@ -539,7 +655,7 @@ def test_fish_kvcache_backend_required_rejects_underestimated_upper_bound(monkey
 
     query = torch.zeros((2, 32, 128), dtype=torch.float16)
     output = torch.zeros_like(query)
-    kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16)
+    kv_cache = _decode_kv_cache()
 
     with pytest.raises(RuntimeError, match="underestimates"):
         model.layers[0].self_attn.attn.impl.forward(
@@ -569,7 +685,7 @@ def test_fish_kvcache_backend_required_rejects_missing_upper_bound(monkeypatch):
 
     query = torch.zeros((2, 32, 128), dtype=torch.float16)
     output = torch.zeros_like(query)
-    kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16)
+    kv_cache = _decode_kv_cache()
 
     with pytest.raises(RuntimeError, match="requires seq_lens_cpu_upper_bound"):
         model.layers[0].self_attn.attn.impl.forward(
@@ -603,7 +719,7 @@ def test_fish_kvcache_backend_required_raises_on_guard_miss(monkeypatch):
 
     query = torch.zeros((2, 32, 128), dtype=torch.float16)
     output = torch.zeros_like(query)
-    kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16)
+    kv_cache = _decode_kv_cache()
 
     with pytest.raises(RuntimeError, match="required"):
         model.layers[0].self_attn.attn.impl.forward(
@@ -629,7 +745,7 @@ def test_fish_kvcache_backend_required_allows_prefill_fallback(monkeypatch):
 
     query = torch.zeros((2, 32, 128), dtype=torch.float16)
     output = torch.zeros_like(query)
-    kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16)
+    kv_cache = _decode_kv_cache()
 
     result = model.layers[0].self_attn.attn.impl.forward(
         None,
@@ -675,7 +791,7 @@ def test_fish_kvcache_backend_falls_back_without_cpu_upper_bound(monkeypatch):
 
     query = torch.zeros((2, 32, 128), dtype=torch.float16)
     output = torch.zeros_like(query)
-    kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16)
+    kv_cache = _decode_kv_cache()
 
     result = model.layers[0].self_attn.attn.impl.forward(
         None,

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.

tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py

@@ -122,9 +122,6 @@ def test_speech_missing_required_fields(omni_server: OmniServer, openai_client:
             {"response_format": "mpeg"}, ("response_format", "literal_error", "wav"), id="response_format_invalid"
         ),
         pytest.param({"stream_format": 1}, ("stream_format", "literal_error", "audio"), id="stream_format_wrong_type"),
-        pytest.param(
-            {"stream_format": "sse"}, ("stream_format", "value_error", "audio"), id="stream_format_sse_blocked"
-        ),
         pytest.param({"stream": "wrong_type"}, ("stream", "bool_parsing", "validation error"), id="stream_wrong_type"),
         pytest.param(
             {

tests/diffusion/models/bagel/test_naive_cache.py

@@ -0,0 +1,223 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for NaiveCache merge/split logic used in batched CFG."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.bagel.bagel_transformer import NaiveCache
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+NUM_LAYERS = 2
+NUM_KV_HEADS = 4
+HEAD_DIM = 8
+
+
+def _make_cache(num_layers, seq_len, num_kv_heads=NUM_KV_HEADS, head_dim=HEAD_DIM, seed=0):
+    """Create a NaiveCache with deterministic random data. seq_len=0 returns an empty cache."""
+    gen = torch.Generator().manual_seed(seed)
+    cache = NaiveCache(num_layers)
+    if seq_len == 0:
+        return cache
+    for layer in range(num_layers):
+        cache.key_cache[layer] = torch.randn(seq_len, num_kv_heads, head_dim, generator=gen)
+        cache.value_cache[layer] = torch.randn(seq_len, num_kv_heads, head_dim, generator=gen)
+    return cache
+
+
+def test_init_creates_none_entries():
+    """Ensure the NaiveCache is initialized with None values per layer."""
+    cache = NaiveCache(NUM_LAYERS)
+    assert cache.num_layers == NUM_LAYERS
+    for layer in range(NUM_LAYERS):
+        assert cache.key_cache[layer] is None
+        assert cache.value_cache[layer] is None
+
+
+@pytest.mark.parametrize("seq_len", [0, 10])
+def test_seq_lens_empty(seq_len):
+    """Ensure that by default, we have 0 seq lens."""
+    cache = _make_cache(NUM_LAYERS, seq_len=seq_len)
+    assert cache.seq_lens == seq_len
+    assert cache.num_layers == NUM_LAYERS
+
+
+### Merge tests
+def test_merge_two_equal_length():
+    """Ensure that we can merge two NaiveCaches that are identically shaped."""
+    c0 = _make_cache(NUM_LAYERS, seq_len=5, seed=0)
+    c1 = _make_cache(NUM_LAYERS, seq_len=5, seed=1)
+    merged = NaiveCache.merge([c0, c1])
+
+    assert merged.key_values_lens == [5, 5]
+    for layer in range(NUM_LAYERS):
+        assert merged.key_cache[layer].shape[0] == 10
+        # the merged cache will just have keys and values per layer concatenated
+        assert torch.equal(merged.key_cache[layer][:5], c0.key_cache[layer])
+        assert torch.equal(merged.key_cache[layer][5:], c1.key_cache[layer])
+        assert torch.equal(merged.value_cache[layer][:5], c0.value_cache[layer])
+        assert torch.equal(merged.value_cache[layer][5:], c1.value_cache[layer])
+
+
+def test_merge_three_zero_len_cache():
+    """Ensure we handle zero length cache correctly in merge."""
+    # NOTE: This is relevant for text_cfg in Bagel, which has a len of 0 by default
+    gen_cache = _make_cache(NUM_LAYERS, seq_len=10, seed=0)
+    text_cfg_cache = _make_cache(NUM_LAYERS, seq_len=0)
+    img_cfg_cache = _make_cache(NUM_LAYERS, seq_len=7, seed=2)
+    merged = NaiveCache.merge([gen_cache, text_cfg_cache, img_cfg_cache])
+
+    assert merged.key_values_lens == [10, 0, 7]
+    for layer in range(NUM_LAYERS):
+        assert merged.key_cache[layer].shape[0] == 17
+        assert torch.equal(merged.key_cache[layer][:10], gen_cache.key_cache[layer])
+        assert torch.equal(merged.key_cache[layer][10:], img_cfg_cache.key_cache[layer])
+
+
+def test_merge_all_empty():
+    """Ensure that merging empty caches is well defined."""
+    caches = [_make_cache(NUM_LAYERS, seq_len=0) for _ in range(3)]
+    merged = NaiveCache.merge(caches)
+
+    assert merged.key_values_lens == [0, 0, 0]
+    for layer in range(NUM_LAYERS):
+        assert merged.key_cache[layer] is None
+        assert merged.value_cache[layer] is None
+
+
+def test_merge_single_cache():
+    """Ensure merging one cache returns an identical cache."""
+    c = _make_cache(NUM_LAYERS, seq_len=8, seed=42)
+    merged = NaiveCache.merge([c])
+
+    assert merged.key_values_lens == [8]
+    for layer in range(NUM_LAYERS):
+        assert torch.equal(merged.key_cache[layer], c.key_cache[layer])
+
+
+def test_merge_preserves_dtype():
+    """Ensure merging doesn't modify dtypes."""
+    cache = NaiveCache(NUM_LAYERS)
+    for layer in range(NUM_LAYERS):
+        cache.key_cache[layer] = torch.randn(5, NUM_KV_HEADS, HEAD_DIM, dtype=torch.bfloat16)
+        cache.value_cache[layer] = torch.randn(5, NUM_KV_HEADS, HEAD_DIM, dtype=torch.bfloat16)
+
+    merged = NaiveCache.merge([cache])
+    assert merged.key_cache[0].dtype == torch.bfloat16
+    assert merged.value_cache[0].dtype == torch.bfloat16
+
+
+### Split tests
+def test_split_all_nonzero():
+    """Ensure NaiveCache splits in the simple case (all lens nonzero)."""
+    t = torch.randn(15, NUM_KV_HEADS, HEAD_DIM)
+    parts = NaiveCache.split_with_zeros(t, [5, 4, 6])
+
+    assert len(parts) == 3
+    assert all(p is not None for p in parts)
+    assert parts[0].shape[0] == 5
+    assert parts[1].shape[0] == 4
+    assert parts[2].shape[0] == 6
+    assert torch.equal(torch.cat(parts), t)
+
+
+def test_split_with_zero():
+    """Ensure NaiveCache split handles zero length correctly (used in Bagel)."""
+    t = torch.randn(17, NUM_KV_HEADS, HEAD_DIM)
+    parts = NaiveCache.split_with_zeros(t, [10, 0, 7])
+
+    assert parts[0].shape[0] == 10
+    assert parts[1] is None
+    assert parts[2].shape[0] == 7
+    assert torch.equal(torch.cat([parts[0], parts[2]]), t)
+
+
+def test_split_wrong_sum_raises():
+    """Ensure NaiveCache raises if splits don't match the sum of dims on axis 0."""
+    t = torch.randn(10, NUM_KV_HEADS, HEAD_DIM)
+    with pytest.raises(ValueError, match="dim 0"):
+        NaiveCache.split_with_zeros(t, [5, 3])
+
+
+def test_split_negative_length_raises():
+    """Ensure NaiveCache raises if splits have any negative values."""
+    t = torch.randn(10, NUM_KV_HEADS, HEAD_DIM)
+    with pytest.raises(ValueError, match="greater than or equal to zero"):
+        NaiveCache.split_with_zeros(t, [5, -1, 6])
+
+
+def test_split_preserves_dtype():
+    """Ensure NaiveCache split preserves dtype."""
+    t = torch.randn(10, NUM_KV_HEADS, HEAD_DIM, dtype=torch.bfloat16)
+    parts = NaiveCache.split_with_zeros(t, [4, 6])
+    assert parts[0].dtype == torch.bfloat16
+    assert parts[1].dtype == torch.bfloat16
+
+
+### from_object tests (for kv cache transfer)
+def test_from_object_passthrough():
+    """Ensure a NaiveCache input is returned as is."""
+    cache = _make_cache(NUM_LAYERS, seq_len=5)
+    assert NaiveCache.from_object(cache) is cache
+
+
+def test_from_object_converts_simple_namespace():
+    """Ensure SimpleNamespace with list-based caches converts to NaiveCache."""
+    keys = [torch.randn(5, NUM_KV_HEADS, HEAD_DIM) for _ in range(NUM_LAYERS)]
+    values = [torch.randn(5, NUM_KV_HEADS, HEAD_DIM) for _ in range(NUM_LAYERS)]
+    ns = SimpleNamespace(key_cache=keys, value_cache=values)
+
+    cache = NaiveCache.from_object(ns)
+
+    assert isinstance(cache, NaiveCache)
+    assert cache.num_layers == NUM_LAYERS
+    for i in range(NUM_LAYERS):
+        assert torch.equal(cache.key_cache[i], keys[i])
+        assert torch.equal(cache.value_cache[i], values[i])
+
+
+def test_from_object_mismatched_lengths_raises():
+    """Ensure mismatched key/value cache lengths raise due to strict=True in zip."""
+    keys = [torch.randn(5, NUM_KV_HEADS, HEAD_DIM) for _ in range(2)]
+    values = [torch.randn(5, NUM_KV_HEADS, HEAD_DIM) for _ in range(3)]
+    ns = SimpleNamespace(key_cache=keys, value_cache=values)
+
+    with pytest.raises(ValueError):
+        NaiveCache.from_object(ns)
+
+
+### End to end test for split / merge
+def test_round_trip_two_populated():
+    """Roundtrip test for merging and resplitting two simple caches."""
+    c0 = _make_cache(NUM_LAYERS, seq_len=5, seed=0)
+    c1 = _make_cache(NUM_LAYERS, seq_len=8, seed=1)
+    merged = NaiveCache.merge([c0, c1])
+
+    for layer in range(NUM_LAYERS):
+        k_parts = NaiveCache.split_with_zeros(merged.key_cache[layer], merged.key_values_lens)
+        v_parts = NaiveCache.split_with_zeros(merged.value_cache[layer], merged.key_values_lens)
+        assert torch.equal(k_parts[0], c0.key_cache[layer])
+        assert torch.equal(k_parts[1], c1.key_cache[layer])
+        assert torch.equal(v_parts[0], c0.value_cache[layer])
+        assert torch.equal(v_parts[1], c1.value_cache[layer])
+
+
+def test_round_trip_three_branches_with_zero_cfg():
+    """Roundtrip test with a zero entry (i.e., same as Bagel's gen/text_cfg/img_cfg case)."""
+    gen_cache = _make_cache(NUM_LAYERS, seq_len=10, seed=0)
+    text_cfg_cache = _make_cache(NUM_LAYERS, seq_len=0)
+    img_cfg_cache = _make_cache(NUM_LAYERS, seq_len=7, seed=2)
+    merged = NaiveCache.merge([gen_cache, text_cfg_cache, img_cfg_cache])
+
+    for layer in range(NUM_LAYERS):
+        k_parts = NaiveCache.split_with_zeros(merged.key_cache[layer], merged.key_values_lens)
+        v_parts = NaiveCache.split_with_zeros(merged.value_cache[layer], merged.key_values_lens)
+        assert torch.equal(k_parts[0], gen_cache.key_cache[layer])
+        assert k_parts[1] is None
+        assert torch.equal(k_parts[2], img_cfg_cache.key_cache[layer])
+        assert torch.equal(v_parts[0], gen_cache.value_cache[layer])
+        assert v_parts[1] is None
+        assert torch.equal(v_parts[2], img_cfg_cache.value_cache[layer])

tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py

@@ -9,6 +9,7 @@ from dataclasses import dataclass
 from types import SimpleNamespace
 from typing import Any
 
+import numpy as np
 import pytest
 import torch
 from PIL import Image
@@ -53,8 +54,10 @@ class StubCosmos3VAE:
             latents_mean=[0.0] * z_dim,
             latents_std=[1.0] * z_dim,
         )
+        self.encode_input_shapes: list[tuple[int, ...]] = []
 
     def encode(self, video: torch.Tensor):
+        self.encode_input_shapes.append(tuple(video.shape))
         latent_frames = (video.shape[2] - 1) // self.config.scale_factor_temporal + 1
         latent_height = video.shape[-2] // self.config.scale_factor_spatial
         latent_width = video.shape[-1] // self.config.scale_factor_spatial
@@ -131,9 +134,12 @@ class StubCosmos3Transformer(nn.Module):
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         token = int(text_ids.reshape(-1)[0].item()) if text_ids.numel() else 0
         sound_latents = kwargs.get("sound_latents")
+        control_latents = kwargs.get("control_latents")
+        control_bonus = 100 if control_latents is not None else 0
         self.calls.append(
             {
                 "token": token,
+                "has_control": control_latents is not None,
                 "timestep": timestep.clone(),
                 "text_mask": text_mask.clone(),
                 "cache_before": self.cached_kv,
@@ -145,7 +151,7 @@ class StubCosmos3Transformer(nn.Module):
             self.cached_kv = [(marker, marker + 100)]
             self.cached_freqs_gen = (marker + 200, marker + 300)
         action_latents = kwargs.get("action_latents")
-        outputs: list[torch.Tensor] = [torch.full_like(hidden_states, float(token))]
+        outputs: list[torch.Tensor] = [torch.full_like(hidden_states, float(token + control_bonus))]
         if action_latents is not None:
             outputs.append(torch.full_like(action_latents, float(token + 20)))
         if sound_latents is not None:
@@ -190,6 +196,7 @@ def make_cosmos3_pipeline():
         pipeline._current_flow_shift = 1.0
         pipeline._guidance_scale = None
         pipeline._num_timesteps = None
+        pipeline._cosmos3_branch_caches = None
         pipeline._cache_dit_requires_paired_cfg = False
         pipeline._sound_tokenizer = None
         pipeline.progress_bar = passthrough_progress_bar
@@ -198,6 +205,13 @@ def make_cosmos3_pipeline():
     return _make
 
 
+@pytest.fixture
+def sequential_cfg_parallel(monkeypatch: pytest.MonkeyPatch) -> None:
+    from vllm_omni.diffusion.distributed import cfg_parallel
+
+    monkeypatch.setattr(cfg_parallel, "get_classifier_free_guidance_world_size", lambda: 1)
+
+
 def make_sampling_params(**overrides: Any) -> SimpleNamespace:
     values = {
         "height": None,
@@ -276,12 +290,20 @@ def stub_real_pipeline_init(monkeypatch: pytest.MonkeyPatch):
             return self
 
     class _StubDiffusersScheduler:
-        config = SimpleNamespace(flow_shift=1.0)
+        from_config_calls: list[dict[str, Any]] = []
+
+        def __init__(self, *, flow_shift: float = 1.0) -> None:
+            self.config = SimpleNamespace(flow_shift=flow_shift)
 
         @classmethod
         def from_pretrained(cls, *args, **kwargs):
             return cls()
 
+        @classmethod
+        def from_config(cls, config, **kwargs):
+            cls.from_config_calls.append({"config": config, "kwargs": dict(kwargs)})
+            return cls(flow_shift=kwargs.get("flow_shift", getattr(config, "flow_shift", 1.0)))
+
     class _StubVideoProcessor:
         def __init__(self, *args, **kwargs) -> None:
             pass
@@ -291,6 +313,7 @@ def stub_real_pipeline_init(monkeypatch: pytest.MonkeyPatch):
     monkeypatch.setattr(pipeline_cosmos3, "UniPCMultistepScheduler", _StubDiffusersScheduler)
     monkeypatch.setattr(pipeline_cosmos3, "VideoProcessor", _StubVideoProcessor)
     monkeypatch.setattr(pipeline_cosmos3, "get_local_device", lambda: torch.device("cpu"))
+    return _StubDiffusersScheduler
 
 
 def _make_od_config(*, sound_gen: bool) -> SimpleNamespace:
@@ -332,6 +355,21 @@ def test_pipeline_init_skips_tokenizer_when_sound_disabled(stub_real_pipeline_in
     assert not hasattr(pipeline.transformer, "audio_proj_out")
 
 
+def test_pipeline_init_rebuilds_scheduler_with_cosmos3_defaults(stub_real_pipeline_init) -> None:
+    from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline
+
+    od_config = _make_od_config(sound_gen=False)
+    od_config.flow_shift = 2.5
+
+    pipeline = Cosmos3OmniDiffusersPipeline(od_config=od_config)
+
+    assert stub_real_pipeline_init.from_config_calls
+    call = stub_real_pipeline_init.from_config_calls[-1]
+    assert getattr(call["config"], "flow_shift") == 1.0
+    assert call["kwargs"] == {"flow_shift": 2.5, "use_karras_sigmas": False}
+    assert pipeline._engine_init_flow_shift == 2.5
+
+
 def test_pipeline_init_passes_tokenizer_attrs_into_transformer(
     stub_real_pipeline_init,
     monkeypatch: pytest.MonkeyPatch,
@@ -395,6 +433,160 @@ def test_preprocess_i2v_image_and_action_video_inputs() -> None:
     assert additional["condition_frame_indexes_vision"] == [0, 1]
 
 
+def test_transfer_config_media_helpers_and_preprocess_budget(monkeypatch: pytest.MonkeyPatch) -> None:
+    from vllm_omni.diffusion.models.cosmos3 import transfer
+    from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import (
+        Cosmos3OmniDiffusersPipeline,
+        get_cosmos3_pre_process_func,
+    )
+
+    cfg = transfer.resolve_transfer_config(make_sampling_params(extra_args={"edge": True}))
+    assert cfg is not None
+    assert list(cfg.hints) == ["edge"]
+    assert cfg.guidance_scale == 3.0
+    assert cfg.control_guidance == 1.5
+    assert cfg.flow_shift == 10.0
+    assert cfg.num_video_frames_per_chunk == 93
+    assert cfg.share_vision_temporal_positions is True
+    # fps omitted (no fps/frame_rate on the sampling params) -> wsm preset default (10) applies.
+    defaulted_fps_cfg = transfer.resolve_transfer_config(make_sampling_params(extra_args={"wsm": True}))
+    assert defaulted_fps_cfg is not None
+    assert defaulted_fps_cfg.fps == 10
+    # fps provided (frame_rate set) -> the user value wins over the preset default.
+    explicit_fps_cfg = transfer.resolve_transfer_config(make_sampling_params(frame_rate=24.0, extra_args={"wsm": True}))
+    assert explicit_fps_cfg is not None
+    assert explicit_fps_cfg.fps == 24.0
+    assert (
+        Cosmos3OmniDiffusersPipeline.reference_video_decode_spec(extra_args={"edge": True, "max_frames": 4}).max_frames
+        == 4
+    )
+    frames_for_pad = torch.arange(3 * 3, dtype=torch.uint8).reshape(1, 3, 1, 3)
+    assert transfer.pad_temporal_frames(frames_for_pad, 5)[0, :, 0, 0].tolist() == [0, 3, 6, 6, 3]
+
+    real_import_module = transfer.importlib.import_module
+
+    def raise_missing_cv2(name: str, *args: Any, **kwargs: Any):
+        if name == "cv2":
+            raise ImportError("missing cv2")
+        return real_import_module(name, *args, **kwargs)
+
+    monkeypatch.setattr(transfer.importlib, "import_module", raise_missing_cv2)
+    with pytest.raises(ImportError, match="opencv-python"):
+        transfer.load_or_compute_control_frames(
+            cfg.hints["edge"],
+            height=8,
+            width=8,
+            max_frames=1,
+            input_frames=torch.zeros(3, 1, 8, 8, dtype=torch.uint8),
+        )
+
+    precomputed = torch.zeros(3, 2, 8, 8, dtype=torch.uint8)
+    precomputed_cfg = transfer.resolve_transfer_config(
+        make_sampling_params(extra_args={"edge": {"control": precomputed}})
+    )
+    assert precomputed_cfg is not None
+    loaded = transfer.load_or_compute_control_frames(
+        precomputed_cfg.hints["edge"],
+        height=8,
+        width=8,
+        max_frames=2,
+        input_frames=None,
+    )
+    assert tuple(loaded.shape) == (3, 2, 8, 8)
+
+    preprocess = get_cosmos3_pre_process_func(SimpleNamespace())
+
+    class FramesWithFps(list):
+        fps = 12.5
+
+    frames = FramesWithFps(Image.new("RGB", (8, 4), color) for color in ("red", "green", "blue", "yellow", "black"))
+    request = SimpleNamespace(
+        prompts=[{"prompt": "transfer", "multi_modal_data": {"video": frames}}],
+        sampling_params=SimpleNamespace(
+            height=16,
+            width=32,
+            extra_args={"edge": True, "max_frames": 4, "resolution": "256"},
+        ),
+    )
+    additional = preprocess(request).prompts[0]["additional_information"]
+    assert (request.sampling_params.height, request.sampling_params.width) == (192, 320)
+    assert tuple(additional["preprocessed_transfer_video"].shape) == (1, 3, 4, 192, 320)
+    assert additional["transfer_input_fps"] == 12.5
+    assert "preprocessed_video" not in additional
+
+
+def test_transfer_fps_matches_resolved_frame_rate_precedence() -> None:
+    """When fps and frame_rate differ, transfer must pick frame_rate -- the same
+    precedence as OmniDiffusionSamplingParams.resolved_frame_rate. Uses the real
+    sampling-params dataclass so the resolved_frame_rate property is exercised."""
+    from vllm_omni.diffusion.models.cosmos3 import transfer
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+    sp = OmniDiffusionSamplingParams(fps=24, frame_rate=12.0, extra_args={"edge": True})
+    assert sp.resolved_frame_rate == 12.0
+    cfg = transfer.resolve_transfer_config(sp)
+    assert cfg is not None
+    # edge has no preset fps default, so cfg.fps comes straight from fps resolution.
+    assert cfg.fps == sp.resolved_frame_rate == 12.0
+
+
+def test_transfer_edge_uses_rgb_canny(monkeypatch: pytest.MonkeyPatch) -> None:
+    from vllm_omni.diffusion.models.cosmos3 import transfer
+
+    class FakeCv2:
+        def __init__(self) -> None:
+            self.canny_inputs: list[np.ndarray] = []
+
+        def Canny(self, image, lower, upper):
+            assert (lower, upper) == (100, 200)
+            self.canny_inputs.append(image.copy())
+            return np.zeros(image.shape[:2], dtype=np.uint8)
+
+    fake_cv2 = FakeCv2()
+    monkeypatch.setattr(transfer, "_import_cv2", lambda _hint_key: fake_cv2)
+
+    frames = torch.zeros(3, 1, 4, 5, dtype=torch.uint8)
+    frames[0] = 255
+    edge = transfer.make_edge_control(frames, "medium")
+
+    assert tuple(edge.shape) == (3, 1, 4, 5)
+    assert len(fake_cv2.canny_inputs) == 1
+    assert fake_cv2.canny_inputs[0].shape == (4, 5, 3)
+
+
+def test_transfer_blur_uses_scaled_bilateral(monkeypatch: pytest.MonkeyPatch) -> None:
+    from vllm_omni.diffusion.models.cosmos3 import transfer
+
+    class FakeCv2:
+        INTER_AREA = 1
+        INTER_LINEAR = 2
+        INTER_CUBIC = 3
+
+        def __init__(self) -> None:
+            self.bilateral_calls: list[tuple[tuple[int, int, int], int, float, float]] = []
+
+        def resize(self, image, size, interpolation):
+            del interpolation
+            width, height = size
+            return np.zeros((height, width, image.shape[2]), dtype=image.dtype)
+
+        def bilateralFilter(self, image, diameter, sigma_color, sigma_space):
+            self.bilateral_calls.append((image.shape, diameter, sigma_color, sigma_space))
+            return image
+
+        def GaussianBlur(self, *args, **kwargs):
+            raise AssertionError("Cosmos3 transfer blur should use bilateralFilter, not GaussianBlur.")
+
+    fake_cv2 = FakeCv2()
+    monkeypatch.setattr(transfer, "_import_cv2", lambda _hint_key: fake_cv2)
+
+    frames = torch.zeros(3, 1, 72, 72, dtype=torch.uint8)
+    blurred = transfer.make_blur_control(frames, "high")
+
+    assert tuple(blurred.shape) == (3, 1, 72, 72)
+    assert fake_cv2.bilateral_calls == [((72, 72, 3), 3, 15.0, 10.0)]
+
+
 def test_postprocess_handles_image_video_audio_and_validation() -> None:
     from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import get_cosmos3_post_process_func
 
@@ -543,7 +735,7 @@ def test_prepare_latents_for_video_image_sound_and_action(make_cosmos3_pipeline)
     latents = pipeline._prepare_latents(16, 24, 5, torch.Generator(device="cpu").manual_seed(0))
     assert latents.shape == (1, 2, 2, 2, 3)
 
-    pipeline._encode_conditioning_video = lambda *args, **kwargs: torch.full((1, 2, 2, 2, 3), 5.0)
+    pipeline._encode_conditioning_image_latent = lambda *args, **kwargs: torch.full((1, 2, 1, 2, 3), 5.0)
     i2v_latents, velocity_mask, image_latent = pipeline._prepare_latents_i2v(
         torch.zeros(1, 3, 16, 24), 16, 24, 5, torch.Generator(device="cpu").manual_seed(0)
     )
@@ -593,6 +785,44 @@ def test_prepare_latents_for_video_image_sound_and_action(make_cosmos3_pipeline)
     torch.testing.assert_close(action, clean)
 
 
+def test_prepare_latents_i2v_encodes_only_conditioning_frame(make_cosmos3_pipeline) -> None:
+    pipeline = make_cosmos3_pipeline()
+    calls: list[tuple[str, tuple[int, ...]]] = []
+
+    def record_to_vae_device(tensor: torch.Tensor, *, pin_cpu: bool = False) -> torch.Tensor:
+        assert pin_cpu is True
+        calls.append(("to_vae_device", tuple(tensor.shape)))
+        return tensor
+
+    class RecordingVAE(StubCosmos3VAE):
+        def encode(self, video: torch.Tensor):
+            calls.append(("encode", tuple(video.shape)))
+            return super().encode(video)
+
+    pipeline._to_vae_device = record_to_vae_device
+    pipeline.vae = RecordingVAE(z_dim=2)
+    generator = torch.Generator(device="cpu").manual_seed(0)
+
+    latents, velocity_mask, image_latent = pipeline._prepare_latents_i2v(
+        torch.zeros(1, 3, 16, 24),
+        16,
+        24,
+        9,
+        generator,
+    )
+
+    assert calls == [
+        ("to_vae_device", (1, 3, 16, 24)),
+        ("encode", (1, 3, 1, 16, 24)),
+    ]
+    assert pipeline.vae.encode_input_shapes[-1] == (1, 3, 1, 16, 24)
+    assert latents.shape == (1, 2, 3, 2, 3)
+    assert image_latent.shape == (1, 2, 1, 2, 3)
+    torch.testing.assert_close(latents[:, :, 0:1], image_latent)
+    assert not torch.allclose(latents[:, :, 1:], image_latent.expand(-1, -1, 2, -1, -1))
+    assert velocity_mask.tolist() == [[[[[0.0]], [[1.0]], [[1.0]]]]]
+
+
 def test_diffuse_covers_cfg_i2v_and_multimodal_steps(make_cosmos3_pipeline) -> None:
     pipeline = make_cosmos3_pipeline()
     latents = torch.zeros(1, 2, 1, 1, 1)
@@ -624,6 +854,9 @@ def test_diffuse_covers_cfg_i2v_and_multimodal_steps(make_cosmos3_pipeline) -> N
         image_latent=torch.full((1, 2, 1, 1, 1), 7.0),
     )
     torch.testing.assert_close(i2v[:, :, 0:1], torch.full((1, 2, 1, 1, 1), 7.0))
+    i2v_noise = pipeline.scheduler.step_calls[-1][0]
+    torch.testing.assert_close(i2v_noise[:, :, 0:1], torch.zeros(1, 2, 1, 1, 1))
+    torch.testing.assert_close(i2v_noise[:, :, 1:2], torch.full((1, 2, 1, 1, 1), 2.0))
 
     pipeline.transformer = pipeline.transformer.__class__(latent_channel_size=2, action_gen=True, action_dim=4)
     video_result, action_result = pipeline.diffuse(
@@ -643,6 +876,267 @@ def test_diffuse_covers_cfg_i2v_and_multimodal_steps(make_cosmos3_pipeline) -> N
     torch.testing.assert_close(action_result, torch.full((), 44.0).expand_as(action_result))
 
 
+def test_diffuse_transfer_applies_control_cfg(make_cosmos3_pipeline, sequential_cfg_parallel) -> None:
+    pipeline = make_cosmos3_pipeline()
+    latents = torch.zeros(1, 2, 1, 1, 1)
+    velocity_mask = torch.ones(1, 1, 1, 1, 1)
+
+    result = pipeline.diffuse_transfer(
+        latents=latents,
+        timesteps=torch.tensor([7]),
+        cond_ids=_ids(2),
+        cond_mask=_mask(),
+        uncond_ids=_ids(1),
+        uncond_mask=_mask(),
+        guidance_scale=3.0,
+        control_guidance=1.5,
+        control_guidance_interval=None,
+        control_latents=[torch.zeros_like(latents)],
+        shared_kwargs={"video_shape": (1, 1, 1), "fps": 24.0, "noisy_frame_mask": velocity_mask},
+        velocity_mask=velocity_mask,
+        condition_latents=torch.zeros_like(latents),
+    )
+
+    assert [(call["token"], call["has_control"]) for call in pipeline.transformer.calls] == [
+        (2, True),
+        (2, False),
+        (1, True),
+    ]
+    torch.testing.assert_close(result, torch.full_like(latents, 254.0))
+
+
+def test_diffuse_transfer_skips_idle_cfg_branches(make_cosmos3_pipeline, sequential_cfg_parallel) -> None:
+    latents = torch.zeros(1, 2, 1, 1, 1)
+    velocity_mask = torch.ones(1, 1, 1, 1, 1)
+
+    control_only = make_cosmos3_pipeline()
+    control_result = control_only.diffuse_transfer(
+        latents=latents,
+        timesteps=torch.tensor([7]),
+        cond_ids=_ids(2),
+        cond_mask=_mask(),
+        uncond_ids=_ids(1),
+        uncond_mask=_mask(),
+        guidance_scale=1.0,
+        control_guidance=1.5,
+        control_guidance_interval=None,
+        control_latents=[torch.zeros_like(latents)],
+        shared_kwargs={"video_shape": (1, 1, 1), "fps": 24.0, "noisy_frame_mask": velocity_mask},
+        velocity_mask=velocity_mask,
+        condition_latents=torch.zeros_like(latents),
+    )
+    assert [(call["token"], call["has_control"]) for call in control_only.transformer.calls] == [
+        (2, True),
+        (2, False),
+    ]
+    torch.testing.assert_close(control_result, torch.full_like(latents, 152.0))
+
+    text_only = make_cosmos3_pipeline()
+    text_result = text_only.diffuse_transfer(
+        latents=latents,
+        timesteps=torch.tensor([7]),
+        cond_ids=_ids(2),
+        cond_mask=_mask(),
+        uncond_ids=_ids(1),
+        uncond_mask=_mask(),
+        guidance_scale=3.0,
+        control_guidance=1.0,
+        control_guidance_interval=None,
+        control_latents=[torch.zeros_like(latents)],
+        shared_kwargs={"video_shape": (1, 1, 1), "fps": 24.0, "noisy_frame_mask": velocity_mask},
+        velocity_mask=velocity_mask,
+        condition_latents=torch.zeros_like(latents),
+    )
+    assert [(call["token"], call["has_control"]) for call in text_only.transformer.calls] == [
+        (2, True),
+        (1, True),
+    ]
+    torch.testing.assert_close(text_result, torch.full_like(latents, 104.0))
+
+
+def test_diffuse_transfer_interval_switches_branch_counts(make_cosmos3_pipeline, sequential_cfg_parallel) -> None:
+    pipeline = make_cosmos3_pipeline()
+    latents = torch.zeros(1, 2, 1, 1, 1)
+    velocity_mask = torch.ones(1, 1, 1, 1, 1)
+
+    result = pipeline.diffuse_transfer(
+        latents=latents,
+        timesteps=torch.tensor([900, 500, 100]),
+        cond_ids=_ids(2),
+        cond_mask=_mask(),
+        uncond_ids=_ids(1),
+        uncond_mask=_mask(),
+        guidance_scale=3.0,
+        control_guidance=1.5,
+        control_guidance_interval=(400.0, 1000.0),
+        control_latents=[torch.zeros_like(latents)],
+        shared_kwargs={"video_shape": (1, 1, 1), "fps": 24.0, "noisy_frame_mask": velocity_mask},
+        velocity_mask=velocity_mask,
+        condition_latents=torch.zeros_like(latents),
+        guidance_interval=(800.0, 1000.0),
+    )
+
+    assert [(call["token"], call["has_control"]) for call in pipeline.transformer.calls] == [
+        (2, True),
+        (2, False),
+        (1, True),
+        (2, True),
+        (2, False),
+        (2, True),
+    ]
+    torch.testing.assert_close(result, torch.full_like(latents, 508.0))
+
+
+@pytest.mark.parametrize(("hint_key", "expected_fps"), [("edge", 8.0), ("wsm", 10.0)])
+def test_forward_transfer_uses_source_fps_except_wsm(make_cosmos3_pipeline, hint_key: str, expected_fps: float) -> None:
+    pipeline = make_cosmos3_pipeline()
+    captured: dict[str, Any] = {}
+
+    def fake_format(prompt, negative_prompt, num_frames, frame_rate, height, width, *args, **kwargs):
+        del prompt, negative_prompt, num_frames, height, width, args, kwargs
+        captured["format_frame_rate"] = frame_rate
+        return _ids(2), _mask(), _ids(1), _mask()
+
+    def fake_encode(video: torch.Tensor) -> torch.Tensor:
+        latent_frames = (video.shape[2] - 1) // pipeline.vae_scale_factor_temporal + 1
+        return torch.ones(
+            1,
+            2,
+            latent_frames,
+            max(1, video.shape[-2] // pipeline.vae_scale_factor_spatial),
+            max(1, video.shape[-1] // pipeline.vae_scale_factor_spatial),
+        )
+
+    def fake_prepare(target_norm, current_conditional_frames, generator):
+        del current_conditional_frames, generator
+        latent = fake_encode(target_norm)
+        velocity_mask = torch.ones(1, 1, latent.shape[2], 1, 1)
+        return torch.zeros_like(latent), velocity_mask, torch.zeros_like(latent)
+
+    def fake_diffuse_transfer(**kwargs):
+        captured["shared_kwargs"] = kwargs["shared_kwargs"]
+        return kwargs["latents"]
+
+    pipeline._transfer_bucket_size = lambda sp, source_hw: (16, 16)
+    pipeline._format_and_tokenize_prompts = fake_format
+    pipeline._encode_video_tensor = fake_encode
+    pipeline._prepare_transfer_latents = fake_prepare
+    pipeline.diffuse_transfer = fake_diffuse_transfer
+    pipeline._set_flow_shift = lambda target: captured.setdefault("flow_shifts", []).append(target)
+    pipeline._decode_latents = lambda latents: torch.zeros(1, 3, 5, 16, 16, device="meta")
+
+    control = torch.zeros(3, 5, 16, 16, dtype=torch.uint8)
+    request = SimpleNamespace(
+        prompts=[
+            {
+                "prompt": "transfer",
+                "modalities": ["video"],
+                "additional_information": {
+                    "preprocessed_transfer_video": torch.zeros(1, 3, 5, 16, 16),
+                    "transfer_input_fps": 8.0,
+                },
+            }
+        ],
+        sampling_params=make_sampling_params(
+            height=16,
+            width=16,
+            # fps omitted (no frame_rate) -> non-wsm uses the source video fps (8), wsm uses its preset (10).
+            extra_args={
+                hint_key: {"control": control},
+                "max_frames": 5,
+                "num_video_frames_per_chunk": 5,
+                "show_control_condition": True,
+            },
+        ),
+    )
+
+    output = pipeline.forward(request)
+
+    assert captured["format_frame_rate"] == expected_fps
+    assert captured["shared_kwargs"]["fps"] == expected_fps
+    assert captured["flow_shifts"] == [10.0]
+    assert output.custom_output["fps"] == expected_fps
+    assert output.output["video"].device.type == "meta"
+
+
+def test_forward_transfer_runs_multichunk_overlap_path(
+    make_cosmos3_pipeline,
+    sequential_cfg_parallel,
+) -> None:
+    pipeline = make_cosmos3_pipeline()
+    captured: dict[str, Any] = {"targets": [], "conditional_frames": []}
+
+    pipeline._transfer_bucket_size = lambda sp, source_hw: (16, 16)
+    pipeline._format_and_tokenize_prompts = lambda *args, **kwargs: (_ids(2), _mask(), _ids(1), _mask())
+    pipeline._set_flow_shift = lambda target: captured.setdefault("flow_shifts", []).append(target)
+
+    original_prepare = pipeline._prepare_transfer_latents
+
+    def recording_prepare(target_norm, current_conditional_frames, generator):
+        captured["targets"].append(target_norm.detach().clone())
+        captured["conditional_frames"].append(current_conditional_frames)
+        return original_prepare(target_norm, current_conditional_frames, generator)
+
+    pipeline._prepare_transfer_latents = recording_prepare
+
+    decoded_chunks = [
+        torch.tensor([-0.6, -0.5, -0.4, -0.3, -0.2], dtype=torch.float32),
+        torch.tensor([0.1, 0.2, 0.3, 0.4, 0.5], dtype=torch.float32),
+    ]
+
+    def fake_decode(latents):
+        chunk_values = decoded_chunks[len(captured.setdefault("decode_calls", []))]
+        captured["decode_calls"].append(latents.detach().clone())
+        return chunk_values.view(1, 1, 5, 1, 1).expand(1, 3, 5, 16, 16).clone()
+
+    pipeline._decode_latents = fake_decode
+
+    transfer_video = torch.zeros(1, 3, 8, 16, 16)
+    transfer_video[:, :, 1] = 1.0
+    control = torch.zeros(3, 8, 16, 16, dtype=torch.uint8)
+    request = SimpleNamespace(
+        prompts=[
+            {
+                "prompt": "transfer",
+                "modalities": ["video"],
+                "additional_information": {
+                    "preprocessed_transfer_video": transfer_video,
+                    "transfer_input_fps": 8.0,
+                },
+            }
+        ],
+        sampling_params=make_sampling_params(
+            height=16,
+            width=16,
+            num_inference_steps=1,
+            guidance_scale=1.0,
+            extra_args={
+                "edge": {"control": control},
+                "control_guidance": 1.0,
+                "max_frames": 8,
+                "num_video_frames_per_chunk": 5,
+                "num_conditional_frames": 1,
+                "num_first_chunk_conditional_frames": 2,
+            },
+        ),
+    )
+
+    output = pipeline.forward(request)
+
+    assert captured["conditional_frames"] == [2, 1]
+    assert len(captured["decode_calls"]) == 2
+    assert output.output["video"].shape == (1, 3, 8, 16, 16)
+    torch.testing.assert_close(
+        output.output["video"][0, 0, :, 0, 0],
+        torch.tensor([-0.6, -0.5, -0.4, -0.3, -0.2, 0.2, 0.3, 0.4]),
+    )
+    assert output.custom_output["transfer_controls"]["edge"].shape == (1, 3, 8, 16, 16)
+    torch.testing.assert_close(captured["targets"][0][:, :, 0], torch.full((1, 3, 16, 16), -1.0))
+    torch.testing.assert_close(captured["targets"][0][:, :, 1], torch.full((1, 3, 16, 16), 1.0))
+    torch.testing.assert_close(captured["targets"][0][:, :, 2:], torch.full((1, 3, 3, 16, 16), 1.0))
+    torch.testing.assert_close(captured["targets"][1][:, :, 0], torch.full((1, 3, 16, 16), -0.2))
+
+
 def test_diffuse_keeps_paired_cfg_when_cache_dit_active(make_cosmos3_pipeline) -> None:
     """With cache-dit active the uncond pass is kept even outside the guidance
     interval (so cache-dit's has_separate_cfg parity stays in phase), and the
@@ -680,7 +1174,18 @@ class TestForwardRouting:
     def _install_forward_stubs(self, pipeline):
         captured: dict[str, object] = {"diffuse_calls": [], "prepare_calls": []}
 
-        def fake_format(prompt, negative_prompt, num_frames, frame_rate, height, width, *args, **kwargs):
+        def fake_format(
+            prompt,
+            negative_prompt,
+            num_frames,
+            frame_rate,
+            height,
+            width,
+            max_sequence_length,
+            sp,
+            use_system_prompt=False,
+            is_t2i=False,
+        ):
             captured["format"] = {
                 "prompt": prompt,
                 "negative_prompt": negative_prompt,
@@ -688,7 +1193,8 @@ class TestForwardRouting:
                 "frame_rate": frame_rate,
                 "height": height,
                 "width": width,
-                "is_t2i": kwargs["is_t2i"],
+                "use_system_prompt": use_system_prompt,
+                "is_t2i": is_t2i,
             }
             return _ids(2), _mask(), _ids(1), _mask()
 
@@ -718,12 +1224,26 @@ class TestForwardRouting:
             (
                 {"prompt": "A painted robot", "modalities": ["image"]},
                 make_sampling_params(num_outputs_per_prompt=2),
-                {"key": "image", "is_t2i": True, "flow": [3.0], "steps": [50, 50], "frames": 1},
+                {
+                    "key": "image",
+                    "is_t2i": True,
+                    "use_system_prompt": True,
+                    "flow": [3.0],
+                    "steps": [50, 50],
+                    "frames": 1,
+                },
             ),
             (
                 "A warehouse robot",
                 make_sampling_params(),
-                {"key": "video", "is_t2i": False, "flow": [1.0], "steps": [35], "frames": 189},
+                {
+                    "key": "video",
+                    "is_t2i": False,
+                    "use_system_prompt": True,
+                    "flow": [1.0],
+                    "steps": [35],
+                    "frames": 189,
+                },
             ),
         ],
     )
@@ -741,6 +1261,7 @@ class TestForwardRouting:
 
         assert expected["key"] in output.output
         assert captured["format"]["is_t2i"] is expected["is_t2i"]
+        assert captured["format"]["use_system_prompt"] is expected["use_system_prompt"]
         assert captured["format"]["num_frames"] == expected["frames"]
         assert captured["flow_shifts"] == expected["flow"]
         assert [call[0] for call in pipeline.scheduler.set_timesteps_calls] == expected["steps"]
@@ -769,6 +1290,7 @@ class TestForwardRouting:
             )
         )
         assert captured["diffuse_calls"][-1]["shared_kwargs"]["noisy_frame_mask"] is velocity_mask
+        assert captured["format"]["use_system_prompt"] is True
 
         video_tensor = torch.zeros(1, 3, 5, 16, 16)
         v2v_condition = torch.full((1, 2, 2, 1, 1), 4.0)
@@ -795,6 +1317,7 @@ class TestForwardRouting:
         )
         assert captured["flow_shifts"][-1] == 10.0
         assert captured["format"]["negative_prompt"] == ""
+        assert captured["format"]["use_system_prompt"] is True
         assert captured["diffuse_calls"][-1]["shared_kwargs"]["noisy_frame_mask"] is v2v_mask
         assert captured["diffuse_calls"][-1]["condition_latents"] is v2v_condition
 
@@ -810,6 +1333,7 @@ class TestForwardRouting:
             )
         )
         assert captured["diffuse_calls"][-1]["sound_latents"] is sound_latents
+        assert captured["format"]["use_system_prompt"] is True
         assert output.output["audio_sample_rate"] == 10
 
         pipeline.transformer = pipeline.transformer.__class__(latent_channel_size=2, action_gen=True, action_dim=4)
@@ -835,6 +1359,7 @@ class TestForwardRouting:
             )
         )
         assert captured["diffuse_calls"][-1]["shared_kwargs"]["action_domain_ids"].tolist() == [7]
+        assert captured["format"]["use_system_prompt"] is False
         assert output.custom_output["action"].shape == (1, 2, 2)
         assert "action_only_output" not in output.custom_output
 
@@ -910,6 +1435,7 @@ class TestForwardRouting:
             "frame_rate": 15.0,
             "height": 16,
             "width": 16,
+            "use_system_prompt": False,
             "is_t2i": False,
         }
         assert "flow_shifts" not in captured
@@ -936,6 +1462,26 @@ class TestForwardRouting:
                 make_sampling_params(),
                 "only for video",
             ),
+            (
+                [{"prompt": "x", "modalities": ["image"]}],
+                make_sampling_params(extra_args={"edge": {"control_path": "/tmp/control.mp4"}}),
+                "transfer inference is supported only for video outputs",
+            ),
+            (
+                [{"prompt": "x", "modalities": ["video"], "generate_sound": True}],
+                make_sampling_params(extra_args={"edge": {"control_path": "/tmp/control.mp4"}}),
+                "cannot be combined with sound generation",
+            ),
+            (
+                [{"prompt": "x", "modalities": ["video"]}],
+                make_sampling_params(
+                    extra_args={
+                        "edge": {"control_path": "/tmp/control.mp4"},
+                        "action_mode": "policy",
+                    }
+                ),
+                "cannot be combined with action generation",
+            ),
         ],
     )
     def test_forward_rejects_invalid_public_requests(

tests/diffusion/models/cosmos3/test_cosmos3_transformer.py

@@ -131,6 +131,38 @@ def test_forward_returns_video_prediction(monkeypatch: pytest.MonkeyPatch) -> No
     assert tuple(output.shape) == (1, 2, 1, 2, 2)
 
 
+def test_forward_accepts_transfer_control_latents(monkeypatch: pytest.MonkeyPatch) -> None:
+    from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
+
+    monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None))
+
+    model = transformer_cosmos3.Cosmos3VFMTransformer(
+        SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32)
+    )
+    hidden_states = torch.zeros(1, 2, 1, 2, 2)
+    output = model(
+        hidden_states=hidden_states,
+        timestep=torch.tensor([1.0]),
+        text_ids=torch.tensor([[1, 2]], dtype=torch.long),
+        text_mask=torch.ones(1, 2, dtype=torch.long),
+        video_shape=(1, 2, 2),
+        fps=24.0,
+        control_latents=[torch.ones_like(hidden_states), torch.full_like(hidden_states, 2.0)],
+    )
+
+    assert tuple(output.shape) == tuple(hidden_states.shape)
+    with pytest.raises(ValueError, match="control latent shape"):
+        model(
+            hidden_states=hidden_states,
+            timestep=torch.tensor([1.0]),
+            text_ids=torch.tensor([[1, 2]], dtype=torch.long),
+            text_mask=torch.ones(1, 2, dtype=torch.long),
+            video_shape=(1, 2, 2),
+            fps=24.0,
+            control_latents=[torch.zeros(1, 2, 2, 2, 2)],
+        )
+
+
 def test_forward_gathers_gen_tokens_before_unpatchify(monkeypatch: pytest.MonkeyPatch) -> None:
     from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
 
@@ -354,3 +386,33 @@ def test_compute_rope_freqs_places_text_video_action_and_sound_positions() -> No
     _, gen_pos = rotary.position_ids
     assert gen_pos.shape == (3, 1, 5)
     assert gen_pos[0, 0].tolist() == [102, 103, 103, 104, 102]
+
+    rotary.position_ids.clear()
+    model._compute_rope_freqs(
+        text_mask=torch.tensor([[1, 1]], dtype=torch.long),
+        t=2,
+        hp=1,
+        wp=1,
+        fps=24.0,
+        device=torch.device("cpu"),
+        dtype=torch.float32,
+        num_vision_items=3,
+        share_vision_temporal_positions=True,
+    )
+    _, shared_gen_pos = rotary.position_ids
+    assert shared_gen_pos[0, 0].tolist() == [102, 103, 102, 103, 102, 103]
+
+    rotary.position_ids.clear()
+    model._compute_rope_freqs(
+        text_mask=torch.tensor([[1, 1]], dtype=torch.long),
+        t=2,
+        hp=1,
+        wp=1,
+        fps=24.0,
+        device=torch.device("cpu"),
+        dtype=torch.float32,
+        num_vision_items=3,
+        share_vision_temporal_positions=False,
+    )
+    _, offset_gen_pos = rotary.position_ids
+    assert offset_gen_pos[0, 0].tolist() == [102, 103, 104, 105, 106, 107]

tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py

@@ -466,11 +466,18 @@ class TestCFGParallelForwardPath:
     """Test the LTX-2.3 CFG-parallel denoising path without loading model weights."""
 
     @pytest.mark.parametrize(("cfg_rank", "expected_prompt_value"), [(0, 1.0), (1, 0.0)])
+    @pytest.mark.parametrize(
+        ("frame_rate_input", "audio_sampling_rate", "expected_frame_rate"),
+        [(1.0, 1, 1.0), (None, 24, 24.0)],
+    )
     def test_forward_cfg_parallel_steps_video_and_audio_scheduler(
         self,
         monkeypatch,
         cfg_rank,
         expected_prompt_value,
+        frame_rate_input,
+        audio_sampling_rate,
+        expected_frame_rate,
     ):
         from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2_3 as ltx23
         from vllm_omni.diffusion.request import OmniDiffusionRequest
@@ -484,7 +491,7 @@ class TestCFGParallelForwardPath:
         pipe.vae_temporal_compression_ratio = 1
         pipe.transformer_spatial_patch_size = 1
         pipe.transformer_temporal_patch_size = 1
-        pipe.audio_sampling_rate = 1
+        pipe.audio_sampling_rate = audio_sampling_rate
         pipe.audio_hop_length = 1
         pipe.audio_vae_temporal_compression_ratio = 1
         pipe.audio_vae_mel_compression_ratio = 1
@@ -553,8 +560,11 @@ class TestCFGParallelForwardPath:
                 assert prompt_embeds.shape[0] == 2
                 return prompt_embeds, prompt_embeds, prompt_attention_mask
 
+        rope_video_fps: list[float] = []
+
         class FakeRope:
             def prepare_video_coords(self, batch_size, num_frames, height, width, device, fps):
+                rope_video_fps.append(fps)
                 return torch.zeros(batch_size, num_frames * height * width, 3, device=device)
 
             def prepare_audio_coords(self, batch_size, num_frames, device):
@@ -611,7 +621,7 @@ class TestCFGParallelForwardPath:
                 height=32,
                 width=32,
                 num_frames=1,
-                frame_rate=1.0,
+                frame_rate=frame_rate_input,
                 num_inference_steps=2,
                 guidance_scale=4.0,
                 latents=video_latents,
@@ -651,6 +661,11 @@ class TestCFGParallelForwardPath:
         torch.testing.assert_close(video_out, (video_latents - 2 * expected_video_noise).reshape(1, 2, 1, 1, 1))
         torch.testing.assert_close(audio_out, (audio_latents - 2 * expected_audio_noise).reshape(1, 1, 1, 2))
 
+        # fps regression guard: an omitted request fps (frame_rate_input=None) must resolve
+        # to the model's own 24.0 default, not crash on None; a provided rate is passed through.
+        assert rope_video_fps
+        assert all(fps == expected_frame_rate for fps in rope_video_fps)
+
 
 class TestRegistryIntegration:
     """Verify all LTX-2.3 pipeline variants are registered."""

tests/diffusion/test_multiproc_engine_concurrency.py

@@ -730,6 +730,39 @@ class TestStageDiffusionClientErrorPropagation:
                 batch_size=1,
             )
 
+    @pytest.mark.asyncio
+    async def test_collective_rpc_async_returns_none_result(self, monkeypatch):
+        client = self._make_client()
+        client._owns_process = False
+        client._proc = None
+        client._encoder.encode.return_value = b"encoded-rpc"
+
+        async def _unexpected_poll(*_, **__):
+            raise AssertionError("collective_rpc_async should not keep polling after a None rpc_result arrives")
+
+        client._response_poller = SimpleNamespace(poll=_unexpected_poll)
+
+        rpc_id = "rpc-none"
+        monkeypatch.setattr(
+            "vllm_omni.diffusion.stage_diffusion_client.uuid.uuid4",
+            lambda: SimpleNamespace(hex=rpc_id),
+        )
+
+        def _drain() -> None:
+            client._rpc_results[rpc_id] = None
+
+        client._drain_responses = _drain
+
+        result = await client.collective_rpc_async(
+            method="profile",
+            timeout=0.01,
+            args=(False, None),
+        )
+
+        assert result is None
+        client._request_socket.send.assert_called_once_with(b"encoded-rpc")
+        assert rpc_id not in client._pending_rpcs
+
 
 # ───────── monitor thread & death sentinel integration tests ─────────
 

tests/diffusion/utils/test_kv_utils.py

@@ -0,0 +1,75 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for kv utils."""
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.utils.kv_utils import left_pad_stack
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_uniform_lengths():
+    """Ensure uniform 3D tensors correctly stack to 4D and have no mask."""
+    tensors = [torch.randn(10, 4, 128) for _ in range(3)]
+    stacked, mask = left_pad_stack(tensors)
+    assert stacked.shape == (3, 10, 4, 128)
+    assert mask is None
+    for i in range(3):
+        assert torch.equal(stacked[i], tensors[i])
+
+
+def test_variable_lengths():
+    """Ensure variable 3D tensors correctly stack to 4D with a mask."""
+    t1 = torch.ones(5, 2, 4)
+    t2 = torch.ones(8, 2, 4)
+    t3 = torch.ones(3, 2, 4)
+    stacked, mask = left_pad_stack([t1, t2, t3])
+
+    assert stacked.shape == (3, 8, 2, 4)
+    assert mask is not None
+    assert mask.shape == (3, 8)
+    # Ensure summing over dim 1 gives our seq lens back
+    assert mask.sum(dim=1).tolist() == [5, 8, 3]
+
+
+def test_single_tensor_is_4d():
+    """Ensure a single 3D tensor is expanded to 4D."""
+    t = torch.randn(7, 4, 128)
+    stacked, mask = left_pad_stack([t])
+    assert stacked.shape == (1, 7, 4, 128)
+    assert mask is None
+    assert torch.equal(stacked[0], t)
+
+
+def test_preserves_device_and_dtype():
+    """Ensure device/dtype is preserved."""
+    t1 = torch.randn(3, 2, dtype=torch.bfloat16)
+    t2 = torch.randn(5, 2, dtype=torch.bfloat16)
+    stacked, mask = left_pad_stack([t1, t2])
+    assert stacked.dtype == torch.bfloat16
+    assert stacked.device == t1.device
+    assert mask.device == t1.device
+
+
+def test_mismatched_trailing_shapes_raises():
+    """Ensure that mismatched dims outside of 0 explodes."""
+    t1 = torch.randn(5, 4, 128)
+    t2 = torch.randn(5, 8, 128)
+    with pytest.raises(ValueError):
+        left_pad_stack([t1, t2])
+
+
+def test_mismatched_ndim_raises():
+    """Ensure that mismatched ndims explodes."""
+    t1 = torch.randn(5, 4)
+    t2 = torch.randn(5, 4, 128)
+    with pytest.raises(ValueError):
+        left_pad_stack([t1, t2])
+
+
+def test_empty_list_raises():
+    """Ensure that tensors must be nonempty."""
+    with pytest.raises(ValueError):
+        left_pad_stack([])

tests/distributed/omni_connectors/test_bagel_shared_memory_connector.py

@@ -7,7 +7,7 @@ End-to-end tests for Bagel with shared memory connector: img2img and text2img.
 - img2img: validates output vs reference pixels within a ±10 tolerance.
 - text2img: validates output vs reference pixels within a ±5 tolerance
   (equivalent to `examples/offline_inference/bagel/end2end.py` with
-  text2img modality and 15 steps).
+  text2img modality and 14 steps).
 """
 
 import os
@@ -46,7 +46,6 @@ REFERENCE_PIXELS = [
     {"position": (256, 256), "rgb": (181, 201, 221)},
 ]
 
-
 # text2img reference pixels (aligned with offline `bagel/end2end.py` text2img, 15 steps)
 # "Generated with seed=52, num_inference_steps=15,
 # prompt='A cute cat'"

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
 

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)

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"

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 <img> placeholders.
 

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
 # =============================================================================

tests/entrypoints/openai_api/test_serving_speech.py

@@ -1826,15 +1826,27 @@ class TestFileValidationFunctions:
 
 
 class TestStreamingProtocolValidation:
-    """Unit tests for the stream field validators in OpenAICreateSpeechRequest."""
+    """Unit tests for streaming validators in OpenAICreateSpeechRequest."""
+
+    def test_default_is_non_streaming(self):
+        req = OpenAICreateSpeechRequest(input="Hello")
+        assert req.stream is False
+        assert req.stream_format is None
+        assert req.is_streaming() is False
 
     def test_stream_validation_errors(self):
-        """stream=True requires response_format not in ('pcm', 'wav') and speed=1.0."""
+        """stream=True requires response_format in ('pcm', 'wav') and speed=1.0."""
         with pytest.raises(ValidationError, match="requires response_format='pcm' or 'wav'"):
             OpenAICreateSpeechRequest(input="Hello", stream=True, response_format="mp3")
         with pytest.raises(ValidationError, match="Speed adjustment is not supported"):
             OpenAICreateSpeechRequest(input="Hello", stream=True, response_format="pcm", speed=2.0)
 
+    def test_stream_format_audio_validation_errors(self):
+        with pytest.raises(ValidationError, match="requires response_format='pcm' or 'wav'"):
+            OpenAICreateSpeechRequest(input="Hello", stream_format="audio", response_format="mp3")
+        with pytest.raises(ValidationError, match="Speed adjustment is not supported"):
+            OpenAICreateSpeechRequest(input="Hello", stream_format="audio", response_format="pcm", speed=2.0)
+
     def test_stream_valid(self):
         """stream=True + response_format in ('pcm', 'wav') + speed=1.0 is accepted."""
         req = OpenAICreateSpeechRequest(input="Hello", stream=True, response_format="pcm")
@@ -1843,10 +1855,17 @@ class TestStreamingProtocolValidation:
         req = OpenAICreateSpeechRequest(input="Hello", stream=True, response_format="wav")
         assert req.stream is True
 
-    def test_sse_stream_format_is_blocked(self):
-        """stream_format='sse' is blocked."""
-        with pytest.raises(ValidationError, match="sse"):
-            OpenAICreateSpeechRequest(input="Hello", stream_format="sse")
+    def test_stream_format_audio_is_valid(self):
+        req = OpenAICreateSpeechRequest(input="Hello", stream_format="audio", response_format="pcm")
+        assert req.stream_format == "audio"
+        assert req.is_raw_audio_stream() is True
+
+    def test_sse_stream_format_is_valid(self):
+        """stream_format='sse' is accepted for /audio/speech."""
+        req = OpenAICreateSpeechRequest(input="Hello", stream_format="sse")
+
+        assert req.stream_format == "sse"
+        assert req.is_sse_stream() is True
 
 
 class TestStreamingResponse:
@@ -1927,6 +1946,146 @@ class TestStreamingResponse:
         assert "audio/pcm" in response.headers["content-type"]
         assert len(response.content) > 0
 
+    def test_stream_format_audio_streaming(self, streaming_app):
+        """stream_format=audio without stream=True returns raw audio/pcm chunks."""
+        client = TestClient(streaming_app)
+        response = client.post(
+            "/v1/audio/speech",
+            json={"input": "Hello", "stream_format": "audio", "response_format": "pcm"},
+        )
+        assert response.status_code == 200
+        assert "audio/pcm" in response.headers["content-type"]
+        assert "text/event-stream" not in response.headers["content-type"]
+        assert len(response.content) > 0
+
+    def test_sse_streaming(self, streaming_app):
+        """stream_format=sse without stream=True returns audio deltas as SSE."""
+        client = TestClient(streaming_app)
+        response = client.post(
+            "/v1/audio/speech",
+            json={"input": "Hello", "stream_format": "sse", "response_format": "pcm"},
+        )
+
+        assert response.status_code == 200
+        assert "text/event-stream" in response.headers["content-type"]
+        body = response.text
+        assert "event: speech.audio.delta" in body
+        assert "event: speech.audio.done" in body
+        data_line = next(line for line in body.splitlines() if line.startswith("data: "))
+        payload = json.loads(data_line.removeprefix("data: "))
+        assert payload["type"] == "speech.audio.delta"
+        assert payload["response_format"] == "pcm"
+        assert base64.b64decode(payload["audio"])
+
+    def test_stream_true_prefers_raw_audio_over_sse(self, streaming_app):
+        """stream=True keeps the existing raw audio stream behavior even with stream_format=sse."""
+        client = TestClient(streaming_app)
+        response = client.post(
+            "/v1/audio/speech",
+            json={"input": "Hello", "stream": True, "stream_format": "sse", "response_format": "pcm"},
+        )
+
+        assert response.status_code == 200
+        assert "audio/pcm" in response.headers["content-type"]
+        assert "text/event-stream" not in response.headers["content-type"]
+        assert len(response.content) > 0
+
+    def test_sse_rejects_unsupported_response_format(self, streaming_app):
+        """stream_format=sse with a non-pcm/wav format must fail before streaming starts."""
+        client = TestClient(streaming_app)
+        response = client.post(
+            "/v1/audio/speech",
+            json={"input": "Hello", "stream_format": "sse", "response_format": "mp3"},
+        )
+
+        assert response.status_code in (400, 422)
+        assert "text/event-stream" not in response.headers.get("content-type", "")
+
+    def test_sse_rejects_speed_adjustment(self, streaming_app):
+        """stream_format=sse with speed != 1.0 must fail before streaming starts."""
+        client = TestClient(streaming_app)
+        response = client.post(
+            "/v1/audio/speech",
+            json={"input": "Hello", "stream_format": "sse", "response_format": "pcm", "speed": 2.0},
+        )
+
+        assert response.status_code in (400, 422)
+        assert "text/event-stream" not in response.headers.get("content-type", "")
+
+    def test_stream_format_audio_rejects_unsupported_response_format(self, streaming_app):
+        client = TestClient(streaming_app)
+        response = client.post(
+            "/v1/audio/speech",
+            json={"input": "Hello", "stream_format": "audio", "response_format": "mp3"},
+        )
+
+        assert response.status_code in (400, 422)
+        assert "audio/" not in response.headers.get("content-type", "")
+
+    def test_stream_format_audio_rejects_speed_adjustment(self, streaming_app):
+        client = TestClient(streaming_app)
+        response = client.post(
+            "/v1/audio/speech",
+            json={"input": "Hello", "stream_format": "audio", "response_format": "pcm", "speed": 2.0},
+        )
+
+        assert response.status_code in (400, 422)
+        assert "audio/" not in response.headers.get("content-type", "")
+
+    @pytest.fixture
+    def erroring_streaming_app(self, mocker: MockerFixture):
+        """Test app whose mock engine raises mid-stream, to exercise the SSE error event."""
+
+        async def mock_generate_streaming(*args, **kwargs):
+            raise RuntimeError("boom: simulated engine failure")
+            yield  # pragma: no cover - generator marker, unreachable
+
+        mock_engine_client = mocker.MagicMock()
+        mock_engine_client.errored = False
+        mock_engine_client.generate = mocker.MagicMock(side_effect=mock_generate_streaming)
+        mock_engine_client.default_sampling_params_list = [{}]
+        mock_models = mocker.MagicMock()
+        mock_models.is_base_model.return_value = True
+
+        speech_server = OmniOpenAIServingSpeech(
+            engine_client=mock_engine_client,
+            models=mock_models,
+            request_logger=mocker.MagicMock(),
+        )
+
+        original_create_speech = speech_server.create_speech
+        sig = signature(original_create_speech)
+        new_parameters = [p for name, p in sig.parameters.items() if name != "raw_request"]
+        new_sig = Signature(parameters=new_parameters, return_annotation=sig.return_annotation)
+
+        async def awaitable_create_speech(*args, **kwargs):
+            return await original_create_speech(*args, **kwargs)
+
+        awaitable_create_speech.__signature__ = new_sig
+        speech_server.create_speech = awaitable_create_speech
+
+        app = FastAPI()
+        app.add_api_route("/v1/audio/speech", speech_server.create_speech, methods=["POST"], response_model=None)
+        return app
+
+    def test_sse_emits_error_event_on_generator_failure(self, erroring_streaming_app):
+        """An exception inside the SSE generator must surface as a speech.audio.error event."""
+        client = TestClient(erroring_streaming_app)
+        response = client.post(
+            "/v1/audio/speech",
+            json={"input": "Hello", "stream_format": "sse", "response_format": "pcm"},
+        )
+
+        assert response.status_code == 200
+        assert "text/event-stream" in response.headers["content-type"]
+        body = response.text
+        assert "event: speech.audio.error" in body
+        error_line = next(line for line in body.splitlines() if line.startswith("data: "))
+        payload = json.loads(error_line.removeprefix("data: "))
+        assert payload["type"] == "speech.audio.error"
+        assert "error" in payload
+        assert payload["error"]["message"]
+
     def test_non_streaming_unchanged(self, streaming_app):
         """Non-streaming path must still return audio/wav."""
         client = TestClient(streaming_app)
@@ -2883,6 +3042,40 @@ class TestTTSAsyncOffloading:
         qwen3_tts_server._build_tts_params.assert_called_once()
         qwen3_tts_server._estimate_prompt_len_async.assert_awaited_once()
 
+    def test_prepare_speech_generation_treats_sse_as_streaming(self, qwen3_tts_server, mocker: MockerFixture):
+        """stream_format=sse should request delta-style multimodal outputs."""
+        qwen3_tts_server._validate_tts_request = mocker.MagicMock(return_value=None)
+        qwen3_tts_server._build_tts_params = mocker.MagicMock(
+            return_value={"text": ["hello"], "task_type": ["CustomVoice"], "speaker": ["Vivian"]}
+        )
+        qwen3_tts_server._estimate_prompt_len_async = mocker.AsyncMock(return_value=512)
+        mock_coerce = mocker.patch(
+            "vllm_omni.entrypoints.openai.serving_speech.coerce_param_message_types",
+            return_value=qwen3_tts_server.engine_client.default_sampling_params_list,
+        )
+        request = OpenAICreateSpeechRequest(input="hello", stream_format="sse")
+
+        asyncio.run(qwen3_tts_server._prepare_speech_generation(request))
+
+        mock_coerce.assert_called_once_with(qwen3_tts_server.engine_client.default_sampling_params_list, True)
+
+    def test_prepare_speech_generation_treats_audio_as_streaming(self, qwen3_tts_server, mocker: MockerFixture):
+        """stream_format=audio should request delta-style multimodal outputs."""
+        qwen3_tts_server._validate_tts_request = mocker.MagicMock(return_value=None)
+        qwen3_tts_server._build_tts_params = mocker.MagicMock(
+            return_value={"text": ["hello"], "task_type": ["CustomVoice"], "speaker": ["Vivian"]}
+        )
+        qwen3_tts_server._estimate_prompt_len_async = mocker.AsyncMock(return_value=512)
+        mock_coerce = mocker.patch(
+            "vllm_omni.entrypoints.openai.serving_speech.coerce_param_message_types",
+            return_value=qwen3_tts_server.engine_client.default_sampling_params_list,
+        )
+        request = OpenAICreateSpeechRequest(input="hello", stream_format="audio", response_format="pcm")
+
+        asyncio.run(qwen3_tts_server._prepare_speech_generation(request))
+
+        mock_coerce.assert_called_once_with(qwen3_tts_server.engine_client.default_sampling_params_list, True)
+
     def test_prepare_speech_generation_qwen3_voicedesign_non_streaming_mode_false(
         self, qwen3_tts_server, mocker: MockerFixture
     ):

tests/entrypoints/openai_api/test_serving_speech_stream.py

@@ -1,4 +1,6 @@
 import asyncio
+import base64
+from types import SimpleNamespace
 
 import pytest
 from fastapi import FastAPI, WebSocket
@@ -9,6 +11,7 @@ from starlette.websockets import WebSocketDisconnect
 from vllm_omni.entrypoints.openai import serving_speech_stream as streaming_speech_module
 from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
 from vllm_omni.entrypoints.openai.serving_speech_stream import OmniStreamingSpeechHandler
+from vllm_omni.utils.forced_aligner import WordTimestamp
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -25,10 +28,11 @@ def _build_test_app(
         speech_service = mocker.MagicMock(spec=OmniOpenAIServingSpeech)
         speech_service._generate_audio_bytes = mocker.AsyncMock(return_value=(b"RIFF" + b"\x00" * 32, "audio/wav"))
         speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req-1", object(), {}))
+        speech_service.forced_aligner_config = None
 
-        async def mock_generate_pcm_chunks(_generator, _request_id):
+        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False):
             for chunk in (b"\x01\x02", b"\x03\x04\x05"):
-                yield chunk
+                yield (chunk, 24000) if include_sample_rate else chunk
 
         speech_service._generate_pcm_chunks = mock_generate_pcm_chunks
         speech_service.engine_client = mocker.MagicMock()
@@ -82,6 +86,7 @@ class TestStreamingSpeechWebSocket:
         speech_service._generate_audio_bytes = mocker.AsyncMock(return_value=(b"", "audio/wav"))
         speech_service.engine_client = mocker.MagicMock()
         speech_service.engine_client.abort = mocker.AsyncMock()
+        speech_service.forced_aligner_config = None
 
         async def mock_prepare_speech_generation(request):
             captured_requests.append(request)
@@ -89,9 +94,9 @@ class TestStreamingSpeechWebSocket:
 
         speech_service._prepare_speech_generation = mock_prepare_speech_generation
 
-        async def mock_generate_pcm_chunks(_generator, _request_id):
+        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False):
             for chunk in (b"\x01\x02", b"\x03\x04\x05", b"\x06"):
-                yield chunk
+                yield (chunk, 24000) if include_sample_rate else chunk
 
         speech_service._generate_pcm_chunks = mock_generate_pcm_chunks
         app, _ = _build_test_app(speech_service)
@@ -130,6 +135,164 @@ class TestStreamingSpeechWebSocket:
         assert captured_requests[0].initial_codec_chunk_frames == 12
         assert speech_service._generate_audio_bytes.await_count == 0
 
+    def test_word_timestamps_requires_configured_aligner(self, mocker: MockerFixture):
+        app, _ = _build_test_app(mocker=mocker)
+
+        with TestClient(app) as client:
+            with client.websocket_connect("/v1/audio/speech/stream") as ws:
+                ws.send_json(
+                    {
+                        "type": "session.config",
+                        "voice": "Vivian",
+                        "stream_audio": True,
+                        "response_format": "pcm",
+                        "word_timestamps": True,
+                    }
+                )
+                ws.send_json({"type": "input.text", "text": "Hello world. "})
+
+                error = ws.receive_json()
+                assert error["type"] == "error"
+                assert "without --forced-aligner" in error["message"]
+
+    def test_word_timestamps_emit_sidecar_json_frame(self, mocker: MockerFixture):
+        captured_requests = []
+        speech_service = mocker.MagicMock(spec=OmniOpenAIServingSpeech)
+        speech_service._generate_audio_bytes = mocker.AsyncMock(return_value=(b"", "audio/wav"))
+        speech_service.engine_client = mocker.MagicMock()
+        speech_service.engine_client.abort = mocker.AsyncMock()
+        speech_service.forced_aligner_config = SimpleNamespace(model="aligner")
+
+        async def mock_prepare_speech_generation(request):
+            captured_requests.append(request)
+            return "req-stream", object(), {}
+
+        speech_service._prepare_speech_generation = mock_prepare_speech_generation
+
+        first_chunk = b"\x01" * 1000
+        second_chunk = b"\x02" * 1000
+
+        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False):
+            for chunk in (first_chunk, second_chunk):
+                yield (chunk, 1000) if include_sample_rate else chunk
+
+        speech_service._generate_pcm_chunks = mock_generate_pcm_chunks
+        # Sentence-level: aligner runs once over the whole sentence audio.
+        mock_align = mocker.AsyncMock(
+            return_value=[
+                WordTimestamp("Hello", 0, 200),
+                WordTimestamp("world", 200, 900),
+            ]
+        )
+        mocker.patch.object(streaming_speech_module, "forced_align", mock_align)
+        app, _ = _build_test_app(speech_service)
+
+        with TestClient(app) as client:
+            with client.websocket_connect("/v1/audio/speech/stream") as ws:
+                ws.send_json(
+                    {
+                        "type": "session.config",
+                        "voice": "Vivian",
+                        "stream_audio": True,
+                        "response_format": "pcm",
+                        "word_timestamps": True,
+                    }
+                )
+                ws.send_json({"type": "input.text", "text": "Hello world. "})
+
+                start = ws.receive_json()
+                assert start["type"] == "audio.start"
+                assert start["word_timestamps"] is True
+
+                # Audio streams first; timestamps are null until the sentence
+                # is fully aligned.
+                chunk = ws.receive_json()
+                assert chunk["type"] == "audio.chunk"
+                assert chunk["chunk_id"] == 0
+                assert chunk["chunk_start_ms"] == 0
+                assert chunk["chunk_end_ms"] == 500
+                assert chunk["sample_rate"] == 1000
+                assert base64.b64decode(chunk["audio_b64"]) == first_chunk
+                assert chunk["timestamps"] is None
+
+                chunk = ws.receive_json()
+                assert chunk["type"] == "audio.chunk"
+                assert chunk["chunk_id"] == 1
+                assert chunk["chunk_start_ms"] == 500
+                assert chunk["chunk_end_ms"] == 1000
+                assert chunk["sample_rate"] == 1000
+                assert base64.b64decode(chunk["audio_b64"]) == second_chunk
+                assert chunk["timestamps"] is None
+
+                # Final frame: empty audio carrying the whole-sentence timestamps.
+                chunk = ws.receive_json()
+                assert chunk["type"] == "audio.chunk"
+                assert chunk["chunk_id"] == 2
+                assert chunk["chunk_start_ms"] == 0
+                assert chunk["chunk_end_ms"] == 1000
+                assert chunk["sample_rate"] == 1000
+                assert base64.b64decode(chunk["audio_b64"]) == b""
+                assert chunk["timestamps"] == [
+                    {"word": "Hello", "start_ms": 0, "end_ms": 200},
+                    {"word": "world", "start_ms": 200, "end_ms": 900},
+                ]
+
+                done = ws.receive_json()
+                assert done == {"type": "audio.done", "sentence_index": 0, "total_bytes": 2000, "error": False}
+
+        assert captured_requests[0].word_timestamps is True
+        assert mock_align.await_count == 1
+
+    def test_word_timestamps_emit_word_dicts(self, mocker: MockerFixture):
+        # The streaming layer forwards the aligner's (already monotonic,
+        # non-overlapping) words as JSON dicts in the trailing frame.
+        speech_service = mocker.MagicMock(spec=OmniOpenAIServingSpeech)
+        speech_service._generate_audio_bytes = mocker.AsyncMock(return_value=(b"", "audio/wav"))
+        speech_service.engine_client = mocker.MagicMock()
+        speech_service.engine_client.abort = mocker.AsyncMock()
+        speech_service.forced_aligner_config = SimpleNamespace(model="aligner")
+        speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req", object(), {}))
+
+        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False):
+            chunk = b"\x01" * 1000
+            yield (chunk, 1000) if include_sample_rate else chunk
+
+        speech_service._generate_pcm_chunks = mock_generate_pcm_chunks
+        mock_align = mocker.AsyncMock(
+            return_value=[
+                WordTimestamp("Hello", 0, 1000),
+                WordTimestamp("world", 1000, 1200),
+            ]
+        )
+        mocker.patch.object(streaming_speech_module, "forced_align", mock_align)
+        app, _ = _build_test_app(speech_service)
+
+        with TestClient(app) as client:
+            with client.websocket_connect("/v1/audio/speech/stream") as ws:
+                ws.send_json(
+                    {
+                        "type": "session.config",
+                        "voice": "Vivian",
+                        "stream_audio": True,
+                        "response_format": "pcm",
+                        "word_timestamps": True,
+                    }
+                )
+                ws.send_json({"type": "input.text", "text": "Hello world. "})
+
+                assert ws.receive_json()["type"] == "audio.start"
+                # Real-time audio chunk (timestamps null), then the timestamp frame.
+                assert ws.receive_json()["timestamps"] is None
+                final = ws.receive_json()
+                assert final["type"] == "audio.chunk"
+                timestamps = final["timestamps"]
+                assert timestamps == [
+                    {"word": "Hello", "start_ms": 0, "end_ms": 1000},
+                    {"word": "world", "start_ms": 1000, "end_ms": 1200},
+                ]
+                for ts in timestamps:
+                    assert ts["end_ms"] >= ts["start_ms"]
+
     def test_flush_on_input_done(self, mocker: MockerFixture):
         app, _ = _build_test_app(mocker=mocker)
 
@@ -249,6 +412,7 @@ class TestStreamingSpeechWebSocket:
         speech_service._generate_pcm_chunks = mocker.AsyncMock()
         speech_service.engine_client = mocker.MagicMock()
         speech_service.engine_client.abort = mocker.AsyncMock()
+        speech_service.forced_aligner_config = None
         app, _ = _build_test_app(speech_service)
 
         with TestClient(app) as client:
@@ -269,8 +433,9 @@ class TestStreamingSpeechWebSocket:
         speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req-stream-err", object(), {}))
         speech_service.engine_client = mocker.MagicMock()
         speech_service.engine_client.abort = mocker.AsyncMock()
+        speech_service.forced_aligner_config = None
 
-        async def mock_generate_pcm_chunks(_generator, _request_id):
+        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False):
             yield b"\x01\x02"
             raise RuntimeError("stream boom")
 
@@ -361,8 +526,9 @@ class TestStreamingSpeechWebSocket:
         speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req-abort", object(), {}))
         speech_service.engine_client = mocker.MagicMock()
         speech_service.engine_client.abort = mocker.AsyncMock()
+        speech_service.forced_aligner_config = None
 
-        async def mock_generate_pcm_chunks(_generator, _request_id):
+        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False):
             yield b"\x01\x02"
 
         speech_service._generate_pcm_chunks = mock_generate_pcm_chunks
@@ -385,7 +551,9 @@ class TestStreamingSpeechWebSocket:
         config.ref_audio = None
         config.ref_text = None
         config.x_vector_only_mode = None
+        config.speaker_embedding = None
         config.stream_audio = True
+        config.word_timestamps = False
 
         with pytest.raises(WebSocketDisconnect):
             asyncio.run(handler._generate_and_send(websocket, config, "Hello world.", 0))

tests/entrypoints/openai_api/test_video_server.py

@@ -423,6 +423,7 @@ def test_decode_video_bytes_can_keep_last_frames():
     )
 
     assert len(frames) == 2
+    assert frames.fps == pytest.approx(8.0)
     red_means = [np.asarray(frame)[:, :, 0].mean() for frame in frames]
     assert red_means[0] > 100
     assert red_means[1] > red_means[0]
@@ -520,11 +521,49 @@ def test_seconds_defaults_fps_and_frames(test_client, mocker: MockerFixture):
     engine = test_client.app.state.openai_serving_video._engine_client
     captured = engine.captured_sampling_params_list[0]
     assert captured.num_frames == 72
-    assert captured.fps == 24
-    assert captured.frame_rate == 24.0
+    # fps omitted -> sampling params carry None (the "not provided" signal); the 24
+    # default is applied only at output encoding.
+    assert captured.fps is None
+    assert captured.frame_rate is None
     assert fps_values == [24]
 
 
+def test_model_reported_fps_wins_when_request_fps_omitted(test_client, mocker: MockerFixture):
+    fps_values = []
+
+    def _fake_encode(video, fps, audio=None, audio_sample_rate=None, **kwargs):
+        del video, audio, audio_sample_rate, kwargs
+        fps_values.append(fps)
+        return b"fake-video"
+
+    mocker.patch(
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        side_effect=_fake_encode,
+    )
+
+    engine = test_client.app.state.openai_serving_video._engine_client
+
+    async def _generate(prompt, request_id, sampling_params_list):
+        engine.captured_prompt = prompt
+        engine.captured_sampling_params_list = sampling_params_list
+        result = MockVideoResult([object()])
+        result.multimodal_output["fps"] = 8
+        yield result
+
+    engine.generate = _generate
+
+    response = test_client.post("/v1/videos", data={"prompt": "source fps"})
+
+    assert response.status_code == 200
+    video_id = response.json()["id"]
+    _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
+    captured = engine.captured_sampling_params_list[0]
+    # fps omitted -> None on the sampling params; the model-reported fps (8) wins for output.
+    assert captured.fps is None
+    assert captured.frame_rate is None
+    assert fps_values == [8]
+
+
 def test_size_param_sets_width_height(test_client, mocker: MockerFixture):
     mocker.patch(
         "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",

tests/fullduplex/test_joyvl_brain.py

@@ -0,0 +1,96 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+
+from vllm_omni.experimental.fullduplex.joyvl.decision.output_parser import parse_action
+from vllm_omni.experimental.fullduplex.joyvl.memory.brain import InteractionBrain
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _FakeSummarizer:
+    def __init__(self):
+        self.chunks = 0
+
+    async def summarize_chunk(self, chunk_index, frame_range, frames):
+        self.chunks += 1
+        return f"summary-of-chunk-{chunk_index}"
+
+    async def compress_to_long_term(self, existing, mids):
+        return (existing + " | " if existing else "") + f"compressed({len(mids)})"
+
+
+class _FakeDelegation:
+    def __init__(self):
+        self._ready = {}
+
+    async def submit(self, question, note, frames):
+        self._ready["t1"] = question
+        return "t1"
+
+    async def poll(self, task_id):
+        from vllm_omni.experimental.fullduplex.joyvl.bridges.delegation import DelegationResult
+
+        return DelegationResult(task_id, "ready", digest=f"answer to {self._ready[task_id]}")
+
+
+def test_query_freshness_and_archive_on_change():
+    b = InteractionBrain()
+    assert b.update_query("count bottles") is True
+    assert b.update_query("count bottles") is False
+    b.record_response("1 bottle")
+    assert b.update_query("now describe") is True
+    assert b.memory.qa_history[-1].query == "count bottles"
+    assert b.memory.qa_history[-1].responses == [("0.0s", "1 bottle")]
+
+
+def test_query_in_frame_on_arrival_then_prefix_on_later_chunks():
+    b = InteractionBrain(frame_seconds=1.0)
+    b.tick(3)
+    b.update_query("alert me if a fire breaks out")
+    assert b.query_in_current_chunk is True
+    assert "alert me if a fire breaks out" not in b.build_prefix()
+    b.record_response("a fire is breaking out")
+    b.close_chunk()
+    assert b.query_in_current_chunk is False
+    prefix = b.build_prefix()
+    assert "alert me if a fire breaks out" in prefix
+    assert "a fire is breaking out" in prefix
+
+
+@pytest.mark.asyncio
+async def test_flush_archives_qa_and_summarizes():
+    summ = _FakeSummarizer()
+    b = InteractionBrain(summarizer=summ, chunk_frames=3, long_term_every_n_chunks=2, frame_seconds=1.0)
+    b.update_query("watch the room")
+    b.tick(3)
+    assert b.should_flush() is True
+    b.record_response("a person enters")
+    await b.flush([("0.0s", "u0"), ("2.0s", "u2")])
+    assert summ.chunks == 1
+    assert b.chunk_index == 2
+    assert b.should_flush() is False
+    assert b.memory.qa_history[-1].responses == [("2.0s", "a person enters")]
+    assert len(b.memory.mid_term_summaries) == 1
+
+
+@pytest.mark.asyncio
+async def test_long_term_compression_after_n_chunks():
+    summ = _FakeSummarizer()
+    b = InteractionBrain(summarizer=summ, chunk_frames=1, long_term_every_n_chunks=2, frame_seconds=1.0)
+    for _ in range(2):
+        await b.flush([("0.0s", "u")])
+    assert b.memory.long_term_memory.startswith("compressed(")
+    assert b.memory.mid_term_summaries == []
+
+
+@pytest.mark.asyncio
+async def test_delegation_submit_and_fold():
+    b = InteractionBrain(delegation=_FakeDelegation())
+    action = parse_action("</response> hold on <delegation> solve this")
+    info = await b.submit_delegation(action, [])
+    assert info["status"] == "submitted"
+    folded = await b.fold_delegations()
+    assert folded["status"] == "ready"
+    assert b.memory.qa_history[-1].responses[0][1] == "answer to solve this"

tests/fullduplex/test_joyvl_policy.py

@@ -0,0 +1,282 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import asyncio
+
+import pytest
+
+from vllm_omni.experimental.fullduplex.joyvl.bridges.delegation import (
+    DelegationResult,
+    ImageEditDelegationBridge,
+    ImageGenDelegationBridge,
+    OpenAIDelegationBridge,
+)
+from vllm_omni.experimental.fullduplex.joyvl.decision.policy import JoyVLPolicy, sample_frames
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_sample_frames_keeps_recent():
+    assert sample_frames(["a", "b"], 4) == ["a", "b"]
+    out = sample_frames([str(i) for i in range(10)], 4)
+    assert out[-1] == "9" and len(out) == 4
+
+
+def test_build_messages_stable_head_and_query():
+    p = JoyVLPolicy(num_frames=4)
+    p.tick(2)
+    p.set_query("count the bottles")
+    msgs, user = p.build_messages([{"type": "image_url", "image_url": {"url": "x"}}])
+    assert msgs[0]["role"] == "system" and "</silence>" in msgs[0]["content"]
+
+    assert any("count the bottles" in part.get("text", "") for part in user["content"])
+
+
+def test_commit_records_only_spoken():
+    p = JoyVLPolicy()
+    p.set_query("q")
+    assert p.commit("</silence>").action.value == "silence"
+    assert p.brain.response_records == []
+    assert p.commit("</response> hello").action.value == "response"
+    assert p.brain.response_records == [("0.0s", "hello")]
+
+
+def test_commit_dedup_default_is_exact_match_only():
+    p = JoyVLPolicy()  # threshold 1.0 == reference behavior
+    assert p.commit("</response> a red ball appears").action.value == "response"
+    assert p.commit("</response> a red ball appears").action.value == "silence"  # exact repeat dropped
+    assert p.commit("</response> a red ball is appearing").action.value == "response"  # near-dup kept
+
+
+def test_commit_dedup_threshold_drops_near_duplicates():
+    p = JoyVLPolicy(response_dedup_threshold=0.8)
+    assert p.commit("</response> a red ball appears").action.value == "response"
+    assert p.commit("</response> a red ball appears!").action.value == "silence"  # near-dup dropped
+    assert p.commit("</response> the sky is now blue").action.value == "response"  # distinct kept
+
+
+def test_commit_dedup_rearms_after_silence():
+    p = JoyVLPolicy(response_dedup_threshold=0.85)
+    assert p.commit("</response> someone raises a hand").action.value == "response"
+    # no gap: an immediate repeat is still suppressed
+    assert p.commit("</response> someone raises a hand").action.value == "silence"
+    # a genuine silent tick (hand lowered) re-arms; the recurrence speaks again
+    assert p.commit("</silence>").action.value == "silence"
+    assert p.commit("</response> someone raises a hand").action.value == "response"
+
+
+class _FakeBrain(OpenAIDelegationBridge):
+    """Override only the network call so submit/poll plumbing is exercised for real."""
+
+    def __init__(self, answer, gate=None):
+        super().__init__("http://unused/v1", "background-model")
+        self._answer = answer
+        self._gate = gate
+
+    async def _complete(self, messages):
+        if self._gate is not None:
+            await self._gate.wait()
+        return self._answer
+
+
+@pytest.mark.asyncio
+async def test_openai_delegation_bridge_is_nonblocking_and_resolves():
+    gate = asyncio.Event()
+    bridge = _FakeBrain("the bottle is on the left shelf", gate=gate)
+    frames = [("0.0s", "data:image/jpeg;base64,AAA")]
+    task_id = await bridge.submit("where is the bottle?", "let me check", frames)
+
+    # poll must not block on the slow background brain
+    assert (await bridge.poll(task_id)).status == "pending"
+
+    gate.set()
+    await asyncio.sleep(0)
+    for _ in range(10):
+        result = await bridge.poll(task_id)
+        if result.status != "pending":
+            break
+        await asyncio.sleep(0)
+    assert result.is_ready
+    assert "bottle is on the left shelf" in result.digest
+    assert (await bridge.poll("missing")).status == "error"
+
+
+class _FakeImageBrain(ImageGenDelegationBridge):
+    async def _call_images(self, payload):
+        # echo the prompt back so the test can assert it was passed through
+        return {"data": [{"b64_json": "QUJD"}]}  # "ABC"
+
+
+@pytest.mark.asyncio
+async def test_image_delegation_bridge_returns_media_not_text():
+    bridge = _FakeImageBrain("http://unused/v1")
+    task_id = await bridge.submit("draw a red dragon", "", [])
+    for _ in range(10):
+        result = await bridge.poll(task_id)
+        if result.status != "pending":
+            break
+        await asyncio.sleep(0)
+    assert result.is_ready
+    # the image rides in media as a data URL; digest is only a short text placeholder
+    assert result.media == "data:image/png;base64,QUJD"
+    assert "QUJD" not in result.digest and "red dragon" in result.digest
+
+
+class _FakeEditBrain(ImageEditDelegationBridge):
+    def __init__(self, base_url, model):
+        super().__init__(base_url, model)
+        self.seen_frame = None
+
+    async def _call_chat(self, frame, instruction):
+        self.seen_frame = frame
+        return {"choices": [{"message": {"content": [{"image_url": {"url": "data:image/png;base64,Q0FSVA=="}}]}}]}
+
+
+@pytest.mark.asyncio
+async def test_image_edit_delegation_conditions_on_frame_and_returns_media():
+    bridge = _FakeEditBrain("http://unused/v1", "Qwen-Image-Edit")
+    frames = [("0.0s", "data:image/jpeg;base64,OLD"), ("1.0s", "data:image/jpeg;base64,NOW")]
+    task_id = await bridge.submit("把画面变成卡通风格", "", frames)
+    for _ in range(10):
+        result = await bridge.poll(task_id)
+        if result.status != "pending":
+            break
+        await asyncio.sleep(0)
+    assert result.is_ready
+    assert bridge.seen_frame == "data:image/jpeg;base64,NOW"  # latest frame, not the old one
+    assert result.media == "data:image/png;base64,Q0FSVA=="
+    assert result.media not in result.digest  # image rides in media, not context
+
+
+@pytest.mark.asyncio
+async def test_image_edit_delegation_errors_without_frame():
+    bridge = _FakeEditBrain("http://unused/v1", "Qwen-Image-Edit")
+    task_id = await bridge.submit("cartoonify", "", [])
+    for _ in range(10):
+        result = await bridge.poll(task_id)
+        if result.status != "pending":
+            break
+        await asyncio.sleep(0)
+    assert result.status == "error"
+
+
+class _NamedBridge:
+    def __init__(self, name):
+        self.name = name
+
+    async def submit(self, question, note, frames):
+        return f"{self.name}-1"
+
+    async def poll(self, task_id):
+        return DelegationResult(task_id, "ready", digest=self.name)
+
+
+@pytest.mark.asyncio
+async def test_routing_delegation_dispatches_by_request():
+    from vllm_omni.experimental.fullduplex.joyvl.bridges.delegation import RoutingDelegationBridge
+
+    chat, image, edit = _NamedBridge("chat"), _NamedBridge("image"), _NamedBridge("edit")
+    r = RoutingDelegationBridge(chat=chat, image=image, edit=edit)
+
+    async def route(q):
+        tid = await r.submit(q, "", [])
+        return (await r.poll(tid)).digest
+
+    assert await route("把你看到的画面变成卡通风格") == "edit"  # restyle current view
+    assert await route("画一只红色的龙") == "image"  # generate a new image
+    assert await route("这道数学题的答案是什么") == "chat"  # hard question -> chat brain
+
+
+class _Delegation:
+    async def submit(self, question, note, frames):
+        return "t1"
+
+    async def poll(self, task_id):
+        return DelegationResult(task_id, "ready", digest="background answer")
+
+
+def test_build_delegation_off_unless_backend_configured():
+    from vllm_omni.experimental.fullduplex.joyvl.serving.config import InteractionConfig
+    from vllm_omni.experimental.fullduplex.joyvl.serving.server import SessionManager
+
+    # No backend URL -> delegation disabled (None), NOT a silent stub that folds fake answers.
+    assert SessionManager._build_delegation(InteractionConfig()) is None
+    assert SessionManager._build_delegation(InteractionConfig(delegation_kind="router")) is None
+    # Stub only on explicit opt-in.
+    assert SessionManager._build_delegation(InteractionConfig(delegation_kind="stub")) is not None
+    # A real chat backend -> a live bridge.
+    cfg = InteractionConfig(delegation_backend_url="http://x/v1", delegation_model="m")
+    assert SessionManager._build_delegation(cfg) is not None
+
+
+@pytest.mark.asyncio
+async def test_delegation_cancel_drops_task():
+    from vllm_omni.experimental.fullduplex.joyvl.bridges.delegation import StubDelegationBridge
+
+    b = StubDelegationBridge(ready_after_ticks=1)
+    task_id = await b.submit("q", "", [])
+    b.cancel(task_id)
+    # after cancel the task is gone -> poll reports it as unknown rather than leaking
+    assert (await b.poll(task_id)).status == "error"
+
+
+@pytest.mark.asyncio
+async def test_delegation_aclose_clears_pending():
+    from vllm_omni.experimental.fullduplex.joyvl.bridges.delegation import StubDelegationBridge
+
+    b = StubDelegationBridge(ready_after_ticks=5)
+    task_id = await b.submit("q", "", [])
+    await b.aclose()
+    assert (await b.poll(task_id)).status == "error"
+
+
+@pytest.mark.asyncio
+async def test_session_reset_cancels_and_awaits_consolidation():
+    from vllm_omni.experimental.fullduplex.joyvl.serving.config import InteractionConfig
+    from vllm_omni.experimental.fullduplex.joyvl.serving.session import InteractionSession
+
+    class _Backend:
+        async def generate(self, *a, **k):
+            return "</silence>", None
+
+        async def aclose(self):
+            pass
+
+    class _Summarizer:
+        async def summarize_chunk(self, *a, **k):
+            await asyncio.sleep(100)  # hang so the consolidation task is in-flight
+            return "x"
+
+        async def compress_to_long_term(self, *a, **k):
+            return ""
+
+        async def aclose(self):
+            pass
+
+    sess = InteractionSession("s", InteractionConfig(), _Backend(), summarizer=_Summarizer())
+    sess._spawn_consolidation(1, [("0.0s", "u")])
+    assert len(sess._consolidating) == 1
+    await sess.reset()  # must cancel AND await the in-flight task, not just clear the set
+    assert len(sess._consolidating) == 0
+
+
+@pytest.mark.asyncio
+async def test_manager_aclose_closes_without_error():
+    from vllm_omni.experimental.fullduplex.joyvl.serving.config import InteractionConfig
+    from vllm_omni.experimental.fullduplex.joyvl.serving.server import SessionManager
+
+    mgr = SessionManager(InteractionConfig(enable_memory=False, enable_delegation=False))
+    await mgr.aclose()
+    await mgr.aclose()  # idempotent
+
+
+@pytest.mark.asyncio
+async def test_delegation_submit_and_fold():
+    p = JoyVLPolicy(delegation=_Delegation())
+    action = p.commit("</response> hold on <delegation> hard question")
+    assert action.action.value == "delegate"
+    info = await p.submit_if_delegate(action)
+    assert info["status"] == "submitted"
+    folded = await p.fold_delegations()
+    assert folded["status"] == "ready"
+    assert p.brain.memory.qa_history[-1].responses[0][1] == "background answer"

tests/fullduplex/test_runtime.py

@@ -0,0 +1,172 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import asyncio
+from collections.abc import AsyncIterator
+
+import pytest
+
+from vllm_omni.experimental.fullduplex.core import protocol as ev
+from vllm_omni.experimental.fullduplex.core.adapter import DuplexCapability, OutputChunk
+from vllm_omni.experimental.fullduplex.core.runtime import DuplexRuntime
+from vllm_omni.experimental.fullduplex.core.session import DuplexSession, DuplexSessionConfig
+from vllm_omni.experimental.fullduplex.joyvl.adapter import JoyVLDuplexAdapter
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+async def _feed(events):
+    for e in events:
+        yield e
+
+
+def _collector():
+    out: list[dict] = []
+
+    async def emit(event: dict) -> None:
+        out.append(event)
+
+    return out, emit
+
+
+class _FakeAdapter:
+    def __init__(self, chunks, barge_after=None):
+        self._chunks = chunks
+        self._barge_after = barge_after
+
+    def capabilities(self):
+        return DuplexCapability(frozenset({"text"}), frozenset({"text"}), proactive=False)
+
+    async def on_input(self, session, modality, data):
+        pass
+
+    def should_respond(self, session):
+        return True
+
+    async def respond(self, session) -> AsyncIterator[OutputChunk]:
+        for i, c in enumerate(self._chunks):
+            if self._barge_after is not None and i == self._barge_after:
+                session.barge_in()
+            yield OutputChunk("text", c)
+
+    async def on_barge_in(self, session):
+        pass
+
+    async def on_playback_ack(self, session, cursor):
+        pass
+
+
+@pytest.mark.asyncio
+async def test_runtime_basic_response():
+    session = DuplexSession("s", DuplexSessionConfig(output_modalities=("text",)))
+    rt = DuplexRuntime(session, _FakeAdapter(["a", "b"]))
+    out, emit = _collector()
+    await rt.run(_feed([{"type": ev.INPUT_COMMIT}, {"type": ev.CLOSE}]), emit)
+    types = [e["type"] for e in out]
+    assert types == [ev.RESPONSE_CREATED, ev.RESPONSE_DELTA, ev.RESPONSE_DELTA, ev.RESPONSE_DONE]
+    assert [e["data"] for e in out if e["type"] == ev.RESPONSE_DELTA] == ["a", "b"]
+
+
+@pytest.mark.asyncio
+async def test_runtime_barge_in_drops_stale_output():
+    session = DuplexSession("s")
+    rt = DuplexRuntime(session, _FakeAdapter(["a", "b", "c"], barge_after=1))
+    out, emit = _collector()
+    await rt.run(_feed([{"type": ev.INPUT_COMMIT}, {"type": ev.CLOSE}]), emit)
+    data = [e["data"] for e in out if e["type"] == ev.RESPONSE_DELTA]
+    assert data == ["a"]
+    assert ev.RESPONSE_DONE not in [e["type"] for e in out]
+
+
+class _SlowAdapter:
+    def capabilities(self):
+        return DuplexCapability(frozenset({"text"}), frozenset({"text"}), proactive=False)
+
+    async def on_input(self, session, modality, data):
+        pass
+
+    def should_respond(self, session):
+        return True
+
+    async def respond(self, session) -> AsyncIterator[OutputChunk]:
+        for c in ["a", "b", "c", "d", "e"]:
+            await asyncio.sleep(0.02)
+            yield OutputChunk("text", c)
+
+    async def on_barge_in(self, session):
+        pass
+
+    async def on_playback_ack(self, session, cursor):
+        pass
+
+
+@pytest.mark.asyncio
+async def test_runtime_cancel_event_interrupts_active_response():
+    session = DuplexSession("s")
+    rt = DuplexRuntime(session, _SlowAdapter())
+    out, emit = _collector()
+
+    async def feed():
+        yield {"type": ev.INPUT_COMMIT}
+        await asyncio.sleep(0.03)
+        yield {"type": ev.RESPONSE_CANCEL}
+        yield {"type": ev.CLOSE}
+
+    await rt.run(feed(), emit)
+    deltas = [e["data"] for e in out if e["type"] == ev.RESPONSE_DELTA]
+    types = [e["type"] for e in out]
+    assert len(deltas) < 5
+    assert ev.RESPONSE_CANCELLED in types
+    assert ev.RESPONSE_DONE not in types
+
+
+@pytest.mark.asyncio
+async def test_runtime_new_response_supersedes_inflight_without_blocking():
+    # A second response trigger arriving while the first is still streaming must
+    # supersede it (barge-in) rather than block the input loop until it finishes.
+    session = DuplexSession("s")
+    rt = DuplexRuntime(session, _SlowAdapter())
+    out, emit = _collector()
+
+    async def feed():
+        yield {"type": ev.INPUT_COMMIT}
+        await asyncio.sleep(0.03)  # let the first response emit ~one chunk
+        yield {"type": ev.INPUT_COMMIT}  # supersedes the in-flight first response
+        await asyncio.sleep(0.2)  # let the second response finish
+        yield {"type": ev.CLOSE}
+
+    await rt.run(feed(), emit)
+    created = [e for e in out if e["type"] == ev.RESPONSE_CREATED]
+    done = [e for e in out if e["type"] == ev.RESPONSE_DONE]
+    # two responses created; only the latest one completes (the first was cancelled)
+    assert len(created) == 2
+    assert len(done) == 1
+    assert done[0]["response_index"] == created[-1]["response_index"]
+
+
+@pytest.mark.asyncio
+async def test_joyvl_adapter_speaks_then_silences():
+    replies = iter(["</response> a fire is breaking out", "</silence>"])
+
+    async def fake_generate(messages):
+        return next(replies)
+
+    adapter = JoyVLDuplexAdapter(fake_generate, num_frames=4)
+    cfg = DuplexSessionConfig(input_modalities=("video", "text"), output_modalities=("text",), proactive=True)
+    rt = DuplexRuntime(DuplexSession("s", cfg), adapter)
+    out, emit = _collector()
+    await rt.run(
+        _feed(
+            [
+                {"type": ev.INPUT_APPEND, "modality": "text", "data": "alert me if a fire breaks out"},
+                {"type": ev.INPUT_APPEND, "modality": "video", "data": "data:image/jpeg;base64,AAA"},
+                {"type": ev.INPUT_APPEND, "modality": "video", "data": "data:image/jpeg;base64,BBB"},
+                {"type": ev.CLOSE},
+            ]
+        ),
+        emit,
+    )
+    deltas = [e["data"] for e in out if e["type"] == ev.RESPONSE_DELTA]
+    assert "a fire is breaking out" in deltas
+
+    assert deltas.count("a fire is breaking out") == 1

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"

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)]) == []

tests/model_extras/test_model_extras.py

@@ -57,6 +57,53 @@ def test_sensenova_extra_registry_declares_request_and_response_params() -> None
     assert should_init_extra_args_for_non_diffusion_stages("SenseNovaU1Pipeline") is False
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_magi_human_extra_registry_declares_request_and_response_params() -> None:
+    assert get_extra_body_params("MagiHumanPipeline") == frozenset(
+        {
+            "seconds",
+            "audio_path",
+            "image_path",
+            "sr_height",
+            "sr_width",
+            "sr_num_inference_steps",
+        }
+    )
+    assert get_extra_output_params("MagiHumanPipeline") == frozenset()
+    assert should_init_extra_args_for_non_diffusion_stages("MagiHumanPipeline") is False
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_helios_extra_registry_declares_request_and_response_params() -> None:
+    expected_body = frozenset(
+        {
+            "is_enable_stage2",
+            "pyramid_num_stages",
+            "pyramid_num_inference_steps_list",
+            "is_amplify_first_chunk",
+            "is_skip_first_chunk",
+            "use_cfg_zero_star",
+            "use_zero_init",
+            "zero_steps",
+            "image",
+            "video",
+            "add_noise_to_image_latents",
+            "image_noise_sigma_min",
+            "image_noise_sigma_max",
+            "add_noise_to_video_latents",
+            "video_noise_sigma_min",
+            "video_noise_sigma_max",
+        }
+    )
+    # Both the base and pyramid class names resolve to the same declaration.
+    for cls in ("HeliosPipeline", "HeliosPyramidPipeline"):
+        assert get_extra_body_params(cls) == expected_body
+        assert get_extra_output_params(cls) == frozenset()
+        assert should_init_extra_args_for_non_diffusion_stages(cls) is False
+
+
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_unknown_pipeline_has_empty_extra_registry() -> None:

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

tests/utils/test_forced_aligner.py

@@ -0,0 +1,165 @@
+import numpy as np
+import pytest
+
+from vllm_omni.utils import forced_aligner
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_decode_timestamps_maps_boundary_bins_to_words():
+    logits = np.zeros((4, 5), dtype=np.float32)
+    logits[0, 0] = 1.0
+    logits[1, 2] = 1.0
+    logits[2, 2] = 1.0
+    logits[3, 4] = 1.0
+
+    timestamps = forced_aligner._decode_timestamps(
+        logits=logits,
+        words=["hello", "world"],
+        timestamp_positions=[0, 1, 2, 3],
+        classify_num=5,
+        audio_duration_ms=1000,
+    )
+
+    assert timestamps == [
+        forced_aligner.WordTimestamp("hello", 0, 400),
+        forced_aligner.WordTimestamp("world", 400, 800),
+    ]
+
+
+def test_decode_timestamps_repairs_non_monotonic_bins():
+    # Word 1's end bin (1) dips below its start bin (2); fix_timestamp should
+    # snap it back so the frame never carries end_ms < start_ms.
+    logits = np.zeros((4, 5), dtype=np.float32)
+    logits[0, 0] = 1.0  # word0 start -> bin 0
+    logits[1, 2] = 1.0  # word0 end   -> bin 2
+    logits[2, 1] = 1.0  # word1 start -> bin 1 (out of order)
+    logits[3, 3] = 1.0  # word1 end   -> bin 3
+
+    timestamps = forced_aligner._decode_timestamps(
+        logits=logits,
+        words=["hello", "world"],
+        timestamp_positions=[0, 1, 2, 3],
+        classify_num=5,
+        timestamp_segment_time_ms=200,
+        audio_duration_ms=1000,
+    )
+
+    for ts in timestamps:
+        assert ts.end_ms >= ts.start_ms
+    # bins [0, 2, 1, 3] -> repaired [0, 2, 2, 3] -> ms x200
+    assert timestamps == [
+        forced_aligner.WordTimestamp("hello", 0, 400),
+        forced_aligner.WordTimestamp("world", 400, 600),
+    ]
+
+
+def test_decode_timestamps_clamps_to_audio_duration_and_stays_ordered():
+    # Bins run past the 1000 ms audio; the decoder must cap each bound at the
+    # audio length and keep words non-overlapping (start <= end, monotonic).
+    logits = np.zeros((4, 5), dtype=np.float32)
+    logits[0, 0] = 1.0  # word0 start -> bin 0
+    logits[1, 3] = 1.0  # word0 end   -> bin 3 (1200 ms, past audio)
+    logits[2, 3] = 1.0  # word1 start -> bin 3 (1200 ms, past audio)
+    logits[3, 4] = 1.0  # word1 end   -> bin 4 (1600 ms, past audio)
+
+    timestamps = forced_aligner._decode_timestamps(
+        logits=logits,
+        words=["a", "b"],
+        timestamp_positions=[0, 1, 2, 3],
+        classify_num=5,
+        timestamp_segment_time_ms=400,
+        audio_duration_ms=1000,
+    )
+
+    assert timestamps == [
+        forced_aligner.WordTimestamp("a", 0, 1000),
+        forced_aligner.WordTimestamp("b", 1000, 1000),
+    ]
+
+
+def test_decode_timestamps_rejects_marker_count_mismatch():
+    logits = np.zeros((2, 5), dtype=np.float32)
+
+    timestamps = forced_aligner._decode_timestamps(
+        logits=logits,
+        words=["hello", "world"],
+        timestamp_positions=[0, 1],
+        classify_num=5,
+        audio_duration_ms=1000,
+    )
+
+    assert timestamps == []
+
+
+def test_build_config_from_yaml(tmp_path):
+    cfg = tmp_path / "forced_aligner.yaml"
+    cfg.write_text(
+        """
+forced_aligner:
+  model: Qwen/Qwen3-ForcedAligner-0.6B
+  gpu_memory_utilization: 0.42
+  dtype: float16
+  max_model_len: 2048
+  trust_remote_code: false
+""",
+        encoding="utf-8",
+    )
+    args = type("Args", (), {"forced_aligner": None, "forced_aligner_config": str(cfg)})()
+
+    out = forced_aligner.build_forced_aligner_config(args)
+
+    assert out == forced_aligner.ForcedAlignerConfig(
+        model="Qwen/Qwen3-ForcedAligner-0.6B",
+        runner="pooling",
+        architecture="Qwen3ASRForcedAlignerForTokenClassification",
+        pooling_task="token_classify",
+        gpu_memory_utilization=0.42,
+        dtype="float16",
+        max_model_len=2048,
+        trust_remote_code=False,
+    )
+
+
+def test_build_config_cli_model_overrides_yaml(tmp_path):
+    cfg = tmp_path / "forced_aligner.yaml"
+    cfg.write_text(
+        "forced_aligner:\n  model: old\n  gpu_memory_utilization: 0.2\n  dtype: float16\n",
+        encoding="utf-8",
+    )
+    args = type(
+        "Args",
+        (),
+        {
+            "forced_aligner": "new",
+            "forced_aligner_config": str(cfg),
+        },
+    )()
+
+    out = forced_aligner.build_forced_aligner_config(args)
+
+    assert out is not None
+    # --forced-aligner overrides the YAML model; gpu_memory_utilization/dtype
+    # come from the user YAML (there is no longer a CLI flag for gpu mem).
+    assert out.model == "new"
+    assert out.gpu_memory_utilization == 0.2
+    assert out.dtype == "float16"
+    assert out.runner == "pooling"
+
+
+def test_build_config_from_cli_model_uses_default_yaml():
+    args = type(
+        "Args",
+        (),
+        {
+            "forced_aligner": "local-aligner",
+            "forced_aligner_config": None,
+        },
+    )()
+
+    out = forced_aligner.build_forced_aligner_config(args)
+
+    assert out is not None
+    assert out.model == "local-aligner"
+    assert out.runner == "pooling"
+    assert out.architecture == "Qwen3ASRForcedAlignerForTokenClassification"

tests/utils/test_qwen3_force_align_processor.py

@@ -0,0 +1,65 @@
+import pytest
+
+from vllm_omni.utils import qwen3_force_align_processor as processor
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_build_prompt_has_boundary_timestamp_markers():
+    prompt = processor.build_prompt(["hello", "world"])
+
+    assert prompt.count("<timestamp>") == 4
+    assert "hello<timestamp><timestamp>world" in prompt
+    assert prompt.endswith("<|im_start|>assistant\n")
+
+
+@pytest.mark.parametrize(
+    "text,expected",
+    [
+        # Punctuation inside a whitespace segment is stripped, not split on:
+        # the hand-rolled splitter used to break these into many fake words.
+        ("U.S.A", ["USA"]),
+        ("hello, world!", ["hello", "world"]),
+        ("don't stop", ["don't", "stop"]),
+        # CJK characters peel off as individual tokens, Latin runs stay whole.
+        ("你好world", ["你", "好", "world"]),
+        ("我爱 China", ["我", "爱", "China"]),
+        ("   spaced   out   ", ["spaced", "out"]),
+    ],
+)
+def test_tokenize_space_lang_matches_official_segmentation(text, expected):
+    assert processor._tokenize_space_lang(text) == expected
+
+
+def test_segment_words_falls_back_to_port_without_qwen_asr(monkeypatch):
+    # With qwen_asr unavailable, segmentation must use the built-in port and
+    # still produce the faithful result for the common (non-JP/KO) case.
+    monkeypatch.setattr(processor, "_get_official_processor", lambda: None)
+
+    assert processor.segment_words("U.S.A is here", "auto") == ["USA", "is", "here"]
+
+
+def test_fix_timestamp_repairs_dip_onto_monotonic_sequence():
+    # 200 dips below the preceding 400; it snaps back up to 400.
+    assert processor.fix_timestamp([0, 400, 200, 800]) == [0, 400, 400, 800]
+
+
+def test_fix_timestamp_passes_through_monotonic():
+    assert processor.fix_timestamp([0, 100, 100, 250]) == [0, 100, 100, 250]
+
+
+def test_resolve_timestamp_token_id_defaults_to_marker_token():
+    # Regression: the default must resolve the same <timestamp> marker that
+    # build_prompt inserts, not None (which would degrade every request).
+    seen = {}
+
+    class FakeTokenizer:
+        def convert_tokens_to_ids(self, token):
+            seen["token"] = token
+            return 151705
+
+    tid = processor.resolve_timestamp_token_id(FakeTokenizer())
+
+    assert tid == 151705
+    assert seen["token"] == processor.TIMESTAMP_TOKEN
+    assert processor.TIMESTAMP_TOKEN in processor.build_prompt(["hello", "world"])

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)

vllm_omni/attention/fish_kvcache_backend.py

@@ -233,7 +233,12 @@ def _forward_with_fish_kvcache(
 
     if attn_metadata is not None and not attn_metadata.use_cascade:
         num_actual_tokens = attn_metadata.num_actual_tokens
-        key_cache, value_cache = kv_cache.unbind(0)
+        # vLLM >=0.23.0 lays out the KV cache as
+        # (num_blocks, 2, block_size, num_kv_heads, head_size), so key/value
+        # live on dim 1 (matches FlashAttentionImpl's kv_cache.unbind(1)).
+        # The pre-0.23 layout had them on dim 0; unbind(0) here would split
+        # along num_blocks and raise "too many values to unpack".
+        key_cache, value_cache = kv_cache.unbind(1)
         q = query[:num_actual_tokens]
         out = output[:num_actual_tokens]
         active_batch = int(q.shape[0])

vllm_omni/benchmarks/data_modules/daily_omni_text_audio.py

@@ -1,255 +0,0 @@
-"""Daily-Omni: optional consistency check between text stream and generated speech.
-
-The benchmark MCQ accuracy uses ``generated_text`` only. When the omni server also
-streams ``modality=audio`` (TTS), this module can transcribe the concatenated WAV
-with Whisper and compare the inferred option letter to the one parsed from text.
-
-Requires ``openai-whisper`` (``pip install openai-whisper``). Enable via env
-``DAILY_OMNI_TEXT_AUDIO_CONSISTENCY=1`` or CLI ``--daily-omni-text-audio-consistency``.
-
-Whisper model name defaults to ``tiny`` (override with ``DAILY_OMNI_WHISPER_MODEL``).
-"""
-
-from __future__ import annotations
-
-import logging
-import os
-import re
-import threading
-from typing import Any
-
-from vllm_omni.benchmarks.data_modules.daily_omni_dataset import DailyOmniSampleRequest
-from vllm_omni.benchmarks.data_modules.daily_omni_eval import extract_predicted_choice
-
-logger = logging.getLogger(__name__)
-
-_whisper_model = None
-_whisper_model_name: str | None = None
-_whisper_lock = threading.Lock()
-
-
-def env_text_audio_check_enabled() -> bool:
-    return os.environ.get("DAILY_OMNI_TEXT_AUDIO_CONSISTENCY", "").lower() in (
-        "1",
-        "true",
-        "yes",
-    )
-
-
-def extract_choice_from_asr_transcript(transcript: str) -> str | None:
-    """Parse A–D from ASR text; extends :func:`extract_predicted_choice` with spoken Chinese phrases."""
-    c = extract_predicted_choice(transcript)
-    if c:
-        return c
-    t = transcript or ""
-    for pat in (
-        r"(?i)选项\s*([ABCD])\b",
-        r"(?i)选\s*([ABCD])\b",
-        r"(?i)答案\s*是\s*([ABCD])\b",
-        r"(?i)答案\s*([ABCD])\b",
-    ):
-        m = re.search(pat, t)
-        if m:
-            return m.group(1).upper()
-    return None
-
-
-def _get_whisper_model(model_name: str):
-    global _whisper_model, _whisper_model_name
-    with _whisper_lock:
-        if _whisper_model is None or _whisper_model_name != model_name:
-            import whisper
-
-            logger.warning(
-                "Loading Whisper model %r for Daily-Omni text/audio consistency (one-time)...",
-                model_name,
-            )
-            _whisper_model = whisper.load_model(model_name)
-            _whisper_model_name = model_name
-        return _whisper_model
-
-
-def transcribe_wav_bytes(
-    wav_bytes: bytes,
-    *,
-    language: str | None = None,
-    model_name: str | None = None,
-) -> tuple[str | None, str | None]:
-    """Transcribe WAV bytes. Returns ``(transcript, error)`` — one of them is set.
-
-    Args:
-        wav_bytes: RIFF WAV file bytes.
-        language: Optional Whisper language code (e.g. ``en``, ``zh``); improves accuracy/latency.
-        model_name: Override model id; else ``DAILY_OMNI_WHISPER_MODEL`` or ``tiny``.
-    """
-    if not wav_bytes:
-        return None, "empty_wav"
-    if model_name is None or not str(model_name).strip():
-        model_name = os.environ.get("DAILY_OMNI_WHISPER_MODEL") or "tiny"
-    model_name = str(model_name).strip() or "tiny"
-    path: str | None = None
-    try:
-        import tempfile
-
-        model = _get_whisper_model(model_name)
-        fd, path = tempfile.mkstemp(suffix=".wav")
-        with os.fdopen(fd, "wb") as fp:
-            fp.write(wav_bytes)
-        kwargs: dict = {}
-        if language:
-            kwargs["language"] = language
-        result = model.transcribe(path, **kwargs)
-        text = (result.get("text") or "").strip()
-        return (text if text else None), None
-    except ImportError:
-        return None, "openai-whisper is not installed (pip install openai-whisper)"
-    except Exception as e:
-        return None, str(e)[:500]
-    finally:
-        if path:
-            try:
-                os.unlink(path)
-            except OSError:
-                pass
-
-
-def compute_daily_omni_text_audio_consistency_metrics(
-    input_requests: list[Any],
-    outputs: list[Any],
-    *,
-    include_per_item: bool = False,
-) -> dict[str, Any] | None:
-    """Compare option letter from ``generated_text`` vs Whisper transcript of output audio.
-
-    Only considers requests where ``outputs[i]`` has ``generated_audio_wav_bytes`` set
-    (populated by the omni benchmark when TA check is enabled).
-    """
-    if not input_requests or len(input_requests) != len(outputs):
-        return None
-    if not all(isinstance(r, DailyOmniSampleRequest) for r in input_requests):
-        return None
-
-    ta_no_wav = 0
-    ta_asr_failed = 0
-    ta_text_unparsed = 0
-    ta_audio_unparsed = 0
-    ta_consistent = 0
-    ta_mismatch = 0
-    ta_both_parsed = 0
-    items: list[dict[str, Any]] = []
-
-    for req, out in zip(input_requests, outputs, strict=True):
-        assert isinstance(req, DailyOmniSampleRequest)
-        rid = req.request_id
-        if not getattr(out, "success", False):
-            if include_per_item:
-                items.append(
-                    {
-                        "request_id": rid,
-                        "skipped": True,
-                        "reason": "request_not_success",
-                    }
-                )
-            continue
-
-        wav = getattr(out, "generated_audio_wav_bytes", None)
-        if not wav:
-            ta_no_wav += 1
-            if include_per_item:
-                items.append(
-                    {
-                        "request_id": rid,
-                        "skipped": False,
-                        "reason": "no_output_audio",
-                        "text_choice": extract_predicted_choice(getattr(out, "generated_text", "") or ""),
-                    }
-                )
-            continue
-
-        transcript, asr_err = transcribe_wav_bytes(wav)
-        if asr_err:
-            ta_asr_failed += 1
-            if include_per_item:
-                items.append(
-                    {
-                        "request_id": rid,
-                        "asr_error": asr_err,
-                        "text_choice": extract_predicted_choice(getattr(out, "generated_text", "") or ""),
-                    }
-                )
-            continue
-
-        text_choice = extract_predicted_choice(getattr(out, "generated_text", "") or "")
-        audio_choice = extract_choice_from_asr_transcript(transcript or "")
-
-        if text_choice is None:
-            ta_text_unparsed += 1
-        if audio_choice is None:
-            ta_audio_unparsed += 1
-
-        if text_choice is not None and audio_choice is not None:
-            ta_both_parsed += 1
-            if text_choice == audio_choice:
-                ta_consistent += 1
-            else:
-                ta_mismatch += 1
-
-        if include_per_item:
-            consistent: bool | None
-            if text_choice is None or audio_choice is None:
-                consistent = None
-            else:
-                consistent = text_choice == audio_choice
-            items.append(
-                {
-                    "request_id": rid,
-                    "text_choice": text_choice,
-                    "audio_choice": audio_choice,
-                    "asr_transcript": (transcript or "")[:500],
-                    "text_audio_consistent": consistent,
-                }
-            )
-
-    comparable = ta_consistent + ta_mismatch
-    rate = (ta_consistent / comparable) if comparable else None
-
-    out: dict[str, Any] = {
-        "daily_omni_ta_enabled": True,
-        "daily_omni_ta_no_output_audio": ta_no_wav,
-        "daily_omni_ta_asr_failed": ta_asr_failed,
-        "daily_omni_ta_text_unparsed": ta_text_unparsed,
-        "daily_omni_ta_audio_unparsed": ta_audio_unparsed,
-        "daily_omni_ta_both_parsed": ta_both_parsed,
-        "daily_omni_ta_consistent": ta_consistent,
-        "daily_omni_ta_mismatch": ta_mismatch,
-        "daily_omni_ta_consistency_rate": rate,
-    }
-    if include_per_item:
-        out["daily_omni_ta_items"] = items
-    return out
-
-
-def print_daily_omni_text_audio_summary(metrics: dict[str, Any]) -> None:
-    if not metrics.get("daily_omni_ta_enabled"):
-        return
-    print("{s:{c}^{n}}".format(s=" Daily-Omni text vs audio (ASR) ", n=50, c="="))
-    print("{:<40} {:<10}".format("No output audio captured:", metrics.get("daily_omni_ta_no_output_audio", 0)))
-    print("{:<40} {:<10}".format("ASR failed:", metrics.get("daily_omni_ta_asr_failed", 0)))
-    print("{:<40} {:<10}".format("Both text+audio letter parsed:", metrics.get("daily_omni_ta_both_parsed", 0)))
-    print("{:<40} {:<10}".format("Consistent (same letter):", metrics.get("daily_omni_ta_consistent", 0)))
-    print("{:<40} {:<10}".format("Mismatch:", metrics.get("daily_omni_ta_mismatch", 0)))
-    r = metrics.get("daily_omni_ta_consistency_rate")
-    if r is not None:
-        print("{:<40} {:<10.4f}".format("Consistency rate (of both parsed):", r))
-    print(
-        "{:<40} {:<10}".format(
-            "Text unparsed (among w/ audio):",
-            metrics.get("daily_omni_ta_text_unparsed", 0),
-        )
-    )
-    print(
-        "{:<40} {:<10}".format(
-            "Audio unparsed (among w/ audio):",
-            metrics.get("daily_omni_ta_audio_unparsed", 0),
-        )
-    )

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",

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