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

@@ -32,8 +32,6 @@ python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
 
 - By default, image tasks talk to `http://<host>:<port>/v1/chat/completions`; video tasks talk to `/v1/videos`.
 - If you run the server on another host or port, pass `--base-url` accordingly.
-- To compare diffusion multi-replica serving, restart the server with different
-  `runtime.num_replicas` values and keep the benchmark parameters unchanged.
 
 ## 2. Supported Datasets
 
@@ -95,35 +93,6 @@ You can point to your own trace using `--dataset-path`.
 
 ## 3. Benchmark Parameters
 
-### Multi-replica sweep
-
-For replica scaling, use the same model, prompt count, image/video size, and
-step count for every run. Change only the serving stage config:
-
-```yaml
-runtime:
-  devices: "0,1,2,3"
-  num_replicas: 4
-```
-
-Then run a fixed benchmark and write one JSON file per replica count:
-
-```bash
-python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
-  --base-url http://localhost:8099 \
-  --model Qwen/Qwen-Image \
-  --task t2i \
-  --dataset random \
-  --num-prompts 32 \
-  --max-concurrency 8 \
-  --request-rate inf \
-  --num-inference-steps 10 \
-  --width 512 --height 512 \
-  --output-file diffusion_replica_4.json
-```
-
-Repeat for `num_replicas=1`, `2`, `3`, and `4`, using matching `runtime.devices`.
-
 ### Basic flags
 
 - `--base-url`: Server address; `--endpoint` selects the path appended to this base URL.

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/diffusion/multi_replica_serving.md

@@ -1,80 +0,0 @@
-# Diffusion Multi-Replica Serving
-
-Diffusion stages can be replicated in online serving by setting `runtime.num_replicas` on a `stage_type: diffusion` stage. Each replica owns an independent diffusion engine process and receives requests through round-robin routing with request affinity.
-
-## When to Use
-
-Use multi-replica serving when one diffusion process does not saturate the available GPUs, or when you want higher request throughput for many independent image/video generation requests. It improves concurrency throughput; it does not make a single request faster.
-
-## Configuration
-
-A replica consumes the devices required by that diffusion stage. For single-device diffusion pipelines, set `runtime.devices` to one device per replica:
-
-```yaml
-stages:
-  - stage_id: 0
-    stage_type: diffusion
-    runtime:
-      devices: "0,1,2,3"
-      num_replicas: 4
-    engine_args:
-      model: Qwen/Qwen-Image
-      parallel_config:
-        world_size: 1
-```
-
-For multi-device diffusion pipelines, `parallel_config.world_size` is treated as the number of devices per replica. For example, four 2-GPU replicas need eight devices:
-
-```yaml
-stages:
-  - stage_id: 0
-    stage_type: diffusion
-    runtime:
-      devices: "0,1,2,3,4,5,6,7"
-      num_replicas: 4
-    engine_args:
-      model: Wan-AI/Wan2.2-T2V-A14B-Diffusers
-      parallel_config:
-        world_size: 2
-```
-
-The runtime splits the device list into contiguous chunks per replica. With the second example, replica 0 uses `0,1`, replica 1 uses `2,3`, replica 2 uses `4,5`, and replica 3 uses `6,7`.
-
-## Serving
-
-Start the server with the stage config:
-
-```bash
-vllm serve Qwen/Qwen-Image \
-  --omni \
-  --port 8099 \
-  --stage-configs-path /path/to/diffusion_multi_replica.yaml
-```
-
-Then send normal OpenAI-compatible diffusion chat requests. The orchestrator routes new requests across replicas and keeps updates/abort calls on the same replica for each request.
-
-## Benchmarking
-
-Benchmark each replica count by changing `runtime.num_replicas` and `runtime.devices`, restarting the server, then running:
-
-```bash
-python benchmarks/diffusion/diffusion_benchmark_serving.py \
-  --base-url http://localhost:8099 \
-  --model Qwen/Qwen-Image \
-  --task t2i \
-  --dataset random \
-  --num-prompts 8 \
-  --request-rate inf \
-  --num-inference-steps 2 \
-  --width 256 --height 256 \
-  --output-file diffusion_replica_4.json
-```
-
-For a quick scaling sweep, test `num_replicas` values `1`, `2`, `3`, and `4` with the same prompt count and generation parameters. Compare throughput and end-to-end latency from the JSON outputs.
-
-## Notes and Limitations
-
-- Multi-replica serving is currently supported for diffusion stages. Single-stage mode still rejects non-diffusion replicas.
-- `runtime.devices` must provide exactly `num_replicas * devices_per_replica` devices.
-- Request-level latency may stay similar or slightly increase because each request still runs on one replica.
-- Throughput improves only when enough concurrent requests are issued and each replica has enough GPU memory.

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/qwen3_omni/end2end_async_chunk.py

