fix: harden omni weight snapshot downloads

Serialize full-repo Hugging Face snapshot materialization with a repo-wide
node lock and retry incomplete snapshots so CI does not race on shared caches.

Co-authored-by: Cursor <cursoragent@cursor.com>

.buildkite/pipeline-intel.yaml

@@ -10,7 +10,7 @@ steps:
           DOCKER_BUILDKIT: "1"
           # Buildkite will automatically replace this with the actual commit hash
           VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
-          VLLM_VERSION: "v0.21.0"
+          VLLM_VERSION: "v0.22.0"
         priority: 100
         timeout_in_minutes: 60
         soft_fail: false

.buildkite/test-amd-merge.yml

@@ -1,15 +1,28 @@
 steps:
 
-- label: "Simple Unit Test"
+- label: "Simple · Diffusion & Model Executor Test"
   agent_pool: mi325_1
   depends_on: amd-build
+  timeout_in_minutes: 50
   mirror_hardwares: [amdproduction]
   grade: Blocking
   commands:
     - export VLLM_ROCM_USE_AITER=0
     # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
     # that is not supported on CPU
-    - "timeout 60m pytest -v -s -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+    - "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+
+- label: "Simple · Other Test"
+  agent_pool: mi325_1
+  depends_on: amd-build
+  timeout_in_minutes: 50
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export VLLM_ROCM_USE_AITER=0
+    # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
+    # that is not supported on CPU
+    - "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
 
 - label: "Diffusion Images API LoRA E2E"
   agent_pool: mi325_1
@@ -49,16 +62,6 @@ steps:
 #     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
 #     - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_expansion.py -m "advanced_model and diffusion and L4" --run-level advanced_model
 
-- label: "Diffusion Cache Backend Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - export VLLM_LOGGING_LEVEL=DEBUG
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - timeout 15m pytest -s -v -m "core_model and cache and diffusion and not distributed_cuda and L4"
-
 - label: "Diffusion Sequence Parallelism Test (Need 4 GPUs)"
   agent_pool: mi325_4
   depends_on: amd-build
@@ -123,7 +126,7 @@ steps:
   grade: Blocking
   commands:
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py tests/e2e/online_serving/test_mimo_audio.py -m "advanced_model" --run-level "advanced_model"
+    - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py -m 'advanced_model and cuda' --run-level 'advanced_model'
 
 - label: "Qwen3-TTS CustomVoice E2E Test"
   agent_pool: mi325_1

.buildkite/test-amd-ready.yaml

@@ -1,6 +1,6 @@
 steps:
 
-- label: "Simple Unit Test"
+- label: "Simple · Diffusion & Model Executor Test"
   agent_pool: mi325_1
   depends_on: amd-build
   mirror_hardwares: [amdproduction]
@@ -9,7 +9,18 @@ steps:
     - export VLLM_ROCM_USE_AITER=0
     # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
     # that is not supported on CPU
-    - "timeout 60m pytest -v -s -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+    - "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+
+- label: "Simple · Other Test"
+  agent_pool: mi325_1
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export VLLM_ROCM_USE_AITER=0
+    # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
+    # that is not supported on CPU
+    - "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
 
 - label: "Voxtral TTS CUDA Unit Test"
   agent_pool: mi325_1
@@ -19,13 +30,13 @@ steps:
   commands:
     - timeout 10m pytest -s -v tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py
 
-- label: "Diffusion Batching Test"
+- label: "Diffusion · Batch Test"
   agent_pool: mi325_1
   depends_on: amd-build
   mirror_hardwares: [amdproduction]
   grade: Blocking
   commands:
-    - timeout 20m pytest -s -v tests/diffusion/batching/test_diffusion_batching.py -m "core_model and diffusion" --run-level "core_model"
+    - timeout 20m pytest -sv tests/diffusion/batching -m 'core_model and cuda' --run-level "core_model"
 
 - label: "Custom Pipeline Test"
   agent_pool: mi325_1

.buildkite/test-merge.yml

@@ -389,6 +389,29 @@ steps:
               volumes:
                 - "/fsx/hf_cache:/fsx/hf_cache"
 
+      - label: "TTS · VoxCPM2 Online Test"
+        timeout_in_minutes: 30
+        commands:
+          - |
+            timeout 30m bash -c "
+              export VLLM_LOGGING_LEVEL=DEBUG
+              export VLLM_WORKER_MULTIPROC_METHOD=spawn
+              pytest -s -v tests/e2e/online_serving/test_voxcpm2_tts.py -m 'advanced_model and cuda' --run-level 'advanced_model'
+            "
+        agents:
+          queue: "gpu_1_queue"
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              always-pull: true
+              propagate-environment: true
+              shm-size: "8gb"
+              environment:
+                - "HF_HOME=/fsx/hf_cache"
+                - "HF_TOKEN"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"
+
       - label: "Omni · Higgs-Audio V2 Test"
         timeout_in_minutes: 30
         commands:

.buildkite/test-nightly.yml

@@ -398,7 +398,7 @@ steps:
         timeout_in_minutes: 90
         commands:
           - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
-          - pytest -s -v tests/e2e/ -m "full_model and L4 and tts" --run-level "full_model" --ignore=tests/e2e/accuracy
+          - pytest -s -v tests/e2e/ -m "full_model and L4 and tts" --run-level "full_model" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_voxcpm2_tts_expansion.py
         agents:
           queue: "gpu_1_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
         plugins:
@@ -413,6 +413,30 @@ steps:
               volumes:
                 - "/fsx/hf_cache:/fsx/hf_cache"
 
+      - label: ":full_moon: TTS · Accuracy Test · VoxCPM2"
+        timeout_in_minutes: 30
+        commands:
+          - |
+            timeout 30m bash -c "
+              export VLLM_LOGGING_LEVEL=DEBUG
+              export VLLM_WORKER_MULTIPROC_METHOD=spawn
+              export VLLM_ALLOW_LONG_MAX_MODEL_LEN=1
+              pytest -s -v tests/e2e/online_serving/test_voxcpm2_tts_expansion.py -m 'full_model and cuda and tts' --run-level 'full_model'
+            "
+        agents:
+          queue: "gpu_1_queue"
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              always-pull: true
+              propagate-environment: true
+              shm-size: "8gb"
+              environment:
+                - "HF_HOME=/fsx/hf_cache"
+                - "HF_TOKEN"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"
+
       - label: ":full_moon: TTS · Perf Test"
         key: nightly-tts-performance
         timeout_in_minutes: 180
@@ -459,8 +483,8 @@ steps:
                       path: /mnt/hf-cache
                       type: DirectoryOrCreate
 
-      - label: ":full_moon: VoxCPM2 · Perf Test"
-        key: nightly-voxcpm2-performance
+      - label: ":full_moon: TTS · VoxCPM2 · Perf Test"
+        key: nightly-tts-performance-voxcpm2
         timeout_in_minutes: 180
         commands:
           - export BENCHMARK_DIR=tests/dfx/perf/results
@@ -695,6 +719,48 @@ steps:
                       path: /mnt/hf-cache
                       type: DirectoryOrCreate
 
+      - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · Multi-GPU DreamZero"
+        timeout_in_minutes: 120
+        commands:
+          - >-
+            pytest -sv
+            tests/e2e/online_serving/test_dreamzero_expansion.py
+            -m "full_model and diffusion and H100 and distributed_cuda"
+            --run-level "full_model"
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 2
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                      - name: HF_TOKEN
+                        valueFrom:
+                          secretKeyRef:
+                            name: hf-token-secret
+                            key: token
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
+                  - name: devshm
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate
+
 
       - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with L4"
         timeout_in_minutes: 60

.buildkite/test-ready.yml

@@ -296,6 +296,29 @@ steps:
               volumes:
                 - "/fsx/hf_cache:/fsx/hf_cache"
 
+      - label: "TTS · VoxCPM2 Online Test"
+        timeout_in_minutes: 30
+        commands:
+          - |
+            timeout 30m bash -c "
+              export VLLM_LOGGING_LEVEL=DEBUG
+              export VLLM_WORKER_MULTIPROC_METHOD=spawn
+              pytest -s -v tests/e2e/online_serving/test_voxcpm2_tts.py -m 'core_model and tts' --run-level 'core_model'
+            "
+        agents:
+          queue: "gpu_1_queue"
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              always-pull: true
+              propagate-environment: true
+              shm-size: "8gb"
+              environment:
+                - "HF_HOME=/fsx/hf_cache"
+                - "HF_TOKEN"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"
+
       - label: "TTS · Qwen3-TTS CustomVoice Test"
         commands:
           - |

.buildkite/test-template-amd-omni.j2

@@ -53,6 +53,9 @@
 {% endif %}
             {{ indented_cmd | safe }}
         priority: 100
+        {% if step.timeout_in_minutes %}
+        timeout_in_minutes: {{ step.timeout_in_minutes }}
+        {% endif %}
         {% if step.grade and step.grade == "Blocking" %}
         soft_fail: false
         {% else %}

.gitignore

@@ -270,3 +270,6 @@ vllm_omni/_version.py
 .worktrees/
 # CI overlay yamls materialized from tests/utils.py:_CI_OVERLAYS at test time
 tests/.ci_generated/
+
+# Local Lance precision-matching dump scripts (not part of the repo)
+lance_compare/

benchmarks/build_dataset/sound_effect_smoke/en/meta.lst

@@ -0,0 +1,10 @@
+se001|ocean waves crashing on a rocky beach|10.0
+se002|heavy rain on a metal roof at night|8.0
+se003|a dog barking in the distance|6.0
+se004|a crackling campfire in a quiet forest|10.0
+se005|footsteps on dry leaves|6.0
+se006|thunder rolling across a wide valley|8.0
+se007|a busy coffee shop with chatter and espresso machines|10.0
+se008|wind howling through tall pine trees|8.0
+se009|a small stream flowing over pebbles|10.0
+se010|a creaking wooden door slowly opening|6.0

benchmarks/build_dataset/ttsd_smoke/en/meta.lst

@@ -0,0 +1,10 @@
+ttsd001|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Hello and welcome to the show today. [S2] Thanks for having me, it's a pleasure to be here. [S1] Let's dive right in.
+ttsd002|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Did you finish the report? [S2] Almost — I'm just polishing the conclusion. [S1] Great, send it over when you're done.
+ttsd003|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Have you tried that new restaurant downtown? [S2] Yes, the pasta is amazing. [S1] I'll book a table for Friday.
+ttsd004|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] What did you think of yesterday's meeting? [S2] Productive, but a bit long. [S1] Agreed, we should set an hour cap next time.
+ttsd005|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] The weather is beautiful today. [S2] Perfect for a hike. [S1] Want to drive out to the trails this afternoon?
+ttsd006|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] How did the demo go? [S2] Better than expected — the client signed off on the prototype. [S1] Fantastic news!
+ttsd007|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Coffee or tea this morning? [S2] Coffee please, with a splash of milk. [S1] Coming right up.
+ttsd008|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Did you read the email about the deadline change? [S2] I did, we now have until next Friday. [S1] That should give us enough buffer.
+ttsd009|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Are you joining the team dinner tonight? [S2] I'll be there around seven. [S1] See you then.
+ttsd010|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Have we backed up the database? [S2] Last backup ran two hours ago. [S1] Good, I'll trigger a fresh snapshot before the migration.

benchmarks/tts/bench_tts.py

@@ -53,6 +53,8 @@ _TASK_TO_DATASET: dict[str, str] = {
     "voice_clone": "seed-tts",
     "default_voice": "seed-tts-text",
     "voice_design": "seed-tts-design",
+    "dialogue": "ttsd",
+    "sound_effect": "sound-effect",
 }
 
 # Default design dataset path (bundled with the repo)
@@ -81,6 +83,7 @@ def build_bench_args(
     output_dir: str | None,
     result_filename: str | None,
     extra_cli_args: list[str],
+    output_len: int | None = None,
 ) -> list[str]:
     """Build the ``vllm bench serve --omni`` command for one (task, concurrency) run."""
     dataset_name = _TASK_TO_DATASET[task]
@@ -139,6 +142,9 @@ def build_bench_args(
     if wer_eval:
         cmd.append("--seed-tts-wer-eval")
 
+    if output_len is not None:
+        cmd += ["--hf-output-len", str(output_len)]
+
     if output_dir or result_filename:
         out_dir = output_dir or "."
         os.makedirs(out_dir, exist_ok=True)
@@ -238,6 +244,15 @@ def main() -> None:
     parser.add_argument("--host", default="localhost")
     parser.add_argument("--port", type=int, default=8000)
     parser.add_argument("--model-configs", default=str(_DEFAULT_MODEL_CONFIGS), help="Path to model_configs.yaml")
+    parser.add_argument(
+        "--output-len",
+        type=int,
+        default=None,
+        help="Override per-request max_new_tokens forwarded to the server "
+        "(seed-tts datasets default to 2048; the MOSS-TTS full talker takes "
+        "this as 'max_new_frames' so smaller values cap audio length per "
+        "request and keep bench runtime tractable).",
+    )
     parser.add_argument("extra", nargs=argparse.REMAINDER, help="Extra args passed directly to vllm bench serve")
     args = parser.parse_args()
 
@@ -292,6 +307,7 @@ def main() -> None:
                 output_dir=args.output_dir,
                 result_filename=result_filename,
                 extra_cli_args=args.extra or [],
+                output_len=args.output_len,
             )
             result = run_one_benchmark(cmd)
             if result is not None:

benchmarks/tts/model_configs.yaml

@@ -43,3 +43,61 @@ models:
       voice_clone: {}
       default_voice:
         voice: default
+
+  # MOSS-TTS family. After the serving_speech.py variant-detection refactor
+  # all five variants now go through OpenAI /v1/audio/speech with the
+  # variant-specific request fields (ambient_sound, ref_audio_2, instructions).
+  OpenMOSS-Team/MOSS-TTS:
+    supported_tasks: [voice_clone]
+    backend: openai-audio-speech
+    endpoint: /v1/audio/speech
+    task_extra_body:
+      voice_clone: {}
+
+  # MOSS-TTS-v1.5: continued-training upgrade of MOSS-TTS 1.0 — same
+  # MossTTSDelay-8B architecture and API, so identical wiring. Adds 31-language
+  # synthesis via the `language` field and inline `[pause Xs]` markers.
+  OpenMOSS-Team/MOSS-TTS-v1.5:
+    supported_tasks: [voice_clone]
+    backend: openai-audio-speech
+    endpoint: /v1/audio/speech
+    task_extra_body:
+      voice_clone: {}
+
+  OpenMOSS-Team/MOSS-TTS-Realtime:
+    supported_tasks: [voice_clone]
+    backend: openai-audio-speech
+    endpoint: /v1/audio/speech
+    task_extra_body:
+      voice_clone: {}
+
+  # MOSS-TTSD synthesizes [S1]/[S2] dialogue. The seed-tts voice_clone dataset
+  # has single-speaker rows; bench piggybacks on it (both speakers share the
+  # same ref_audio). For real dialogue evaluation use the ttsd dataset.
+  OpenMOSS-Team/MOSS-TTSD-v1.0:
+    supported_tasks: [voice_clone, dialogue]
+    backend: openai-audio-speech
+    endpoint: /v1/audio/speech
+    task_extra_body:
+      voice_clone: {}
+      dialogue: {}
+
+  # MOSS-SoundEffect ignores text/ref_audio; expects ambient_sound +
+  # duration_seconds. The sound_effect dataset module emits those.
+  OpenMOSS-Team/MOSS-SoundEffect:
+    supported_tasks: [sound_effect]
+    backend: openai-audio-speech
+    endpoint: /v1/audio/speech
+    task_extra_body:
+      sound_effect:
+        duration_seconds: 10
+
+  # MOSS-VoiceGenerator: instructions-driven (no ref_audio). seed-tts-design
+  # dataset already emits {instructions, target_text}, so voice_design works
+  # out of the box once the serving validator accepts instructions.
+  OpenMOSS-Team/MOSS-VoiceGenerator:
+    supported_tasks: [voice_design]
+    backend: openai-audio-speech
+    endpoint: /v1/audio/speech
+    task_extra_body:
+      voice_design: {}

docker/Dockerfile.cuda

@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=vllm/vllm-openai:v0.21.0
+ARG BASE_IMAGE=vllm/vllm-openai:v0.22.0
 FROM ${BASE_IMAGE}
 
 ARG COMMON_WORKDIR=/app

docker/Dockerfile.rocm

@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.21.0
+ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.22.0
 FROM ${BASE_IMAGE} AS base
 
 # Declare a variable to know if we want to use the nightly build or the stable build.

docker/Dockerfile.xpu

@@ -88,7 +88,7 @@ ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
 ENV UV_INDEX_STRATEGY="unsafe-best-match"
 ENV UV_LINK_MODE="copy"
 
-ARG VLLM_VERSION=v0.21.0
+ARG VLLM_VERSION=v0.22.0
 RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm
 WORKDIR /workspace/vllm
 

docs/design/index.md

@@ -13,6 +13,10 @@ This section contains design documents and architecture specifications for vLLM-
 - [Adding Step Execution Support for Diffusion Pipelines](feature/diffusion_step_execution.md)
 - [Continuous Batching for Step-Wise Diffusion](feature/diffusion_continuous_batching.md)
 
+## Infrastructure Design Documents
+
+- [Prometheus Metrics](metrics.md)
+
 ## Module Design Documents
 
 - [AR Module](module/ar_module.md)

docs/design/metrics.md

@@ -0,0 +1,208 @@
+# Prometheus Metrics Design
+
+This document describes how vLLM-Omni exposes Prometheus metrics for multi-stage pipelines, the constraints that shaped the design, and how the pipeline-level metrics coexist with upstream vLLM per-engine metrics. This iteration covers the text and audio surface; diffusion / image / video families land in a follow-up.
+
+## Objectives
+
+- Expose pipeline-level request and latency metrics that span the full multi-stage execution (orchestrator scope).
+- Preserve all upstream vLLM per-engine metrics (`vllm:*`) for stages backed by an AR LLM engine, and reshape their `engine` label into `stage` + `replica` so multi-replica deployments gain per-replica visibility automatically.
+- Expose per-modality SLO metrics that the upstream `vllm:*` families do not capture — audio TTFP / RTF / duration / frames / streaming continuity / silent-loss.
+- Expose per-replica-edge cross-stage transfer metrics so the slack between E2E latency and the sum of per-stage `gen_time` (queueing, serialization, network) becomes attributable.
+- Keep the metrics collection overhead low enough that it does not regress TTFA or throughput.
+
+## Background
+
+### Upstream vLLM Metrics
+
+Upstream vLLM defines ~37 Prometheus metric families under the `vllm:` prefix. These are registered by `PrometheusStatLogger` and cover engine-level state: KV cache usage, running/waiting request counts, token throughput, TTFT, inter-token latency, e2e latency, and so on. They are served via the `/metrics` HTTP endpoint provided by `prometheus_fastapi_instrumentator` and the default `prometheus_client` WSGI handler.
+
+vLLM's `unregister_vllm_metrics()` function strips every `prometheus_client` collector whose `_name` attribute contains the substring `"vllm"`. This runs during engine initialization to clean up stale collectors from prior instantiations within the same process.
+
+### The Problem
+
+vLLM-Omni runs multiple engine instances (stages × replicas) within a single process, coordinated by an Orchestrator. The pipeline needs its own metrics — aggregate request counts, end-to-end latency across all stages, per-modality SLO signals, and cross-stage transfer attribution — that do not exist in upstream vLLM. All pipeline-level metrics use the `vllm:omni_` prefix to distinguish them from upstream per-engine metrics. At import time (see `vllm_omni/patch.py`) we replace `unregister_vllm_metrics()` with a scoped version that still strips upstream `vllm:*` collectors before each new `PrometheusStatLogger` registers (so multi-engine processes don't crash on duplicate registration), but preserves anything prefixed `vllm:omni_` / `vllm_omni`.
+
+Upstream per-engine metrics retain the `vllm:` prefix but are now registered by `OmniPrometheusStatLogger`, a thin subclass of upstream's `PrometheusStatLogger` that reshapes the single `engine` label into a `stage` + `replica` pair (see "OmniPrometheusStatLogger wrap" below).
+
+## Architecture
+
+### Component Overview
+
+```text
+                       +------------------------+
+                       |  API Server (FastAPI)  |
+                       |   GET /metrics         |
+                       +-----------+------------+
+                                   |
+                  prometheus_client default registry
+                                   |
+        +--------+--------+--------+--------+--------+
+        |                                            |
+   vllm:omni_*                                    vllm:*
+   collectors                                  collectors
+        |                                            |
+   +----+--------+   +-----------+   +----------+   +-----------+
+   | OmniPromet- |   | OmniMod-  |   | OmniTra- |   | OmniProm- |
+   | heusMetrics |   | alityMet- |   | nsferMe- |   | etheusSt- |
+   |             |   | rics      |   | trics    |   | atLogger  |
+   +----+--------+   +-----+-----+   +----+-----+   +----+------+
+        |                  |              |              |
+     OmniBase           OmniBase     Orchestrator    Orchestrator
+   (request life-     (finalize +   (record_trans-  (per-(stage,
+    cycle, success/    streaming     fer_tx/rx        replica)
+    fail counter)      hooks via     hooks via        scheduler/
+                       observe_*     emit hook in     iteration
+                       APIs)         OrchestratorAg-  stats)
+                                     gregator)
+```
+
+### Data Flow
+
+There are four independent paths for metric collection.
+
+**Path 1: Pipeline-level metrics (`vllm:omni_*`)**
+
+`OmniPrometheusMetrics` registers the Gauge / Counter / Histogram collectors at import time. It is instantiated once per entrypoint, labeled with the model name. The entrypoint calls its methods as requests progress:
+
+- `set_running(n)` / `set_waiting(n)` — updated after each request completes. The running count comes from `OmniRequestCounter`, a simple counter incremented/decremented by the Orchestrator as it tracks requests. Waiting is derived as `total - running`.
+- `request_succeeded(e2e_seconds, finished_reason="stop")` — recorded when a request finishes at the final stage. `finished_reason` is extracted from `engine_outputs.outputs[0].finish_reason` (vLLM `CompletionOutput` convention) and increments `vllm:omni_requests_success_total{finished_reason}`.
+- `request_failed()` — recorded by the cleanup path when a request exits without natural completion. Internally maps to `finished_reason="abort"` so a single Counter family covers both natural and aborted completion.
+
+**Path 2: Audio modality metrics (`vllm:omni_audio_*`)**
+
+`OmniModalityMetrics` registers seven audio families with `{model_name, stage, replica}` (plus an extra `threshold_ms` / `reason` label on the two extra-cardinality Counters). Three observation sites:
+
+- `observe_modality_at_finalize(...)` — called from `omni_base._process_single_result` inside the existing `e2e_done` finalize guard. For `output_type == "audio"` it emits `audio_frames_total`, `audio_duration_s`, `audio_rtf` (or `audio_skipped_requests_total{reason="no_audio_data"}` when no audio was produced). Sample rate is resolved from `engine_outputs.multimodal_output` via `definitions.resolve_audio_sample_rate(...)` (fallback chain mirrors `serving_chat.py`'s audio response path).
+- `observe_audio_first_packet(...)` — called from the OpenAI SSE audio branch in `serving_chat.py` on the first audio packet for a request. The once-per-request guard is held by `ClientRequestState.first_audio_ts`. The `request_arrival_ts` anchor is stored in `ClientRequestState` by `async_omni.generate()`, computed at request entry.
+- `observe_audio_streaming_finalize(...)` — called from `serving_chat.py` after the streaming chunk loop exhausts. It runs the per-chunk player simulation from `vllm_omni/benchmarks/audio_continuity.py` to compute the worst-case underrun and emits `audio_underrun_s` plus (when the request stayed below the threshold) `audio_continuity_ok_total{threshold_ms}`. Per-chunk PCM byte counts and arrival timestamps are recorded by the same audio branch that updates `first_audio_ts`.
+
+**Path 3: Cross-stage transfer metrics (`vllm:omni_transfer_*`)**
+
+`OmniTransferMetrics` registers four Histogram families with `{model_name, from_stage, from_replica, to_stage, to_replica}` labels. Each observation corresponds to one physical transfer hop (one chunk between adjacent stages), not the per-request accumulated total — so the histograms track per-transfer distribution.
+
+The hook lives in `OrchestratorAggregator.record_transfer_tx` and `record_transfer_rx`. After the existing `TransferEdgeStats` accumulation, the aggregator calls `_emit_transfer_tx` / `_emit_transfer_rx`. Those:
+
+1. Resolve `from_replica` / `to_replica` via a `replica_resolver` callback supplied by `async_omni.py`. The resolver delegates to `stage_pool.get_bound_replica_id(request_id)` — i.e. the orchestrator's existing sticky-routing binding is the source of truth.
+2. Convert the underlying `_ms` accumulators to seconds and call the `_s`-suffixed observe methods on `OmniTransferMetrics`.
+
+Defensive fail-safe: if `transfer_emitter` or `replica_resolver` is missing, or the resolver returns `None` for either side, the emit is skipped silently (the underlying `TransferEdgeStats` accumulation is unaffected).
+
+**Path 4: Per-engine metrics (`vllm:*`, stage/replica wrap)**
+
+The Orchestrator instantiates `OmniPrometheusStatLogger` (a thin subclass of upstream `vllm.v1.metrics.loggers.PrometheusStatLogger`) and feeds it scheduler stats and iteration stats after processing each batch of engine outputs. This populates the standard ~37 vLLM metric families (TTFT, ITL, TPOT, KV cache usage, etc.) using the same upstream code path — but with the `engine` label reshaped into `stage` + `replica` so multi-replica deployments produce distinct series per replica. See the next section for the wrap mechanics.
+
+### Shared State Between Threads
+
+The Orchestrator runs in a background thread. The API server (OmniBase) runs in the asyncio event loop thread. `OmniRequestCounter` bridges them — a plain Python object with an `int` field. The Orchestrator increments/decrements it; the entrypoint reads it for gauge updates. No lock is needed because the counter is advisory (a stale read by one Prometheus scrape interval is acceptable). It is created by `AsyncOmniEngine.__init__()` and passed to the Orchestrator at construction time.
+
+### Metric Registration and Lifecycle
+
+All `vllm:omni_*` collectors are registered once when their owning class (`OmniPrometheusMetrics` / `OmniModalityMetrics` / `OmniTransferMetrics`) is imported. Per-`(stage, replica)` labels are bound lazily on first observation to avoid registering label sets for combinations that never produce data.
+
+The `prometheus_client` default registry holds all collectors. FastAPI's `/metrics` endpoint serves the default registry, so `vllm:omni_*` and the wrapped `vllm:*` metrics appear in the same scrape response alongside `http_*` and `process_*` metrics from the instrumentator and the Python client runtime.
+
+## OmniPrometheusStatLogger Wrap
+
+Upstream `PrometheusStatLogger.__init__` hard-codes `labelnames = ["model_name", "engine"]` as a local variable, references it across ~37 metric-family construction sites, and uses the `engine` label value in five different `.labels()` call shapes (kwarg with int engine, kwarg with str engine, positional with str engine in the middle, plus a `metrics_info["engine"] = str(...)` dict pattern). To reshape `engine` into `stage` + `replica` without forking the entire upstream `__init__`, the wrap uses three coordinated mechanisms:
+
+1. **Class-level metric class slot overrides.** `OmniPrometheusStatLogger` overrides `_gauge_cls`, `_counter_cls`, `_histogram_cls` (which upstream calls via `self._gauge_cls(...)` etc.) with `_RelabelGauge` / `_RelabelCounter` / `_RelabelHistogram` wrapper classes. These intercept the `labelnames` kwarg at metric family creation time and replace `engine` with `("stage", "replica")`.
+2. **Property descriptor for `per_engine_labelvalues`.** Upstream builds `self.per_engine_labelvalues = {idx: [model_name, str(idx)]}` inside `__init__` and then captures it into a local variable for `create_metric_per_engine` calls. By making `per_engine_labelvalues` a Python property on the subclass, the setter intercepts upstream's assignment and rewrites each 2-tuple into a 3-tuple `[model_name, stage, replica]` using the `stage_replica_map` supplied at construction time. The captured local then sees the rewritten dict.
+3. **Override of `.labels()` on the wrapper classes.** For the five call sites that pass `engine` directly (kwarg or positional, int or str), `_RelabelMixin.labels()` translates the engine value back to `(stage, replica)` via a process-level `_ENGINE_INDEX_MAP` populated by `OmniPrometheusStatLogger.__init__`. This handles `gauge_engine_sleep_state.labels(engine=idx, ...)`, `counter_request_success_base.labels(model_name, str(idx), str(reason))`, `info_gauge.labels(**metrics_info)`, etc.
+
+The three sub-helpers that upstream `PrometheusStatLogger.__init__` constructs (`spec_decoding_prom` / `kv_connector_prom` / `perf_metrics_prom`) use their own `_counter_cls` / `_gauge_cls` / `_histogram_cls` slots and would otherwise build families with the raw 2-element labelnames. `_OmniPerfMetricsProm` / `_OmniSpecDecodingProm` / `_OmniKVConnectorProm` subclass each helper to route the same relabel mixin through their internal family construction.
+
+The `Orchestrator` constructs `stage_replica_map` from the static `stage_pools` configuration at startup:
+
+```python
+stage_replica_map = {
+    flat_idx: (str(stage_id), str(replica_id))
+    for flat_idx, (stage_id, replica_id) in enumerate(
+        (s, r)
+        for s, pool in enumerate(stage_pools)
+        for r in range(pool.num_replicas)
+    )
+}
+```
+
+A reverse map `(stage_id, replica_id) -> flat_idx` is maintained on the Orchestrator so the per-replica `record(engine_idx=...)` call site can look up the right flat index.
+
+> Dynamic add/remove of replicas at runtime is intentionally out of scope — the upstream `PrometheusStatLogger` materializes per-engine_idx child metrics at init time, and supporting hot-add would require non-trivial intervention into upstream's per-family child dictionaries.
+
+## Gating: `--log-stats`
+
+All metrics — both the 15 `vllm:omni_*` families and the ~65 upstream `vllm:*` wrap families — are gated by the user's `--log-stats` CLI flag (default off). The flag is plumbed from `OmniBase.__init__(log_stats=...)` through `AsyncOmniEngine` to the stage-spawn helpers and to the three Prometheus metric classes (`OmniPrometheusMetrics` / `OmniModalityMetrics` / `OmniTransferMetrics`), and is forwarded to `Orchestrator._init_metrics_state(...)`.
+
+Behavior with `--log-stats=off` (default):
+
+- The three `Omni*Metrics` classes register their module-level Gauge / Counter / Histogram families at import time (prometheus_client requires up-front registration), but each `observe / inc / set` method early-returns. The per-label child series for `vllm:omni_*` stay materialized but never have data written to them.
+- `OmniPrometheusStatLogger` is not constructed in `_init_metrics_state`, so the ~65 upstream `vllm:*` wrap families are not registered in the default registry at all.
+- The engine core's `Scheduler.make_stats()` also short-circuits inside upstream (`if not self.log_stats: return None`), so no `SchedulerStats` is produced per step — the per-iteration cost is bounded by the existing upstream gate.
+
+Behavior with `--log-stats=on`: all metric paths fire normally; the orchestrator's per-replica recording is bounded only by `OmniSchedulerMixin.make_stats()`'s per-scheduler 1 Hz throttle (see next section).
+
+The overhead with the flag on is small enough that an A/B benchmark on Qwen3-Omni-30B single replica (30 sequential audio requests) showed a mean latency delta of +0.6% (Welch's t = 0.318, n=30, not statistically significant at α=0.05); the `/metrics` line count drops from 1358 to 124 lines when the flag is off.
+
+## Throttling: `make_stats()` Override
+
+Upstream vLLM's `Scheduler.make_stats()` runs on every AR generation step, returning a SchedulerStats object for the orchestrator. Under vLLM's architecture, this is fine. But since vLLM-Omni requires that the object be serialized and transferred over ZMQ, receiving a SchedulerStats object on every step can introduce unacceptable overhead to the system.
+
+`OmniSchedulerMixin.make_stats()` (in `vllm_omni/core/sched/omni_scheduler_mixin.py`) throttles stats emission to at most once per second **per scheduler** — i.e. per `(stage, replica)` since each replica owns its own scheduler instance. Between intervals it returns `None`, which the engine core skips serializing. This keeps gauges fresh enough for Prometheus scrapes (typically 15-30s intervals) while eliminating the per-step overhead.
+
+The orchestrator side does not add its own throttle on top: the per-replica recording loop gates only on `raw_outputs.scheduler_stats is not None` (i.e. this replica's scheduler passed its own 1 Hz gate). A previous global `_last_stats_ts` on the orchestrator was removed because it starved every replica other than the first to emit within each second.
+
+## Metric Definitions
+
+### Pipeline (4)
+
+| Metric | Type | Labels | Description |
+|--------|------|--------|-------------|
+| `vllm:omni_num_requests_running` | Gauge | `model_name` | Requests currently executing across all stages |
+| `vllm:omni_num_requests_waiting` | Gauge | `model_name` | Requests queued but not yet scheduled |
+| `vllm:omni_requests_success_total` | Counter | `model_name`, `finished_reason` | Total requests by completion reason ({stop, length, abort, ...}); aborts cover client-disconnect / cancellation paths in addition to upstream `FinishReason.ABORT` |
+| `vllm:omni_e2e_request_latency_s` | Histogram | `model_name` | Pipeline-global end-to-end latency in seconds |
+
+### Audio (7)
+
+Labels: `{model_name, stage, replica}` plus the listed extra label.
+
+| Metric | Type | Extra label | Description |
+|--------|------|-------------|-------------|
+| `vllm:omni_audio_ttfp_s` | Histogram | — | Time from request arrival to first audio packet/frame |
+| `vllm:omni_audio_duration_s` | Histogram | — | Audio content duration (`audio_frames / sample_rate`) |
+| `vllm:omni_audio_rtf` | Histogram | — | Real-time factor `stage_gen_time_s / audio_duration_s` (SLO `< 1`); uses `RTF_BUCKETS` |
+| `vllm:omni_audio_frames_total` | Counter | — | Cumulative audio frames generated |
+| `vllm:omni_audio_underrun_s` | Histogram | — | Per-request worst-case player deficit; `> 0` indicates listener heard silent gaps |
+| `vllm:omni_audio_continuity_ok_total` | Counter | `threshold_ms` | Incremented when the request's worst underrun stayed below `threshold_ms` |
+| `vllm:omni_audio_skipped_requests_total` | Counter | `reason` | Silent-loss counter — code2wav rejected malformed codec input and returned `200 OK` with empty audio |
+
+### Cross-stage transfer (4)
+
+Labels: `{model_name, from_stage, from_replica, to_stage, to_replica}`.
+
+| Metric | Type | Description |
+|--------|------|-------------|
+| `vllm:omni_transfer_size_bytes` | Histogram | Per-transfer payload size in bytes |
+| `vllm:omni_transfer_tx_s` | Histogram | Sender-side time (serialize + submit to connector) |
+| `vllm:omni_transfer_rx_s` | Histogram | Receiver-side time (recv + deserialize) |
+| `vllm:omni_transfer_in_flight_s` | Histogram | Network in-flight time (TX done → RX recv start) |
+
+### LLM stage-level (wrapped `vllm:*`)
+
+After the wrap, every upstream `vllm:*` family — TTFT, ITL, TPOT, e2e latency, KV cache usage, scheduler running/waiting, request success counts, etc. — carries `{model_name, stage, replica}` labels. For the full upstream catalog see [the vLLM docs](https://github.com/vllm-project/vllm/blob/main/docs/usage/metrics.md); note that metrics depending on features unsupported in vLLM-Omni (e.g. speculative decoding, LoRA) will not be available.
+
+## Naming Convention
+
+- All time-bearing metrics use the `_s` suffix (values in seconds). Two bucket families are used:
+  - `SECONDS_BUCKETS` (0.05 s – 300 s) for e2e / generation / TTFP style values.
+  - `SECONDS_FAST_BUCKETS` (0.001 s – 60 s) for fine-grained cross-stage transfer and audio-underrun values that need millisecond-level resolution.
+- Counters use the `_total` suffix (auto-appended by `prometheus_client`).
+- Sizes use the `_bytes` suffix.
+- All omni-specific families are prefixed `vllm:omni_`. The upstream `unregister_vllm_metrics()` function is monkey-patched to a scoped version that still strips upstream `vllm:*` collectors (so multi-engine init within one process does not crash on duplicate registration) but preserves anything prefixed `vllm:omni_` / `vllm_omni`.
+
+## Logging vs. Prometheus
+
+`OrchestratorAggregator` (in `vllm_omni/metrics/stats.py`) is the logging-oriented metrics path. It collects detailed per-request, per-stage, and per-transfer statistics and prints formatted tables to the `INFO` log. This is designed for development and debugging — individual request traces, transfer bandwidth, inter-stage timing.
+
+`OmniPrometheusMetrics` / `OmniModalityMetrics` / `OmniTransferMetrics` form the Prometheus-oriented path. They record aggregate counters, gauges, and histograms suitable for time-series monitoring and alerting. Both paths share the same source data (`StageRequestStats`, `TransferEdgeStats`) — `OrchestratorAggregator.record_transfer_tx/rx` in particular calls both the existing accumulator code and the Prometheus emit hook in the same method body. The two consumption models can run simultaneously without coupling.
+
+The separation follows upstream vLLM's pattern of `LoggingStatLogger` vs. `PrometheusStatLogger` — same underlying data, different consumption models.

docs/design/qwen3_omni_tts_performance_optimization.md

@@ -411,6 +411,8 @@ Notes:
 - `runtime.max_batch_size` controls stage-level batching.
 - Thinker/Talker commonly use `enforce_eager: false` for CUDA Graph paths.
 - Code2Wav often remains eager (`enforce_eager: true`) depending on runtime behavior.
+- Qwen3-Omni defaults `VLLM_USE_FLASHINFER_MOE_FP16=0`. The Triton has been more stable & faster
+ than the FlashInfer CUTLASS unquantized MoE backend on recent vLLM rebases.
 
 #### 2) Enable async chunk
 

docs/getting_started/installation/gpu.md

@@ -34,7 +34,7 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
 
 ### Pre-built wheels
 
-Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, and 0.21.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
+Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, and 0.22.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
 
 === "NVIDIA CUDA"
 

docs/getting_started/installation/gpu/cuda.inc.md

@@ -20,7 +20,7 @@ Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new**
 
 vLLM-Omni is built based on vLLM. Please install it with command below.
 ```bash
-uv pip install vllm==0.21.0 --torch-backend=auto
+uv pip install vllm==0.22.0 --torch-backend=auto
 ```
 
 #### Installation of vLLM-Omni
@@ -39,13 +39,13 @@ uv pip install 'vllm-omni[demo]'
 # --8<-- [start:build-wheel-from-source]
 
 #### Installation of vLLM
-If you do not need to modify source code of vLLM, you can directly install the stable 0.21.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.22.0 release version of the library
 
 ```bash
-uv pip install vllm==0.21.0 --torch-backend=auto
+uv pip install vllm==0.22.0 --torch-backend=auto
 ```
 
-The 0.21.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
+The 0.22.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
 
 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -66,12 +66,12 @@ If you want to check, modify or debug with source code of vLLM, install the libr
 ```bash
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
-git checkout v0.21.0
+git checkout v0.22.0
 ```
 Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file.
 ```bash
-#For CUDA 13.0 (the default for v0.21.0; the wheel filename has no `+cu130` suffix)
-export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.21.0/vllm-0.21.0-cp38-abi3-manylinux_2_35_x86_64.whl
+#For CUDA 13.0 (the default for v0.22.0; the wheel filename has no `+cu130` suffix)
+export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.22.0/vllm-0.22.0-cp38-abi3-manylinux_2_35_x86_64.whl
 ```
 Install vllm with command below (If you have no existing PyTorch).
 ```bash
@@ -102,7 +102,7 @@ docker run --runtime nvidia --gpus 2 \
     --env "HF_TOKEN=$HF_TOKEN" \
     -p 8091:8091 \
     --ipc=host \
-    vllm/vllm-omni:v0.21.0 \
+    vllm/vllm-omni:v0.22.0 \
     vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```
 

docs/getting_started/installation/gpu/rocm.inc.md

@@ -13,7 +13,7 @@ vLLM-Omni current recommends the steps in under setup through Docker Images.
 
 vLLM-Omni is built based on vLLM. Please install it with command below.
 ```bash
-uv pip install vllm==0.21.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.21.0/rocm722
+uv pip install vllm==0.22.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.22.0/rocm722
 ```
 
 #### Installation of vLLM-Omni
@@ -34,13 +34,13 @@ uv pip install onnxruntime-rocm
 # --8<-- [start:build-wheel-from-source]
 
 #### Installation of vLLM
-If you do not need to modify source code of vLLM, you can directly install the stable 0.21.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.22.0 release version of the library
 
 ```bash
-uv pip install vllm==0.21.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.21.0/rocm722
+uv pip install vllm==0.22.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.22.0/rocm722
 ```
 
-The pre-built 0.21.0 vLLM wheel targets ROCm 7.2.2. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
+The pre-built 0.22.0 vLLM wheel targets ROCm 7.2.2. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
 
 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -58,7 +58,7 @@ If you want to check, modify or debug with source code of vLLM, install the libr
 ```bash
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
-git checkout v0.21.0
+git checkout v0.22.0
 python3 -m pip install -r requirements/rocm.txt
 python3 setup.py develop
 ```
@@ -130,7 +130,7 @@ docker run --rm \
   -v ~/.cache/huggingface:/root/.cache/huggingface \
   --env "HF_TOKEN=$HF_TOKEN" \
   -p 8091:8091 \
-  vllm/vllm-omni-rocm:v0.21.0 \
+  vllm/vllm-omni-rocm:v0.22.0 \
   --model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```
 
@@ -149,7 +149,7 @@ docker run --rm -it \
   -v ~/.cache/huggingface:/root/.cache/huggingface \
   --env "HF_TOKEN=$HF_TOKEN" \
   --entrypoint bash \
-  vllm/vllm-omni-rocm:v0.21.0
+  vllm/vllm-omni-rocm:v0.22.0
 ```
 
 # --8<-- [end:pre-built-images]

docs/getting_started/quickstart.md

@@ -19,10 +19,10 @@ uv venv --python 3.12 --seed
 source .venv/bin/activate
 
 # On CUDA
-uv pip install vllm==0.21.0 --torch-backend=auto
+uv pip install vllm==0.22.0 --torch-backend=auto
 
 # On ROCm
-uv pip install vllm==0.21.0+rocm721 --extra-index-url https://wheels.vllm.ai/rocm/0.21.0/rocm721
+uv pip install vllm==0.22.0+rocm721 --extra-index-url https://wheels.vllm.ai/rocm/0.22.0/rocm721
 
 git clone https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni
@@ -35,7 +35,7 @@ For additional installation methods — please see the [installation guide](inst
 !!! note
     It is important to install the same major & minor version of vLLM and vLLM Omni, otherwise things may not work as expected. If the versions are misaligned, you will see a warning when you import vLLM Omni.
 
-    If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.21.0` and vLLM Omni `0.21.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue.
+    If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.22.0` and vLLM Omni `0.22.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue.
 
 ## Offline Inference
 

docs/mkdocs/hooks/generate_argparse.py

@@ -119,9 +119,10 @@ def extract_omni_serve_subparser_init():
                     exec_globals = {
                         "_FlexibleArgumentParser": _FlexibleArgumentParser,
                         "FlexibleArgumentParser": _FlexibleArgumentParser,
+                        "TrackingArgumentParser": _FlexibleArgumentParser,  # only needs add_argument, not tracking
+                        "TrackingNamespace": None,  # type annotation only
                         "make_arg_parser": lambda parser: parser,  # no-op for doc
                         "_ensure_vllm_platform": lambda: None,  # no-op for doc
-                        "nullify_stage_engine_defaults": lambda parser: None,  # no-op for doc
                         "VLLM_SUBCMD_PARSER_EPILOG": "",
                         "logger": logger,
                         "DummySubparsers": DummySubparsers,

docs/models/supported_models.md

@@ -33,6 +33,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, T2V with sound | `nvidia/Cosmos3-Nano` | ✅︎ | | | |
 | `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` | ✅︎ | ✅︎ | | |
@@ -41,6 +42,7 @@ th {
 | `LTX2ImageToVideoTwoStagesPipeline` | LTX-2-I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
 | `LTX23Pipeline` | LTX-2.3-T2V | `dg845/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
 | `LTX23ImageToVideoPipeline` | LTX-2.3-I2V | `dg845/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
+| `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | | | |
 | `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | |
 | `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | |
 | `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ |

docs/usage/metrics.md

@@ -0,0 +1,98 @@
+# Production Metrics
+
+vLLM-Omni exposes Prometheus metrics via the `/metrics` endpoint on the OpenAI-compatible API server. This page covers the text and audio surface; diffusion / image / video metrics are tracked in a follow-up PR.
+
+```bash
+vllm-omni serve Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8000 --log-stats
+curl http://localhost:8000/metrics
+```
+
+**`--log-stats` is required to populate metric data.** Without the flag, the endpoint still returns `200 OK` but the upstream `vllm:*` wrap is not registered at all, and the 15 `vllm:omni_*` families are registered as placeholders with no data written to them. This keeps the runtime cost essentially zero for deployments that don't need monitoring. With the flag, all ~80 families populate.
+
+## Metric Namespaces
+
+| Prefix | Source | Present when |
+|--------|--------|--------------|
+| `vllm:omni_` | vLLM-Omni orchestrator / audio modality / cross-stage transfer | Pipeline-dependent |
+| `vllm:` | Upstream vLLM engine, wrapped by `OmniPrometheusStatLogger` to expose `{stage, replica}` | Pipeline includes an LLM (AR) stage |
+| `http_` / `process_` | Uvicorn / Python runtime | Always |
+
+## Pipeline-Level Metrics (`vllm:omni_`)
+
+Defined in `vllm_omni/metrics/prometheus.py`. Track request lifecycle across the full multi-stage pipeline.
+
+### Request counts
+
+| Metric | Type | Labels | Description |
+|--------|------|--------|-------------|
+| `vllm:omni_num_requests_running` | Gauge | `model_name` | Pipeline-global in-flight requests (dispatched to engine, not yet finalized) |
+| `vllm:omni_num_requests_waiting` | Gauge | `model_name` | Requests waiting in the Orchestrator queue |
+| `vllm:omni_requests_success_total` | Counter | `model_name`, `finished_reason` | Total requests by completion reason. `finished_reason` ∈ {`stop`, `length`, `abort`, ...} mirroring upstream `vllm:request_success_total`; aborts cover client disconnect / cancellation paths in addition to upstream `FinishReason.ABORT` |
+
+### Latency
+
+| Metric | Type | Labels | Description |
+|--------|------|--------|-------------|
+| `vllm:omni_e2e_request_latency_s` | Histogram | `model_name` | Pipeline-global end-to-end request latency in seconds |
+
+## Audio Modality Metrics (`vllm:omni_`)
+
+Emitted at request finalize, except for `audio_ttfp_s` (streaming-hook at the first audio packet) and `audio_underrun_s` / `audio_continuity_ok_total` (streaming finalize, after the chunk stream is exhausted). All carry `{model_name, stage, replica}` plus the listed extra label.
+
+| Metric | Type | Extra label | Description |
+|--------|------|-------------|-------------|
+| `vllm:omni_audio_ttfp_s` | Histogram | — | Time from request arrival to first audio packet/frame |
+| `vllm:omni_audio_duration_s` | Histogram | — | Audio content duration (`audio_frames / sample_rate`) |
+| `vllm:omni_audio_rtf` | Histogram | — | Real-time factor (`stage_gen_time_s / audio_duration_s`); streaming TTS SLO red line `< 1`; uses `RTF_BUCKETS` |
+| `vllm:omni_audio_frames_total` | Counter | — | Cumulative audio frame count; throughput via `rate()` |
+| `vllm:omni_audio_underrun_s` | Histogram | — | Per-request worst-case player deficit; `> 0` indicates listener heard silent gaps |
+| `vllm:omni_audio_continuity_ok_total` | Counter | `threshold_ms` | Incremented when the request's worst underrun stayed below `threshold_ms` |
+| `vllm:omni_audio_skipped_requests_total` | Counter | `reason` | Silent-loss counter — code2wav rejected malformed codec input and returned `200 OK` with empty audio |
+
+The continuity math comes from `vllm_omni/benchmarks/audio_continuity.py::compute_continuity_stats` so the server-side observation aligns with the bench-side definition.
+
+## Cross-Stage Transfer Metrics (`vllm:omni_`)
+
+Per-physical-transfer histograms tracking the data hop between adjacent stages. Labels `{model_name, from_stage, from_replica, to_stage, to_replica}` let dashboards attribute latency to specific replica edges. `from_replica` / `to_replica` are resolved from the orchestrator's sticky-routing binding (`stage_pool.get_bound_replica_id(request_id)`), so no extra plumbing through `TransferEdgeStats` is needed.
+
+| Metric | Type | Description |
+|--------|------|-------------|
+| `vllm:omni_transfer_size_bytes` | Histogram | Per-transfer payload size in bytes |
+| `vllm:omni_transfer_tx_s` | Histogram | Sender-side time (serialize + submit to connector) |
+| `vllm:omni_transfer_rx_s` | Histogram | Receiver-side time (recv + deserialize) |
+| `vllm:omni_transfer_in_flight_s` | Histogram | Network in-flight time (TX done → RX recv start) |
+
+## vLLM Engine Metrics (`vllm:`)
+
+When the pipeline includes an LLM stage, the upstream vLLM engine exposes its full set of ~37 metric families under the `vllm:` prefix.
+
+vLLM-Omni wraps the upstream `vllm.v1.metrics.loggers.PrometheusStatLogger` with `OmniPrometheusStatLogger` so that the original `engine` single label is reshaped into `stage` + `replica`. Every `vllm:*` family — TTFT, ITL, TPOT, e2e latency, KV cache usage, scheduler running/waiting, request success counts, etc. — therefore gains per-`(stage, replica)` visibility automatically. No omni-side duplicate is needed for the text path.
+
+```text
+# Before wrap:
+vllm:num_requests_running{model_name="...", engine="1"}              3.0
+
+# After wrap:
+vllm:num_requests_running{model_name="...", stage="1", replica="0"}  2.0
+vllm:num_requests_running{model_name="...", stage="1", replica="1"}  1.0
+```
+
+For the full list of upstream metrics, see [the vLLM docs](https://github.com/vllm-project/vllm/blob/main/docs/usage/metrics.md).
+
+## Metric Availability by Pipeline Type
+
+| Metric group | Multi-stage LLM (Qwen3-Omni) |
+|---|---|
+| `vllm:omni_` request tracking + latency | With `--log-stats` |
+| `vllm:omni_` audio modality | With `--log-stats`, if pipeline has a talker stage |
+| `vllm:omni_` transfer | With `--log-stats`, if pipeline has ≥ 2 stages |
+| `vllm:` engine metrics (per `(stage, replica)`) | With `--log-stats` |
+| `vllm:` MFU metrics | With `--log-stats --enable-mfu-metrics` |
+
+## Naming Convention
+
+- All time-bearing metrics use the `_s` suffix (values in seconds). Buckets are `SECONDS_BUCKETS` for e2e / generation-style values and `SECONDS_FAST_BUCKETS` (1 ms → 60 s) for the fine-grained transfer and audio-underrun values.
+- Counters use the `_total` suffix (auto-appended by `prometheus_client`).
+- Sizes use the `_bytes` suffix.
+- All omni-specific families are prefixed `vllm:omni_`. The upstream `unregister_vllm_metrics()` function is monkey-patched (see `vllm_omni/patch.py`) to a scoped version that still strips upstream `vllm:*` collectors so multi-engine init within one process does not crash on duplicate registration, but preserves anything prefixed `vllm:omni_` / `vllm_omni`.
+- Text and audio first-output use distinct families (`vllm:time_to_first_token_seconds` reused from upstream for text; `vllm:omni_audio_ttfp_s` for audio) rather than a single metric with a `modality` label.

docs/user_guide/diffusion_features.md

@@ -133,10 +133,12 @@ The following tables show which models support each feature:
 | **Stable-Diffusion3.5**  |     ❌     |     ✅      |           ❌           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **Z-Image**              |     ✅     |     ✅      |           ✅           |       ❓        |   ✅ (TP=2 only)   |          ❌          |   ✅    |             ❌             |      ✅ (decode)      |       ✅        |        ❌         |
 | **ERNIE-Image**          |     ❌     |     ✅      |           ✅           |       ❓        |         ✅         |          ❌          |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
+| **Cosmos3**              |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ✅        |        ❌         |
 
 > Notes:
 > 1. Nextstep_1(T2I) does not support cache acceleration methods such as TeaCache or Cache-DiT.
 > 2. `Tongyi-MAI/Z-Image-Turbo` and `SII-GAIR/daVinci-MagiHuman-Base-1080p` are distilled models with minimal NFEs; CFG-Parallel is not necessary.
+> 3. Cosmos3 T2I uses `Cosmos3OmniDiffusersPipeline` with `modalities=["image"]`. Model-level CPU offload is not supported; use layerwise offload.
 
 ### VideoGen
 
@@ -149,6 +151,8 @@ The following tables show which models support each feature:
 | **Helios**                   |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **HunyuanVideo-1.5 T2V I2V** |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |  ✅ (encode/decode)   |       ✅        |        ❌         |
 | **DreamID-Omni**             |     ❌     |     ❌      |           ❌           |       ✅        |         ❌         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
+| **Cosmos3**                  |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |  ✅ (encode/decode)   |       ✅        |        ❌         |
+
 
 **Frame Interpolation Support**
 

docs/user_guide/examples/online_serving/qwen3_omni.md

@@ -17,6 +17,9 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 
 The default deployment configuration situated at `vllm_omni/deploy/qwen3_omni_moe.yaml` is resolved and loaded
 automatically via the model registry, obviating the necessity for the `--deploy-config` flag in standard deployment topologies.
+The bundled Qwen3-Omni setup defaults `VLLM_USE_FLASHINFER_MOE_FP16=0`. This keeps the Thinker & Talker on vLLM's
+Triton unquantized MoE path and avoids the performance regression observed with the FlashInfer CUTLASS unquantized MoE
+backend.
 Asynchronous chunk streaming is **enabled by default** within the bundled configuration.
 
 To explicitly utilize a custom deployment YAML, specify the configuration path:
@@ -72,6 +75,12 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
     --stage-overrides '{"1": {"gpu_memory_utilization": 0.5}}'
 ```
 
+To experiment with the FlashInfer FP16 MoE path, set `VLLM_USE_FLASHINFER_MOE_FP16=1` before launching the server:
+```bash
+VLLM_USE_FLASHINFER_MOE_FP16=1 \
+vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
+```
+
 For the stage-based CLI, you usually do **not** need `--stage-overrides` for
 that kind of change. Since each command launches one stage, just pass the knob
 directly on that stage command:

docs/user_guide/examples/online_serving/text_to_speech.md

@@ -447,6 +447,7 @@ python qwen3_tts/gradio_fastrtc_demo.py --api-base http://localhost:8000
 - Base voice cloning has uniproc-vs-mp tradeoffs depending on per-request reference audio cost; see the executor-backend section above.
 - With async chunking, Qwen3-TTS Base voice cloning sends the full reference context in the first Code2Wav packet, then caches that prefix on the Code2Wav stage for follow-up chunks in the same request.
 - `vllm_omni/deploy/qwen3_tts.yaml` is the default deploy config (loaded by HF `model_type`); per-stage runtime overrides are available via `--stage-N-<field> <value>`.
+- Under vocoder-bound overload (single-stream `rtf_p99 ≥ 1` at the target concurrency), set `active_stream_window: 2` at the top of the deploy yaml to cap simultaneously active Stage 1 streams. Off by default; trades TTFP for streaming continuity. See [#3592](https://github.com/vllm-project/vllm-omni/pull/3592) for the mechanism and tradeoff numbers.
 
 ---
 

examples/offline_inference/bagel/README.md

@@ -195,6 +195,101 @@ stages:
     devices: "0,1"
 ```
 
+### VAE Patch Parallelism
+
+[VAE Patch Parallelism](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism/vae_patch_parallel.html) splits Bagel VAE **decode/encode** tiles across multiple GPUs on the **DiT stage**, reducing **per-GPU peak memory during VAE decode**. Use it when high-resolution `text2img` or `img2img` hits VAE OOM or large decode spikes.
+
+**Bagel-specific notes:**
+
+- Implemented in `BagelPipeline` via `DistributedAutoEncoder` (DiT stage only).
+- **Single-stage** is the simplest path: one DiT process with TP + VAE patch parallel.
+- **Two-stage**: enable on **stage 1 (DiT)** only; stage 0 (Thinker) keeps encoder-only `VAEEncoder` and does not use VAE patch parallel.
+- You need a DiT `world_size` ≥ `vae_patch_parallel_size` (typically `tensor_parallel_size=2` on that stage). VAE PP reuses the DiT process group; it is not a standalone second-GPU VAE worker.
+
+**Single-stage via deploy YAML** (recommended for `end2end.py`):
+
+```yaml
+pipeline: bagel_single_stage
+async_chunk: false
+
+stages:
+  - stage_id: 0
+    max_num_batched_tokens: 32768
+    max_num_seqs: 1
+    enforce_eager: true
+    trust_remote_code: true
+    enable_prefix_caching: false
+    devices: "0,1"
+    vae_use_tiling: true
+    parallel_config:
+      tensor_parallel_size: 2
+      vae_patch_parallel_size: 2
+    default_sampling_params:
+      seed: 52
+```
+
+```bash
+cd examples/offline_inference/bagel
+
+CUDA_VISIBLE_DEVICES=0,1 python end2end.py \
+    --model /path/to/BAGEL-7B-MoT \
+    --deploy-config /path/to/bagel_single_stage_vae_pp.yaml \
+    --modality text2img \
+    --prompts "A cute cat" \
+    --steps 10 \
+    --output ./out_vae_pp
+```
+
+**Single-stage via `Omni` kwargs** (same flags as online serving):
+
+```python
+from vllm_omni.entrypoints.omni import Omni
+
+omni = Omni(
+    model="ByteDance-Seed/BAGEL-7B-MoT",
+    deploy_config="vllm_omni/deploy/bagel_single_stage.yaml",
+    tensor_parallel_size=2,
+    vae_patch_parallel_size=2,
+    vae_use_tiling=True,
+)
+# Then call omni.generate(...) as in end2end.py
+```
+
+**Two-stage (VAE PP on DiT only):**
+
+```yaml
+stages:
+  - stage_id: 0
+    devices: "0"
+    # AR Thinker — no vae_patch_parallel here
+
+  - stage_id: 1
+    devices: "0,1"
+    vae_use_tiling: true
+    parallel_config:
+      tensor_parallel_size: 2
+      vae_patch_parallel_size: 2
+```
+
+```bash
+python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
+    --deploy-config /path/to/bagel_vae_pp.yaml \
+    --modality text2img \
+    --prompts "A cute cat"
+```
+
+**Startup log checks:**
+
+```text
+INFO ... vae_patch_parallel_size=2 requires vae_use_tiling; automatically enabling it.
+```
+
+| Setting | Role |
+| :------ | :--- |
+| `parallel_config.tensor_parallel_size` | DiT world size / TP (must be ≥ `vae_patch_parallel_size`) |
+| `parallel_config.vae_patch_parallel_size` | Number of ranks for distributed VAE tiles (`1` = off) |
+| `vae_use_tiling` | Enable spatial tiling (auto-enabled when `vae_patch_parallel_size > 1`) |
+
 #### Hybrid Sharded Data Parallel (HSDP)
 
 For larger Bagel deployments on multiple GPUs, you can enable HSDP (Hybrid Sharded Data Parallel) by modifying the stage configuration (for example, [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml)). HSDP shards transformer weights across GPUs to reduce per-GPU memory usage.

examples/offline_inference/bagel/end2end.py

@@ -121,9 +121,6 @@ def parse_args():
         help="Temperature for text generation sampling (default: 0.3).",
     )
 
-    from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
-
-    nullify_stage_engine_defaults(parser)
     args = parser.parse_args()
     return args
 
@@ -154,6 +151,7 @@ def main():
     from vllm_omni.entrypoints.omni import Omni
 
     omni_kwargs = vars(args).copy()
+    omni_kwargs.pop("model", None)
     deploy_config = args.deploy_config
     if args.think and deploy_config is None:
         deploy_config = "vllm_omni/deploy/bagel_think.yaml"
@@ -164,9 +162,7 @@ def main():
     if args.quantization:
         omni_kwargs["quantization_config"] = args.quantization
 
-    # Override CLI --model with the derived model_name.
-    omni_kwargs["model"] = model_name
-    omni = Omni(**omni_kwargs)
+    omni = Omni(model=model_name, **omni_kwargs)
 
     formatted_prompts = []
     for p in prompts:

examples/offline_inference/dreamzero/export_prediction_video.py

@@ -0,0 +1,304 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import argparse
+import uuid
+from pathlib import Path
+
+import cv2
+import numpy as np
+import torch
+from PIL import Image
+
+from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+
+WORKER_EXTENSION = "vllm_omni.diffusion.models.dreamzero.video_export_worker.DreamZeroVideoExportWorkerExtension"
+ASSET_REPO_ID = "YangshenDeng/vllm-omni-dreamzero-assets"
+DEFAULT_SESSION_PREFIX = "dreamzero-export"
+RELATIVE_OFFSETS = [-23, -16, -8, 0]
+ACTION_HORIZON = 24
+CAMERA_FILES = {
+    "observation/exterior_image_0_left": "exterior_image_1_left.mp4",
+    "observation/exterior_image_1_left": "exterior_image_2_left.mp4",
+    "observation/wrist_image_left": "wrist_image_left.mp4",
+}
+
+
+def _parse_args() -> argparse.Namespace:
+    repo_root = Path(__file__).resolve().parents[3]
+    parser = argparse.ArgumentParser(description="Export DreamZero prediction video from downloaded example inputs.")
+    parser.add_argument("--model", default="GEAR-Dreams/DreamZero-DROID")
+    parser.add_argument("--deploy-config", type=Path, required=True)
+    parser.add_argument(
+        "--video-dir",
+        type=Path,
+        default=repo_root / "outputs" / "dreamzero" / "assets",
+        help="Directory containing the three camera MP4 files.",
+    )
+    parser.add_argument(
+        "--output-dir", type=Path, default=repo_root / "outputs" / "dreamzero" / "generated_predictions"
+    )
+    parser.add_argument("--output-stem", default="dreamzero_prediction")
+    parser.add_argument(
+        "--prompt",
+        default="Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan",
+    )
+    parser.add_argument("--session-id", default=None)
+    parser.add_argument("--save-input-video", action="store_true")
+    parser.add_argument("--save-gif", action="store_true")
+    parser.add_argument("--save-actions", action="store_true")
+    parser.add_argument("--fps", type=int, default=5)
+    return parser.parse_args()
+
+
+def _load_all_frames(video_path: Path) -> np.ndarray:
+    cap = cv2.VideoCapture(str(video_path))
+    frames = []
+    while True:
+        ok, frame = cap.read()
+        if not ok:
+            break
+        frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+    cap.release()
+    if not frames:
+        raise RuntimeError(f"No frames loaded from {video_path}")
+    return np.stack(frames, axis=0)
+
+
+def _load_camera_frames(video_dir: Path) -> dict[str, np.ndarray]:
+    camera_frames: dict[str, np.ndarray] = {}
+    for camera_key, file_name in CAMERA_FILES.items():
+        video_path = video_dir / file_name
+        if not video_path.exists():
+            raise FileNotFoundError(
+                f"Missing DreamZero example asset: {video_path}. "
+                "Download the example videos with: "
+                f"`hf download {ASSET_REPO_ID} --repo-type dataset --local-dir {video_dir}`"
+            )
+        camera_frames[camera_key] = _load_all_frames(video_path)
+    return camera_frames
+
+
+def _build_frame_schedule(total_frames: int, num_chunks: int) -> list[list[int]]:
+    chunks: list[list[int]] = []
+    current_frame = 23
+    for _ in range(num_chunks):
+        indices = [max(current_frame + offset, 0) for offset in RELATIVE_OFFSETS]
+        if indices[-1] >= total_frames:
+            break
+        chunks.append(indices)
+        current_frame += ACTION_HORIZON
+    return chunks
+
+
+def _make_obs_from_video(
+    camera_frames: dict[str, np.ndarray],
+    frame_indices: list[int],
+    *,
+    prompt: str,
+    session_id: str,
+) -> dict:
+    obs: dict = {}
+    for camera_key, all_frames in camera_frames.items():
+        selected = all_frames[frame_indices]
+        obs[camera_key] = selected[0] if len(frame_indices) == 1 else selected
+
+    obs["observation/joint_position"] = np.zeros(7, dtype=np.float32)
+    obs["observation/cartesian_position"] = np.zeros(6, dtype=np.float32)
+    obs["observation/gripper_position"] = np.zeros(1, dtype=np.float32)
+    obs["prompt"] = prompt
+    obs["session_id"] = session_id
+    return obs
+
+
+def _build_observations(video_dir: Path, prompt: str, session_id: str) -> tuple[dict[str, np.ndarray], list[dict]]:
+    camera_frames = _load_camera_frames(video_dir)
+    total_frames = min(frames.shape[0] for frames in camera_frames.values())
+    chunks = _build_frame_schedule(total_frames, 1)
+    observations = [
+        _make_obs_from_video(camera_frames, [0], prompt=prompt, session_id=session_id),
+    ]
+    if chunks:
+        observations.append(
+            _make_obs_from_video(
+                camera_frames,
+                chunks[0],
+                prompt=prompt,
+                session_id=session_id,
+            )
+        )
+    if len(observations) < 2:
+        raise RuntimeError("Need at least two DreamZero example observations to export a prediction video.")
+    return camera_frames, observations[:2]
+
+
+def _extract_latents(output: OmniRequestOutput) -> torch.Tensor:
+    if not isinstance(output, OmniRequestOutput):
+        raise TypeError(f"Expected OmniRequestOutput, got {type(output)!r}")
+    if not output.images:
+        raise RuntimeError("DreamZero output does not contain video latents in `images`.")
+
+    latents = output.images[0]
+    if not isinstance(latents, torch.Tensor):
+        raise TypeError(f"Expected tensor latents, got {type(latents)!r}")
+
+    latents = latents.detach().cpu()
+    if latents.dim() == 4:
+        latents = latents.unsqueeze(0)
+    if latents.dim() != 5:
+        raise ValueError(f"Unexpected latent shape: {tuple(latents.shape)}")
+
+    if latents.shape[1] < latents.shape[2]:
+        latents = latents.transpose(1, 2).contiguous()
+    return latents
+
+
+def _write_mp4(path: Path, frames: np.ndarray, fps: int) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    height, width = frames.shape[1:3]
+    writer = cv2.VideoWriter(
+        str(path),
+        cv2.VideoWriter_fourcc(*"mp4v"),
+        float(fps),
+        (width, height),
+    )
+    if not writer.isOpened():
+        raise RuntimeError(f"Failed to open video writer for {path}")
+    try:
+        for frame in frames:
+            writer.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
+    finally:
+        writer.release()
+
+
+def _write_gif(path: Path, frames: np.ndarray, fps: int) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    images = [Image.fromarray(frame) for frame in frames]
+    duration_ms = max(int(round(1000 / max(fps, 1))), 1)
+    images[0].save(
+        path,
+        save_all=True,
+        append_images=images[1:],
+        duration=duration_ms,
+        loop=0,
+    )
+
+
+def _stitch_input_frames(camera_frames: dict[str, np.ndarray]) -> np.ndarray:
+    total_frames = min(frames.shape[0] for frames in camera_frames.values())
+    stitched = []
+    for frame_index in range(total_frames):
+        left = camera_frames["observation/exterior_image_0_left"][frame_index]
+        right = camera_frames["observation/exterior_image_1_left"][frame_index]
+        wrist = camera_frames["observation/wrist_image_left"][frame_index]
+        pad = np.zeros((left.shape[0], left.shape[1], 3), dtype=np.uint8)
+        canvas = np.concatenate([left, right], axis=1)
+        bottom = np.concatenate([wrist, pad], axis=1)
+        stitched.append(np.concatenate([canvas, bottom], axis=0))
+    return np.stack(stitched, axis=0)
+
+
+def _run_generation(
+    model: str, deploy_config_path: Path, observations: list[dict]
+) -> tuple[Omni, list[OmniRequestOutput]]:
+    omni = Omni(
+        model=model,
+        deploy_config=str(deploy_config_path),
+        enforce_eager=True,
+        worker_extension_cls=WORKER_EXTENSION,
+    )
+
+    outputs: list[OmniRequestOutput] = []
+    for index, obs in enumerate(observations):
+        sampling_params = OmniDiffusionSamplingParams(
+            extra_args={
+                "reset": index == 0,
+                "session_id": obs["session_id"],
+                "robot_obs": obs,
+            }
+        )
+        result = omni.generate(obs["prompt"], sampling_params_list=[sampling_params])
+        if not result:
+            raise RuntimeError(f"No output returned for DreamZero request {index}")
+        outputs.append(result[0])
+    return omni, outputs
+
+
+def _decode_with_worker(omni: Omni, full_latents: torch.Tensor) -> np.ndarray:
+    stage_client = omni.engine.stage_clients[0]
+    engine = getattr(stage_client, "_engine", None)
+    if engine is None:
+        raise RuntimeError("DreamZero export requires inline diffusion stage access.")
+
+    decoded = engine.executor.collective_rpc(
+        "decode_video_latents_to_uint8",
+        args=(full_latents,),
+        unique_reply_rank=0,
+        exec_all_ranks=True,
+    )
+    if isinstance(decoded, torch.Tensor):
+        decoded = decoded.numpy()
+    if not isinstance(decoded, np.ndarray):
+        raise TypeError(f"Unexpected decoded output type: {type(decoded)!r}")
+    return decoded
+
+
+def main() -> None:
+    args = _parse_args()
+    session_id = args.session_id or f"{DEFAULT_SESSION_PREFIX}-{uuid.uuid4()}"
+
+    camera_frames, observations = _build_observations(
+        video_dir=args.video_dir,
+        prompt=args.prompt,
+        session_id=session_id,
+    )
+
+    args.output_dir.mkdir(parents=True, exist_ok=True)
+
+    if args.save_input_video:
+        input_frames = _stitch_input_frames(camera_frames)
+        _write_mp4(args.output_dir / f"{args.output_stem}_input.mp4", input_frames, fps=15)
+        if args.save_gif:
+            _write_gif(args.output_dir / f"{args.output_stem}_input.gif", input_frames[::3], fps=5)
+
+    omni = None
+    try:
+        omni, outputs = _run_generation(
+            model=args.model,
+            deploy_config_path=args.deploy_config,
+            observations=observations,
+        )
+        latent_steps = [_extract_latents(output) for output in outputs]
+        full_latents = torch.cat(latent_steps, dim=2)
+        frames = _decode_with_worker(omni, full_latents)
+    finally:
+        if omni is not None:
+            omni.close()
+
+    mp4_path = args.output_dir / f"{args.output_stem}.mp4"
+
+    _write_mp4(mp4_path, frames, fps=args.fps)
+    print(f"SAVED_MP4={mp4_path}")
+
+    if args.save_gif:
+        gif_path = args.output_dir / f"{args.output_stem}.gif"
+        _write_gif(gif_path, frames, fps=args.fps)
+        print(f"SAVED_GIF={gif_path}")
+
+    if args.save_actions:
+        npz_path = args.output_dir / f"{args.output_stem}_actions.npz"
+        np.savez(
+            npz_path,
+            step0=np.asarray(outputs[0].multimodal_output.get("actions")),
+            step1=np.asarray(outputs[1].multimodal_output.get("actions")),
+        )
+        print(f"SAVED_ACTIONS={npz_path}")
+
+
+if __name__ == "__main__":
+    main()

examples/offline_inference/dynin_omni/README.md

@@ -42,7 +42,7 @@ pip install \
 ## 3. Hardware and VRAM Requirements
 
 This example uses a 3-stage pipeline on one GPU by default
-([`dynin_omni.yaml`](../../../vllm_omni/model_executor/stage_configs/dynin_omni.yaml)):
+([`dynin_omni.yaml`](../../../vllm_omni/deploy/dynin_omni.yaml)):
 
 - Stage-0 (`token2text`): `gpu_memory_utilization: 0.5`
 - Stage-1 (`token2image`): `gpu_memory_utilization: 0.1`

examples/offline_inference/dynin_omni/end2end.py

@@ -18,8 +18,6 @@ import numpy as np
 import torch
 from PIL import Image
 
-from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
-
 TASK_CHOICES = ("t2t", "t2i", "t2s", "i2i", "i2t", "s2t", "v2t")
 
 TASK_DEFAULT_RUNTIME = {
@@ -880,8 +878,8 @@ def parse_args(repo_root: Path) -> argparse.Namespace:
     parser.add_argument(
         "--stage-config-path",
         type=str,
-        default=str(repo_root / "vllm_omni/model_executor/stage_configs/dynin_omni.yaml"),
-        help="Path to stage config yaml.",
+        default=str(repo_root / "vllm_omni/deploy/dynin_omni.yaml"),
+        help="Path to deploy config yaml.",
     )
     parser.add_argument(
         "--dynin-config-path",
@@ -972,8 +970,6 @@ def parse_args(repo_root: Path) -> argparse.Namespace:
     parser.add_argument("--vq-model-audio-local-files-only", action=argparse.BooleanOptionalAction, default=None)
 
     parser.add_argument("--disable-hf-xet", action=argparse.BooleanOptionalAction, default=True)
-
-    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/hunyuan_image3/end2end.py

@@ -107,9 +107,6 @@ def parse_args():
         ),
     )
 
-    from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
-
-    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/lance/README.md

@@ -0,0 +1,52 @@
+# Lance (ByteDance) — offline inference
+
+[Lance](https://huggingface.co/bytedance-research/Lance) is a 3B unified
+autoregressive + diffusion multimodal model on a Qwen2.5-VL backbone. It is
+**BAGEL-lineage** (ByteDance Mixture-of-Transformers): the released `Lance_3B`
+checkpoint uses the same `*_moe_gen` MoT weight layout as BAGEL, so vLLM-Omni
+implements it by reusing the BAGEL transformer core and specializing only the
+ViT (Qwen2.5-VL vision), the VAE (Wan2.2) and the checkpoint layout.
+
+This example covers all six Lance modalities from the upstream HF model card:
+`t2i`, `t2v`, `image_edit`, `video_edit`, `x2t_image` (image understanding) and
+`x2t_video` (video understanding).
+
+## Hardware
+
+Single NVIDIA GPU with 16 GB+ VRAM in BF16 (we test on B300 / A100). CUDA ≥ 12.4.
+
+## Run
+
+```bash
+# Text-to-image
+python examples/offline_inference/lance/end2end.py \
+    --model bytedance-research/Lance \
+    --prompts "a corgi astronaut on the moon, cinematic" \
+    --steps 30 --cfg-text-scale 4.0 --timestep-shift 3.5 \
+    --height 1024 --width 1024 \
+    --output ./out
+
+# Text-to-video (uses the Lance_3B_Video subfolder; see ``--modality``
+# choices for all six task variants)
+python examples/offline_inference/lance/end2end.py \
+    --model bytedance-research/Lance/Lance_3B_Video --modality text2video \
+    --num-frames 25 --video-height 480 --video-width 768 \
+    --prompts "a cat playing piano, cinematic" \
+    --steps 30 --fps 8 --output ./out
+```
+
+`video_edit` requires `--model bytedance-research/Lance/Lance_3B_Video` so
+the 3-D `latent_pos_embed` table is loaded; the other paths can point at
+the top-level `bytedance-research/Lance` repo and resolve the right
+sub-checkpoint automatically.
+
+The HF repo bundles everything (`Lance_3B/`, `Lance_3B_Video/`,
+`Qwen2.5-VL-ViT/`, `Wan2.2_VAE.pth`); no separate downloads are required.
+
+## Defaults
+
+Matches upstream `inference_lance.sh`: 30 denoising steps, timestep-shift 3.5,
+text CFG 4.0, seed 42, 1024×1024 (override with `--height` / `--width`).  For
+the understanding paths (`img2text` / `video2text`), sampling is enabled by
+default at `--text-temperature 0.8` because Lance's greedy decoder emits an
+immediate EOS for many prompts.

examples/offline_inference/lance/end2end.py

@@ -0,0 +1,339 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Lance-3B offline inference (single-stage).
+
+Lance is BAGEL-lineage; this example covers the single-stage paths wired in
+this PR:
+
+    text2img  — image generation from text
+    img2text  — image understanding (caption / VQA)
+
+Examples:
+
+    python examples/offline_inference/lance/end2end.py \
+        --model bytedance-research/Lance \
+        --prompts "a corgi astronaut on the moon, cinematic" \
+        --output ./out
+
+    python examples/offline_inference/lance/end2end.py \
+        --model bytedance-research/Lance --modality img2text \
+        --image-path /path/to/photo.jpg \
+        --prompts "Describe this image in detail."
+"""
+
+import argparse
+import os
+
+
+def parse_args():
+    p = argparse.ArgumentParser()
+    p.add_argument("--model", default="bytedance-research/Lance", help="HF repo or local path.")
+    p.add_argument("--prompts", nargs="+", default=None, help="Input text prompts.")
+    p.add_argument("--txt-prompts", type=str, default=None, help="File with one prompt per line.")
+    p.add_argument(
+        "--modality",
+        default="text2img",
+        choices=["text2img", "img2text", "text2video", "video2text", "img2img", "video2video", "image2video"],
+        help="Lance single-stage modality.",
+    )
+    p.add_argument("--image-path", type=str, default=None, help="Input image for img2text.")
+    p.add_argument("--video-path", type=str, default=None, help="Input video file for video2text.")
+    p.add_argument("--num-frames", type=int, default=25, help="Number of RGB frames for text2video (max 121).")
+    p.add_argument("--video-height", type=int, default=480, help="Video frame height.")
+    p.add_argument("--video-width", type=int, default=768, help="Video frame width.")
+    p.add_argument("--height", type=int, default=None, help="Image height (t2i). Default = max_hw (1024).")
+    p.add_argument("--width", type=int, default=None, help="Image width (t2i). Default = max_hw (1024).")
+    p.add_argument(
+        "--fps", type=int, default=12, help="Output video FPS when saving MP4 (matches upstream Lance's save_fps=12)."
+    )
+    p.add_argument("--output", type=str, default=".", help="Output directory.")
+    p.add_argument("--steps", type=int, default=30, help="Denoising steps (Lance default 30).")
+    p.add_argument("--cfg-text-scale", type=float, default=4.0, help="Text CFG scale (Lance default 4.0).")
+    p.add_argument("--timestep-shift", type=float, default=3.5, help="Flow-match timestep shift (Lance default 3.5).")
+    p.add_argument("--negative-prompt", type=str, default=None)
+    p.add_argument("--seed", type=int, default=42)
+    p.add_argument("--max-text-tokens", type=int, default=512, help="Max tokens for img2text generation.")
+    p.add_argument(
+        "--do-sample",
+        action="store_true",
+        default=True,
+        help="Sample (vs greedy) for x2t generation; greedy frequently emits an immediate EOS for Lance.",
+    )
+    p.add_argument("--no-sample", dest="do_sample", action="store_false", help="Disable sampling (greedy decoding).")
+    p.add_argument(
+        "--text-temperature",
+        type=float,
+        default=0.8,
+        help="Sampling temperature for x2t generation. Lance frequently emits an immediate EOS below ~0.7; 0.8 is a good default.",
+    )
+    p.add_argument(
+        "--system-prompt",
+        type=str,
+        default=None,
+        help=(
+            "Override the task system prompt (used by x2t_image / x2t_video for "
+            'per-example QA instructions like "Look at the image carefully and '
+            'answer the question."; without it x2t falls back to a '
+            "caption-style default and the model describes instead of answering)."
+        ),
+    )
+    # Lance is single-stage diffusion — no deploy YAML needed.  Required
+    # engine knobs (``pipeline``, ``enforce_eager``, ``trust_remote_code``,
+    # ``max_num_seqs=1`` …) are passed as flat kwargs to ``Omni`` below
+    # and ``create_default_diffusion`` materializes the stage config.
+
+    from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
+
+    nullify_stage_engine_defaults(p)
+    return p.parse_args()
+
+
+from vllm_omni.diffusion.models.lance.prompts import (  # noqa: E402
+    VIDEO_PAD,
+    VISION_END,
+    VISION_START,
+    render_lance_prompt,
+)
+
+# Visual block matching upstream Lance: <|vision_start|><|video_pad|><|vision_end|>
+# (Lance's renderer uses ``<|video_pad|>`` even for image inputs by default.)
+_VISION_BLOCK = f"{VISION_START}{VIDEO_PAD}{VISION_END}"
+
+
+def _format_text2img_prompts(prompts):
+    return [{"prompt": render_lance_prompt("t2i", p), "modalities": ["image"]} for p in prompts]
+
+
+def _format_text2video_prompts(prompts, num_frames, video_h, video_w):
+    extra = {"num_frames": num_frames, "video_height": video_h, "video_width": video_w}
+    return [{"prompt": render_lance_prompt("t2v", p), "modalities": ["video"], "extra_args": extra} for p in prompts]
+
+
+def _format_image2video_prompts(prompts, image_path, num_frames, video_h, video_w):
+    """Image-to-Video (no first-frame pin): image + text → long video.
+
+    Passes the input image via ``multi_modal_data.first_frame``.  The pipeline
+    treats it as a 1-frame reference (VAE+ViT prefill) and generates a fresh
+    multi-frame video at the requested ``num_frames`` × ``video_h`` ×
+    ``video_w`` shape.
+    """
+    import os as _os
+
+    if not image_path or not _os.path.exists(image_path):
+        raise ValueError(f"image2video requires --image-path pointing to an existing file, got: {image_path}")
+    from PIL import Image as _Image
+
+    img = _Image.open(image_path).convert("RGB")
+    extra = {"num_frames": num_frames, "video_height": video_h, "video_width": video_w}
+    return [
+        {
+            "prompt": render_lance_prompt("i2v", p, vision_token=_VISION_BLOCK),
+            "multi_modal_data": {"first_frame": img},
+            "modalities": ["video"],
+            "extra_args": extra,
+        }
+        for p in prompts
+    ]
+
+
+def _format_img2img_prompts(prompts, image_path):
+    import os as _os
+
+    if not image_path or not _os.path.exists(image_path):
+        raise ValueError(f"img2img requires --image-path pointing to an existing file, got: {image_path}")
+    from PIL import Image as _Image
+
+    img = _Image.open(image_path).convert("RGB")
+    return [
+        {
+            "prompt": render_lance_prompt("image_edit", p, vision_token=_VISION_BLOCK),
+            "multi_modal_data": {"img2img": img},
+            "modalities": ["image"],
+        }
+        for p in prompts
+    ]
+
+
+def _format_video2video_prompts(prompts, video_path, num_frames, video_h, video_w):
+    import os as _os
+
+    if not video_path or not _os.path.exists(video_path):
+        raise ValueError(f"video2video requires --video-path pointing to an existing file, got: {video_path}")
+
+    extra = {"num_frames": num_frames, "video_height": video_h, "video_width": video_w}
+    return [
+        {
+            "prompt": render_lance_prompt("video_edit", p, vision_token=_VISION_BLOCK),
+            "multi_modal_data": {"video": video_path},
+            "modalities": ["video"],
+            "extra_args": extra,
+        }
+        for p in prompts
+    ]
+
+
+def _format_video2text_prompts(prompts, video_path, system_prompt=None):
+    import os as _os
+
+    if not video_path or not _os.path.exists(video_path):
+        raise ValueError(f"video2text requires --video-path pointing to an existing file, got: {video_path}")
+    import imageio.v3 as iio
+
+    video = iio.imread(video_path)
+    return [
+        {
+            "prompt": render_lance_prompt("x2t_video", p, vision_token=_VISION_BLOCK, system_prompt=system_prompt),
+            "multi_modal_data": {"video": video},
+            "modalities": ["text"],
+        }
+        for p in prompts
+    ]
+
+
+def _format_img2text_prompts(prompts, image_path, system_prompt=None):
+    from PIL import Image
+
+    if not image_path or not os.path.exists(image_path):
+        raise ValueError(f"img2text requires --image-path pointing to an existing file, got: {image_path}")
+    img = Image.open(image_path).convert("RGB")
+    return [
+        {
+            "prompt": render_lance_prompt("x2t_image", p, vision_token=_VISION_BLOCK, system_prompt=system_prompt),
+            "multi_modal_data": {"image": img},
+            "modalities": ["text"],
+        }
+        for p in prompts
+    ]
+
+
+def main():
+    args = parse_args()
+    os.makedirs(args.output, exist_ok=True)
+
+    # Lance aligns more naturally with upstream when using FLASH_ATTN
+    # (upstream uses flash_attn_varlen_func; SDPA accumulates ~5x more
+    # numerical drift through the 36-layer Qwen2 stack on B300).  Default
+    # to flash_attn when the user hasn't explicitly set a backend.
+    os.environ.setdefault("DIFFUSION_ATTENTION_BACKEND", "FLASH_ATTN")
+
+    if args.txt_prompts:
+        with open(args.txt_prompts, encoding="utf-8") as f:
+            prompts = [ln.strip() for ln in f if ln.strip()]
+    elif args.modality == "img2text":
+        prompts = args.prompts or ["Describe this image in detail."]
+    else:
+        prompts = args.prompts or ["A cute cat, studio lighting, 4k"]
+
+    from vllm_omni.entrypoints.omni import Omni
+
+    # Video modalities require the Lance_3B_Video checkpoint (3D
+    # latent_pos_embed shaped (126976, 2048)).  Without it, vllm-omni silently
+    # loads Lance_3B (4096, 2048) — the image-only table — and any t_lat >= 1
+    # immediately indexes out of bounds.  Rewrite the model path so
+    # LancePipeline._select_video_variant picks the video subfolder.
+    resolved_model = args.model
+    if args.modality in ("text2video", "video2video", "video2text", "image2video"):
+        if not str(resolved_model).rstrip("/").endswith("Lance_3B_Video"):
+            candidate = os.path.join(resolved_model, "Lance_3B_Video")
+            if os.path.isdir(candidate):
+                resolved_model = candidate
+
+    omni_kwargs = vars(args).copy()
+    omni_kwargs["model"] = resolved_model
+    # Lance single-stage defaults (formerly in vllm_omni/deploy/lance.yaml):
+    omni_kwargs.setdefault("pipeline", "lance")
+    omni_kwargs.setdefault("max_num_batched_tokens", 32768)
+    omni_kwargs.setdefault("max_num_seqs", 1)
+    omni_kwargs.setdefault("enforce_eager", True)
+    omni_kwargs.setdefault("trust_remote_code", True)
+    omni_kwargs.setdefault("enable_prefix_caching", False)
+    omni_kwargs.setdefault("async_chunk", False)
+    omni = Omni(**omni_kwargs)
+
+    if args.modality == "img2text":
+        formatted = _format_img2text_prompts(prompts, args.image_path, system_prompt=args.system_prompt)
+    elif args.modality == "text2video":
+        formatted = _format_text2video_prompts(prompts, args.num_frames, args.video_height, args.video_width)
+    elif args.modality == "video2text":
+        formatted = _format_video2text_prompts(prompts, args.video_path, system_prompt=args.system_prompt)
+    elif args.modality == "img2img":
+        formatted = _format_img2img_prompts(prompts, args.image_path)
+    elif args.modality == "video2video":
+        formatted = _format_video2video_prompts(
+            prompts, args.video_path, args.num_frames, args.video_height, args.video_width
+        )
+    elif args.modality == "image2video":
+        formatted = _format_image2video_prompts(
+            prompts, args.image_path, args.num_frames, args.video_height, args.video_width
+        )
+    else:
+        formatted = _format_text2img_prompts(prompts)
+
+    params_list = omni.default_sampling_params_list
+    diffusion_params = params_list[0]  # single-stage: one param set
+    diffusion_params.num_inference_steps = args.steps  # type: ignore
+    if args.seed is not None:
+        diffusion_params.seed = args.seed  # type: ignore
+    if args.height is not None:
+        diffusion_params.height = args.height  # type: ignore
+    if args.width is not None:
+        diffusion_params.width = args.width  # type: ignore
+    extra = getattr(diffusion_params, "extra_args", {}) or {}
+    extra["cfg_text_scale"] = args.cfg_text_scale
+    extra["timestep_shift"] = args.timestep_shift
+    if args.modality in ("img2text", "video2text"):
+        extra["max_think_tokens"] = args.max_text_tokens
+        extra["do_sample"] = args.do_sample
+        extra["text_temperature"] = args.text_temperature
+    if args.negative_prompt is not None:
+        extra["negative_prompt"] = args.negative_prompt
+    diffusion_params.extra_args = extra  # type: ignore
+
+    outputs = list(omni.generate(prompts=formatted, sampling_params_list=params_list))
+
+    if args.modality in ("img2text", "video2text"):
+        for i, req_output in enumerate(outputs):
+            text = getattr(req_output, "output", None) or getattr(req_output, "text", None)
+            if text is None and hasattr(req_output, "custom_output"):
+                text = (req_output.custom_output or {}).get("text_output")
+            print(f"[Output {i}] {text}")
+        return
+
+    if args.modality in ("text2video", "video2video", "image2video"):
+        for i, req_output in enumerate(outputs):
+            custom = getattr(req_output, "custom_output", None) or {}
+            frames = custom.get("video_frames")
+            if not frames:
+                print(f"[Output {i}] no video frames returned")
+                continue
+            mp4_path = os.path.join(args.output, f"lance_{i}.mp4")
+            try:
+                import imageio.v3 as iio
+                import numpy as np
+
+                arr = np.stack([np.asarray(f) for f in frames], axis=0)
+                iio.imwrite(mp4_path, arr, fps=args.fps, codec="libx264", quality=8)
+                print(f"[Output {i}] Saved MP4 ({len(frames)} frames, {arr.shape[1]}x{arr.shape[2]}) to {mp4_path}")
+            except Exception as exc:
+                print(f"[Output {i}] MP4 encode failed ({exc!r}); writing per-frame ONGs instead")
+                frame_dir = os.path.join(args.output, f"lance_{i}_frames")
+                os.makedirs(frame_dir, exist_ok=True)
+                for j, frame in enumerate(frames):
+                    frame.save(os.path.join(frame_dir, f"{j:04d}.png"))
+                print(f"  → {frame_dir} ({len(frames)} frames)")
+        return
+
+    idx = 0
+    for req_output in outputs:
+        images = getattr(req_output, "images", None)
+        if not images:
+            continue
+        for j, img in enumerate(images):
+            path = os.path.join(args.output, f"lance_{idx}_{j}.png")
+            img.save(path)
+            print(f"[Output] Saved image to {path}")
+        idx += 1
+
+
+if __name__ == "__main__":
+    main()

examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py

@@ -19,7 +19,6 @@ from vllm import SamplingParams
 from vllm.multimodal.image import convert_image_mode
 
 from vllm_omni import Omni
-from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 
 DEFAULT_SYSTEM = "You are a helpful assistant."
 DEFAULT_QUESTION = "Please summarize the content of this image."
@@ -49,7 +48,6 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable diffusion pipeline profiler to display stage durations.",
     )
-    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_t2i.py

@@ -29,7 +29,6 @@ from PIL import Image
 from vllm.sampling_params import SamplingParams
 
 from vllm_omni import Omni
-from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 logger = logging.getLogger(__name__)
@@ -118,7 +117,6 @@ def parse_args() -> argparse.Namespace:
     )
     p.add_argument("--out", type=str, default="output.png", help="Path to save the generated image.")
     p.add_argument("--trust-remote-code", action="store_true", help="Trust remote code when loading the model.")
-    nullify_stage_engine_defaults(p)
     args = p.parse_args()
     if not args.prompt:
         args.prompt = ["A stylish woman with sunglasses riding a motorcycle in NYC."]

examples/offline_inference/mimo_audio/end2end.py

@@ -21,11 +21,10 @@ from message_convert import (
     to_prompt,
 )
 from vllm import SamplingParams
-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.inputs.data import OmniTokensPrompt
+from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 
 SEED = 42
 
@@ -343,7 +342,7 @@ def main(args):
 
 
 def parse_args():
-    parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
+    parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
     parser.add_argument(
         "--model-name",
         "-m",
@@ -439,7 +438,6 @@ def parse_args():
         "vllm_omni/deploy/mimo_audio.yaml based on the HF model_type.",
     )
 
-    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/ming_flash_omni/README.md

@@ -53,7 +53,7 @@ python examples/offline_inference/ming_flash_omni/end2end.py \
     --image-path ./3_0.png
 ```
 
-### Omni-Speech (Thinker + Talker)
+### Omni-Speech (Thinker + Talker) {#omni-speech-thinker--talker}
 
 The default deploy YAML already runs thinker+talker, so spoken output only requires requesting `audio` (or `text,audio`) modalities.
 The thinker processes your multimodal input, generates text, then the talker synthesises the response as speech.

examples/offline_inference/ming_flash_omni/end2end.py

@@ -17,14 +17,13 @@ from vllm.assets.image import ImageAsset
 from vllm.assets.video import VideoAsset, video_to_ndarrays
 from vllm.multimodal.image import convert_image_mode
 from vllm.multimodal.media.audio import load_audio
-from vllm.utils.argparse_utils import FlexibleArgumentParser
 
 import vllm_omni
-from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 from vllm_omni.entrypoints.omni import Omni
 
 # Imports the processor also registers itself
 from vllm_omni.transformers_utils.processors.ming import MingFlashOmniProcessor  # noqa: F401
+from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 
 SEED = 42
 MODEL_NAME = "Jonathan1909/Ming-flash-omni-2.0"
@@ -397,7 +396,7 @@ def main(args):
 
 
 def parse_args():
-    parser = FlexibleArgumentParser(description="Ming-flash-omni 2.0 offline inference example")
+    parser = TrackingArgumentParser(description="Ming-flash-omni 2.0 offline inference example")
     parser.add_argument(
         "--query-type",
         "-q",
@@ -496,7 +495,6 @@ def parse_args():
         help="Output directory for results.",
     )
 
-    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/qwen2_5_omni/end2end.py

@@ -19,10 +19,9 @@ from vllm.assets.video import VideoAsset, video_to_ndarrays
 from vllm.multimodal.image import convert_image_mode
 from vllm.multimodal.media.audio import load_audio
 from vllm.sampling_params import SamplingParams
-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.utils.tracking_parser import TrackingArgumentParser
 
 SEED = 42
 
@@ -432,7 +431,7 @@ def main(args):
 
 
 def parse_args():
-    parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
+    parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
     parser.add_argument(
         "--model",
         type=str,
@@ -554,7 +553,6 @@ def parse_args():
         default=False,
         help="Use py_generator mode. The returned type of Omni.generate() is a Python Generator object.",
     )
-    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/qwen3_omni/end2end.py

@@ -19,10 +19,9 @@ from vllm.assets.image import ImageAsset
 from vllm.assets.video import VideoAsset, video_to_ndarrays
 from vllm.multimodal.image import convert_image_mode
 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.utils.tracking_parser import TrackingArgumentParser
 
 SEED = 42
 
@@ -422,7 +421,7 @@ def main(args):
 
 
 def parse_args():
-    parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
+    parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
     parser.add_argument(
         "--model",
         type=str,
@@ -561,7 +560,6 @@ def parse_args():
         help="Model dtype (auto, half, float16, bfloat16, float, float32).",
     )
 
-    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/qwen3_omni/end2end_async_chunk.py

@@ -39,10 +39,9 @@ from vllm.assets.image import ImageAsset
 from vllm.assets.video import VideoAsset, video_to_ndarrays
 from vllm.multimodal.image import convert_image_mode
 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.async_omni import AsyncOmni
+from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 
 logger = logging.getLogger(__name__)
 
@@ -383,7 +382,9 @@ async def run_all(args):
     print(f"[Info] Creating AsyncOmni with deploy_config={args.deploy_config}")
     async_omni = None
     try:
-        async_omni = AsyncOmni(**vars(args))
+        # ``from_cli_args`` forwards only explicitly-passed CLI args so
+        # argparse defaults do not silently override deploy YAML values.
+        async_omni = AsyncOmni.from_cli_args(args, model=args.model)
 
         # Use default sampling params from stage config (they are pre-configured
         # in the YAML for each stage).
@@ -453,7 +454,7 @@ async def run_all(args):
 
 
 def parse_args():
-    parser = FlexibleArgumentParser(
+    parser = TrackingArgumentParser(
         description=(
             "Offline inference with async_chunk enabled via AsyncOmni. "
             "Downstream stages start before upstream stages finish, "
@@ -589,7 +590,6 @@ def parse_args():
         default=16000,
         help="Sampling rate for audio loading.",
     )
-    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/sensenova_u1/end2end.py

@@ -29,7 +29,6 @@ import os
 
 from PIL import Image
 
-from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
@@ -168,7 +167,6 @@ def parse_args():
         help="Enable cache-dit summary logging after diffusion forward passes.",
     )
 
-    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/text_to_image/text_to_image.py

@@ -307,9 +307,6 @@ def parse_args() -> argparse.Namespace:
         help="Supplementary auxiliary text encoder parameters model name or path (especially for Hidream-l1-full).",
     )
     current_omni_platform.pre_register_and_update(parser)
-    from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
-
-    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/text_to_speech/cosyvoice3/end2end.py

@@ -10,7 +10,6 @@ import soundfile as sf
 from vllm import SamplingParams
 from vllm.multimodal.media.audio import load_audio
 
-from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer
 from vllm_omni.model_executor.models.cosyvoice3.utils import extract_text_token
@@ -65,7 +64,6 @@ def run_e2e():
         required=True,
         help="Path to tokenizer directory (e.g., <model_path>/CosyVoice-BlankEN).",
     )
-    nullify_stage_engine_defaults(parser)
     args = parser.parse_args()
     # Ensure tokenizer directory exists
     if not os.path.exists(args.tokenizer):

examples/offline_inference/text_to_speech/fish_speech/end2end.py

@@ -26,16 +26,14 @@ import torch
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
-from vllm.utils.argparse_utils import FlexibleArgumentParser
-
 from vllm_omni import AsyncOmni, Omni
-from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 from vllm_omni.model_executor.models.fish_speech.dac_utils import DAC_HOP_LENGTH, DAC_SAMPLE_RATE
 from vllm_omni.model_executor.models.fish_speech.prompt_utils import (
     build_fish_text_only_prompt_ids,
     estimate_fish_voice_clone_prompt_len_from_normalized,
     normalize_fish_voice_clone_texts,
 )
+from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 
 logger = logging.getLogger(__name__)
 
@@ -211,7 +209,7 @@ async def main_streaming(args):
 
 
 def parse_args():
-    parser = FlexibleArgumentParser(
+    parser = TrackingArgumentParser(
         description="Fish Speech S2 Pro offline inference with vLLM Omni",
     )
     parser.add_argument(
@@ -267,7 +265,6 @@ def parse_args():
         default=False,
         help="Stream audio chunks as they arrive via AsyncOmni.",
     )
-    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py

@@ -30,10 +30,9 @@ from pathlib import Path
 import numpy as np
 import soundfile as sf
 import torch
-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.tracking_parser import TrackingArgumentParser
 
 SAMPLE_RATE = 24_000
 DEFAULT_TEXTS = (
@@ -43,7 +42,7 @@ DEFAULT_TEXTS = (
 
 
 def parse_args():
-    parser = FlexibleArgumentParser(description="Offline higgs-audio v2 inference")
+    parser = TrackingArgumentParser(description="Offline higgs-audio v2 inference")
     parser.add_argument(
         "--model",
         type=str,
@@ -88,7 +87,6 @@ def parse_args():
         default=None,
         help="Transcript of the reference clip. Required when --ref-audio is set.",
     )
-    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/text_to_speech/ming_flash_omni_tts/end2end.py

@@ -5,11 +5,11 @@ from typing import Any
 
 import soundfile as sf
 import torch
-from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
-from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniTokensPrompt
 from vllm_omni.model_executor.models.ming_flash_omni.prompt_utils import (
@@ -68,7 +68,7 @@ def save_audio(mm: dict[str, Any], output_path: str) -> None:
 
 
 def parse_args():
-    parser = FlexibleArgumentParser(description="Ming-flash-omni standalone talker offline e2e example")
+    parser = TrackingArgumentParser(description="Ming-flash-omni standalone talker offline e2e example")
     parser.add_argument("--model", type=str, default=MODEL_NAME, help="Model name or local path.")
     parser.add_argument(
         "--deploy-config",
@@ -89,7 +89,6 @@ def parse_args():
     parser.add_argument("--init-timeout", type=int, default=600, help="Engine init timeout in seconds.")
     parser.add_argument("--stage-init-timeout", type=int, default=300, help="Single stage init timeout in seconds.")
 
-    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/text_to_speech/moss_tts/end2end.py

@@ -0,0 +1,377 @@
+"""Offline inference example for MOSS-TTS variants.
+
+Covers all five models:
+  - OpenMOSS-Team/MOSS-TTS           (8B, general TTS)
+  - OpenMOSS-Team/MOSS-TTSD-v1.0     (8B, dialogue TTS)
+  - OpenMOSS-Team/MOSS-SoundEffect   (8B, sound effect synthesis)
+  - OpenMOSS-Team/MOSS-TTS-Realtime  (1.7B, low-latency streaming)
+  - OpenMOSS-Team/MOSS-VoiceGenerator(1.7B, voice design)
+
+Usage
+-----
+# Standard TTS with voice cloning:
+python end2end.py \\
+    --model OpenMOSS-Team/MOSS-TTS \\
+    --text "Hello, this is a MOSS-TTS test." \\
+    --ref-audio path/to/reference.wav \\
+    --output output.wav
+"""
+
+from __future__ import annotations
+
+import argparse
+import gc
+import sys
+
+import soundfile as sf
+import torch
+
+
+def _load_ref_audio(path: str, target_sr: int = 24000) -> torch.Tensor:
+    """Load and resample reference audio to (1, T) at target_sr."""
+    data, sr = sf.read(path, always_2d=True)  # (T, C)
+    wav = torch.from_numpy(data.T.astype("float32"))  # (C, T)
+    if wav.shape[0] > 1:
+        wav = wav.mean(dim=0, keepdim=True)
+    if sr != target_sr:
+        import torchaudio
+
+        wav = torchaudio.functional.resample(wav, sr, target_sr)
+    return wav  # (1, T)
+
+
+def _build_unified_codes(
+    model_id: str,
+    text: str,
+    ref_audio: torch.Tensor | None,
+    *,
+    instruction: str | None = None,
+    tokens: int | None = None,
+    quality: str | None = None,
+    sound_event: str | None = None,
+    ambient_sound: str | None = None,
+    language: str | None = None,
+) -> tuple[list[int], torch.Tensor, int]:
+    """Build the upstream MossTTSDelay unified-codes prompt.
+
+    Returns (text_token_ids, audio_codes_LxNQ, n_vq) where
+    ``text_token_ids`` is the prefill prompt for the talker and
+    ``audio_codes_LxNQ`` is the matching delay-pattern audio code grid
+    (``audio_pad_code`` everywhere except inside the reference-audio block).
+    """
+    from transformers import AutoProcessor
+
+    proc = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
+
+    reference: list[torch.Tensor] | None = None
+    if ref_audio is not None:
+        # Pre-encode so the processor doesn't need torchaudio.load (which
+        # now requires torchcodec on recent torchaudio releases).
+        # encode_audios_from_wav expects (C, T) per item.
+        wav_2d = ref_audio if ref_audio.dim() == 2 else ref_audio.unsqueeze(0)
+        codes_list = proc.encode_audios_from_wav([wav_2d], sampling_rate=24000, n_vq=proc.model_config.n_vq)
+        reference = [codes_list[0]]
+
+    user_msg_kwargs: dict = {"text": text}
+    if reference is not None:
+        user_msg_kwargs["reference"] = reference
+    if instruction is not None:
+        user_msg_kwargs["instruction"] = instruction
+    if tokens is not None:
+        user_msg_kwargs["tokens"] = tokens
+    if quality is not None:
+        user_msg_kwargs["quality"] = quality
+    if sound_event is not None:
+        user_msg_kwargs["sound_event"] = sound_event
+    if ambient_sound is not None:
+        user_msg_kwargs["ambient_sound"] = ambient_sound
+    if language is not None:
+        user_msg_kwargs["language"] = language
+
+    user_msg = proc.build_user_message(**user_msg_kwargs)
+    batch = proc(conversations=[[user_msg]], mode="generation")
+    unified = batch["input_ids"][0]  # (L, 1+n_vq)
+    n_vq = int(unified.shape[-1] - 1)
+    text_ids = unified[:, 0].tolist()
+    audio_codes = unified[:, 1:].contiguous().to(torch.int64)  # (L, n_vq)
+
+    # Free the processor (incl. ~1.6B-param audio_tokenizer on CPU) before
+    # we instantiate Omni — we only needed it for the one-shot encode.
+    del proc
+    gc.collect()
+
+    return text_ids, audio_codes, n_vq
+
+
+def _build_realtime_prompt(
+    model_id: str,
+    text: str,
+    ref_audio: torch.Tensor | None,
+) -> tuple[list[int], torch.Tensor, int]:
+    """Build the prefill prompt for MOSS-TTS-Realtime.
+
+    Realtime uses a different prompt format than the delay variants: a system
+    prompt + (optional) voice-clone context with the encoded reference audio
+    inlined, followed by the user text and an assistant header. The format is
+    a ``(L, channels+1=17)`` int grid where col 0 is text/special tokens and
+    cols 1..16 are RVQ codebook entries (``audio_channel_pad`` outside the
+    ref-audio block).
+
+    The upstream ``MossTTSRealtimeProcessor`` is shipped via
+    ``trust_remote_code`` but is not auto-discovered by ``AutoProcessor``
+    (no ``processor_config.json``) — we import its module directly from the
+    snapshot directory and use its ``make_ensemble``/``make_user_prompt``.
+    """
+    import importlib.util
+    import sys as _sys
+    from pathlib import Path
+
+    import numpy as np
+    import torch
+    from huggingface_hub import snapshot_download
+    from transformers import AutoModel, AutoTokenizer
+
+    snap_dir = Path(snapshot_download(repo_id=model_id))
+    proc_module_path = snap_dir / "processing_mossttsrealtime.py"
+    if not proc_module_path.exists():
+        raise FileNotFoundError(f"realtime processor module missing at {proc_module_path}")
+
+    # Load the processor module without polluting sys.path globally.
+    spec = importlib.util.spec_from_file_location("_moss_tts_realtime_proc", proc_module_path)
+    mod = importlib.util.module_from_spec(spec)  # type: ignore[arg-type]
+    _sys.modules[spec.name] = mod  # type: ignore[union-attr]
+    spec.loader.exec_module(mod)  # type: ignore[union-attr]
+
+    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+    realtime_processor = mod.MossTTSRealtimeProcessor(tokenizer=tokenizer)
+
+    # Encode reference audio via the standalone MOSS Audio Tokenizer (CPU is
+    # fine — it runs once per request).
+    audio_tokens = None
+    if ref_audio is not None:
+        codec = AutoModel.from_pretrained("OpenMOSS-Team/MOSS-Audio-Tokenizer", trust_remote_code=True).eval()
+        wav_2d = ref_audio if ref_audio.dim() == 2 else ref_audio.unsqueeze(0)
+        with torch.no_grad():
+            enc = codec.batch_encode([wav_2d.squeeze(0)], num_quantizers=16)
+        # enc.audio_codes: (NQ, B, T); we want (T, NQ) numpy.
+        codes = enc.audio_codes[:, 0, : int(enc.audio_codes_lengths[0].item())]
+        audio_tokens = codes.transpose(0, 1).contiguous().cpu().numpy()  # (T, 16)
+        del codec
+        gc.collect()
+
+    # Build the inference-style prefill (mirrors upstream
+    # ``MossTTSRealtimeInference._build_prefill_batch``):
+    #   system_prompt [+ voice_clone_block]
+    #   <|im_start|>assistant\n
+    #   <text_tokens...>            (col 0 = text, cols 1..16 = audio_channel_pad)
+    #   <last text token>           (col 1 = audio_bos at the last text position)
+    # Generation continues from here; the local transformer emits real audio
+    # codes for codebook 0 directly (no BOS frame at gen_step=0).
+    system_grid = realtime_processor.make_ensemble(prompt_audio_tokens=audio_tokens)  # (L_sys, 17)
+
+    # Append "<|im_start|>assistant\n" header (the snapshot's make_ensemble
+    # omits it; inferencer.py's version includes it. We add it here.)
+    assistant_header_ids = tokenizer.encode("<|im_start|>assistant\n", add_special_tokens=False)
+    assistant_grid = np.full((len(assistant_header_ids), 17), 1024, dtype=np.int64)
+    assistant_grid[:, 0] = assistant_header_ids
+
+    # Append the FIRST 12 text tokens with audio_bos (1025) at the last
+    # included text position (audio column 0 == grid col 1). The remaining
+    # text tokens are fed one-per-step during generation via the talker's
+    # compute_logits → vLLM sampler pipeline (see ``ids.all`` in the
+    # additional_information). This matches upstream
+    # ``MossTTSRealtimeInference._build_prefill_batch`` (prefill_max_text=12).
+    text_ids_only = tokenizer.encode(text, add_special_tokens=False)
+    if not text_ids_only:
+        raise ValueError("Empty text after tokenisation")
+    PREFILL_MAX_TEXT = 12
+    cur_len = min(len(text_ids_only), PREFILL_MAX_TEXT)
+    text_grid = np.full((cur_len, 17), 1024, dtype=np.int64)
+    text_grid[:, 0] = text_ids_only[:cur_len]
+    text_grid[-1, 1] = 1025  # audio_bos_token at last *prefilled* text token
+
+    grid = np.concatenate([system_grid, assistant_grid, text_grid], axis=0)
+    text_ids = grid[:, 0].tolist()
+    audio_codes = torch.from_numpy(grid[:, 1:].astype(np.int64).copy())  # (L, 16)
+
+    # ``remaining_text_ids`` will be fed one-per-step by the talker.
+    remaining_text_ids = list(text_ids_only[cur_len:])
+    return text_ids, audio_codes, 16, remaining_text_ids
+
+
+def run_tts(args: argparse.Namespace) -> None:
+    from vllm import SamplingParams
+
+    from vllm_omni import Omni
+
+    deploy_map = {
+        "OpenMOSS-Team/MOSS-TTS": "moss_tts",
+        "OpenMOSS-Team/MOSS-TTSD-v1.0": "moss_ttsd",
+        "OpenMOSS-Team/MOSS-SoundEffect": "moss_sound_effect",
+        "OpenMOSS-Team/MOSS-TTS-Realtime": "moss_tts_realtime",
+        "OpenMOSS-Team/MOSS-VoiceGenerator": "moss_voice_generator",
+    }
+    deploy_config = deploy_map.get(args.model, "moss_tts")
+
+    is_sound_effect = "SoundEffect" in args.model
+    is_voice_gen = "VoiceGenerator" in args.model
+    is_realtime = "Realtime" in args.model
+
+    ref_audio = None
+    builder_kwargs: dict = {}
+
+    remaining_text_ids: list[int] | None = None
+    if is_realtime:
+        # MOSS-TTS-Realtime uses a different prompt format and has its own
+        # per-step depth transformer; the delay-style ``MossTTSDelayProcessor``
+        # path doesn't apply.
+        if not args.ref_audio:
+            sys.exit("MOSS-TTS-Realtime currently requires --ref-audio (voice clone).")
+        ref_audio = _load_ref_audio(args.ref_audio)
+        text = args.text or ""
+        text_ids, audio_codes, n_vq, remaining_text_ids = _build_realtime_prompt(args.model, text, ref_audio)
+        print(
+            f"Prefill prompt: {len(text_ids)} tokens, {n_vq}-quantizer audio block (realtime); "
+            f"{len(remaining_text_ids)} text tokens to stream"
+        )
+    elif is_sound_effect:
+        if not args.ambient_sound:
+            sys.exit("MOSS-SoundEffect requires --ambient-sound")
+        builder_kwargs["ambient_sound"] = args.ambient_sound
+        # MOSS-SoundEffect emits ~12.5 codec frames per second; ``tokens``
+        # tells the AR loop the target duration.
+        if args.duration_tokens:
+            builder_kwargs["tokens"] = args.duration_tokens
+        elif args.duration_seconds:
+            builder_kwargs["tokens"] = max(1, int(float(args.duration_seconds) * 12.5))
+        text = args.text or ""
+    elif is_voice_gen:
+        # MOSS-VoiceGenerator needs an instruction (voice description) plus
+        # the text to synthesise; no reference audio.
+        if not args.instruction:
+            sys.exit("MOSS-VoiceGenerator requires --instruction (voice description)")
+        builder_kwargs["instruction"] = args.instruction
+        text = args.text or ""
+    else:
+        if not args.ref_audio:
+            sys.exit("Voice cloning requires --ref-audio")
+        ref_audio = _load_ref_audio(args.ref_audio)
+        text = args.text or ""
+
+    if not is_realtime:
+        text_ids, audio_codes, n_vq = _build_unified_codes(args.model, text, ref_audio, **builder_kwargs)
+        print(f"Prefill prompt: {len(text_ids)} tokens, {n_vq}-quantizer audio block")
+
+    omni = Omni(model=args.model, deploy_config=deploy_config, stage_init_timeout=600)
+
+    talker_sp = SamplingParams(
+        temperature=1.7,
+        top_p=0.8,
+        top_k=25,
+        max_tokens=2048,
+        seed=42,
+        detokenize=False,
+    )
+    codec_sp = SamplingParams(
+        temperature=0.0,
+        top_p=1.0,
+        top_k=-1,
+        max_tokens=65536,
+        seed=42,
+        detokenize=True,
+    )
+
+    additional_info: dict = {
+        # Reference-audio codes consumed by the talker prefill (additive
+        # embedding); shape is (L, n_vq) where pad_code marks non-audio
+        # positions. The OmniPayload schema reserves ``codes.ref`` for
+        # exactly this kind of conditioning input.
+        "codes": {"ref": audio_codes},
+    }
+    if remaining_text_ids:
+        # Realtime variant: feed the remaining text tokens (everything past
+        # the first ``PREFILL_MAX_TEXT``) one-per-step during decode via the
+        # talker's compute_logits → vLLM sampler one-hot trick.
+        additional_info["ids"] = {"all": remaining_text_ids}
+
+    request = {
+        "prompt_token_ids": text_ids,
+        "additional_information": additional_info,
+    }
+
+    import time as _time
+
+    chunks: list[torch.Tensor] = []
+    t_gen_start = _time.monotonic()
+    t_first_audio: float | None = None
+    for stage_outputs in omni.generate(request, [talker_sp, codec_sp]):
+        # Each yielded ``stage_outputs`` is an OmniRequestOutput; the codec
+        # stage attaches its waveform under ``multimodal_output["model_outputs"]``.
+        # Avoid ``a or b`` when ``a`` may be a tensor (truthiness raises).
+        mm = stage_outputs.multimodal_output or {}
+        wav = mm.get("audio")
+        if wav is None:
+            wav = mm.get("model_outputs")
+        if wav is None:
+            continue
+        if isinstance(wav, list):
+            non_empty = [w.reshape(-1) for w in wav if isinstance(w, torch.Tensor) and w.numel() > 0]
+            if non_empty and t_first_audio is None:
+                t_first_audio = _time.monotonic()
+            chunks.extend(non_empty)
+        elif isinstance(wav, torch.Tensor) and wav.numel() > 0:
+            if t_first_audio is None:
+                t_first_audio = _time.monotonic()
+            chunks.append(wav.reshape(-1))
+
+    t_gen_end = _time.monotonic()
+    if not chunks:
+        print("No audio in output — check model loading logs.")
+        return
+
+    wav_t = torch.cat(chunks).float().cpu()
+    sf.write(args.output, wav_t.numpy(), 24000)
+    print(f"Saved {len(wav_t) / 24000:.2f}s of audio to {args.output}")
+    ttfa_ms = (t_first_audio - t_gen_start) * 1000.0 if t_first_audio is not None else float("nan")
+    gen_total_ms = (t_gen_end - t_gen_start) * 1000.0
+    print(f"TTFA_MS={ttfa_ms:.1f} GEN_TOTAL_MS={gen_total_ms:.1f} AUDIO_S={len(wav_t) / 24000:.2f}")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="MOSS-TTS offline inference example")
+    parser.add_argument(
+        "--model",
+        default="OpenMOSS-Team/MOSS-TTS",
+        help="HuggingFace model ID",
+    )
+    parser.add_argument("--text", default="", help="Text to synthesise (TTS variants)")
+    parser.add_argument("--ref-audio", default=None, help="Reference audio for voice cloning")
+    parser.add_argument(
+        "--ambient-sound",
+        default=None,
+        help="Sound description for MOSS-SoundEffect",
+    )
+    parser.add_argument(
+        "--duration-tokens",
+        type=int,
+        default=None,
+        help="Target duration in tokens (1 s ≈ 12.5 tokens) for MOSS-SoundEffect",
+    )
+    parser.add_argument(
+        "--duration-seconds",
+        type=float,
+        default=None,
+        help="Target duration in seconds (converted via ~12.5 tokens/s) for MOSS-SoundEffect",
+    )
+    parser.add_argument(
+        "--instruction",
+        default=None,
+        help="Voice description for MOSS-VoiceGenerator (e.g. 'a young woman with a clear voice')",
+    )
+    parser.add_argument("--output", default="output.wav", help="Output WAV path")
+    args = parser.parse_args()
+    run_tts(args)
+
+
+if __name__ == "__main__":
+    main()

examples/offline_inference/text_to_speech/moss_tts_nano/end2end.py

@@ -38,7 +38,8 @@ from pathlib import Path
 import soundfile as sf
 import torch
 from vllm import SamplingParams
-from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 
 # Prevent multiprocessing from re-importing CUDA in the wrong context.
 os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
@@ -155,7 +156,7 @@ def main(args) -> None:
 
 
 def parse_args():
-    parser = FlexibleArgumentParser(description="MOSS-TTS-Nano offline inference")
+    parser = TrackingArgumentParser(description="MOSS-TTS-Nano offline inference")
     parser.add_argument("--text", default="Hello, this is MOSS-TTS-Nano speaking.", help="Text to synthesize.")
     parser.add_argument(
         "--ref-audio",

examples/offline_inference/text_to_speech/omnivoice/end2end.py

@@ -21,7 +21,6 @@ import os
 import numpy as np
 import soundfile as sf
 
-from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
@@ -86,7 +85,6 @@ def run_e2e():
         default=None,
         help="Random seed for generation",
     )
-    nullify_stage_engine_defaults(parser)
     args = parser.parse_args()
 
     if not os.path.exists(args.stage_config):

examples/offline_inference/text_to_speech/qwen3_tts/end2end.py

@@ -15,10 +15,8 @@ import torch
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
-from vllm.utils.argparse_utils import FlexibleArgumentParser
-
 from vllm_omni import AsyncOmni, Omni
-from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
+from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 
 logger = logging.getLogger(__name__)
 
@@ -435,7 +433,7 @@ async def main_streaming(args):
 
 
 def parse_args():
-    parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
+    parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
     parser.add_argument(
         "--query-type",
         "-q",
@@ -548,7 +546,6 @@ def parse_args():
         help="Number of prompts per batch (default: 1, sequential).",
     )
 
-    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/text_to_speech/voxcpm2/end2end.py

@@ -12,17 +12,16 @@ from pathlib import Path
 
 import soundfile as sf
 import torch
-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.tracking_parser import TrackingArgumentParser
 
 REPO_ROOT = Path(__file__).resolve().parents[4]
 SAMPLE_RATE = 48_000
 
 
 def parse_args():
-    parser = FlexibleArgumentParser(description="Offline VoxCPM2 native AR inference")
+    parser = TrackingArgumentParser(description="Offline VoxCPM2 native AR inference")
     parser.add_argument(
         "--model",
         type=str,
@@ -60,7 +59,6 @@ def parse_args():
         default=None,
         help="Optional transcript of --ref-audio (enables continuation mode).",
     )
-    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/text_to_speech/voxtral_tts/end2end.py

@@ -27,11 +27,10 @@ except ImportError:
         "pip install -e ."
     )
 from vllm import SamplingParams
-from vllm.utils.argparse_utils import FlexibleArgumentParser
 
 from vllm_omni import AsyncOmni
-from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 
 logger = logging.getLogger(__name__)
 
@@ -230,7 +229,7 @@ def run_non_streaming(inputs, sampling_params_list, model_name, args, output_dir
 
 
 def parse_args() -> Namespace:
-    parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with Voxtral TTS")
+    parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with Voxtral TTS")
     parser.add_argument(
         "--model",
         type=str,
@@ -313,7 +312,6 @@ def parse_args() -> Namespace:
         default=None,
         help="Quantization method (e.g. 'fp8'). Applied to the language model only.",
     )
-    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/online_serving/bagel/README.md

@@ -68,6 +68,67 @@ vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --tensor-parallel-size
 
 Or set `tensor_parallel_size` per stage in a custom deploy YAML.
 
+### VAE Patch Parallelism
+
+[VAE Patch Parallelism](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism/vae_patch_parallel.html) distributes Bagel VAE **decode/encode** across multiple GPUs by splitting latent tiles. It lowers **per-GPU peak memory during VAE decode**, which helps high-resolution `text2img` / `img2img` when VAE becomes a bottleneck.
+
+**Scope for Bagel:**
+
+| Topology | VAE patch parallel |
+| :------- | :----------------- |
+| **Single-stage** (DiT only) | Supported on stage 0 (`BagelPipeline` + `DistributedAutoEncoder`) |
+| **Two-stage** | Supported on **stage 1 (DiT)** only; stage 0 (Thinker) uses encoder-only VAE and is unrelated |
+
+**Requirements:**
+
+- `vae_patch_parallel_size > 1` and a distributed VAE (`DistributedAutoEncoder` on the DiT pipeline).
+- The DiT process group must have at least `vae_patch_parallel_size` ranks. In practice this means the diffusion stage `world_size` must be ≥ 2 (commonly `tensor_parallel_size=2` on that stage).
+- `vae_use_tiling` must be enabled. If you set `vae_patch_parallel_size > 1` and omit tiling, the registry auto-enables `vae_use_tiling` at startup.
+
+VAE patch parallel **reuses the DiT process group** (`dit_group`); it does not create a separate VAE-only worker pool. It is not a substitute for single-GPU VAE tiling (`vae_pp=1`).
+
+**Online serving (single-stage, 2 GPUs):**
+
+```bash
+CUDA_VISIBLE_DEVICES=0,1 vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 \
+    --deploy-config vllm_omni/deploy/bagel_single_stage.yaml \
+    --tensor-parallel-size 2 \
+    --vae-patch-parallel-size 2 \
+    --vae-use-tiling
+```
+
+**Online serving (two-stage, VAE PP on DiT stage 1):** use a custom deploy YAML, for example:
+
+```yaml
+stages:
+  - stage_id: 0
+    devices: "0"
+    # Thinker (AR) — no VAE patch parallel here
+
+  - stage_id: 1
+    devices: "0,1"
+    vae_use_tiling: true
+    parallel_config:
+      tensor_parallel_size: 2
+      vae_patch_parallel_size: 2
+```
+
+```bash
+vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 \
+    --deploy-config /path/to/bagel_vae_pp.yaml
+```
+
+**Verify it is active** (check server logs at startup):
+
+```text
+INFO ... vae_patch_parallel_size=2 requires vae_use_tiling; automatically enabling it.
+```
+
+| CLI flag | Default | Description |
+| :------- | :------ | :---------- |
+| `--vae-patch-parallel-size` | `1` | Number of DiT ranks used for VAE tile parallelism. Set to `2` or higher to enable. Should be ≤ DiT process group size (typically match `--tensor-parallel-size` on the diffusion stage). |
+| `--vae-use-tiling` | off | Enable VAE spatial tiling. Required for VAE patch parallel (auto-enabled when `vae_patch_parallel_size > 1`). |
+
 #### Hybrid Sharded Data Parallel (HSDP)
 
 For larger Bagel deployments on multiple GPUs, you can enable HSDP (Hybrid Sharded Data Parallel) by modifying the stage configuration (for example, [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml)). HSDP shards transformer weights across GPUs to reduce per-GPU memory usage.
@@ -323,6 +384,8 @@ python openai_chat_client.py \
 | `stages[].gpu_memory_utilization` | per-stage | Fraction of GPU memory to use |
 | `stages[].enforce_eager` | per-stage | Disable CUDA graphs |
 | `stages[].tensor_parallel_size` | per-stage | TP degree for this stage |
+| `stages[].parallel_config.vae_patch_parallel_size` | per-stage (DiT) | VAE tile parallelism degree (DiT stage only) |
+| `stages[].vae_use_tiling` | per-stage (DiT) | Enable VAE tiling (required for VAE patch parallel) |
 | `connectors` | top-level | Define available connector instances (SHM, Mooncake) |
 | `platforms` | top-level | Platform-specific overrides (e.g. `xpu`) |
 

examples/online_serving/dreamzero/README.md

@@ -0,0 +1,132 @@
+# DreamZero
+
+DreamZero is the robot-policy serving example for the OpenPI-compatible `/v1/realtime/robot/openpi` endpoint.
+
+## Supported checkpoints
+
+| Model | HuggingFace repo |
+|---|---|
+| DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` |
+
+## Quick start
+
+### Start the server
+
+```bash
+bash examples/online_serving/dreamzero/run_server.sh
+```
+
+By default this launches:
+
+```bash
+vllm serve GEAR-Dreams/DreamZero-DROID --omni --port 8000 \
+    --served-model-name dreamzero-droid \
+    --deploy-config vllm_omni/deploy/dreamzero_tp1_cfg2.yaml \
+    --enforce-eager --disable-log-stats
+```
+
+Override `MODEL`, `PORT`, `HOST`, `DEPLOY_CONFIG`, or `SERVED_MODEL_NAME` through the script environment if needed.
+
+### Install example dependencies
+
+The core `pip install -e .` setup does not include the extra packages used by the DreamZero example scripts.
+
+- `openpi_client.py`:
+  `openpi-client`, `websockets`, `opencv-python`
+- `droid_sim_eval_client.py`:
+  `mediapy`, `websockets`, `openpi-client`
+
+The DROID sim-eval script also needs an Isaac Lab environment that provides `isaaclab`, `isaaclab_tasks`, `sim_evals`, and `gymnasium`.
+
+If you run the DROID client on Python < 3.12, also install `typing-extensions`.
+
+### Configure TP and CFG parallelism
+
+The bundled DreamZero configs intentionally keep only:
+
+| Config | Purpose |
+|---|---|
+| `vllm_omni/deploy/dreamzero.yaml` | Default TP=1, CFG parallel disabled |
+| `vllm_omni/deploy/dreamzero_tp1_cfg2.yaml` | TP=1, CFG parallel size=2 |
+
+For other topologies, use CLI parallelism flags and update stage 0 `devices` with `--stage-overrides`. The number of listed devices must match `tensor_parallel_size * cfg_parallel_size`.
+
+TP=2 with CFG parallel disabled:
+
+```bash
+vllm serve GEAR-Dreams/DreamZero-DROID --omni --port 8000 \
+    --served-model-name dreamzero-droid \
+    --deploy-config vllm_omni/deploy/dreamzero.yaml \
+    --tensor-parallel-size 2 \
+    --stage-overrides '{"0": {"devices": "0,1"}}' \
+    --enforce-eager --disable-log-stats
+```
+
+TP=2 with CFG parallel size=2:
+
+```bash
+vllm serve GEAR-Dreams/DreamZero-DROID --omni --port 8000 \
+    --served-model-name dreamzero-droid \
+    --deploy-config vllm_omni/deploy/dreamzero.yaml \
+    --tensor-parallel-size 2 \
+    --cfg-parallel-size 2 \
+    --stage-overrides '{"0": {"devices": "0,1,2,3"}}' \
+    --enforce-eager --disable-log-stats
+```
+
+### Download example assets
+
+The OpenPI client and DROID sim-eval example expect the three camera MP4 files in `outputs/dreamzero/assets`.
+
+```bash
+hf download YangshenDeng/vllm-omni-dreamzero-assets --repo-type dataset --local-dir outputs/dreamzero/assets
+```
+
+### Run the OpenPI client
+
+```bash
+python examples/online_serving/dreamzero/openpi_client.py \
+    --host 127.0.0.1 \
+    --port 8000 \
+    --video-dir outputs/dreamzero/assets
+```
+
+This client uses downloaded example videos and talks to the OpenPI websocket server.
+
+### Run DROID sim eval
+
+```bash
+${ISAACLAB_LAUNCHER} -p examples/online_serving/dreamzero/droid_sim_eval_client.py \
+    --host 127.0.0.1 \
+    --port 8000 \
+    --scene 1 \
+    --episodes 1 \
+    --headless \
+    --device cuda:0
+```
+
+Set `ISAACLAB_LAUNCHER=path/to/isaaclab.sh` from the vLLM-Omni repository root before running the command.
+This launches Isaac Lab / sim-evals and runs the DROID benchmark loop against the same websocket endpoint.
+
+### Export comparison videos offline
+
+```bash
+python examples/offline_inference/dreamzero/export_prediction_video.py \
+    --deploy-config vllm_omni/deploy/dreamzero_tp1_cfg2.yaml \
+    --save-input-video \
+    --save-gif
+```
+
+The export script uses local `Omni` inference, not the websocket server. It writes the input rollout video and the predicted output video artifacts for side-by-side inspection.
+
+### MolmoSpace demo
+
+```bash
+python examples/online_serving/dreamzero/molmospace_dreamzero_eval_demo.py \
+    --host 127.0.0.1 \
+    --port 8000 \
+    --benchmark_dir /path/to/benchmark \
+    --output_dir /path/to/output
+```
+
+This demo adapts DreamZero to the MolmoSpace-style remote policy eval loop.

examples/online_serving/dreamzero/droid_sim_eval_client.py

@@ -0,0 +1,819 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Run one or more DROID sim-eval rollouts against the vLLM DreamZero server.
+
+This script is the vLLM/OpenPI adaptation of the upstream DreamZero
+`eval_utils/run_sim_eval.py` sim-eval client.
+
+Behavior intentionally kept close to upstream:
+- same DROID observation extraction (`external_cam`, `external_cam_2`,
+  `wrist_cam`, joint position, gripper position)
+- same resize-with-pad preprocessing to `(180, 320)`
+- same `open_loop_horizon=8`
+- same gripper binarization rule (`> 0.5 -> 1`, else `0`)
+- same per-scene language prompts
+
+Unlike upstream DreamZero, vLLM serves the compatible websocket policy endpoint
+at `/v1/realtime/robot/openpi`, so this script includes the path suffix in the
+client URI.
+
+Run this script through Isaac Lab's launcher from the vLLM-Omni repository
+root, for example:
+
+    "${ISAACLAB_LAUNCHER}" -p \
+      examples/online_serving/dreamzero/droid_sim_eval_client.py \
+      --host 127.0.0.1 \
+      --port 8000 \
+      --scene 1 \
+      --episodes 1 \
+      --headless \
+      --device cuda:1
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+import sys
+import time
+import uuid
+from dataclasses import dataclass
+from datetime import UTC, datetime
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+
+try:
+    import mediapy
+except ImportError as exc:  # pragma: no cover - runtime dependency guard
+    raise ImportError("Optional DROID sim-eval client requires `mediapy`.") from exc
+
+try:
+    from typing import override
+except ImportError:
+    try:
+        from typing_extensions import override
+    except ImportError as exc:  # pragma: no cover - runtime dependency guard
+        raise ImportError("Optional DROID sim-eval client requires `typing-extensions` on Python < 3.12.") from exc
+
+try:
+    import websockets.sync.client
+except ImportError as exc:  # pragma: no cover - runtime dependency guard
+    raise ImportError("Optional DROID sim-eval client requires `websockets`.") from exc
+
+# NOTE:
+# This directory already contains a local file named `openpi_client.py`.
+# However, what we want here is the *installed* `openpi_client` package from
+# upstream OpenPI, not the sibling example file. When a script is executed
+# directly, Python often puts the script directory at `sys.path[0]`, which
+# would cause `import openpi_client` to resolve to the local example file and
+# create a circular import.
+#
+# To avoid that ambiguity, temporarily remove the current example directory
+# from the front of `sys.path`, import the real package, and then restore the
+# path afterwards.
+example_dir = str(Path(__file__).resolve().parent)
+removed_path = False
+if sys.path and sys.path[0] == example_dir:
+    sys.path.pop(0)
+    removed_path = True
+try:
+    from openpi_client import image_tools, msgpack_numpy
+    from openpi_client.base_policy import BasePolicy
+except ImportError as exc:  # pragma: no cover - runtime dependency guard
+    raise ImportError("Optional DROID sim-eval client requires the `openpi-client` package.") from exc
+finally:
+    if removed_path:
+        sys.path.insert(0, example_dir)
+
+# -----------------------------------------------------------------------------
+# Constants
+# -----------------------------------------------------------------------------
+#
+# These values intentionally mirror the upstream DreamZero sim-eval client
+# where possible. The important distinction is:
+#
+# - ACTION_HORIZON = 24
+#     The model returns 24 future actions per inference call.
+# - DEFAULT_OPEN_LOOP_HORIZON = 8
+#     The sim client only executes the first 8 actions locally before asking
+#     the server to replan from a fresh observation.
+#
+# So a single server call predicts 24x8 actions, but the rollout consumes only
+# 8 of them before replanning.
+PING_INTERVAL_SECS = 60
+PING_TIMEOUT_SECS = 600
+DEFAULT_PATH = "/v1/realtime/robot/openpi"
+DEFAULT_OPEN_LOOP_HORIZON = 8
+ACTION_HORIZON = 24
+ACTION_DIM = 8
+DEFAULT_OUTPUT_ROOT = Path("runs") / "dreamzero_sim_eval"
+SCENE_PROMPTS = {
+    1: "put the cube in the bowl",
+    2: "pick up the can and put it in the mug",
+    3: "put the banana in the bin",
+}
+
+
+def _decode_action_response(response: bytes | str) -> np.ndarray:
+    if isinstance(response, str):
+        raise RuntimeError(f"Error in inference server:\n{response}")
+    decoded = msgpack_numpy.unpackb(response)
+    if isinstance(decoded, dict) and decoded.get("type") == "error":
+        message = decoded.get("message", decoded)
+        raise RuntimeError(f"Error in inference server:\n{message}")
+    return np.asarray(decoded, dtype=np.float32)
+
+
+@dataclass(frozen=True)
+class StepRecord:
+    """One fully materialized rollout step for later JSON export.
+
+    The `episode_00.json` artifact is intended to be human-readable and
+    post-process-friendly. Instead of keeping raw tensors around, each step is
+    flattened into plain Python types so it can be serialized directly.
+    """
+
+    # Index of this control step within the episode.
+    step_index: int
+    # Whether this step triggered a fresh model call (as opposed to reusing the
+    # cached open-loop chunk from the previous server response).
+    used_server_call: bool
+    # End-to-end latency of the server call that produced the current chunk.
+    # This is `None` on steps that only reuse cached actions.
+    chunk_latency_s: float | None
+    # The concrete 8-D action sent into the simulator at this step.
+    action: list[float]
+    # Observed 7-DoF arm joint positions before the next environment step.
+    joint_position: list[float]
+    # Observed gripper scalar before the next environment step.
+    gripper_position: list[float]
+    # Reward and termination signals directly returned by the simulator.
+    reward: float
+    terminated: bool
+    truncated: bool
+    # Optional scene object positions for downstream debugging / success
+    # heuristics. This may be empty if the environment does not expose them.
+    scene_objects: dict[str, list[float]]
+
+
+class OpenPIWebsocketClientPolicy(BasePolicy):
+    """Minimal websocket client for the DreamZero/OpenPI policy protocol.
+
+    Protocol shape:
+    - connect -> server immediately sends a metadata dict
+    - infer   -> send msgpack observation, receive action chunk
+    - reset   -> send msgpack reset command, receive confirmation string
+
+    This class intentionally stays very small because the more interesting
+    DreamZero-specific behavior lives one layer above, in
+    `DreamZeroJointPosClient`.
+    """
+
+    def __init__(
+        self,
+        host: str = "127.0.0.1",
+        port: int = 8000,
+        path: str = DEFAULT_PATH,
+    ) -> None:
+        # vLLM serves the robot endpoint under `/v1/realtime/robot/openpi`.
+        self._uri = f"ws://{host}:{port}{path}"
+        # Upstream protocol uses msgpack with numpy support, not JSON.
+        self._packer = msgpack_numpy.Packer()
+        # Connect immediately and cache the server handshake metadata.
+        self._ws, self._server_metadata = self._wait_for_server()
+
+    def get_server_metadata(self) -> dict[str, Any]:
+        """Return a copy of the server handshake metadata."""
+
+        return dict(self._server_metadata)
+
+    def _wait_for_server(self):
+        """Connect to the websocket server and read the initial metadata frame."""
+
+        logging.info("Connecting to %s", self._uri)
+        conn = websockets.sync.client.connect(
+            self._uri,
+            compression=None,
+            max_size=None,
+            ping_interval=PING_INTERVAL_SECS,
+            ping_timeout=PING_TIMEOUT_SECS,
+        )
+        metadata = msgpack_numpy.unpackb(conn.recv())
+        if not isinstance(metadata, dict):
+            raise TypeError(f"Expected dict metadata from server, got {type(metadata)!r}")
+        return conn, metadata
+
+    @override
+    def infer(self, obs: dict[str, Any]) -> np.ndarray:
+        """Send an inference request and return the decoded action chunk."""
+
+        # Keep the upstream DreamZero/OpenPI convention that the request itself
+        # tells the server which logical endpoint is being called.
+        payload = dict(obs)
+        payload["endpoint"] = "infer"
+        self._ws.send(self._packer.pack(payload))
+        response = self._ws.recv()
+        return _decode_action_response(response)
+
+    @override
+    def reset(self, reset_info: dict[str, Any] | None = None) -> str:
+        """Tell the server to reset its session-side state."""
+
+        payload = dict(reset_info or {})
+        payload["endpoint"] = "reset"
+        self._ws.send(self._packer.pack(payload))
+        response = self._ws.recv()
+        if isinstance(response, str):
+            return response
+        decoded = msgpack_numpy.unpackb(response)
+        if not isinstance(decoded, dict) or decoded.get("status") != "reset successful":
+            raise RuntimeError(f"Unexpected reset response: {decoded!r}")
+        return str(decoded["status"])
+
+    def close(self) -> None:
+        """Close the websocket connection explicitly."""
+
+        self._ws.close()
+
+
+class DreamZeroJointPosClient:
+    """DROID sim-eval client that talks to the vLLM OpenPI websocket server.
+
+    This is the main compatibility layer between:
+    - Isaac Lab DROID observations (`obs["policy"][...]`)
+    - DreamZero/OpenPI websocket payloads
+    - local open-loop action reuse across several simulator steps
+
+    In other words:
+        simulator obs -> websocket request -> action chunk -> one action per step
+    """
+
+    def __init__(
+        self,
+        remote_host: str = "127.0.0.1",
+        remote_port: int = 8000,
+        path: str = DEFAULT_PATH,
+        open_loop_horizon: int = DEFAULT_OPEN_LOOP_HORIZON,
+    ) -> None:
+        # Low-level transport client.
+        self.client = OpenPIWebsocketClientPolicy(remote_host, remote_port, path=path)
+        # Number of actions to execute locally before replanning.
+        self.open_loop_horizon = open_loop_horizon
+        # Cursor into the currently cached action chunk.
+        self.actions_from_chunk_completed = 0
+        # Most recent `(ACTION_HORIZON, ACTION_DIM)` server response.
+        self.pred_action_chunk: np.ndarray | None = None
+        # Session id is part of the DreamZero serving contract. Changing it
+        # causes the server side to treat the rollout as a fresh episode.
+        self.session_id = str(uuid.uuid4())
+        # Simple runtime stats for reporting.
+        self.server_calls = 0
+        self.last_chunk_latency_s: float | None = None
+        self.last_used_server_call = False
+
+    def metadata(self) -> dict[str, Any]:
+        """Expose the server metadata to callers / logs."""
+
+        return self.client.get_server_metadata()
+
+    def reset(self) -> str:
+        """Reset local chunk state and remote session state.
+
+        Local reset:
+        - drop cached action chunk
+        - rewind chunk cursor
+        - allocate a fresh session id
+
+        Remote reset:
+        - send a websocket `reset` message so the server can clear any
+          request/session-side state it associates with this client
+        """
+
+        self.actions_from_chunk_completed = 0
+        self.pred_action_chunk = None
+        self.session_id = str(uuid.uuid4())
+        self.last_chunk_latency_s = None
+        self.last_used_server_call = False
+        return self.client.reset({})
+
+    def infer(self, obs: dict[str, Any], instruction: str) -> dict[str, Any]:
+        """Turn one simulator observation into one executable 8-D action.
+
+        Key behavior:
+        - call the server only when the local chunk cache is empty/exhausted
+        - otherwise, keep consuming the cached chunk open-loop
+        - always return exactly one 8-D action for the current simulator step
+        """
+
+        # Convert Isaac Lab observation structure into a plain numpy-friendly
+        # record that is easier to serialize and visualize.
+        curr_obs = self._extract_observation(obs)
+        self.last_used_server_call = False
+
+        # Replan if:
+        # 1. this is the first step of a rollout / chunk
+        # 2. we already consumed `open_loop_horizon` actions from the current chunk
+        # 3. no cached chunk is currently available
+        if (
+            self.actions_from_chunk_completed == 0
+            or self.actions_from_chunk_completed >= self.open_loop_horizon
+            or self.pred_action_chunk is None
+        ):
+            self.actions_from_chunk_completed = 0
+            # Build the exact DreamZero/OpenPI payload expected by the server.
+            #
+            # Notes:
+            # - images are resized/padded to the serving contract's 180x320
+            # - proprio is cast to float64 to match upstream client behavior
+            # - cartesian_position is currently unused by DreamZero DROID, so
+            #   a dummy zero vector is sent for protocol completeness
+            request_data = {
+                "observation/exterior_image_0_left": image_tools.resize_with_pad(curr_obs["right_image"], 180, 320),
+                "observation/exterior_image_1_left": image_tools.resize_with_pad(curr_obs["left_image"], 180, 320),
+                "observation/wrist_image_left": image_tools.resize_with_pad(curr_obs["wrist_image"], 180, 320),
+                "observation/joint_position": curr_obs["joint_position"].astype(np.float64),
+                "observation/cartesian_position": np.zeros((6,), dtype=np.float64),
+                "observation/gripper_position": curr_obs["gripper_position"].astype(np.float64),
+                "prompt": instruction,
+                "session_id": self.session_id,
+            }
+
+            # Measure end-to-end server latency for this chunk request.
+            start = time.perf_counter()
+            actions = self.client.infer(request_data)
+            self.last_chunk_latency_s = time.perf_counter() - start
+            self.last_used_server_call = True
+            self.server_calls += 1
+
+            # DreamZero DROID serving is expected to return an action chunk with
+            # 24 future actions, each action being 8-D.
+            if actions.ndim != 2:
+                raise AssertionError(f"Expected 2D action array, got shape {actions.shape}")
+            if actions.shape != (ACTION_HORIZON, ACTION_DIM):
+                raise AssertionError(f"Expected action shape {(ACTION_HORIZON, ACTION_DIM)}, got {actions.shape}")
+            self.pred_action_chunk = actions
+
+        # Consume exactly one action row from the cached chunk for this
+        # simulator step.
+        action = np.array(self.pred_action_chunk[self.actions_from_chunk_completed], copy=True)
+        self.actions_from_chunk_completed += 1
+
+        # Upstream DreamZero sim-eval binarizes the gripper command.
+        action[-1] = 1.0 if action[-1].item() > 0.5 else 0.0
+
+        # Produce a human-friendly visualization strip for videos:
+        # right external | wrist | left external
+        img1 = image_tools.resize_with_pad(curr_obs["right_image"], 224, 224)
+        img2 = image_tools.resize_with_pad(curr_obs["wrist_image"], 224, 224)
+        img3 = image_tools.resize_with_pad(curr_obs["left_image"], 224, 224)
+        viz = np.concatenate([img1, img2, img3], axis=1)
+
+        # Return both the executable action and auxiliary debug info.
+        return {
+            "action": action,
+            "viz": viz,
+            "joint_position": curr_obs["joint_position"],
+            "gripper_position": curr_obs["gripper_position"],
+            "used_server_call": self.last_used_server_call,
+            "chunk_latency_s": self.last_chunk_latency_s if self.last_used_server_call else None,
+        }
+
+    @staticmethod
+    def _extract_observation(obs_dict: dict[str, Any]) -> dict[str, np.ndarray]:
+        """Extract the pieces DreamZero cares about from Isaac Lab observations.
+
+        `sim-evals` exposes camera frames and robot state inside the
+        `obs["policy"]` group. This helper converts those tensors into numpy
+        arrays so they can be fed into image preprocessing / websocket packing.
+        """
+
+        policy = obs_dict["policy"]
+        # Isaac Lab stores camera observations as batched tensors; use env 0.
+        right_image = policy["external_cam"][0].clone().detach().cpu().numpy()
+        left_image = policy["external_cam_2"][0].clone().detach().cpu().numpy()
+        wrist_image = policy["wrist_cam"][0].clone().detach().cpu().numpy()
+        # Robot proprioception.
+        joint_position = policy["arm_joint_pos"].clone().detach().cpu().numpy()
+        gripper_position = policy["gripper_pos"].clone().detach().cpu().numpy()
+
+        return {
+            "right_image": right_image,
+            "left_image": left_image,
+            "wrist_image": wrist_image,
+            "joint_position": joint_position,
+            "gripper_position": gripper_position,
+        }
+
+
+def _scene_instruction(scene: int) -> str:
+    """Map a numeric scene id onto the fixed language prompt used for rollout."""
+
+    try:
+        return SCENE_PROMPTS[scene]
+    except KeyError as exc:
+        raise ValueError(f"Unsupported scene {scene}. Available scenes: {sorted(SCENE_PROMPTS)}") from exc
+
+
+def _capture_scene_objects(env: Any) -> dict[str, list[float]]:
+    """Best-effort extraction of scene object root positions.
+
+    This is only for debugging / reporting. The rollout logic itself does not
+    depend on these positions.
+    """
+
+    objects: dict[str, list[float]] = {}
+    scene = getattr(env, "scene", None)
+    if scene is None:
+        return objects
+
+    # Skip non-task entities such as cameras, the robot, and lighting.
+    for name in scene.keys():
+        if name in {"robot", "external_cam", "external_cam_2", "wrist_cam", "sphere_light", "scene"}:
+            continue
+        entity = scene[name]
+        data = getattr(entity, "data", None)
+        root_pos_w = getattr(data, "root_pos_w", None)
+        if root_pos_w is None:
+            continue
+        # Convert tensors to plain lists for JSON serialization.
+        value = root_pos_w[0].detach().cpu().to(torch.float32).tolist()
+        objects[str(name)] = [float(x) for x in value]
+    return objects
+
+
+def _maybe_infer_success(scene: int, final_objects: dict[str, list[float]]) -> dict[str, Any]:
+    """Best-effort geometric heuristic.
+
+    The simulator itself does not expose a built-in success term; this function
+    provides a transparent fallback for human-readable reporting only.
+    """
+
+    task_pairs = {
+        1: ("cube", "bowl"),
+        2: ("can", "mug"),
+        3: ("banana", "bin"),
+    }
+    source_name, target_name = task_pairs.get(scene, (None, None))
+    if source_name not in final_objects or target_name not in final_objects:
+        return {
+            "has_builtin_success": False,
+            "heuristic_success": None,
+            "reason": "scene object names unavailable for heuristic",
+        }
+
+    source = np.asarray(final_objects[source_name], dtype=np.float32)
+    target = np.asarray(final_objects[target_name], dtype=np.float32)
+    xy_distance = float(np.linalg.norm(source[:2] - target[:2]))
+    z_delta = float(source[2] - target[2])
+    heuristic_success = bool(xy_distance < 0.12 and z_delta > -0.08)
+    return {
+        "has_builtin_success": False,
+        "heuristic_success": heuristic_success,
+        "xy_distance": xy_distance,
+        "z_delta": z_delta,
+        "source_object": source_name,
+        "target_object": target_name,
+        "reason": "no built-in env success flag; using final object pose heuristic",
+    }
+
+
+def _dump_json(path: Path, payload: dict[str, Any]) -> None:
+    """Write a JSON file with stable UTF-8 formatting."""
+
+    path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
+
+
+def _scalar_from_env_value(value: Any) -> float:
+    """Normalize simulator scalar outputs into a plain Python float.
+
+    Isaac Lab / Gym values may come back as tensors, numpy arrays, tuples, or
+    direct Python scalars depending on the wrapper stack. Centralizing the
+    conversion here makes the rollout loop cleaner and more robust.
+    """
+
+    if isinstance(value, torch.Tensor):
+        return float(value.reshape(-1)[0].detach().cpu().item())
+    if isinstance(value, np.ndarray):
+        return float(value.reshape(-1)[0])
+    if isinstance(value, (list, tuple)):
+        return float(np.asarray(value).reshape(-1)[0])
+    return float(value)
+
+
+def _bool_from_env_value(value: Any) -> bool:
+    """Normalize simulator boolean-like outputs into a plain Python bool."""
+
+    if isinstance(value, torch.Tensor):
+        return bool(value.reshape(-1)[0].detach().cpu().item())
+    if isinstance(value, np.ndarray):
+        return bool(value.reshape(-1)[0])
+    if isinstance(value, (list, tuple)):
+        return bool(np.asarray(value).reshape(-1)[0])
+    return bool(value)
+
+
+def _make_output_dir(output_root: Path, scene: int) -> Path:
+    """Create a timestamped output directory for one sim-eval run."""
+
+    timestamp = datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ")
+    output_dir = output_root / f"scene{scene}_{timestamp}"
+    output_dir.mkdir(parents=True, exist_ok=True)
+    return output_dir
+
+
+def main() -> None:
+    """Entry point for one or more DROID simulation rollouts.
+
+    High-level flow:
+    1. parse command-line flags
+    2. bootstrap Isaac Lab / sim-evals imports
+    3. create the DROID environment
+    4. connect the DreamZero websocket client
+    5. run `episodes` rollouts
+    6. export videos + JSON summaries
+    """
+
+    # Script-level arguments.
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--episodes", type=int, default=1, help="Number of episodes to run.")
+    parser.add_argument("--scene", type=int, default=1, help="DROID scene id (1/2/3).")
+    parser.add_argument("--host", type=str, default="127.0.0.1", help="vLLM DreamZero server host.")
+    parser.add_argument("--port", type=int, default=8000, help="vLLM DreamZero server port.")
+    parser.add_argument("--path", type=str, default=DEFAULT_PATH, help="Websocket path suffix.")
+    parser.add_argument(
+        "--open-loop-horizon",
+        type=int,
+        default=DEFAULT_OPEN_LOOP_HORIZON,
+        help="How many actions to consume locally before requesting the next chunk.",
+    )
+    parser.add_argument(
+        "--output-root",
+        type=Path,
+        default=DEFAULT_OUTPUT_ROOT,
+        help="Directory where videos and trajectory logs are stored.",
+    )
+
+    try:
+        from isaaclab.app import AppLauncher
+    except ImportError as exc:  # pragma: no cover - runtime dependency guard
+        raise ImportError(
+            "Optional DROID sim-eval client requires Isaac Lab (`isaaclab`). "
+            "Launch it from an Isaac Lab environment, e.g. via `isaaclab.sh -p`."
+        ) from exc
+
+    # Let Isaac Lab inject its own runtime flags (e.g. headless, device).
+    AppLauncher.add_app_launcher_args(parser)
+    args = parser.parse_args()
+
+    # DreamZero sim-eval always needs camera observations enabled.
+    args.enable_cameras = True
+    # Boot Isaac Sim / Isaac Lab.
+    app_launcher = AppLauncher(args)
+    simulation_app = app_launcher.app
+    # Set defaults so the `finally` block can clean up safely even if an
+    # earlier step fails.
+    env = None
+    client = None
+
+    # Import simulator modules only *after* the app is launched. This matches
+    # Isaac Lab's required import ordering.
+    try:
+        import gymnasium as gym
+    except ImportError as exc:  # pragma: no cover - runtime dependency guard
+        raise ImportError("Optional DROID sim-eval client requires `gymnasium`.") from exc
+
+    try:
+        import sim_evals.environments  # noqa: F401
+    except ImportError as exc:  # pragma: no cover - runtime dependency guard
+        raise ImportError(
+            "Optional DROID sim-eval client requires the external `sim-evals` package or checkout to be importable."
+        ) from exc
+
+    try:
+        from isaaclab_tasks.utils import parse_env_cfg
+    except ImportError as exc:  # pragma: no cover - runtime dependency guard
+        raise ImportError("Optional DROID sim-eval client requires `isaaclab_tasks`.") from exc
+
+    # Resolve output location and scene prompt.
+    output_dir = _make_output_dir(args.output_root.expanduser().resolve(), args.scene)
+    instruction = _scene_instruction(args.scene)
+
+    # Build the DROID environment configuration from `sim-evals`.
+    env_cfg = parse_env_cfg(
+        "DROID",
+        device=args.device,
+        num_envs=1,
+        use_fabric=True,
+    )
+    # Select one of the pre-authored scenes/tasks.
+    env_cfg.set_scene(args.scene)
+    env = gym.make("DROID", cfg=env_cfg)
+
+    # Upstream sim-evals resets twice so materials / cameras are fully ready.
+    obs, _ = env.reset()
+    obs, _ = env.reset()
+
+    # Connect the websocket policy client.
+    client = DreamZeroJointPosClient(
+        remote_host=args.host,
+        remote_port=args.port,
+        path=args.path,
+        open_loop_horizon=args.open_loop_horizon,
+    )
+
+    # Aggregated per-run results.
+    all_episode_summaries: list[dict[str, Any]] = []
+    max_steps = int(env.env.max_episode_length)
+    logging.info("DreamZero metadata: %s", client.metadata())
+    logging.info("Scene %s prompt: %s", args.scene, instruction)
+    logging.info("Writing outputs to %s", output_dir)
+
+    try:
+        # No gradients are needed in inference-only rollout mode.
+        with torch.no_grad():
+            for episode_index in range(args.episodes):
+                # Per-episode collectors.
+                frames: list[np.ndarray] = []
+                step_records: list[StepRecord] = []
+                episode_start = time.perf_counter()
+                server_time_s = 0.0
+                final_reward = 0.0
+                terminated = False
+                truncated = False
+
+                for step_index in range(max_steps):
+                    # Ask the policy for the next action. Internally this may or
+                    # may not trigger a real server request depending on whether
+                    # the local chunk cache has been exhausted.
+                    logging.debug("Episode %d step %d: requesting action", episode_index, step_index)
+                    result = client.infer(obs, instruction)
+                    logging.debug(
+                        "Episode %d step %d: got action (server_call=%s latency=%s)",
+                        episode_index,
+                        step_index,
+                        result["used_server_call"],
+                        result["chunk_latency_s"],
+                    )
+                    # Save one visualization frame per simulator step.
+                    frames.append(result["viz"])
+
+                    # Isaac Lab expects batched actions, hence `[None]`.
+                    action_tensor = torch.tensor(result["action"], dtype=torch.float32)[None]
+                    logging.debug("Episode %d step %d: stepping env", episode_index, step_index)
+                    obs, reward, term, trunc, info = env.step(action_tensor)
+                    logging.debug("Episode %d step %d: env.step returned", episode_index, step_index)
+                    logging.debug("Episode %d step %d: parsing reward/flags", episode_index, step_index)
+                    logging.debug(
+                        "Episode %d step %d: raw types reward=%s term=%s trunc=%s",
+                        episode_index,
+                        step_index,
+                        type(reward).__name__,
+                        type(term).__name__,
+                        type(trunc).__name__,
+                    )
+                    # Normalize environment outputs into plain Python scalars so
+                    # the rest of the code does not depend on wrapper-specific types.
+                    reward_value = _scalar_from_env_value(reward)
+                    term_value = _bool_from_env_value(term)
+                    trunc_value = _bool_from_env_value(trunc)
+                    logging.debug(
+                        "Episode %d step %d: parsed reward=%s term=%s trunc=%s",
+                        episode_index,
+                        step_index,
+                        reward_value,
+                        term_value,
+                        trunc_value,
+                    )
+                    # Keep scene-object capture optional. It is useful for
+                    # debugging / success heuristics, but the rollout should not
+                    # fail if the environment does not expose object roots.
+                    scene_objects = _capture_scene_objects(env)
+
+                    # Accumulate total server-side time only on steps that
+                    # triggered a fresh chunk inference.
+                    if result["chunk_latency_s"] is not None:
+                        server_time_s += float(result["chunk_latency_s"])
+
+                    # Materialize one JSON-serializable trajectory record.
+                    logging.debug("Episode %d step %d: appending trajectory", episode_index, step_index)
+                    step_records.append(
+                        StepRecord(
+                            step_index=step_index,
+                            used_server_call=bool(result["used_server_call"]),
+                            chunk_latency_s=(
+                                float(result["chunk_latency_s"]) if result["chunk_latency_s"] is not None else None
+                            ),
+                            action=[float(x) for x in np.asarray(result["action"], dtype=np.float32).tolist()],
+                            joint_position=[
+                                float(x) for x in np.asarray(result["joint_position"], dtype=np.float32).tolist()
+                            ],
+                            gripper_position=[
+                                float(x) for x in np.asarray(result["gripper_position"], dtype=np.float32).tolist()
+                            ],
+                            reward=reward_value,
+                            terminated=term_value,
+                            truncated=trunc_value,
+                            scene_objects=scene_objects,
+                        )
+                    )
+                    logging.debug("Episode %d step %d: trajectory appended", episode_index, step_index)
+
+                    # Track final status for summary export.
+                    final_reward = reward_value
+                    terminated = term_value
+                    truncated = trunc_value
+                    if term_value or trunc_value:
+                        # End the rollout early if the environment terminates.
+                        break
+
+                # Episode-level timing and video export.
+                episode_wall_time_s = time.perf_counter() - episode_start
+                video_path = output_dir / f"episode_{episode_index:02d}.mp4"
+                logging.info("Episode %d: writing video to %s", episode_index, video_path)
+                mediapy.write_video(video_path, frames, fps=15)
+
+                # Reset the policy server between episodes.
+                logging.info("Episode %d: sending reset", episode_index)
+                reset_response = client.reset()
+                final_objects = step_records[-1].scene_objects if step_records else {}
+                success_report = _maybe_infer_success(args.scene, final_objects)
+
+                # Assemble the per-episode summary that is written to
+                # `episode_XX.json`.
+                episode_summary = {
+                    "episode_index": episode_index,
+                    "prompt": instruction,
+                    "video_path": str(video_path),
+                    "steps_executed": len(step_records),
+                    "max_steps": max_steps,
+                    "terminated": terminated,
+                    "truncated": truncated,
+                    "final_reward": final_reward,
+                    "server_calls": client.server_calls,
+                    "server_time_s": server_time_s,
+                    "episode_wall_time_s": episode_wall_time_s,
+                    "avg_server_time_per_call_s": (
+                        server_time_s / client.server_calls if client.server_calls else None
+                    ),
+                    "reset_response": reset_response,
+                    "success_report": success_report,
+                    "server_metadata": client.metadata(),
+                    "trajectory": [record.__dict__ for record in step_records],
+                }
+                _dump_json(output_dir / f"episode_{episode_index:02d}.json", episode_summary)
+                all_episode_summaries.append(episode_summary)
+
+                logging.info(
+                    "Episode %d done: steps=%d wall=%.2fs server_calls=%d heuristic_success=%s",
+                    episode_index,
+                    len(step_records),
+                    episode_wall_time_s,
+                    client.server_calls,
+                    success_report.get("heuristic_success"),
+                )
+
+                # Reset per-episode counters while keeping the client alive.
+                client.server_calls = 0
+
+        # Top-level run summary across all episodes.
+        summary = {
+            "scene": args.scene,
+            "prompt": instruction,
+            "episodes": args.episodes,
+            "host": args.host,
+            "port": args.port,
+            "path": args.path,
+            "device": args.device,
+            "output_dir": str(output_dir),
+            "server_metadata": client.metadata(),
+            "episodes_summary": all_episode_summaries,
+        }
+        _dump_json(output_dir / "summary.json", summary)
+        # Also print the summary to stdout so the caller can capture it in logs.
+        print(json.dumps(summary, ensure_ascii=False, indent=2))
+    finally:
+        # Best-effort cleanup. Avoid masking the main error if cleanup fails.
+        try:
+            if client is not None:
+                client.client.close()
+        except Exception:
+            pass
+        if env is not None:
+            env.close()
+        simulation_app.close()
+
+
+if __name__ == "__main__":
+    # Keep the script-level logs readable. Per-step rollout details are still
+    # available via `DEBUG` if needed, but websocket / asyncio internals are
+    # usually too noisy for normal usage.
+    logging.basicConfig(level=logging.INFO)
+    logging.getLogger("websockets").setLevel(logging.WARNING)
+    logging.getLogger("asyncio").setLevel(logging.WARNING)
+    main()

examples/online_serving/dreamzero/molmospace_dreamzero_eval_demo.py

@@ -0,0 +1,86 @@
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+from pathlib import Path
+
+os.environ.setdefault("MUJOCO_GL", "egl")
+os.environ.setdefault("PYOPENGL_PLATFORM", "egl")
+
+# Import base configs at module top level so the subclasses below are pickle-
+# resolvable (worker processes import this module fresh via __main__).
+from molmo_spaces.configs.policy_configs_baselines import (  # noqa: E402
+    DreamZeroPolicyConfig,
+)
+from molmo_spaces.evaluation.configs.evaluation_configs import (  # noqa: E402
+    DreamZeroPolicyEvalConfig,
+)
+
+
+# We only need to change the backend host and port to the vllm-host!
+class DreamZeroVllmOmniPolicyConfig(DreamZeroPolicyConfig):
+    remote_config: dict = dict(host="127.0.0.1", port=8000)
+
+
+class DreamZeroVllmOmniEvalConfig(DreamZeroPolicyEvalConfig):
+    policy_config: DreamZeroVllmOmniPolicyConfig = DreamZeroVllmOmniPolicyConfig()
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--host", default="127.0.0.1")
+    parser.add_argument("--port", type=int, default=8000)
+    parser.add_argument(
+        "--benchmark_dir",
+        required=True,
+        help=(
+            "Path to a MolmoSpaces benchmark directory, for example "
+            "$MOLMOSPACES_BENCHMARK_DIR/20260327/ithor/FrankaCloseHardBench/"
+            "FrankaCloseHardBench_20260206_json_benchmark"
+        ),
+    )
+    parser.add_argument("--max_episodes", type=int, default=1)
+    parser.add_argument("--task_horizon_steps", type=int, default=80)
+    parser.add_argument(
+        "--output_dir",
+        required=True,
+        help="Directory to write evaluation outputs (created if missing).",
+    )
+    parser.add_argument("--episode_idx", type=int, default=None)
+    args = parser.parse_args()
+
+    policy_config = DreamZeroVllmOmniPolicyConfig(remote_config=dict(host=args.host, port=args.port))
+    DreamZeroVllmOmniPolicyConfig.model_fields["remote_config"].default = policy_config.remote_config
+    DreamZeroVllmOmniEvalConfig.model_fields["policy_config"].default = policy_config
+
+    # Import after env vars are set so MuJoCo picks EGL.
+    from molmo_spaces.evaluation import run_evaluation
+
+    cfg_cls = DreamZeroVllmOmniEvalConfig
+
+    output_dir = args.output_dir
+    Path(output_dir).mkdir(parents=True, exist_ok=True)
+
+    print(f"[eval] benchmark_dir={args.benchmark_dir}")
+    print(f"[eval] max_episodes={args.max_episodes} task_horizon_steps={args.task_horizon_steps}")
+    print(f"[eval] remote policy: ws://{args.host}:{args.port}/v1/realtime/robot/openpi")
+
+    results = run_evaluation(
+        eval_config_cls=cfg_cls,
+        benchmark_dir=Path(args.benchmark_dir),
+        max_episodes=args.max_episodes,
+        task_horizon_steps=args.task_horizon_steps,
+        num_workers=1,
+        use_wandb=False,
+        output_dir=output_dir,
+        episode_idx=args.episode_idx,
+    )
+
+    print(f"[eval] success={results.success_count}/{results.total_count} ({results.success_rate:.1%})")
+    print(f"[eval] output_dir={results.output_dir}")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

examples/online_serving/dreamzero/openpi_client.py

@@ -0,0 +1,311 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+import sys
+import uuid
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+
+try:
+    import cv2
+except ImportError as exc:  # pragma: no cover - runtime dependency guard
+    raise ImportError("DreamZero OpenPI example requires `opencv-python`.") from exc
+
+try:
+    example_dir = str(Path(__file__).resolve().parent)
+    removed_path = False
+    if sys.path and sys.path[0] == example_dir:
+        sys.path.pop(0)
+        removed_path = True
+    try:
+        from openpi_client.websocket_client_policy import WebsocketClientPolicy
+    finally:
+        if removed_path:
+            sys.path.insert(0, example_dir)
+except ImportError as exc:  # pragma: no cover - runtime dependency guard
+    raise ImportError("DreamZero OpenPI example requires `openpi-client`.") from exc
+
+DEFAULT_HOST = "127.0.0.1"
+DEFAULT_PORT = 8000
+DEFAULT_PATH = "/v1/realtime/robot/openpi"
+DEFAULT_PROMPT = "Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan"
+ACTION_HORIZON = 24
+DEFAULT_ACTION_DIM = 8
+RELATIVE_OFFSETS = [-23, -16, -8, 0]
+REPO_ROOT = Path(__file__).resolve().parents[3]
+ASSET_REPO_ID = "YangshenDeng/vllm-omni-dreamzero-assets"
+DEFAULT_VIDEO_DIR = REPO_ROOT / "outputs" / "dreamzero" / "assets"
+CAMERA_FILES = {
+    "observation/exterior_image_0_left": "exterior_image_1_left.mp4",
+    "observation/exterior_image_1_left": "exterior_image_2_left.mp4",
+    "observation/wrist_image_left": "wrist_image_left.mp4",
+}
+
+
+def _as_action_array(response: Any) -> np.ndarray:
+    if isinstance(response, dict) and response.get("type") == "error":
+        message = response.get("message", response)
+        raise RuntimeError(f"Inference failed: {message}")
+    return np.asarray(response, dtype=np.float32)
+
+
+@dataclass(frozen=True)
+class DreamZeroServerMetadata:
+    image_resolution: tuple[int, int]
+    n_external_cameras: int
+    needs_wrist_camera: bool
+    needs_stereo_camera: bool
+    needs_session_id: bool
+    action_space: str
+
+    @classmethod
+    def from_dict(cls, payload: dict[str, Any]) -> DreamZeroServerMetadata:
+        required_keys = (
+            "image_resolution",
+            "n_external_cameras",
+            "needs_wrist_camera",
+            "needs_stereo_camera",
+            "needs_session_id",
+            "action_space",
+        )
+        missing_keys = [key for key in required_keys if key not in payload]
+        if missing_keys:
+            raise ValueError(f"Missing DreamZero metadata keys: {missing_keys}")
+
+        image_resolution = payload["image_resolution"]
+        if not isinstance(image_resolution, (list, tuple)) or len(image_resolution) != 2:
+            raise ValueError(f"Invalid image_resolution: {image_resolution!r}")
+
+        return cls(
+            image_resolution=(int(image_resolution[0]), int(image_resolution[1])),
+            n_external_cameras=int(payload["n_external_cameras"]),
+            needs_wrist_camera=bool(payload["needs_wrist_camera"]),
+            needs_stereo_camera=bool(payload["needs_stereo_camera"]),
+            needs_session_id=bool(payload["needs_session_id"]),
+            action_space=str(payload["action_space"]),
+        )
+
+
+def _openpi_uri(host: str, port: int, path: str) -> str:
+    if host.startswith(("ws://", "wss://")):
+        return host
+    return f"ws://{host}:{port}{path}"
+
+
+def _close_policy_client(client: WebsocketClientPolicy) -> None:
+    ws = getattr(client, "_ws", None)
+    if ws is not None:
+        ws.close()
+
+
+def load_all_frames(video_path: Path) -> np.ndarray:
+    cap = cv2.VideoCapture(str(video_path))
+    frames = []
+    while True:
+        ok, frame = cap.read()
+        if not ok:
+            break
+        frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+    cap.release()
+    if not frames:
+        raise RuntimeError(f"No frames loaded from {video_path}")
+    return np.stack(frames, axis=0)
+
+
+def load_camera_frames(video_dir: Path) -> dict[str, np.ndarray]:
+    camera_frames: dict[str, np.ndarray] = {}
+    for camera_key, file_name in CAMERA_FILES.items():
+        video_path = video_dir / file_name
+        if not video_path.exists():
+            raise FileNotFoundError(
+                f"Missing DreamZero example asset: {video_path}. "
+                "Download the example videos with: "
+                f"`hf download {ASSET_REPO_ID} --repo-type dataset --local-dir {video_dir}`"
+            )
+        camera_frames[camera_key] = load_all_frames(video_path)
+    return camera_frames
+
+
+def build_frame_schedule(total_frames: int, num_chunks: int) -> list[list[int]]:
+    chunks: list[list[int]] = []
+    current_frame = 23
+    for _ in range(num_chunks):
+        indices = [max(current_frame + offset, 0) for offset in RELATIVE_OFFSETS]
+        if indices[-1] >= total_frames:
+            break
+        chunks.append(indices)
+        current_frame += ACTION_HORIZON
+    return chunks
+
+
+def make_obs_from_video(
+    camera_frames: dict[str, np.ndarray],
+    frame_indices: list[int],
+    *,
+    prompt: str,
+    session_id: str,
+) -> dict[str, Any]:
+    obs: dict[str, Any] = {}
+    for camera_key, all_frames in camera_frames.items():
+        selected = all_frames[frame_indices]
+        obs[camera_key] = selected[0] if len(frame_indices) == 1 else selected
+
+    obs["observation/joint_position"] = np.zeros(7, dtype=np.float32)
+    obs["observation/cartesian_position"] = np.zeros(6, dtype=np.float32)
+    obs["observation/gripper_position"] = np.zeros(1, dtype=np.float32)
+    obs["prompt"] = prompt
+    obs["session_id"] = session_id
+    return obs
+
+
+def build_demo_observations(
+    camera_frames: dict[str, np.ndarray],
+    *,
+    prompt: str,
+    session_id: str,
+    num_chunks: int = 2,
+) -> list[dict[str, Any]]:
+    if num_chunks < 1:
+        raise ValueError("num_chunks must be at least 1")
+
+    total_frames = min(frames.shape[0] for frames in camera_frames.values())
+    observations = [
+        make_obs_from_video(
+            camera_frames,
+            [0],
+            prompt=prompt,
+            session_id=session_id,
+        )
+    ]
+    for indices in build_frame_schedule(total_frames, num_chunks - 1):
+        observations.append(
+            make_obs_from_video(
+                camera_frames,
+                indices,
+                prompt=prompt,
+                session_id=session_id,
+            )
+        )
+    return observations
+
+
+def validate_session_result(
+    result: dict[str, Any],
+    *,
+    expected_action_horizon: int = ACTION_HORIZON,
+    expected_action_dim: int = DEFAULT_ACTION_DIM,
+    expected_action_count: int | None = None,
+) -> None:
+    metadata = DreamZeroServerMetadata.from_dict(result["metadata"])
+    if metadata.image_resolution != (180, 320):
+        raise AssertionError(f"Unexpected image_resolution: {metadata.image_resolution}")
+    if metadata.n_external_cameras != 2:
+        raise AssertionError(f"Unexpected n_external_cameras: {metadata.n_external_cameras}")
+    if not metadata.needs_wrist_camera:
+        raise AssertionError("DreamZero example expects wrist camera metadata")
+    if metadata.action_space != "joint_position":
+        raise AssertionError(f"Unexpected action_space: {metadata.action_space}")
+
+    actions = result["actions"]
+    if expected_action_count is not None and len(actions) != expected_action_count:
+        raise AssertionError(f"Expected {expected_action_count} action tensors, got {len(actions)}")
+    if not actions:
+        raise AssertionError("Expected at least one action tensor")
+    for index, action in enumerate(actions):
+        if action.shape != (expected_action_horizon, expected_action_dim):
+            raise AssertionError(
+                f"Action {index} shape mismatch: expected "
+                f"{(expected_action_horizon, expected_action_dim)}, got {action.shape}"
+            )
+        if not np.isfinite(action).all():
+            raise AssertionError(f"Action {index} contains non-finite values")
+
+
+def run_policy_session(
+    *,
+    host: str = DEFAULT_HOST,
+    port: int = DEFAULT_PORT,
+    path: str = DEFAULT_PATH,
+    video_dir: Path = DEFAULT_VIDEO_DIR,
+    prompt: str = DEFAULT_PROMPT,
+    session_id: str | None = None,
+    num_chunks: int = 2,
+) -> dict[str, Any]:
+    session_id = session_id or str(uuid.uuid4())
+    camera_frames = load_camera_frames(video_dir)
+    observations = build_demo_observations(
+        camera_frames,
+        prompt=prompt,
+        session_id=session_id,
+        num_chunks=num_chunks,
+    )
+
+    uri = _openpi_uri(host, port, path)
+    logging.info("Connecting to %s", uri)
+    client = WebsocketClientPolicy(host=uri)
+    try:
+        metadata = dict(client.get_server_metadata())
+        actions = [_as_action_array(client.infer(obs)) for obs in observations]
+        return {
+            "metadata": metadata,
+            "actions": actions,
+            "session_id": session_id,
+        }
+    finally:
+        _close_policy_client(client)
+
+
+def format_action_summary(index: int, action: np.ndarray) -> str:
+    return (
+        f"Action {index}: shape={tuple(action.shape)} dtype={action.dtype} "
+        f"min={action.min():.6f} max={action.max():.6f}"
+    )
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="DreamZero OpenPI client example with downloaded real videos.")
+    parser.add_argument("--host", default=DEFAULT_HOST)
+    parser.add_argument("--port", type=int, default=DEFAULT_PORT)
+    parser.add_argument("--path", default=DEFAULT_PATH)
+    parser.add_argument(
+        "--video-dir", type=Path, default=DEFAULT_VIDEO_DIR, help="Directory containing the three camera MP4 files."
+    )
+    parser.add_argument("--prompt", default=DEFAULT_PROMPT)
+    parser.add_argument("--session-id", default=None)
+    parser.add_argument("--num-chunks", type=int, default=2)
+    return parser.parse_args()
+
+
+def main() -> int:
+    args = parse_args()
+    logging.basicConfig(level=logging.INFO)
+
+    result = run_policy_session(
+        host=args.host,
+        port=args.port,
+        path=args.path,
+        video_dir=args.video_dir,
+        prompt=args.prompt,
+        session_id=args.session_id,
+        num_chunks=args.num_chunks,
+    )
+    validate_session_result(result, expected_action_count=args.num_chunks)
+
+    print("Server metadata:", json.dumps(result["metadata"], sort_keys=True))
+    for index, action in enumerate(result["actions"]):
+        print(format_action_summary(index, action))
+    print("Session ID:", result["session_id"])
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

examples/online_serving/dreamzero/run_server.sh

@@ -0,0 +1,24 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+MODEL="${MODEL:-GEAR-Dreams/DreamZero-DROID}"
+HOST="${HOST:-127.0.0.1}"
+PORT="${PORT:-8000}"
+DEPLOY_CONFIG="${DEPLOY_CONFIG:-vllm_omni/deploy/dreamzero_tp1_cfg2.yaml}"
+SERVED_MODEL_NAME="${SERVED_MODEL_NAME:-dreamzero-droid}"
+
+args=(
+  serve
+  "$MODEL"
+  --omni
+  --host "$HOST"
+  --port "$PORT"
+  --served-model-name "$SERVED_MODEL_NAME"
+  --deploy-config "$DEPLOY_CONFIG"
+  --enforce-eager
+  --disable-log-stats
+)
+
+ATTENTION_BACKEND="${ATTENTION_BACKEND:-torch}" \
+DIFFUSION_ATTENTION_BACKEND="${DIFFUSION_ATTENTION_BACKEND:-TORCH_SDPA}" \
+vllm "${args[@]}"

examples/online_serving/dynin_omni/README.md

@@ -22,13 +22,14 @@ export PYTHONPATH=<transformers_modules_path>:$PYTHONPATH
 export HF_MODULES_CACHE=<transformers_modules_path>
 ```
 
-Run from repository root:
+Run from repository root. The 3-stage pipeline is resolved automatically from
+the model type (`vllm_omni/deploy/dynin_omni.yaml`), so no config path is
+required. Pass `--deploy-config` only to override the defaults:
 
 ```bash
 vllm-omni serve snu-aidas/Dynin-Omni \
   --omni \
-  --port 8091 \
-  --stage-configs-path "$(pwd)/vllm_omni/model_executor/stage_configs/dynin_omni.yaml"
+  --port 8091
 ```
 
 If `vllm-omni` is not in PATH, run:
@@ -36,8 +37,7 @@ If `vllm-omni` is not in PATH, run:
 ```bash
 PYTHONPATH="$(pwd)" python -m vllm_omni.entrypoints.cli.main serve snu-aidas/Dynin-Omni \
   --omni \
-  --port 8091 \
-  --stage-configs-path "$(pwd)/vllm_omni/model_executor/stage_configs/dynin_omni.yaml"
+  --port 8091
 ```
 
 Wait until the server logs show both `All stages initialized successfully` and

examples/online_serving/lance/README.md

@@ -0,0 +1,40 @@
+# Lance — online serving
+
+OpenAI-compatible chat completions API for Lance, served by `vllm-omni`.
+
+## Start the server
+
+```bash
+bash examples/online_serving/lance/run_server.sh
+# or, with overrides
+MODEL=bytedance-research/Lance \
+DEPLOY_CONFIG=vllm_omni/deploy/lance.yaml \
+PORT=8091 \
+    bash examples/online_serving/lance/run_server.sh
+```
+
+The default deploy config is the single-stage `lance.yaml`; for `video_edit`
+or `text2video` workloads, point ``DEPLOY_CONFIG`` at a video-checkpoint
+config (or pass ``--model bytedance-research/Lance/Lance_3B_Video`` to
+`run_server.sh`).
+
+## Send a request
+
+```bash
+# text-to-image
+python examples/online_serving/lance/openai_chat_client.py \
+    --prompt "A cute corgi astronaut on the moon, cinematic" \
+    --modality text2img \
+    --output corgi.png
+
+# image edit
+python examples/online_serving/lance/openai_chat_client.py \
+    --prompt "Convert this into a vibrant cartoon-style illustration" \
+    --modality img2img \
+    --image-url path/to/photo.png \
+    --output edited.png
+```
+
+The client is shared with BAGEL — same OpenAI message format, same
+``modalities`` and ``num_inference_steps`` / ``seed`` / ``height`` /
+``width`` knobs.

examples/online_serving/lance/openai_chat_client.py

@@ -0,0 +1,184 @@
+#!/usr/bin/env python3
+"""
+Lance OpenAI-compatible chat client for image generation and multimodal tasks.
+
+Usage:
+    python examples/online_serving/lance/openai_chat_client.py --prompt "A cute cat" --output output.png
+    python examples/online_serving/lance/openai_chat_client.py --prompt "Describe this image" --image-url https://example.com/image.png
+"""
+
+import argparse
+import base64
+from pathlib import Path
+
+import requests
+
+
+def generate_image(
+    prompt: str,
+    server_url: str = "http://localhost:8091",
+    image_url: str | None = None,
+    height: int | None = None,
+    width: int | None = None,
+    steps: int | None = None,
+    seed: int | None = None,
+    negative_prompt: str | None = None,
+    modality: str = "text2img",  # "text2img" (default), "img2img", "img2text", "text2text"
+) -> bytes | str | None:
+    """Generate an image or text using the chat completions API.
+
+    Args:
+        prompt: Text description or prompt
+        server_url: Server URL
+        image_url: URL or path to input image (for img2img/img2text)
+        height: Image height in pixels
+        width: Image width in pixels
+        steps: Number of inference steps
+        seed: Random seed
+        negative_prompt: Negative prompt
+        modality: Task modality hint
+
+    Returns:
+        Image bytes (for image outputs) or Text string (for text outputs) or None if failed
+    """
+
+    # Construct Message Content
+    content = [{"type": "text", "text": f"<|im_start|>{prompt}<|im_end|>"}]
+
+    if image_url:
+        # Check if local file
+        if Path(image_url).exists():
+            with open(image_url, "rb") as f:
+                b64_data = base64.b64encode(f.read()).decode("utf-8")
+                final_image_url = f"data:image/jpeg;base64,{b64_data}"
+        else:
+            final_image_url = image_url
+
+        content.append({"type": "image_url", "image_url": {"url": final_image_url}})
+
+    messages = [{"role": "user", "content": content}]
+
+    # Build request payload with all parameters at top level
+    # Note: vLLM ignores "extra_body", so we put parameters directly in the payload
+    payload = {"messages": messages}
+
+    # Set output modalities at top level
+    if modality == "text2img" or modality == "img2img":
+        payload["modalities"] = ["image"]
+    elif modality == "img2text" or modality == "text2text":
+        payload["modalities"] = ["text"]
+
+    # Add generation parameters directly to payload
+    if height is not None:
+        payload["height"] = height
+    if width is not None:
+        payload["width"] = width
+    if steps is not None:
+        payload["num_inference_steps"] = steps
+    if seed is not None:
+        payload["seed"] = seed
+    if negative_prompt:
+        payload["negative_prompt"] = negative_prompt
+
+    # Send request
+    try:
+        print(f"Sending request to {server_url} with modality {modality}...")
+        response = requests.post(
+            f"{server_url}/v1/chat/completions",
+            headers={"Content-Type": "application/json"},
+            json=payload,
+            timeout=300,
+        )
+        response.raise_for_status()
+        data = response.json()
+
+        # Extract content - check ALL choices since server may return multiple
+        # (e.g., text in choices[0], image in choices[1])
+        choices = data.get("choices", [])
+
+        # First pass: look for image output in any choice
+        for choice in choices:
+            choice_content = choice.get("message", {}).get("content")
+
+            # Handle Image Output
+            if isinstance(choice_content, list) and len(choice_content) > 0:
+                first_item = choice_content[0]
+                if isinstance(first_item, dict) and "image_url" in first_item:
+                    img_url_str = first_item["image_url"].get("url", "")
+                    if img_url_str.startswith("data:image"):
+                        _, b64_data = img_url_str.split(",", 1)
+                        return base64.b64decode(b64_data)
+
+        # Second pass: look for text output if no image found
+        for choice in choices:
+            choice_content = choice.get("message", {}).get("content")
+            if isinstance(choice_content, str) and choice_content:
+                return choice_content
+
+        print(f"Unexpected response format: {choices}")
+        return None
+
+    except Exception as e:
+        print(f"Error: {e}")
+        return None
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Bagel multimodal chat client")
+    parser.add_argument("--prompt", "-p", default="A cute cat", help="Text prompt")
+    parser.add_argument("--output", "-o", default="bagel_output.png", help="Output file (for image results)")
+    parser.add_argument("--server", "-s", default="http://localhost:8091", help="Server URL")
+
+    # Modality Control
+    parser.add_argument("--image-url", "-i", type=str, help="Input image URL or local path")
+    parser.add_argument(
+        "--modality",
+        "-m",
+        default="text2img",
+        choices=["text2img", "img2img", "img2text", "text2text"],
+        help="Task modality",
+    )
+
+    # Generation Params
+    parser.add_argument("--height", type=int, default=512, help="Image height")
+    parser.add_argument("--width", type=int, default=512, help="Image width")
+    parser.add_argument("--steps", type=int, default=25, help="Inference steps")
+    parser.add_argument("--seed", type=int, default=42, help="Random seed")
+    parser.add_argument("--negative", help="Negative prompt")
+
+    args = parser.parse_args()
+
+    print(f"Mode: {args.modality}")
+    if args.image_url:
+        print(f"Input Image: {args.image_url}")
+
+    result = generate_image(
+        prompt=args.prompt,
+        server_url=args.server,
+        image_url=args.image_url,
+        height=args.height,
+        width=args.width,
+        steps=args.steps,
+        seed=args.seed,
+        negative_prompt=args.negative,
+        modality=args.modality,
+    )
+
+    if result:
+        if isinstance(result, bytes):
+            # It's an image
+            output_path = Path(args.output)
+            output_path.write_bytes(result)
+            print(f"Image saved to: {output_path}")
+            print(f"Size: {len(result) / 1024:.1f} KB")
+        elif isinstance(result, str):
+            # It's text
+            print("Response:")
+            print(result)
+    else:
+        print("Failed to generate response")
+        exit(1)
+
+
+if __name__ == "__main__":
+    main()

examples/online_serving/lance/run_server.sh

@@ -0,0 +1,15 @@
+#!/bin/bash
+# Lance online serving startup script
+
+MODEL="${MODEL:-bytedance-research/Lance}"
+PORT="${PORT:-8091}"
+DEPLOY_CONFIG="${DEPLOY_CONFIG:-vllm_omni/deploy/lance.yaml}"
+
+echo "Starting Lance server..."
+echo "Model:         $MODEL"
+echo "Deploy config: $DEPLOY_CONFIG"
+echo "Port:          $PORT"
+
+vllm serve "$MODEL" --omni \
+    --deploy-config "$DEPLOY_CONFIG" \
+    --port "$PORT"

examples/online_serving/mimo_audio/openai_chat_completion_client_for_multimodal_generation.py

@@ -12,7 +12,8 @@ from typing import Any
 import requests
 from openai import OpenAI
 from vllm.assets.audio import AudioAsset
-from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 
 # Modify OpenAI's API key and API base to use vLLM's API server.
 openai_api_key = "EMPTY"
@@ -493,7 +494,7 @@ def run_multimodal_generation(args) -> None:
 
 
 def parse_args():
-    parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
+    parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
     parser.add_argument(
         "--query-type",
         "-q",

examples/online_serving/minicpmo/README.md

@@ -0,0 +1,65 @@
+# vLLM-Omni · MiniCPM-o 4.5 Online Demo
+
+Gradio-based web UI for **MiniCPM-o 4.5** served via `vllm-omni`'s
+OpenAI-compatible endpoints.
+
+The UI supports:
+
+- **Inputs**: text prompt + optional image, audio (file or mic), video.
+- **Outputs**: text + speech (WAV player).
+
+## 1. Start the backend server
+
+The deploy config auto-loads via `--omni`; the default
+`vllm_omni/deploy/minicpmo_4_5.yaml` targets a 2-GPU layout (thinker on
+GPU 0, talker + t2w sharing GPU 1).  For other hardware layouts pick
+one of the deploy variants below.
+
+| deploy config | GPUs | Notes |
+|---|---|---|
+| `minicpmo_4_5.yaml` (default) | 2 | Thinker on GPU0, talker+t2w on GPU1. |
+| `minicpmo_4_5_3gpu.yaml` | 3 | Thinker 2-way TP on GPU0/1, talker+t2w share GPU2. |
+| `minicpmo_4_5_8x4090.yaml` | 8 | Full 8x4090 layout. |
+
+Default (2-GPU):
+
+```bash
+vllm serve openbmb/MiniCPM-o-4_5 --omni \
+    --trust-remote-code \
+    --host 0.0.0.0 --port 8099
+```
+
+Other layouts via `--deploy-config`:
+
+```bash
+vllm serve openbmb/MiniCPM-o-4_5 --omni \
+    --deploy-config vllm_omni/deploy/minicpmo_4_5_8x4090.yaml \
+    --trust-remote-code \
+    --host 0.0.0.0 --port 8099
+```
+
+## 2. Launch the Gradio demo
+
+```bash
+bash examples/online_serving/minicpmo/run_gradio_demo.sh
+
+# Or run the python entry point directly:
+python examples/online_serving/minicpmo/gradio_demo.py \
+    --minicpmo45-api-base http://localhost:8099/v1 \
+    --minicpmo45-model openbmb/MiniCPM-o-4_5 \
+    --port 7862
+```
+
+Open `http://<host>:7862` in a browser.
+
+## Notes
+
+- **TTS trigger**: the demo sets
+  `extra_body.chat_template_kwargs.use_tts_template=True`, which appends
+  `<|tts_bos|>` to the assistant prefix.
+- Uncheck **"Generate speech output (TTS)"** to get text-only responses
+  (faster).
+- The audio output is the raw WAV returned by the stage-1 talker +
+  Token2Wav; sample rate is 24 kHz.
+- Video input is forwarded as a base64 `video_url` entry; the server needs
+  decord/torchvision to decode it.

examples/online_serving/minicpmo/gradio_demo.py

@@ -0,0 +1,580 @@
+"""Gradio demo for MiniCPM-o 4.5 online serving through vllm-omni.
+
+Supports:
+- Text / Image / Audio / Video inputs (OpenAI chat-completions multimodal format).
+- Text + audio outputs, returned side-by-side in the UI.
+
+Usage:
+    # run a 4.5 server on :8099
+    python gradio_demo.py \
+        --minicpmo45-api-base http://localhost:8099/v1 \
+        --minicpmo45-model openbmb/MiniCPM-o-4_5 \
+        --port 7862
+"""
+
+from __future__ import annotations
+
+import argparse
+import base64
+import io
+import os
+from collections.abc import Iterable
+from pathlib import Path
+from typing import Any
+
+import gradio as gr
+import numpy as np
+import soundfile as sf
+from openai import OpenAI
+from PIL import Image
+
+MINICPMO45 = "MiniCPM-o-4.5"
+
+# Default reference audio per model, relative to the model path. Used to build
+# the audio_assistant system prompt so that the LLM adopts the "speak in this
+# voice" persona (matches the official MiniCPM-o demo).
+_DEFAULT_REF_AUDIO = {
+    MINICPMO45: "assets/HT_ref_audio.wav",
+}
+
+# audio_assistant prompt text, mirrors MiniCPMO.get_sys_prompt("audio_assistant")
+# in the MiniCPM-o-Demo reference implementation.
+_AUDIO_ASSISTANT_PROMPT = {
+    "zh": {
+        "prefix": "模仿音频样本的音色并生成新的内容。",
+        "suffix": (
+            "你的任务是用这种声音模式来当一个助手。请认真、高质量地回复用户的问题。"
+            "请用高自然度的方式和用户聊天。你是由面壁智能开发的人工智能助手：面壁小钢炮。"
+        ),
+    },
+    "en": {
+        "prefix": "Use the voice in the audio prompt to synthesize new content.",
+        "suffix": "You are a helpful assistant with the above voice style.",
+    },
+}
+
+
+# ---------------------------------------------------------------------------
+# Media helpers
+# ---------------------------------------------------------------------------
+
+
+def image_to_base64_data_url(image: Image.Image) -> str:
+    buf = io.BytesIO()
+    if image.mode != "RGB":
+        image = image.convert("RGB")
+    image.save(buf, format="JPEG", quality=92)
+    b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
+    return f"data:image/jpeg;base64,{b64}"
+
+
+def audio_to_base64_data_url(audio_np: np.ndarray, sample_rate: int) -> str:
+    if audio_np.dtype != np.int16:
+        if audio_np.dtype in (np.float32, np.float64):
+            audio_np = np.clip(audio_np, -1.0, 1.0)
+            audio_np = (audio_np * 32767).astype(np.int16)
+        else:
+            audio_np = audio_np.astype(np.int16)
+    buf = io.BytesIO()
+    sf.write(buf, audio_np, sample_rate, format="WAV")
+    b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
+    return f"data:audio/wav;base64,{b64}"
+
+
+def ref_audio_to_data_url(path: str, target_sr: int = 16000, max_s: float | None = 8.0) -> str:
+    """Load a file, downmix to mono, resample to 16kHz, optionally truncate, and
+    return a data-URL usable as audio_url content in the chat API.
+
+    MiniCPM-o expects 16kHz mono audio input; long reference audios can waste
+    tokens so we cap at max_s seconds by default (3-10s is the recommended range).
+    """
+    data, sr = sf.read(path)
+    if data.ndim > 1:
+        data = data[:, 0]
+    data = data.astype(np.float32)
+    if sr != target_sr:
+        try:
+            # Demo client only; librosa here resamples audio before sending to
+            # the server.  The fallback below covers environments without it.
+            import librosa  # noqa: TID251
+
+            data = librosa.resample(data, orig_sr=sr, target_sr=target_sr)
+        except ImportError:
+            from math import gcd
+
+            g = gcd(sr, target_sr)
+            up, down = target_sr // g, sr // g
+            idx = (np.arange(len(data) * up) / up).astype(np.int64)
+            data = data[idx][::down]
+        sr = target_sr
+    if max_s is not None and len(data) > int(max_s * sr):
+        data = data[: int(max_s * sr)]
+    return audio_to_base64_data_url(data, sr)
+
+
+def video_to_base64_data_url(path: str) -> str:
+    p = Path(path)
+    if not p.exists():
+        raise FileNotFoundError(f"Video not found: {path}")
+    ext = p.suffix.lower()
+    mime = {
+        ".mp4": "video/mp4",
+        ".webm": "video/webm",
+        ".mov": "video/quicktime",
+        ".avi": "video/x-msvideo",
+        ".mkv": "video/x-matroska",
+    }.get(ext, "video/mp4")
+    b64 = base64.b64encode(p.read_bytes()).decode("utf-8")
+    return f"data:{mime};base64,{b64}"
+
+
+def process_audio_input(audio_input: Any | None) -> tuple[np.ndarray, int] | None:
+    """Normalize Gradio audio input to (np.ndarray mono float32, sample_rate).
+
+    Gradio's `gr.Audio(type="numpy")` yields (sr, ndarray); we also accept
+    filesystem paths and a few nested tuple variants that appear in different
+    Gradio versions.
+    """
+    if audio_input is None:
+        return None
+
+    def _from_path(p: str) -> tuple[np.ndarray, int] | None:
+        if not p:
+            return None
+        fp = Path(p)
+        if not fp.exists():
+            return None
+        data, sr = sf.read(fp)
+        if data.ndim > 1:
+            data = data[:, 0]
+        return data.astype(np.float32), int(sr)
+
+    audio_np: np.ndarray | None = None
+    sr: int | None = None
+
+    if isinstance(audio_input, tuple) and len(audio_input) == 2:
+        a, b = audio_input
+        if isinstance(a, int | float) and isinstance(b, np.ndarray):
+            sr, audio_np = int(a), b
+        elif isinstance(a, str):
+            loaded = _from_path(a)
+            if loaded is not None:
+                audio_np, sr = loaded
+    elif isinstance(audio_input, str):
+        loaded = _from_path(audio_input)
+        if loaded is not None:
+            audio_np, sr = loaded
+
+    if audio_np is None or sr is None:
+        return None
+    if audio_np.ndim > 1:
+        audio_np = audio_np[:, 0]
+    return audio_np.astype(np.float32), sr
+
+
+# ---------------------------------------------------------------------------
+# Per-model inference config
+# ---------------------------------------------------------------------------
+
+
+class ModelEndpoint:
+    """Endpoint + TTS-trigger convention for one MiniCPM-o variant."""
+
+    def __init__(self, name: str, api_base: str, model_path: str):
+        self.name = name
+        self.api_base = api_base
+        self.model_path = model_path
+        self.client = OpenAI(api_key="EMPTY", base_url=api_base)
+        self._ref_audio_data_url: str | None = None
+
+    @property
+    def ref_audio_path(self) -> str | None:
+        rel = _DEFAULT_REF_AUDIO.get(self.name)
+        if not rel:
+            return None
+        full = os.path.join(self.model_path, rel)
+        return full if os.path.exists(full) else None
+
+    def get_ref_audio_data_url(self) -> str | None:
+        """Lazily load and cache the default reference audio as a 16kHz mono
+        data URL. Returns None if the reference audio file is missing."""
+        if self._ref_audio_data_url is not None:
+            return self._ref_audio_data_url
+        p = self.ref_audio_path
+        if not p:
+            print(f"[{self.name}] no default reference audio found; falling back to text-only system prompt")
+            return None
+        try:
+            self._ref_audio_data_url = ref_audio_to_data_url(p)
+            print(f"[{self.name}] loaded reference audio: {p}")
+        except Exception as e:
+            print(f"[{self.name}] failed to load reference audio {p}: {e}")
+            return None
+        return self._ref_audio_data_url
+
+    def build_extras(self) -> dict[str, Any]:
+        """Extra kwargs passed to chat.completions.create for this model.
+
+        4.5: use_tts_template=True chat-template flag appends <|tts_bos|>.
+        """
+        if self.name == MINICPMO45:
+            return {
+                "extra_body": {
+                    "chat_template_kwargs": {"use_tts_template": True},
+                    "modalities": ["text", "audio"],
+                },
+            }
+        return {}
+
+
+# ---------------------------------------------------------------------------
+# System prompt (matches MiniCPMO.get_sys_prompt("audio_assistant") in the
+# official MiniCPM-o-Demo). When TTS is disabled we fall back to a plain
+# text-only system prompt.
+# ---------------------------------------------------------------------------
+
+
+def build_audio_assistant_system(
+    ref_audio_data_url: str | None,
+    language: str = "zh",
+) -> dict[str, Any]:
+    """Return an OpenAI-style system message carrying the reference audio."""
+    txt = _AUDIO_ASSISTANT_PROMPT.get(language, _AUDIO_ASSISTANT_PROMPT["zh"])
+    if ref_audio_data_url is None:
+        return {
+            "role": "system",
+            "content": (
+                "You are MiniCPM-o, a helpful multimodal assistant that can "
+                "understand images, audio and video, and respond in text and speech."
+            ),
+        }
+    return {
+        "role": "system",
+        "content": [
+            {"type": "text", "text": txt["prefix"]},
+            {"type": "audio_url", "audio_url": {"url": ref_audio_data_url}},
+            {"type": "text", "text": txt["suffix"]},
+        ],
+    }
+
+
+# ---------------------------------------------------------------------------
+# Core inference function
+# ---------------------------------------------------------------------------
+
+
+def run_inference(
+    endpoint: ModelEndpoint,
+    user_prompt: str,
+    audio_tuple: tuple[np.ndarray, int] | None,
+    image_pil: Image.Image | None,
+    video_path: str | None,
+    enable_tts: bool,
+    max_tokens: int,
+    temperature: float,
+) -> Iterable[tuple[str, tuple[int, np.ndarray] | None]]:
+    user_prompt = user_prompt or ""
+    if not user_prompt.strip() and not audio_tuple and not image_pil and not video_path:
+        yield "Please provide a text prompt or at least one multimodal input.", None
+        return
+
+    try:
+        content: list[dict[str, Any]] = []
+
+        if audio_tuple is not None:
+            audio_np, sr = audio_tuple
+            content.append(
+                {"type": "audio_url", "audio_url": {"url": audio_to_base64_data_url(audio_np, sr)}},
+            )
+
+        if image_pil is not None:
+            img = image_pil if image_pil.mode == "RGB" else image_pil.convert("RGB")
+            content.append(
+                {"type": "image_url", "image_url": {"url": image_to_base64_data_url(img)}},
+            )
+
+        if video_path:
+            content.append(
+                {"type": "video_url", "video_url": {"url": video_to_base64_data_url(video_path)}},
+            )
+
+        if user_prompt.strip():
+            content.append({"type": "text", "text": user_prompt})
+
+        if enable_tts:
+            ref_url = endpoint.get_ref_audio_data_url()
+            system_msg = build_audio_assistant_system(ref_url, language="zh")
+        else:
+            system_msg = {
+                "role": "system",
+                "content": (
+                    "You are MiniCPM-o, a helpful multimodal assistant that can "
+                    "understand images, audio and video, and respond in text and speech."
+                ),
+            }
+
+        messages = [
+            system_msg,
+            {"role": "user", "content": content},
+        ]
+
+        kwargs: dict[str, Any] = {
+            "model": endpoint.model_path,
+            "messages": messages,
+            "max_tokens": int(max_tokens),
+            "temperature": float(temperature),
+        }
+        if enable_tts:
+            extras = endpoint.build_extras()
+            # Merge extra_body dicts cleanly.
+            extra_body = dict(extras.get("extra_body", {}))
+            kwargs["extra_body"] = extra_body
+        else:
+            # Text-only: suppress audio output.
+            kwargs["extra_body"] = {"modalities": ["text"]}
+
+        completion = endpoint.client.chat.completions.create(**kwargs)
+
+        text_parts: list[str] = []
+        audio_out: tuple[int, np.ndarray] | None = None
+
+        for choice in completion.choices:
+            msg = choice.message
+            if getattr(msg, "content", None):
+                text_parts.append(msg.content)
+            audio_obj = getattr(msg, "audio", None)
+            if audio_obj and getattr(audio_obj, "data", None):
+                raw = base64.b64decode(audio_obj.data)
+                if len(raw) > 80:
+                    data, sr = sf.read(io.BytesIO(raw))
+                    if data.ndim > 1:
+                        data = data[:, 0]
+                    audio_out = (int(sr), data.astype(np.float32))
+
+        text_response = "\n\n".join(t for t in text_parts if t) or "(no text returned)"
+        yield text_response, audio_out
+    except Exception as exc:  # noqa: BLE001
+        yield f"Inference failed: {type(exc).__name__}: {exc}", None
+
+
+# ---------------------------------------------------------------------------
+# UI
+# ---------------------------------------------------------------------------
+
+DEFAULT_PROMPTS = [
+    "你好，请用一句话介绍北京。",
+    "Describe what you see in the image.",
+    "Transcribe the audio and translate it into English.",
+    "What is happening in this video?",
+]
+
+
+def build_interface(endpoints: dict[str, ModelEndpoint], default_model: str) -> gr.Blocks:
+    model_choices = list(endpoints.keys())
+
+    css = """
+    #generate-btn button { width: 100%; font-size: 1.05rem; }
+    .media-row > div { min-height: 260px; }
+    .model-banner { color: #555; font-size: 0.9rem; }
+    """
+
+    with gr.Blocks(css=css, title="vLLM-Omni · MiniCPM-o Demo") as demo:
+        gr.Markdown(
+            "# vLLM-Omni · MiniCPM-o Online Demo\n"
+            "Multimodal chat (text / image / audio / video in; text + speech out) for "
+            "MiniCPM-o 4.5 served by vllm-omni."
+        )
+
+        with gr.Row():
+            model_dd = gr.Dropdown(
+                label="Model",
+                choices=model_choices,
+                value=default_model,
+                scale=2,
+            )
+            endpoint_info = gr.Markdown(
+                elem_classes="model-banner",
+                value=f"`{endpoints[default_model].api_base}`  →  `{endpoints[default_model].model_path}`",
+            )
+
+        with gr.Row():
+            prompt_box = gr.Textbox(
+                label="Text Prompt",
+                placeholder="Ask the model anything. Combine with image / audio / video below.",
+                lines=3,
+                value=DEFAULT_PROMPTS[0],
+                scale=3,
+            )
+            with gr.Column(scale=1):
+                tts_checkbox = gr.Checkbox(label="Generate speech output (TTS)", value=True)
+                max_tokens_slider = gr.Slider(label="Max tokens", minimum=32, maximum=4096, value=1024, step=32)
+                temperature_slider = gr.Slider(label="Temperature", minimum=0.0, maximum=1.5, value=0.7, step=0.05)
+
+        with gr.Row(elem_classes="media-row"):
+            image_input = gr.Image(label="Image (optional)", type="pil", sources=["upload"])
+            audio_input = gr.Audio(
+                label="Audio (optional)",
+                type="numpy",
+                sources=["upload", "microphone"],
+            )
+            video_input = gr.Video(label="Video (optional)", sources=["upload"])
+
+        with gr.Row():
+            generate_btn = gr.Button("Generate", variant="primary", elem_id="generate-btn")
+            clear_btn = gr.Button("Clear")
+
+        with gr.Row():
+            text_output = gr.Textbox(label="Text Output", lines=10, scale=2)
+            audio_output = gr.Audio(label="Audio Output", interactive=False, scale=1)
+
+        gr.Examples(
+            examples=[
+                ["你好，请用一句话介绍北京。"],
+                ["Describe what you see in detail."],
+                ["请先识别这段语音的内容，再用中文回复说话人。"],
+                ["What is happening in this video?"],
+            ],
+            inputs=[prompt_box],
+            label="Quick prompts",
+        )
+
+        def _on_model_change(name: str) -> str:
+            ep = endpoints[name]
+            return f"`{ep.api_base}`  →  `{ep.model_path}`"
+
+        model_dd.change(_on_model_change, inputs=[model_dd], outputs=[endpoint_info])
+
+        def _run(
+            model_name: str,
+            prompt: str,
+            image_pil: Image.Image | None,
+            audio_in: Any,
+            video_path: str | None,
+            enable_tts: bool,
+            max_tokens: int,
+            temperature: float,
+        ):
+            ep = endpoints[model_name]
+            audio_tuple = process_audio_input(audio_in)
+            yield from run_inference(
+                ep,
+                prompt,
+                audio_tuple,
+                image_pil,
+                video_path,
+                enable_tts,
+                max_tokens,
+                temperature,
+            )
+
+        generate_btn.click(
+            _run,
+            inputs=[
+                model_dd,
+                prompt_box,
+                image_input,
+                audio_input,
+                video_input,
+                tts_checkbox,
+                max_tokens_slider,
+                temperature_slider,
+            ],
+            outputs=[text_output, audio_output],
+        )
+
+        def _clear():
+            return "", None, None, None, "", None
+
+        clear_btn.click(
+            _clear,
+            outputs=[
+                prompt_box,
+                image_input,
+                audio_input,
+                video_input,
+                text_output,
+                audio_output,
+            ],
+        )
+
+        demo.queue()
+    return demo
+
+
+# ---------------------------------------------------------------------------
+# Entry point
+# ---------------------------------------------------------------------------
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("--minicpmo45-api-base", default=os.environ.get("MINICPMO45_API_BASE", "http://localhost:8099/v1"))
+    p.add_argument(
+        "--minicpmo45-model",
+        default=os.environ.get("MINICPMO45_MODEL", "openbmb/MiniCPM-o-4_5"),
+    )
+    p.add_argument("--host", default="0.0.0.0")
+    p.add_argument("--port", type=int, default=7862)
+    p.add_argument("--share", action="store_true")
+    p.add_argument(
+        "--ssl-certfile",
+        default=os.environ.get("GRADIO_SSL_CERTFILE", ""),
+        help="Path to TLS certificate (PEM). Browsers require HTTPS to access the microphone.",
+    )
+    p.add_argument(
+        "--ssl-keyfile",
+        default=os.environ.get("GRADIO_SSL_KEYFILE", ""),
+        help="Path to TLS private key (PEM).",
+    )
+    p.add_argument(
+        "--ssl-verify",
+        action="store_true",
+        default=os.environ.get("GRADIO_SSL_VERIFY", "0") == "1",
+        help="Require browser to verify the certificate (off by default; self-signed certs should stay off).",
+    )
+    return p.parse_args()
+
+
+def _ping(api_base: str, timeout: float = 3.0) -> bool:
+    import urllib.request
+
+    try:
+        with urllib.request.urlopen(api_base.rstrip("/").replace("/v1", "/health"), timeout=timeout):
+            return True
+    except Exception:
+        return False
+
+
+def main() -> None:
+    args = parse_args()
+    endpoints: dict[str, ModelEndpoint] = {}
+
+    if args.minicpmo45_api_base and args.minicpmo45_model:
+        ok = _ping(args.minicpmo45_api_base)
+        status = "reachable" if ok else "not reachable (will error on generate)"
+        print(f"[4.5] {args.minicpmo45_api_base}  ({status})")
+        endpoints[MINICPMO45] = ModelEndpoint(MINICPMO45, args.minicpmo45_api_base, args.minicpmo45_model)
+
+    if not endpoints:
+        raise SystemExit("No endpoints configured. Pass --minicpmo45-api-base/--minicpmo45-model.")
+
+    default_model = next(iter(endpoints.keys()))
+    demo = build_interface(endpoints, default_model)
+
+    launch_kwargs: dict = {
+        "server_name": args.host,
+        "server_port": args.port,
+        "share": args.share,
+    }
+    if args.ssl_certfile and args.ssl_keyfile:
+        launch_kwargs.update(
+            ssl_certfile=args.ssl_certfile,
+            ssl_keyfile=args.ssl_keyfile,
+            ssl_verify=args.ssl_verify,
+        )
+        print(f"[tls] cert={args.ssl_certfile} key={args.ssl_keyfile} verify={args.ssl_verify}")
+    demo.launch(**launch_kwargs)
+
+
+if __name__ == "__main__":
+    main()

examples/online_serving/minicpmo/run_gradio_demo.sh

@@ -0,0 +1,36 @@
+#!/bin/bash
+# Launch the MiniCPM-o 4.5 gradio demo.
+#
+# Prereq:
+#   Start a vllm-omni OpenAI server for MiniCPM-o 4.5 on :8099 (see the
+#   8x4090 stage config under vllm_omni/model_executor/stage_configs).
+set -e
+
+HERE="$(cd "$(dirname "$0")" && pwd)"
+
+: "${MINICPMO45_API_BASE:=http://localhost:8099/v1}"
+: "${MINICPMO45_MODEL:=openbmb/MiniCPM-o-4_5}"
+: "${GRADIO_HOST:=0.0.0.0}"
+: "${GRADIO_PORT:=7862}"
+# HTTPS (browsers require a secure context for microphone access).
+# Set GRADIO_SSL_CERTFILE / GRADIO_SSL_KEYFILE to enable TLS.
+: "${GRADIO_SSL_CERTFILE:=}"
+: "${GRADIO_SSL_KEYFILE:=}"
+
+export MINICPMO45_API_BASE MINICPMO45_MODEL
+
+SSL_ARGS=()
+if [ -n "$GRADIO_SSL_CERTFILE" ] && [ -n "$GRADIO_SSL_KEYFILE" ] \
+   && [ -f "$GRADIO_SSL_CERTFILE" ] && [ -f "$GRADIO_SSL_KEYFILE" ]; then
+  SSL_ARGS=(--ssl-certfile "$GRADIO_SSL_CERTFILE" --ssl-keyfile "$GRADIO_SSL_KEYFILE")
+  echo "HTTPS enabled: cert=$GRADIO_SSL_CERTFILE key=$GRADIO_SSL_KEYFILE"
+else
+  echo "HTTPS disabled (cert/key not found). Microphone won't work on remote browsers."
+fi
+
+exec python "$HERE/gradio_demo.py" \
+    --minicpmo45-api-base "$MINICPMO45_API_BASE" \
+    --minicpmo45-model "$MINICPMO45_MODEL" \
+    --host "$GRADIO_HOST" \
+    --port "$GRADIO_PORT" \
+    "${SSL_ARGS[@]}"

examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py

@@ -6,7 +6,8 @@ from typing import NamedTuple
 import requests
 from openai import OpenAI
 from vllm.assets.audio import AudioAsset
-from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 
 SEED = 42
 
@@ -197,7 +198,14 @@ def _build_prompt_for_query_type(
         return query_func(custom_prompt=custom_prompt)
     if query_type == "use_audio_in_video":
         return query_func(video_path=video_path, custom_prompt=custom_prompt)
-    # use_mixed_modalities / use_multi_audios
+    if query_type == "use_mixed_modalities":
+        return query_func(
+            video_path=video_path,
+            image_path=image_path,
+            audio_path=audio_path,
+            custom_prompt=custom_prompt,
+        )
+    # use_multi_audios
     return query_func(custom_prompt=custom_prompt)
 
 
@@ -478,7 +486,7 @@ def run_multimodal_generation(args, client: OpenAI) -> None:
 
 
 def parse_args():
-    parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
+    parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
     parser.add_argument(
         "--query-type",
         "-q",

pyproject.toml

@@ -76,6 +76,18 @@ local = [
     "FlagEmbedding",
 ]
 
+# Optional runtime dependencies for the MiniCPM-o 4.5 talker pipeline.
+# The talker loads MiniCPMTTS via `trust_remote_code` from the model repo,
+# and that module performs `from stepaudio2 import Token2wav` against the
+# MiniCPM-o-flavored Token2wav vocoder (PyPI: `stepaudio2-minicpmo`, NOT
+# upstream `stepfun-ai/Step-Audio2` — the upstream signature rejects
+# `n_timesteps` and crashes at construction). A clean `vllm-omni` install
+# does not pull this in by default to avoid forcing the dep onto users
+# who only need other models; opt in with `pip install vllm-omni[minicpmo]`.
+minicpmo = [
+    "stepaudio2-minicpmo",
+]
+
 docs = [
     "mkdocs>=1.5.0",
     "mkdocs-api-autonav",
@@ -101,6 +113,13 @@ Documentation = "https://vllm-omni.readthedocs.io"
 [project.scripts]
 vllm-omni = "vllm_omni.entrypoints.cli.main:main"
 
+# Register vllm-omni's model registrations on every vllm subprocess. vllm
+# auto-loads anything advertised under `vllm.general_plugins` in its own
+# `load_general_plugins()`, so worker subprocesses spawned by vllm (which
+# only import vllm, not vllm_omni) will still pick up the omni archs.
+[project.entry-points."vllm.general_plugins"]
+vllm_omni_register_models = "vllm_omni.engine.arg_utils:register_omni_models_to_vllm"
+
 
 [tool.setuptools.packages.find]
 where = ["."]
@@ -149,6 +168,10 @@ ignore = [
 [tool.ruff.lint.per-file-ignores]
 "examples/**" = ["E501"]  # Allow long lines in examples
 "tests/**" = ["E501"]  # Allow long lines in tests
+# MiniCPM-o reference modeling code vendored from upstream HF/transformers/MiniCPM-o.
+# 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"]
 
 [tool.ruff.lint.flake8-tidy-imports.banned-api]
 "librosa".msg = "The librosa module is banned, use vllm.multimodal helpers instead"

recipes/ByteDance/Lance.md

@@ -0,0 +1,177 @@
+# Lance
+
+> Unified autoregressive + diffusion multimodal (text / image / video)
+
+## Summary
+
+- Vendor: ByteDance
+- Model: [`bytedance-research/Lance`](https://huggingface.co/bytedance-research/Lance) (`Lance_3B`, `Lance_3B_Video`)
+- Task: text2img, text2video, img2img (image edit), video2video (video edit),
+  img2text (image understanding), video2text (video understanding)
+- Mode: Offline inference, Online serving (OpenAI-compatible API)
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe to run Lance — a 3B unified AR + diffusion multimodal model
+on a Qwen2.5-VL backbone — via vLLM-Omni. Lance is **BAGEL-lineage**: the
+released checkpoint uses the same `*_moe_gen` Mixture-of-Transformers weight
+layout as BAGEL, so vLLM-Omni reuses BAGEL's transformer core and specializes
+only the ViT (Qwen2.5-VL vision), the VAE (Wan2.2) and the checkpoint layout.
+
+All six single-stage modalities are supported on the `Lance_3B` checkpoint;
+`text2video` and `video2video` additionally need `Lance_3B_Video` for the 3-D
+`latent_pos_embed` table. The two-stage AR + DiT topology is deferred and
+needs `LanceConfig` / `LanceProcessor` registered in the `vllm` package (a
+separate upstream PR).
+
+## References
+
+- Offline example:
+  [`examples/offline_inference/lance/end2end.py`](../../examples/offline_inference/lance/end2end.py)
+- Online serving:
+  [`examples/online_serving/lance/`](../../examples/online_serving/lance/)
+- Deploy config:
+  [`vllm_omni/deploy/lance.yaml`](../../vllm_omni/deploy/lance.yaml)
+- E2E tests:
+  [`tests/e2e/online_serving/test_lance.py`](../../tests/e2e/online_serving/test_lance.py)
+- HuggingFace model page:
+  [bytedance-research/Lance](https://huggingface.co/bytedance-research/Lance)
+- PR: [vllm-project/vllm-omni#3710](https://github.com/vllm-project/vllm-omni/pull/3710)
+  (closes [#3697](https://github.com/vllm-project/vllm-omni/issues/3697))
+
+## Hardware Support
+
+## GPU
+
+### 1x NVIDIA B300 / A100 80GB
+
+#### Environment
+
+- OS: Linux
+- Python: 3.12
+- Driver / runtime: CUDA ≥ 12.4
+- vLLM-Omni version: 0.18.x.dev
+
+#### Command — text-to-image (default)
+
+```bash
+python examples/offline_inference/lance/end2end.py \
+    --model bytedance-research/Lance \
+    --prompts "a corgi astronaut on the moon, cinematic" \
+    --steps 30 --cfg-text-scale 4.0 --timestep-shift 3.5 \
+    --height 1024 --width 1024 \
+    --seed 42 --output ./out
+```
+
+Defaults match upstream `inference_lance.sh`: 30 denoising steps,
+`timestep-shift 3.5`, text CFG 4.0, seed 42, 1024×1024.
+
+#### Command — image edit (img2img)
+
+```bash
+python examples/offline_inference/lance/end2end.py \
+    --model bytedance-research/Lance --modality img2img \
+    --image-path /path/to/input.png \
+    --prompts "Convert this into a vibrant cartoon-style illustration" \
+    --steps 30 --cfg-text-scale 4.0 --timestep-shift 3.5 \
+    --output ./out
+```
+
+The Lance-native VAE prefill scatters Wan2.2 latents into the LLM query
+sequence; no separate image encoder is needed.
+
+#### Command — text-to-video
+
+```bash
+python examples/offline_inference/lance/end2end.py \
+    --model bytedance-research/Lance/Lance_3B_Video --modality text2video \
+    --num-frames 25 --video-height 480 --video-width 768 \
+    --prompts "a cat playing piano, cinematic" \
+    --steps 30 --fps 8 --output ./out
+```
+
+Use the `Lance_3B_Video` subfolder for any video path so the 3-D
+`latent_pos_embed` table is loaded; image / understanding paths can point at
+the top-level repo and resolve the right sub-checkpoint automatically.
+
+#### Command — image / video understanding
+
+```bash
+# Image → text (caption / VQA)
+python examples/offline_inference/lance/end2end.py \
+    --model bytedance-research/Lance --modality img2text \
+    --image-path /path/to/photo.jpg \
+    --prompts "Describe this image in detail." \
+    --do-sample --text-temperature 0.8
+
+# Video → text
+python examples/offline_inference/lance/end2end.py \
+    --model bytedance-research/Lance --modality video2text \
+    --video-path /path/to/clip.mp4 \
+    --prompts "What is happening in this video?"
+```
+
+Sampling is enabled by default at `--text-temperature 0.8` for the
+understanding paths because Lance's greedy decoder emits an immediate EOS for
+many prompts.
+
+#### Verification
+
+```bash
+pytest -s -v tests/e2e/online_serving/test_lance.py
+```
+
+#### Notes
+
+- BF16 footprint: ~7 GB LLM + Qwen2.5-VL ViT + Wan2.2 VAE; comfortably fits
+  on a single 16 GB+ GPU for `Lance_3B`.
+- `rope_scaling = {"type": "mrope", "mrope_section": [16, 24, 24]}` is wired
+  through `BagelRotaryEmbedding`. `text2img` uses scalar position ids
+  (BAGEL-equivalent); `img2text` / `video` / edit paths thread per-axis 3-D
+  position ids.
+- `video_edit` quality is more abstract than `text2video` at the same
+  resolution (known position-id offset between VAE-prefill and gen-latent
+  blocks). Functionally correct end-to-end.
+
+## Online Serving
+
+Lance supports all single-stage modalities via the OpenAI-compatible
+`/v1/chat/completions` API.
+
+### Launch
+
+```bash
+bash examples/online_serving/lance/run_server.sh
+# or, with overrides
+MODEL=bytedance-research/Lance \
+DEPLOY_CONFIG=vllm_omni/deploy/lance.yaml \
+PORT=8091 \
+    bash examples/online_serving/lance/run_server.sh
+```
+
+For `text2video` / `video2video`, set
+`MODEL=bytedance-research/Lance/Lance_3B_Video`.
+
+### Send requests
+
+```bash
+# Text-to-image
+python examples/online_serving/lance/openai_chat_client.py \
+    --prompt "A cute corgi astronaut on the moon, cinematic" \
+    --modality text2img --output corgi.png
+
+# Image edit
+python examples/online_serving/lance/openai_chat_client.py \
+    --prompt "Convert this into a vibrant cartoon-style illustration" \
+    --modality img2img --image-url path/to/photo.png \
+    --output edited.png
+
+# Image understanding
+python examples/online_serving/lance/openai_chat_client.py \
+    --prompt "Describe this image" \
+    --modality img2text --image-url photo.jpg
+```
+
+The client is shared with BAGEL — same OpenAI message format, same
+`modalities` / `num_inference_steps` / `seed` / `height` / `width` knobs.

recipes/OpenBMB/MiniCPM-o-4_5.md

@@ -0,0 +1,229 @@
+# MiniCPM-o 4.5
+
+> Online serving for omni multimodal chat (text / image / audio / video → text + 24 kHz speech)
+
+## Summary
+
+- Vendor: OpenBMB
+- Model: [`openbmb/MiniCPM-o-4_5`](https://huggingface.co/openbmb/MiniCPM-o-4_5)
+- Task: Omni multimodal chat — accepts text / image / audio / video input;
+  emits text and 24 kHz mono speech in the same response
+- Mode: Online serving via the OpenAI-compatible `/v1/chat/completions`
+  API, plus a bundled Gradio demo (text + speech UI)
+- Maintainer: [`@tc-mb`](https://github.com/tc-mb) (MiniCPM-V / MiniCPM-o team)
+
+## When to use this recipe
+
+Use this recipe as a known-good starting point for serving
+`openbmb/MiniCPM-o-4_5` on vLLM-Omni. MiniCPM-o 4.5 is the omni member
+of the MiniCPM-o family — it pairs a multimodal-understanding thinker
+LLM with a streaming `MiniCPMTTS + Token2Wav` talker so a single
+`/v1/chat/completions` call can return text and 24 kHz speech in one
+shot. The recipe covers three shipped GPU layouts (2 / 3 / 8 GPUs)
+selected via `--deploy-config`.
+
+## References
+
+- Default deploy configs (auto-loaded by HF `model_type=minicpmo` +
+  `hf_config.version="4.5"`):
+  - 2-GPU layout (default):
+    [`vllm_omni/deploy/minicpmo_4_5.yaml`](../../vllm_omni/deploy/minicpmo_4_5.yaml)
+  - 3-GPU layout (thinker TP=2):
+    [`vllm_omni/deploy/minicpmo_4_5_3gpu.yaml`](../../vllm_omni/deploy/minicpmo_4_5_3gpu.yaml)
+  - 8x RTX 4090 layout:
+    [`vllm_omni/deploy/minicpmo_4_5_8x4090.yaml`](../../vllm_omni/deploy/minicpmo_4_5_8x4090.yaml)
+- Online example + Gradio demo:
+  [`examples/online_serving/minicpmo/`](../../examples/online_serving/minicpmo/)
+- Pipeline / talker source:
+  [`vllm_omni/model_executor/models/minicpmo_4_5/`](../../vllm_omni/model_executor/models/minicpmo_4_5/)
+- Stage-input processor (thinker → talker bridge):
+  [`vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py`](../../vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py)
+- Upstream model card:
+  [`openbmb/MiniCPM-o-4_5`](https://huggingface.co/openbmb/MiniCPM-o-4_5)
+- Integration PR:
+  [vllm-project/vllm-omni#3642](https://github.com/vllm-project/vllm-omni/pull/3642)
+
+## Hardware Support
+
+Three GPU layouts ship with default deploy configs. Pick the layout that
+matches your hardware and pass it via `--deploy-config`; the talker
+(`MiniCPMTTS + Token2Wav`) always lives on its own GPU because of the
+in-process vocoder, and the thinker is the part that scales out via TP.
+
+| Layout | Thinker | Talker + Token2Wav | Typical hardware |
+| --- | --- | --- | --- |
+| 2-GPU (default) | GPU 0 | GPU 1 | 2x A100/H100/H200 80GB |
+| 3-GPU (thinker TP=2) | GPU 0,1 (TP=2) | GPU 2 | 3x mid-tier GPUs |
+| 8x RTX 4090 24GB | GPU 0–3 (TP=4) | GPU 4 | 8x RTX 4090 consumer |
+
+## GPU
+
+### 2 x GPU (default — single command)
+
+The default
+[`vllm_omni/deploy/minicpmo_4_5.yaml`](../../vllm_omni/deploy/minicpmo_4_5.yaml)
+puts the thinker on GPU 0 (`~70 %` memory, `enforce_eager: true`,
+`max_num_seqs: 1`) and the talker + Token2Wav vocoder on GPU 1
+(`~75 %` memory). This is the recommended starting layout — works on
+any pair of 80GB-class GPUs (A100, H100, H200) and on most 40GB+
+pairs as long as the thinker model weights fit.
+
+#### Environment
+
+- OS: Linux
+- Python: 3.10+
+- vLLM / vLLM-Omni: >= 0.21.0 (or current `main`)
+- Optional Talker dep: `stepaudio2-minicpmo` (see Notes for why this is
+  required and how to install it)
+
+#### Command
+
+```bash
+vllm serve openbmb/MiniCPM-o-4_5 --omni \
+    --trust-remote-code \
+    --host 0.0.0.0 --port 8099
+```
+
+The deploy config is auto-loaded by the model registry — no
+`--deploy-config` flag needed for this default 2-GPU layout.
+
+#### Verification
+
+**Quick smoke test (text-only output)**:
+
+```bash
+curl http://localhost:8099/v1/chat/completions \
+    -H "Content-Type: application/json" \
+    -d '{
+        "model": "openbmb/MiniCPM-o-4_5",
+        "messages": [{"role": "user", "content": "Briefly introduce yourself."}],
+        "modalities": ["text"]
+    }'
+```
+
+**Text + speech in one response** (the headline 4.5 feature). The TTS
+path is gated by a Jinja flag on the chat template, passed through
+`extra_body.chat_template_kwargs.use_tts_template=true`:
+
+```bash
+curl http://localhost:8099/v1/chat/completions \
+    -H "Content-Type: application/json" \
+    -d '{
+        "model": "openbmb/MiniCPM-o-4_5",
+        "messages": [{"role": "user", "content": "Say hello, then introduce vLLM in one sentence."}],
+        "modalities": ["text", "audio"],
+        "extra_body": {"chat_template_kwargs": {"use_tts_template": true}}
+    }'
+```
+
+Response carries text in `choices[0].message.content` and base64 WAV
+in `choices[0].message.audio.data` (24 kHz mono, see Notes).
+
+**Gradio demo (text + image + audio + video UI)**:
+
+```bash
+bash examples/online_serving/minicpmo/run_gradio_demo.sh
+# or run the python entry point directly:
+python examples/online_serving/minicpmo/gradio_demo.py \
+    --minicpmo45-api-base http://localhost:8099/v1 \
+    --minicpmo45-model openbmb/MiniCPM-o-4_5 \
+    --port 7862
+```
+
+Open `http://<host>:7862` and try a text prompt with the **"Generate
+speech output (TTS)"** checkbox on / off.
+
+#### Notes
+
+- Memory budget: thinker weights occupy GPU 0 at `gpu_memory_utilization:
+  0.7`; talker + Token2Wav vocoder share GPU 1 at `0.75`.
+- `--trust-remote-code` is required — the HF repo ships a custom
+  `MiniCPMO` config / model class.
+- Pin: `enforce_eager: true` on both stages (CUDA graph capture is off
+  by design for the talker's Token2Wav path).
+- Stage 1 (talker) is hard-capped to `max_num_seqs: 1`: the talker
+  only consumes `runtime_additional_information[0]`, so any value > 1
+  makes concurrent requests share request-0's audio. This is the same
+  cap baked into the deploy config.
+
+### 3 x GPU (thinker TP=2)
+
+Use
+[`vllm_omni/deploy/minicpmo_4_5_3gpu.yaml`](../../vllm_omni/deploy/minicpmo_4_5_3gpu.yaml)
+when you have a third GPU available and want the thinker on 2-way
+tensor parallel for higher throughput; the talker stays on its own
+GPU (talker has its own in-process Token2Wav vocoder, so co-locating
+it with the thinker risks OOM under load).
+
+#### Command
+
+```bash
+vllm serve openbmb/MiniCPM-o-4_5 --omni \
+    --deploy-config vllm_omni/deploy/minicpmo_4_5_3gpu.yaml \
+    --trust-remote-code \
+    --host 0.0.0.0 --port 8099
+```
+
+Verification and Notes mirror the 2-GPU section; thinker latency
+roughly halves under load thanks to TP=2.
+
+### 8 x RTX 4090 24GB (consumer-GPU layout)
+
+Use
+[`vllm_omni/deploy/minicpmo_4_5_8x4090.yaml`](../../vllm_omni/deploy/minicpmo_4_5_8x4090.yaml)
+on an 8x RTX 4090 host. Thinker uses 4-way TP across GPUs 0–3
+(`~85 %` mem each ≈ 20.4 GiB/card), talker + Token2Wav lives on GPU 4
+(`~90 %` mem). GPUs 5–7 are left free.
+
+#### Command
+
+```bash
+vllm serve openbmb/MiniCPM-o-4_5 --omni \
+    --deploy-config vllm_omni/deploy/minicpmo_4_5_8x4090.yaml \
+    --trust-remote-code \
+    --host 0.0.0.0 --port 8099
+```
+
+#### Notes
+
+- `max_model_len` is capped at 4096 in this layout — 8192 still OOMs on
+  4090s. Raise it if your cards have more headroom (e.g. 4090 D /
+  custom 32 GB SKUs), but verify with a long-prompt run before
+  promoting.
+- All other knobs match the 2-GPU section; the only difference is the
+  per-card memory pressure on the thinker shards.
+
+## Notes (applies to all layouts)
+
+- **Talker dependency**: the `MiniCPM-o 4.5` talker calls
+  `from stepaudio2 import Token2wav` against the MiniCPM-o-flavored
+  vocoder (PyPI package `stepaudio2-minicpmo` — NOT the upstream
+  `stepfun-ai/Step-Audio2`, whose `Token2wav.__init__` signature
+  rejects `n_timesteps`). Install via the published extra:
+
+  ```bash
+  pip install 'vllm-omni[minicpmo]'
+  ```
+
+  Equivalent direct install: `pip install stepaudio2-minicpmo`. A
+  missing dep raises `ImportError` at first request with the same
+  install hint instead of silently emitting empty audio.
+
+- **TTS trigger**: speech output is only emitted when the client passes
+  `extra_body.chat_template_kwargs.use_tts_template=true`. Without it,
+  the response is text-only (which is also faster).
+
+- **Output audio**: 24 kHz mono WAV inside the OpenAI-style
+  `message.audio.data` (base64). The Gradio demo's WAV player decodes
+  this automatically.
+
+- **Routing**: MiniCPM-o 4.5 and 2.6 both ship `architectures=
+  ["MiniCPMO"]` in HF config; routing is disambiguated by
+  `hf_config.version == "4.5"` via the
+  `hf_config_predicate` on the 4.5 pipeline. A 2.6 checkpoint loaded
+  with this recipe's `--deploy-config` will be rejected at startup
+  rather than silently misrouted.
+
+- **Async chunking**: disabled in all three deploy configs
+  (`async_chunk: false`) — the talker batches a single full thinker
+  output, not chunks.

recipes/OpenBMB/VoxCPM2.md

@@ -0,0 +1,249 @@
+# VoxCPM2
+
+> Online + offline TTS with native AR pipeline (48 kHz, 30+ languages)
+
+## Summary
+
+- Vendor: OpenBMB
+- Model: `openbmb/VoxCPM2`
+- Task: Text-to-speech (zero-shot synthesis and reference-audio voice cloning)
+- Mode: Online serving via the OpenAI-compatible `/v1/audio/speech` API, plus
+  an offline end-to-end example
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe as a known-good starting point for serving `openbmb/VoxCPM2`
+on vLLM-Omni on consumer-class GPUs. VoxCPM2 is a 2B-parameter native AR
+TTS model that emits 48 kHz mono audio and supports 30+ languages out of
+the box. It runs as a single-stage pipeline (`MiniCPM4 base LM → FSQ →
+MiniCPM4 residual LM → LocDiT CFM solver → AudioVAE`) and supports both
+zero-shot synthesis and reference-audio voice cloning, with optional
+gapless streaming via the bundled Gradio AudioWorklet demo.
+
+Because the model is only 2B params, a single 24 GB consumer GPU is
+sufficient — no tensor parallelism is required.
+
+## References
+
+- Example guide (online):
+  [`examples/online_serving/text_to_speech/README.md`](../../examples/online_serving/text_to_speech/README.md#voxcpm2)
+- Online client / Gradio demo:
+  [`examples/online_serving/text_to_speech/voxcpm2/`](../../examples/online_serving/text_to_speech/voxcpm2/)
+- Offline end-to-end script:
+  [`examples/offline_inference/text_to_speech/voxcpm2/end2end.py`](../../examples/offline_inference/text_to_speech/voxcpm2/end2end.py)
+- Default deploy config:
+  [`vllm_omni/deploy/voxcpm2.yaml`](../../vllm_omni/deploy/voxcpm2.yaml)
+- Talker / pipeline source:
+  [`vllm_omni/model_executor/models/voxcpm2/`](../../vllm_omni/model_executor/models/voxcpm2/)
+- Speech API reference:
+  [`docs/serving/speech_api.md`](../../docs/serving/speech_api.md#voxcpm2)
+- Related issue or discussion:
+  [RFC: add recipes folder](https://github.com/vllm-project/vllm-omni/issues/2645)
+
+## Hardware Support
+
+This recipe documents one tested 24 GB consumer-GPU configuration.
+Larger-VRAM (H20 / H100 / A100) and other vendor sections (ROCm, NPU) are
+welcome as community validation lands.
+
+## GPU
+
+### 1 x RTX 4090 24GB (Single GPU, Minimum Recommended)
+
+VoxCPM2 (2B params in bfloat16) fits comfortably on a single 24 GB GPU.
+The bundled default config at
+[`vllm_omni/deploy/voxcpm2.yaml`](../../vllm_omni/deploy/voxcpm2.yaml)
+(`gpu_memory_utilization: 0.9`, `max_num_seqs: 4`, `enforce_eager: true`)
+loads cleanly on a 4090 with ~4.9 GiB for model weights, ~15.2 GiB for
+KV cache, and ~2 GiB for the talker's CUDA-Graph / CFM / VAE buffers,
+for a total resident footprint of about **22 GiB / 24 GiB** in steady
+state.
+
+#### Environment
+
+- OS: Linux (Ubuntu 22.04)
+- Python: 3.12
+- Driver / runtime: NVIDIA driver 575.57.08, CUDA 12.9 runtime
+- torch: 2.11.0+cu128
+- vLLM: 0.21.0
+- vLLM-Omni: 0.20.1.dev (current `main`)
+- Native VoxCPM2 source: `voxcpm` >= 2.0
+
+#### Command
+
+Start the server from the repository root. On a clean 4090 with no other
+tenant on the GPU, the defaults from `voxcpm2.yaml` work as-is:
+
+```bash
+vllm serve openbmb/VoxCPM2 --omni \
+    --host 0.0.0.0 --port 8000
+```
+
+The deploy config at
+[`vllm_omni/deploy/voxcpm2.yaml`](../../vllm_omni/deploy/voxcpm2.yaml)
+is loaded automatically by the model registry (HF `model_type=voxcpm2`).
+Pass `--deploy-config <path>` to override.
+
+#### Verification
+
+**Server cold-start**: ~60 s from `vllm serve` to `Application startup
+complete` (subprocess fork + vLLM 0.21 init + model load + flashinfer JIT
++ torch.compile of LocDiT / feat_encoder / AudioVAE + CUDA-Graph warmup).
+The first request after startup pays a small additional cost; steady-state
+requests are much faster.
+
+**T1 — zero-shot synthesis (curl, WAV)**:
+
+```bash
+curl -X POST http://127.0.0.1:8000/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
+        "model": "voxcpm2",
+        "input": "Hello, this is VoxCPM2 served by vLLM-Omni on a single RTX 4090.",
+        "voice": "default",
+        "response_format": "wav"
+    }' --output output.wav
+```
+
+Observed: `output.wav` 491 KB, 5.12 s @ 48 kHz mono, peak amplitude 0.996.
+Server-to-client wall time: 2.6 s (server-side RTF ~0.5 after warmup).
+
+**T2 — bundled Python client (zero-shot)**:
+
+```bash
+python examples/online_serving/text_to_speech/voxcpm2/openai_speech_client.py \
+    --text "This audio was generated by the bundled openai_speech_client.py script against a vLLM-Omni server." \
+    --api-base http://127.0.0.1:8000 \
+    --output output.wav
+```
+
+Observed: `output.wav` 968 KB, 10.08 s @ 48 kHz mono, peak amplitude 0.949.
+
+**T3 — voice cloning from a reference WAV**:
+
+```bash
+python examples/online_serving/text_to_speech/voxcpm2/openai_speech_client.py \
+    --text "This should sound like the reference speaker, but saying a different sentence on a single RTX 4090." \
+    --ref-audio /path/to/reference.wav \
+    --api-base http://127.0.0.1:8000 \
+    --output cloned.wav
+```
+
+Observed (using the T2 output `output.wav` as `--ref-audio`):
+`cloned.wav` 568 KB, 5.92 s @ 48 kHz mono, peak amplitude 0.990;
+end-to-end wall time 3.16 s.
+
+`--ref-audio` accepts local file paths (auto-base64), HTTP/HTTPS URLs, or
+`data:audio/wav;base64,...` data URIs. The `voice` field is required by
+the OpenAI schema but ignored by VoxCPM2 unless it refers to an
+uploaded voice — voice cloning is driven entirely by `ref_audio`.
+
+**T4 — streaming PCM (raw, 48 kHz s16le)**:
+
+```bash
+curl -X POST http://127.0.0.1:8000/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
+        "model": "voxcpm2",
+        "input": "Streaming PCM test on a single RTX 4090 with VoxCPM2.",
+        "voice": "default",
+        "stream": true,
+        "response_format": "pcm"
+    }' --no-buffer | play -t raw -r 48000 -e signed -b 16 -c 1 -
+```
+
+Observed: 230400 int16 samples = 4.80 s of PCM at 48 kHz, peak 32525 /
+32767 (~99 % full scale). The player sample rate is **48 kHz**, not
+24 kHz.
+
+**T5 — gapless streaming (Gradio AudioWorklet demo)**:
+
+```bash
+python examples/online_serving/text_to_speech/voxcpm2/gradio_demo.py \
+    --api-base http://127.0.0.1:8000
+```
+
+Observed: the demo binds to `http://0.0.0.0:7860` in ~4 s after launch.
+The UI ("VoxCPM2 Streaming Demo") exposes three modes — **Voice Design**
+(control instruction only), **Controllable Cloning** (ref audio +
+optional style control), and **Ultimate Cloning** (ref audio +
+transcript for audio continuation) — plus a `Stream (gapless)`
+checkbox (default on) that drives the AudioWorklet streaming path. The
+right pane includes runnable example prompts and a `Ready` status
+indicator. End-to-end browser playback through `Generate Speech` was
+verified working on the same vLLM-Omni server backing T1–T4 above.
+
+**T6 — offline inference (no server, single process per request)**:
+
+```bash
+python examples/offline_inference/text_to_speech/voxcpm2/end2end.py \
+    --text "Hello, this is VoxCPM2 native AR synthesis running on vLLM-Omni and a single RTX 4090." \
+    --output-dir output_audio
+```
+
+`end2end.py` is a one-shot script (init → one `generate()` → exit), so
+its reported RTF includes per-process compile + CUDA-Graph capture +
+first-request runtime warmup. Note that the default `voxcpm2.yaml` sets
+`enforce_eager: true`, so vLLM's engine-level CUDA graphs are off by
+design; the "CUDA-Graph capture" referenced throughout this recipe is the
+talker's own `_CapturedGraph` over the CFM solver / AudioVAE, not the LM
+decode loop. Flipping `enforce_eager=false` does **not** remove this cost.
+To separate the warmup cost from the
+talker's true per-step throughput, the table below reports **5
+consecutive `engine.generate([prompt])` calls in the same process**
+after a one-off ~28 s engine init:
+
+| call | duration | inference | RTF   | notes                                            |
+| ---- | -------- | --------- | ----- | ------------------------------------------------ |
+| #1   | 6.72 s   | 11.97 s   | 1.782 | cold: torch.compile + CUDA-Graph capture         |
+| #2   | 6.24 s   | 11.43 s   | 1.831 | still runtime warmup                             |
+| #3   | 6.88 s   | 0.82 s    | 0.120 | ⚡ steady-state                                   |
+| #4   | 6.56 s   | 0.78 s    | 0.119 | steady-state                                     |
+| #5   | 5.76 s   | 0.70 s    | 0.121 | steady-state                                     |
+
+**Steady-state RTF (mean of calls #3–#5): 0.120**, i.e. ~8× real-time on
+a single 4090. The high call-#1 / call-#2 numbers are one-time warmup
+costs that **amortise to zero after the second request** — they are not
+the model's per-step throughput.
+
+A literal one-shot `python end2end.py` invocation in a fresh process
+only sees call #1 (~RTF 1.8). For repeated workloads, either reuse the
+same engine across calls (as in the table above) or use the online
+serving path (T1–T5); a long-lived server amortises the warmup across
+all subsequent requests (T1 wall-time RTF was 0.51 with a warm server).
+
+Voice cloning + continuation mode (`ref_audio` + transcript):
+
+```bash
+python examples/offline_inference/text_to_speech/voxcpm2/end2end.py \
+    --text "Hello, this is the offline voice cloning test on a single RTX 4090." \
+    --ref-audio /path/to/reference.wav \
+    --ref-text "Original transcript of the reference audio."
+```
+
+Same 5-call methodology as the zero-shot table, this time with a
+10.08 s reference WAV supplied via `--ref-audio` (and matching
+`--ref-text`):
+
+| call | duration | inference | RTF   | notes                                            |
+| ---- | -------- | --------- | ----- | ------------------------------------------------ |
+| #1   | 5.44 s   | 12.38 s   | 2.276 | cold: compile + CUDA-Graph capture + ref encode  |
+| #2   | 5.12 s   | 2.47 s    | 0.482 | most warmup done                                 |
+| #3   | 5.44 s   | 0.76 s    | 0.139 | ⚡ steady-state                                   |
+| #4   | 4.96 s   | 0.68 s    | 0.137 | steady-state                                     |
+| #5   | 5.28 s   | 0.71 s    | 0.134 | steady-state                                     |
+
+**Steady-state cloning RTF (mean of calls #3–#5): 0.137**, i.e. only
+**+0.017 RTF** above the zero-shot steady-state of 0.120. That delta is
+the per-request `AudioVAE.encode` over the reference clip (~85 ms for a
+10 s ref WAV) plus the ref-continuation prompt build; the AR decode
+steps themselves run at the same per-step rate as zero-shot.
+
+#### Notes
+
+- Prerequisites: `pip install voxcpm soundfile httpx ninja`. `ninja` must be on `PATH` because `flashinfer` JIT-compiles a top-k sampling kernel during the first profile pass (without it, startup dies with `FileNotFoundError: 'ninja'`).
+- Memory usage: ~4.9 GiB model weights + ~15.2 GiB KV cache ≈ 22 GiB / 24 GiB resident with the default `voxcpm2.yaml`. On a shared GPU, pass `--gpu-memory-utilization 0.75` (or lower) if startup fails the free-memory check.
+- Key flags: `--omni` is required. `--trust-remote-code` is not needed — the HF config registers under `model_type=voxcpm2`.
+- Output: 48 kHz mono. When streaming PCM to a player, use `-r 48000`.
+- Cold start: ~60 s server boot + ~25 s first-inference overhead (torch.compile + flashinfer JIT + the talker's own `_CapturedGraph` capture over CFM / AudioVAE — not vLLM engine CUDA graphs, which `enforce_eager: true` disables). Steady-state inference RTF is **~0.12** (~8× real-time on a single 4090; see T6 table for the per-call breakdown); online server requests run at ~RTF 0.5 wall-time (includes HTTP round-trip). Keep the server warm for interactive use.

recipes/OpenMOSS/MOSS-TTS.md

@@ -0,0 +1,145 @@
+# MOSS-TTS
+
+## Summary
+
+- Vendor: OpenMOSS
+- Models: `OpenMOSS-Team/MOSS-TTS` (8B), `OpenMOSS-Team/MOSS-TTS-v1.5` (8B),
+  `OpenMOSS-Team/MOSS-TTS-Realtime` (1.7B), `OpenMOSS-Team/MOSS-TTSD-v1.0` (8B),
+  `OpenMOSS-Team/MOSS-SoundEffect` (8B), `OpenMOSS-Team/MOSS-VoiceGenerator` (1.7B)
+- Task: Text-to-speech synthesis, sound effect generation, zero-shot voice design
+- Mode: Online serving via the OpenAI-compatible `/v1/audio/speech` API; offline inference
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe for 24 kHz multilingual TTS with voice cloning (20 languages including
+Chinese and English). Choose a variant based on your latency and quality requirements:
+
+| Model | Params | Use case |
+|---|---|---|
+| MOSS-TTS | 8B | General TTS, highest quality |
+| MOSS-TTS-v1.5 | 8B | General TTS upgrade of 1.0: 31 languages, steadier cloning, `[pause Xs]` markers (set `language` for best results); same `MossTTSDelay` API |
+| MOSS-TTS-Realtime | 1.7B | Lowest latency (TTFB ~180 ms), streaming-first |
+| MOSS-TTSD-v1.0 | 8B | Multi-turn dialogue TTS |
+| MOSS-SoundEffect | 8B | Sound effect synthesis from text description |
+| MOSS-VoiceGenerator | 1.7B | Zero-shot voice design |
+
+All variants share the same codec (`OpenMOSS-Team/MOSS-Audio-Tokenizer`, ~7 GB) and
+output 24 kHz mono audio.
+
+## References
+
+- Offline inference example: [`examples/offline_inference/text_to_speech/moss_tts/`](../../examples/offline_inference/text_to_speech/moss_tts/)
+- Deploy configs: [`vllm_omni/deploy/moss_tts.yaml`](../../vllm_omni/deploy/moss_tts.yaml) and variants
+- HuggingFace org: <https://huggingface.co/OpenMOSS-Team>
+
+## Hardware Support
+
+### GPU
+
+#### 1x H100 80GB — MOSS-TTS (8B)
+
+##### Environment
+
+- OS: Linux
+- Python: 3.11+
+- CUDA 12.8
+- vLLM-Omni version: see `vllm_omni/__version__.py`
+
+##### Command
+
+```bash
+# The codec is loaded automatically from OpenMOSS-Team/MOSS-Audio-Tokenizer.
+# Override the path with MOSS_TTS_CODEC_PATH if you have a local copy.
+vllm serve OpenMOSS-Team/MOSS-TTS --omni --port 8091
+```
+
+##### Verification
+
+Voice cloning (provide a reference audio clip):
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
+        "model": "OpenMOSS-Team/MOSS-TTS",
+        "input": "Hello, this is a voice cloning test.",
+        "voice": "default",
+        "ref_audio": "https://raw.githubusercontent.com/OpenMOSS/MOSS-TTS/main/assets/audio/zh_1.wav",
+        "response_format": "wav"
+    }' --output output.wav
+```
+
+##### Notes
+
+- Peak GPU memory: ~18 GB for the talker (8B) + ~8 GB for the codec decoder on the same device.
+  Use `gpu_memory_utilization: 0.85` in `moss_tts.yaml` (default).
+- Output: 24 kHz mono WAV.
+- The `MOSS_TTS_CODEC_PATH` environment variable overrides the codec checkpoint location.
+
+---
+
+#### 1x A10G 24GB — MOSS-TTS-Realtime (1.7B)
+
+##### Environment
+
+- OS: Linux
+- Python: 3.11+
+- CUDA 12.8
+
+##### Command
+
+```bash
+vllm serve OpenMOSS-Team/MOSS-TTS-Realtime --omni --port 8091
+```
+
+##### Verification
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
+        "model": "OpenMOSS-Team/MOSS-TTS-Realtime",
+        "input": "This is a low-latency streaming TTS test.",
+        "voice": "default",
+        "ref_audio": "https://raw.githubusercontent.com/OpenMOSS/MOSS-TTS/main/assets/audio/zh_1.wav",
+        "response_format": "wav",
+        "stream": true
+    }' --output output.wav
+```
+
+##### Notes
+
+- Peak GPU memory: ~6 GB for the talker (1.7B) + ~8 GB for the codec decoder.
+- First-audio latency (TTFB): ~180 ms on A10G.
+- `codec_chunk_frames: 15` in `moss_tts_realtime.yaml` for lower TTFA than the 8B variant.
+
+---
+
+#### 1x A10G 24GB — MOSS-SoundEffect (8B, sound effect synthesis)
+
+##### Command
+
+```bash
+vllm serve OpenMOSS-Team/MOSS-SoundEffect --omni --port 8091
+```
+
+##### Verification
+
+Sound effect synthesis takes a text description instead of reference audio:
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
+        "model": "OpenMOSS-Team/MOSS-SoundEffect",
+        "input": "Thunder rumbling, rain pattering on a tin roof.",
+        "response_format": "wav"
+    }' --output thunder.wav
+```
+
+##### Notes
+
+- No `ref_audio` required or accepted for MOSS-SoundEffect.
+- Input field maps to the `ambient_sound` parameter in the upstream processor.
+- Rate: ~12.5 tokens per second; longer descriptions produce longer audio.

recipes/README.md

@@ -29,12 +29,17 @@ recipes/
 |--------|------|----------|
 | [`audiox/AudioX.md`](./audiox/AudioX.md) | Offline + online unified text/video→audio diffusion | 1x L4 24GB |
 | [`Baidu/ERNIE-Image.md`](./Baidu/ERNIE-Image.md) | Text-to-image online serving (ERNIE-Image 8B) | 1x or 2x RTX 4090 24GB |
+| [`ByteDance/Lance.md`](./ByteDance/Lance.md) | Unified AR+diffusion: text/img/video gen + understanding (Lance 3B) | 1x B300 / A100 80GB |
 | [`fishaudio/Fish-Speech-S2-Pro.md`](./fishaudio/Fish-Speech-S2-Pro.md) | Online serving for TTS | 1x A800 80GB |
 | [`Helios/Helios.md`](./Helios/Helios.md) | Text-to-video, image-to-video, and video-to-video generation | 1x NVIDIA H20 |
 | [`inclusionAI/Ming-flash-omni-2.0.md`](./inclusionAI/Ming-flash-omni-2.0.md) | Online serving for multimodal chat + standalone TTS | 4x H100 / 1x H100 80GB |
 | [`LTX/LTX-2.md`](./LTX/LTX-2.md) | Text-to-video and image-to-video serving | 1x H200 141GB |
 | [`LTX/LTX-2.3.md`](./LTX/LTX-2.3.md) | Text-to-video with audio generation (22B) | 1x GPU (96GB VRAM) |
 | [`mistralai/Voxtral-TTS.md`](./mistralai/Voxtral-TTS.md) | Online serving for TTS | 1x RTX 4090 24GB |
+| [`nvidia/Cosmos3-Nano.md`](./nvidia/Cosmos3-Nano.md) | Text-to-image, text-to-video, image-to-video generation, text to video with sound  | 1x H200 141GB / B300 |
+| [`nvidia/Cosmos3-Super.md`](./nvidia/Cosmos3-Super.md) | 64B T2I / T2V / I2V generation (+ optional audio) | 8x H200/H100/A100 / 2x H200 / B300 |
+| [`OpenBMB/MiniCPM-o-4_5.md`](./OpenBMB/MiniCPM-o-4_5.md) | Online serving for omni multimodal chat (text / image / audio / video → text + 24 kHz speech) | 2x A100/H100 80GB / 3x mid-tier GPU / 8x RTX 4090 24GB |
+| [`OpenBMB/VoxCPM2.md`](./OpenBMB/VoxCPM2.md) | Online + offline TTS with native AR pipeline (48 kHz, 30+ languages) | 1x RTX 4090 24GB |
 | [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt mixed FP8/NVFP4 | 1x A100 80GB / 2x B200 |
 | [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay | 1x A100 80GB |
 | [`Qwen/Qwen-Image-Edit.md`](./Qwen/Qwen-Image-Edit.md) | Text-guided single-image editing | 1x or 2x H200 141GB |

recipes/nvidia/Cosmos3-Nano.md

@@ -0,0 +1,245 @@
+# Cosmos3-Nano
+
+> Text-to-image, text-to-video, and image-to-video serving
+
+## Summary
+
+- Vendor: NVIDIA
+- Model: `nvidia/Cosmos3-Nano`
+- Task: Text-to-image (T2I), text-to-video (T2V), and image-to-video (I2V) generation, with optional synchronized audio (video + sound)
+- Mode: Online serving with the OpenAI-compatible image/video APIs, plus offline generation via the `Omni` API
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe to deploy `nvidia/Cosmos3-Nano` for image and video generation.
+A single pipeline class (`Cosmos3OmniDiffusersPipeline`) serves all three modes;
+the mode is selected per request:
+
+- **T2I** — `POST /v1/images/generations` (or a prompt carrying `modalities=["image"]`).
+- **T2V** — `POST /v1/videos/sync` with `num_frames > 1` and no reference image.
+- **I2V** — `POST /v1/videos/sync` with a reference image (`input_reference` file
+  upload, or `image_reference` JSON).
+- **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.
+
+## References
+
+- Model card (authoritative usage + example assets): <https://huggingface.co/nvidia/Cosmos3-Nano>
+- Example inputs/outputs live in the repo's `assets/` (`example_t2v_prompt.json`,
+  `example_i2v_prompt.json`, `example_i2v_input.jpg`, `negative_prompt.json`;
+  audio examples: `example_t2vs_prompt.json`, `example_t2vs_output.mp4`,
+  `example_i2vs_output.mp4`).
+- Prompt upsampling (recommended for quality): the model expects JSON-upsampled
+  structured prompts; see NVIDIA's `cosmos-framework` prompt-upsampling docs.
+- Pipeline: [`vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py`](../../vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py)
+- Smoke tests (canonical request formats): [`tests/e2e/accuracy/test_cosmos3_similarity.py`](../../tests/e2e/accuracy/test_cosmos3_similarity.py)
+
+## Hardware Support
+
+## GPU
+
+### 1x H200 141GB / B300 (Online serving)
+
+#### Environment
+
+- OS: Ubuntu 22.04+
+- Python: 3.12+
+- Driver / runtime: NVIDIA CUDA environment
+- vLLM version: match the repository requirements from your current checkout
+- vLLM-Omni version or commit: use the commit you are deploying from
+
+#### Command
+
+Requires the `vllm-omni` package (or the `vllm/vllm-omni:cosmos3` container),
+which provides the `vllm serve … --omni` entrypoint used below.
+
+Safety guardrails are **on by default** (NVIDIA Open Model License). They load
+the **gated** `nvidia/Cosmos-1.0-Guardrail` model, so to keep them on you must:
+
+1. `pip install cosmos-guardrail`
+2. Accept the license at <https://huggingface.co/nvidia/Cosmos-1.0-Guardrail>
+3. Export a token with access: `export HF_TOKEN=hf_...`
+
+Then launch the recommended server:
+
+```bash
+vllm serve nvidia/Cosmos3-Nano \
+  --omni \
+  --host 0.0.0.0 --port 8000 \
+  --init-timeout 1800
+```
+
+To run **without** guardrails (you are responsible for license compliance),
+add `--no-guardrails` (no token/`cosmos-guardrail` needed). For extra GPUs use
+`--ulysses-degree N` (context parallel) or `--tensor-parallel-size N`;
+`--enable-layerwise-offload` reduces VRAM on smaller GPUs. The pipeline
+auto-resolves from `model_index.json`; pass
+`--model-class-name Cosmos3OmniDiffusersPipeline` to force it explicitly.
+
+#### Verification
+
+Best quality uses the JSON-upsampled prompts from `assets/` (download with
+`hf download nvidia/Cosmos3-Nano assets/ --local-dir Cosmos3-Nano`). Minimal
+self-contained examples:
+
+```bash
+curl http://localhost:8000/v1/models
+
+# Text-to-image -> /v1/images/generations  (1024x1024, 50 steps; base64 PNG)
+curl -sS -X POST http://localhost:8000/v1/images/generations \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "nvidia/Cosmos3-Nano",
+    "prompt": "A photorealistic red sports car on a city street at golden hour, cinematic lighting.",
+    "negative_prompt": "blurry, distorted, low quality",
+    "size": "1024x1024", "n": 1, "response_format": "b64_json",
+    "num_inference_steps": 50, "guidance_scale": 7.0, "seed": 42
+  }' | python -c "import sys,json,base64; open('cosmos3_t2i.png','wb').write(base64.b64decode(json.load(sys.stdin)['data'][0]['b64_json']))"
+
+# Text-to-video -> /v1/videos/sync  (720p, 189 frames @ 24fps; official params)
+curl -sS -X POST http://localhost:8000/v1/videos/sync \
+  -H "Accept: video/mp4" \
+  -F "model=nvidia/Cosmos3-Nano" \
+  -F "prompt=A robot arm is cleaning a plate in the kitchen" \
+  -F "negative_prompt=blurry, distorted, low quality, jittery, deformed" \
+  -F "size=1280x720" -F "num_frames=189" -F "fps=24" \
+  -F "num_inference_steps=35" -F "guidance_scale=6.0" \
+  -F "max_sequence_length=4096" -F "flow_shift=10.0" \
+  -F 'extra_params={"use_resolution_template":false,"use_duration_template":false,"guardrails":true}' \
+  -F "seed=123" \
+  -o cosmos3_t2v.mp4
+
+# Image-to-video -> /v1/videos/sync with an uploaded reference image
+curl -sS -X POST http://localhost:8000/v1/videos/sync \
+  -H "Accept: video/mp4" \
+  -F "model=nvidia/Cosmos3-Nano" \
+  -F "prompt=The scene comes to life with smooth, natural motion." \
+  -F "negative_prompt=blurry, distorted, low quality" \
+  -F "size=1280x720" -F "num_frames=189" -F "fps=24" \
+  -F "num_inference_steps=35" -F "guidance_scale=6.0" \
+  -F "max_sequence_length=4096" -F "flow_shift=10.0" \
+  -F 'extra_params={"use_resolution_template":false,"use_duration_template":false,"guardrails":true}' \
+  -F "seed=1111" \
+  -F "input_reference=@/path/to/reference.jpg;type=image/jpeg" \
+  -o cosmos3_i2v.mp4
+
+
+# Text-to-video-with-sound
+curl -sS -X POST http://localhost:8000/v1/videos/sync \
+  -H "Accept: video/mp4" \
+  -F "model=nvidia/Cosmos3-Nano" \
+  -F "prompt=The video opens with a view of a well-lit indoor fruit display. A robotic arm picks up a pear, an orange, and a carambola one by one, placing each into a plastic bag in a shopping cart with red handles. The video is 7.875 seconds long, 24 FPS, and 1280x720. Audio description: soft servo whirs, gentle fruit thuds, plastic bag rustling, and a faint refrigeration hum." \
+  -F "negative_prompt=blurry, distorted, low quality" \
+  -F "size=1280x720" \
+  -F "num_frames=189" \
+  -F "fps=24" \
+  -F "num_inference_steps=35" \
+  -F "guidance_scale=6.0" \
+  -F "max_sequence_length=4096" \
+  -F "flow_shift=10.0" \
+  -F "seed=0" \
+  -F "generate_sound=true" \
+  -F "sound_duration=7.875" \
+  -F 'extra_params={"use_resolution_template":false,"use_duration_template":false,"guardrails":true}' \
+  -o cosmos3_t2v_with_sound.mp4
+```
+
+#### Notes
+
+- **Measured latency (1x B300, bf16, guardrails off):**
+  - T2I 1024² — 10 / 25 / 50 steps → ~0.4 / 0.7 / **1.3 s**
+  - T2V 1280×720 @ 35 steps — 25 / 49 / 93 / **189** frames → ~7 / 15 / 33 / **~93 s**
+  - I2V 1280×720, 189 frames @ 35 steps → ~**99 s**
+  - Guardrails-on overhead: ~8% on T2I, negligible on video.
+- **Memory:** transformer ~17 GiB (bf16); peak ~46 GiB for 720p video on 1 GPU;
+  full repo (transformer + Wan VAE + Qwen3-VL vision encoder + audio tokenizer)
+  ~33 GB on disk.
+- **Determinism:** identical seed reproduces identical output on the same
+  hardware; outputs are not bit-identical across different GPU types.
+- **Supported sizes (per model card):** 256p / 480p / 720p at 16:9, 4:3, 1:1,
+  3:4, 9:16. Defaults: T2I 1024², 50 steps, guidance 7.0; T2V/I2V 1280×720,
+  189 frames, 35 steps, guidance 6.0, `flow_shift=10.0`.
+- **Key flags / params:** `--no-guardrails` (server) or
+  `extra_params={"guardrails":false}` (per request) toggles safety. The
+  per-request flag only takes effect when the server was launched **with**
+  guardrails enabled (it cannot re-enable them on a `--no-guardrails` server).
+  `use_resolution_template` / `use_duration_template` are off by default and only
+  needed when not using upsampled prompts that already encode resolution/duration.
+- **Known limitations:**
+  - Guardrails-on requires `cosmos-guardrail` **and** access to the gated
+    `nvidia/Cosmos-1.0-Guardrail` repo (accept license + `HF_TOKEN`); otherwise
+    the server fails at pipeline build with a gated-repo / safety-checker error.
+  - A guardrail-blocked prompt currently returns HTTP 500
+    (`"Guardrail blocked prompt"`).
+  - Action (policy / forward- / inverse-dynamics) modalities are not part of
+    this integration yet.
+
+### 1x GPU (Offline generation)
+
+#### Environment
+
+- OS: Ubuntu 22.04+
+- Python: 3.12+
+- Driver / runtime: NVIDIA CUDA environment
+- vLLM-Omni version or commit: use the commit you are deploying from
+
+#### Command
+
+```python
+# cosmos3_offline.py  —  run with:  python cosmos3_offline.py
+import torch
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+
+def main():
+    omni = Omni(
+        model="nvidia/Cosmos3-Nano",
+        model_class_name="Cosmos3OmniDiffusersPipeline",
+        trust_remote_code=True,
+        enforce_eager=True,
+        # Guardrails are disabled here for a quick local run; install
+        # cosmos-guardrail + gated-repo access and drop this to enable them.
+        model_config={"guardrails": False},
+    )
+    gen = torch.Generator(device="cpu").manual_seed(42)
+
+    # Text-to-image (modalities=["image"]). For T2V use modalities=["video"]
+    # plus num_frames/fps; for I2V add multi_modal_data={"image": <PIL.Image>}.
+    outputs = omni.generate(
+        {
+            "prompt": "A photorealistic red sports car at golden hour, cinematic lighting.",
+            "negative_prompt": "blurry, distorted, low quality",
+            "modalities": ["image"],
+        },
+        OmniDiffusionSamplingParams(
+            height=1024, width=1024, generator=gen,
+            guidance_scale=7.0, num_inference_steps=50, num_outputs_per_prompt=1,
+        ),
+    )
+    outputs[0].request_output.images[0].save("cosmos3_t2i.png")
+    omni.close()
+
+
+if __name__ == "__main__":
+    main()
+```
+
+#### Verification
+
+```bash
+python cosmos3_offline.py
+python -c "from PIL import Image; im=Image.open('cosmos3_t2i.png'); print('image', im.size, im.mode)"
+```
+
+#### Notes
+
+- Same `Cosmos3OmniDiffusersPipeline` as online; mode is chosen by
+  `prompt["modalities"]` (`["image"]` → T2I, `["video"]` → T2V) plus
+  `num_frames`/`fps`, and `multi_modal_data={"image": ...}` for I2V. For video,
+  frames are returned in `outputs[0].request_output.images` as an
+  `(B, F, H, W, 3)` array.
+- The offline entry must be guarded by `if __name__ == "__main__":` — the engine
+  spawns workers with the `spawn` start method.

recipes/nvidia/Cosmos3-Super.md

@@ -0,0 +1,108 @@
+# Cosmos3-Super
+
+> Frontier 64B world model: text-to-image, text-to-video, image-to-video (+ optional audio)
+
+## Summary
+
+- Vendor: NVIDIA
+- Model: `nvidia/Cosmos3-Super` (64B; also `Cosmos3-Super-Text2Image`, `Cosmos3-Super-Image2Video`)
+- Task: T2I, T2V, I2V generation, with optional synchronized audio (video + sound)
+- Mode: Online serving with the OpenAI-compatible image/video APIs
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe to deploy the 64B `nvidia/Cosmos3-Super` for the highest-quality
+Cosmos3 generation. It shares the same `Cosmos3OmniDiffusersPipeline` and request
+formats as [Cosmos3-Nano](./Cosmos3-Nano.md) — only the checkpoint size and the
+recommended parallelism differ. Mode is selected per request (T2I →
+`/v1/images/generations`; T2V/I2V → `/v1/videos/sync`; add `generate_sound=true`
+for audio).
+
+## References
+
+- Model card (authoritative usage + example assets): <https://huggingface.co/nvidia/Cosmos3-Super>
+- Nano recipe (same APIs/params): [`Cosmos3-Nano.md`](./Cosmos3-Nano.md)
+- Pipeline: [`vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py`](../../vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py)
+
+## Hardware Support
+
+## GPU
+
+Requires the `vllm-omni` package (or the `vllm/vllm-omni:cosmos3` container),
+which provides the `vllm serve … --omni` entrypoint used below.
+
+### 8x H200/H100/A100 (recommended, per model card)
+
+```bash
+vllm serve nvidia/Cosmos3-Super \
+  --omni \
+  --host 0.0.0.0 --port 8000 \
+  --cfg-parallel-size 2 \
+  --ulysses-degree 4 \
+  --use-hsdp --hsdp-shard-size 8 \
+  --init-timeout 1800
+```
+
+### 2x H200 / B300 (minimum)
+
+```bash
+vllm serve nvidia/Cosmos3-Super \
+  --omni \
+  --host 0.0.0.0 --port 8000 \
+  --cfg-parallel-size 2 \
+  --use-hsdp --hsdp-shard-size 2 \
+  --init-timeout 1800
+```
+
+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.
+
+#### Verification
+
+Requests are identical to Nano (see [`Cosmos3-Nano.md`](./Cosmos3-Nano.md) for full
+T2I/T2V/I2V/T2VS curls); official params: `size=1280x720, num_frames=189, fps=24,
+num_inference_steps=35, guidance_scale=6.0, flow_shift=10.0, max_sequence_length=4096`.
+
+```bash
+curl http://localhost:8000/v1/models
+# T2V (official prompt assets give best quality)
+curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \
+  -F "model=nvidia/Cosmos3-Super" -F "prompt=A robot arm is cleaning a plate in the kitchen" \
+  -F "size=1280x720" -F "num_frames=189" -F "fps=24" -F "num_inference_steps=35" \
+  -F "guidance_scale=6.0" -F "max_sequence_length=4096" -F "flow_shift=10.0" \
+  -F 'extra_params={"use_resolution_template":false,"use_duration_template":false,"guardrails":true}' \
+  -F "seed=17" -o cosmos3_super_t2v.mp4
+
+# I2V — add an uploaded reference image
+curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \
+  -F "model=nvidia/Cosmos3-Super" -F "prompt=The scene comes to life with smooth, natural motion." \
+  -F "size=1280x720" -F "num_frames=189" -F "fps=24" -F "num_inference_steps=35" \
+  -F "guidance_scale=6.0" -F "max_sequence_length=4096" -F "flow_shift=10.0" \
+  -F 'extra_params={"use_resolution_template":false,"use_duration_template":false,"guardrails":true}' \
+  -F "seed=1111" -F "input_reference=@/path/to/reference.jpg;type=image/jpeg" \
+  -o cosmos3_super_i2v.mp4
+
+# T2V + sound — add generate_sound/sound_duration (output muxes AAC 48 kHz stereo)
+curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \
+  -F "model=nvidia/Cosmos3-Super" -F "prompt=A robot arm is cleaning a plate in the kitchen" \
+  -F "size=1280x720" -F "num_frames=189" -F "fps=24" -F "num_inference_steps=35" \
+  -F "guidance_scale=6.0" -F "max_sequence_length=4096" -F "flow_shift=10.0" \
+  -F "generate_sound=true" -F "sound_duration=7.875" \
+  -F 'extra_params={"use_resolution_template":false,"use_duration_template":false,"guardrails":true}' \
+  -F "seed=17" -o cosmos3_super_t2vs.mp4
+```
+
+#### Notes
+
+- **Measured (2x B300, bf16, guardrails off, official 2-GPU config above):**
+  - T2I 1024², 50 steps → **~6 s**
+  - T2V 1280×720, 189 frames, 35 steps → **~197 s**
+  - I2V 1280×720, 189 frames, 35 steps → **~200 s**
+  - T2V + sound (189 frames, 35 steps) → **~198 s**, output muxes **AAC 48 kHz stereo**
+  - (NVIDIA's reference: 8×H200 @ 50 steps ≈ 55 s/video; 2×H200 @ 35 steps ≈ 3 min/video.)
+- **Memory:** ~61.5 GiB per GPU when sharded across 2 GPUs (HSDP shard 2); repo ~135 GB on disk.
+- Same generation defaults, supported sizes, and `generate_sound`/`sound_duration`
+  semantics as Nano. Action (policy / forward- / inverse-dynamics) modalities are
+  not part of this integration yet.

tests/comfyui/test_comfyui_integration.py

@@ -30,12 +30,12 @@ from PIL import Image
 from pytest_mock import MockerFixture
 from vllm import SamplingParams
 from vllm.outputs import CompletionOutput, RequestOutput
-from vllm.utils.argparse_utils import FlexibleArgumentParser
 
 from vllm_omni.entrypoints.async_omni import AsyncOmni as RealAsyncOmni
 from vllm_omni.entrypoints.cli.serve import OmniServeCommand
 from vllm_omni.inputs.data import OmniSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
+from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -444,7 +444,7 @@ def mock_async_omni(
 @pytest.fixture
 def api_server(unused_tcp_port_factory, server_case: ServerCase, mock_async_omni):
     """Set up a API server in background process from command line with parametrized model name and mocked AsyncOmni."""
-    parser = FlexibleArgumentParser()
+    parser = TrackingArgumentParser()
     subparsers = parser.add_subparsers(dest="command")
     cmd = OmniServeCommand()
     cmd.subparser_init(subparsers)

tests/core/test_prefix_cache.py

@@ -157,6 +157,75 @@ def test_update_no_multimodal():
         assert torch.all(slot_states == new_states)
 
 
+def test_update_uses_precomputed_hidden_states_cpu():
+    """Precomputed CPU staging should be the only hidden-state cache source."""
+    cache = get_omni_pcache()
+
+    num_tokens_unpadded = 8
+    slot_offset = 8
+    slot_mapping = torch.arange(slot_offset, slot_offset + num_tokens_unpadded)
+    hidden_states = torch.zeros((num_tokens_unpadded, HIDDEN_SIZE), dtype=DTYPE)
+    hidden_states_cpu = torch.rand((num_tokens_unpadded, HIDDEN_SIZE), dtype=DTYPE)
+
+    cache.update_omni_tensor_prefix_cache(
+        hidden_states=hidden_states,
+        hidden_states_cpu=hidden_states_cpu,
+        multimodal_outputs=None,
+        num_tokens_unpadded=num_tokens_unpadded,
+        slot_mapping=slot_mapping,
+    )
+
+    hs_rows = cache.hidden_states_cache.view(NUM_BLOCKS * BLOCK_SIZE, HIDDEN_SIZE)
+    for slot_idx, new_states in zip(slot_mapping, hidden_states_cpu):
+        slot_states = hs_rows[slot_idx]
+        assert torch.all(slot_states == new_states)
+
+
+@pytest.mark.parametrize(
+    "hidden_states_cpu",
+    [
+        torch.rand((4, HIDDEN_SIZE), dtype=DTYPE),
+        torch.rand((8, HIDDEN_SIZE), dtype=OTHER_DTYPE),
+        torch.rand((8, HIDDEN_SIZE + 1), dtype=DTYPE),
+        torch.rand((HIDDEN_SIZE, 8), dtype=DTYPE).t(),
+    ],
+)
+def test_update_rejects_invalid_hidden_states_cpu(hidden_states_cpu):
+    cache = get_omni_pcache()
+
+    num_tokens_unpadded = 8
+    slot_mapping = torch.arange(num_tokens_unpadded)
+    hidden_states = torch.rand((num_tokens_unpadded, HIDDEN_SIZE), dtype=DTYPE)
+
+    with pytest.raises(RuntimeError):
+        cache.update_omni_tensor_prefix_cache(
+            hidden_states=hidden_states,
+            hidden_states_cpu=hidden_states_cpu,
+            multimodal_outputs=None,
+            num_tokens_unpadded=num_tokens_unpadded,
+            slot_mapping=slot_mapping,
+        )
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+def test_update_rejects_gpu_hidden_states_cpu():
+    cache = get_omni_pcache()
+
+    num_tokens_unpadded = 8
+    slot_mapping = torch.arange(num_tokens_unpadded)
+    hidden_states = torch.rand((num_tokens_unpadded, HIDDEN_SIZE), dtype=DTYPE)
+    hidden_states_cpu = torch.rand((num_tokens_unpadded, HIDDEN_SIZE), dtype=DTYPE, device="cuda")
+
+    with pytest.raises(RuntimeError):
+        cache.update_omni_tensor_prefix_cache(
+            hidden_states=hidden_states,
+            hidden_states_cpu=hidden_states_cpu,
+            multimodal_outputs=None,
+            num_tokens_unpadded=num_tokens_unpadded,
+            slot_mapping=slot_mapping,
+        )
+
+
 @pytest.mark.parametrize(
     "feat_dims",
     [
@@ -359,6 +428,76 @@ def test_get_merged_hidden_states(num_tokens_padded, mocker):
     assert torch.all(req2_merged_states == req2_new_states)
 
 
+def test_get_merged_hidden_states_uses_precomputed_hidden_states_cpu(mocker):
+    cache = get_omni_pcache()
+
+    orig_num_tokens_unpadded = 8
+    slot_offset = 8
+    orig_slot_mapping = torch.arange(slot_offset, slot_offset + orig_num_tokens_unpadded)
+    orig_hidden_states = torch.rand((orig_num_tokens_unpadded, HIDDEN_SIZE), dtype=DTYPE)
+
+    cache.update_omni_tensor_prefix_cache(
+        hidden_states=orig_hidden_states,
+        multimodal_outputs=None,
+        num_tokens_unpadded=orig_num_tokens_unpadded,
+        slot_mapping=orig_slot_mapping,
+    )
+
+    num_new_toks_req1 = 3
+    num_new_toks_req2 = 2
+    cache.add_prefix_cached_new_req_id("req1")
+    num_scheduled_tokens = {
+        "req1": num_new_toks_req1,
+        "req2": num_new_toks_req2,
+    }
+    new_hidden_states = torch.zeros(
+        (num_new_toks_req1 + num_new_toks_req2, HIDDEN_SIZE),
+        dtype=DTYPE,
+    )
+    hidden_states_cpu = torch.rand_like(new_hidden_states)
+    req1_new_states = hidden_states_cpu[:num_new_toks_req1]
+    req2_new_states = hidden_states_cpu[-num_new_toks_req2:]
+    input_batch = MockInputBatch(num_computed_tokens_cpu=torch.Tensor([orig_num_tokens_unpadded, 0]))
+
+    mocker.patch(
+        "vllm_omni.core.prefix_cache.OmniTensorPrefixCache._get_cached_block_ids",
+        new=fake_get_cached_block_ids,
+    )
+    merged_states = cache.get_merged_hidden_states(
+        query_start_loc=[0, num_new_toks_req1],
+        input_batch=input_batch,
+        hidden_states=new_hidden_states,
+        hidden_states_cpu=hidden_states_cpu,
+        num_scheduled_tokens=num_scheduled_tokens,
+    )
+
+    assert torch.all(merged_states["req1"][:orig_num_tokens_unpadded] == orig_hidden_states)
+    assert torch.all(merged_states["req1"][-num_new_toks_req1:] == req1_new_states)
+    assert torch.all(merged_states["req2"] == req2_new_states)
+
+
+def test_get_merged_hidden_states_rejects_short_hidden_states_cpu(mocker):
+    cache = get_omni_pcache()
+    num_scheduled_tokens = {
+        "req1": 3,
+        "req2": 2,
+    }
+    input_batch = MockInputBatch(num_computed_tokens_cpu=torch.Tensor([0, 0]))
+
+    mocker.patch(
+        "vllm_omni.core.prefix_cache.OmniTensorPrefixCache._get_cached_block_ids",
+        new=fake_get_cached_block_ids,
+    )
+    with pytest.raises(RuntimeError):
+        cache.get_merged_hidden_states(
+            query_start_loc=[0, 3],
+            input_batch=input_batch,
+            hidden_states=torch.rand((5, HIDDEN_SIZE), dtype=DTYPE),
+            hidden_states_cpu=torch.rand((4, HIDDEN_SIZE), dtype=DTYPE),
+            num_scheduled_tokens=num_scheduled_tokens,
+        )
+
+
 @pytest.mark.parametrize("num_tokens_padded", [None, 16])
 @pytest.mark.parametrize(
     "feat_dims",

tests/diffusion/cache/test_cache_dit.py

@@ -25,6 +25,7 @@ SEPARATE_CFG_ENABLERS = [
     cd_backend.enable_cache_for_helios,
     cd_backend.enable_cache_for_wan22,
     cd_backend.enable_cache_for_longcat_image,
+    cd_backend.enable_cache_for_cosmos3,
 ]
 
 SAMPLE_CACHE_CONFIG = DiffusionCacheConfig()
@@ -47,6 +48,24 @@ def test_separate_cfg(mock_cache_dit, mock_block_adapter, enabler):
     assert adapter_kwargs["has_separate_cfg"] is True
 
 
+@patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter")
+@patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit")
+def test_cosmos3_cache_dit_wraps_gen_layers(mock_cache_dit, mock_block_adapter):
+    """Cosmos3 should cache only the repeated GEN pathway blocks."""
+    mock_pipeline = Mock()
+    gen_layers = object()
+    mock_pipeline.transformer.gen_layers = gen_layers
+
+    cd_backend.enable_cache_for_cosmos3(mock_pipeline, SAMPLE_CACHE_CONFIG)
+
+    mock_cache_dit.enable_cache.assert_called_once()
+    adapter_kwargs = mock_block_adapter.call_args.kwargs
+    assert adapter_kwargs["transformer"] is mock_pipeline.transformer
+    assert adapter_kwargs["blocks"] == [gen_layers]
+    assert adapter_kwargs["has_separate_cfg"] is True
+    assert adapter_kwargs["check_forward_pattern"] is False
+
+
 # This test is skipped on ROCm since rocm_unquantized_gemm doesn't support CPU backend
 @pytest.mark.skipif(
     current_omni_platform.is_rocm(),

tests/diffusion/models/bagel/test_trajectory_recording.py

@@ -19,8 +19,9 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 NUM_TOKENS = 8
 HIDDEN_DIM = 16
 NUM_TIMESTEPS = 5
-# generate_image uses timesteps[:-1], so actual steps = NUM_TIMESTEPS - 1
-EXPECTED_STEPS = NUM_TIMESTEPS - 1
+# generate_image samples (num_timesteps + 1) points then drops the last, matching
+# upstream lance: see Bagel.generate_image in bagel_transformer.py.
+EXPECTED_STEPS = NUM_TIMESTEPS
 
 
 def _make_mock_bagel(mocker: MockerFixture):
@@ -178,7 +179,7 @@ class TestTrajectoryRecording:
         bagel, args = bagel_and_args
 
         # Recompute the expected timestep schedule (mirrors generate_image logic)
-        ts = torch.linspace(1, 0, args["num_timesteps"])
+        ts = torch.linspace(1, 0, args["num_timesteps"] + 1)
         shift = args.get("timestep_shift", 1.0)
         ts = shift * ts / (1 + (shift - 1) * ts)
         expected_timesteps = ts[:-1]  # last element is dropped
@@ -196,7 +197,7 @@ class TestTrajectoryRecording:
         bagel, args = bagel_and_args
 
         # Compute the number of denoising steps the same way generate_image does
-        ts = torch.linspace(1, 0, args["num_timesteps"])
+        ts = torch.linspace(1, 0, args["num_timesteps"] + 1)
         shift = args.get("timestep_shift", 1.0)
         ts = shift * ts / (1 + (shift - 1) * ts)
         num_steps = len(ts) - 1  # timesteps = ts[:-1]

tests/diffusion/models/cosmos3/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py

@@ -0,0 +1,677 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import sys
+import types
+from types import SimpleNamespace
+from typing import Any
+
+import pytest
+import torch
+from PIL import Image
+from torch import nn
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+class StubScheduler:
+    def __init__(self, timesteps: list[int] | None = None, *, flow_shift: float = 1.0) -> None:
+        self.timesteps = torch.tensor(timesteps or [9, 3], dtype=torch.int64)
+        self.config = SimpleNamespace(num_train_timesteps=1000, flow_shift=flow_shift)
+        self.set_timesteps_calls: list[tuple[int, torch.device]] = []
+        self.step_calls: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = []
+
+    def set_timesteps(self, num_steps: int, device: torch.device) -> None:
+        self.set_timesteps_calls.append((num_steps, device))
+        self.timesteps = torch.arange(num_steps, 0, -1, dtype=torch.int64, device=device)
+
+    def step(self, noise_pred: torch.Tensor, timestep: torch.Tensor, latents: torch.Tensor, **kwargs):
+        del kwargs
+        self.step_calls.append((noise_pred.clone(), timestep.clone(), latents.clone()))
+        return (latents + noise_pred,)
+
+
+class _ModeLatentDist:
+    def __init__(self, latents: torch.Tensor) -> None:
+        self._latents = latents
+
+    def mode(self) -> torch.Tensor:
+        return self._latents
+
+
+class StubCosmos3VAE:
+    dtype = torch.float32
+
+    def __init__(self, z_dim: int = 2, *, temporal: int = 4, spatial: int = 8) -> None:
+        self.config = SimpleNamespace(
+            z_dim=z_dim,
+            scale_factor_temporal=temporal,
+            scale_factor_spatial=spatial,
+            latents_mean=[0.0] * z_dim,
+            latents_std=[1.0] * z_dim,
+        )
+
+    def encode(self, video: torch.Tensor):
+        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
+        latents = torch.ones(
+            video.shape[0],
+            self.config.z_dim,
+            latent_frames,
+            latent_height,
+            latent_width,
+            dtype=video.dtype,
+            device=video.device,
+        )
+        return SimpleNamespace(latent_dist=_ModeLatentDist(latents))
+
+    def decode(self, latents: torch.Tensor, return_dict: bool = False):
+        del return_dict
+        return (latents,)
+
+
+class StubCosmos3AVAE:
+    def __init__(self, **kwargs: Any) -> None:
+        self.kwargs = kwargs
+        self.sample_rate = int(kwargs["sample_rate"])
+        self.audio_channels = int(kwargs["audio_channels"])
+        self.latent_ch = int(kwargs["io_channels"])
+        self.temporal_compression_factor = int(kwargs["hop_size"])
+
+    def get_latent_num_samples(self, num_audio_samples: int) -> int:
+        return int(num_audio_samples) // self.temporal_compression_factor
+
+    def get_audio_num_samples(self, num_latent_samples: int) -> int:
+        return int(num_latent_samples) * self.temporal_compression_factor
+
+    def decode(self, latents: torch.Tensor) -> torch.Tensor:
+        return torch.zeros(latents.shape[0], self.audio_channels, 8)
+
+
+class StubCosmos3Transformer(nn.Module):
+    def __init__(
+        self,
+        *,
+        latent_channel_size: int = 2,
+        sound_gen: bool = False,
+        sound_dim: int = 3,
+        sound_latent_fps: float = 25.0,
+    ) -> None:
+        super().__init__()
+        self.latent_channel_size = latent_channel_size
+        self.sound_gen = sound_gen
+        self.sound_dim = sound_dim
+        self.sound_latent_fps = sound_latent_fps
+        self.cached_kv: Any | None = None
+        self.cached_freqs_gen: Any | None = None
+        self.calls: list[dict[str, Any]] = []
+        self.reset_calls = 0
+
+    def reset_cache(self) -> None:
+        self.reset_calls += 1
+        self.cached_kv = None
+        self.cached_freqs_gen = None
+
+    def forward(
+        self,
+        *,
+        hidden_states: torch.Tensor,
+        timestep: torch.Tensor,
+        text_ids: torch.Tensor,
+        text_mask: torch.Tensor,
+        **kwargs: Any,
+    ) -> 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")
+        self.calls.append(
+            {
+                "token": token,
+                "timestep": timestep.clone(),
+                "text_mask": text_mask.clone(),
+                "cache_before": self.cached_kv,
+                "kwargs": dict(kwargs),
+            }
+        )
+        if self.cached_kv is None:
+            marker = torch.tensor([token], dtype=torch.float32)
+            self.cached_kv = [(marker, marker + 100)]
+            self.cached_freqs_gen = (marker + 200, marker + 300)
+        outputs: list[torch.Tensor] = [torch.full_like(hidden_states, float(token))]
+        if sound_latents is not None:
+            outputs.append(torch.full_like(sound_latents, float(token + 10)))
+        return outputs[0] if len(outputs) == 1 else tuple(outputs)
+
+
+def passthrough_progress_bar(iterable):
+    return iterable
+
+
+@pytest.fixture(autouse=True)
+def fake_cosmos3_guardrails(monkeypatch: pytest.MonkeyPatch):
+    module = types.ModuleType("vllm_omni.diffusion.models.cosmos3.guardrails")
+    module.is_guardrails_enabled = lambda od_config, sampling_params=None: False
+    module.ensure_initialized = lambda od_config: None
+    module.check_text_safety = lambda text: None
+    module.check_video_safety = lambda video: video
+    monkeypatch.setitem(sys.modules, module.__name__, module)
+    return module
+
+
+@pytest.fixture
+def make_cosmos3_pipeline():
+    def _make():
+        from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import (
+            Cosmos3OmniDiffusersPipeline,
+        )
+
+        pipeline = object.__new__(Cosmos3OmniDiffusersPipeline)
+        nn.Module.__init__(pipeline)
+        pipeline.od_config = SimpleNamespace()
+        pipeline.device = torch.device("cpu")
+        pipeline.dtype = torch.float32
+        pipeline.transformer = StubCosmos3Transformer(latent_channel_size=2)
+        pipeline.vae = StubCosmos3VAE(z_dim=2)
+        pipeline.vae_scale_factor_temporal = 4
+        pipeline.vae_scale_factor_spatial = 8
+        pipeline.scheduler = StubScheduler([9, 3], flow_shift=1.0)
+        pipeline._base_scheduler_config = pipeline.scheduler.config
+        pipeline._engine_init_flow_shift = 1.0
+        pipeline._current_flow_shift = 1.0
+        pipeline._guidance_scale = None
+        pipeline._num_timesteps = None
+        pipeline._cache_dit_requires_paired_cfg = False
+        pipeline._sound_tokenizer = None
+        pipeline.progress_bar = passthrough_progress_bar
+        return pipeline
+
+    return _make
+
+
+def make_sampling_params(**overrides: Any) -> SimpleNamespace:
+    values = {
+        "height": None,
+        "width": None,
+        "num_frames": None,
+        "num_inference_steps": None,
+        "guidance_scale": None,
+        "generator": None,
+        "seed": 123,
+        "num_outputs_per_prompt": 1,
+        "frame_rate": None,
+        "resolved_frame_rate": None,
+        "max_sequence_length": None,
+        "extra_args": {},
+    }
+    values.update(overrides)
+    return SimpleNamespace(**values)
+
+
+def _ids(value: int) -> torch.Tensor:
+    return torch.tensor([[value]], dtype=torch.long)
+
+
+def _mask() -> torch.Tensor:
+    return torch.ones(1, 1, dtype=torch.long)
+
+
+def test_pipeline_registered_and_exported() -> None:
+    from vllm_omni.diffusion.cache.cache_dit_backend import CUSTOM_DIT_ENABLERS
+    from vllm_omni.diffusion.models import cosmos3
+    from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline
+    from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
+    from vllm_omni.diffusion.registry import (
+        _DIFFUSION_MODELS,
+        _DIFFUSION_POST_PROCESS_FUNCS,
+        _DIFFUSION_PRE_PROCESS_FUNCS,
+    )
+
+    assert issubclass(Cosmos3OmniDiffusersPipeline, nn.Module)
+    assert issubclass(Cosmos3OmniDiffusersPipeline, ProgressBarMixin)
+    assert Cosmos3OmniDiffusersPipeline.support_image_input is True
+    assert _DIFFUSION_MODELS["Cosmos3OmniDiffusersPipeline"] == (
+        "cosmos3",
+        "pipeline_cosmos3",
+        "Cosmos3OmniDiffusersPipeline",
+    )
+    assert _DIFFUSION_PRE_PROCESS_FUNCS["Cosmos3OmniDiffusersPipeline"] == "get_cosmos3_pre_process_func"
+    assert _DIFFUSION_POST_PROCESS_FUNCS["Cosmos3OmniDiffusersPipeline"] == "get_cosmos3_post_process_func"
+    assert "Cosmos3OmniDiffusersPipeline" in CUSTOM_DIT_ENABLERS
+    assert "Cosmos3OmniDiffusersPipeline" in cosmos3.__all__
+
+
+@pytest.fixture
+def stub_real_pipeline_init(monkeypatch: pytest.MonkeyPatch):
+    from vllm_omni.diffusion.models.cosmos3 import pipeline_cosmos3
+
+    class _StubAutoTokenizer:
+        @classmethod
+        def from_pretrained(cls, *args, **kwargs):
+            return SimpleNamespace()
+
+    class _StubDiffusersVAE:
+        config = SimpleNamespace(scale_factor_temporal=4, scale_factor_spatial=8)
+
+        @classmethod
+        def from_pretrained(cls, *args, **kwargs):
+            return cls()
+
+        def to(self, _device):
+            return self
+
+    class _StubDiffusersScheduler:
+        config = SimpleNamespace(flow_shift=1.0)
+
+        @classmethod
+        def from_pretrained(cls, *args, **kwargs):
+            return cls()
+
+    class _StubVideoProcessor:
+        def __init__(self, *args, **kwargs) -> None:
+            pass
+
+    monkeypatch.setattr(pipeline_cosmos3, "AutoTokenizer", _StubAutoTokenizer)
+    monkeypatch.setattr(pipeline_cosmos3, "DistributedAutoencoderKLWan", _StubDiffusersVAE)
+    monkeypatch.setattr(pipeline_cosmos3, "UniPCMultistepScheduler", _StubDiffusersScheduler)
+    monkeypatch.setattr(pipeline_cosmos3, "VideoProcessor", _StubVideoProcessor)
+    monkeypatch.setattr(pipeline_cosmos3, "get_local_device", lambda: torch.device("cpu"))
+
+
+def _make_od_config(*, sound_gen: bool) -> SimpleNamespace:
+    tf_model_config = {
+        "hidden_size": 8,
+        "num_hidden_layers": 0,
+        "num_attention_heads": 2,
+        "num_key_value_heads": 2,
+        "head_dim": 4,
+        "intermediate_size": 16,
+        "vocab_size": 32,
+        "latent_patch_size": 1,
+        "latent_channel": 2,
+        "rope_scaling": {"mrope_section": [1, 1, 0]},
+    }
+    if sound_gen:
+        tf_model_config["sound_gen"] = True
+    return SimpleNamespace(
+        enable_cpu_offload=False,
+        enable_diffusion_pipeline_profiler=False,
+        model="/nonexistent/model/path",
+        dtype=torch.float32,
+        flow_shift=None,
+        quantization_config=None,
+        custom_pipeline_args={},
+        model_config={},
+        tf_model_config=tf_model_config,
+    )
+
+
+def test_pipeline_init_skips_tokenizer_when_sound_disabled(stub_real_pipeline_init) -> None:
+    from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline
+
+    pipeline = Cosmos3OmniDiffusersPipeline(od_config=_make_od_config(sound_gen=False))
+
+    assert pipeline._sound_tokenizer is None
+    assert pipeline.transformer.sound_gen is False
+    assert not hasattr(pipeline.transformer, "audio_proj_in")
+    assert not hasattr(pipeline.transformer, "audio_proj_out")
+
+
+def test_pipeline_init_passes_tokenizer_attrs_into_transformer(
+    stub_real_pipeline_init,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    from vllm_omni.diffusion.models.cosmos3 import sound_tokenizer
+    from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline
+
+    stub_tokenizer = sound_tokenizer.Cosmos3SoundTokenizer(
+        StubCosmos3AVAE(sample_rate=32000, audio_channels=2, io_channels=5, hop_size=800)
+    )
+    monkeypatch.setattr(
+        sound_tokenizer.Cosmos3SoundTokenizer,
+        "from_config",
+        classmethod(lambda cls, od_config: stub_tokenizer),
+    )
+
+    pipeline = Cosmos3OmniDiffusersPipeline(od_config=_make_od_config(sound_gen=True))
+
+    assert pipeline._sound_tokenizer is stub_tokenizer
+    assert pipeline.transformer.sound_gen is True
+    assert pipeline.transformer.sound_dim == pipeline._sound_tokenizer.latent_ch == 5
+    assert pipeline.transformer.sound_latent_fps == pipeline._sound_tokenizer.latent_fps == 40.0
+    assert pipeline.transformer.audio_proj_in.in_features == 5
+    assert pipeline.transformer.audio_proj_out.out_features == 5
+
+
+def test_preprocess_i2v_image_input() -> None:
+    from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import get_cosmos3_pre_process_func
+
+    preprocess = get_cosmos3_pre_process_func(SimpleNamespace())
+    i2v = SimpleNamespace(
+        prompts=[{"prompt": "A slow camera push.", "multi_modal_data": {"image": Image.new("RGB", (320, 160))}}],
+        sampling_params=SimpleNamespace(height=None, width=None, extra_args={}),
+    )
+
+    result = preprocess(i2v)
+    assert (result.sampling_params.height, result.sampling_params.width) == (672, 1344)
+    assert tuple(result.prompts[0]["additional_information"]["preprocessed_image"].shape[-2:]) == (672, 1344)
+
+
+def test_postprocess_handles_image_video_audio_and_validation() -> None:
+    from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import get_cosmos3_post_process_func
+
+    func = get_cosmos3_post_process_func(SimpleNamespace())
+    video = torch.zeros(1, 3, 1, 4, 4)
+
+    assert func(video, output_type="latent") is video
+    assert func({"image": video})[0].size == (4, 4)
+    # Video-only postprocess returns the bare processed video (not a dict),
+    # matching the image/latent branches and peer audio-capable pipelines.
+    assert not isinstance(func({"video": video}), dict)
+    assert (
+        func(
+            {"video": video, "audio": torch.ones(1, 2, 16), "audio_sample_rate": 48000},
+            sampling_params=SimpleNamespace(extra_args={"resolved_frame_rate": 12}),
+        )["audio_sample_rate"]
+        == 48000
+    )
+
+    with pytest.raises(ValueError, match="text-to-image postprocess expects"):
+        func({"image": torch.zeros(1, 3, 2, 4, 4)})
+    with pytest.raises(ValueError, match="both image and video"):
+        func({"image": video, "video": video})
+
+
+def test_prompt_formatting_and_checkpoint_key_remap(make_cosmos3_pipeline) -> None:
+    from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline
+
+    pipeline = make_cosmos3_pipeline()
+    captured: list[str] = []
+    pipeline._tokenize_prompt = lambda text, *args, **kwargs: (captured.append(text) or _ids(len(captured)), _mask())
+
+    pipeline._format_and_tokenize_prompts(
+        "A robot",
+        "bad",
+        num_frames=48,
+        frame_rate=24,
+        height=720,
+        width=1280,
+        max_sequence_length=32,
+        sp=SimpleNamespace(
+            extra_args={
+                "negative_metadata_mode": "inverse",
+                # Duration/resolution metadata templates are off by default
+                # (commit "Change resolution and duration templates to off by
+                # default"); enable them explicitly to exercise the formatting.
+                "use_duration_template": True,
+                "use_resolution_template": True,
+            }
+        ),
+        use_system_prompt=True,
+        is_t2i=False,
+    )
+    assert "The video is 2.0 seconds long" in captured[0]
+    assert "The video is not 2.0 seconds long" in captured[1]
+
+    remaps = {
+        "embed_tokens.weight": "transformer.language_model.embed_tokens.weight",
+        "model.embed_tokens.weight": "transformer.language_model.embed_tokens.weight",
+        "norm.weight": "transformer.language_model.norm.weight",
+        "norm_moe_gen.weight": "transformer.norm_moe_gen.weight",
+        "proj_in.weight": "transformer.proj_in.weight",
+        "proj_out.bias": "transformer.proj_out.bias",
+        "layers.3.self_attn.to_q.weight": "transformer.language_model.layers.3.self_attn.to_q.weight",
+        "layers.3.self_attn.to_out.weight": "transformer.language_model.layers.3.self_attn.to_out.weight",
+        "layers.3.self_attn.norm_q.weight": "transformer.language_model.layers.3.self_attn.norm_q.weight",
+        "layers.3.self_attn.add_q_proj.weight": "transformer.gen_layers.3.cross_attention.to_q.weight",
+        "layers.3.self_attn.to_add_out.weight": "transformer.gen_layers.3.cross_attention.to_out.weight",
+        "layers.3.self_attn.norm_added_q.weight": "transformer.gen_layers.3.cross_attention.norm_q.weight",
+        "transformer.model.layers.3.self_attn.add_k_proj.weight": (
+            "transformer.gen_layers.3.cross_attention.to_k.weight"
+        ),
+    }
+    assert {key: Cosmos3OmniDiffusersPipeline._remap_ckpt_key(key) for key in remaps} == remaps
+
+
+def test_prepare_latents_for_video_image_and_sound(make_cosmos3_pipeline) -> None:
+    pipeline = 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)
+    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)
+    )
+    torch.testing.assert_close(i2v_latents[:, :, 0], torch.full((1, 2, 2, 3), 5.0))
+    assert velocity_mask.tolist() == [[[[[0.0]], [[1.0]]]]]
+    assert image_latent.shape == (1, 2, 1, 2, 3)
+
+    pipeline.transformer = pipeline.transformer.__class__(latent_channel_size=2, sound_gen=True, sound_dim=3)
+    pipeline._sound_tokenizer = SimpleNamespace(
+        sample_rate=10,
+        latent_ch=3,
+        hop_size=4,
+        decode=lambda x: torch.ones(x.shape[0], 2, 24),
+    )
+    assert pipeline._resolve_sound_target_samples(SimpleNamespace(extra_args={"sound_duration": 2.0}), 9, 3.0) == (
+        20,
+        2.0,
+        10,
+    )
+    sound_latents, latent_frames = pipeline._prepare_sound_latents(21, torch.Generator(device="cpu").manual_seed(0))
+    assert (sound_latents.shape, latent_frames) == (torch.Size([1, 3, 6]), 6)
+    assert pipeline._decode_sound_latents(torch.zeros(1, 3, 6), target_audio_samples=21).shape == (1, 2, 21)
+
+
+def test_diffuse_covers_cfg_i2v_and_sound_steps(make_cosmos3_pipeline) -> None:
+    pipeline = make_cosmos3_pipeline()
+    latents = torch.zeros(1, 2, 1, 1, 1)
+
+    result = pipeline.diffuse(
+        latents=latents,
+        timesteps=torch.tensor([900, 100]),
+        cond_ids=_ids(2),
+        cond_mask=_mask(),
+        uncond_ids=_ids(1),
+        uncond_mask=_mask(),
+        guidance_scale=3.0,
+        shared_kwargs={"video_shape": (1, 1, 1), "fps": 24.0},
+        guidance_interval=(500.0, 1000.0),
+    )
+    assert [call["token"] for call in pipeline.transformer.calls] == [2, 1, 2]
+    torch.testing.assert_close(result, torch.full_like(latents, 6.0))
+
+    i2v = pipeline.diffuse(
+        latents=torch.zeros(1, 2, 2, 1, 1),
+        timesteps=torch.tensor([7]),
+        cond_ids=_ids(2),
+        cond_mask=_mask(),
+        uncond_ids=_ids(1),
+        uncond_mask=_mask(),
+        guidance_scale=1.0,
+        shared_kwargs={"video_shape": (2, 1, 1), "fps": 24.0},
+        velocity_mask=torch.tensor([[[[[0.0]], [[1.0]]]]]),
+        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))
+
+    pipeline.transformer = pipeline.transformer.__class__(latent_channel_size=2, sound_gen=True, sound_dim=3)
+    video_result, sound_result = pipeline.diffuse(
+        latents=latents,
+        sound_latents=torch.zeros(1, 3, 4),
+        timesteps=torch.tensor([7, 3]),
+        cond_ids=_ids(2),
+        cond_mask=_mask(),
+        uncond_ids=_ids(1),
+        uncond_mask=_mask(),
+        guidance_scale=1.0,
+        shared_kwargs={"video_shape": (1, 1, 1), "fps": 24.0},
+    )
+    torch.testing.assert_close(video_result, torch.full_like(latents, 4.0))
+    torch.testing.assert_close(sound_result, torch.full((), 24.0).expand_as(sound_result))
+
+
+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
+    output is numerically identical to the skip path.
+
+    Contrast with ``test_diffuse_covers_cfg_and_i2v_steps`` (no marker), where
+    the same inputs skip the out-of-interval uncond pass: calls == [2, 1, 2].
+    """
+    pipeline = make_cosmos3_pipeline()
+    # Marker normally set by ``enable_cache_for_cosmos3`` when cache-dit is on.
+    pipeline._cache_dit_requires_paired_cfg = True
+    latents = torch.zeros(1, 2, 1, 1, 1)
+
+    result = pipeline.diffuse(
+        latents=latents,
+        timesteps=torch.tensor([900, 100]),
+        cond_ids=_ids(2),
+        cond_mask=_mask(),
+        uncond_ids=_ids(1),
+        uncond_mask=_mask(),
+        guidance_scale=3.0,
+        shared_kwargs={"video_shape": (1, 1, 1), "fps": 24.0},
+        guidance_interval=(500.0, 1000.0),
+    )
+
+    # t=900 is inside the interval (cond+uncond); t=100 is outside but the
+    # uncond pass is still issued -> paired cond/uncond at every step.
+    assert [call["token"] for call in pipeline.transformer.calls] == [2, 1, 2, 1]
+    # Identical result to the skip path: out-of-interval combine uses scale=1.0,
+    # so combine_cfg_noise(cond=2, uncond=1, 1.0) == 2 == the skipped cond value.
+    torch.testing.assert_close(result, torch.full_like(latents, 6.0))
+
+
+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):
+            captured["format"] = {
+                "prompt": prompt,
+                "negative_prompt": negative_prompt,
+                "num_frames": num_frames,
+                "frame_rate": frame_rate,
+                "height": height,
+                "width": width,
+                "is_t2i": kwargs["is_t2i"],
+            }
+            return _ids(2), _mask(), _ids(1), _mask()
+
+        def fake_prepare(height, width, num_frames, generator):
+            captured["prepare_calls"].append((height, width, num_frames, generator.initial_seed()))
+            return torch.zeros(1, 2, 1, 1, 1)
+
+        def fake_diffuse(**kwargs):
+            captured["diffuse_calls"].append(kwargs)
+            outputs = [kwargs["latents"] + len(captured["diffuse_calls"])]
+            if kwargs.get("sound_latents") is not None:
+                outputs.append(kwargs["sound_latents"] + 2.0)
+            return outputs[0] if len(outputs) == 1 else tuple(outputs)
+
+        pipeline._format_and_tokenize_prompts = fake_format
+        pipeline._prepare_latents = fake_prepare
+        pipeline._set_flow_shift = lambda target: captured.setdefault("flow_shifts", []).append(target)
+        pipeline.diffuse = fake_diffuse
+        pipeline._decode_latents = lambda latents: latents
+        return captured
+
+    @pytest.mark.parametrize(
+        ("prompt", "sampling_params", "expected"),
+        [
+            (
+                {"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},
+            ),
+            (
+                "A warehouse robot",
+                make_sampling_params(),
+                {"key": "video", "is_t2i": False, "flow": [1.0], "steps": [35], "frames": 189},
+            ),
+        ],
+    )
+    def test_forward_defaults_and_mode_selection(
+        self,
+        make_cosmos3_pipeline,
+        prompt,
+        sampling_params,
+        expected,
+    ) -> None:
+        pipeline = make_cosmos3_pipeline()
+        captured = self._install_forward_stubs(pipeline)
+
+        output = pipeline.forward(SimpleNamespace(prompts=[prompt], sampling_params=sampling_params))
+
+        assert expected["key"] in output.output
+        assert captured["format"]["is_t2i"] is expected["is_t2i"]
+        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"]
+
+    def test_forward_i2v_and_sound_routes(self, make_cosmos3_pipeline) -> None:
+        pipeline = make_cosmos3_pipeline()
+        captured = self._install_forward_stubs(pipeline)
+        image_tensor = torch.zeros(1, 3, 16, 16)
+        velocity_mask = torch.ones(1, 1, 1, 1, 1)
+
+        pipeline._prepare_latents_i2v = lambda *args, **kwargs: (
+            torch.zeros(1, 2, 1, 1, 1),
+            velocity_mask,
+            torch.zeros(1, 2, 1, 1, 1),
+        )
+        pipeline.forward(
+            SimpleNamespace(
+                prompts=[
+                    {
+                        "prompt": "move",
+                        "modalities": ["video"],
+                        "additional_information": {"preprocessed_image": image_tensor},
+                    }
+                ],
+                sampling_params=make_sampling_params(height=16, width=16, num_frames=5),
+            )
+        )
+        assert captured["diffuse_calls"][-1]["shared_kwargs"]["noisy_frame_mask"] is velocity_mask
+
+        pipeline.transformer = pipeline.transformer.__class__(latent_channel_size=2, sound_gen=True, sound_dim=3)
+        sound_latents = torch.zeros(1, 3, 4)
+        pipeline._resolve_sound_target_samples = lambda *args: (20, 2.0, 10)
+        pipeline._prepare_sound_latents = lambda *args: (sound_latents, 4)
+        pipeline._decode_sound_latents = lambda *args: torch.ones(1, 2, 20)
+        output = pipeline.forward(
+            SimpleNamespace(
+                prompts=[{"prompt": "A robot", "modalities": ["video"], "generate_sound": True}],
+                sampling_params=make_sampling_params(num_frames=9, frame_rate=3.0),
+            )
+        )
+        assert captured["diffuse_calls"][-1]["sound_latents"] is sound_latents
+        assert output.output["audio_sample_rate"] == 10
+
+    @pytest.mark.parametrize(
+        ("prompt", "sampling_params", "message"),
+        [
+            (["one", "two"], make_sampling_params(), "single prompt"),
+            ([{"prompt": "one", "modalities": ["image", "video"]}], make_sampling_params(), "both image and video"),
+            (
+                [{"prompt": "x", "modalities": ["image"], "generate_sound": True}],
+                make_sampling_params(),
+                "only for video",
+            ),
+        ],
+    )
+    def test_forward_rejects_invalid_public_requests(
+        self,
+        make_cosmos3_pipeline,
+        prompt,
+        sampling_params,
+        message,
+    ) -> None:
+        pipeline = make_cosmos3_pipeline()
+        pipeline.transformer = pipeline.transformer.__class__(latent_channel_size=2, sound_gen=True, sound_dim=3)
+
+        with pytest.raises(ValueError, match=message):
+            pipeline.forward(SimpleNamespace(prompts=prompt, sampling_params=sampling_params))

tests/diffusion/models/cosmos3/test_cosmos3_sound_tokenizer.py

@@ -0,0 +1,236 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+DIFFUSERS_SOUND_TOKENIZER_CHECKPOINT_NAME = "diffusion_pytorch_model.safetensors"
+
+
+class _FakeAVAEAudioTokenizer:
+    def __init__(self, **kwargs) -> None:
+        self.kwargs = kwargs
+        self.sample_rate = int(kwargs["sample_rate"])
+        self.audio_channels = int(kwargs["audio_channels"])
+        self.latent_ch = int(kwargs["io_channels"])
+        self.temporal_compression_factor = int(kwargs["hop_size"])
+
+    def get_latent_num_samples(self, num_audio_samples: int) -> int:
+        return int(num_audio_samples) // self.temporal_compression_factor
+
+    def get_audio_num_samples(self, num_latent_samples: int) -> int:
+        return int(num_latent_samples) * self.temporal_compression_factor
+
+    def decode(self, latents: torch.Tensor) -> torch.Tensor:
+        return torch.zeros(latents.shape[0], self.audio_channels, 8)
+
+
+def _write_component(root: Path, config: dict | None = None, checkpoint_name: str | None = None) -> Path:
+    tokenizer_dir = root / "sound_tokenizer"
+    tokenizer_dir.mkdir(parents=True)
+    if checkpoint_name:
+        (tokenizer_dir / checkpoint_name).write_bytes(b"stub")
+    (tokenizer_dir / "config.json").write_text(json.dumps(config or {}), encoding="utf-8")
+    return tokenizer_dir
+
+
+def _patch_fake_avae(monkeypatch: pytest.MonkeyPatch, created: dict) -> None:
+    from vllm_omni.diffusion.models.cosmos3 import sound_tokenizer
+
+    class FakeAVAE(_FakeAVAEAudioTokenizer):
+        def __init__(self, **kwargs) -> None:
+            created.update(kwargs)
+            super().__init__(**kwargs)
+
+    monkeypatch.setattr(sound_tokenizer, "Cosmos3AVAEAudioTokenizer", FakeAVAE)
+    monkeypatch.setattr(sound_tokenizer, "get_local_device", lambda: torch.device("cpu"))
+
+
+def test_from_config_loads_local_diffusers_component(tmp_path, monkeypatch: pytest.MonkeyPatch) -> None:
+    from vllm_omni.diffusion.models.cosmos3 import sound_tokenizer
+
+    model_dir = tmp_path / "model"
+    tokenizer_dir = _write_component(model_dir, checkpoint_name=DIFFUSERS_SOUND_TOKENIZER_CHECKPOINT_NAME)
+    created = {}
+    _patch_fake_avae(monkeypatch, created)
+
+    tokenizer = sound_tokenizer.Cosmos3SoundTokenizer.from_config(
+        SimpleNamespace(
+            model=str(model_dir),
+            custom_pipeline_args={"sound_sample_rate": 32000, "sound_hop_size": 800, "sound_dim": 3},
+            dtype=torch.float32,
+        )
+    )
+
+    assert created["checkpoint_path"] == str(tokenizer_dir / DIFFUSERS_SOUND_TOKENIZER_CHECKPOINT_NAME)
+    assert created["config_path"] == str(tokenizer_dir / "config.json")
+    assert (tokenizer.sample_rate, tokenizer.latent_ch, tokenizer.hop_size, tokenizer.latent_fps) == (
+        32000,
+        3,
+        800,
+        40.0,
+    )
+
+
+def test_from_config_downloads_component_from_hf_repo(tmp_path, monkeypatch: pytest.MonkeyPatch) -> None:
+    import huggingface_hub
+
+    from vllm_omni.diffusion.models.cosmos3 import sound_tokenizer
+
+    cache_dir = tmp_path / "hf"
+    _write_component(cache_dir, checkpoint_name=DIFFUSERS_SOUND_TOKENIZER_CHECKPOINT_NAME)
+    calls = []
+    created = {}
+    _patch_fake_avae(monkeypatch, created)
+
+    def fake_snapshot_download(repo_id: str, *, revision: str | None, allow_patterns: list[str]) -> str:
+        calls.append((repo_id, revision, allow_patterns))
+        return str(cache_dir)
+
+    monkeypatch.setattr(huggingface_hub, "snapshot_download", fake_snapshot_download)
+
+    sound_tokenizer.Cosmos3SoundTokenizer.from_config(
+        SimpleNamespace(
+            model="nvidia/cosmos3",
+            revision="test-rev",
+            custom_pipeline_args={"sound_sample_rate": 32000, "sound_hop_size": 800, "sound_dim": 3},
+            dtype=torch.float32,
+        )
+    )
+
+    assert created["checkpoint_path"].endswith(DIFFUSERS_SOUND_TOKENIZER_CHECKPOINT_NAME)
+    assert calls == [
+        (
+            "nvidia/cosmos3",
+            "test-rev",
+            ["sound_tokenizer/config.json", f"sound_tokenizer/{DIFFUSERS_SOUND_TOKENIZER_CHECKPOINT_NAME}"],
+        )
+    ]
+
+
+@pytest.mark.parametrize(
+    ("checkpoint_name", "message"),
+    [
+        (None, "no AVAE sound tokenizer checkpoint"),
+        ("model.safetensors", DIFFUSERS_SOUND_TOKENIZER_CHECKPOINT_NAME),
+    ],
+)
+def test_default_component_requires_diffusers_checkpoint_name(tmp_path, checkpoint_name, message) -> None:
+    from vllm_omni.diffusion.models.cosmos3 import sound_tokenizer
+
+    model_dir = tmp_path / "model"
+    _write_component(model_dir, checkpoint_name=checkpoint_name)
+
+    with pytest.raises(ValueError, match=message):
+        sound_tokenizer.Cosmos3SoundTokenizer.from_config(
+            SimpleNamespace(model=str(model_dir), custom_pipeline_args={}, dtype=torch.float32)
+        )
+
+
+def test_component_config_precedence_and_conflict_detection(tmp_path, monkeypatch: pytest.MonkeyPatch) -> None:
+    from vllm_omni.diffusion.models.cosmos3 import sound_tokenizer
+
+    component_config = {
+        "sampling_rate": 48000,
+        "dec_out_channels": 2,
+        "vocoder_input_dim": 64,
+        "hop_size": 1920,
+    }
+    model_dir = tmp_path / "model"
+    _write_component(model_dir, component_config, DIFFUSERS_SOUND_TOKENIZER_CHECKPOINT_NAME)
+    created = {}
+    _patch_fake_avae(monkeypatch, created)
+
+    tokenizer = sound_tokenizer.Cosmos3SoundTokenizer.from_config(
+        SimpleNamespace(
+            model=str(model_dir),
+            custom_pipeline_args={
+                "sound_normalize_latents": True,
+                "sound_normalization_type": "tanh",
+                "sound_tanh_input_scale": 2.0,
+            },
+            model_config={
+                "sound_tokenizer": {
+                    "sample_rate": 32000,
+                    "audio_channels": 1,
+                    "io_channels": 3,
+                    "hop_size": 800,
+                    "normalize_latents": False,
+                    "normalization_type": "none",
+                }
+            },
+            dtype=torch.float32,
+        )
+    )
+
+    assert (created["sample_rate"], created["audio_channels"], created["io_channels"], created["hop_size"]) == (
+        48000,
+        2,
+        64,
+        1920,
+    )
+    assert (created["normalize_latents"], created["normalization_type"], created["tanh_input_scale"]) == (
+        True,
+        "tanh",
+        2.0,
+    )
+    assert (tokenizer.sample_rate, tokenizer.latent_ch, tokenizer.hop_size, tokenizer.latent_fps) == (
+        48000,
+        64,
+        1920,
+        25.0,
+    )
+
+    with pytest.raises(ValueError, match=r"sample_rate.*48000.*32000"):
+        sound_tokenizer.Cosmos3SoundTokenizer.from_config(
+            SimpleNamespace(
+                model=str(model_dir),
+                custom_pipeline_args={"sound_sample_rate": 32000},
+                dtype=torch.float32,
+            )
+        )
+
+
+def test_avae_uses_diffusers_decoder_state_dict_layout(tmp_path) -> None:
+    from safetensors.torch import save_file
+
+    from vllm_omni.diffusion.models.cosmos3.audio_tokenizer import avae
+
+    config = {
+        "sampling_rate": 8000,
+        "hop_size": 2,
+        "dec_dim": 4,
+        "dec_c_mults": [1],
+        "dec_strides": [2],
+        "dec_out_channels": 1,
+        "vocoder_input_dim": 2,
+        "normalization_type": "none",
+    }
+    checkpoint_path = tmp_path / DIFFUSERS_SOUND_TOKENIZER_CHECKPOINT_NAME
+    config_path = tmp_path / "config.json"
+    config_path.write_text(json.dumps(config), encoding="utf-8")
+
+    decoder = avae.OobleckDecoder(4, 2, 1, [2], [1])
+    save_file({f"decoder.{key}": value for key, value in decoder.state_dict().items()}, str(checkpoint_path))
+
+    tokenizer = avae.Cosmos3AVAEAudioTokenizer(
+        checkpoint_path=checkpoint_path,
+        config_path=config_path,
+        dtype=torch.float32,
+        device="cpu",
+    )
+
+    keys = set(tokenizer.state_dict())
+    assert {"decoder.conv1.weight_g", "decoder.block.0.conv_t1.weight_g", "decoder.conv2.weight_g"} <= keys
+    assert not any(key.startswith(("decoder.layers.", "model.decoder.")) for key in keys)
+    assert tokenizer.decode(torch.zeros(1, 2, 3)).shape == (1, 1, 6)
+    with pytest.raises(NotImplementedError, match="decoder-only"):
+        tokenizer.encode(torch.zeros(1, 1, 6))

tests/diffusion/models/cosmos3/test_cosmos3_transformer.py

@@ -0,0 +1,283 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+from typing import Any
+
+import pytest
+import torch
+from torch import nn
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def _tiny_cosmos3_config(**overrides):
+    config = {
+        "hidden_size": 8,
+        "num_hidden_layers": 0,
+        "num_attention_heads": 2,
+        "num_key_value_heads": 2,
+        "head_dim": 4,
+        "intermediate_size": 16,
+        "vocab_size": 32,
+        "latent_patch_size": 1,
+        "latent_channel": 2,
+        "rope_scaling": {"mrope_section": [1, 1, 0]},
+    }
+    config.update(overrides)
+    return config
+
+
+def test_mrope_position_ids_cover_text_video_and_sound() -> None:
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import (
+        compute_mrope_position_ids_sound,
+        compute_mrope_position_ids_text,
+        compute_mrope_position_ids_vision,
+    )
+
+    text_ids, text_offset = compute_mrope_position_ids_text(num_tokens=3, temporal_offset=5)
+    assert text_ids.tolist() == [[5, 6, 7], [5, 6, 7], [5, 6, 7]]
+    assert text_offset == 8
+
+    vision_ids, vision_offset = compute_mrope_position_ids_vision(2, 2, 3, temporal_offset=10, fps=None)
+    assert vision_ids.shape == (3, 12)
+    assert vision_ids[0].tolist() == [10] * 6 + [11] * 6
+    assert vision_offset == 12
+
+    modulated_ids, modulated_offset = compute_mrope_position_ids_vision(
+        2,
+        1,
+        1,
+        temporal_offset=10,
+        fps=12.0,
+        base_fps=24.0,
+        temporal_compression_factor=4,
+    )
+    torch.testing.assert_close(modulated_ids[0], torch.tensor([10.0, 12.0]))
+    assert modulated_offset == 13
+
+    sound_ids, sound_offset = compute_mrope_position_ids_sound(3, temporal_offset=10, sound_latent_fps=25.0)
+    torch.testing.assert_close(sound_ids[0], torch.tensor([10.0, 10.96, 11.92]))
+    assert sound_offset == 12
+
+
+@pytest.mark.parametrize(
+    ("key", "value"),
+    [
+        ("qk_norm_for_diffusion", False),
+        ("qk_norm_for_text", False),
+        ("position_embedding_type", "rotary"),
+        ("unified_3d_mrope_reset_spatial_ids", False),
+        ("joint_attn_implementation", "one_way"),
+    ],
+)
+def test_validate_supported_config_rejects_unsupported_flags(key: str, value) -> None:
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer
+
+    with pytest.raises(ValueError, match=f"{key}="):
+        Cosmos3VFMTransformer._validate_supported_config({key: value})
+    Cosmos3VFMTransformer._validate_supported_config({})
+    Cosmos3VFMTransformer._validate_supported_config(None)
+
+
+def test_transformer_sharding_offload_and_patch_round_trip_contracts() -> None:
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer
+
+    model = object.__new__(Cosmos3VFMTransformer)
+    nn.Module.__init__(model)
+    model.language_model = nn.Module()
+    model.language_model.layers = nn.ModuleList([nn.Linear(2, 2) for _ in range(2)])
+    model.gen_layers = nn.ModuleList([nn.Linear(2, 2)])
+    model.norm_moe_gen = nn.LayerNorm(2)
+
+    matched = [
+        name
+        for name, module in model.named_modules()
+        if any(condition(name, module) for condition in model._hsdp_shard_conditions)
+    ]
+    assert matched == ["language_model.layers.0", "language_model.layers.1", "gen_layers.0"]
+    assert Cosmos3VFMTransformer._layerwise_offload_blocks_attrs == ["gen_layers"]
+    assert Cosmos3VFMTransformer._repeated_blocks == ["Cosmos3GenDecoderLayer"]
+
+    model.latent_patch_size = 2
+    model.latent_channel_size = 3
+    latents = torch.arange(1 * 3 * 1 * 3 * 5, dtype=torch.float32).reshape(1, 3, 1, 3, 5)
+    torch.testing.assert_close(model.unpatchify(model.patchify(latents, t=1, h=3, w=5), t=1, h=3, w=5), latents)
+
+
+def test_forward_returns_video_prediction(monkeypatch: pytest.MonkeyPatch) -> None:
+    from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
+
+    monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None))
+
+    output = transformer_cosmos3.Cosmos3VFMTransformer(
+        SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32)
+    )(
+        hidden_states=torch.zeros(1, 2, 1, 2, 2),
+        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,
+    )
+
+    assert tuple(output.shape) == (1, 2, 1, 2, 2)
+
+
+def test_sound_modules_follow_injected_sound_dim() -> None:
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer
+
+    tiny = _tiny_cosmos3_config()
+    no_modal = Cosmos3VFMTransformer(SimpleNamespace(tf_model_config=tiny, dtype=torch.float32))
+    with_sound = Cosmos3VFMTransformer(
+        SimpleNamespace(tf_model_config=tiny, dtype=torch.float32),
+        sound_gen=True,
+        sound_dim=5,
+        sound_latent_fps=40.0,
+    )
+
+    assert no_modal.sound_gen is False
+    assert not hasattr(no_modal, "audio_proj_in")
+    assert with_sound.sound_dim == 5
+    assert with_sound.sound_latent_fps == 40.0
+    assert with_sound.audio_proj_in.in_features == 5
+
+
+@pytest.mark.parametrize(
+    "kwargs",
+    [
+        {"sound_gen": True},
+        {"sound_gen": True, "sound_dim": 5},
+        {"sound_gen": True, "sound_latent_fps": 40.0},
+    ],
+)
+def test_transformer_requires_sound_dim_and_fps_when_sound_gen_true(kwargs: dict[str, Any]) -> None:
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer
+
+    with pytest.raises(ValueError, match=r"requires an explicit sound_dim and sound_latent_fps"):
+        Cosmos3VFMTransformer(
+            SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32),
+            **kwargs,
+        )
+
+
+def test_sound_pack_unpack_validate_shapes() -> None:
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer
+
+    model = object.__new__(Cosmos3VFMTransformer)
+    nn.Module.__init__(model)
+    model.sound_dim = 3
+
+    sound = torch.arange(2 * 3 * 4, dtype=torch.float32).reshape(2, 3, 4)
+    torch.testing.assert_close(model.unpack_sound(model.pack_sound(sound)), sound)
+
+    with pytest.raises(ValueError, match="channel mismatch"):
+        model.pack_sound(torch.zeros(1, 4, 2))
+
+
+def test_forward_returns_video_and_sound_predictions(monkeypatch: pytest.MonkeyPatch) -> None:
+    from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
+
+    monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None))
+
+    output = transformer_cosmos3.Cosmos3VFMTransformer(
+        SimpleNamespace(
+            tf_model_config=_tiny_cosmos3_config(),
+            dtype=torch.float32,
+        ),
+        sound_gen=True,
+        sound_dim=3,
+        sound_latent_fps=40.0,
+    )(
+        hidden_states=torch.zeros(1, 2, 1, 2, 2),
+        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,
+        sound_latents=torch.zeros(1, 3, 4),
+    )
+
+    assert isinstance(output, tuple)
+    assert [tuple(tensor.shape) for tensor in output] == [(1, 2, 1, 2, 2), (1, 3, 4)]
+
+
+def test_forward_with_sound_ulysses_error_mentions_combined_sequence(monkeypatch: pytest.MonkeyPatch) -> None:
+    import vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 as cosmos3_module
+
+    model = cosmos3_module.Cosmos3VFMTransformer(
+        SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32),
+        sound_gen=True,
+        sound_dim=3,
+        sound_latent_fps=40.0,
+    )
+    monkeypatch.setattr(cosmos3_module, "_get_ulysses_state", lambda: (2, 0, None))
+
+    with pytest.raises(ValueError, match=r"GEN sequence length \(3 = video tokens 2 \+ sound tokens 1\)"):
+        model(
+            hidden_states=torch.zeros(1, 2, 1, 1, 2),
+            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, 1, 2),
+            fps=24.0,
+            sound_latents=torch.zeros(1, 3, 1),
+        )
+
+
+def test_compute_rope_freqs_places_text_video_and_sound_positions() -> None:
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer
+
+    class FakeRotary:
+        def __init__(self) -> None:
+            self.position_ids: list[torch.Tensor] = []
+
+        def __call__(self, x, position_ids):
+            del x
+            self.position_ids.append(position_ids.detach().cpu())
+            batch, seq = position_ids.shape[1], position_ids.shape[2]
+            return torch.zeros(batch, seq, 4), torch.ones(batch, seq, 4)
+
+    rotary = FakeRotary()
+    model = object.__new__(Cosmos3VFMTransformer)
+    nn.Module.__init__(model)
+    model.language_model = SimpleNamespace(rotary_emb=rotary)
+    model.temporal_modality_margin = 100
+    model.base_fps = 24.0
+    model.temporal_compression_factor = 4
+    model.temporal_compression_factor_sound = 1
+    model.sound_latent_fps = 25.0
+    model.enable_fps_modulation = False
+
+    freqs_und, freqs_gen = model._compute_rope_freqs(
+        text_mask=torch.tensor([[1, 1, 0], [1, 0, 0]], dtype=torch.long),
+        t=2,
+        hp=1,
+        wp=1,
+        fps=24.0,
+        device=torch.device("cpu"),
+        dtype=torch.float32,
+    )
+    text_pos, vision_pos = rotary.position_ids
+    assert text_pos[:, 0, :].tolist() == [[0, 1, 0], [0, 1, 0], [0, 1, 0]]
+    assert vision_pos[0, 0].tolist() == [102, 103]
+    assert freqs_und[0].shape == (2, 3, 1, 4)
+    assert freqs_gen[0].shape == (2, 2, 1, 4)
+
+    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,
+        t_sound=1,
+    )
+
+    _, gen_pos = rotary.position_ids
+    assert gen_pos.shape == (3, 1, 3)
+    assert gen_pos[0, 0].tolist() == [102, 103, 102]

tests/diffusion/test_diffusion_engine.py

@@ -10,6 +10,9 @@ from types import SimpleNamespace
 from typing import Any
 
 import pytest
+import torch
+
+from vllm_omni.diffusion.diffusion_engine import _move_tensor_tree_to_cpu
 
 
 @dataclass
@@ -63,6 +66,80 @@ class MockScheduler:
         return [req.request_id for req in sched_output.scheduled_new_reqs]
 
 
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.cpu
+def test_move_tensor_tree_keeps_cpu_tensor_identity() -> None:
+    tensor = torch.arange(8, dtype=torch.float32)
+
+    moved = _move_tensor_tree_to_cpu(tensor)
+
+    assert moved is tensor
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.cpu
+def test_move_tensor_tree_preserves_nested_structure_without_mutating_input() -> None:
+    tensor = torch.arange(4, dtype=torch.float32)
+    nested_tensor = torch.arange(6, dtype=torch.float32).reshape(2, 3)
+    sentinel = object()
+    payload = {
+        "tensor": tensor,
+        "list": [nested_tensor, sentinel],
+        "tuple": ({"inner": tensor}, "metadata"),
+        "scalar": 3,
+    }
+
+    moved = _move_tensor_tree_to_cpu(payload)
+
+    assert moved is not payload
+    assert set(moved) == {"tensor", "list", "tuple", "scalar"}
+    assert moved["list"] is not payload["list"]
+    assert moved["tuple"] is not payload["tuple"]
+    assert moved["tuple"][0] is not payload["tuple"][0]
+    assert moved["tensor"] is tensor
+    assert moved["list"][0] is nested_tensor
+    assert moved["list"][1] is sentinel
+    assert moved["tuple"][0]["inner"] is tensor
+    assert moved["tuple"][1] == "metadata"
+    assert moved["scalar"] == 3
+    assert payload["list"][0] is nested_tensor
+    assert payload["list"][1] is sentinel
+    assert payload["tuple"][0]["inner"] is tensor
+    assert payload["tuple"][1] == "metadata"
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.cpu
+def test_move_tensor_tree_returns_non_tensor_values_unchanged() -> None:
+    value = object()
+
+    moved = _move_tensor_tree_to_cpu(value)
+
+    assert moved is value
+
+
+@pytest.mark.diffusion
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+def test_move_tensor_tree_moves_nested_cuda_tensors_to_cpu() -> None:
+    tensor = torch.arange(8, dtype=torch.float32, device="cuda")
+    other = torch.arange(4, dtype=torch.int64, device="cuda")
+    payload = {"tensor": tensor, "items": [other, ("keep", tensor)]}
+
+    moved = _move_tensor_tree_to_cpu(payload)
+
+    assert moved["tensor"].device.type == "cpu"
+    assert moved["items"][0].device.type == "cpu"
+    assert moved["items"][1][1].device.type == "cpu"
+    torch.testing.assert_close(moved["tensor"], tensor.cpu())
+    torch.testing.assert_close(moved["items"][0], other.cpu())
+    torch.testing.assert_close(moved["items"][1][1], tensor.cpu())
+    assert moved["items"][1][0] == "keep"
+
+
 @pytest.mark.asyncio
 async def test_async_add_req_and_wait_for_response():
     from vllm_omni.diffusion.diffusion_engine import DiffusionEngine

tests/diffusion/test_diffusion_ipc.py

@@ -0,0 +1,183 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import contextlib
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.data import DiffusionOutput
+from vllm_omni.diffusion.ipc import (
+    _SHM_TENSOR_THRESHOLD,
+    _pack_value_if_large,
+    _unpack_if_shm_handle,
+    pack_diffusion_output_shm,
+    unpack_diffusion_output_shm,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def _large_numel(dtype: torch.dtype) -> int:
+    return (_SHM_TENSOR_THRESHOLD // torch.empty((), dtype=dtype).element_size()) + 1
+
+
+def _cleanup_shm_handle(value: object) -> None:
+    if isinstance(value, dict) and value.get("__tensor_shm__"):
+        with contextlib.suppress(FileNotFoundError):
+            _unpack_if_shm_handle(value)
+
+
+def test_diffusion_output_dict_tensors_round_trip_through_shm() -> None:
+    image = torch.arange(300_000, dtype=torch.float32)
+    video = torch.arange(300_000, dtype=torch.float32) * 2
+    output = DiffusionOutput(output={"image": image, "video": video, "metadata": {"keep": "inline"}})
+
+    pack_diffusion_output_shm(output)
+
+    assert output.output["image"]["__tensor_shm__"] is True
+    assert output.output["video"]["__tensor_shm__"] is True
+    assert output.output["metadata"] == {"keep": "inline"}
+
+    unpack_diffusion_output_shm(output)
+
+    torch.testing.assert_close(output.output["image"], image)
+    torch.testing.assert_close(output.output["video"], video)
+    assert output.output["metadata"] == {"keep": "inline"}
+
+
+def test_diffusion_output_tuple_tensors_round_trip_through_shm() -> None:
+    # LTX2 / DreamID return (video, audio) tuples as DiffusionOutput.output.
+    video = torch.arange(300_000, dtype=torch.float32)
+    audio = torch.arange(300_000, dtype=torch.float32) * 3
+    output = DiffusionOutput(output=(video, audio))
+
+    pack_diffusion_output_shm(output)
+
+    assert isinstance(output.output, tuple)
+    assert output.output[0]["__tensor_shm__"] is True
+    assert output.output[1]["__tensor_shm__"] is True
+
+    unpack_diffusion_output_shm(output)
+
+    assert isinstance(output.output, tuple)
+    torch.testing.assert_close(output.output[0], video)
+    torch.testing.assert_close(output.output[1], audio)
+
+
+def test_diffusion_output_list_tensors_round_trip_through_shm() -> None:
+    frames = [torch.arange(300_000, dtype=torch.float32), torch.arange(300_000, dtype=torch.float32) + 1]
+    output = DiffusionOutput(output=list(frames))
+
+    pack_diffusion_output_shm(output)
+
+    assert isinstance(output.output, list)
+    assert all(isinstance(item, dict) and item["__tensor_shm__"] is True for item in output.output)
+
+    unpack_diffusion_output_shm(output)
+
+    assert isinstance(output.output, list)
+    torch.testing.assert_close(output.output[0], frames[0])
+    torch.testing.assert_close(output.output[1], frames[1])
+
+
+def test_pack_value_keeps_tensor_at_threshold_inline() -> None:
+    tensor = torch.arange(
+        _SHM_TENSOR_THRESHOLD // torch.empty((), dtype=torch.float32).element_size(),
+        dtype=torch.float32,
+    )
+
+    packed = _pack_value_if_large(tensor)
+
+    assert packed is tensor
+
+
+def test_pack_value_packs_large_tensor_and_round_trips() -> None:
+    tensor = torch.arange(_large_numel(torch.float32), dtype=torch.float32)
+    packed = _pack_value_if_large(tensor)
+
+    try:
+        assert isinstance(packed, dict)
+        assert packed["__tensor_shm__"] is True
+        assert packed["shape"] == [tensor.numel()]
+        assert packed["torch_dtype"] == "torch.float32"
+
+        unpacked = _unpack_if_shm_handle(packed)
+        assert isinstance(unpacked, torch.Tensor)
+        torch.testing.assert_close(unpacked, tensor)
+    finally:
+        _cleanup_shm_handle(packed)
+
+
+def test_pack_value_recurses_nested_dicts_and_lists_without_mutating_inline_values() -> None:
+    large = torch.arange(_large_numel(torch.float32), dtype=torch.float32)
+    small = torch.arange(8, dtype=torch.float32)
+    list_tensor = torch.arange(_large_numel(torch.float32), dtype=torch.float32)
+    payload = {
+        "media": {
+            "large": large,
+            "small": small,
+        },
+        "list_value": [list_tensor],
+        "metadata": {"prompt": "keep inline"},
+    }
+
+    packed = _pack_value_if_large(payload)
+
+    try:
+        assert packed is not payload
+        assert packed["media"] is not payload["media"]
+        assert packed["media"]["large"]["__tensor_shm__"] is True
+        assert packed["media"]["small"] is small
+        # Lists are recursed too: the large tensor inside is packed and a new
+        # list is returned, while the input payload is left untouched.
+        assert packed["list_value"] is not payload["list_value"]
+        assert packed["list_value"][0]["__tensor_shm__"] is True
+        assert payload["list_value"][0] is list_tensor
+        assert packed["metadata"] == {"prompt": "keep inline"}
+
+        torch.testing.assert_close(_unpack_if_shm_handle(packed["media"]["large"]), large)
+        torch.testing.assert_close(_unpack_if_shm_handle(packed["list_value"][0]), list_tensor)
+    finally:
+        if isinstance(packed, dict):
+            _cleanup_shm_handle(packed.get("media", {}).get("large"))
+            list_value = packed.get("list_value")
+            if isinstance(list_value, list) and list_value:
+                _cleanup_shm_handle(list_value[0])
+
+
+def test_pack_value_preserves_dtype_shape_and_values_for_bfloat16() -> None:
+    tensor = torch.arange(_large_numel(torch.bfloat16), dtype=torch.float32).to(torch.bfloat16).reshape(1, -1)
+    packed = _pack_value_if_large(tensor)
+
+    try:
+        assert isinstance(packed, dict)
+        assert packed["__tensor_shm__"] is True
+        assert packed["shape"] == list(tensor.shape)
+        assert packed["torch_dtype"] == "torch.bfloat16"
+        assert packed["numpy_dtype"] == "float32"
+
+        unpacked = _unpack_if_shm_handle(packed)
+        assert isinstance(unpacked, torch.Tensor)
+        assert unpacked.dtype == torch.bfloat16
+        torch.testing.assert_close(unpacked, tensor)
+    finally:
+        _cleanup_shm_handle(packed)
+
+
+def test_pack_value_packs_non_contiguous_large_tensor_values() -> None:
+    tensor = torch.arange(_large_numel(torch.float32) * 2, dtype=torch.float32).reshape(-1, 2)[:, 0]
+    assert not tensor.is_contiguous()
+
+    packed = _pack_value_if_large(tensor)
+
+    try:
+        assert isinstance(packed, dict)
+        assert packed["__tensor_shm__"] is True
+        assert packed["shape"] == list(tensor.shape)
+
+        unpacked = _unpack_if_shm_handle(packed)
+        assert isinstance(unpacked, torch.Tensor)
+        torch.testing.assert_close(unpacked, tensor)
+    finally:
+        _cleanup_shm_handle(packed)

tests/diffusion/test_diffusion_scheduler.py

@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-import asyncio
 import queue
 import threading
 from types import SimpleNamespace
@@ -447,11 +446,7 @@ class TestDiffusionEngine:
     @pytest.mark.asyncio
     async def test_step_raises_aborted_error(self, mocker: MockerFixture) -> None:
         engine = DiffusionEngine.__new__(DiffusionEngine)
-        engine._closed = False
-        engine._loop_started = True
-        engine._init_lock = asyncio.Lock()
-        engine.main_loop = asyncio.get_running_loop()
-        engine.stop_event = threading.Event()
+        engine._check_and_start_background_loop = mocker.AsyncMock()
         engine.pre_process_func = None
         engine.async_add_req_and_wait_for_response = mocker.AsyncMock(
             return_value=DiffusionOutput(aborted=True, abort_message="Request req-abort aborted.")
@@ -526,6 +521,52 @@ class TestDiffusionEngine:
         with pytest.raises(RuntimeError, match="Dummy run failed: boom"):
             engine._dummy_run()
 
+    @pytest.mark.asyncio
+    async def test_step_multi_request_reuses_multimodal_slice_logic(self, mocker: MockerFixture) -> None:
+        engine = DiffusionEngine.__new__(DiffusionEngine)
+        engine.od_config = SimpleNamespace(
+            model_class_name="mock_model",
+            enable_cpu_offload=False,
+        )
+        engine.pre_process_func = None
+        engine.post_process_func = None
+        engine._check_and_start_background_loop = mocker.AsyncMock()
+        engine.async_add_req_and_wait_for_response = mocker.AsyncMock(
+            return_value=DiffusionOutput(
+                output={
+                    "video": ["frame-0", "frame-1"],
+                    "audio": ["audio-0", "audio-1"],
+                    "actions": torch.tensor([[1.0, 2.0], [3.0, 4.0]]),
+                }
+            )
+        )
+
+        request = OmniDiffusionRequest(
+            prompts=["prompt-0", "prompt-1"],
+            sampling_params=OmniDiffusionSamplingParams(
+                num_inference_steps=1,
+                num_outputs_per_prompt=1,
+            ),
+            request_id="req-batch",
+        )
+
+        mocker.patch("vllm_omni.diffusion.diffusion_engine.supports_audio_output", return_value=False)
+        outputs = await engine.step(request)
+
+        assert len(outputs) == 2
+        assert outputs[0].images == ["frame-0"]
+        assert outputs[1].images == ["frame-1"]
+        assert outputs[0].multimodal_output["audio"] == "audio-0"
+        assert outputs[1].multimodal_output["audio"] == "audio-1"
+        torch.testing.assert_close(
+            outputs[0].multimodal_output["actions"],
+            torch.tensor([1.0, 2.0]),
+        )
+        torch.testing.assert_close(
+            outputs[1].multimodal_output["actions"],
+            torch.tensor([3.0, 4.0]),
+        )
+
 
 class TestStepScheduler:
     def setup_method(self) -> None:

tests/diffusion/test_omni_diffusion_config_master_port.py

@@ -11,26 +11,54 @@ from vllm_omni.diffusion.data import OmniDiffusionConfig
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
+def _get_available_port() -> int:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+        sock.bind(("", 0))
+        return sock.getsockname()[1]
+
+
+def _bind_available_port_below(max_port: int, port_inc: int) -> socket.socket:
+    for _ in range(100):
+        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        try:
+            sock.bind(("", 0))
+            port = sock.getsockname()[1]
+            if port >= max_port - port_inc:
+                sock.close()
+                continue
+            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as next_sock:
+                next_sock.bind(("", port + port_inc))
+            return sock
+        except OSError:
+            sock.close()
+    raise RuntimeError(f"Failed to reserve an available port below {max_port - port_inc}")
+
+
 class TestOmniDiffusionConfigMasterPort:
     def test_honors_master_port_env(self, monkeypatch: pytest.MonkeyPatch) -> None:
-        monkeypatch.setenv("MASTER_PORT", "43210")
+        port = _get_available_port()
+        monkeypatch.setenv("MASTER_PORT", str(port))
         config = OmniDiffusionConfig(model="test")
-        assert config.master_port == 43210
+        assert config.master_port == port
 
     def test_honors_explicit_master_port(self, monkeypatch: pytest.MonkeyPatch) -> None:
         monkeypatch.delenv("MASTER_PORT", raising=False)
-        config = OmniDiffusionConfig(model="test", master_port=40000)
-        assert config.master_port == 40000
+        port = _get_available_port()
+        config = OmniDiffusionConfig(model="test", master_port=port)
+        assert config.master_port == port
 
     def test_master_port_env_takes_precedence_over_kwarg(self, monkeypatch: pytest.MonkeyPatch) -> None:
-        monkeypatch.setenv("MASTER_PORT", "43210")
-        config = OmniDiffusionConfig(model="test", master_port=40000)
-        assert config.master_port == 43210
+        env_port = _get_available_port()
+        kwarg_port = _get_available_port()
+        monkeypatch.setenv("MASTER_PORT", str(env_port))
+        config = OmniDiffusionConfig(model="test", master_port=kwarg_port)
+        assert config.master_port == env_port
 
     def test_explicit_master_port_is_stable_without_jitter(self, monkeypatch: pytest.MonkeyPatch) -> None:
         monkeypatch.delenv("MASTER_PORT", raising=False)
-        ports = {OmniDiffusionConfig(model="test", master_port=40123).master_port for _ in range(5)}
-        assert ports == {40123}
+        port = _get_available_port()
+        ports = {OmniDiffusionConfig(model="test", master_port=port).master_port for _ in range(5)}
+        assert ports == {port}
 
     def test_default_master_port_is_available(self, monkeypatch: pytest.MonkeyPatch) -> None:
         monkeypatch.delenv("MASTER_PORT", raising=False)
@@ -42,8 +70,7 @@ class TestOmniDiffusionConfigMasterPort:
 
     def test_settle_port_steps_when_port_is_busy(self, monkeypatch: pytest.MonkeyPatch) -> None:
         monkeypatch.delenv("MASTER_PORT", raising=False)
-        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
-            sock.bind(("", 0))
+        with _bind_available_port_below(max_port=60000, port_inc=37) as sock:
             busy_port = sock.getsockname()[1]
             config = OmniDiffusionConfig(model="test", master_port=busy_port)
             assert config.master_port == busy_port + 37

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -50,6 +50,7 @@ def build_adapter(monkeypatch, mocker: MockerFixture):
         stage_id: int = 1,
         model_mode: str = "ar",
         max_num_seqs: int = 2,
+        active_stream_window: int = 0,
         connector_extra: dict | None = None,
     ):
         connector = mocker.MagicMock()
@@ -76,7 +77,11 @@ def build_adapter(monkeypatch, mocker: MockerFixture):
             classmethod(lambda cls, _model_config: connector),
         )
 
-        model_config = SimpleNamespace(worker_type=model_mode)
+        model_config = SimpleNamespace(
+            worker_type=model_mode,
+            max_num_seqs=max_num_seqs,
+            active_stream_window=active_stream_window,
+        )
         scheduler_config = SimpleNamespace(max_num_seqs=max_num_seqs)
         adapter = OmniChunkTransferAdapter(
             SimpleNamespace(model_config=model_config, scheduler_config=scheduler_config)
@@ -396,6 +401,75 @@ def test_process_and_restore_queues(build_adapter):
     assert adapter.waiting_for_chunk_running_requests == deque()
 
 
+def test_fifo_promotion(build_adapter):
+    adapter, _ = build_adapter(stage_id=1, max_num_seqs=2, active_stream_window=2)
+    reqs = [_req(f"req-{idx}", RequestStatus.WAITING) for idx in range(1, 5)]
+    waiting_queue = DummyWaitingQueue(reqs)
+    running_queue = []
+
+    adapter.process_pending_chunks(waiting_queue, running_queue)
+
+    assert list(adapter._active_streams) == ["req-1", "req-2"]
+    assert waiting_queue == reqs[2:]
+    assert reqs[0].status == RequestStatus.WAITING_FOR_CHUNK
+    assert reqs[1].status == RequestStatus.WAITING_FOR_CHUNK
+    assert reqs[2].status == RequestStatus.WAITING
+
+    adapter.finished_requests.add("req-1")
+    # Eviction is deferred to postprocess_scheduler_output in the runtime path
+    # (commit c4d95fd9 — otherwise the terminal chunk deadlocks at c=8 K=2).
+    # Simulate it here so promotion can pick up the freed slot.
+    adapter._evict_finished_active_streams({"req-1"})
+    adapter.process_pending_chunks(waiting_queue, running_queue)
+
+    assert list(adapter._active_streams) == ["req-2", "req-3"]
+    # Promotion + chunk processing happen in the same call, so req-3 is
+    # already WAITING_FOR_CHUNK by the time we check.
+    assert reqs[2].status == RequestStatus.WAITING_FOR_CHUNK
+
+
+def test_legacy_k0(build_adapter):
+    adapter, _ = build_adapter(stage_id=1, max_num_seqs=1, active_stream_window=0)
+    waiting_req = _req("waiting", RequestStatus.WAITING)
+    running_req_1 = _req("running-1", RequestStatus.RUNNING)
+    running_req_2 = _req("running-2", RequestStatus.RUNNING)
+    waiting_queue = DummyWaitingQueue([waiting_req])
+    running_queue = [running_req_1, running_req_2]
+
+    adapter.requests_with_ready_chunks.update({"running-1", "running-2"})
+    adapter.process_pending_chunks(waiting_queue, running_queue)
+
+    assert waiting_req.status == RequestStatus.WAITING_FOR_CHUNK
+    assert adapter.waiting_for_chunk_waiting_requests == deque([waiting_req])
+    assert running_queue == [running_req_1]
+    assert waiting_queue == [running_req_2]
+    assert running_req_2.status == RequestStatus.PREEMPTED
+    assert adapter._active_streams == {}
+
+
+def test_finished_releases_slot(build_adapter):
+    adapter, _ = build_adapter(stage_id=1, max_num_seqs=1, active_stream_window=1)
+    req_1 = _req("req-1", RequestStatus.WAITING)
+    req_2 = _req("req-2", RequestStatus.WAITING)
+    waiting_queue = DummyWaitingQueue([req_1, req_2])
+    running_queue = []
+
+    adapter.process_pending_chunks(waiting_queue, running_queue)
+    assert list(adapter._active_streams) == ["req-1"]
+    assert waiting_queue == [req_2]
+
+    adapter.finished_requests.add("req-1")
+    # Eviction is deferred to postprocess_scheduler_output in the runtime path
+    # (commit c4d95fd9). Simulate it so promotion can pick up the freed slot.
+    adapter._evict_finished_active_streams({"req-1"})
+    adapter.process_pending_chunks(waiting_queue, running_queue)
+
+    assert list(adapter._active_streams) == ["req-2"]
+    # Promotion + chunk processing happen in the same call, so req-2 is
+    # already WAITING_FOR_CHUNK by the time we check.
+    assert req_2.status == RequestStatus.WAITING_FOR_CHUNK
+
+
 def test_postprocess_scheduler_output(build_adapter):
     adapter, _ = build_adapter()
     adapter.requests_with_ready_chunks = {"new-ready", "cached-ready", "leftover"}
@@ -422,6 +496,7 @@ def test_postprocess_scheduler_output(build_adapter):
 def _populate_adapter_state(adapter, req_id="req-1", ext_id="ext-1"):
     """Fill every per-request structure so cleanup can be verified."""
     adapter.finished_requests.add(req_id)
+    adapter._active_streams[req_id] = SimpleNamespace(request_id=req_id)
     adapter.get_req_chunk[req_id] = 3
     adapter.requests_with_ready_chunks.add(req_id)
     adapter.request_ids_mapping[req_id] = ext_id
@@ -442,6 +517,7 @@ def test_cleanup_clears_all_state(build_adapter):
     adapter.cleanup(req_id, ext_id)
 
     assert req_id not in adapter.finished_requests
+    assert req_id not in adapter._active_streams
     assert req_id not in adapter.get_req_chunk
     assert req_id not in adapter.requests_with_ready_chunks
     assert req_id not in adapter.request_ids_mapping

tests/dreamzero/openpi_client_helper.py

@@ -0,0 +1,307 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import uuid
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+
+try:
+    import cv2
+except ImportError:  # pragma: no cover - optional e2e dependency
+    cv2 = None
+
+try:
+    import websockets.sync.client as websockets_client
+except ImportError:  # pragma: no cover - optional e2e dependency
+    websockets_client = None
+
+try:
+    from openpi_client import msgpack_numpy
+except ImportError:  # pragma: no cover - optional e2e dependency
+    msgpack_numpy = None
+
+PING_INTERVAL_SECS = 300
+PING_TIMEOUT_SECS = 3600
+DEFAULT_HOST = "127.0.0.1"
+DEFAULT_PORT = 8000
+DEFAULT_PATH = "/v1/realtime/robot/openpi"
+DEFAULT_PROMPT = "Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan"
+ACTION_HORIZON = 24
+DEFAULT_ACTION_DIM = 8
+RELATIVE_OFFSETS = [-23, -16, -8, 0]
+CAMERA_FILES = {
+    "observation/exterior_image_0_left": "exterior_image_1_left.mp4",
+    "observation/exterior_image_1_left": "exterior_image_2_left.mp4",
+    "observation/wrist_image_left": "wrist_image_left.mp4",
+}
+
+
+def _decode_action_response(response: bytes | str) -> np.ndarray:
+    if isinstance(response, str):
+        raise RuntimeError(f"Inference failed: {response}")
+    decoded = msgpack_numpy.unpackb(response)
+    if isinstance(decoded, dict) and decoded.get("type") == "error":
+        message = decoded.get("message", decoded)
+        raise RuntimeError(f"Inference failed: {message}")
+    return np.asarray(decoded, dtype=np.float32)
+
+
+def require_dependencies() -> None:
+    missing = []
+    if cv2 is None:
+        missing.append("opencv-python")
+    if websockets_client is None:
+        missing.append("websockets")
+    if msgpack_numpy is None:
+        missing.append("openpi-client")
+    if missing:
+        raise ModuleNotFoundError(f"DreamZero OpenPI test dependencies are missing: {', '.join(missing)}")
+
+
+@dataclass(frozen=True)
+class DreamZeroServerMetadata:
+    image_resolution: tuple[int, int]
+    n_external_cameras: int
+    needs_wrist_camera: bool
+    needs_stereo_camera: bool
+    needs_session_id: bool
+    action_space: str
+
+    @classmethod
+    def from_dict(cls, payload: dict[str, Any]) -> DreamZeroServerMetadata:
+        required_keys = (
+            "image_resolution",
+            "n_external_cameras",
+            "needs_wrist_camera",
+            "needs_stereo_camera",
+            "needs_session_id",
+            "action_space",
+        )
+        missing_keys = [key for key in required_keys if key not in payload]
+        if missing_keys:
+            raise ValueError(f"Missing DreamZero metadata keys: {missing_keys}")
+
+        image_resolution = payload["image_resolution"]
+        if not isinstance(image_resolution, (list, tuple)) or len(image_resolution) != 2:
+            raise ValueError(f"Invalid image_resolution: {image_resolution!r}")
+
+        return cls(
+            image_resolution=(int(image_resolution[0]), int(image_resolution[1])),
+            n_external_cameras=int(payload["n_external_cameras"]),
+            needs_wrist_camera=bool(payload["needs_wrist_camera"]),
+            needs_stereo_camera=bool(payload["needs_stereo_camera"]),
+            needs_session_id=bool(payload["needs_session_id"]),
+            action_space=str(payload["action_space"]),
+        )
+
+
+class OpenPIWebsocketClient:
+    def __init__(
+        self,
+        *,
+        host: str = DEFAULT_HOST,
+        port: int = DEFAULT_PORT,
+        path: str = DEFAULT_PATH,
+    ) -> None:
+        require_dependencies()
+        self._uri = f"ws://{host}:{port}{path}"
+        self._packer = msgpack_numpy.Packer()
+        self._ws, self._server_metadata = self._connect()
+
+    def _connect(self):
+        conn = websockets_client.connect(
+            self._uri,
+            compression=None,
+            max_size=None,
+            ping_interval=PING_INTERVAL_SECS,
+            ping_timeout=PING_TIMEOUT_SECS,
+        )
+        metadata = msgpack_numpy.unpackb(conn.recv())
+        if not isinstance(metadata, dict):
+            raise TypeError(f"Expected dict metadata from server, got {type(metadata)!r}")
+        return conn, metadata
+
+    def get_server_metadata(self) -> dict[str, Any]:
+        return dict(self._server_metadata)
+
+    def infer(self, obs: dict[str, Any]) -> np.ndarray:
+        payload = dict(obs)
+        payload["endpoint"] = "infer"
+        self._ws.send(self._packer.pack(payload))
+        response = self._ws.recv()
+        return _decode_action_response(response)
+
+    def reset(self, reset_info: dict[str, Any] | None = None) -> str:
+        payload = dict(reset_info or {})
+        payload["endpoint"] = "reset"
+        self._ws.send(self._packer.pack(payload))
+        response = self._ws.recv()
+        if isinstance(response, str):
+            return response
+        decoded = msgpack_numpy.unpackb(response)
+        if not isinstance(decoded, dict) or decoded.get("status") != "reset successful":
+            raise RuntimeError(f"Unexpected reset response: {decoded!r}")
+        return str(decoded["status"])
+
+    def close(self) -> None:
+        self._ws.close()
+
+
+def load_all_frames(video_path: Path) -> np.ndarray:
+    require_dependencies()
+    cap = cv2.VideoCapture(str(video_path))
+    frames = []
+    while True:
+        ok, frame = cap.read()
+        if not ok:
+            break
+        frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+    cap.release()
+    if not frames:
+        raise RuntimeError(f"No frames loaded from {video_path}")
+    return np.stack(frames, axis=0)
+
+
+def load_camera_frames(video_dir: Path) -> dict[str, np.ndarray]:
+    camera_frames: dict[str, np.ndarray] = {}
+    for camera_key, file_name in CAMERA_FILES.items():
+        video_path = video_dir / file_name
+        if not video_path.exists():
+            raise FileNotFoundError(f"Missing DreamZero test asset: {video_path}")
+        camera_frames[camera_key] = load_all_frames(video_path)
+    return camera_frames
+
+
+def build_frame_schedule(total_frames: int, num_chunks: int) -> list[list[int]]:
+    chunks: list[list[int]] = []
+    current_frame = 23
+    for _ in range(num_chunks):
+        indices = [max(current_frame + offset, 0) for offset in RELATIVE_OFFSETS]
+        if indices[-1] >= total_frames:
+            break
+        chunks.append(indices)
+        current_frame += ACTION_HORIZON
+    return chunks
+
+
+def make_obs_from_video(
+    camera_frames: dict[str, np.ndarray],
+    frame_indices: list[int],
+    *,
+    prompt: str,
+    session_id: str,
+) -> dict[str, Any]:
+    obs: dict[str, Any] = {}
+    for camera_key, all_frames in camera_frames.items():
+        selected = all_frames[frame_indices]
+        obs[camera_key] = selected[0] if len(frame_indices) == 1 else selected
+
+    obs["observation/joint_position"] = np.zeros(7, dtype=np.float32)
+    obs["observation/cartesian_position"] = np.zeros(6, dtype=np.float32)
+    obs["observation/gripper_position"] = np.zeros(1, dtype=np.float32)
+    obs["prompt"] = prompt
+    obs["session_id"] = session_id
+    return obs
+
+
+def build_demo_observations(
+    camera_frames: dict[str, np.ndarray],
+    *,
+    prompt: str,
+    session_id: str,
+    num_chunks: int = 2,
+) -> list[dict[str, Any]]:
+    if num_chunks < 1:
+        raise ValueError("num_chunks must be at least 1")
+
+    total_frames = min(frames.shape[0] for frames in camera_frames.values())
+    observations = [
+        make_obs_from_video(
+            camera_frames,
+            [0],
+            prompt=prompt,
+            session_id=session_id,
+        )
+    ]
+    for indices in build_frame_schedule(total_frames, num_chunks - 1):
+        observations.append(
+            make_obs_from_video(
+                camera_frames,
+                indices,
+                prompt=prompt,
+                session_id=session_id,
+            )
+        )
+    return observations
+
+
+def validate_session_result(
+    result: dict[str, Any],
+    *,
+    expected_action_horizon: int = ACTION_HORIZON,
+    expected_action_dim: int = DEFAULT_ACTION_DIM,
+) -> None:
+    metadata = DreamZeroServerMetadata.from_dict(result["metadata"])
+    if metadata.image_resolution != (180, 320):
+        raise AssertionError(f"Unexpected image_resolution: {metadata.image_resolution}")
+    if metadata.n_external_cameras != 2:
+        raise AssertionError(f"Unexpected n_external_cameras: {metadata.n_external_cameras}")
+    if not metadata.needs_wrist_camera:
+        raise AssertionError("DreamZero test expects wrist camera metadata")
+    if metadata.action_space != "joint_position":
+        raise AssertionError(f"Unexpected action_space: {metadata.action_space}")
+
+    actions = result["actions"]
+    if len(actions) != 3:
+        raise AssertionError(f"Expected 3 action tensors, got {len(actions)}")
+    for index, action in enumerate(actions):
+        if action.shape != (expected_action_horizon, expected_action_dim):
+            raise AssertionError(
+                f"Action {index} shape mismatch: expected "
+                f"{(expected_action_horizon, expected_action_dim)}, got {action.shape}"
+            )
+        if not np.isfinite(action).all():
+            raise AssertionError(f"Action {index} contains non-finite values")
+
+    if result["reset_status"] != "reset successful":
+        raise AssertionError(f"Unexpected reset status: {result['reset_status']!r}")
+
+
+def run_policy_session(
+    *,
+    host: str = DEFAULT_HOST,
+    port: int = DEFAULT_PORT,
+    path: str = DEFAULT_PATH,
+    video_dir: Path,
+    prompt: str = DEFAULT_PROMPT,
+    session_id: str | None = None,
+    num_chunks: int = 2,
+) -> dict[str, Any]:
+    session_id = session_id or str(uuid.uuid4())
+    camera_frames = load_camera_frames(video_dir)
+    observations = build_demo_observations(
+        camera_frames,
+        prompt=prompt,
+        session_id=session_id,
+        num_chunks=num_chunks,
+    )
+
+    client = OpenPIWebsocketClient(host=host, port=port, path=path)
+    try:
+        metadata = client.get_server_metadata()
+        actions = [client.infer(obs) for obs in observations]
+        reset_status = client.reset({})
+        actions.append(client.infer(observations[0]))
+        return {
+            "metadata": metadata,
+            "actions": actions,
+            "reset_status": reset_status,
+            "session_id": session_id,
+        }
+    finally:
+        client.close()

tests/dreamzero/test_dreamzero_utils.py

@@ -0,0 +1,24 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+
+from vllm_omni.diffusion.models.dreamzero.utils import (
+    DEFAULT_CFG_SCALE,
+    DEFAULT_EMBODIMENT_NAME_TO_ID,
+    DEFAULT_NEGATIVE_PROMPT,
+    DEFAULT_NUM_INFERENCE_STEPS,
+    DEFAULT_SEED,
+    DEFAULT_SIGMA_SHIFT,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_dreamzero_default_constants_match_source_baseline():
+    assert DEFAULT_NUM_INFERENCE_STEPS == 16
+    assert DEFAULT_CFG_SCALE == 5.0
+    assert DEFAULT_SIGMA_SHIFT == 5.0
+    assert DEFAULT_SEED == 1140
+    assert "worst quality" in DEFAULT_NEGATIVE_PROMPT
+    assert DEFAULT_EMBODIMENT_NAME_TO_ID["oxe_droid"] == 17

tests/dreamzero/test_openpi_client_helper.py

@@ -0,0 +1,46 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import numpy as np
+import pytest
+
+from tests.dreamzero import openpi_client_helper
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class FakeMsgpackNumpy:
+    payloads = {}
+
+    @staticmethod
+    def packb(obj):
+        key = str(len(FakeMsgpackNumpy.payloads)).encode()
+        FakeMsgpackNumpy.payloads[key] = obj
+        return key
+
+    @staticmethod
+    def unpackb(data):
+        return FakeMsgpackNumpy.payloads[data]
+
+
+def test_decode_action_response_surfaces_structured_error(monkeypatch):
+    monkeypatch.setattr(openpi_client_helper, "msgpack_numpy", FakeMsgpackNumpy)
+    payload = FakeMsgpackNumpy.packb(
+        {
+            "type": "error",
+            "message": "Internal inference error",
+        }
+    )
+
+    with pytest.raises(RuntimeError, match="Internal inference error"):
+        openpi_client_helper._decode_action_response(payload)
+
+
+def test_decode_action_response_converts_action_payload_to_float32(monkeypatch):
+    monkeypatch.setattr(openpi_client_helper, "msgpack_numpy", FakeMsgpackNumpy)
+    payload = FakeMsgpackNumpy.packb(np.asarray([[1.0, 2.0]], dtype=np.float64))
+
+    actions = openpi_client_helper._decode_action_response(payload)
+
+    assert actions.dtype == np.float32
+    np.testing.assert_array_equal(actions, np.asarray([[1.0, 2.0]], dtype=np.float32))

tests/dreamzero/test_pipeline_state.py

@@ -0,0 +1,59 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections import OrderedDict
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
+from vllm_omni.diffusion.models.dreamzero.state_dreamzero import DreamZeroState
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _empty_pipeline() -> DreamZeroPipeline:
+    pipeline = DreamZeroPipeline.__new__(DreamZeroPipeline)
+    pipeline._states = OrderedDict()
+    pipeline._max_session_states = 2
+    return pipeline
+
+
+def test_dreamzero_pipeline_state_is_session_keyed() -> None:
+    pipeline = _empty_pipeline()
+
+    session_a = pipeline._get_or_create_state("session-a")
+    session_b = pipeline._get_or_create_state("session-b")
+    session_a.call_count = 7
+    session_b.call_count = 3
+
+    assert pipeline._get_or_create_state("session-a") is session_a
+    assert pipeline._get_or_create_state("session-b") is session_b
+    assert session_a.call_count == 7
+    assert session_b.call_count == 3
+
+
+def test_dreamzero_pipeline_state_lru_caps_retained_sessions() -> None:
+    pipeline = _empty_pipeline()
+
+    session_a = pipeline._get_or_create_state("session-a")
+    pipeline._get_or_create_state("session-b")
+    assert pipeline._get_or_create_state("session-a") is session_a
+
+    pipeline._get_or_create_state("session-c")
+
+    assert list(pipeline._states) == ["session-a", "session-c"]
+    assert "session-b" not in pipeline._states
+
+
+def test_dreamzero_state_cache_access_requires_initialization() -> None:
+    state = DreamZeroState()
+
+    with pytest.raises(RuntimeError, match="KV caches not initialized"):
+        state.get_kv_caches()
+
+    with pytest.raises(RuntimeError, match="Cross-attn caches not initialized"):
+        state.get_crossattn_caches()
+
+    with pytest.raises(RuntimeError, match="create_kv_caches first"):
+        state.update_kv_cache(0, torch.empty(0))

tests/dreamzero/upstream/openpi_test_client_ar.py

@@ -0,0 +1,381 @@
+#!/usr/bin/env python3
+"""
+Copied from the DreamZero repository's `test_client_AR.py`:
+https://github.com/dreamzero0/dreamzero/blob/main/test_client_AR.py
+
+Kept here for end-to-end compatibility / parity testing against the vLLM
+OpenPI server.
+
+Test client for AR_droid policy server using roboarena interface.
+
+Sends real video frames from debug_image/ directory instead of zero dummy images.
+
+Frame schedule (matching debug_inference.py):
+  - Step 0 (initial): send frame [0]             (1 frame, H W 3)
+  - Step 1: send frames [0, 7, 15, 23]           (4 frames, 4 H W 3)
+  - Step 2: send frames [24, 31, 39, 47]         (4 frames)
+  - Step 3: send frames [48, 55, 63, 71]         (4 frames)
+  - ...
+
+Expected server configuration:
+    - image_resolution: (180, 320)
+    - n_external_cameras: 2
+    - needs_wrist_camera: True
+    - action_space: "joint_position"
+
+Usage:
+    # Start server with roboarena interface:
+    torchrun --nproc_per_node=8 socket_test_optimized_AR.py --port 8000
+
+    # Run against the original DreamZero websocket server (default path):
+    python test_client_AR.py --host <server_host> --port 8000
+
+    # Run against the vLLM OpenPI server:
+    python test_client_AR.py --host <server_host> --port 8000 --path /v1/realtime/robot/openpi
+
+    # Use zero images instead of real video (old behavior):
+    python test_client_AR.py --host <server_host> --port 8000 --use-zero-images
+"""
+
+import argparse
+import logging
+import os
+import sys
+import time
+import uuid
+from pathlib import Path
+
+import cv2
+import numpy as np
+from openpi_client import msgpack_numpy
+
+_DREAMZERO_REPO_ENV = os.environ.get("DREAMZERO_REPO")
+DREAMZERO_REPO = Path(_DREAMZERO_REPO_ENV).expanduser() if _DREAMZERO_REPO_ENV else None
+
+
+def _import_upstream_policy_modules():
+    if DREAMZERO_REPO is None:
+        raise ImportError("Set DREAMZERO_REPO to an upstream DreamZero checkout before using this helper.")
+    if DREAMZERO_REPO.exists() and str(DREAMZERO_REPO) not in sys.path:
+        sys.path.insert(0, str(DREAMZERO_REPO))
+
+    import eval_utils.policy_server as policy_server
+    from eval_utils.policy_client import WebsocketClientPolicy
+
+    return policy_server, WebsocketClientPolicy
+
+
+policy_server, WebsocketClientPolicy = _import_upstream_policy_modules()
+
+VIDEO_DIR = os.environ.get(
+    "DREAMZERO_VIDEO_DIR",
+    str(DREAMZERO_REPO / "debug_image") if DREAMZERO_REPO is not None else "debug_image",
+)
+
+# roboarena key -> video filename
+CAMERA_FILES = {
+    "observation/exterior_image_0_left": "exterior_image_1_left.mp4",
+    "observation/exterior_image_1_left": "exterior_image_2_left.mp4",
+    "observation/wrist_image_left": "wrist_image_left.mp4",
+}
+
+# Frame schedule constants (matching debug_inference.py)
+RELATIVE_OFFSETS = [-23, -16, -8, 0]
+ACTION_HORIZON = 24
+DEFAULT_WEBSOCKET_PATH = ""
+
+
+class OpenPIWebsocketClientPolicy(WebsocketClientPolicy):
+    """DreamZero websocket client with a configurable path suffix.
+
+    The original DreamZero client connects to ``ws://host:port``.
+    vLLM serves the compatible robot policy endpoint at
+    ``/v1/realtime/robot/openpi`` when ``path`` is set accordingly.
+    """
+
+    def __init__(
+        self,
+        host: str = "0.0.0.0",
+        port: int = 8000,
+        path: str = DEFAULT_WEBSOCKET_PATH,
+    ) -> None:
+        self._uri = f"ws://{host}:{port}{path}"
+        self._packer = msgpack_numpy.Packer()
+        self._ws, self._server_metadata = self._wait_for_server()
+
+
+def load_all_frames(video_path: str) -> np.ndarray:
+    """Load all frames from a video file. Returns (N, H, W, 3) uint8 array (RGB)."""
+    cap = cv2.VideoCapture(video_path)
+    frames = []
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            break
+        frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+    cap.release()
+    if not frames:
+        raise RuntimeError(f"No frames loaded from {video_path}")
+    return np.stack(frames, axis=0)
+
+
+def load_camera_frames() -> dict[str, np.ndarray]:
+    """Load all video frames for each camera from the debug_image/ directory.
+
+    Returns:
+        Dict mapping roboarena camera keys to (N, H, W, 3) uint8 arrays.
+    """
+    camera_frames: dict[str, np.ndarray] = {}
+    for cam_key, fname in CAMERA_FILES.items():
+        path = os.path.join(VIDEO_DIR, fname)
+        camera_frames[cam_key] = load_all_frames(path)
+        logging.info(f"Loaded {cam_key}: {camera_frames[cam_key].shape}")
+    return camera_frames
+
+
+def build_frame_schedule(total_frames: int, num_chunks: int) -> list[list[int]]:
+    """Build the frame index schedule for multi-frame chunks.
+
+    Returns a list of frame-index lists. Each inner list has 4 indices.
+    """
+    chunks: list[list[int]] = []
+    current_frame = 23  # first anchor frame
+    for _ in range(num_chunks):
+        indices = [max(current_frame + off, 0) for off in RELATIVE_OFFSETS]
+        if indices[-1] >= total_frames:
+            logging.info(f"Frame {indices[-1]} >= {total_frames}, stopping at {len(chunks)} chunks")
+            break
+        chunks.append(indices)
+        current_frame += ACTION_HORIZON
+    return chunks
+
+
+def _make_obs_from_video(
+    camera_frames: dict[str, np.ndarray],
+    frame_indices: list[int],
+    prompt: str,
+    session_id: str,
+) -> dict:
+    """Build an observation dict from real video frames.
+
+    For 1 frame: each image key is (H, W, 3).
+    For 4 frames: each image key is (4, H, W, 3).
+    """
+    obs: dict = {}
+    for cam_key, all_frames in camera_frames.items():
+        selected = all_frames[frame_indices]  # (T, H, W, 3)
+        if len(frame_indices) == 1:
+            selected = selected[0]  # (H, W, 3)
+        obs[cam_key] = selected
+
+    obs["observation/joint_position"] = np.zeros(7, dtype=np.float32)
+    obs["observation/cartesian_position"] = np.zeros(6, dtype=np.float32)
+    obs["observation/gripper_position"] = np.zeros(1, dtype=np.float32)
+    obs["prompt"] = prompt
+    obs["session_id"] = session_id
+    return obs
+
+
+def _make_zero_observation(
+    server_config: policy_server.PolicyServerConfig,
+    prompt: str = "pick up the object",
+    session_id: str | None = None,
+) -> dict:
+    """Create a dummy observation matching AR_droid expectations.
+
+    AR_droid expects:
+        - 2 external cameras (exterior_image_0_left, exterior_image_1_left)
+        - 1 wrist camera (wrist_image_left)
+        - Image resolution: 180x320 (H x W)
+        - joint_position: 7 DoF
+        - gripper_position: 1 DoF
+    """
+    obs = {}
+
+    # Determine image resolution
+    if server_config.image_resolution is not None:
+        h, w = server_config.image_resolution
+    else:
+        # Default for AR_droid
+        h, w = 180, 320
+
+    # External cameras (0-indexed in roboarena)
+    for i in range(server_config.n_external_cameras):
+        obs[f"observation/exterior_image_{i}_left"] = np.zeros((h, w, 3), dtype=np.uint8)
+        if server_config.needs_stereo_camera:
+            obs[f"observation/exterior_image_{i}_right"] = np.zeros((h, w, 3), dtype=np.uint8)
+
+    # Wrist camera
+    if server_config.needs_wrist_camera:
+        obs["observation/wrist_image_left"] = np.zeros((h, w, 3), dtype=np.uint8)
+        if server_config.needs_stereo_camera:
+            obs["observation/wrist_image_right"] = np.zeros((h, w, 3), dtype=np.uint8)
+
+    # Session ID - should be passed in to ensure consistency within a session
+    if server_config.needs_session_id:
+        import uuid
+
+        # Generate unique session ID if not provided
+        obs["session_id"] = session_id if session_id else str(uuid.uuid4())
+
+    # State observations (AR_droid: 7 DoF arm + 1 gripper)
+    obs["observation/joint_position"] = np.zeros(7, dtype=np.float32)
+    obs["observation/cartesian_position"] = np.zeros(6, dtype=np.float32)
+    obs["observation/gripper_position"] = np.zeros(1, dtype=np.float32)
+
+    # Language prompt
+    obs["prompt"] = prompt
+
+    return obs
+
+
+def test_ar_droid_policy_server(
+    host: str = "localhost",
+    port: int = 8000,
+    path: str = DEFAULT_WEBSOCKET_PATH,
+    num_chunks: int = 15,
+    prompt: str = "Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan",
+    use_zero_images: bool = False,
+):
+    """Test the AR_droid policy server with roboarena interface.
+
+    When use_zero_images is False (default), loads real video frames from
+    debug_image/ and follows the frame schedule from debug_inference.py.
+    """
+    logging.info(f"Connecting to AR_droid server at ws://{host}:{port}{path} ...")
+
+    client = OpenPIWebsocketClientPolicy(host=host, port=port, path=path)
+
+    # Validate server metadata
+    metadata = client.get_server_metadata()
+    logging.info(f"Server metadata: {metadata}")
+    assert isinstance(metadata, dict), "Metadata should be a dict"
+
+    try:
+        server_config = policy_server.PolicyServerConfig(**metadata)
+    except Exception as e:
+        logging.error(f"Error parsing metadata: {e}")
+        raise e
+
+    # Validate expected AR_droid configuration
+    logging.info(f"Server config: {server_config}")
+    assert server_config.n_external_cameras == 2, f"Expected 2 external cameras, got {server_config.n_external_cameras}"
+    assert server_config.needs_wrist_camera, "Expected wrist camera to be enabled"
+    assert server_config.action_space == "joint_position", (
+        f"Expected joint_position action space, got {server_config.action_space}"
+    )
+
+    logging.info("Server configuration validated for AR_droid")
+
+    # Generate unique session ID for this test run
+    session_id = str(uuid.uuid4())
+    logging.info(f"Session ID: {session_id}")
+
+    # ── Zero-image fallback mode ──────────────────────────────────────
+    if use_zero_images:
+        logging.info("Using ZERO dummy images (legacy mode)")
+        for i in range(num_chunks):
+            obs = _make_zero_observation(server_config, prompt=prompt, session_id=session_id)
+            logging.info(f"Inference {i + 1}/{num_chunks}: prompt='{prompt}'")
+            t0 = time.time()
+            actions = client.infer(obs)
+            dt = time.time() - t0
+            _log_action(actions, dt)
+
+        logging.info("Sending reset...")
+        client.reset({})
+        logging.info("Done (zero-image mode).")
+        return
+
+    # ── Real video frame mode ─────────────────────────────────────────
+    logging.info("Loading real video frames from debug_image/ directory")
+    camera_frames = load_camera_frames()
+
+    total_frames = min(v.shape[0] for v in camera_frames.values())
+    logging.info(f"Total frames available: {total_frames}")
+
+    # Build frame schedule
+    chunks = build_frame_schedule(total_frames, num_chunks)
+
+    logging.info("Frame schedule:")
+    logging.info("  Initial: [0]")
+    for i, indices in enumerate(chunks):
+        logging.info(f"  Chunk {i}: {indices}")
+
+    # Step 0: initial single frame
+    logging.info("=== Initial: frame [0] ===")
+    obs = _make_obs_from_video(camera_frames, [0], prompt, session_id)
+    t0 = time.time()
+    actions = client.infer(obs)
+    dt = time.time() - t0
+    _log_action(actions, dt)
+
+    # Subsequent chunks: send 4 frames at a time
+    for chunk_idx, frame_indices in enumerate(chunks):
+        logging.info(f"=== Chunk {chunk_idx}: frames {frame_indices} ===")
+        obs = _make_obs_from_video(camera_frames, frame_indices, prompt, session_id)
+        t0 = time.time()
+        actions = client.infer(obs)
+        dt = time.time() - t0
+        _log_action(actions, dt)
+
+    # Reset triggers video save on the server
+    logging.info("Sending reset to save video...")
+    client.reset({})
+
+    logging.info("Done.")
+
+
+def _log_action(actions: np.ndarray, dt: float) -> None:
+    """Pretty-print action shape, range, and timing."""
+    assert isinstance(actions, np.ndarray), f"Expected numpy array, got {type(actions)}"
+    assert actions.ndim == 2, f"Expected 2D array, got shape {actions.shape}"
+    assert actions.shape[-1] == 8, f"Expected 8 action dims (7 joints + 1 gripper), got {actions.shape[-1]}"
+    logging.info(f"  Action shape: {actions.shape}, range: [{actions.min():.4f}, {actions.max():.4f}], time: {dt:.2f}s")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Test AR_droid policy server with real video frames from debug_image/")
+    parser.add_argument("--host", default="localhost", help="Server hostname")
+    parser.add_argument("--port", type=int, default=8000, help="Server port")
+    parser.add_argument(
+        "--path",
+        default=DEFAULT_WEBSOCKET_PATH,
+        help="WebSocket path suffix (default: empty string for the original DreamZero server)",
+    )
+    parser.add_argument(
+        "--num-chunks",
+        type=int,
+        default=15,
+        help="Number of 4-frame chunks to send after the initial frame (default: 15)",
+    )
+    parser.add_argument(
+        "--prompt",
+        default="Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan",
+        help="Language prompt for the policy",
+    )
+    parser.add_argument(
+        "--use-zero-images",
+        action="store_true",
+        help="Use zero dummy images instead of real video frames (legacy mode)",
+    )
+
+    args = parser.parse_args()
+
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s [%(levelname)s] %(message)s",
+    )
+
+    test_ar_droid_policy_server(
+        host=args.host,
+        port=args.port,
+        path=args.path,
+        num_chunks=args.num_chunks,
+        prompt=args.prompt,
+        use_zero_images=args.use_zero_images,
+    )
+
+
+if __name__ == "__main__":
+    main()

tests/dreamzero/upstream/test_openpi_e2e_source_parity.py

@@ -0,0 +1,358 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Formal OpenPI end-to-end parity: upstream DreamZero server vs `vllm serve`.
+
+This test uses DreamZero's own client-side observation builders from
+`${DREAMZERO_REPO}/test_client_AR.py`, and client-side websocket protocol from
+`${DREAMZERO_REPO}/eval_utils/policy_client.py`.
+
+The only client-side adaptation for vLLM is the websocket path:
+DreamZero's upstream server serves at `/`, while vLLM serves OpenPI at
+`/v1/realtime/robot/openpi`.
+
+Current scope for this test:
+- default two-GPU run (`nproc_per_node=2` on upstream, `--cfg-parallel-size 2` on `vllm serve`)
+- non-`torch.compile` (upstream launched through
+  `upstream_socket_server_no_compile.py`, vLLM with `--enforce-eager`)
+- non-DiT-cache / non-skip-schedule (`NUM_DIT_STEPS=16`)
+
+Serving contract locked by this test:
+- upstream DreamZero still boots from the local checkpoint directory
+- vLLM boots from the official DreamZero HF repo name (`GEAR-Dreams/DreamZero-DROID`)
+  rather than a prepared local bundle path
+"""
+
+from __future__ import annotations
+
+import os
+import shutil
+import subprocess
+import sys
+import time
+from pathlib import Path
+
+import numpy as np
+import pytest
+import torch
+
+from tests.helpers.runtime import get_open_port
+
+msgpack_numpy = pytest.importorskip("openpi_client.msgpack_numpy")
+
+_DREAMZERO_REPO_ENV = os.environ.get("DREAMZERO_REPO")
+DREAMZERO_REPO = Path(_DREAMZERO_REPO_ENV).expanduser() if _DREAMZERO_REPO_ENV else None
+if DREAMZERO_REPO is not None and str(DREAMZERO_REPO) not in sys.path:
+    sys.path.insert(0, str(DREAMZERO_REPO))
+
+try:
+    import test_client_AR as dreamzero_client
+    from eval_utils.policy_client import WebsocketClientPolicy
+except Exception:  # pragma: no cover - guarded by pytest skip below
+    dreamzero_client = None
+    WebsocketClientPolicy = None
+_BaseWebsocketClientPolicy = WebsocketClientPolicy if WebsocketClientPolicy is not None else object
+
+CHECKPOINT_DIR = DREAMZERO_REPO / "checkpoints" / "dreamzero" if DREAMZERO_REPO is not None else None
+VLLM_MODEL = os.environ.get("VLLM_DREAMZERO_MODEL", "GEAR-Dreams/DreamZero-DROID")
+SERVICE_READY_TIMEOUT_S = int(os.environ.get("OPENPI_SERVICE_READY_TIMEOUT_S", "900"))
+PROMPT = "Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan"
+SESSION_ID = "openpi-e2e-parity-session"
+
+pytestmark = [
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="GPU required"),
+    pytest.mark.skipif(
+        dreamzero_client is None or WebsocketClientPolicy is None,
+        reason="DreamZero client modules are required on PYTHONPATH",
+    ),
+    pytest.mark.skipif(
+        DREAMZERO_REPO is None or not DREAMZERO_REPO.exists(),
+        reason="DreamZero source repo is required at DREAMZERO_REPO",
+    ),
+    pytest.mark.skipif(
+        CHECKPOINT_DIR is None or not CHECKPOINT_DIR.exists(), reason="DreamZero local checkpoint is required"
+    ),
+]
+
+
+class OpenPIWebsocketClientPolicy(_BaseWebsocketClientPolicy):
+    """DreamZero client protocol with an OpenPI websocket path suffix."""
+
+    def __init__(
+        self,
+        host: str = "127.0.0.1",
+        port: int = 8000,
+        path: str = "/v1/realtime/robot/openpi",
+    ) -> None:
+        self._uri = f"ws://{host}:{port}{path}"
+        self._packer = msgpack_numpy.Packer()
+        self._ws, self._server_metadata = self._wait_for_server()
+
+
+def _vllm_executable() -> str:
+    fallback = Path(sys.executable).with_name("vllm")
+    if fallback.exists():
+        return str(fallback)
+    exe = shutil.which("vllm")
+    if exe:
+        return exe
+    raise FileNotFoundError("Unable to locate `vllm` executable in current environment.")
+
+
+def _cfg_parallel_size() -> int:
+    return int(os.environ.get("OPENPI_E2E_CFG_PARALLEL_SIZE", "2"))
+
+
+def _pick_test_gpus() -> list[str]:
+    cfg_parallel_size = _cfg_parallel_size()
+    override = os.environ.get("OPENPI_E2E_GPUS") or os.environ.get("OPENPI_E2E_GPU")
+    if override is not None:
+        gpus = [part.strip() for part in override.split(",") if part.strip()]
+        if not gpus:
+            raise ValueError("OPENPI_E2E_GPUS is set but empty.")
+        if len(gpus) < cfg_parallel_size:
+            raise RuntimeError(f"Need {cfg_parallel_size} GPUs, but OPENPI_E2E_GPUS only provided {gpus}.")
+        return gpus
+
+    query = subprocess.check_output(
+        [
+            "nvidia-smi",
+            "--query-gpu=index,memory.used",
+            "--format=csv,noheader,nounits",
+        ],
+        text=True,
+    )
+    gpu_rows = []
+    for line in query.strip().splitlines():
+        gpu_index, used_mb = [part.strip() for part in line.split(",", maxsplit=1)]
+        gpu_rows.append((int(used_mb), gpu_index))
+    gpu_rows.sort()
+    gpus = [gpu_index for _, gpu_index in gpu_rows[: max(cfg_parallel_size, 1)]]
+    if len(gpus) < cfg_parallel_size:
+        raise RuntimeError(
+            f"Need {cfg_parallel_size} GPUs for cfg_parallel_size={cfg_parallel_size}, "
+            f"but found only {len(gpus)} candidates."
+        )
+    return gpus
+
+
+def _torchrun_argv(script: str, port: int) -> list[str]:
+    return [
+        sys.executable,
+        "-m",
+        "torch.distributed.run",
+        "--standalone",
+        "--nproc_per_node",
+        str(_cfg_parallel_size()),
+        script,
+        "--port",
+        str(port),
+        "--model_path",
+        str(CHECKPOINT_DIR),
+    ]
+
+
+def _run_upstream_service(port: int, log_path: Path) -> subprocess.Popen[str]:
+    env = os.environ.copy()
+    env.setdefault("PYTHONPATH", "")
+    env["PYTHONPATH"] = f"{Path.cwd()}:{DREAMZERO_REPO}:{env['PYTHONPATH']}".rstrip(":")
+    env["CUDA_VISIBLE_DEVICES"] = ",".join(_pick_test_gpus())
+    env.setdefault("NO_ALBUMENTATIONS_UPDATE", "1")
+    env["ATTENTION_BACKEND"] = "torch"
+    env.setdefault("ENABLE_TENSORRT", "false")
+    env["ENABLE_DIT_CACHE"] = "false"
+    env["NUM_DIT_STEPS"] = "16"
+    env["DYNAMIC_CACHE_SCHEDULE"] = "false"
+    argv = _torchrun_argv(
+        str(Path("tests/dreamzero/upstream/upstream_socket_server_no_compile.py")),
+        port,
+    )
+    log_file = log_path.open("w")
+    proc = subprocess.Popen(
+        argv,
+        stdout=log_file,
+        stderr=subprocess.STDOUT,
+        text=True,
+        cwd=str(Path.cwd()),
+        env=env,
+    )
+    proc._codex_log_file = log_file  # type: ignore[attr-defined]
+    return proc
+
+
+def _run_vllm_service(port: int, log_path: Path) -> subprocess.Popen[str]:
+    env = os.environ.copy()
+    gpus = _pick_test_gpus()
+    cfg_parallel_size = _cfg_parallel_size()
+    if cfg_parallel_size > len(gpus):
+        raise RuntimeError(
+            f"cfg_parallel_size={cfg_parallel_size} requires at least {cfg_parallel_size} GPUs, but only got {gpus}."
+        )
+    env["CUDA_VISIBLE_DEVICES"] = ",".join(gpus[:cfg_parallel_size])
+    env.setdefault("ATTENTION_BACKEND", "torch")
+    env.setdefault("DIFFUSION_ATTENTION_BACKEND", "TORCH_SDPA")
+    env.setdefault("MASTER_PORT", str(get_open_port()))
+    argv = [
+        _vllm_executable(),
+        "serve",
+        VLLM_MODEL,
+        "--omni",
+        "--host",
+        "127.0.0.1",
+        "--port",
+        str(port),
+        "--served-model-name",
+        "dreamzero-droid",
+        "--enforce-eager",
+    ]
+    if cfg_parallel_size > 1:
+        argv.extend(["--cfg-parallel-size", str(cfg_parallel_size)])
+    log_file = log_path.open("w")
+    proc = subprocess.Popen(
+        argv,
+        stdout=log_file,
+        stderr=subprocess.STDOUT,
+        text=True,
+        env=env,
+        cwd=str(Path.cwd()),
+    )
+    proc._codex_log_file = log_file  # type: ignore[attr-defined]
+    return proc
+
+
+def _stop_process(proc: subprocess.Popen[str]) -> None:
+    log_file = getattr(proc, "_codex_log_file", None)
+    if proc.poll() is None:
+        proc.terminate()
+        try:
+            proc.wait(timeout=30)
+        except subprocess.TimeoutExpired:  # pragma: no cover - cleanup path
+            proc.kill()
+            proc.wait(timeout=10)
+    if log_file is not None:
+        log_file.close()
+
+
+def _build_obs_sequence() -> tuple[dict, dict]:
+    camera_frames = dreamzero_client.load_camera_frames()
+    chunks = dreamzero_client.build_frame_schedule(
+        min(v.shape[0] for v in camera_frames.values()),
+        1,
+    )
+    obs0 = dreamzero_client._make_obs_from_video(camera_frames, [0], PROMPT, SESSION_ID)
+    obs1 = dreamzero_client._make_obs_from_video(camera_frames, chunks[0], PROMPT, SESSION_ID)
+    return obs0, obs1
+
+
+def _wait_for_client_ready(client_factory, timeout_s: float, proc=None, log_path: Path | None = None):
+    deadline = time.time() + timeout_s
+    last_err: Exception | None = None
+    while time.time() < deadline:
+        if proc is not None and proc.poll() is not None:
+            details = ""
+            if log_path is not None and log_path.exists():
+                details = log_path.read_text(errors="replace")[-8000:]
+            raise RuntimeError(f"Service exited before becoming ready with code {proc.returncode}.\n{details}")
+        try:
+            return client_factory()
+        except Exception as exc:  # pragma: no cover - retry path
+            last_err = exc
+            time.sleep(1)
+    raise TimeoutError(f"Timed out waiting for websocket service: {last_err}")
+
+
+def _collect_outputs_with_client(client) -> tuple[dict, list[np.ndarray]]:
+    metadata = client.get_server_metadata()
+    obs0, obs1 = _build_obs_sequence()
+    outputs = [
+        client.infer(dict(obs0)),
+        client.infer(dict(obs1)),
+    ]
+    assert _normalize_reset_response(client.reset({})) == "reset successful"
+    outputs.append(client.infer(dict(obs0)))
+    client._ws.close()
+    return metadata, outputs
+
+
+def _normalize_reset_response(response) -> str:
+    if isinstance(response, str):
+        return response
+    decoded = msgpack_numpy.unpackb(response)
+    if isinstance(decoded, dict):
+        return str(decoded.get("status"))
+    return str(decoded)
+
+
+def _normalize_metadata(metadata: dict) -> dict:
+    normalized = dict(metadata)
+    if isinstance(normalized.get("image_resolution"), tuple):
+        normalized["image_resolution"] = list(normalized["image_resolution"])
+    return normalized
+
+
+def _assert_logs_clean(log_path: Path) -> None:
+    text = log_path.read_text(errors="replace")
+    if "SignalException: Process" in text and "got signal: 15" in text:
+        text = text.split("Traceback (most recent call last):", 1)[0]
+    assert "Traceback" not in text, text
+    assert "RuntimeError:" not in text, text
+
+
+def _assert_upstream_log_matches_vllm_baseline(log_path: Path) -> None:
+    text = log_path.read_text(errors="replace")
+    assert "DIT Compute Steps 8 steps" not in text, text
+    assert "DIT Compute Steps 16 steps" in text, text
+
+
+def test_openpi_service_matches_upstream_server_noncompile(tmp_path: Path) -> None:
+    expected_metadata = {
+        "image_resolution": [180, 320],
+        "n_external_cameras": 2,
+        "needs_wrist_camera": True,
+        "needs_stereo_camera": False,
+        "needs_session_id": True,
+        "action_space": "joint_position",
+    }
+
+    upstream_port = get_open_port()
+    upstream_log = tmp_path / "dreamzero_upstream.log"
+    upstream_proc = _run_upstream_service(upstream_port, upstream_log)
+    try:
+        upstream_client = _wait_for_client_ready(
+            lambda: WebsocketClientPolicy(host="127.0.0.1", port=upstream_port),
+            timeout_s=SERVICE_READY_TIMEOUT_S,
+            proc=upstream_proc,
+            log_path=upstream_log,
+        )
+        upstream_metadata, upstream_outputs = _collect_outputs_with_client(upstream_client)
+    finally:
+        _stop_process(upstream_proc)
+    _assert_logs_clean(upstream_log)
+    _assert_upstream_log_matches_vllm_baseline(upstream_log)
+
+    vllm_port = get_open_port()
+    vllm_log = tmp_path / "vllm_openpi.log"
+    vllm_proc = _run_vllm_service(vllm_port, vllm_log)
+    try:
+        vllm_client = _wait_for_client_ready(
+            lambda: OpenPIWebsocketClientPolicy(host="127.0.0.1", port=vllm_port),
+            timeout_s=SERVICE_READY_TIMEOUT_S,
+            proc=vllm_proc,
+            log_path=vllm_log,
+        )
+        vllm_metadata, vllm_outputs = _collect_outputs_with_client(vllm_client)
+    finally:
+        _stop_process(vllm_proc)
+    _assert_logs_clean(vllm_log)
+
+    assert _normalize_metadata(upstream_metadata) == expected_metadata
+    assert _normalize_metadata(vllm_metadata) == expected_metadata
+
+    for idx, (actual, expected) in enumerate(zip(vllm_outputs, upstream_outputs, strict=True)):
+        np.testing.assert_allclose(
+            actual,
+            expected,
+            rtol=1e-2,
+            atol=1e-3,
+            err_msg=f"OpenPI step {idx} output mismatch",
+        )

tests/dreamzero/upstream/upstream_socket_server_no_compile.py

@@ -0,0 +1,184 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Launch the upstream DreamZero websocket server with `torch.compile` disabled.
+
+This wrapper is meant for formal parity tests against `vllm serve --omni`.
+It monkeypatches `torch.compile` before importing DreamZero modules so that
+all import-time decorators and `post_initialize()` compile calls become eager.
+
+For the current DreamZero port baseline we also disable all upstream DiT cache
+and step-skipping behavior so the reference server matches vLLM's current
+eager/no-skip implementation:
+- `ENABLE_DIT_CACHE=false`
+- `NUM_DIT_STEPS=16`
+- `DYNAMIC_CACHE_SCHEDULE=false`
+
+In our CI/dev environment we also do not install Transformer Engine or
+FlashAttention, so upstream's default `socket_test_optimized_AR.main()` path
+(`ATTENTION_BACKEND=TE` and then fallback to `FA2`) cannot execute. To keep the
+formal server-vs-server test runnable without pulling in those heavyweight
+optional deps, this wrapper reproduces upstream `main()` but pins attention to
+PyTorch SDPA (`ATTENTION_BACKEND=torch`) for this subprocess only.
+
+Usage:
+    PYTHONPATH="${DREAMZERO_REPO}" \\
+      .venv/bin/python -m torch.distributed.run --standalone --nproc_per_node=2 \\
+      tests/dreamzero/upstream/upstream_socket_server_no_compile.py --port 18081 \\
+      --model_path "${DREAMZERO_REPO}/checkpoints/dreamzero"
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+
+
+def _identity_compile(*args, **kwargs):
+    if args and callable(args[0]) and len(args) == 1 and not kwargs:
+        return args[0]
+
+    def deco(fn):
+        return fn
+
+    return deco
+
+
+torch.compile = _identity_compile
+
+os.environ.setdefault("NO_ALBUMENTATIONS_UPDATE", "1")
+
+DREAMZERO_REPO_ENV = os.environ.get("DREAMZERO_REPO")
+if not DREAMZERO_REPO_ENV:
+    raise RuntimeError("Set DREAMZERO_REPO to an upstream DreamZero checkout before launching this helper.")
+DREAMZERO_REPO = Path(DREAMZERO_REPO_ENV).expanduser()
+if str(DREAMZERO_REPO) not in sys.path:
+    sys.path.insert(0, str(DREAMZERO_REPO))
+
+import socket_test_optimized_AR as upstream  # noqa: E402
+import tyro  # noqa: E402
+from groot.vla.model.dreamzero.modules import attention as upstream_attention  # noqa: E402
+from groot.vla.model.dreamzero.modules import wan2_1_submodule as upstream_submodule  # noqa: E402
+from groot.vla.model.dreamzero.modules import wan_video_dit as upstream_wan_video_dit  # noqa: E402
+
+
+def _torch_varlen_flash_attention(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    q_lens: torch.Tensor | None = None,
+    k_lens: torch.Tensor | None = None,
+    dropout_p: float = 0.0,
+    causal: bool = False,
+    dtype: torch.dtype = torch.bfloat16,
+    **_: object,
+) -> torch.Tensor:
+    if q_lens is not None or k_lens is not None:
+        upstream_attention.warnings.warn(
+            "Padding mask is disabled in the test-only SDPA fallback.",
+        )
+    out_dtype = q.dtype
+    q = q.transpose(1, 2).to(dtype)
+    k = k.transpose(1, 2).to(dtype)
+    v = v.transpose(1, 2).to(dtype)
+    out = F.scaled_dot_product_attention(
+        q,
+        k,
+        v,
+        attn_mask=None,
+        is_causal=causal,
+        dropout_p=dropout_p,
+    )
+    return out.transpose(1, 2).contiguous().to(out_dtype)
+
+
+upstream_attention.flash_attention = _torch_varlen_flash_attention
+upstream_submodule.flash_attention = _torch_varlen_flash_attention
+upstream_wan_video_dit.FLASH_ATTN_3_AVAILABLE = False
+upstream_wan_video_dit.FLASH_ATTN_2_AVAILABLE = False
+upstream_wan_video_dit.SAGE_ATTN_AVAILABLE = False
+
+
+def main(args: upstream.Args) -> None:
+    os.environ["ENABLE_DIT_CACHE"] = "false"
+    os.environ["ATTENTION_BACKEND"] = "torch"
+    os.environ["NUM_DIT_STEPS"] = "16"
+    os.environ["DYNAMIC_CACHE_SCHEDULE"] = "false"
+    torch._dynamo.config.recompile_limit = 800
+
+    embodiment_tag = "oxe_droid"
+    model_path = args.model_path
+
+    device_mesh = upstream.init_mesh()
+    rank = upstream.dist.get_rank()
+
+    timeout_delta = upstream.datetime.timedelta(seconds=args.timeout_seconds)
+    signal_group = upstream.dist.new_group(backend="gloo", timeout=timeout_delta)
+    upstream.logger.info("Rank %s initialized signal_group (gloo)", rank)
+
+    policy = upstream.GrootSimPolicy(
+        embodiment_tag=upstream.EmbodimentTag(embodiment_tag),
+        model_path=model_path,
+        device="cuda" if torch.cuda.is_available() else "cpu",
+        device_mesh=device_mesh,
+    )
+
+    hostname = upstream.socket.gethostname()
+    local_ip = upstream.socket.gethostbyname(hostname)
+
+    if rank == 0:
+        upstream.logging.info("Creating server (host: %s, ip: %s)", hostname, local_ip)
+        parent_dir = os.path.dirname(model_path)
+        date_suffix = upstream.datetime.datetime.now().strftime("%Y%m%d")
+        checkpoint_name = os.path.basename(model_path)
+        output_dir = os.path.join(
+            parent_dir,
+            f"real_world_eval_gen_{date_suffix}_{args.index}",
+            checkpoint_name,
+        )
+        os.makedirs(output_dir, exist_ok=True)
+        upstream.logging.info("Videos will be saved to: %s", output_dir)
+    else:
+        output_dir = None
+        upstream.logging.info("Rank %s starting as worker for distributed inference...", rank)
+
+    wrapper_policy = upstream.ARDroidRoboarenaPolicy(
+        groot_policy=policy,
+        signal_group=signal_group,
+        output_dir=output_dir,
+    )
+
+    server_config = upstream.PolicyServerConfig(
+        image_resolution=(180, 320),
+        needs_wrist_camera=True,
+        n_external_cameras=2,
+        needs_stereo_camera=False,
+        needs_session_id=True,
+        action_space="joint_position",
+    )
+
+    if rank == 0:
+        upstream.logging.info("Using roboarena policy server interface")
+        upstream.logging.info("Server config: %s", server_config)
+        roboarena_server = upstream.RoboarenaServer(
+            policy=wrapper_policy,
+            server_config=server_config,
+            host="0.0.0.0",
+            port=args.port,
+        )
+        roboarena_server.serve_forever()
+    else:
+        worker = upstream.DistributedWorker(
+            policy=policy,
+            signal_group=signal_group,
+        )
+        upstream.asyncio.run(worker.run())
+
+
+if __name__ == "__main__":
+    upstream.logging.basicConfig(level=upstream.logging.INFO, force=True)
+    main(tyro.cli(upstream.Args))