@@ -257,11 +257,8 @@ async def run_single_request(
                         new_chunks = audio_data[audio_list_consumed:]
                         audio_list_consumed = len(audio_data)
                     elif isinstance(audio_data, torch.Tensor):
-                        if audio_data is not audio_last_tensor:
-                            new_chunks = [audio_data]
-                            audio_last_tensor = audio_data
-                        else:
-                            new_chunks = []
+                        new_chunks = [audio_data]
+                        audio_last_tensor = audio_data
                     else:
                         new_chunks = []
 

examples/offline_inference/qwen3_omni/run_tests.sh

@@ -1,140 +0,0 @@
-#!/usr/bin/env bash
-# =============================================================================
-# Qwen3-Omni end2end functional test script
-# Uses default built-in assets, iterates over query-type x output-modalities combinations
-# =============================================================================
-
-set -euo pipefail
-
-export VLLM_OMNI_USE_V2_RUNNER=1
-
-MODEL_PATH="/models/Qwen3-Omni-30B-A3B-Instruct"
-SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
-END2END="$SCRIPT_DIR/end2end.py"
-LOG_DIR="$SCRIPT_DIR/test_logs"
-SUMMARY_FILE="$LOG_DIR/summary.log"
-
-rm -rf "$LOG_DIR"
-mkdir -p "$LOG_DIR"
-> "$SUMMARY_FILE"
-
-PASS=0
-FAIL=0
-TOTAL=0
-
-run_test() {
-    local name="$1"
-    shift
-    local log_file="$LOG_DIR/${name}.log"
-    local output_dir="$LOG_DIR/${name}_output"
-
-    TOTAL=$((TOTAL + 1))
-    echo "========================================"
-    echo "[${TOTAL}] Running: ${name}"
-    echo "  Command: python $END2END $*"
-    echo "========================================"
-
-    local start_time
-    start_time=$(date +%s)
-
-    if python "$END2END" "$@" --output-dir "$output_dir" > "$log_file" 2>&1; then
-        local end_time
-        end_time=$(date +%s)
-        local elapsed=$((end_time - start_time))
-        echo "  => PASS (${elapsed}s)"
-        echo "[PASS] ${name} (${elapsed}s)" >> "$SUMMARY_FILE"
-        PASS=$((PASS + 1))
-    else
-        local exit_code=$?
-        local end_time
-        end_time=$(date +%s)
-        local elapsed=$((end_time - start_time))
-        echo "  => FAIL (exit code: ${exit_code}, ${elapsed}s)"
-        echo "[FAIL] ${name} (exit code: ${exit_code}, ${elapsed}s)" >> "$SUMMARY_FILE"
-        # Append last 30 lines of log to summary
-        echo "  --- last 30 lines of log ---" >> "$SUMMARY_FILE"
-        tail -n 30 "$log_file" >> "$SUMMARY_FILE"
-        echo "  --- end ---" >> "$SUMMARY_FILE"
-        FAIL=$((FAIL + 1))
-    fi
-    echo ""
-}
-
-# =============================================================================
-# Test cases
-# =============================================================================
-
-# --- 1. Text-only input ---
-run_test "text_to_text" \
-    --model "$MODEL_PATH" --query-type text --modalities text
-
-run_test "text_to_text_audio" \
-    --model "$MODEL_PATH" --query-type text --modalities text,audio
-
-run_test "text_to_audio" \
-    --model "$MODEL_PATH" --query-type text --modalities audio
-
-# --- 2. Image input ---
-run_test "image_to_text" \
-    --model "$MODEL_PATH" --query-type use_image --modalities text
-
-run_test "image_to_text_audio" \
-    --model "$MODEL_PATH" --query-type use_image --modalities text,audio
-
-# --- 3. Audio input ---
-run_test "audio_to_text" \
-    --model "$MODEL_PATH" --query-type use_audio --modalities text
-
-run_test "audio_to_text_audio" \
-    --model "$MODEL_PATH" --query-type use_audio --modalities text,audio
-
-# --- 4. Video input ---
-run_test "video_to_text" \
-    --model "$MODEL_PATH" --query-type use_video --modalities text
-
-run_test "video_to_text_audio" \
-    --model "$MODEL_PATH" --query-type use_video --modalities text,audio
-
-# --- 5. Multi-audio input ---
-run_test "multi_audios_to_text" \
-    --model "$MODEL_PATH" --query-type use_multi_audios --modalities text
-
-run_test "multi_audios_to_text_audio" \
-    --model "$MODEL_PATH" --query-type use_multi_audios --modalities text,audio
-
-# --- 6. Mixed modality input (audio+image+video) ---
-run_test "mixed_to_text" \
-    --model "$MODEL_PATH" --query-type use_mixed_modalities --modalities text
-
-run_test "mixed_to_text_audio" \
-    --model "$MODEL_PATH" --query-type use_mixed_modalities --modalities text,audio
-
-# --- 7. Audio from video ---
-run_test "audio_in_video_to_text" \
-    --model "$MODEL_PATH" --query-type use_audio_in_video --modalities text
-
-run_test "audio_in_video_to_text_audio" \
-    --model "$MODEL_PATH" --query-type use_audio_in_video --modalities text,audio
-
-# --- 8. py-generator mode ---
-run_test "text_to_text_generator" \
-    --model "$MODEL_PATH" --query-type text --modalities text --py-generator
-
-run_test "audio_to_text_audio_generator" \
-    --model "$MODEL_PATH" --query-type use_audio --modalities text,audio --py-generator
-
-# --- 9. Multi-prompt batch inference ---
-run_test "text_to_text_batch3" \
-    --model "$MODEL_PATH" --query-type text --modalities text --num-prompts 3
-
-# =============================================================================
-# Summary
-# =============================================================================
-echo "========================================" | tee -a "$SUMMARY_FILE"
-echo "Tests done: total ${TOTAL}, passed ${PASS}, failed ${FAIL}" | tee -a "$SUMMARY_FILE"
-echo "========================================" | tee -a "$SUMMARY_FILE"
-echo ""
-echo "Detailed logs: $LOG_DIR"
-echo "Summary report: $SUMMARY_FILE"
-
-exit $FAIL

examples/offline_inference/qwen3_omni/run_tests_async_chunk.sh

@@ -1,143 +0,0 @@
-#!/usr/bin/env bash
-# =============================================================================
-# Qwen3-Omni end2end_async_chunk functional test script
-# Uses default built-in assets, iterates over AsyncOmni query-type x output-modalities combinations
-# =============================================================================
-
-set -euo pipefail
-
-export VLLM_OMNI_USE_V2_RUNNER=1
-
-MODEL_PATH="/models/Qwen3-Omni-30B-A3B-Instruct"
-SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
-END2END="$SCRIPT_DIR/end2end_async_chunk.py"
-STAGE_CONFIGS="$SCRIPT_DIR/../../../vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml"
-LOG_DIR="$SCRIPT_DIR/test_logs_async_chunk"
-SUMMARY_FILE="$LOG_DIR/summary.log"
-
-rm -rf "$LOG_DIR"
-mkdir -p "$LOG_DIR"
-> "$SUMMARY_FILE"
-
-PASS=0
-FAIL=0
-TOTAL=0
-
-run_test() {
-    local name="$1"
-    shift
-    local log_file="$LOG_DIR/${name}.log"
-    local output_dir="$LOG_DIR/${name}_output"
-
-    TOTAL=$((TOTAL + 1))
-    echo "========================================"
-    echo "[${TOTAL}] Running: ${name}"
-    echo "  Command: python $END2END $*"
-    echo "========================================"
-
-    local start_time
-    start_time=$(date +%s)
-
-    if python "$END2END" "$@" --output-dir "$output_dir" > "$log_file" 2>&1; then
-        local end_time
-        end_time=$(date +%s)
-        local elapsed=$((end_time - start_time))
-        echo "  => PASS (${elapsed}s)"
-        echo "[PASS] ${name} (${elapsed}s)" >> "$SUMMARY_FILE"
-        PASS=$((PASS + 1))
-    else
-        local exit_code=$?
-        local end_time
-        end_time=$(date +%s)
-        local elapsed=$((end_time - start_time))
-        echo "  => FAIL (exit code: ${exit_code}, ${elapsed}s)"
-        echo "[FAIL] ${name} (exit code: ${exit_code}, ${elapsed}s)" >> "$SUMMARY_FILE"
-        # Append last 30 lines of log to summary
-        echo "  --- last 30 lines of log ---" >> "$SUMMARY_FILE"
-        tail -n 30 "$log_file" >> "$SUMMARY_FILE"
-        echo "  --- end ---" >> "$SUMMARY_FILE"
-        FAIL=$((FAIL + 1))
-    fi
-    echo ""
-}
-
-# =============================================================================
-# Test cases
-# =============================================================================
-
-# --- 1. Text-only input ---
-run_test "text_to_text" \
-    --model "$MODEL_PATH" --stage-configs-path "$STAGE_CONFIGS" \
-    --query-type text --modalities text
-
-run_test "text_to_text_audio" \
-    --model "$MODEL_PATH" --stage-configs-path "$STAGE_CONFIGS" \
-    --query-type text --modalities text,audio
-
-run_test "text_to_audio" \
-    --model "$MODEL_PATH" --stage-configs-path "$STAGE_CONFIGS" \
-    --query-type text --modalities audio
-
-# --- 2. Image input ---
-run_test "image_to_text" \
-    --model "$MODEL_PATH" --stage-configs-path "$STAGE_CONFIGS" \
-    --query-type use_image --modalities text
-
-run_test "image_to_text_audio" \
-    --model "$MODEL_PATH" --stage-configs-path "$STAGE_CONFIGS" \
-    --query-type use_image --modalities text,audio
-
-# --- 3. Audio input ---
-run_test "audio_to_text" \
-    --model "$MODEL_PATH" --stage-configs-path "$STAGE_CONFIGS" \
-    --query-type use_audio --modalities text
-
-run_test "audio_to_text_audio" \
-    --model "$MODEL_PATH" --stage-configs-path "$STAGE_CONFIGS" \
-    --query-type use_audio --modalities text,audio
-
-# --- 4. Video input ---
-run_test "video_to_text" \
-    --model "$MODEL_PATH" --stage-configs-path "$STAGE_CONFIGS" \
-    --query-type use_video --modalities text
-
-run_test "video_to_text_audio" \
-    --model "$MODEL_PATH" --stage-configs-path "$STAGE_CONFIGS" \
-    --query-type use_video --modalities text,audio
-
-# --- 5. Multi-prompt batch inference ---
-run_test "text_to_text_batch3" \
-    --model "$MODEL_PATH" --stage-configs-path "$STAGE_CONFIGS" \
-    --query-type text --modalities text --num-prompts 3
-
-# --- 6. Concurrent requests (max-in-flight) ---
-run_test "text_to_text_audio_concurrent2" \
-    --model "$MODEL_PATH" --stage-configs-path "$STAGE_CONFIGS" \
-    --query-type text --modalities text,audio --num-prompts 2 --max-in-flight 2
-
-# --- 7. stream-audio-to-disk mode ---
-run_test "audio_to_text_audio_stream" \
-    --model "$MODEL_PATH" --stage-configs-path "$STAGE_CONFIGS" \
-    --query-type use_audio --modalities text,audio --stream-audio-to-disk
-
-# --- 8. request-timeout test (long enough timeout, should PASS) ---
-run_test "text_to_text_with_timeout" \
-    --model "$MODEL_PATH" --stage-configs-path "$STAGE_CONFIGS" \
-    --query-type text --modalities text --request-timeout-s 600
-
-# --- 9. batch-timeout test (long enough timeout, should PASS) ---
-run_test "text_to_text_audio_batch_timeout" \
-    --model "$MODEL_PATH" --stage-configs-path "$STAGE_CONFIGS" \
-    --query-type text --modalities text,audio --batch-timeout-s 600
-
-# =============================================================================
-# Summary
-# =============================================================================
-echo "========================================" | tee -a "$SUMMARY_FILE"
-echo "Tests done: total ${TOTAL}, passed ${PASS}, failed ${FAIL}" | tee -a "$SUMMARY_FILE"
-echo "========================================" | tee -a "$SUMMARY_FILE"
-echo ""
-echo "Detailed logs: $LOG_DIR"
-echo "Summary report: $SUMMARY_FILE"
-
-exit $FAIL

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/end2end.py

@@ -135,18 +135,17 @@ def _estimate_prompt_len(
         return 2048
 
 
-def get_custom_voice_query(use_batch_sample: bool = False, model_name: str | None = None) -> QueryResult:
+def get_custom_voice_query(use_batch_sample: bool = False) -> QueryResult:
     """Build CustomVoice sample inputs.
 
     Args:
         use_batch_sample: When True, return a batch of prompts; otherwise a single prompt.
-        model_name: Override the default HF model path (e.g. with a local path).
 
     Returns:
         QueryResult with Omni inputs and the CustomVoice model path.
     """
     task_type = "CustomVoice"
-    model_name = model_name or "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"
+    model_name = "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"
     if use_batch_sample:
         texts = [
             "其实我真的有发现，我是一个特别善于观察别人情绪的人。",
@@ -197,18 +196,17 @@ def get_custom_voice_query(use_batch_sample: bool = False, model_name: str | Non
     )
 
 
-def get_voice_design_query(use_batch_sample: bool = False, model_name: str | None = None) -> QueryResult:
+def get_voice_design_query(use_batch_sample: bool = False) -> QueryResult:
     """Build VoiceDesign sample inputs.
 
     Args:
         use_batch_sample: When True, return a batch of prompts; otherwise a single prompt.
-        model_name: Override the default HF model path (e.g. with a local path).
 
     Returns:
         QueryResult with Omni inputs and the VoiceDesign model path.
     """
     task_type = "VoiceDesign"
-    model_name = model_name or "Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign"
+    model_name = "Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign"
     if use_batch_sample:
         texts = [
             "哥哥，你回来啦，人家等了你好久好久了，要抱抱！",
@@ -257,19 +255,18 @@ def get_voice_design_query(use_batch_sample: bool = False, model_name: str | Non
     )
 
 
-def get_base_query(use_batch_sample: bool = False, mode_tag: str = "icl", model_name: str | None = None) -> QueryResult:
+def get_base_query(use_batch_sample: bool = False, mode_tag: str = "icl") -> QueryResult:
     """Build Base (voice clone) sample inputs.
 
     Args:
         use_batch_sample: When True, return a batch of prompts (Case 2).
         mode_tag: "icl" or "xvec_only" to control x_vector_only_mode behavior.
-        model_name: Override the default HF model path (e.g. with a local path).
 
     Returns:
         QueryResult with Omni inputs and the Base model path.
     """
     task_type = "Base"
-    model_name = model_name or "Qwen/Qwen3-TTS-12Hz-1.7B-Base"
+    model_name = "Qwen/Qwen3-TTS-12Hz-1.7B-Base"
     ref_audio_path_1 = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-TTS-Repo/clone_2.wav"
     ref_audio_single = ref_audio_path_1
     ref_text_single = (
@@ -338,17 +335,14 @@ def _build_inputs(args) -> tuple[str, list]:
         )
 
     query_func = query_map[args.query_type]
-    model_override = args.model if args.model else None
     if args.query_type in {"CustomVoice", "VoiceDesign"}:
-        query_result = query_func(use_batch_sample=args.use_batch_sample, model_name=model_override)
+        query_result = query_func(use_batch_sample=args.use_batch_sample)
     elif args.query_type == "Base":
-        query_result = query_func(
-            use_batch_sample=args.use_batch_sample, mode_tag=args.mode_tag, model_name=model_override
-        )
+        query_result = query_func(use_batch_sample=args.use_batch_sample, mode_tag=args.mode_tag)
     else:
         query_result = query_func()
 
-    model_name = args.model if args.model else query_result.model_name
+    model_name = query_result.model_name
 
     if args.txt_prompts:
         with open(args.txt_prompts) as f:
@@ -365,10 +359,7 @@ def _build_inputs(args) -> tuple[str, list]:
             for t in lines
         ]
     else:
-        if isinstance(query_result.inputs, list):
-            inputs = query_result.inputs * args.num_prompts
-        else:
-            inputs = [query_result.inputs] * args.num_prompts
+        inputs = query_result.inputs if isinstance(query_result.inputs, list) else [query_result.inputs]
 
     return model_name, inputs
 
@@ -443,12 +434,6 @@ async def main_streaming(args):
 
 def parse_args():
     parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
-    parser.add_argument(
-        "--model",
-        type=str,
-        default=None,
-        help="Override the default model path for the selected query type.",
-    )
     parser.add_argument(
         "--query-type",
         "-q",

examples/offline_inference/text_to_speech/qwen3_tts/end2end_benchmark_script.py

@@ -1,640 +0,0 @@
-"""Offline benchmark script for Qwen3 TTS via vLLM Omni.
-
-Supports warmup + test rounds with a single Omni instance to avoid
-repeated model loading. Prints per-round wall time summary at the end.
-"""
-
-import asyncio
-import logging
-import os
-import time
-from typing import Any, NamedTuple
-
-import soundfile as sf
-import torch
-
-os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
-
-from vllm.utils.argparse_utils import FlexibleArgumentParser
-
-from vllm_omni import AsyncOmni, Omni
-
-logger = logging.getLogger(__name__)
-
-
-class QueryResult(NamedTuple):
-    """Container for a prepared Omni request."""
-
-    inputs: dict
-    model_name: str
-
-
-_model_cache: dict[str, Any] = {}
-
-
-def _estimate_prompt_len(
-    additional_information: dict[str, Any],
-    model_name: str,
-) -> int:
-    """Estimate prompt_token_ids placeholder length for the Talker stage."""
-    try:
-        from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import Qwen3TTSConfig
-        from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import (
-            Qwen3TTSTalkerForConditionalGeneration,
-        )
-
-        if model_name not in _model_cache:
-            from transformers import AutoTokenizer
-
-            tok = AutoTokenizer.from_pretrained(
-                model_name, trust_remote_code=True, padding_side="left", fix_mistral_regex=True
-            )
-            cfg = Qwen3TTSConfig.from_pretrained(model_name, trust_remote_code=True)
-
-            speech_tok = None
-            try:
-                from transformers.utils import cached_file
-
-                from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_tokenizer import Qwen3TTSTokenizer
-
-                st_cfg_path = cached_file(model_name, "speech_tokenizer/config.json")
-                if st_cfg_path:
-                    speech_tok = Qwen3TTSTokenizer.from_pretrained(
-                        # os.path.dirname(st_cfg_path), torch_dtype=torch.bfloat16
-                        os.path.dirname(st_cfg_path),
-                        torch_dtype=torch.bfloat16,
-                        fix_mistral_regex=True,
-                    )
-                    logger.info("Loaded speech tokenizer for exact ref_code_len estimation")
-            except Exception as e:
-                logger.debug("Could not load speech tokenizer: %s", e)
-
-            _model_cache[model_name] = (tok, getattr(cfg, "talker_config", None), speech_tok)
-
-        tok, tcfg, speech_tok = _model_cache[model_name]
-        task_type = (additional_information.get("task_type") or ["CustomVoice"])[0]
-
-        def _estimate_ref_code_len(ref_audio: object) -> int | None:
-            if not isinstance(ref_audio, (str, list)):
-                return None
-            audio_path = ref_audio[0] if isinstance(ref_audio, list) else ref_audio
-            if not isinstance(audio_path, str) or not audio_path.strip():
-                return None
-            try:
-                from vllm.multimodal.media import MediaConnector
-
-                connector = MediaConnector(allowed_local_media_path="/")
-                audio, sr = connector.fetch_audio(audio_path)
-                import numpy as np
-
-                wav_np = np.asarray(audio, dtype=np.float32)
-
-                if speech_tok is not None:
-                    enc = speech_tok.encode(wav_np, sr=int(sr), return_dict=True)
-                    ref_code = getattr(enc, "audio_codes", None)
-                    if isinstance(ref_code, list):
-                        ref_code = ref_code[0] if ref_code else None
-                    if ref_code is not None and hasattr(ref_code, "shape"):
-                        shape = ref_code.shape
-                        return int(shape[0]) if len(shape) == 2 else int(shape[1]) if len(shape) == 3 else None
-
-                codec_hz = getattr(tcfg, "codec_frame_rate", None) or 12
-                return int(len(audio) / sr * codec_hz)
-            except Exception:
-                return None
-
-        return Qwen3TTSTalkerForConditionalGeneration.estimate_prompt_len_from_additional_information(
-            additional_information=additional_information,
-            task_type=task_type,
-            tokenize_prompt=lambda t: tok(t, padding=False)["input_ids"],
-            codec_language_id=getattr(tcfg, "codec_language_id", None),
-            spk_is_dialect=getattr(tcfg, "spk_is_dialect", None),
-            estimate_ref_code_len=_estimate_ref_code_len,
-        )
-    except Exception as exc:
-        logger.warning("Failed to estimate prompt length, using fallback 2048: %s", exc)
-        return 2048
-
-
-def get_custom_voice_query(use_batch_sample: bool = False) -> QueryResult:
-    task_type = "CustomVoice"
-    model_name = "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"
-    if use_batch_sample:
-        texts = [
-            "其实我真的有发现，我是一个特别善于观察别人情绪的人。",
-            "She said she would be here by noon.",
-            "I like you very much.",
-            "Really, you do?",
-            "Yes, absolutely.",
-        ]
-        instructs = ["", "Very happy.", "Very happy.", "Very happy.", "Very happy."]
-        languages = ["Chinese", "English", "English", "English", "English"]
-        speakers = ["Vivian", "Ryan", "Ryan", "Ryan", "Ryan"]
-        inputs = []
-        for text, instruct, language, speaker in zip(texts, instructs, languages, speakers):
-            additional_information = {
-                "task_type": [task_type],
-                "text": [text],
-                "instruct": [instruct],
-                "language": [language],
-                "speaker": [speaker],
-                "max_new_tokens": [2048],
-            }
-            inputs.append(
-                {
-                    "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name),
-                    "additional_information": additional_information,
-                }
-            )
-    else:
-        text = "其实我真的有发现，我是一个特别善于观察别人情绪的人。"
-        language = "Chinese"
-        speaker = "Vivian"
-        instruct = "用特别愤怒的语气说"
-        additional_information = {
-            "task_type": [task_type],
-            "text": [text],
-            "language": [language],
-            "speaker": [speaker],
-            "instruct": [instruct],
-            "max_new_tokens": [2048],
-        }
-        inputs = {
-            "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name),
-            "additional_information": additional_information,
-        }
-    return QueryResult(inputs=inputs, model_name=model_name)
-
-
-def get_voice_design_query(use_batch_sample: bool = False) -> QueryResult:
-    task_type = "VoiceDesign"
-    model_name = "Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign"
-    if use_batch_sample:
-        texts = [
-            "哥哥，你回来啦，人家等了你好久好久了，要抱抱！",
-            "It's in the top drawer... wait, it's empty? No way, that's impossible! I'm sure I put it there!",
-        ]
-        instructs = [
-            "体现撒娇稚嫩的萝莉女声，音调偏高且起伏明显，营造出黏人、做作又刻意卖萌的听觉效果。",
-            "Speak in an incredulous tone, but with a hint of panic beginning to creep into your voice.",
-        ]
-        languages = ["Chinese", "English"]
-        inputs = []
-        for text, instruct, language in zip(texts, instructs, languages):
-            additional_information = {
-                "task_type": [task_type],
-                "text": [text],
-                "language": [language],
-                "instruct": [instruct],
-                "max_new_tokens": [2048],
-                "non_streaming_mode": [True],
-            }
-            inputs.append(
-                {
-                    "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name),
-                    "additional_information": additional_information,
-                }
-            )
-    else:
-        text = "哥哥，你回来啦，人家等了你好久好久了，要抱抱！"
-        instruct = "体现撒娇稚嫩的萝莉女声，音调偏高且起伏明显，营造出黏人、做作又刻意卖萌的听觉效果。"
-        language = "Chinese"
-        additional_information = {
-            "task_type": [task_type],
-            "text": [text],
-            "language": [language],
-            "instruct": [instruct],
-            "max_new_tokens": [2048],
-            "non_streaming_mode": [True],
-        }
-        inputs = {
-            "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name),
-            "additional_information": additional_information,
-        }
-    return QueryResult(inputs=inputs, model_name=model_name)
-
-
-def get_base_query(use_batch_sample: bool = False, mode_tag: str = "icl") -> QueryResult:
-    task_type = "Base"
-    model_name = "Qwen/Qwen3-TTS-12Hz-1.7B-Base"
-    ref_audio_path_1 = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-TTS-Repo/clone_2.wav"
-    ref_audio_single = ref_audio_path_1
-    ref_text_single = (
-        "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
-    )
-    syn_text_single = "Good one. Okay, fine, I'm just gonna leave this sock monkey here. Goodbye."
-    syn_lang_single = "Auto"
-    x_vector_only_mode = mode_tag == "xvec_only"
-    if use_batch_sample:
-        syn_text_batch = [
-            "Good one. Okay, fine, I'm just gonna leave this sock monkey here. Goodbye.",
-            "其实我真的有发现，我是一个特别善于观察别人情绪的人。",
-        ]
-        syn_lang_batch = ["Chinese", "English"]
-        inputs = []
-        for text, language in zip(syn_text_batch, syn_lang_batch):
-            additional_information = {
-                "task_type": [task_type],
-                "ref_audio": [ref_audio_single],
-                "ref_text": [ref_text_single],
-                "text": [text],
-                "language": [language],
-                "x_vector_only_mode": [x_vector_only_mode],
-                "max_new_tokens": [2048],
-            }
-            inputs.append(
-                {
-                    "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name),
-                    "additional_information": additional_information,
-                }
-            )
-    else:
-        additional_information = {
-            "task_type": [task_type],
-            "ref_audio": [ref_audio_single],
-            "ref_text": [ref_text_single],
-            "text": [syn_text_single],
-            "language": [syn_lang_single],
-            "x_vector_only_mode": [x_vector_only_mode],
-            "max_new_tokens": [2048],
-        }
-        inputs = {
-            "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name),
-            "additional_information": additional_information,
-        }
-    return QueryResult(inputs=inputs, model_name=model_name)
-
-
-query_map = {
-    "CustomVoice": get_custom_voice_query,
-    "VoiceDesign": get_voice_design_query,
-    "Base": get_base_query,
-}
-
-
-def _build_inputs(args) -> tuple[str, list]:
-    """Resolve model name and inputs list from CLI args."""
-    if args.batch_size < 1 or (args.batch_size & (args.batch_size - 1)) != 0:
-        raise ValueError(
-            f"--batch-size must be a power of two (got {args.batch_size}); "
-            "non-power-of-two values do not align with CUDA graph capture sizes "
-            "of Code2Wav."
-        )
-
-    query_func = query_map[args.query_type]
-    if args.query_type in {"CustomVoice", "VoiceDesign"}:
-        query_result = query_func(use_batch_sample=args.use_batch_sample)
-    elif args.query_type == "Base":
-        query_result = query_func(use_batch_sample=args.use_batch_sample, mode_tag=args.mode_tag)
-    else:
-        query_result = query_func()
-
-    model_name = query_result.model_name
-
-    if args.txt_prompts:
-        with open(args.txt_prompts) as f:
-            lines = [line.strip() for line in f if line.strip()]
-        if not lines:
-            raise ValueError(f"No valid prompts found in {args.txt_prompts}")
-        template = query_result.inputs if not isinstance(query_result.inputs, list) else query_result.inputs[0]
-        template_info = template["additional_information"]
-        inputs = [
-            {
-                "prompt_token_ids": [0] * _estimate_prompt_len({**template_info, "text": [t]}, model_name),
-                "additional_information": {**template_info, "text": [t]},
-            }
-            for t in lines
-        ]
-    else:
-        inputs = query_result.inputs if isinstance(query_result.inputs, list) else [query_result.inputs]
-
-    return model_name, inputs
-
-
-def _save_wav(output_dir: str, request_id: str, mm: dict) -> None:
-    """Concatenate audio chunks and write to a wav file."""
-    audio_data = mm["audio"]
-    sr_raw = mm["sr"]
-    sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw
-    sr = 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
-    out_wav = os.path.join(output_dir, f"output_{request_id}.wav")
-    sf.write(out_wav, audio_tensor.float().cpu().numpy().flatten(), samplerate=sr, format="WAV")
-    logger.info("Saved audio to %s", out_wav)
-
-
-class _SummaryCollector:
-    """Monkey-patches omni._log_summary_and_cleanup to capture summary dicts.
-
-    NOTE: This relies on a private method ``_log_summary_and_cleanup`` which may
-    break if the upstream Omni implementation is refactored.  If Omni exposes a
-    public hook for summary collection in the future, prefer that instead.
-    """
-
-    def __init__(self, omni):
-        self.omni = omni
-        self.summaries: list[dict[str, Any]] = []
-        self.collecting = False
-        self._orig = omni._log_summary_and_cleanup
-
-    def install(self):
-        from pprint import pformat as _pformat
-
-        collector = self
-
-        def _patched(request_id: str) -> None:
-            req_state = collector.omni.request_states.get(request_id)
-            try:
-                if req_state is None or req_state.metrics is None:
-                    return
-                summary = req_state.metrics.build_and_log_summary()
-                logger.info("[Summary] %s", _pformat(summary, sort_dicts=False))
-                if collector.collecting and summary:
-                    collector.summaries.append(summary)
-            except Exception:
-                logger.exception("Failed to build/log summary for req=%s", request_id)
-            finally:
-                collector.omni.request_states.pop(request_id, None)
-
-        collector.omni._log_summary_and_cleanup = _patched
-
-
-def _print_summary(round_times: list[float], summaries: list[dict[str, Any]] | None = None) -> None:
-    """Print per-round wall time and optional detailed pipeline summary."""
-    print("\n" + "=" * 60)
-    print("BENCHMARK SUMMARY")
-    print("=" * 60)
-    for i, t in enumerate(round_times):
-        print(f"  Test round {i}: {t:.3f}s")
-    print("-" * 60)
-    avg = sum(round_times) / len(round_times)
-    print(f"  Rounds:  {len(round_times)}")
-    print(f"  Total:   {sum(round_times):.3f}s")
-    print(f"  Average: {avg:.3f}s")
-    print(f"  Min:     {min(round_times):.3f}s")
-    print(f"  Max:     {max(round_times):.3f}s")
-    print("=" * 60)
-
-    if not summaries:
-        print()
-        return
-
-    # --- Aggregate from captured summaries ---
-    overall_list = [s["overall_summary"] for s in summaries if "overall_summary" in s]
-    if not overall_list:
-        print()
-        return
-
-    n = len(overall_list)
-
-    # Average Wall time
-    avg_wall = sum(d["e2e_wall_time_ms"] for d in overall_list) / n
-    print(f"\n{'=' * 60}")
-    print("DETAILED PIPELINE SUMMARY (test rounds average)")
-    print("=" * 60)
-    print(f"  Avg Wall Time:           {avg_wall:.3f} ms")
-
-    # Per-stage timing
-    stage_keys = sorted(k for k in overall_list[0] if k.startswith("e2e_stage_") and k.endswith("_wall_time_ms"))
-    if stage_keys:
-        print("-" * 60)
-        print("  Per-stage timing (avg):")
-        for sk in stage_keys:
-            stage_avg = sum(d.get(sk, 0.0) for d in overall_list) / n
-            # extract stage index from key like e2e_stage_0_wall_time_ms
-            stage_idx = sk.replace("e2e_stage_", "").replace("_wall_time_ms", "")
-            print(f"    Stage {stage_idx} wall time:   {stage_avg:.3f} ms")
-
-    # Per-stage gen time from stage_table
-    all_stage_gen: dict[int, list[float]] = {}
-    for s in summaries:
-        for req_entry in s.get("stage_table", []):
-            for stage_info in req_entry.get("stages", []):
-                sid = stage_info.get("stage_id", -1)
-                gen_ms = stage_info.get("stage_gen_time_ms", 0.0)
-                all_stage_gen.setdefault(sid, []).append(gen_ms)
-    if all_stage_gen:
-        print("  Per-stage gen time (avg):")
-        for sid in sorted(all_stage_gen):
-            vals = all_stage_gen[sid]
-            print(f"    Stage {sid} gen time:     {sum(vals) / len(vals):.3f} ms")
-
-    # Per-prompt throughput
-    avg_tokens_per_s = sum(d.get("e2e_avg_tokens_per_s", 0.0) for d in overall_list) / n
-    avg_time_per_req = sum(d.get("e2e_avg_time_per_request_ms", 0.0) for d in overall_list) / n
-    total_tokens_list = [d.get("e2e_total_tokens", 0) for d in overall_list]
-    avg_total_tokens = sum(total_tokens_list) / n
-    print("-" * 60)
-    print("  Per-prompt throughput (avg):")
-    print(f"    Tokens per request:    {avg_total_tokens:.1f}")
-    print(f"    Time per request:      {avg_time_per_req:.3f} ms")
-    print(f"    Tokens/s:              {avg_tokens_per_s:.3f}")
-    print("=" * 60 + "\n")
-
-
-def main(args):
-    """Run sync benchmark: warmup rounds then test rounds."""
-    model_name, inputs = _build_inputs(args)
-    output_dir = args.output_dir
-    os.makedirs(output_dir, exist_ok=True)
-
-    omni = Omni(
-        model=model_name,
-        stage_configs_path=args.stage_configs_path,
-        log_stats=args.log_stats,
-        stage_init_timeout=args.stage_init_timeout,
-        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
-    )
-
-    collector = _SummaryCollector(omni)
-    collector.install()
-
-    batch_size = args.batch_size
-
-    # --- Warmup ---
-    for r in range(args.warmup_rounds):
-        print(f"[Warmup {r}/{args.warmup_rounds}] running...")
-        for batch_start in range(0, len(inputs), batch_size):
-            batch = inputs[batch_start : batch_start + batch_size]
-            for _ in omni.generate(batch):
-                pass
-        print(f"[Warmup {r}/{args.warmup_rounds}] done")
-
-    # --- Test ---
-    collector.collecting = True
-    round_times = []
-    for r in range(args.test_rounds):
-        print(f"[Test {r}/{args.test_rounds}] running...")
-        t_start = time.perf_counter()
-        for batch_start in range(0, len(inputs), batch_size):
-            batch = inputs[batch_start : batch_start + batch_size]
-            for stage_outputs in omni.generate(batch):
-                output = stage_outputs.request_output
-                _save_wav(
-                    output_dir,
-                    f"round{r}_{output.request_id}",
-                    output.outputs[0].multimodal_output,
-                )
-        elapsed = time.perf_counter() - t_start
-        round_times.append(elapsed)
-        print(f"[Test {r}/{args.test_rounds}] done in {elapsed:.3f}s")
-    collector.collecting = False
-
-    _print_summary(round_times, collector.summaries if args.log_stats else None)
-
-
-async def main_streaming(args):
-    """Run async streaming benchmark: warmup rounds then test rounds."""
-    model_name, inputs = _build_inputs(args)
-    output_dir = args.output_dir
-    os.makedirs(output_dir, exist_ok=True)
-
-    omni = AsyncOmni(
-        model=model_name,
-        stage_configs_path=args.stage_configs_path,
-        log_stats=args.log_stats,
-        stage_init_timeout=args.stage_init_timeout,
-        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
-    )
-
-    collector = _SummaryCollector(omni)
-    collector.install()
-
-    # --- Warmup ---
-    for r in range(args.warmup_rounds):
-        print(f"[Warmup {r}/{args.warmup_rounds}] running...")
-        for i, prompt in enumerate(inputs):
-            async for _ in omni.generate(prompt, request_id=f"warmup_{r}_{i}"):
-                pass
-        print(f"[Warmup {r}/{args.warmup_rounds}] done")
-
-    # --- Test ---
-    collector.collecting = True
-    round_times = []
-    for r in range(args.test_rounds):
-        print(f"[Test {r}/{args.test_rounds}] running...")
-        t_start = time.perf_counter()
-        for i, prompt in enumerate(inputs):
-            request_id = f"round{r}_{i}"
-            t_req_start = time.perf_counter()
-            t_prev = t_req_start
-            chunk_idx = 0
-            async for stage_output in omni.generate(prompt, request_id=request_id):
-                mm = stage_output.request_output.outputs[0].multimodal_output
-                if not stage_output.finished:
-                    t_now = time.perf_counter()
-                    audio = mm.get("audio")
-                    n = len(audio) if isinstance(audio, list) else (0 if audio is None else 1)
-                    if chunk_idx == 0:
-                        ttfa_ms = (t_now - t_req_start) * 1000
-                        logger.info("Request %s: chunk %d samples=%d TTFA=%.1fms", request_id, chunk_idx, n, ttfa_ms)
-                    else:
-                        dt_ms = (t_now - t_prev) * 1000
-                        logger.info(
-                            "Request %s: chunk %d samples=%d inter_chunk=%.1fms", request_id, chunk_idx, n, dt_ms
-                        )
-                    t_prev = t_now
-                    chunk_idx += 1
-                else:
-                    t_end = time.perf_counter()
-                    total_ms = (t_end - t_req_start) * 1000
-                    logger.info("Request %s: done total=%.1fms chunks=%d", request_id, total_ms, chunk_idx)
-                    _save_wav(output_dir, request_id, mm)
-        elapsed = time.perf_counter() - t_start
-        round_times.append(elapsed)
-        print(f"[Test {r}/{args.test_rounds}] done in {elapsed:.3f}s")
-    collector.collecting = False
-
-    _print_summary(round_times, collector.summaries if args.log_stats else None)
-
-
-def parse_args():
-    parser = FlexibleArgumentParser(description="Benchmark script for Qwen3 TTS via vLLM Omni")
-    parser.add_argument(
-        "--query-type",
-        "-q",
-        type=str,
-        default="CustomVoice",
-        choices=query_map.keys(),
-        help="Query type.",
-    )
-    parser.add_argument(
-        "--warmup-rounds",
-        type=int,
-        default=3,
-        help="Number of warmup rounds (results discarded, default: 3).",
-    )
-    parser.add_argument(
-        "--test-rounds",
-        type=int,
-        default=7,
-        help="Number of test rounds (results saved and timed, default: 7).",
-    )
-    parser.add_argument(
-        "--log-stats",
-        action="store_true",
-        default=False,
-        help="Enable writing detailed statistics (default: disabled).",
-    )
-    parser.add_argument(
-        "--stage-init-timeout",
-        type=int,
-        default=300,
-        help="Timeout for initializing a single stage in seconds (default: 300).",
-    )
-    parser.add_argument(
-        "--output-dir",
-        default="output_audio",
-        help="Output directory for generated wav files (default: output_audio).",
-    )
-    parser.add_argument(
-        "--txt-prompts",
-        type=str,
-        default=None,
-        help="Path to a .txt file with one prompt per line.",
-    )
-    parser.add_argument(
-        "--stage-configs-path",
-        type=str,
-        default=None,
-        help="Path to a stage configs file.",
-    )
-    parser.add_argument(
-        "--use-batch-sample",
-        action="store_true",
-        default=False,
-        help="Use batch input sample for CustomVoice/VoiceDesign/Base query.",
-    )
-    parser.add_argument(
-        "--mode-tag",
-        type=str,
-        default="icl",
-        choices=["icl", "xvec_only"],
-        help="Mode tag for Base query x_vector_only_mode (default: icl).",
-    )
-    parser.add_argument(
-        "--streaming",
-        action="store_true",
-        default=False,
-        help="Stream audio chunks via AsyncOmni.",
-    )
-    parser.add_argument(
-        "--batch-size",
-        type=int,
-        default=1,
-        help="Number of prompts per batch (default: 1).",
-    )
-    parser.add_argument(
-        "--enable-diffusion-pipeline-profiler",
-        action="store_true",
-        help="Enable diffusion pipeline profiler to display stage durations.",
-    )
-    return parser.parse_args()
-
-
-if __name__ == "__main__":
-    args = parse_args()
-    if args.streaming:
-        asyncio.run(main_streaming(args))
-    else:
-        main(args)

examples/offline_inference/text_to_speech/qwen3_tts/run_tests.sh

@@ -1,173 +0,0 @@
-#!/usr/bin/env bash
-# =============================================================================
-# Qwen3-TTS end2end functional test script
-# Iterates over query-type and mode combinations for qwen3_tts/end2end.py
-#
-# Usage:
-#   ./run_tests.sh --custom-voice-model /path/a --voice-design-model /path/b --base-model /path/c
-#   ./run_tests.sh --custom-voice-model /path/a                # only run CustomVoice tests
-#   ./run_tests.sh --base-model /path/c                        # only run Base tests
-# =============================================================================
-
-set -euo pipefail
-
-export VLLM_OMNI_USE_V2_RUNNER=1
-# export VLLM_ALLOW_LONG_MAX_MODEL_LEN=1
-
-# ---- Parse named arguments ----
-CUSTOM_VOICE_MODEL=""
-VOICE_DESIGN_MODEL=""
-BASE_MODEL=""
-
-usage() {
-    echo "Usage: $0 [--custom-voice-model PATH] [--voice-design-model PATH] [--base-model PATH]"
-    echo ""
-    echo "Pass one or more model paths. Only the query types with a model provided will be tested."
-    exit 1
-}
-
-if [[ $# -eq 0 ]]; then
-    usage
-fi
-
-while [[ $# -gt 0 ]]; do
-    case "$1" in
-        --custom-voice-model) CUSTOM_VOICE_MODEL="$2"; shift 2 ;;
-        --voice-design-model) VOICE_DESIGN_MODEL="$2"; shift 2 ;;
-        --base-model)         BASE_MODEL="$2";         shift 2 ;;
-        -h|--help)            usage ;;
-        *) echo "Unknown option: $1"; usage ;;
-    esac
-done
-
-if [[ -z "$CUSTOM_VOICE_MODEL" && -z "$VOICE_DESIGN_MODEL" && -z "$BASE_MODEL" ]]; then
-    echo "Error: at least one model path must be provided."
-    usage
-fi
-
-SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
-END2END="$SCRIPT_DIR/end2end.py"
-LOG_DIR="$SCRIPT_DIR/test_logs"
-SUMMARY_FILE="$LOG_DIR/summary.log"
-
-rm -rf "$LOG_DIR"
-mkdir -p "$LOG_DIR"
-> "$SUMMARY_FILE"
-
-PASS=0
-FAIL=0
-TOTAL=0
-
-run_test() {
-    local name="$1"
-    shift
-    local log_file="$LOG_DIR/${name}.log"
-    local output_dir="$LOG_DIR/${name}_output"
-
-    TOTAL=$((TOTAL + 1))
-    echo "========================================"
-    echo "[${TOTAL}] Running: ${name}"
-    echo "  Command: python $END2END $*"
-    echo "========================================"
-
-    local start_time
-    start_time=$(date +%s)
-
-    if python "$END2END" "$@" --output-dir "$output_dir" > "$log_file" 2>&1; then
-        local end_time
-        end_time=$(date +%s)
-        local elapsed=$((end_time - start_time))
-        echo "  => PASS (${elapsed}s)"
-        echo "[PASS] ${name} (${elapsed}s)" >> "$SUMMARY_FILE"
-        PASS=$((PASS + 1))
-    else
-        local exit_code=$?
-        local end_time
-        end_time=$(date +%s)
-        local elapsed=$((end_time - start_time))
-        echo "  => FAIL (exit code: ${exit_code}, ${elapsed}s)"
-        echo "[FAIL] ${name} (exit code: ${exit_code}, ${elapsed}s)" >> "$SUMMARY_FILE"
-        echo "  --- last 30 lines of log ---" >> "$SUMMARY_FILE"
-        tail -n 30 "$log_file" >> "$SUMMARY_FILE"
-        echo "  --- end ---" >> "$SUMMARY_FILE"
-        FAIL=$((FAIL + 1))
-    fi
-    echo ""
-}
-
-# =============================================================================
-# Test cases — only run sections whose model path was provided
-# =============================================================================
-
-# --- 1. CustomVoice ---
-if [[ -n "$CUSTOM_VOICE_MODEL" ]]; then
-    run_test "custom_voice" \
-        --model "$CUSTOM_VOICE_MODEL" --query-type CustomVoice
-
-    run_test "custom_voice_batch_sample" \
-        --model "$CUSTOM_VOICE_MODEL" --query-type CustomVoice --use-batch-sample
-
-    run_test "custom_voice_py_generator" \
-        --model "$CUSTOM_VOICE_MODEL" --query-type CustomVoice --py-generator
-
-    run_test "custom_voice_streaming" \
-        --model "$CUSTOM_VOICE_MODEL" --query-type CustomVoice --streaming
-
-    run_test "custom_voice_batch4" \
-        --model "$CUSTOM_VOICE_MODEL" --query-type CustomVoice --batch-size 4
-
-    run_test "custom_voice_batch4_batch_sample" \
-        --model "$CUSTOM_VOICE_MODEL" --query-type CustomVoice --batch-size 4 --use-batch-sample --stage-configs-path "$SCRIPT_DIR/../../../vllm_omni/model_executor/stage_configs/qwen3_tts_batch.yaml"
-
-    run_test "custom_voice_num_prompts3" \
-        --model "$CUSTOM_VOICE_MODEL" --query-type CustomVoice --num-prompts 3
-fi
-
-# --- 2. VoiceDesign ---
-if [[ -n "$VOICE_DESIGN_MODEL" ]]; then
-    run_test "voice_design" \
-        --model "$VOICE_DESIGN_MODEL" --query-type VoiceDesign
-
-    run_test "voice_design_batch_sample" \
-        --model "$VOICE_DESIGN_MODEL" --query-type VoiceDesign --use-batch-sample
-
-    run_test "voice_design_py_generator" \
-        --model "$VOICE_DESIGN_MODEL" --query-type VoiceDesign --py-generator
-
-    run_test "voice_design_streaming" \
-        --model "$VOICE_DESIGN_MODEL" --query-type VoiceDesign --streaming
-fi
-
-# --- 3. Base (icl mode) ---
-if [[ -n "$BASE_MODEL" ]]; then
-    run_test "base_icl" \
-        --model "$BASE_MODEL" --query-type Base --mode-tag icl
-
-    run_test "base_icl_batch_sample" \
-        --model "$BASE_MODEL" --query-type Base --mode-tag icl --use-batch-sample
-
-    run_test "base_icl_py_generator" \
-        --model "$BASE_MODEL" --query-type Base --mode-tag icl --py-generator
-
-    run_test "base_icl_streaming" \
-        --model "$BASE_MODEL" --query-type Base --mode-tag icl --streaming
-
-    # --- 4. Base (xvec_only mode) ---
-    run_test "base_xvec_only" \
-        --model "$BASE_MODEL" --query-type Base --mode-tag xvec_only
-
-    run_test "base_xvec_only_batch_sample" \
-        --model "$BASE_MODEL" --query-type Base --mode-tag xvec_only --use-batch-sample
-fi
-
-# =============================================================================
-# Summary
-# =============================================================================
-echo "========================================" | tee -a "$SUMMARY_FILE"
-echo "Tests done: total ${TOTAL}, passed ${PASS}, failed ${FAIL}" | tee -a "$SUMMARY_FILE"
-echo "========================================" | tee -a "$SUMMARY_FILE"
-echo ""
-echo "Detailed logs: $LOG_DIR"
-echo "Summary report: $SUMMARY_FILE"
-
-exit $FAIL

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/qwen2_5_omni/run_curl_multimodal_generation.sh

@@ -1,9 +1,6 @@
 #!/usr/bin/env bash
 set -euo pipefail
 
-# Default model name
-MODEL="${MODEL:-Qwen/Qwen2.5-Omni-7B}"
-
 # Default query type
 QUERY_TYPE="${1:-mixed_modalities}"
 
@@ -165,7 +162,7 @@ echo ""
 
 request_body=$(cat <<EOF
 {
-  "model": "$MODEL",
+  "model": "Qwen/Qwen2.5-Omni-7B",
   "sampling_params_list": $sampling_params_list,
   "mm_processor_kwargs": $mm_processor_kwargs,
   "modalities": $MODALITIES,

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/comfyui/test_comfyui_integration.py

@@ -523,7 +523,6 @@ def api_server(unused_tcp_port_factory, server_case: ServerCase, mock_async_omni
             "Qwen/Qwen-Image-Edit",
             True,
             id="image-to-image-dalle-endpoint",
-            marks=pytest.mark.skip(reason="Temporarily disabled due to failure."),
         ),
         pytest.param(
             ServerCase(

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/core/sched/test_generation_scheduler_finish_condition.py

@@ -1,133 +0,0 @@
-"""Tests for OmniGenerationScheduler update_from_output finish conditions.
-
-The three-way finish condition in update_from_output determines when a
-generation request (e.g. Code2Wav) is marked FINISHED_STOPPED:
-
-  1. No adapter (sync mode): finish when all prompt tokens are computed.
-  2. Adapter present, adapter done: finish when computed AND adapter signals done.
-  3. Adapter present, adapter NOT done: do NOT finish — wait for next chunk.
-
-These cases are the core correctness guarantee for async_chunk batch mode.
-Without condition #3, a request finishes prematurely after processing
-only the first chunk, producing truncated audio.
-"""
-
-import unittest
-
-from vllm.v1.request import RequestStatus
-
-
-class FakeRequest:
-    """Minimal request for testing finish conditions."""
-
-    def __init__(self, request_id: str, prompt_len: int, num_computed: int):
-        self.request_id = request_id
-        self.prompt_token_ids = list(range(prompt_len))
-        self.num_computed_tokens = num_computed
-        self.status = RequestStatus.RUNNING
-        self.stop_reason = None
-        self.sampling_params = None
-        self.num_cached_tokens = 0
-        self.num_output_tokens = 0
-        self.num_output_placeholders = 0
-        self.num_external_computed_tokens = 0
-        self.trace_headers = None
-        self.num_nans_in_logits = 0
-        self.client_index = 0
-
-    def is_finished(self):
-        return self.status in (
-            RequestStatus.FINISHED_STOPPED,
-            RequestStatus.FINISHED_LENGTH_CAPPED,
-            RequestStatus.FINISHED_ABORTED,
-            RequestStatus.FINISHED_ERROR,
-        )
-
-    def get_finished_reason(self):
-        return "stop"
-
-    def take_events(self):
-        return None
-
-
-class FakeChunkAdapter:
-    """Minimal chunk transfer adapter mock."""
-
-    def __init__(self, finished_request_ids: set[str] | None = None):
-        self.finished_requests = finished_request_ids or set()
-
-    def cleanup(self, req_id, external_req_id=None):
-        pass
-
-
-def _evaluate_finish_condition(request, chunk_transfer_adapter):
-    """Extract the three-way finish logic from update_from_output."""
-    _all_computed = request.num_computed_tokens >= len(request.prompt_token_ids)
-    _adapter_done = (
-        chunk_transfer_adapter is not None and request.request_id in chunk_transfer_adapter.finished_requests
-    )
-    return (
-        request.status == RequestStatus.FINISHED_STOPPED
-        or (_all_computed and chunk_transfer_adapter is None)
-        or (_all_computed and _adapter_done)
-    )
-
-
-class TestFinishConditionSyncMode(unittest.TestCase):
-    """Case 1: No adapter (sync mode) — finish when all tokens computed."""
-
-    def test_finishes_when_all_computed(self):
-        req = FakeRequest("r1", prompt_len=100, num_computed=100)
-        assert _evaluate_finish_condition(req, chunk_transfer_adapter=None)
-
-    def test_does_not_finish_when_partially_computed(self):
-        req = FakeRequest("r1", prompt_len=100, num_computed=50)
-        assert not _evaluate_finish_condition(req, chunk_transfer_adapter=None)
-
-
-class TestFinishConditionAsyncAdapterDone(unittest.TestCase):
-    """Case 2: Adapter present AND adapter signals done — finish."""
-
-    def test_finishes_when_computed_and_adapter_done(self):
-        req = FakeRequest("r1", prompt_len=100, num_computed=100)
-        adapter = FakeChunkAdapter(finished_request_ids={"r1"})
-        assert _evaluate_finish_condition(req, adapter)
-
-    def test_does_not_finish_when_computed_but_adapter_not_done(self):
-        """Case 3: Adapter present but NOT done — must NOT finish."""
-        req = FakeRequest("r1", prompt_len=100, num_computed=100)
-        adapter = FakeChunkAdapter(finished_request_ids=set())
-        assert not _evaluate_finish_condition(req, adapter)
-
-
-class TestFinishConditionAsyncNotComputed(unittest.TestCase):
-    """Adapter present, tokens not all computed — never finishes."""
-
-    def test_does_not_finish_partial_compute_adapter_done(self):
-        req = FakeRequest("r1", prompt_len=100, num_computed=50)
-        adapter = FakeChunkAdapter(finished_request_ids={"r1"})
-        assert not _evaluate_finish_condition(req, adapter)
-
-    def test_does_not_finish_partial_compute_adapter_not_done(self):
-        req = FakeRequest("r1", prompt_len=100, num_computed=50)
-        adapter = FakeChunkAdapter(finished_request_ids=set())
-        assert not _evaluate_finish_condition(req, adapter)
-
-
-class TestFinishConditionAlreadyStopped(unittest.TestCase):
-    """Request already FINISHED_STOPPED — always returns True."""
-
-    def test_already_stopped_no_adapter(self):
-        req = FakeRequest("r1", prompt_len=100, num_computed=0)
-        req.status = RequestStatus.FINISHED_STOPPED
-        assert _evaluate_finish_condition(req, chunk_transfer_adapter=None)
-
-    def test_already_stopped_with_adapter(self):
-        req = FakeRequest("r1", prompt_len=100, num_computed=0)
-        req.status = RequestStatus.FINISHED_STOPPED
-        adapter = FakeChunkAdapter(finished_request_ids=set())
-        assert _evaluate_finish_condition(req, adapter)
-
-
-if __name__ == "__main__":
-    unittest.main()

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/e2e/accuracy/qwen3_omni/test_qwen3_omni.py

@@ -61,7 +61,6 @@ models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
 pytestmark = [pytest.mark.full_model, pytest.mark.omni]
 
 _CI_DEPLOY = get_deploy_config_path("ci/qwen3_omni_moe.yaml")
-_CI_MIN_DAILY_OMNI_ACCURACY = "0.685"
 
 
 def get_chunk_config(config_path: str | None = None):
@@ -120,11 +119,9 @@ def test_qwen3_omni_daily_omni_accuracy_bench(omni_server) -> None:
     _require_vllm_cli()
     pytest.importorskip("datasets")
     pytest.importorskip("huggingface_hub")
-    argv = build_acc_benchmark_cli_argv(omni_server, skip_seed=True, skip_daily=False)
-    # Daily-Omni's full-set score can move by a few MCQ items across GPU and
-    # batching schedules; keep CI below the nominal 0.69 CLI baseline.
-    argv.extend(["--min-daily-omni-accuracy", _CI_MIN_DAILY_OMNI_ACCURACY])
-    ns = _acc_bench.parse_acc_benchmark_args(argv)
+    ns = _acc_bench.parse_acc_benchmark_args(
+        build_acc_benchmark_cli_argv(omni_server, skip_seed=True, skip_daily=False)
+    )
     assert _acc_bench.run_acc_benchmark(ns) == 0
 
 

tests/e2e/offline_inference/test_audiox_model.py

@@ -1,75 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import os
-import sys
-from pathlib import Path
-
-import numpy as np
-import pytest
-import torch
-
-from tests.helpers.mark import hardware_test
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-from vllm_omni.outputs import OmniRequestOutput
-from vllm_omni.platforms import current_omni_platform
-
-# ruff: noqa: E402
-REPO_ROOT = Path(__file__).resolve().parents[2]
-if str(REPO_ROOT) not in sys.path:
-    sys.path.insert(0, str(REPO_ROOT))
-
-from vllm_omni import Omni
-
-# Prefer a tiny/random checkpoint for CI.
-# Override in CI if needed: AUDIOX_TEST_MODEL=<model-or-local-path>
-models = [os.environ.get("AUDIOX_TEST_MODEL", "zhangj1an/audiox_random")]
-
-
-@pytest.mark.core_model
-@pytest.mark.diffusion
-@hardware_test(res={"cuda": "L4", "xpu": "B60"})
-@pytest.mark.parametrize("model_name", models)
-def test_audiox_model(model_name: str):
-    m = Omni(model=model_name, model_class_name="AudioXPipeline")
-
-    # Keep runtime short for CI.
-    seconds_total = 2.0
-    # AudioXPipeline always emits 44.1 kHz stereo (advertised via class-level
-    # ``audio_sample_rate``); the trimmed output should match this rate.
-    sample_rate = 44100
-
-    outputs = m.generate(
-        prompts={"prompt": "A dog barking in a quiet park."},
-        sampling_params_list=OmniDiffusionSamplingParams(
-            num_inference_steps=4,
-            guidance_scale=6.0,
-            generator=torch.Generator(current_omni_platform.device_type).manual_seed(42),
-            num_outputs_per_prompt=1,
-            extra_args={
-                "audiox_task": "t2a",
-                "seconds_start": 0.0,
-                "seconds_total": seconds_total,
-            },
-        ),
-    )
-
-    assert outputs is not None
-    first_output = outputs[0]
-    assert first_output.final_output_type == "audio"
-    assert hasattr(first_output, "request_output") and first_output.request_output
-
-    req_out = first_output.request_output
-    assert isinstance(req_out, OmniRequestOutput)
-    assert req_out.final_output_type == "audio"
-    assert hasattr(req_out, "multimodal_output") and req_out.multimodal_output
-
-    audio = req_out.multimodal_output.get("audio")
-    assert isinstance(audio, np.ndarray)
-    # audio shape: (batch, channels, samples)
-    assert audio.ndim == 3
-    assert audio.shape[0] == 1
-    assert audio.shape[1] == 2
-    assert audio.shape[2] > 0
-    expected_samples = int(seconds_total * sample_rate)
-    assert abs(audio.shape[2] - expected_samples) <= 2 * 1024

tests/e2e/offline_inference/test_ernie_image.py

@@ -1,58 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""
-End-to-end test for ErnieImage text-to-image generation.
-
-Equivalent to running:
-    vllm serve baidu/ERNIE-Image --omni
-"""
-
-import pytest
-from PIL import Image
-
-from tests.helpers.runtime import OmniRunner
-from vllm_omni.diffusion.data import DiffusionParallelConfig
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-MODEL = "baidu/ERNIE-Image"
-
-
-@pytest.mark.core_model
-@pytest.mark.diffusion
-def test_ernie_image_text_to_image():
-    with OmniRunner(
-        MODEL,
-        parallel_config=DiffusionParallelConfig(
-            tensor_parallel_size=2,
-        ),
-        enable_cpu_offload=True,
-    ) as runner:
-        omni_outputs = list(
-            runner.omni.generate(
-                prompts=["A photo of a cat sitting on a laptop"],
-                sampling_params_list=OmniDiffusionSamplingParams(
-                    height=512,
-                    width=512,
-                    num_inference_steps=2,
-                    guidance_scale=4.0,
-                    seed=42,
-                ),
-            )
-        )
-
-    assert len(omni_outputs) > 0
-    output = omni_outputs[0]
-    images = None
-    if output.images:
-        images = output.images
-    elif hasattr(output, "request_output") and output.request_output:
-        for stage_out in output.request_output:
-            if hasattr(stage_out, "images") and stage_out.images:
-                images = stage_out.images
-                break
-
-    assert images is not None
-    assert len(images) > 0
-    assert isinstance(images[0], Image.Image)
-    assert images[0].size == (512, 512)

tests/e2e/offline_inference/test_internvla_a1.py

@@ -1,90 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""Offline E2E test for InternVLA-A1 open-loop inference."""
-
-from __future__ import annotations
-
-import json
-import os
-import subprocess
-import sys
-import tempfile
-from pathlib import Path
-
-import pytest
-
-from tests.helpers.mark import hardware_test
-
-os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
-
-REPO_ROOT = Path(__file__).resolve().parents[3]
-EXAMPLE_SCRIPT = REPO_ROOT / "examples" / "offline_inference" / "internvla_a1" / "end2end.py"
-
-
-def _required_env(name: str) -> str:
-    value = os.getenv(name)
-    if not value:
-        pytest.skip(f"{name} is not set")
-    return value
-
-
-@pytest.mark.core_model
-@pytest.mark.advanced_model
-@pytest.mark.diffusion
-@hardware_test(res={"cuda": "H100", "rocm": "MI325"})
-def test_internvla_a1_offline_open_loop(run_level: str) -> None:
-    if run_level != "advanced_model":
-        pytest.skip("InternVLA-A1 offline evaluation requires real local checkpoints and dataset.")
-
-    model_dir = _required_env("INTERNVLA_A1_MODEL_DIR")
-    dataset_dir = _required_env("INTERNVLA_A1_DATASET_DIR")
-    processor_dir = _required_env("INTERNVLA_A1_PROCESSOR_DIR")
-    cosmos_dir = _required_env("INTERNVLA_A1_COSMOS_DIR")
-
-    env = os.environ.copy()
-    env["INTERNVLA_A1_PROCESSOR_DIR"] = processor_dir
-    env["INTERNVLA_A1_COSMOS_DIR"] = cosmos_dir
-
-    with tempfile.TemporaryDirectory(prefix="internvla_a1_e2e_") as tmpdir:
-        output_dir = Path(tmpdir) / "outputs"
-        cmd = [
-            sys.executable,
-            str(EXAMPLE_SCRIPT),
-            "--model-dir",
-            model_dir,
-            "--dataset-dir",
-            dataset_dir,
-            "--output-dir",
-            str(output_dir),
-            "--num-samples",
-            "1",
-            "--num-episodes",
-            "1",
-            "--dtype",
-            "float32",
-            "--attn-implementation",
-            "eager",
-            "--skip-plots",
-        ]
-        subprocess.run(cmd, check=True, env=env)
-
-        summary_path = output_dir / "summary.json"
-        assert summary_path.exists(), f"Missing summary file: {summary_path}"
-
-        with open(summary_path, encoding="utf-8") as f:
-            summary = json.load(f)
-
-        assert summary["mode"] == "registry"
-        assert summary["attn_implementation"] == "eager"
-        assert summary["dtype"] == "float32"
-        assert len(summary["results"]) == 1
-        assert "registry" in summary["eval_summaries"]
-
-        eval_summary = summary["eval_summaries"]["registry"]
-        assert eval_summary["num_episodes"] == 1
-        assert eval_summary["average_mse"] is not None
-        assert eval_summary["average_mae"] is not None
-
-        log_path = output_dir / "registry" / "log.json"
-        assert log_path.exists(), f"Missing evaluation log: {log_path}"

tests/e2e/offline_inference/test_moss_tts_nano.py

@@ -1,164 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""E2E offline inference tests for MOSS-TTS-Nano single-stage pipeline.
-
-Every request needs a reference audio clip. We test upstream's
-recommended ``voice_clone`` mode (no transcript needed). We fetch the
-upstream sample (assets/audio/zh_1.wav, ~50 KB) once per test session
-and reuse it for all cases.
-"""
-
-from __future__ import annotations
-
-import os
-import urllib.request
-
-import pytest
-import torch
-from vllm import SamplingParams
-
-from tests.helpers.mark import hardware_test
-from vllm_omni import Omni
-
-MODEL_NAME = "OpenMOSS-Team/MOSS-TTS-Nano"
-SAMPLE_RATE = 48000
-REF_AUDIO_URL = "https://raw.githubusercontent.com/OpenMOSS/MOSS-TTS-Nano/main/assets/audio/zh_1.wav"
-
-DEFAULT_SAMPLING = SamplingParams(
-    temperature=1.0,
-    top_p=1.0,
-    top_k=50,
-    max_tokens=4096,
-    seed=42,
-    detokenize=False,
-)
-
-
-@pytest.fixture(scope="session")
-def ref_audio_path(tmp_path_factory) -> str:
-    """Download the upstream reference clip once per test session.
-
-    Bounded with a 30 s timeout (urlretrieve has no timeout kwarg, so we use
-    urlopen + write ourselves). Fetch failures are hard failures, not skips,
-    so a broken network path can't silently mask regressions. Set
-    ``MOSS_TTS_NANO_SKIP_ON_NET_FAIL=1`` to opt into skipping for air-gapped
-    environments.
-    """
-    cache_dir = tmp_path_factory.mktemp("moss_tts_nano_ref")
-    target = cache_dir / "zh_1.wav"
-    try:
-        with urllib.request.urlopen(REF_AUDIO_URL, timeout=30) as resp:
-            data = resp.read()
-        target.write_bytes(data)
-    except Exception as e:
-        msg = f"Cannot fetch upstream reference clip {REF_AUDIO_URL}: {e}"
-        if os.environ.get("MOSS_TTS_NANO_SKIP_ON_NET_FAIL"):
-            pytest.skip(msg)
-        pytest.fail(msg)
-    if not target.exists() or os.path.getsize(target) == 0:
-        pytest.fail(f"Reference clip empty after download: {target}")
-    return str(target)
-
-
-def _build_request(
-    text: str,
-    prompt_audio_path: str,
-    prompt_text: str | None = None,
-    mode: str = "voice_clone",
-    max_new_frames: int = 100,  # short for tests
-    seed: int = 42,
-) -> dict:
-    """Build a MOSS-TTS-Nano offline request.
-
-    Upstream forbids ``prompt_text`` in ``voice_clone`` mode; only forward
-    it when explicitly supplied (and typically with ``mode='continuation'``).
-    """
-    additional: dict = {
-        "text": [text],
-        "mode": [mode],
-        "prompt_audio_path": [prompt_audio_path],
-        "max_new_frames": [max_new_frames],
-        "seed": [seed],
-    }
-    if prompt_text is not None:
-        additional["prompt_text"] = [prompt_text]
-    return {
-        "prompt": "<|im_start|>assistant\n",
-        "additional_information": additional,
-    }
-
-
-def _collect_audio(omni: Omni, request: dict) -> tuple[torch.Tensor, int]:
-    """Run a single request and return (waveform, sample_rate)."""
-    for stage_outputs in omni.generate(request, DEFAULT_SAMPLING):
-        for req_output in stage_outputs.request_output:
-            mm = req_output.outputs[0].multimodal_output
-            assert mm is not None, "Expected multimodal_output to be non-None"
-            audio = mm.get("audio")
-            sr = mm.get("sr")
-            assert audio is not None, "Expected 'audio' key in multimodal_output"
-            assert isinstance(audio, torch.Tensor), f"audio should be Tensor, got {type(audio)}"
-            return audio.cpu(), int(sr.item()) if sr is not None else SAMPLE_RATE
-    raise AssertionError("No stage outputs received")
-
-
-@pytest.fixture(scope="module")
-def omni_engine():
-    """Module-scoped Omni engine to avoid re-loading the model for each test."""
-    return Omni(model=MODEL_NAME, stage_init_timeout=180)
-
-
-@pytest.mark.omni
-@hardware_test(res={"cuda": "L4"})
-def test_moss_tts_nano_english(omni_engine, ref_audio_path):
-    """English voice_clone mode (no prompt_text) produces non-empty 48 kHz audio."""
-    req = _build_request("Hello, this is a short voice cloning demo for testing.", ref_audio_path)
-    audio, sr = _collect_audio(omni_engine, req)
-
-    assert sr == SAMPLE_RATE, f"Expected sample_rate={SAMPLE_RATE}, got {sr}"
-    assert audio.numel() > 0, "Audio tensor should not be empty"
-    assert not torch.all(audio == 0), "Audio should not be all-zeros (silence)"
-
-
-@pytest.mark.omni
-@hardware_test(res={"cuda": "L4"})
-def test_moss_tts_nano_chinese(omni_engine, ref_audio_path):
-    """Chinese voice_clone mode produces non-empty audio."""
-    req = _build_request("你好，这是语音合成测试。", ref_audio_path)
-    audio, sr = _collect_audio(omni_engine, req)
-
-    assert sr == SAMPLE_RATE
-    assert audio.numel() > 0
-    assert not torch.all(audio == 0)
-
-
-@pytest.mark.omni
-@hardware_test(res={"cuda": "L4"})
-def test_moss_tts_nano_deterministic(omni_engine, ref_audio_path):
-    """Same seed produces identical waveforms."""
-    req = _build_request("Reproducible output test.", ref_audio_path, seed=123)
-    audio1, _ = _collect_audio(omni_engine, req)
-    audio2, _ = _collect_audio(omni_engine, req)
-
-    assert audio1.shape == audio2.shape, "Waveform shapes should match with same seed"
-    assert torch.allclose(audio1, audio2, atol=1e-4), "Waveforms should match with same seed"
-
-
-@pytest.mark.omni
-@hardware_test(res={"cuda": "L4"})
-def test_moss_tts_nano_batch(omni_engine, ref_audio_path):
-    """Batch of two requests returns audio for each."""
-    requests = [
-        _build_request("First request.", ref_audio_path),
-        _build_request("Second request.", ref_audio_path),
-    ]
-    results = []
-    for stage_outputs in omni_engine.generate(requests, [DEFAULT_SAMPLING] * 2):
-        for req_output in stage_outputs.request_output:
-            mm = req_output.outputs[0].multimodal_output
-            assert mm is not None
-            results.append(mm["audio"].cpu())
-
-    assert len(results) == 2, f"Expected 2 outputs, got {len(results)}"
-    for i, audio in enumerate(results):
-        assert audio.numel() > 0, f"Audio {i} is empty"

tests/e2e/offline_inference/test_moss_tts_nano_expansion.py

@@ -32,7 +32,6 @@ _OMNI_RUNNER_PARAM = (
 )
 
 pytestmark = [
-    pytest.mark.skip(reason="issue#4361"),
     pytest.mark.full_model,
     pytest.mark.tts,
     pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True),

tests/e2e/online_serving/test_audiox_online.py

@@ -1,55 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-E2E online serving test for AudioX text-to-audio diffusion.
-
-Mirrors `tests/e2e/online_serving/test_sd3_expansion.py` for image diffusion: spin up
-`vllm-omni` with `--model-class-name AudioXPipeline` and validate via the standard
-`send_diffusion_request` helper (now audio-aware).
-"""
-
-import os
-
-import pytest
-
-from tests.helpers.mark import hardware_marks
-from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data
-
-# Tiny / random checkpoint usable in CI; override to a real bundle locally.
-AUDIOX_TEST_MODEL = os.environ.get("AUDIOX_TEST_MODEL", "zhangj1an/audiox_random")
-T2A_PROMPT = "A quiet living room with soft fabric rustle and gentle cat breathing."
-
-SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "L4"})
-
-
-def _audiox_server_cases(model: str):
-    return [
-        pytest.param(
-            OmniServerParams(
-                model=model,
-                server_args=["--model-class-name", "AudioXPipeline"],
-            ),
-            id="t2a",
-            marks=SINGLE_CARD_FEATURE_MARKS,
-        ),
-    ]
-
-
-@pytest.mark.core_model
-@pytest.mark.diffusion
-@pytest.mark.parametrize("omni_server", _audiox_server_cases(AUDIOX_TEST_MODEL), indirect=True)
-def test_audiox_t2a_online(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
-    """AudioX text-to-audio: chat completion returns a non-empty WAV in `message.audio.data`."""
-    request_config = {
-        "model": omni_server.model,
-        "messages": dummy_messages_from_mix_data(content_text=T2A_PROMPT),
-        "extra_body": {
-            "num_inference_steps": 4,
-            "guidance_scale": 6.0,
-            "seed": 42,
-            "audiox_task": "t2a",
-            "seconds_start": 0.0,
-            "seconds_total": 2.0,
-        },
-    }
-    openai_client.send_diffusion_request(request_config)

tests/e2e/online_serving/test_ming_flash_omni_expansion.py

@@ -26,12 +26,6 @@ _SKIP_NEED_4_H100_NOT_CI = pytest.mark.skip(
     reason="Requires 4x H100 GPUs; skipped in CI for now.",
 )
 
-pytestmark = [pytest.mark.omni, pytest.mark.full_model]
-
-_SKIP_NEED_4_H100_NOT_CI = pytest.mark.skip(
-    reason="Requires 4x H100 GPUs; skipped in CI for now.",
-)
-
 models = ["Jonathan1909/Ming-flash-omni-2.0"]
 
 # Use thinker-only topology to test text outputs

tests/e2e/online_serving/test_moss_tts_nano.py

@@ -1,185 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-E2E Online tests for MOSS-TTS-Nano model with text input and audio output.
-
-These tests verify the /v1/audio/speech endpoint works correctly with
-actual model inference, not mocks. The server uses upstream's
-``voice_clone`` mode (no transcript needed); ``ref_text`` and ``voice``
-are accepted in the OpenAI schema but ignored. One case includes a
-non-empty ``ref_text`` to verify it does not break the pipeline. The
-reference clip is fetched from the upstream repo (assets/audio/zh_1.wav,
-~50 KB).
-"""
-
-import base64
-import os
-import urllib.request
-
-os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
-os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
-
-import pytest
-
-from tests.helpers.mark import hardware_test
-from tests.helpers.runtime import OmniServerParams
-from tests.helpers.stage_config import get_deploy_config_path
-
-MODEL = "OpenMOSS-Team/MOSS-TTS-Nano"
-REF_AUDIO_URL = "https://raw.githubusercontent.com/OpenMOSS/MOSS-TTS-Nano/main/assets/audio/zh_1.wav"
-REF_AUDIO_TRANSCRIPT = "欢迎关注模思智能、上海创智学院与复旦大学自然语言处理实验室。"
-
-
-@pytest.fixture(scope="session")
-def ref_audio_data_url() -> str:
-    """Fetch the upstream sample clip and return it as a base64 data URL.
-
-    The fetch failure is escalated to a hard failure (not pytest.skip) so that
-    a broken network path does not silently mask regressions in
-    /v1/audio/speech. Set ``MOSS_TTS_NANO_SKIP_ON_NET_FAIL=1`` to opt into
-    skipping in air-gapped environments.
-    """
-    try:
-        with urllib.request.urlopen(REF_AUDIO_URL, timeout=30) as resp:
-            data = resp.read()
-    except Exception as e:
-        msg = f"Cannot fetch upstream reference clip {REF_AUDIO_URL}: {e}"
-        if os.environ.get("MOSS_TTS_NANO_SKIP_ON_NET_FAIL"):
-            pytest.skip(msg)
-        pytest.fail(msg)
-    if not data:
-        pytest.fail(f"Reference clip empty: {REF_AUDIO_URL}")
-    return f"data:audio/wav;base64,{base64.b64encode(data).decode('ascii')}"
-
-
-def get_prompt(prompt_type="text"):
-    """Text prompt for text-to-audio tests.
-
-    Avoid the model's own name ("MOSS-TTS-Nano") in the test input — the
-    codec consistently mispronounces it ("MOS's TTS Nano", "Was TTS Nano",
-    etc.), which trips the transcript-similarity assertion in CI without
-    indicating a real regression. Use plain natural-sounding sentences.
-    """
-    prompts = {
-        "text": "Hello, this is a short voice cloning demo for testing.",
-        "chinese": "你好，这是一段简单的语音合成测试。",
-    }
-    return prompts.get(prompt_type, prompts["text"])
-
-
-tts_server_params = [
-    pytest.param(
-        OmniServerParams(
-            model=MODEL,
-            stage_config_path=get_deploy_config_path("moss_tts_nano.yaml"),
-            server_args=["--disable-log-stats"],
-        ),
-        id="moss_tts_nano",
-    )
-]
-
-
-@pytest.mark.advanced_model
-@pytest.mark.omni
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
-@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
-def test_text_to_audio_001(omni_server, openai_client, ref_audio_data_url) -> None:
-    """
-    Test voice_clone mode (no ref_text) via /v1/audio/speech.
-    Deploy Setting: default yaml
-    Input Modal: text + reference audio (no transcript)
-    Output Modal: audio (48 kHz, WAV)
-    Input Setting: stream=False
-    Datasets: single request
-    """
-    request_config = {
-        "model": omni_server.model,
-        "input": get_prompt(),
-        "stream": False,
-        "response_format": "wav",
-        "ref_audio": ref_audio_data_url,
-    }
-
-    openai_client.send_audio_speech_request(request_config)
-
-
-@pytest.mark.advanced_model
-@pytest.mark.omni
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
-@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
-def test_text_to_audio_002(omni_server, openai_client, ref_audio_data_url) -> None:
-    """
-    Test streaming voice_clone via /v1/audio/speech.
-    Deploy Setting: default yaml
-    Input Modal: text + reference audio (no transcript)
-    Output Modal: audio (48 kHz, PCM stream)
-    Input Setting: stream=True
-    Datasets: single request
-
-    NOTE: ``min_hnr_db=-5.0`` overrides the conftest default of 1.0 dB.
-    MOSS-TTS-Nano's voice_clone output is intrinsically noisy by that
-    metric — verified on H20, both streaming and non-streaming clips
-    measure HNR around -2 dB even when the audio sounds correct. We
-    keep the PCM-stream check in place to catch catastrophic decoder
-    failures (which produce HNR much further below 0), just with a
-    threshold MOSS can clear in normal operation.
-    """
-    request_config = {
-        "model": omni_server.model,
-        "input": get_prompt(),
-        "stream": True,
-        "response_format": "pcm",
-        "ref_audio": ref_audio_data_url,
-        "min_hnr_db": -5.0,
-    }
-
-    openai_client.send_audio_speech_request(request_config)
-
-
-@pytest.mark.advanced_model
-@pytest.mark.omni
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
-@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
-def test_text_to_audio_003(omni_server, openai_client, ref_audio_data_url) -> None:
-    """
-    Test Chinese voice_clone via /v1/audio/speech.
-    Deploy Setting: default yaml
-    Input Modal: text (Chinese) + reference audio (no transcript)
-    Output Modal: audio (48 kHz, WAV)
-    Input Setting: stream=False
-    Datasets: single request
-    """
-    request_config = {
-        "model": omni_server.model,
-        "input": get_prompt("chinese"),
-        "stream": False,
-        "response_format": "wav",
-        "ref_audio": ref_audio_data_url,
-    }
-
-    openai_client.send_audio_speech_request(request_config)
-
-
-@pytest.mark.advanced_model
-@pytest.mark.omni
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
-@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
-def test_text_to_audio_004_ref_text_ignored(omni_server, openai_client, ref_audio_data_url) -> None:
-    """
-    Sending ``ref_text`` alongside ``ref_audio`` is accepted but ignored.
-
-    The OpenAI ``ref_text`` field is part of the schema, so clients may
-    send it. The server uses upstream's voice_clone mode regardless and
-    drops the transcript; this case verifies the request still completes
-    and produces audio.
-    """
-    request_config = {
-        "model": omni_server.model,
-        "input": get_prompt("chinese"),
-        "stream": False,
-        "response_format": "wav",
-        "ref_audio": ref_audio_data_url,
-        "ref_text": REF_AUDIO_TRANSCRIPT,
-    }
-
-    openai_client.send_audio_speech_request(request_config)

tests/e2e/online_serving/test_qwen3_omni_realtime_websocket.py

@@ -1,206 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-E2E online tests for Qwen3-Omni /v1/realtime WebSocket (streaming PCM in, audio out).
-"""
-
-from __future__ import annotations
-
-import asyncio
-import base64
-import io
-import json
-import os
-import wave
-
-import pytest
-import websockets
-
-from tests.helpers.mark import hardware_test
-from tests.helpers.media import (
-    convert_audio_bytes_to_text,
-    cosine_similarity_text,
-    generate_synthetic_audio,
-)
-from tests.helpers.runtime import OmniServerParams
-from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
-
-os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
-
-MODEL = "Qwen/Qwen3-Omni-30B-A3B-Instruct"
-
-# The new-schema CI overlay bakes in async_chunk: False and covers CUDA/ROCm/XPU
-# via its ``platforms:`` section, so one path serves all three.
-default_stage_config = get_deploy_config_path("ci/qwen3_omni_moe.yaml")
-
-
-def _realtime_stage_config_path() -> str:
-    """CI omni layout without async_chunk; stage 0 thinker max_tokens=10."""
-    return modify_stage_config(
-        default_stage_config,
-        updates={"stages": {0: {"default_sampling_params.max_tokens": 10}}},
-    )
-
-
-realtime_server_params = [
-    pytest.param(
-        OmniServerParams(
-            model=MODEL,
-            stage_config_path=_realtime_stage_config_path(),
-            use_stage_cli=True,
-        ),
-        id="thinker_max_tokens_10",
-    ),
-]
-
-
-def _pcm16_mono_16k_from_wav_bytes(wav_bytes: bytes) -> bytes:
-    with wave.open(io.BytesIO(wav_bytes), "rb") as wf:
-        if wf.getnchannels() != 1:
-            raise ValueError(f"Expected mono WAV, got {wf.getnchannels()} channels")
-        if wf.getsampwidth() != 2:
-            raise ValueError(f"Expected 16-bit PCM, sampwidth={wf.getsampwidth()}")
-        if wf.getframerate() != 16000:
-            raise ValueError(f"Expected 16 kHz input for /v1/realtime, got {wf.getframerate()} Hz")
-        if wf.getcomptype() != "NONE":
-            raise ValueError(f"Expected uncompressed PCM, comptype={wf.getcomptype()!r}")
-        return wf.readframes(wf.getnframes())
-
-
-def _wav_bytes_from_pcm16(pcm: bytes, sample_rate_hz: int) -> bytes:
-    buf = io.BytesIO()
-    with wave.open(buf, "wb") as wf:
-        wf.setnchannels(1)
-        wf.setsampwidth(2)
-        wf.setframerate(sample_rate_hz)
-        wf.writeframes(pcm)
-    return buf.getvalue()
-
-
-async def _run_realtime_audio_roundtrip(
-    host: str,
-    port: int,
-    model: str,
-    pcm16: bytes,
-    *,
-    chunk_ms: int = 100,
-) -> dict:
-    uri = f"ws://{host}:{port}/v1/realtime"
-    incremental: list[bytes] = []
-    output_sr = 24000
-    text_chunks: list[str] = []
-    final_text = ""
-    delta_events = 0
-
-    bytes_per_ms = 16000 * 2 // 1000
-    chunk_bytes = max(bytes_per_ms * chunk_ms, 2)
-
-    async with websockets.connect(uri, max_size=64 * 1024 * 1024) as ws:
-        await ws.send(json.dumps({"type": "session.update", "model": model}))
-        await ws.send(json.dumps({"type": "input_audio_buffer.commit", "final": False}))
-
-        for i in range(0, len(pcm16), chunk_bytes):
-            chunk = pcm16[i : i + chunk_bytes]
-            await ws.send(
-                json.dumps(
-                    {
-                        "type": "input_audio_buffer.append",
-                        "audio": base64.b64encode(chunk).decode("utf-8"),
-                    }
-                )
-            )
-
-        await ws.send(json.dumps({"type": "input_audio_buffer.commit", "final": True}))
-
-        while True:
-            message = await asyncio.wait_for(ws.recv(), timeout=600)
-            if isinstance(message, bytes):
-                continue
-
-            event = json.loads(message)
-            event_type = event.get("type")
-
-            if event_type == "session.created":
-                continue
-
-            if event_type == "response.audio.delta":
-                delta_events += 1
-                sr = event.get("sample_rate_hz")
-                if isinstance(sr, int) and sr > 0:
-                    output_sr = sr
-                audio_b64 = event.get("audio", "")
-                if audio_b64:
-                    incremental.append(base64.b64decode(audio_b64))
-                continue
-
-            if event_type == "transcription.delta":
-                d = event.get("delta", "")
-                if d:
-                    text_chunks.append(d)
-                continue
-
-            if event_type == "transcription.done":
-                final_text = event.get("text", "") or "".join(text_chunks)
-                continue
-
-            if event_type == "response.audio.done":
-                break
-
-            if event_type == "error":
-                raise AssertionError(f"WebSocket error: {event}")
-
-            raise AssertionError(f"Unexpected WebSocket event: {event}")
-
-    out_pcm = b"".join(incremental)
-    return {
-        "output_pcm": out_pcm,
-        "output_sample_rate": output_sr,
-        "transcription_text": final_text if final_text else "".join(text_chunks),
-        "delta_events": delta_events,
-    }
-
-
-class TestQwen3OmniRealtimeWebSocket:
-    @pytest.mark.advanced_model
-    @pytest.mark.omni
-    @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-    @pytest.mark.parametrize("omni_server", realtime_server_params, indirect=True)
-    def test_streaming_audio_input_pcm_output(self, omni_server) -> None:
-        """
-        Short streamed 16 kHz mono PCM16 input; expect streamed PCM16 audio deltas and
-        transcription. Verify Whisper(output audio) aligns with model text (same idea
-        as multimodal omni e2e).
-        """
-        syn = generate_synthetic_audio(10, 1, sample_rate=16000)
-        wav_bytes = base64.b64decode(syn["base64"])
-        pcm16 = _pcm16_mono_16k_from_wav_bytes(wav_bytes)
-
-        result = asyncio.run(
-            _run_realtime_audio_roundtrip(
-                omni_server.host,
-                omni_server.port,
-                omni_server.model,
-                pcm16,
-                chunk_ms=100,
-            )
-        )
-
-        out_pcm = result["output_pcm"]
-        assert result["delta_events"] >= 1
-        assert out_pcm, "No output PCM from response.audio.delta"
-        assert len(out_pcm) % 2 == 0
-        assert len(out_pcm) >= 4096, "Output audio unexpectedly small"
-        assert result["output_sample_rate"] > 0
-
-        final_text = (result["transcription_text"] or "").strip()
-        assert final_text, "Expected non-empty transcription (model text stream)"
-
-        wav_out = _wav_bytes_from_pcm16(out_pcm, result["output_sample_rate"])
-        whisper_text = convert_audio_bytes_to_text(wav_out).strip()
-        assert whisper_text, "Whisper returned empty string for synthesized output audio"
-
-        sim = cosine_similarity_text(whisper_text.lower(), final_text.lower())
-        assert sim > 0.9, (
-            f"Output audio transcript should match model text (sim={sim:.3f}): "
-            f"whisper={whisper_text!r}, model_text={final_text!r}"
-        )