Merge branch 'main' into feat/diffusion-prefetch-protection

.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:
           - |

docs/models/supported_models.md

@@ -33,7 +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 | `nvidia/Cosmos3-Nano` | ✅︎ | | | |
+| `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` | ✅︎ | ✅︎ | | |
@@ -42,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` | ✅︎ | ✅︎ | | ✅︎ |

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

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

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

@@ -198,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)
 
 

recipes/README.md

@@ -36,7 +36,8 @@ recipes/
 | [`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, and image-to-video generation | 1x H200 141GB / B300 |
+| [`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 |

recipes/nvidia/Cosmos3-Nano.md

@@ -6,7 +6,7 @@
 
 - Vendor: NVIDIA
 - Model: `nvidia/Cosmos3-Nano`
-- Task: Text-to-image (T2I), text-to-video (T2V), and image-to-video (I2V) generation
+- 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
 
@@ -20,12 +20,17 @@ the mode is selected per request:
 - **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`).
+  `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)
@@ -47,6 +52,9 @@ the mode is selected per request:
 
 #### 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:
 
@@ -116,6 +124,26 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync \
   -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
@@ -134,7 +162,9 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync \
   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;
+  `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:**
@@ -143,8 +173,8 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync \
     the server fails at pipeline build with a gated-repo / safety-checker error.
   - A guardrail-blocked prompt currently returns HTTP 500
     (`"Guardrail blocked prompt"`).
-  - Video + audio, and action (policy / forward- / inverse-dynamics) modalities
-    are not part of this integration yet.
+  - Action (policy / forward- / inverse-dynamics) modalities are not part of
+    this integration yet.
 
 ### 1x GPU (Offline generation)
 
@@ -170,8 +200,8 @@ def main():
         model_class_name="Cosmos3OmniDiffusersPipeline",
         trust_remote_code=True,
         enforce_eager=True,
-        # Keep guardrails on by installing cosmos-guardrail + gated-repo access;
-        # this disables them for a quick local run.
+        # 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)

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/diffusion/models/cosmos3/test_cosmos3_pipeline.py

@@ -73,14 +73,38 @@ class StubCosmos3VAE:
         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]] = []
@@ -99,8 +123,9 @@ class StubCosmos3Transformer(nn.Module):
         text_ids: torch.Tensor,
         text_mask: torch.Tensor,
         **kwargs: Any,
-    ) -> torch.Tensor:
+    ) -> 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,
@@ -114,7 +139,10 @@ class StubCosmos3Transformer(nn.Module):
             marker = torch.tensor([token], dtype=torch.float32)
             self.cached_kv = [(marker, marker + 100)]
             self.cached_freqs_gen = (marker + 200, marker + 300)
-        return torch.full_like(hidden_states, float(token))
+        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):
@@ -155,6 +183,7 @@ def make_cosmos3_pipeline():
         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
 
@@ -213,6 +242,108 @@ def test_pipeline_registered_and_exported() -> None:
     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
 
@@ -227,7 +358,7 @@ def test_preprocess_i2v_image_input() -> None:
     assert tuple(result.prompts[0]["additional_information"]["preprocessed_image"].shape[-2:]) == (672, 1344)
 
 
-def test_postprocess_handles_image_video_and_validation() -> None:
+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())
@@ -235,6 +366,16 @@ def test_postprocess_handles_image_video_and_validation() -> None:
 
     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)})
@@ -293,7 +434,7 @@ def test_prompt_formatting_and_checkpoint_key_remap(make_cosmos3_pipeline) -> No
     assert {key: Cosmos3OmniDiffusersPipeline._remap_ckpt_key(key) for key in remaps} == remaps
 
 
-def test_prepare_latents_for_video_and_image(make_cosmos3_pipeline) -> None:
+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)
@@ -306,8 +447,24 @@ def test_prepare_latents_for_video_and_image(make_cosmos3_pipeline) -> None:
     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_and_i2v_steps(make_cosmos3_pipeline) -> None:
+
+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)
 
@@ -339,6 +496,21 @@ def test_diffuse_covers_cfg_and_i2v_steps(make_cosmos3_pipeline) -> None:
     )
     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
@@ -395,7 +567,10 @@ class TestForwardRouting:
 
         def fake_diffuse(**kwargs):
             captured["diffuse_calls"].append(kwargs)
-            return kwargs["latents"] + len(captured["diffuse_calls"])
+            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
@@ -437,7 +612,7 @@ class TestForwardRouting:
         assert captured["flow_shifts"] == expected["flow"]
         assert [call[0] for call in pipeline.scheduler.set_timesteps_calls] == expected["steps"]
 
-    def test_forward_i2v_route(self, make_cosmos3_pipeline) -> None:
+    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)
@@ -462,11 +637,30 @@ class TestForwardRouting:
         )
         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(
@@ -477,6 +671,7 @@ class TestForwardRouting:
         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

@@ -4,6 +4,7 @@
 from __future__ import annotations
 
 from types import SimpleNamespace
+from typing import Any
 
 import pytest
 import torch
@@ -29,8 +30,9 @@ def _tiny_cosmos3_config(**overrides):
     return config
 
 
-def test_mrope_position_ids_cover_text_and_video() -> None:
+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,
     )
@@ -56,6 +58,10 @@ def test_mrope_position_ids_cover_text_and_video() -> None:
     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"),
@@ -120,7 +126,108 @@ def test_forward_returns_video_prediction(monkeypatch: pytest.MonkeyPatch) -> No
     assert tuple(output.shape) == (1, 2, 1, 2, 2)
 
 
-def test_compute_rope_freqs_places_text_and_video_positions() -> None:
+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:
@@ -140,6 +247,8 @@ def test_compute_rope_freqs_places_text_and_video_positions() -> None:
     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(
@@ -156,3 +265,19 @@ def test_compute_rope_freqs_places_text_and_video_positions() -> None:
     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_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/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))

tests/e2e/online_serving/test_dreamzero_expansion.py

@@ -0,0 +1,90 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""E2E online serving test for DreamZero OpenPI websocket serving."""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+from tests.dreamzero import openpi_client_helper as openpi_client
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServerParams, get_open_port
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
+
+MODEL = "GEAR-Dreams/DreamZero-DROID"
+
+
+test_params = [
+    OmniServerParams(
+        model=MODEL,
+        port=8091,
+        server_args=[
+            "--deploy-config",
+            "vllm_omni/deploy/dreamzero_tp1_cfg2.yaml",
+            "--enforce-eager",
+            "--disable-log-stats",
+        ],
+        env_dict={
+            "ATTENTION_BACKEND": "torch",
+            "DIFFUSION_ATTENTION_BACKEND": "TORCH_SDPA",
+            "VLLM_DISABLE_COMPILE_CACHE": "1",
+            "MASTER_PORT": str(get_open_port()),
+        },
+    )
+]
+
+
+def _write_synthetic_video(path: Path, cv2_module, *, channel: int) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    height, width, num_frames = 180, 320, 24
+    writer = cv2_module.VideoWriter(str(path), cv2_module.VideoWriter_fourcc(*"mp4v"), 15.0, (width, height))
+    if not writer.isOpened():
+        raise RuntimeError(f"Failed to open video writer for {path}")
+    try:
+        for frame_idx in range(num_frames):
+            frame = np.zeros((height, width, 3), dtype=np.uint8)
+            frame[..., channel] = (frame_idx * 7) % 255
+            frame[..., (channel + 1) % 3] = 64
+            writer.write(cv2_module.cvtColor(frame, cv2_module.COLOR_RGB2BGR))
+    finally:
+        writer.release()
+
+
+def _write_synthetic_dreamzero_videos(client_mod, video_dir: Path) -> None:
+    for channel, file_name in enumerate(client_mod.CAMERA_FILES.values()):
+        _write_synthetic_video(video_dir / file_name, client_mod.cv2, channel=channel)
+
+
+@pytest.mark.full_model
+@pytest.mark.diffusion
+@pytest.mark.distributed_cuda
+@hardware_test(res={"cuda": "H100"}, num_cards=2)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_dreamzero_openpi_online(omni_server, tmp_path: Path) -> None:
+    try:
+        openpi_client.require_dependencies()
+    except ModuleNotFoundError as exc:
+        pytest.skip(str(exc))
+
+    video_dir = tmp_path / "dreamzero_videos"
+    _write_synthetic_dreamzero_videos(openpi_client, video_dir)
+    result = openpi_client.run_policy_session(
+        host=omni_server.host,
+        port=omni_server.port,
+        video_dir=video_dir,
+        session_id="dreamzero-online-e2e",
+    )
+
+    openpi_client.validate_session_result(result)
+
+    metadata = result["metadata"]
+    assert metadata["needs_session_id"] is True
+    assert metadata["needs_stereo_camera"] is False
+    assert tuple(metadata["image_resolution"]) == (180, 320)

tests/e2e/online_serving/test_voxcpm2_tts.py

@@ -0,0 +1,97 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+E2E Online tests for VoxCPM2 with text input and audio output.
+
+These tests verify the /v1/audio/speech endpoint works correctly with
+actual model inference, not mocks. VoxCPM2 zero-shot needs only ``input``
+and a placeholder ``voice`` (``default``); no ``task_type`` / ``ref_text``.
+"""
+
+import os
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+import pytest
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServerParams
+from tests.helpers.stage_config import get_deploy_config_path
+
+MODEL = "openbmb/VoxCPM2"
+DEFAULT_AUDIO_SPEECH_TIMEOUT_S = 300.0
+MAX_CONCURRENT = 4
+
+# ~0.5 s of 48 kHz mono PCM_16 in WAV (~48k payload + header).
+_MIN_AUDIO_BYTES = 40_000
+
+
+def get_prompt(prompt_type="text"):
+    """English prompt for zero-shot TTS (Whisper-friendly for advanced_model accuracy)."""
+    prompts = {
+        "text": "The weather is nice today, perfect for a walk in the park.",
+    }
+    return prompts.get(prompt_type, prompts["text"])
+
+
+tts_server_params = [
+    pytest.param(
+        OmniServerParams(
+            model=MODEL,
+            stage_config_path=get_deploy_config_path("voxcpm2.yaml"),
+            server_args=["--trust-remote-code", "--disable-log-stats"],
+        ),
+        id="voxcpm2",
+    )
+]
+
+
+@pytest.mark.advanced_model
+@pytest.mark.core_model
+@pytest.mark.tts
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+def test_text_to_audio_001(omni_server, openai_client) -> None:
+    """
+    Test zero-shot text-to-audio via OpenAI API.
+    Deploy Setting: voxcpm2.yaml
+    Input Modal: text
+    Output Modal: audio
+    Input Setting: stream=False
+    Datasets: few requests (max_num_seqs=4)
+    """
+    request_config = {
+        "model": omni_server.model,
+        "input": get_prompt(),
+        "stream": False,
+        "timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S,
+        "response_format": "wav",
+        "voice": "default",
+        "min_audio_bytes": _MIN_AUDIO_BYTES,
+    }
+    openai_client.send_audio_speech_request(request_config, request_num=MAX_CONCURRENT)
+
+
+@pytest.mark.advanced_model
+@pytest.mark.tts
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+def test_text_to_audio_002(omni_server, openai_client) -> None:
+    """
+    Test zero-shot streaming text-to-audio via OpenAI API.
+    Deploy Setting: voxcpm2.yaml
+    Input Modal: text
+    Output Modal: audio
+    Input Setting: stream=True
+    Datasets: single request
+    """
+    request_config = {
+        "model": omni_server.model,
+        "input": get_prompt(),
+        "stream": True,
+        "timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S,
+        "response_format": "wav",
+        "voice": "default",
+        "min_audio_bytes": _MIN_AUDIO_BYTES,
+    }
+    openai_client.send_audio_speech_request(request_config)

tests/e2e/online_serving/test_voxcpm2_tts_expansion.py

@@ -0,0 +1,95 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+E2E Online expansion tests for VoxCPM2 (voice clone, PCM, concurrency).
+
+Voice clone uses ``ref_audio`` only (no ``ref_text``). Reference clip is
+vendored under tests/assets/qwen3_tts/clone_2.wav — see test_qwen3_tts_base.py.
+"""
+
+import os
+
+import pytest
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import load_test_audio_data_url
+from tests.helpers.runtime import OmniServerParams
+from tests.helpers.stage_config import get_deploy_config_path
+
+pytestmark = [pytest.mark.full_model, pytest.mark.tts]
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+MODEL = "openbmb/VoxCPM2"
+DEFAULT_AUDIO_SPEECH_TIMEOUT_S = 300.0
+_MIN_AUDIO_BYTES = 40_000
+MAX_CONCURRENT = 4
+
+REF_AUDIO_URL = load_test_audio_data_url("qwen3_tts/clone_2.wav")
+
+
+def get_prompt(prompt_type="text"):
+    prompts = {
+        "text": "The weather is nice today, perfect for a walk in the park.",
+    }
+    return prompts.get(prompt_type, prompts["text"])
+
+
+tts_server_params = [
+    pytest.param(
+        OmniServerParams(
+            model=MODEL,
+            stage_config_path=get_deploy_config_path("voxcpm2.yaml"),
+            server_args=["--trust-remote-code", "--disable-log-stats"],
+        ),
+        id="voxcpm2",
+    )
+]
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+def test_voice_clone_streaming_001(omni_server, openai_client) -> None:
+    """
+    Test voice-clone text-to-audio via OpenAI API.
+    Deploy Setting: voxcpm2.yaml
+    Input Modal: text + ref_audio
+    Output Modal: audio
+    Input Setting: stream=True
+    Datasets: few requests (max_num_seqs=4)
+    """
+    request_config = {
+        "model": omni_server.model,
+        "input": get_prompt(),
+        "stream": True,
+        "timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S,
+        "response_format": "wav",
+        "voice": "default",
+        "ref_audio": REF_AUDIO_URL,
+        "min_audio_bytes": _MIN_AUDIO_BYTES,
+    }
+    openai_client.send_audio_speech_request(request_config, request_num=MAX_CONCURRENT)
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+def test_response_format_001(omni_server, openai_client) -> None:
+    """
+    Test voice-clone non-stream PCM output via OpenAI API.
+    Deploy Setting: voxcpm2.yaml
+    Input Modal: text + ref_audio
+    Output Modal: audio (pcm)
+    Input Setting: stream=False
+    Datasets: single request
+    """
+    request_config = {
+        "model": omni_server.model,
+        "input": get_prompt(),
+        "response_format": "pcm",
+        "stream": False,
+        "timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S,
+        "voice": "default",
+        "ref_audio": REF_AUDIO_URL,
+        "min_audio_bytes": _MIN_AUDIO_BYTES,
+    }
+    openai_client.send_audio_speech_request(request_config)

tests/entrypoints/openai_api/test_video_server.py

@@ -399,6 +399,8 @@ def test_sampling_params_pass_through(test_client, mocker: MockerFixture):
             "true_cfg_scale": "4.0",
             "boundary_ratio": "0.7",
             "flow_shift": "0.25",
+            "generate_sound": "true",
+            "sound_duration": "2.5",
         },
     )
 
@@ -413,6 +415,8 @@ def test_sampling_params_pass_through(test_client, mocker: MockerFixture):
     assert captured.true_cfg_scale == 4.0
     assert captured.boundary_ratio == 0.7
     assert captured.extra_args["flow_shift"] == 0.25
+    assert captured.extra_args["generate_sound"] is True
+    assert captured.extra_args["sound_duration"] == 2.5
 
 
 def test_frame_interpolation_params_pass_to_diffusion_sampling_params(test_client, mocker: MockerFixture):
@@ -756,6 +760,9 @@ def test_invalid_uploaded_input_reference_returns_400(test_client):
 def test_video_request_validation():
     req = VideoGenerationRequest(prompt="test")
     assert req.prompt == "test"
+    assert req.generate_sound is False
+    assert req.sound_duration is None
+    assert VideoGenerationRequest(prompt="test", generate_sound=True, sound_duration=1.5).generate_sound is True
     with pytest.raises(ValueError):
         VideoGenerationRequest(prompt="test", size="invalid")
 
@@ -768,6 +775,8 @@ def test_video_request_validation():
         VideoGenerationRequest(prompt="test", frame_interpolation_exp=0)
     with pytest.raises(ValueError):
         VideoGenerationRequest(prompt="test", frame_interpolation_scale=0)
+    with pytest.raises(ValueError):
+        VideoGenerationRequest(prompt="test", sound_duration=0)
 
 
 def test_list_videos_supports_order_after_and_limit(test_client, mocker: MockerFixture):

tests/entrypoints/test_omni_entrypoints.py

@@ -413,6 +413,19 @@ def test_openai_serving_models_can_consume_async_omni_compat_attrs():
 
 
 def test_get_diffusion_od_config_returns_diffusion_stage_config():
+    diffusion_od_config = object()
+    omni = object.__new__(AsyncOmni)
+    omni.engine = SimpleNamespace(
+        stage_clients=[
+            SimpleNamespace(stage_type="llm"),
+            SimpleNamespace(stage_type="diffusion", od_config=diffusion_od_config),
+        ]
+    )
+
+    assert omni.get_diffusion_od_config() is diffusion_od_config
+
+
+def test_get_diffusion_od_config_falls_back_to_inner_engine():
     diffusion_od_config = object()
     omni = object.__new__(AsyncOmni)
     omni.engine = SimpleNamespace(

tests/entrypoints/test_resolve_dreamzero_config.py

@@ -0,0 +1,56 @@
+import pytest
+
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.stage_diffusion_proc import StageDiffusionProc
+from vllm_omni.entrypoints.utils import load_stage_configs_from_model, resolve_model_config_path
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_dreamzero_vla_resolves_to_dreamzero_config(monkeypatch):
+    monkeypatch.setattr(
+        "vllm_omni.entrypoints.utils.get_config",
+        lambda _model, trust_remote_code=True: type("Cfg", (), {"model_type": "vla"})(),
+    )
+    monkeypatch.setattr(
+        "vllm_omni.entrypoints.utils._looks_like_dreamzero",
+        lambda _model: True,
+    )
+    result = resolve_model_config_path("GEAR-Dreams/DreamZero-DROID")
+
+    assert result is not None
+    assert result.endswith("vllm_omni/deploy/dreamzero.yaml")
+
+
+def test_dreamzero_config_sets_model_class_and_policy_config(monkeypatch):
+    monkeypatch.setattr(
+        "vllm_omni.config.stage_config.StageConfigFactory._auto_detect_model_type",
+        classmethod(lambda _cls, _model, trust_remote_code=True: ("vla", None)),
+    )
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.utils.hf_utils._looks_like_dreamzero",
+        lambda _model: True,
+    )
+
+    stage_configs = load_stage_configs_from_model("GEAR-Dreams/DreamZero-DROID")
+    engine_args = stage_configs[0].engine_args
+
+    assert engine_args.model_class_name == "DreamZeroPipeline"
+    assert engine_args.model_config.policy_server_config.action_space == "joint_position"
+
+
+def test_dreamzero_enrich_config_preserves_explicit_model_class_name(monkeypatch):
+    monkeypatch.setattr(
+        "vllm.transformers_utils.config.get_hf_file_to_dict",
+        lambda path, _model: None if path == "model_index.json" else {"model_type": "vla", "architectures": ["VLA"]},
+    )
+
+    od_config = OmniDiffusionConfig(
+        model="GEAR-Dreams/DreamZero-DROID",
+        model_class_name="DreamZeroPipeline",
+    )
+    proc = StageDiffusionProc(od_config.model, od_config)
+
+    proc._enrich_config()
+
+    assert od_config.model_class_name == "DreamZeroPipeline"

vllm_omni/config/pipeline_registry.py

@@ -17,11 +17,10 @@ Adding a new pipeline:
        ``vllm_omni/.../pipeline.py``.
     2. Add one line to ``_OMNI_PIPELINES`` below.
 
-Single-stage diffusion models continue to use the
-``_create_default_diffusion_stage_cfg`` fallback in
-``async_omni_engine.py`` — they don't need a registry entry. The empty
-``_DIFFUSION_PIPELINES`` placeholder previously here (#2915) was removed
-once #2987 (which would have populated it) was deferred.
+Plain single-stage diffusion models continue to use the
+``_create_default_diffusion_stage_cfg`` fallback in ``async_omni_engine.py``.
+The empty ``_DIFFUSION_PIPELINES`` placeholder previously here (#2915) was
+removed once #2987 (which would have populated it) was deferred.
 
 ``register_pipeline(config)`` in ``stage_config`` is still supported for
 out-of-tree plugins and tests that create pipelines at runtime; those override
@@ -70,6 +69,10 @@ _OMNI_PIPELINES: dict[str, tuple[str, str]] = {
         "vllm_omni.model_executor.models.lance.pipeline",
         "LANCE_PIPELINE",
     ),
+    "dreamzero": (
+        "vllm_omni.model_executor.models.dreamzero.pipeline",
+        "DREAMZERO_PIPELINE",
+    ),
     "glm_image": (
         "vllm_omni.model_executor.models.glm_image.pipeline",
         "GLM_IMAGE_PIPELINE",
@@ -146,4 +149,8 @@ _OMNI_PIPELINES: dict[str, tuple[str, str]] = {
         "vllm_omni.model_executor.models.higgs_audio_v2.pipeline",
         "HIGGS_AUDIO_V2_PIPELINE",
     ),
+    "dynin_omni": (
+        "vllm_omni.model_executor.models.dynin_omni.pipeline",
+        "DYNIN_OMNI_PIPELINE",
+    ),
 }

vllm_omni/config/stage_config.py

@@ -1165,6 +1165,11 @@ class StageConfigFactory:
 
         # --- New path: check pipeline registry by model_type first ---
         model_type, hf_config = cls._auto_detect_model_type(model, trust_remote_code=trust_remote_code)
+        if model_type == "vla":
+            from vllm_omni.diffusion.utils.hf_utils import _looks_like_dreamzero
+
+            if _looks_like_dreamzero(model):
+                model_type = "dreamzero"
         if model_type and model_type in _PIPELINE_REGISTRY:
             return cls._create_from_registry(model_type, cli_overrides, deploy_config_path)
 

vllm_omni/deploy/dreamzero.yaml

@@ -0,0 +1,25 @@
+# DreamZero-DROID deploy: single diffusion stage.
+#
+# Topology is declared in vllm_omni/model_executor/models/dreamzero/pipeline.py.
+# This default uses one GPU with TP=1 and CFG parallel disabled.
+
+pipeline: dreamzero
+async_chunk: false
+distributed_executor_backend: mp
+dtype: bfloat16
+
+stages:
+  - stage_id: 0
+    devices: "0"
+    max_num_seqs: 1
+    enforce_eager: true
+    model_class_name: DreamZeroPipeline
+    model_config:
+      default_robot_embodiment: roboarena
+      policy_server_config:
+        image_resolution: [180, 320]
+        n_external_cameras: 2
+        needs_wrist_camera: true
+        needs_stereo_camera: false
+        needs_session_id: true
+        action_space: joint_position

vllm_omni/deploy/dreamzero_tp1_cfg2.yaml

@@ -0,0 +1,25 @@
+# DreamZero-DROID deploy: TP=1, CFG parallel size=2.
+
+pipeline: dreamzero
+async_chunk: false
+distributed_executor_backend: mp
+dtype: bfloat16
+
+stages:
+  - stage_id: 0
+    devices: "0,1"
+    max_num_seqs: 1
+    enforce_eager: true
+    model_class_name: DreamZeroPipeline
+    parallel_config:
+      tensor_parallel_size: 1
+      cfg_parallel_size: 2
+    model_config:
+      default_robot_embodiment: roboarena
+      policy_server_config:
+        image_resolution: [180, 320]
+        n_external_cameras: 2
+        needs_wrist_camera: true
+        needs_stereo_camera: false
+        needs_session_id: true
+        action_space: joint_position

vllm_omni/deploy/dynin_omni.yaml

@@ -0,0 +1,38 @@
+# Dynin-Omni 3-stage deploy: token2text → token2image → token2audio.
+#
+# Topology (stages, processor functions, model_arch) is frozen in
+# vllm_omni/model_executor/models/dynin_omni/pipeline.py. This file carries
+# only the deploy-time knobs (devices, GPU memory, batched tokens).
+#
+# async_chunk=false: dynin hands the full token payload across stages through
+# the worker-connector data plane (``*_full_payload`` producers), not chunked
+# streaming. Pipeline-wide ``trust_remote_code`` is required by the custom
+# Dynin HF modeling code.
+async_chunk: false
+trust_remote_code: true
+enable_prefix_caching: false
+
+stages:
+  - stage_id: 0
+    devices: "0"
+    max_num_seqs: 1
+    gpu_memory_utilization: 0.5
+    max_num_batched_tokens: 32768
+    enforce_eager: true
+    async_scheduling: false
+
+  - stage_id: 1
+    devices: "0"
+    max_num_seqs: 1
+    gpu_memory_utilization: 0.1
+    max_num_batched_tokens: 32768
+    enforce_eager: true
+    async_scheduling: false
+
+  - stage_id: 2
+    devices: "0"
+    max_num_seqs: 1
+    gpu_memory_utilization: 0.1
+    max_num_batched_tokens: 32768
+    enforce_eager: true
+    async_scheduling: false

vllm_omni/deploy/dynin_omni_ci.yaml

@@ -1,81 +1,35 @@
-# stage config for running dynin_omni with a 3-stage architecture.
-# this config is intended for e2e smoke tests.
+# Dynin-Omni deploy for e2e smoke tests (3-stage, one GPU).
+#
+# Topology is frozen in models/dynin_omni/pipeline.py; ``pipeline: dynin_omni``
+# pins it explicitly so this CI YAML resolves regardless of model_type
+# auto-detection. Differs from the default deploy only in reduced GPU memory
+# and a smaller max_num_batched_tokens so the pipeline fits L4-class CI GPUs.
+pipeline: dynin_omni
+async_chunk: false
+trust_remote_code: true
+enable_prefix_caching: false
 
-stage_args:
+stages:
   - stage_id: 0
-    stage_type: llm
-    runtime:
-      process: true
-      devices: "0"
-      max_batch_size: 1
-    engine_args:
-      model_stage: token2text
-      model_arch: DyninOmniForConditionalGeneration
-      worker_type: generation
-      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
-      engine_output_type: latent
-      custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.dynin_omni.token2text_to_token2image_full_payload
-      trust_remote_code: true
-      gpu_memory_utilization: 0.5
-      enforce_eager: true
-      enable_prefix_caching: false
-      async_scheduling: false
-      max_num_batched_tokens: 4096
-    is_comprehension: true
-    final_output: true
-    final_output_type: text
+    devices: "0"
+    max_num_seqs: 1
+    gpu_memory_utilization: 0.5
+    max_num_batched_tokens: 4096
+    enforce_eager: true
+    async_scheduling: false
 
   - stage_id: 1
-    stage_type: llm
-    runtime:
-      process: true
-      devices: "0"
-      max_batch_size: 1
-    engine_args:
-      model_stage: token2image
-      model_arch: DyninOmniForConditionalGeneration
-      worker_type: generation
-      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
-      engine_output_type: latent
-      custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.dynin_omni.token2image_to_token2audio_full_payload
-      trust_remote_code: true
-      gpu_memory_utilization: 0.2
-      enforce_eager: true
-      enable_prefix_caching: false
-      async_scheduling: false
-      max_num_batched_tokens: 4096
-    engine_input_source: [0]
-    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.dynin_omni.token2text_to_token2image
-    final_output: true
-    final_output_type: image
+    devices: "0"
+    max_num_seqs: 1
+    gpu_memory_utilization: 0.2
+    max_num_batched_tokens: 4096
+    enforce_eager: true
+    async_scheduling: false
 
   - stage_id: 2
-    stage_type: llm
-    runtime:
-      process: true
-      devices: "0"
-      max_batch_size: 1
-    engine_args:
-      model_stage: token2audio
-      model_arch: DyninOmniForConditionalGeneration
-      worker_type: generation
-      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
-      engine_output_type: latent
-      trust_remote_code: true
-      gpu_memory_utilization: 0.2
-      enforce_eager: true
-      enable_prefix_caching: false
-      async_scheduling: false
-      max_num_batched_tokens: 4096
-    engine_input_source: [1]
-    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.dynin_omni.token2image_to_token2audio
-    final_output: true
-    final_output_type: audio
-
-runtime:
-  enabled: true
-  edges:
-    - from: 0
-      to: 1
-    - from: 1
-      to: 2
+    devices: "0"
+    max_num_seqs: 1
+    gpu_memory_utilization: 0.2
+    max_num_batched_tokens: 4096
+    enforce_eager: true
+    async_scheduling: false

vllm_omni/deploy/dynin_omni_multiconnector.yaml

@@ -0,0 +1,65 @@
+# Dynin-Omni 3-stage deploy with cross-node connectors (Mooncake demo).
+#
+# Topology is frozen in models/dynin_omni/pipeline.py (``pipeline: dynin_omni``).
+# This variant wires a MooncakeConnector for each stage hand-off in place of
+# the default SharedMemoryConnector. Update the ``extra`` host/master/metadata
+# fields to match your Mooncake deployment before use.
+pipeline: dynin_omni
+async_chunk: false
+trust_remote_code: false
+enable_prefix_caching: false
+
+connectors:
+  # Mooncake connector for cross-node/intra-node communication.
+  mooncake_connector:
+    name: MooncakeConnector
+    extra:
+      host: "127.0.0.1"
+      metadata_server: "http://10.90.67.86:8080/metadata"
+      master: "10.90.67.86:50051"
+      segment: 512000000    # 512MB
+      localbuf: 64000000    # 64MB
+      proto: "tcp"
+
+  # Yuanrong connector for cross-node/intra-node communication.
+  yuanrong_connector:
+    name: YuanrongConnector
+    extra:
+      host: "127.0.0.1"
+      port: "35000"
+
+  # Alternative SHM connector with an explicit threshold.
+  shared_memory_connector:
+    name: SharedMemoryConnector
+    extra:
+      shm_threshold_bytes: 65536    # 64KB threshold
+
+stages:
+  - stage_id: 0
+    devices: "0"
+    max_num_seqs: 1
+    max_num_batched_tokens: 32768
+    enforce_eager: true
+    async_scheduling: false
+    output_connectors:
+      to_stage_1: mooncake_connector
+
+  - stage_id: 1
+    devices: "0"
+    max_num_seqs: 1
+    max_num_batched_tokens: 32768
+    enforce_eager: true
+    async_scheduling: false
+    input_connectors:
+      from_stage_0: mooncake_connector
+    output_connectors:
+      to_stage_2: mooncake_connector
+
+  - stage_id: 2
+    devices: "0"
+    max_num_seqs: 1
+    max_num_batched_tokens: 32768
+    enforce_eager: true
+    async_scheduling: false
+    input_connectors:
+      from_stage_1: mooncake_connector

vllm_omni/diffusion/attention/backends/flash_attn.py

@@ -11,6 +11,7 @@ from vllm_omni.diffusion.attention.backends.abstract import (
     AttentionImpl,
     AttentionMetadata,
 )
+from vllm_omni.diffusion.attention.backends.sdpa import _maybe_reshape_attn_mask
 from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import (
     piecewise_attn,
 )
@@ -309,6 +310,13 @@ class FlashAttentionImpl(AttentionImpl):
                 "Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA"
             )
         attention_mask = attn_metadata.attn_mask if attn_metadata else None
+
+        # NPU aclnnFlashAttentionScore requires mask shape to be one of:
+        # [B, N, Sq, Skv], [B, 1, Sq, Skv], [1, 1, Sq, Skv], or [Sq, Skv]
+        # But the incoming mask is 2D [B, S] — reshape to [B, 1, 1, S]
+        # So reuse SDPA's mask reshape logic: [B, S] -> [B, 1, Sq, Skv]
+        attention_mask = _maybe_reshape_attn_mask(query, key, attention_mask, mask_mode="full_qk")
+
         layout = self.qkv_layout or "BNSD"
         return attention_forward(
             query,

vllm_omni/diffusion/data.py

@@ -972,8 +972,24 @@ class OmniDiffusionConfig:
                         self.model_class_name = "WanS2VPipeline"
                     self.tf_model_config = TransformerConfig()
                     self.update_multimodal_support()
+                elif model_type == "vla":
+                    from vllm_omni.diffusion.utils.hf_utils import _looks_like_dreamzero
+
+                    if _looks_like_dreamzero(self.model):
+                        self.model_class_name = "DreamZeroPipeline"
+                        self.set_tf_model_config(TransformerConfig())
+                        self.update_multimodal_support()
+                    else:
+                        raise
                 elif architectures and len(architectures) == 1:
-                    self.model_class_name = architectures[0]
+                    architecture = architectures[0]
+                    from vllm_omni.diffusion.registry import DiffusionModelRegistry
+
+                    if (
+                        self.model_class_name is None
+                        or DiffusionModelRegistry._try_load_model_cls(architecture) is not None
+                    ):
+                        self.model_class_name = architecture
                 else:
                     raise
 
@@ -1043,10 +1059,10 @@ class DiffusionOutput:
     """
 
     # Fields may be replaced with SHM handle dicts by ipc.pack_diffusion_output_shm
-    output: torch.Tensor | dict | None = None
-    trajectory_timesteps: torch.Tensor | dict | None = None
-    trajectory_latents: torch.Tensor | dict | None = None
-    trajectory_log_probs: torch.Tensor | dict | None = None
+    output: torch.Tensor | tuple[Any, ...] | dict[str, Any] | None = None
+    trajectory_timesteps: torch.Tensor | dict[str, Any] | None = None
+    trajectory_latents: torch.Tensor | dict[str, Any] | None = None
+    trajectory_log_probs: torch.Tensor | dict[str, Any] | None = None
     trajectory_decoded: list[Image.Image] | None = None
     error: str | None = None
     aborted: bool = False

vllm_omni/diffusion/diffusion_engine.py

@@ -266,8 +266,10 @@ class DiffusionEngine:
         custom_output = output.custom_output or {}
         model_audio_sample_rate = None
         model_fps = None
+        action_payload = None
         if isinstance(outputs, dict):
             audio_payload = outputs.get("audio")
+            action_payload = outputs.get("actions")
             custom_output.update(outputs.get("custom_output") or {})
             model_audio_sample_rate = outputs.get("audio_sample_rate")
             model_fps = outputs.get("fps")
@@ -361,6 +363,8 @@ class DiffusionEngine:
                     mm_output["audio_sample_rate"] = model_audio_sample_rate
                 if model_fps is not None:
                     mm_output["fps"] = model_fps
+                if action_payload is not None:
+                    mm_output["actions"] = action_payload
                 return [
                     OmniRequestOutput.from_diffusion(
                         request_id=request_id,
@@ -430,6 +434,18 @@ class DiffusionEngine:
                         mm_output["audio_sample_rate"] = model_audio_sample_rate
                     if model_fps is not None:
                         mm_output["fps"] = model_fps
+                    if action_payload is not None:
+                        sliced_actions = action_payload
+                        if isinstance(action_payload, (list, tuple)):
+                            sliced_actions = action_payload[start_idx:end_idx]
+                            if len(sliced_actions) == 1:
+                                sliced_actions = sliced_actions[0]
+                        elif hasattr(action_payload, "shape") and getattr(action_payload, "shape", None) is not None:
+                            if len(action_payload.shape) > 0 and action_payload.shape[0] >= end_idx:
+                                sliced_actions = action_payload[start_idx:end_idx]
+                                if num_outputs == 1:
+                                    sliced_actions = sliced_actions[0]
+                        mm_output["actions"] = sliced_actions
                     results.append(
                         OmniRequestOutput.from_diffusion(
                             request_id=request_id,

vllm_omni/diffusion/models/cosmos3/audio_tokenizer/__init__.py

@@ -0,0 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from .avae import Cosmos3AVAEAudioTokenizer
+
+__all__ = ["Cosmos3AVAEAudioTokenizer"]

vllm_omni/diffusion/models/cosmos3/audio_tokenizer/avae.py

@@ -0,0 +1,323 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Diffusers-format AVAE audio tokenizer used by Cosmos3 sound generation."""
+
+from __future__ import annotations
+
+import json
+import math
+from pathlib import Path
+from typing import Any
+
+import torch
+from torch import nn
+from torch.nn.utils import weight_norm
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.models.progress_bar import _is_rank_zero
+
+logger = init_logger(__name__)
+
+
+def _default_avae_config(
+    *,
+    sample_rate: int,
+    audio_channels: int,
+    io_channels: int,
+    hop_size: int,
+) -> dict[str, Any]:
+    return {
+        "sampling_rate": sample_rate,
+        "hop_size": hop_size,
+        "dec_dim": 320,
+        "dec_c_mults": [1, 2, 4, 8, 16],
+        "dec_strides": [2, 4, 5, 6, 8],
+        "dec_out_channels": audio_channels,
+        "vocoder_input_dim": io_channels,
+        "normalization_type": "none",
+        "normalize_latents": False,
+        "tanh_input_scale": 1.5,
+        "tanh_output_scale": 3.5,
+        "tanh_clamp": 0.995,
+    }
+
+
+def _config_get(config: dict[str, Any], *keys: str, default: Any = None) -> Any:
+    for key in keys:
+        value = config.get(key)
+        if value is not None:
+            return value
+    return default
+
+
+def _load_config(
+    config_path: str | Path | None,
+    *,
+    sample_rate: int,
+    audio_channels: int,
+    io_channels: int,
+    hop_size: int,
+) -> dict[str, Any]:
+    if config_path:
+        with open(config_path, encoding="utf-8") as f:
+            config = json.load(f)
+        if not isinstance(config, dict):
+            raise TypeError(f"Cosmos3 AVAE config must be a JSON object, got {type(config)!r}.")
+        return config
+    return _default_avae_config(
+        sample_rate=sample_rate,
+        audio_channels=audio_channels,
+        io_channels=io_channels,
+        hop_size=hop_size,
+    )
+
+
+def _load_checkpoint(path: str | Path, map_location: torch.device | str) -> dict[str, torch.Tensor]:
+    path = Path(path)
+    if path.suffix == ".safetensors":
+        try:
+            from safetensors.torch import load_file
+        except ImportError as exc:
+            raise ImportError("Loading AVAE .safetensors checkpoints requires safetensors.") from exc
+        checkpoint = load_file(str(path), device=str(map_location))
+    else:
+        checkpoint = torch.load(path, map_location=map_location)
+
+    if not isinstance(checkpoint, dict):
+        raise TypeError(f"AVAE checkpoint must be a flat state dict, got {type(checkpoint)!r}.")
+    if not all(isinstance(value, torch.Tensor) for value in checkpoint.values()):
+        raise TypeError("AVAE checkpoint must be a flat tensor state dict.")
+    return checkpoint
+
+
+def _validate_diffusers_state_dict(state_dict: dict[str, torch.Tensor]) -> None:
+    if not state_dict:
+        raise RuntimeError("AVAE checkpoint is empty.")
+
+    if not any(key.startswith("decoder.") for key in state_dict):
+        raise RuntimeError("Cosmos3 AVAE checkpoint must contain diffusers-format decoder.* keys.")
+
+
+class Snake1d(nn.Module):
+    """One-dimensional Snake activation matching diffusers' Oobleck layout."""
+
+    def __init__(self, hidden_dim: int, logscale: bool = True) -> None:
+        super().__init__()
+        self.alpha = nn.Parameter(torch.zeros(1, hidden_dim, 1))
+        self.beta = nn.Parameter(torch.zeros(1, hidden_dim, 1))
+        self.logscale = logscale
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        shape = hidden_states.shape
+        alpha = torch.exp(self.alpha) if self.logscale else self.alpha
+        beta = torch.exp(self.beta) if self.logscale else self.beta
+        hidden_states = hidden_states.reshape(shape[0], shape[1], -1)
+        hidden_states = hidden_states + (beta + 1e-9).reciprocal() * torch.sin(alpha * hidden_states).pow(2)
+        return hidden_states.reshape(shape)
+
+
+class OobleckResidualUnit(nn.Module):
+    """Residual unit used by the diffusers Oobleck decoder."""
+
+    def __init__(self, dimension: int = 16, dilation: int = 1) -> None:
+        super().__init__()
+        pad = ((7 - 1) * dilation) // 2
+        self.snake1 = Snake1d(dimension)
+        self.conv1 = weight_norm(nn.Conv1d(dimension, dimension, kernel_size=7, dilation=dilation, padding=pad))
+        self.snake2 = Snake1d(dimension)
+        self.conv2 = weight_norm(nn.Conv1d(dimension, dimension, kernel_size=1))
+
+    def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
+        output_tensor = self.conv1(self.snake1(hidden_state))
+        output_tensor = self.conv2(self.snake2(output_tensor))
+        padding = (hidden_state.shape[-1] - output_tensor.shape[-1]) // 2
+        if padding > 0:
+            hidden_state = hidden_state[..., padding:-padding]
+        return hidden_state + output_tensor
+
+
+class OobleckDecoderBlock(nn.Module):
+    """Decoder block used by the diffusers Oobleck decoder."""
+
+    def __init__(self, input_dim: int, output_dim: int, stride: int = 1, output_padding: int = 0) -> None:
+        super().__init__()
+        self.snake1 = Snake1d(input_dim)
+        self.conv_t1 = weight_norm(
+            nn.ConvTranspose1d(
+                input_dim,
+                output_dim,
+                kernel_size=2 * stride,
+                stride=stride,
+                padding=math.ceil(stride / 2),
+                output_padding=output_padding,
+            )
+        )
+        self.res_unit1 = OobleckResidualUnit(output_dim, dilation=1)
+        self.res_unit2 = OobleckResidualUnit(output_dim, dilation=3)
+        self.res_unit3 = OobleckResidualUnit(output_dim, dilation=9)
+
+    def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
+        hidden_state = self.snake1(hidden_state)
+        hidden_state = self.conv_t1(hidden_state)
+        hidden_state = self.res_unit1(hidden_state)
+        hidden_state = self.res_unit2(hidden_state)
+        return self.res_unit3(hidden_state)
+
+
+class OobleckDecoder(nn.Module):
+    """Diffusers-compatible Oobleck decoder for Cosmos3 AVAE latents."""
+
+    def __init__(
+        self,
+        channels: int,
+        input_channels: int,
+        audio_channels: int,
+        upsampling_ratios: list[int],
+        channel_multiples: list[int],
+    ) -> None:
+        super().__init__()
+        strides = upsampling_ratios
+        channel_multiples = [1] + channel_multiples
+
+        self.conv1 = weight_norm(nn.Conv1d(input_channels, channels * channel_multiples[-1], kernel_size=7, padding=3))
+
+        block = []
+        for stride_index, stride in enumerate(strides):
+            block.append(
+                OobleckDecoderBlock(
+                    input_dim=channels * channel_multiples[len(strides) - stride_index],
+                    output_dim=channels * channel_multiples[len(strides) - stride_index - 1],
+                    stride=stride,
+                    output_padding=stride % 2,
+                )
+            )
+        self.block = nn.ModuleList(block)
+        self.snake1 = Snake1d(channels)
+        self.conv2 = weight_norm(nn.Conv1d(channels, audio_channels, kernel_size=7, padding=3, bias=False))
+
+    def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
+        hidden_state = self.conv1(hidden_state)
+        for layer in self.block:
+            hidden_state = layer(hidden_state)
+        hidden_state = self.snake1(hidden_state)
+        return self.conv2(hidden_state)
+
+
+class Cosmos3AVAEAudioTokenizer(nn.Module):
+    """Decoder-only AVAE tokenizer for Cosmos3 audio latents."""
+
+    def __init__(
+        self,
+        *,
+        checkpoint_path: str | Path,
+        config_path: str | Path | None = None,
+        sample_rate: int = 48000,
+        audio_channels: int = 2,
+        io_channels: int = 64,
+        hop_size: int = 1920,
+        normalize_latents: bool = False,
+        normalization_type: str = "none",
+        tanh_input_scale: float = 1.5,
+        tanh_output_scale: float = 3.5,
+        tanh_clamp: float = 0.995,
+        dtype: torch.dtype = torch.bfloat16,
+        device: torch.device | str = "cuda",
+    ) -> None:
+        super().__init__()
+        self.dtype = dtype
+        self.device = torch.device(device)
+
+        config = _load_config(
+            config_path,
+            sample_rate=sample_rate,
+            audio_channels=audio_channels,
+            io_channels=io_channels,
+            hop_size=hop_size,
+        )
+        self.sample_rate = int(_config_get(config, "sampling_rate", "sample_rate", default=sample_rate))
+        self.audio_channels = int(
+            _config_get(
+                config,
+                "dec_out_channels",
+                "audio_channels",
+                default=2 if bool(config.get("stereo", audio_channels == 2)) else 1,
+            )
+        )
+        self.latent_ch = int(_config_get(config, "vocoder_input_dim", "io_channels", "latent_ch", default=io_channels))
+        dec_strides = [int(stride) for stride in _config_get(config, "dec_strides", default=[2, 4, 5, 6, 8])]
+        self.hop_size = int(
+            _config_get(config, "hop_size", default=math.prod(dec_strides) if dec_strides else hop_size)
+        )
+        dec_stride_product = math.prod(dec_strides)
+        if dec_stride_product != self.hop_size:
+            raise ValueError(
+                "Cosmos3 AVAE config dec_strides product must equal hop_size "
+                f"for correct latent/audio duration math: product={dec_stride_product}, hop_size={self.hop_size}."
+            )
+
+        normalization_type = str(_config_get(config, "normalization_type", default=normalization_type))
+        normalize_latents = bool(_config_get(config, "normalize_latents", default=normalize_latents))
+        if normalization_type == "none" and normalize_latents:
+            normalization_type = "tanh"
+        self.normalization_type = normalization_type
+        self.tanh_input_scale = float(_config_get(config, "tanh_input_scale", default=tanh_input_scale))
+        self.tanh_output_scale = float(_config_get(config, "tanh_output_scale", default=tanh_output_scale))
+        self.tanh_clamp = float(_config_get(config, "tanh_clamp", default=tanh_clamp))
+
+        self.decoder = OobleckDecoder(
+            channels=int(_config_get(config, "dec_dim", default=320)),
+            input_channels=self.latent_ch,
+            audio_channels=self.audio_channels,
+            upsampling_ratios=list(reversed(dec_strides)),
+            channel_multiples=list(_config_get(config, "dec_c_mults", default=[1, 2, 4, 8, 16])),
+        )
+        state_dict = _load_checkpoint(checkpoint_path, self.device)
+        _validate_diffusers_state_dict(state_dict)
+
+        # The checkpoint also contains encoder weights, which we do not support here, hence strict=False
+        self.load_state_dict(state_dict, strict=False)
+
+        self.eval()
+        for param in self.parameters():
+            param.requires_grad = False
+        self.to(device=self.device, dtype=self.dtype)
+        if _is_rank_zero():
+            logger.info("Loaded diffusers-format Cosmos3 AVAE checkpoint from %s", checkpoint_path)
+
+    @property
+    def temporal_compression_factor(self) -> int:
+        return self.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 _denormalize_latent(self, latent: torch.Tensor) -> torch.Tensor:
+        if self.normalization_type == "tanh":
+            in_dtype = latent.dtype
+            latent = torch.clamp(
+                latent.float() / self.tanh_output_scale,
+                -self.tanh_clamp,
+                self.tanh_clamp,
+            )
+            return (torch.atanh(latent) * self.tanh_input_scale).to(in_dtype)
+        if self.normalization_type != "none":
+            raise ValueError(f"Unsupported AVAE normalization_type={self.normalization_type!r}.")
+        return latent
+
+    @torch.no_grad()
+    def encode(self, audio: torch.Tensor, force_pad: bool = False) -> torch.Tensor:
+        del audio, force_pad
+        raise NotImplementedError("Cosmos3AVAEAudioTokenizer is decoder-only for diffusers-format sound_tokenizer/.")
+
+    @torch.no_grad()
+    def decode(self, latent: torch.Tensor) -> torch.Tensor:
+        in_dtype = latent.dtype
+        squeeze = latent.ndim == 2
+        if squeeze:
+            latent = latent.unsqueeze(0)
+        z = self._denormalize_latent(latent.to(self.device)).to(self.dtype)
+        audio = self.decoder(z).clamp(-1.0, 1.0).to(in_dtype)
+        return audio.squeeze(0) if squeeze else audio

vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py

@@ -15,6 +15,7 @@ runtime by:
 
 from __future__ import annotations
 
+import math
 import os
 import time
 from collections.abc import Iterable
@@ -45,7 +46,7 @@ from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPi
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
-from .transformer_cosmos3 import Cosmos3VFMTransformer
+from .transformer_cosmos3 import Cosmos3VFMTransformer, resolve_sound_gen
 
 logger = init_logger(__name__)
 
@@ -160,6 +161,28 @@ def get_cosmos3_post_process_func(od_config: OmniDiffusionConfig):
 
     video_processor = VideoProcessor(vae_scale_factor=16)
 
+    def _sampling_param(sampling_params, key: str, default=None):
+        extra = getattr(sampling_params, "extra_args", None)
+        if isinstance(extra, dict) and extra.get(key) is not None:
+            return extra[key]
+        value = getattr(sampling_params, key, None)
+        return default if value is None else value
+
+    def _resolve_output_fps(sampling_params):
+        fps = (
+            _sampling_param(sampling_params, "resolved_frame_rate")
+            or _sampling_param(sampling_params, "frame_rate")
+            or _sampling_param(sampling_params, "fps")
+            or 24.0
+        )
+        try:
+            fps_value = float(fps)
+        except (TypeError, ValueError):
+            fps_value = 24.0
+        if fps_value <= 0:
+            fps_value = 24.0
+        return int(fps_value) if fps_value.is_integer() else fps_value
+
     def post_process_func(
         output: torch.Tensor | dict[str, torch.Tensor] | tuple,
         output_type: str = "np",
@@ -168,6 +191,8 @@ def get_cosmos3_post_process_func(od_config: OmniDiffusionConfig):
         if output_type == "latent":
             return output
 
+        audio = None
+        audio_sample_rate = None
         if isinstance(output, dict):
             if "image" in output and "video" in output:
                 raise ValueError("Cosmos3 output cannot contain both image and video payloads.")
@@ -177,10 +202,23 @@ def get_cosmos3_post_process_func(od_config: OmniDiffusionConfig):
                 video = output["video"]
             else:
                 raise ValueError("Cosmos3 postprocess expected an 'image' or 'video' output payload.")
+            audio = output.get("audio")
+            audio_sample_rate = output.get("audio_sample_rate")
+        elif isinstance(output, tuple):
+            if len(output) == 3:
+                video, audio, audio_sample_rate = output
+            elif len(output) == 2:
+                video, audio = output
+            else:
+                raise ValueError(
+                    "Cosmos3 postprocess expects output tensor, output dict, or (video, audio[, sample_rate]) tuple."
+                )
         else:
             video = output
 
         if isinstance(output, dict) and "image" in output:
+            if audio is not None:
+                raise ValueError("Cosmos3 text-to-image postprocess does not support audio output.")
             if video.ndim != 5 or video.shape[2] != 1:
                 raise ValueError(
                     "Cosmos3 text-to-image postprocess expects decoded output "
@@ -194,7 +232,19 @@ def get_cosmos3_post_process_func(od_config: OmniDiffusionConfig):
             return video_processor.postprocess(image, output_type="pil")
         if is_guardrails_enabled(od_config, sampling_params):
             video = check_video_safety(video)
-        return video_processor.postprocess_video(video, output_type=output_type)
+        processed_video = video_processor.postprocess_video(video, output_type=output_type)
+        if audio is None:
+            return processed_video
+        if isinstance(audio, torch.Tensor):
+            audio = audio.detach().cpu()
+        result = {
+            "video": processed_video,
+            "audio": audio,
+            "fps": _resolve_output_fps(sampling_params),
+        }
+        if audio_sample_rate is not None:
+            result["audio_sample_rate"] = int(audio_sample_rate)
+        return result
 
     return post_process_func
 
@@ -274,10 +324,22 @@ class Cosmos3OmniDiffusersPipeline(
         self.vae_scale_factor_temporal = int(self.vae.config.scale_factor_temporal)
         self.vae_scale_factor_spatial = getattr(self.vae.config, "scale_factor_spatial", 16)
 
+        sound_gen = resolve_sound_gen(od_config)
+        sound_dim = None
+        sound_latent_fps = None
+        self._sound_tokenizer = None
+        if sound_gen:
+            self._sound_tokenizer = self._get_sound_tokenizer()
+            sound_dim = self._sound_tokenizer.latent_ch
+            sound_latent_fps = self._sound_tokenizer.latent_fps
+
         # --- Transformer (weights loaded later via weights_sources) ---
         self.transformer = Cosmos3VFMTransformer(
             od_config=od_config,
             temporal_compression_factor=self.vae_scale_factor_temporal,
+            sound_gen=sound_gen,
+            sound_dim=sound_dim,
+            sound_latent_fps=sound_latent_fps,
         )
 
         # --- Scheduler ---
@@ -357,9 +419,13 @@ class Cosmos3OmniDiffusersPipeline(
                 "proj_in.",
                 "proj_out.",
                 "time_embedder.",
+                "audio_proj_in.",
+                "audio_proj_out.",
             )
         ):
             return f"transformer.{k}"
+        if k in ("audio_modality_embed", "audio_modality_embed.weight"):
+            return "transformer.audio_modality_embed"
 
         # Skip lm_head
         if k.startswith("lm_head."):
@@ -453,12 +519,22 @@ class Cosmos3OmniDiffusersPipeline(
         loaded = loader.load_weights(_remapped_weights())
         self.transformer.post_load_weights()
         self.transformer.eval()
+        if getattr(self.transformer, "sound_gen", False):
+            sound_markers = ("audio_proj_in.", "audio_proj_out.", "audio_modality_embed")
+            missing = [marker.rstrip(".") for marker in sound_markers if not any(marker in name for name in loaded)]
+            if missing:
+                raise ValueError(
+                    "Cosmos3 transformer config enables sound generation, but "
+                    f"the checkpoint is missing sound weights for {missing}. "
+                    "Use a sound-capable transformer checkpoint."
+                )
         return loaded
 
     def predict_noise(self, **kwargs) -> torch.Tensor | tuple[torch.Tensor, ...]:
         """Override CFGParallelMixin.predict_noise for Cosmos3.
 
-        The transformer returns the raw video noise prediction.
+        The transformer returns the raw prediction: video-only as a tensor,
+        or a tuple in video, sound order for sound generation.
         """
         return self.transformer(**kwargs)
 
@@ -509,6 +585,39 @@ class Cosmos3OmniDiffusersPipeline(
             return val
         return default
 
+    @staticmethod
+    def _truthy(value) -> bool:
+        if isinstance(value, str):
+            return value.strip().lower() in {"1", "true", "yes", "on"}
+        return bool(value)
+
+    @classmethod
+    def _get_prompt_param(cls, prompt_data, key: str, default=None):
+        if not isinstance(prompt_data, dict):
+            return default
+        if prompt_data.get(key) is not None:
+            return prompt_data[key]
+        additional = prompt_data.get("additional_information")
+        if isinstance(additional, dict) and additional.get(key) is not None:
+            return additional[key]
+        return default
+
+    @classmethod
+    def _is_sound_request(cls, prompt_data, sp) -> bool:
+        for key in ("generate_sound", "sound_gen"):
+            if cls._truthy(cls._get_prompt_param(prompt_data, key, None)):
+                return True
+            if cls._truthy(cls._get_sp_param(sp, key, None)):
+                return True
+        return False
+
+    def _get_sound_tokenizer(self):
+        if self._sound_tokenizer is None:
+            from .sound_tokenizer import Cosmos3SoundTokenizer
+
+            self._sound_tokenizer = Cosmos3SoundTokenizer.from_config(self.od_config)
+        return self._sound_tokenizer
+
     @staticmethod
     def _is_t2i_request(req: OmniDiffusionRequest) -> bool:
         """Detect text-to-image mode from request-level prompt modalities."""
@@ -721,6 +830,47 @@ class Cosmos3OmniDiffusersPipeline(
         )
         return randn_tensor(shape, generator=generator, device=self.device, dtype=self.dtype)
 
+    def _prepare_sound_latents(
+        self,
+        target_audio_samples: int,
+        generator: torch.Generator,
+    ) -> tuple[torch.Tensor, int]:
+        sound_tokenizer = self._get_sound_tokenizer()
+        hop_size = int(
+            getattr(sound_tokenizer, "hop_size", None) or getattr(sound_tokenizer, "temporal_compression_factor")
+        )
+        latent_frames = max(1, math.ceil(max(1, int(target_audio_samples)) / hop_size))
+        sound_dim = int(getattr(sound_tokenizer, "latent_ch", 64))
+        transformer_sound_dim = int(getattr(self.transformer, "sound_dim", sound_dim))
+        if sound_dim != transformer_sound_dim:
+            raise ValueError(
+                "Cosmos3 sound tokenizer latent channels do not match transformer "
+                f"sound_dim: tokenizer={sound_dim}, transformer={transformer_sound_dim}."
+            )
+        latents = randn_tensor(
+            (1, sound_dim, latent_frames),
+            generator=generator,
+            device=self.device,
+            dtype=self.dtype,
+        )
+        return latents, latent_frames
+
+    def _resolve_sound_target_samples(
+        self,
+        sp,
+        num_frames: int,
+        frame_rate: float,
+    ) -> tuple[int, float, int]:
+        sound_tokenizer = self._get_sound_tokenizer()
+        duration = self._get_sp_param(sp, "sound_duration", None)
+        if duration is None:
+            duration = self._get_sp_param(sp, "audio_duration", None)
+        if duration is None:
+            duration = num_frames / frame_rate
+        duration = max(float(duration), 1.0 / max(float(frame_rate), 1.0))
+        sample_rate = int(getattr(sound_tokenizer, "sample_rate", 48000))
+        return max(1, int(round(duration * sample_rate))), duration, sample_rate
+
     # -- VAE decode ----------------------------------------------------------
 
     def _decode_latents(self, latents: torch.Tensor) -> torch.Tensor:
@@ -742,6 +892,19 @@ class Cosmos3OmniDiffusersPipeline(
         video = self.vae.decode(latents, return_dict=False)[0]
         return video
 
+    def _decode_sound_latents(
+        self,
+        sound_latents: torch.Tensor,
+        target_audio_samples: int,
+    ) -> torch.Tensor:
+        sound_tokenizer = self._get_sound_tokenizer()
+        audio = sound_tokenizer.decode(sound_latents.to(self.dtype))
+        if audio.shape[-1] > target_audio_samples:
+            audio = audio[..., :target_audio_samples]
+        elif audio.shape[-1] < target_audio_samples:
+            audio = torch.nn.functional.pad(audio, (0, target_audio_samples - audio.shape[-1]))
+        return audio.detach().cpu()
+
     # -- Prompt formatting + tokenization (shared by T2V and I2V) ------------
 
     def _format_and_tokenize_prompts(
@@ -903,11 +1066,12 @@ class Cosmos3OmniDiffusersPipeline(
         guidance_scale: float,
         shared_kwargs: dict,
         *,
+        sound_latents: torch.Tensor | None = None,
         velocity_mask: torch.Tensor | None = None,
         image_latent: torch.Tensor | None = None,
         condition_latents: torch.Tensor | None = None,
         guidance_interval: tuple[float, float] | None = None,
-    ) -> torch.Tensor:
+    ) -> torch.Tensor | tuple[torch.Tensor, ...]:
         """Denoising loop with 3-mode CFG support (parallel, sequential, none).
 
         Cosmos3's UND pathway is text-dependent, so CFG needs separate K/V
@@ -946,21 +1110,85 @@ class Cosmos3OmniDiffusersPipeline(
             lo, hi = guidance_interval
             return lo <= t_scalar <= hi
 
+        # Joint scheduler step over multiple modalities. Safe for flow-matching schedulers
+        # because the update is linear per element; revisit this if Cosmos3 adopts a
+        # scheduler with cross-element dependencies (e.g. per-modality timestep).
+        def _pack_joint(
+            video_tensor: torch.Tensor,
+            sound_tensor: torch.Tensor | None = None,
+        ):
+            batch = video_tensor.shape[0]
+            tensors = [video_tensor]
+            if sound_tensor is not None:
+                tensors.append(sound_tensor)
+            flats = [tensor.reshape(batch, -1) for tensor in tensors]
+            return torch.cat(flats, dim=1), [tensor.shape for tensor in tensors], [flat.shape[1] for flat in flats]
+
+        def _unpack_joint(
+            packed: torch.Tensor,
+            shapes: list[torch.Size],
+            numels: list[int],
+        ) -> tuple[torch.Tensor, ...]:
+            outputs = []
+            offset = 0
+            for shape, numel in zip(shapes, numels, strict=True):
+                outputs.append(packed[:, offset : offset + numel].reshape(shape))
+                offset += numel
+            return tuple(outputs)
+
+        def _split_noise_pred(
+            noise_pred: torch.Tensor | tuple[torch.Tensor, ...],
+        ) -> tuple[torch.Tensor, torch.Tensor | None]:
+            has_sound = sound_latents is not None
+            if not has_sound:
+                if isinstance(noise_pred, tuple):
+                    raise ValueError("Cosmos3 video-only diffusion received tuple predictions.")
+                return noise_pred, None
+            if not isinstance(noise_pred, tuple):
+                raise ValueError("Cosmos3 multimodal diffusion expects transformer predictions as a tuple.")
+            if len(noise_pred) != 2:
+                raise ValueError(f"Cosmos3 sound diffusion expected 2 predictions, got {len(noise_pred)}.")
+            return noise_pred[0], noise_pred[1]
+
         def _step(
-            noise_pred: torch.Tensor,
+            noise_pred: torch.Tensor | tuple[torch.Tensor, ...],
             t: torch.Tensor,
             latents: torch.Tensor,
-        ) -> torch.Tensor:
-            if isinstance(noise_pred, tuple):
-                raise ValueError("Cosmos3 noise prediction must be a single tensor; got a tuple.")
+            sound_latents: torch.Tensor | None,
+        ) -> torch.Tensor | tuple[torch.Tensor, ...]:
+            video_pred, sound_pred = _split_noise_pred(noise_pred)
             if velocity_mask is not None:
-                noise_pred = noise_pred * velocity_mask
-            latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+                video_pred = video_pred * velocity_mask
+            if sound_latents is None:
+                latents = self.scheduler.step(video_pred, t, latents, return_dict=False)[0]
+            else:
+                packed_noise, shapes, numels = _pack_joint(video_pred, sound_pred)
+                packed_latents, _, _ = _pack_joint(latents, sound_latents)
+                packed_next = self.scheduler.step(packed_noise, t, packed_latents, return_dict=False)[0]
+                unpacked = _unpack_joint(packed_next, shapes, numels)
+                latents = unpacked[0]
+                if sound_latents is not None:
+                    sound_latents = unpacked[1]
             if condition_latents is not None and velocity_mask is not None:
                 latents = velocity_mask * latents + (1.0 - velocity_mask) * condition_latents
             elif image_latent is not None:
                 latents[:, :, 0:1, :, :] = image_latent
-            return latents
+            outputs = [latents]
+            if sound_latents is not None:
+                outputs.append(sound_latents)
+            return outputs[0] if len(outputs) == 1 else tuple(outputs)
+
+        def _assign_step_out(step_out: torch.Tensor | tuple[torch.Tensor, ...]) -> None:
+            nonlocal latents, sound_latents
+            if sound_latents is None:
+                assert isinstance(step_out, torch.Tensor)
+                latents = step_out
+                return
+            if not isinstance(step_out, tuple):
+                raise ValueError("Cosmos3 multimodal diffusion step returned a non-tuple result.")
+            latents = step_out[0]
+            if sound_latents is not None:
+                sound_latents = step_out[1]
 
         if cfg_parallel:
             for t in self.progress_bar(timesteps):
@@ -978,6 +1206,7 @@ class Cosmos3OmniDiffusersPipeline(
                         timestep=timestep,
                         text_ids=cond_ids,
                         text_mask=cond_mask,
+                        sound_latents=sound_latents,
                         **shared_kwargs,
                     ),
                     negative_kwargs=dict(
@@ -985,11 +1214,12 @@ class Cosmos3OmniDiffusersPipeline(
                         timestep=timestep,
                         text_ids=uncond_ids,
                         text_mask=uncond_mask,
+                        sound_latents=sound_latents,
                         **shared_kwargs,
                     ),
                     cfg_normalize=False,
                 )
-                latents = _step(noise_pred, t, latents)
+                _assign_step_out(_step(noise_pred, t, latents, sound_latents))
 
         elif do_cfg:
             cond_cache: tuple = (None, None)
@@ -1007,6 +1237,7 @@ class Cosmos3OmniDiffusersPipeline(
                     timestep=timestep,
                     text_ids=cond_ids,
                     text_mask=cond_mask,
+                    sound_latents=sound_latents,
                     **shared_kwargs,
                 )
                 if cond_cache[0] is None:
@@ -1019,6 +1250,7 @@ class Cosmos3OmniDiffusersPipeline(
                         timestep=timestep,
                         text_ids=uncond_ids,
                         text_mask=uncond_mask,
+                        sound_latents=sound_latents,
                         **shared_kwargs,
                     )
                     if uncond_cache[0] is None:
@@ -1031,7 +1263,7 @@ class Cosmos3OmniDiffusersPipeline(
                 else:
                     noise_pred = noise_cond
 
-                latents = _step(noise_pred, t, latents)
+                _assign_step_out(_step(noise_pred, t, latents, sound_latents))
 
         else:
             for t in self.progress_bar(timesteps):
@@ -1041,11 +1273,15 @@ class Cosmos3OmniDiffusersPipeline(
                     timestep=timestep,
                     text_ids=cond_ids,
                     text_mask=cond_mask,
+                    sound_latents=sound_latents,
                     **shared_kwargs,
                 )
-                latents = _step(noise_pred, t, latents)
+                _assign_step_out(_step(noise_pred, t, latents, sound_latents))
 
-        return latents
+        outputs = [latents]
+        if sound_latents is not None:
+            outputs.append(sound_latents)
+        return outputs[0] if len(outputs) == 1 else tuple(outputs)
 
     # -- Forward (main generation entry point) -------------------------------
 
@@ -1072,6 +1308,18 @@ class Cosmos3OmniDiffusersPipeline(
 
         sp = req.sampling_params
         is_t2i = self._is_t2i_request(req)
+        sound_enabled = self._is_sound_request(prompt_data, sp)
+        if sound_enabled and is_t2i:
+            raise ValueError(
+                "Cosmos3 sound generation is supported only for video outputs in "
+                "this phase; text-to-image with sound is unsupported."
+            )
+        if sound_enabled and not getattr(self.transformer, "sound_gen", False):
+            raise ValueError(
+                "Cosmos3 sound generation was requested, but the transformer was "
+                "initialized without sound modules. Check that the checkpoint config "
+                "enables sound_gen or defines sound_dim and includes sound weights."
+            )
         if negative_prompt is None:
             negative_prompt = ""
 
@@ -1163,6 +1411,13 @@ class Cosmos3OmniDiffusersPipeline(
             image_latent = None
             condition_latents = None
 
+        sound_latents = None
+        target_audio_samples = None
+        sound_sample_rate = None
+        if sound_enabled:
+            target_audio_samples, _, sound_sample_rate = self._resolve_sound_target_samples(sp, num_frames, frame_rate)
+            sound_latents, _ = self._prepare_sound_latents(target_audio_samples, generator)
+
         T_latent = latents.shape[2]
         H_latent = latents.shape[3]
         W_latent = latents.shape[4]
@@ -1184,6 +1439,7 @@ class Cosmos3OmniDiffusersPipeline(
                 uncond_mask=uncond_mask,
                 guidance_scale=guidance_scale,
                 shared_kwargs=shared_kwargs,
+                sound_latents=sound_latents,
                 velocity_mask=velocity_mask,
                 image_latent=image_latent,
                 condition_latents=condition_latents,
@@ -1204,7 +1460,11 @@ class Cosmos3OmniDiffusersPipeline(
                 samples.append(_run_diffusion(next_latents))
             latents = torch.cat(samples, dim=0)
         else:
-            latents = _run_diffusion(latents)
+            diffusion_output = _run_diffusion(latents)
+            if sound_enabled:
+                latents, sound_latents = diffusion_output
+            else:
+                latents = diffusion_output
 
         # --- Decode ---
         if _is_rank_zero():
@@ -1215,4 +1475,12 @@ class Cosmos3OmniDiffusersPipeline(
             logger.info("Video decoded in %.2fs", time.time() - decode_start)
             logger.info("Total pipeline time: %.2fs", time.time() - pipeline_start)
 
+        if sound_enabled:
+            if sound_latents is None or target_audio_samples is None or sound_sample_rate is None:
+                raise ValueError("Cosmos3 sound generation finished without sound latents.")
+            if _is_rank_zero():
+                logger.info("Decoding sound...")
+            audio = self._decode_sound_latents(sound_latents, target_audio_samples)
+            return DiffusionOutput(output={"video": video, "audio": audio, "audio_sample_rate": sound_sample_rate})
+
         return DiffusionOutput(output={"image": video} if is_t2i else {"video": video})

vllm_omni/diffusion/models/cosmos3/sound_tokenizer.py

@@ -0,0 +1,504 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Cosmos3 sound tokenizer integration."""
+
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+from typing import Any
+
+import torch
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.utils import get_local_device
+from vllm_omni.diffusion.models.progress_bar import _is_rank_zero
+
+from .audio_tokenizer import Cosmos3AVAEAudioTokenizer
+
+logger = init_logger(__name__)
+
+DEFAULT_SOUND_SAMPLE_RATE = 48000
+DEFAULT_SOUND_CHANNELS = 2
+DEFAULT_SOUND_DIM = 64
+DEFAULT_SOUND_HOP_SIZE = 1920
+DEFAULT_SOUND_LATENT_FPS = DEFAULT_SOUND_SAMPLE_RATE / DEFAULT_SOUND_HOP_SIZE
+DEFAULT_SOUND_NORMALIZE_LATENTS = False
+DEFAULT_SOUND_NORMALIZATION_TYPE = "none"
+DEFAULT_SOUND_TANH_INPUT_SCALE = 1.5
+DEFAULT_SOUND_TANH_OUTPUT_SCALE = 3.5
+DEFAULT_SOUND_TANH_CLAMP = 0.995
+SOUND_TOKENIZER_COMPONENT_NAME = "sound_tokenizer"
+SOUND_TOKENIZER_CHECKPOINT_NAME = "diffusion_pytorch_model.safetensors"
+
+
+def _pipeline_args(od_config: OmniDiffusionConfig) -> dict[str, Any]:
+    return dict(getattr(od_config, "custom_pipeline_args", None) or {})
+
+
+def _config_get(config: Any, key: str, default: Any = None) -> Any:
+    if config is None:
+        return default
+    if isinstance(config, dict):
+        return config.get(key, default)
+    if hasattr(config, "get"):
+        value = config.get(key, None)
+        return default if value is None else value
+    return getattr(config, key, default)
+
+
+def _config_path_get(config: Any, *keys: str) -> Any:
+    value = config
+    for key in keys:
+        value = _config_get(value, key, None)
+        if value is None:
+            return None
+    return value
+
+
+def _sound_tokenizer_config_from(config: Any) -> Any:
+    """Return nested ``sound_tokenizer`` config from Cosmos3 config shapes."""
+    for path in (
+        ("sound_tokenizer",),
+        ("model", "config", "sound_tokenizer"),
+        ("config", "sound_tokenizer"),
+        ("model_config", "sound_tokenizer"),
+    ):
+        value = _config_path_get(config, *path)
+        if value is not None:
+            return value
+    return None
+
+
+def _nested_sound_tokenizer_configs(od_config: OmniDiffusionConfig | None) -> tuple[Any, ...]:
+    if od_config is None:
+        return ()
+    configs = []
+    for source in (
+        getattr(od_config, "model_config", None),
+        getattr(od_config, "tf_model_config", None),
+    ):
+        config = _sound_tokenizer_config_from(source)
+        if config is not None:
+            configs.append(config)
+    return tuple(configs)
+
+
+def _first_value_from_configs(configs: tuple[Any, ...], keys: tuple[str, ...]) -> Any:
+    for config in configs:
+        for key in keys:
+            value = _config_get(config, key, None)
+            if value is not None:
+                return value
+    return None
+
+
+def _top_level_model_value(od_config: OmniDiffusionConfig | None, keys: tuple[str, ...]) -> Any:
+    if od_config is None:
+        return None
+    for source in (
+        getattr(od_config, "model_config", None),
+        getattr(od_config, "tf_model_config", None),
+    ):
+        for key in keys:
+            for path in ((key,), ("model", "config", key), ("config", key), ("model_config", key)):
+                value = _config_path_get(source, *path)
+                if value is not None:
+                    return value
+    return None
+
+
+def _custom_arg_value(args: dict[str, Any], keys: tuple[str, ...]) -> Any:
+    for key in keys:
+        value = args.get(key)
+        if value is not None:
+            return value
+    return None
+
+
+def _as_bool(value: Any) -> bool:
+    if isinstance(value, str):
+        return value.strip().lower() in {"1", "true", "yes", "on"}
+    return bool(value)
+
+
+def _as_audio_channels(value: Any) -> int:
+    if isinstance(value, bool):
+        return 2 if value else 1
+    if isinstance(value, str) and value.strip().lower() in {
+        "1",
+        "0",
+        "true",
+        "false",
+        "yes",
+        "no",
+        "on",
+        "off",
+    }:
+        return 2 if _as_bool(value) else 1
+    return int(value)
+
+
+def _resolve_model_file(path: Any, model_root: str | None) -> str | None:
+    if not path:
+        return None
+    path = str(path)
+    if "://" in path or os.path.isabs(path) or os.path.exists(path) or not model_root:
+        return path
+    return str(Path(model_root) / path)
+
+
+def _load_sound_tokenizer_component_config(config_path: str | None) -> dict[str, Any]:
+    if not config_path:
+        return {}
+    with open(config_path, encoding="utf-8") as f:
+        config = json.load(f)
+    if not isinstance(config, dict):
+        raise TypeError(f"Cosmos3 sound tokenizer config must be a JSON object, got {type(config)!r}.")
+    return config
+
+
+def _component_audio_channels(config: dict[str, Any]) -> Any:
+    if config.get("dec_out_channels") is not None:
+        return config["dec_out_channels"]
+    if config.get("audio_channels") is not None:
+        return config["audio_channels"]
+    if config.get("stereo") is not None:
+        return 2 if _as_bool(config["stereo"]) else 1
+    return None
+
+
+def _component_arch_values(config: dict[str, Any]) -> dict[str, Any]:
+    values = {
+        "sample_rate": config.get("sampling_rate", config.get("sample_rate")),
+        "audio_channels": _component_audio_channels(config),
+        "io_channels": config.get("vocoder_input_dim", config.get("io_channels", config.get("latent_ch"))),
+        "hop_size": config.get("hop_size"),
+    }
+    return {key: value for key, value in values.items() if value is not None}
+
+
+def _resolve_arch_value(
+    od_config: OmniDiffusionConfig,
+    args: dict[str, Any],
+    component_values: dict[str, Any],
+    *,
+    field: str,
+    custom_keys: tuple[str, ...],
+    nested_keys: tuple[str, ...],
+    top_level_keys: tuple[str, ...],
+    default: Any,
+    cast,
+) -> Any:
+    custom_value = _custom_arg_value(args, custom_keys)
+    component_value = component_values.get(field)
+    if component_value is not None:
+        resolved = cast(component_value)
+        if custom_value is not None and cast(custom_value) != resolved:
+            raise ValueError(
+                "Conflicting Cosmos3 sound tokenizer architecture override for "
+                f"{field}: component config has {resolved!r}, custom args have {cast(custom_value)!r}."
+            )
+        return resolved
+
+    if custom_value is not None:
+        return cast(custom_value)
+
+    nested_value = _first_value_from_configs(_nested_sound_tokenizer_configs(od_config), nested_keys)
+    if nested_value is not None:
+        return cast(nested_value)
+
+    top_value = _top_level_model_value(od_config, top_level_keys)
+    if top_value is not None:
+        return cast(top_value)
+
+    return cast(default)
+
+
+def _resolve_normalization_value(
+    od_config: OmniDiffusionConfig,
+    args: dict[str, Any],
+    *,
+    name: str,
+    default: Any,
+    aliases: tuple[str, ...] = (),
+) -> Any:
+    keys = (f"sound_{name}", name, *aliases)
+    custom_value = _custom_arg_value(args, keys)
+    if custom_value is not None:
+        return custom_value
+    nested_value = _first_value_from_configs(_nested_sound_tokenizer_configs(od_config), (name, *aliases))
+    return default if nested_value is None else nested_value
+
+
+def get_sound_config_value(
+    od_config: OmniDiffusionConfig,
+    name: str,
+    default: Any,
+    aliases: tuple[str, ...] = (),
+) -> Any:
+    # Backward-compatible generic accessor.  Prefer the more specific helpers
+    # below for Cosmos3 sound tokenizer fields so precedence stays explicit.
+    keys = (name, *aliases)
+    for config in (
+        _pipeline_args(od_config),
+        getattr(od_config, "model_config", None),
+        getattr(od_config, "tf_model_config", None),
+    ):
+        if config is None:
+            continue
+        for key in keys:
+            if hasattr(config, "get"):
+                value = config.get(key, None)
+            else:
+                value = getattr(config, key, None)
+            if value is not None:
+                return value
+    return default
+
+
+def get_sound_sample_rate(od_config: OmniDiffusionConfig) -> int:
+    args = _pipeline_args(od_config)
+    return _resolve_arch_value(
+        od_config,
+        args,
+        {},
+        field="sample_rate",
+        custom_keys=("sound_sample_rate", "sample_rate"),
+        nested_keys=("sample_rate", "sampling_rate"),
+        top_level_keys=("sound_sample_rate", "sample_rate"),
+        default=DEFAULT_SOUND_SAMPLE_RATE,
+        cast=int,
+    )
+
+
+def get_sound_channels(od_config: OmniDiffusionConfig) -> int:
+    args = _pipeline_args(od_config)
+    return _resolve_arch_value(
+        od_config,
+        args,
+        {},
+        field="audio_channels",
+        custom_keys=("sound_audio_channels", "audio_channels", "stereo"),
+        nested_keys=("audio_channels", "dec_out_channels", "stereo"),
+        top_level_keys=("sound_audio_channels", "audio_channels", "stereo"),
+        default=DEFAULT_SOUND_CHANNELS,
+        cast=_as_audio_channels,
+    )
+
+
+def get_sound_hop_size(od_config: OmniDiffusionConfig) -> int:
+    args = _pipeline_args(od_config)
+    return _resolve_arch_value(
+        od_config,
+        args,
+        {},
+        field="hop_size",
+        custom_keys=("sound_hop_size", "hop_size"),
+        nested_keys=("hop_size",),
+        top_level_keys=("sound_hop_size", "hop_size"),
+        default=DEFAULT_SOUND_HOP_SIZE,
+        cast=int,
+    )
+
+
+class Cosmos3SoundTokenizer:
+    """Thin adapter around the local AVAE tokenizer implementation."""
+
+    def __init__(self, tokenizer: Any) -> None:
+        self.tokenizer = tokenizer
+        self.sample_rate = int(getattr(tokenizer, "sample_rate", DEFAULT_SOUND_SAMPLE_RATE))
+        self.audio_channels = int(getattr(tokenizer, "audio_channels", DEFAULT_SOUND_CHANNELS))
+        self.latent_ch = int(getattr(tokenizer, "latent_ch", DEFAULT_SOUND_DIM))
+        self.hop_size = int(getattr(tokenizer, "temporal_compression_factor", DEFAULT_SOUND_HOP_SIZE))
+        if self.hop_size <= 0:
+            raise ValueError(f"Cosmos3 sound tokenizer hop_size must be positive, got {self.hop_size}.")
+        self.latent_fps = float(self.sample_rate) / float(self.hop_size)
+
+    @classmethod
+    def from_config(cls, od_config: OmniDiffusionConfig) -> Cosmos3SoundTokenizer:
+        args = _pipeline_args(od_config)
+        model_path = getattr(od_config, "model", None)
+        explicit_avae_path = (
+            args.get("sound_tokenizer_path")
+            or args.get("avae_path")
+            or args.get("cosmos3_avae_path")
+            or os.environ.get("COSMOS3_SOUND_TOKENIZER_PATH")
+        )
+        explicit_config_path = args.get("sound_tokenizer_config_path") or os.environ.get(
+            "COSMOS3_SOUND_TOKENIZER_CONFIG_PATH"
+        )
+
+        model_root = str(model_path) if model_path and os.path.isdir(model_path) else None
+        if model_root is None and model_path and not explicit_avae_path:
+            from huggingface_hub import snapshot_download
+
+            model_root = snapshot_download(
+                repo_id=str(model_path),
+                revision=getattr(od_config, "revision", None),
+                allow_patterns=[
+                    f"{SOUND_TOKENIZER_COMPONENT_NAME}/config.json",
+                    f"{SOUND_TOKENIZER_COMPONENT_NAME}/{SOUND_TOKENIZER_CHECKPOINT_NAME}",
+                ],
+            )
+
+        if explicit_avae_path:
+            avae_path = _resolve_model_file(explicit_avae_path, model_root)
+        else:
+            tokenizer_dir = Path(model_root) / SOUND_TOKENIZER_COMPONENT_NAME if model_root else None
+            candidate = tokenizer_dir / SOUND_TOKENIZER_CHECKPOINT_NAME if tokenizer_dir else None
+            avae_path = str(candidate) if candidate and candidate.exists() else None
+
+        if not avae_path:
+            raise ValueError(
+                "Cosmos3 sound generation was requested, but no AVAE sound "
+                "tokenizer checkpoint was provided. Set "
+                "custom_pipeline_args['sound_tokenizer_path'] or "
+                "COSMOS3_SOUND_TOKENIZER_PATH, or include "
+                f"{SOUND_TOKENIZER_COMPONENT_NAME}/{SOUND_TOKENIZER_CHECKPOINT_NAME} under the model path."
+            )
+
+        config_path = _resolve_model_file(explicit_config_path, model_root)
+        if config_path is None and model_root:
+            candidate = Path(model_root) / SOUND_TOKENIZER_COMPONENT_NAME / "config.json"
+            config_path = str(candidate) if candidate.exists() else None
+        component_config = _load_sound_tokenizer_component_config(config_path)
+        component_values = _component_arch_values(component_config)
+
+        sample_rate = _resolve_arch_value(
+            od_config,
+            args,
+            component_values,
+            field="sample_rate",
+            custom_keys=("sound_sample_rate", "sample_rate"),
+            nested_keys=("sample_rate", "sampling_rate"),
+            top_level_keys=("sound_sample_rate", "sample_rate"),
+            default=DEFAULT_SOUND_SAMPLE_RATE,
+            cast=int,
+        )
+        audio_channels = _resolve_arch_value(
+            od_config,
+            args,
+            component_values,
+            field="audio_channels",
+            custom_keys=("sound_audio_channels", "audio_channels", "stereo"),
+            nested_keys=("audio_channels", "dec_out_channels", "stereo"),
+            top_level_keys=("sound_audio_channels", "audio_channels", "stereo"),
+            default=DEFAULT_SOUND_CHANNELS,
+            cast=_as_audio_channels,
+        )
+        sound_dim = _resolve_arch_value(
+            od_config,
+            args,
+            component_values,
+            field="io_channels",
+            custom_keys=("sound_dim", "io_channels", "latent_ch"),
+            nested_keys=("io_channels", "vocoder_input_dim", "latent_ch"),
+            top_level_keys=("sound_dim",),
+            default=DEFAULT_SOUND_DIM,
+            cast=int,
+        )
+        hop_size = _resolve_arch_value(
+            od_config,
+            args,
+            component_values,
+            field="hop_size",
+            custom_keys=("sound_hop_size", "hop_size"),
+            nested_keys=("hop_size",),
+            top_level_keys=("sound_hop_size", "hop_size"),
+            default=DEFAULT_SOUND_HOP_SIZE,
+            cast=int,
+        )
+        normalize_latents = _as_bool(
+            _resolve_normalization_value(
+                od_config,
+                args,
+                name="normalize_latents",
+                default=DEFAULT_SOUND_NORMALIZE_LATENTS,
+            )
+        )
+        normalization_type = str(
+            _resolve_normalization_value(
+                od_config,
+                args,
+                name="normalization_type",
+                default=DEFAULT_SOUND_NORMALIZATION_TYPE,
+            )
+        )
+        tanh_input_scale = float(
+            _resolve_normalization_value(
+                od_config,
+                args,
+                name="tanh_input_scale",
+                default=DEFAULT_SOUND_TANH_INPUT_SCALE,
+            )
+        )
+        tanh_output_scale = float(
+            _resolve_normalization_value(
+                od_config,
+                args,
+                name="tanh_output_scale",
+                default=DEFAULT_SOUND_TANH_OUTPUT_SCALE,
+            )
+        )
+        tanh_clamp = float(
+            _resolve_normalization_value(
+                od_config,
+                args,
+                name="tanh_clamp",
+                default=DEFAULT_SOUND_TANH_CLAMP,
+            )
+        )
+        tokenizer = Cosmos3AVAEAudioTokenizer(
+            checkpoint_path=str(avae_path),
+            config_path=config_path,
+            sample_rate=sample_rate,
+            audio_channels=audio_channels,
+            io_channels=sound_dim,
+            hop_size=hop_size,
+            normalize_latents=normalize_latents,
+            normalization_type=normalization_type,
+            tanh_input_scale=tanh_input_scale,
+            tanh_output_scale=tanh_output_scale,
+            tanh_clamp=tanh_clamp,
+            dtype=getattr(od_config, "dtype", torch.bfloat16),
+            device=get_local_device(),
+        )
+        if _is_rank_zero():
+            logger.info(
+                "Loaded Cosmos3 AVAE sound tokenizer from %s "
+                "(sr=%d, channels=%d, latent_ch=%d, hop=%d, latent_fps=%.3f)",
+                avae_path,
+                sample_rate,
+                audio_channels,
+                sound_dim,
+                hop_size,
+                float(sample_rate) / float(hop_size),
+            )
+        return cls(tokenizer)
+
+    def get_latent_num_samples(self, num_audio_samples: int) -> int:
+        return int(self.tokenizer.get_latent_num_samples(num_audio_samples))
+
+    def get_audio_num_samples(self, num_latent_samples: int) -> int:
+        return int(self.tokenizer.get_audio_num_samples(num_latent_samples))
+
+    @torch.no_grad()
+    def decode(self, latents: torch.Tensor) -> torch.Tensor:
+        """Decode sound latents.
+
+        Args:
+            latents: ``[B, C, T]`` or ``[C, T]`` tensor.
+
+        Returns:
+            ``[B, audio_channels, N]`` tensor for batched input, or
+            ``[audio_channels, N]`` for unbatched input.
+        """
+        squeeze = latents.ndim == 2
+        if squeeze:
+            latents = latents.unsqueeze(0)
+        audio = self.tokenizer.decode(latents)
+        audio = audio.clamp(-1.0, 1.0)
+        return audio.squeeze(0) if squeeze else audio

vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py

@@ -76,6 +76,66 @@ def _tf_config_get(config: Any, key: str, default: Any) -> Any:
     return getattr(config, key, default)
 
 
+def _nested_get(value: Any, key: str) -> Any:
+    if isinstance(value, dict):
+        if key in value:
+            return value[key]
+        for child in value.values():
+            found = _nested_get(child, key)
+            if found is not None:
+                return found
+    elif isinstance(value, list | tuple):
+        for child in value:
+            found = _nested_get(child, key)
+            if found is not None:
+                return found
+    return None
+
+
+def _od_config_get(od_config: Any, key: str, default: Any = None) -> Any:
+    """Read Cosmos3 options from runtime, model, or transformer config."""
+    if od_config is None:
+        return default
+    for attr in ("custom_pipeline_args", "model_config"):
+        source = getattr(od_config, attr, None) or {}
+        if isinstance(source, dict):
+            if key in source:
+                return source[key]
+            found = _nested_get(source, key)
+            if found is not None:
+                return found
+    tf_model_config = getattr(od_config, "tf_model_config", None)
+    if isinstance(tf_model_config, dict):
+        if key in tf_model_config:
+            return tf_model_config[key]
+        found = _nested_get(tf_model_config, key)
+        if found is not None:
+            return found
+    value = _tf_config_get(tf_model_config, key, None)
+    return default if value is None else value
+
+
+def _as_bool(value: Any) -> bool:
+    if isinstance(value, str):
+        return value.strip().lower() in {"1", "true", "yes", "on"}
+    return bool(value)
+
+
+def resolve_sound_gen(od_config: Any) -> bool:
+    """Capability gate shared by the pipeline and transformer.
+
+    Explicit ``sound_gen`` flag wins (including an explicit False);
+    otherwise infer from the presence of any sound-width key in od_config.
+    """
+    sound_gen_value = _od_config_get(od_config, "sound_gen", None)
+    if sound_gen_value is not None:
+        return _as_bool(sound_gen_value)
+    for key in ("sound_dim", "io_channels", "vocoder_input_dim", "latent_ch"):
+        if _od_config_get(od_config, key, None) is not None:
+            return True
+    return False
+
+
 # ---------------------------------------------------------------------------
 # Rotary Position Embeddings (mRoPE)
 # ---------------------------------------------------------------------------
@@ -138,6 +198,28 @@ def compute_mrope_position_ids_vision(
     return mrope_ids, next_offset
 
 
+def compute_mrope_position_ids_sound(
+    grid_t: int,
+    temporal_offset: int | float,
+    sound_latent_fps: float,
+    base_fps: float = 24.0,
+    temporal_compression_factor_sound: int = 1,
+    enable_fps_modulation: bool = True,
+) -> tuple[torch.Tensor, int | float]:
+    """Generate mRoPE IDs for sound tokens as a (T, 1, 1) grid."""
+    return compute_mrope_position_ids_vision(
+        grid_t=grid_t,
+        grid_h=1,
+        grid_w=1,
+        temporal_offset=temporal_offset,
+        fps=sound_latent_fps,
+        base_fps=base_fps,
+        temporal_compression_factor=temporal_compression_factor_sound,
+        base_temporal_compression_factor=temporal_compression_factor_sound,
+        enable_fps_modulation=enable_fps_modulation,
+    )
+
+
 class Qwen3VLTextRotaryEmbedding(nn.Module):
     """Multi-dimensional rotary position embedding for Qwen3-VL."""
 
@@ -839,6 +921,9 @@ class Cosmos3VFMTransformer(nn.Module):
         od_config: OmniDiffusionConfig,
         *,
         temporal_compression_factor: int | None = None,
+        sound_gen: bool = False,
+        sound_dim: int | None = None,
+        sound_latent_fps: float | None = None,
     ) -> None:
         super().__init__()
         model_config = od_config.tf_model_config
@@ -859,9 +944,22 @@ class Cosmos3VFMTransformer(nn.Module):
         self.latent_channel_size = int(_tf_config_get(model_config, "latent_channel", 48))
         self.timestep_scale = float(_tf_config_get(model_config, "timestep_scale", 0.001))
         self.base_fps = float(_tf_config_get(model_config, "base_fps", 24.0))
+        self.sound_gen = sound_gen
+        self.sound_dim = sound_dim
+        self.sound_latent_fps = sound_latent_fps
+
+        if self.sound_gen and (sound_dim is None or sound_latent_fps is None):
+            raise ValueError(
+                "Cosmos3VFMTransformer requires an explicit sound_dim and sound_latent_fps when sound_gen is True; "
+                "the pipeline must pass Cosmos3SoundTokenizer.latent_ch so the audio projection "
+                "layers are sized from the authoritative AVAE latent width."
+            )
         if temporal_compression_factor is None:
             temporal_compression_factor = _tf_config_get(model_config, "temporal_compression_factor", 4)
         self.temporal_compression_factor = int(temporal_compression_factor)
+        self.temporal_compression_factor_sound = int(
+            _tf_config_get(model_config, "temporal_compression_factor_sound", 1)
+        )
         self.enable_fps_modulation = bool(_tf_config_get(model_config, "enable_fps_modulation", True))
         self.temporal_modality_margin = int(
             _tf_config_get(
@@ -894,6 +992,10 @@ class Cosmos3VFMTransformer(nn.Module):
         self.proj_in = nn.Linear(self.patch_latent_dim, self.hidden_size)
         self.proj_out = nn.Linear(self.hidden_size, self.patch_latent_dim)
         self.time_embedder = TimestepEmbedder(self.hidden_size, target_dtype=dtype)
+        if self.sound_gen:
+            self.audio_proj_in = nn.Linear(self.sound_dim, self.hidden_size)
+            self.audio_proj_out = nn.Linear(self.hidden_size, self.sound_dim)
+            self.audio_modality_embed = nn.Parameter(torch.zeros(self.hidden_size))
 
         self.gen_layers = nn.ModuleList(
             [
@@ -962,6 +1064,21 @@ class Cosmos3VFMTransformer(nn.Module):
             x = x[:, :, :, :h, :w]
         return x
 
+    def pack_sound(self, sound_latents: torch.Tensor) -> torch.Tensor:
+        """[B, C_sound, T_sound] -> [B, T_sound, C_sound]."""
+        if sound_latents.ndim != 3:
+            raise ValueError(f"Cosmos3 sound latents must have shape [B, C, T], got {tuple(sound_latents.shape)}.")
+        if sound_latents.shape[1] != self.sound_dim:
+            raise ValueError(
+                f"Cosmos3 sound latent channel mismatch: expected {self.sound_dim}, got {sound_latents.shape[1]}."
+            )
+        return sound_latents.permute(0, 2, 1).contiguous()
+
+    @staticmethod
+    def unpack_sound(tokens: torch.Tensor) -> torch.Tensor:
+        """[B, T_sound, C_sound] -> [B, C_sound, T_sound]."""
+        return tokens.permute(0, 2, 1).contiguous()
+
     # -- RoPE computation ----------------------------------------------------
 
     def _compute_rope_freqs(
@@ -973,12 +1090,14 @@ class Cosmos3VFMTransformer(nn.Module):
         fps: float | None,
         device: torch.device,
         dtype: torch.dtype,
+        t_sound: int | None = None,
     ) -> tuple[tuple[torch.Tensor, torch.Tensor], tuple[torch.Tensor, torch.Tensor]]:
         """Compute mRoPE cos/sin for UND text and GEN media pathways."""
         B = text_mask.shape[0]
         S_text = text_mask.shape[1]
         text_lengths = text_mask.sum(dim=1).long()
         effective_fps = fps if fps is not None and t > 1 else None
+        sound_frames = int(t_sound or 0)
 
         text_pos_list = []
         gen_pos_list = []
@@ -996,6 +1115,21 @@ class Cosmos3VFMTransformer(nn.Module):
                 temporal_compression_factor=self.temporal_compression_factor,
                 enable_fps_modulation=self.enable_fps_modulation,
             )
+            gen_positions = [v_pos]
+            if sound_frames > 0:
+                s_pos, _ = compute_mrope_position_ids_sound(
+                    sound_frames,
+                    temporal_offset=media_temporal_offset,
+                    sound_latent_fps=self.sound_latent_fps,
+                    base_fps=self.base_fps,
+                    temporal_compression_factor_sound=getattr(self, "temporal_compression_factor_sound", 1),
+                    enable_fps_modulation=self.enable_fps_modulation,
+                )
+                gen_positions.append(s_pos)
+            pos_dtype = gen_positions[0].dtype
+            for pos in gen_positions[1:]:
+                pos_dtype = torch.promote_types(pos_dtype, pos.dtype)
+            v_pos = torch.cat([pos.to(pos_dtype) for pos in gen_positions], dim=1)
             if real_len < S_text:
                 t_pos = torch.cat(
                     [t_pos, torch.zeros(3, S_text - real_len, dtype=t_pos.dtype)],
@@ -1026,16 +1160,31 @@ class Cosmos3VFMTransformer(nn.Module):
     def _validate_gen_sequence_parallel(
         *,
         s_gen: int,
+        s_video: int,
+        s_sound: int,
+        has_sound: bool,
         ulysses_size: int,
     ) -> None:
         if ulysses_size <= 1 or s_gen % ulysses_size == 0:
             return
 
+        detail_parts = [f"video tokens {s_video}"]
+        if has_sound:
+            detail_parts.append(f"sound tokens {s_sound}")
+        detail = " = " + " + ".join(detail_parts) if len(detail_parts) > 1 else ""
         adjust_detail = (
-            "Adjust the spatial resolution so that t * ceil(h/patch) * ceil(w/patch) is a multiple of ulysses_degree."
+            "Adjust the spatial resolution, frame count, sound duration, "
+            "or sound latent FPS so the combined media sequence is a "
+            "multiple of ulysses_degree."
+            if has_sound
+            else (
+                "Adjust the spatial resolution so that "
+                "t * ceil(h/patch) * ceil(w/patch) is a multiple "
+                "of ulysses_degree."
+            )
         )
         raise ValueError(
-            f"GEN sequence length ({s_gen} video tokens) must be divisible by "
+            f"GEN sequence length ({s_gen}{detail}) must be divisible by "
             f"ulysses_degree ({ulysses_size}). {adjust_detail}"
         )
 
@@ -1049,9 +1198,10 @@ class Cosmos3VFMTransformer(nn.Module):
         text_mask: torch.Tensor,
         video_shape: tuple[int, int, int],
         fps: float | None = None,
+        sound_latents: torch.Tensor | None = None,
         noisy_frame_mask: torch.Tensor | None = None,
         **kwargs,
-    ) -> torch.Tensor:
+    ) -> torch.Tensor | tuple[torch.Tensor, ...]:
         """
         Args:
             hidden_states: [B, C, t, h, w] noisy latents
@@ -1060,13 +1210,15 @@ class Cosmos3VFMTransformer(nn.Module):
             text_mask: [B, S_text] attention mask (1=real, 0=pad)
             video_shape: (t, h, w) in latent space
             fps: video frame rate for temporal mRoPE modulation
+            sound_latents: Optional [B, C_sound, T_sound] noisy sound latents.
             noisy_frame_mask: Optional [B, 1, t, 1, 1] mask where 1=noisy (add
                 timestep embedding, predict velocity) and 0=conditioned (clean
                 context, skip timestep embedding).  None means all frames noisy
                 (T2V mode).
 
         Returns:
-            [B, C, t, h, w] velocity prediction.
+            [B, C, t, h, w] velocity prediction, or
+            tuple outputs in video, sound order when sound latents are provided.
         """
         t, h, w = video_shape
         hp, wp, _, _ = self._pad_to_patch_size(h, w)
@@ -1078,12 +1230,31 @@ class Cosmos3VFMTransformer(nn.Module):
                 f"Cosmos3 requires identical real text lengths within a batch "
                 f"(got min={min_real_len}, max={max_real_len})."
             )
+        has_sound = sound_latents is not None
+        if has_sound and not self.sound_gen:
+            raise ValueError(
+                "Cosmos3 sound generation was requested, but this transformer "
+                "was initialized without sound modules. Check that the "
+                "transformer config enables sound_gen or defines sound_dim."
+            )
 
         # Query Ulysses state at runtime
         ulysses_size, _, _ = _get_ulysses_state()
 
         # Patchify latents and project to hidden space
         hidden_video = self.proj_in(self.patchify(hidden_states, t, h, w))
+        s_video = hidden_video.shape[1]
+        s_sound = 0
+        hidden_sound = None
+        if sound_latents is not None:
+            if sound_latents.shape[0] != hidden_states.shape[0]:
+                raise ValueError(
+                    "Cosmos3 sound and video batch sizes must match: "
+                    f"video={hidden_states.shape[0]}, sound={sound_latents.shape[0]}."
+                )
+            hidden_sound = self.audio_proj_in(self.pack_sound(sound_latents))
+            hidden_sound = hidden_sound + self.audio_modality_embed.to(hidden_sound.dtype)
+            s_sound = hidden_sound.shape[1]
 
         # Timestep embedding (fp32 for precision).
         # For I2V: only add to noisy tokens, not conditioned ones.
@@ -1106,7 +1277,12 @@ class Cosmos3VFMTransformer(nn.Module):
         else:
             hidden_video = hidden_video + time_embed.unsqueeze(1)
 
-        hidden_gen = hidden_video
+        if hidden_sound is not None:
+            hidden_sound = hidden_sound + time_embed.unsqueeze(1)
+        hidden_parts = [hidden_video]
+        if hidden_sound is not None:
+            hidden_parts.append(hidden_sound)
+        hidden_gen = torch.cat(hidden_parts, dim=1)
 
         # Run UND pathway once and cache K/V (replicated across all ranks)
         if self.cached_kv is None:
@@ -1118,6 +1294,7 @@ class Cosmos3VFMTransformer(nn.Module):
                 fps,
                 hidden_states.device,
                 hidden_states.dtype,
+                t_sound=s_sound,
             )
             cached_kv_full = self.language_model(text_ids, freqs_und)
             self.cached_freqs_gen = freqs_gen
@@ -1133,6 +1310,9 @@ class Cosmos3VFMTransformer(nn.Module):
             raise RuntimeError("Cosmos3 GEN cache was not initialized before running GEN layers.")
         self._validate_gen_sequence_parallel(
             s_gen=hidden_gen.shape[1],
+            s_video=s_video,
+            s_sound=s_sound,
+            has_sound=has_sound,
             ulysses_size=ulysses_size,
         )
         freqs_cos, freqs_sin = self.cached_freqs_gen
@@ -1166,7 +1346,21 @@ class Cosmos3VFMTransformer(nn.Module):
 
         # Final norm and project back to latent space
         hidden_gen = self.norm_moe_gen(hidden_gen)
-        return self.unpatchify(self.proj_out(hidden_gen), t, h, w)
+        if not has_sound:
+            return self.unpatchify(self.proj_out(hidden_gen), t, h, w)
+
+        split_sizes = [s_video]
+        if has_sound:
+            split_sizes.append(s_sound)
+        split_hidden = hidden_gen.split(split_sizes, dim=1)
+        hidden_video = split_hidden[0]
+        video_pred = self.unpatchify(self.proj_out(hidden_video), t, h, w)
+        outputs: list[torch.Tensor] = [video_pred]
+        split_idx = 1
+        if has_sound:
+            hidden_sound = split_hidden[split_idx]
+            outputs.append(self.unpack_sound(self.audio_proj_out(hidden_sound)))
+        return tuple(outputs)
 
     def post_load_weights(self) -> None:
         """Post-load processing: ensure correct dtypes."""

vllm_omni/diffusion/models/dreamzero/__init__.py

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

vllm_omni/diffusion/models/dreamzero/action_encoder.py

@@ -0,0 +1,100 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Action encoder/decoder for DreamZero."""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def swish(x: torch.Tensor) -> torch.Tensor:
+    """swish activation: x * sigmoid(x)"""
+    return x * torch.sigmoid(x)
+
+
+class SinusoidalPositionalEncoding(nn.Module):
+    """Sinusoidal encoding: (B, T) timesteps → (B, T, dim)"""
+
+    def __init__(self, embedding_dim: int) -> None:
+        super().__init__()
+        self.embedding_dim = embedding_dim
+
+    def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
+        timesteps = timesteps.float()
+        half_dim = self.embedding_dim // 2
+        exponent = -torch.arange(half_dim, dtype=torch.float, device=timesteps.device) * (
+            torch.log(torch.tensor(10000.0)) / half_dim
+        )
+        freqs = timesteps.unsqueeze(-1) * exponent.exp()
+        return torch.cat([torch.sin(freqs), torch.cos(freqs)], dim=-1)
+
+
+class CategorySpecificLinear(nn.Module):
+    """Per-category linear: W[cat_id] @ x + b[cat_id]
+
+    Attributes:
+        W: (num_categories, input_dim, hidden_dim)  — note: 0.02 * randn init
+        b: (num_categories, hidden_dim)              — zero init
+    """
+
+    def __init__(self, num_categories: int, input_dim: int, hidden_dim: int) -> None:
+        super().__init__()
+        self.W = nn.Parameter(0.02 * torch.randn(num_categories, input_dim, hidden_dim))
+        self.b = nn.Parameter(torch.zeros(num_categories, hidden_dim))
+
+    def forward(self, x: torch.Tensor, cat_ids: torch.Tensor) -> torch.Tensor:
+        selected_W = self.W[cat_ids]
+        selected_b = self.b[cat_ids]
+        return torch.bmm(x, selected_W) + selected_b.unsqueeze(1)
+
+
+class CategorySpecificMLP(nn.Module):
+    """Two-layer MLP: layer1 (relu) → layer2"""
+
+    def __init__(self, num_categories: int, input_dim: int, hidden_dim: int, output_dim: int) -> None:
+        super().__init__()
+        self.layer1 = CategorySpecificLinear(num_categories, input_dim, hidden_dim)
+        self.layer2 = CategorySpecificLinear(num_categories, hidden_dim, output_dim)
+
+    def forward(self, x: torch.Tensor, cat_ids: torch.Tensor) -> torch.Tensor:
+        hidden = F.relu(self.layer1(x, cat_ids))
+        return self.layer2(hidden, cat_ids)
+
+
+class MultiEmbodimentActionEncoder(nn.Module):
+    """Encode actions with embodiment-specific weights + sinusoidal timestep.
+
+    Flow: actions → W1 → concat(a_emb, pos_enc(timesteps)) → W2 (swish) → W3
+
+    Args:
+        action_dim: action vector dimension (e.g. 32)
+        hidden_size: output/hidden dimension (e.g. 5120 = model dim)
+        num_embodiments: number of robot types (e.g. 32)
+    """
+
+    def __init__(self, action_dim: int, hidden_size: int, num_embodiments: int) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.W1 = CategorySpecificLinear(num_embodiments, action_dim, hidden_size)
+        self.W2 = CategorySpecificLinear(num_embodiments, 2 * hidden_size, hidden_size)
+        self.W3 = CategorySpecificLinear(num_embodiments, hidden_size, hidden_size)
+        self.pos_encoding = SinusoidalPositionalEncoding(hidden_size)
+
+    def forward(self, actions: torch.Tensor, timesteps: torch.Tensor, cat_ids: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            actions:   (B, T, action_dim)
+            timesteps: (B, T) — per-token timestep
+            cat_ids:   (B,)   — embodiment id per sample
+        Returns:
+            (B, T, hidden_size)
+        """
+        a_emb = self.W1(actions, cat_ids)
+        tau_emb = self.pos_encoding(timesteps).to(dtype=a_emb.dtype)
+        x = torch.cat([a_emb, tau_emb], dim=-1)
+        x = swish(self.W2(x, cat_ids))
+        x = self.W3(x, cat_ids)
+        return x

vllm_omni/diffusion/models/dreamzero/causal_wan_model.py

@@ -0,0 +1,954 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""CausalWanModel — 40-layer DiT with causal attention and KV cache.
+
+Key differences from WanTransformer3DModel:
+- Causal self-attention (new frames only see history)
+- KV cache for streaming inference
+- Action/state token support (appended after video tokens)
+- Extended RoPE with action/state-specific frequencies
+- Inference-only forward with KV cache
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Any
+
+import torch
+import torch.nn as nn
+from vllm.distributed import (
+    get_tensor_model_parallel_rank,
+    get_tensor_model_parallel_world_size,
+    get_tp_group,
+)
+from vllm.model_executor.layers.conv import Conv3dLayer
+from vllm.model_executor.layers.linear import (
+    ColumnParallelLinear,
+    RowParallelLinear,
+)
+from vllm.model_executor.utils import set_weight_attrs
+
+from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.models.dreamzero.action_encoder import (
+    CategorySpecificMLP,
+    MultiEmbodimentActionEncoder,
+)
+
+# ── RoPE utilities ──────────────────────────────────────────────────
+
+
+def sinusoidal_embedding_1d(dim: int, position: torch.Tensor) -> torch.Tensor:
+    """Sinusoidal positional embedding for timesteps."""
+    if dim % 2 != 0:
+        raise ValueError(f"dim must be even, got {dim}.")
+    half = dim // 2
+    position = position.type(torch.float64)
+    sinusoid = torch.outer(
+        position,
+        torch.pow(10000, -torch.arange(half, dtype=position.dtype, device=position.device).div(half)),
+    )
+    x = torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1)
+    return x
+
+
+def rope_params(max_seq_len: int, dim: int) -> torch.Tensor:
+    """Precompute complex-valued RoPE frequencies (polar form).
+    Returns: complex tensor [max_seq_len, dim // 2]
+    """
+    if dim % 2 != 0:
+        raise ValueError(f"dim must be even, got {dim}.")
+    freqs = torch.outer(
+        torch.arange(max_seq_len),
+        1.0 / torch.pow(10000, torch.arange(0, dim, 2).to(torch.float64).div(dim)),
+    )
+    freqs = torch.polar(torch.ones_like(freqs), freqs)
+    return freqs
+
+
+def rope_apply(x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
+    """Apply RoPE to x using precomputed complex freqs."""
+    B, seq_len, n, _ = x.shape
+    x = torch.view_as_complex(x.to(torch.float64).reshape(B, seq_len, n, -1, 2))
+    freqs = freqs.unsqueeze(0)
+    x = torch.view_as_real(x * freqs).flatten(3)
+    return x
+
+
+def rope_action_apply(
+    x: torch.Tensor,
+    freqs: torch.Tensor,
+    freqs_action: torch.Tensor,
+    freqs_state: torch.Tensor,
+    action_register_length: int | None,
+    num_action_per_block: int = 32,
+    num_state_per_block: int = 1,
+) -> torch.Tensor:
+    """RoPE with action/state frequency tables for multi-step sequences."""
+    B, seq_len, n, _ = x.shape
+    x = torch.view_as_complex(x.to(torch.float64).reshape(B, seq_len, n, -1, 2))
+    if action_register_length is not None:
+        if num_action_per_block is None:
+            raise ValueError("num_action_per_block is required when action_register_length is set.")
+        if num_state_per_block is None:
+            raise ValueError("num_state_per_block is required when action_register_length is set.")
+        chunk_size = action_register_length // (num_action_per_block + num_state_per_block)
+        freqs_1d_action = freqs_action[: chunk_size * num_action_per_block].view(
+            chunk_size * num_action_per_block, 1, -1
+        )
+        freqs_1d_state = freqs_state[: chunk_size * num_state_per_block].view(chunk_size * num_state_per_block, 1, -1)
+        freqs = torch.cat([freqs, freqs_1d_action, freqs_1d_state], dim=0)
+    freqs = freqs.unsqueeze(0)
+    x = torch.view_as_real(x * freqs).flatten(3)
+    return x
+
+
+def causal_rope_action_apply(
+    x: torch.Tensor,
+    freqs: torch.Tensor,
+    freqs_action: torch.Tensor,
+    freqs_state: torch.Tensor,
+    action_register_length: int | None,
+    num_action_per_block: int,
+    num_state_per_block: int,
+    action_state_index: int,
+) -> torch.Tensor:
+    """RoPE for single inference step (causal / KV-cache mode)."""
+    B, seq_len, n, _ = x.shape
+    x = torch.view_as_complex(x.to(torch.float64).reshape(B, seq_len, n, -1, 2))
+    if action_register_length is not None:
+        expected_length = num_action_per_block + num_state_per_block
+        if action_register_length != expected_length:
+            raise ValueError(
+                f"action_register_length must equal num_action_per_block + num_state_per_block "
+                f"({expected_length}), got {action_register_length}."
+            )
+        freqs_action = freqs_action[
+            action_state_index * num_action_per_block : (action_state_index + 1) * num_action_per_block
+        ]
+        freqs_state = freqs_state[
+            action_state_index * num_state_per_block : (action_state_index + 1) * num_state_per_block
+        ]
+        freqs_1d = torch.cat([freqs_action, freqs_state], dim=0).view(action_register_length, 1, -1)
+        freqs = torch.cat([freqs, freqs_1d], dim=0)
+    freqs = freqs.unsqueeze(0)
+    x = torch.view_as_real(x * freqs).flatten(3)
+    return x
+
+
+# ── Normalization ───────────────────────────────────────────────────
+
+
+class WanLayerNorm(nn.LayerNorm):
+    """LayerNorm wrapper used by DreamZero blocks."""
+
+    def __init__(self, dim: int, eps: float = 1e-6, elementwise_affine: bool = False) -> None:
+        super().__init__(dim, eps=eps, elementwise_affine=elementwise_affine)
+
+
+class DistributedRMSNorm(nn.Module):
+    """RMSNorm that computes global RMS across tensor parallel ranks."""
+
+    def __init__(self, hidden_size: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        set_weight_attrs(self.weight, {"weight_loader": self.weight_loader})
+
+    def weight_loader(self, param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
+        if param.shape == loaded_weight.shape:
+            param.data.copy_(loaded_weight)
+            return
+
+        tp_size = get_tensor_model_parallel_world_size()
+        if loaded_weight.shape[0] % tp_size != 0:
+            raise ValueError(
+                f"Cannot shard RMSNorm weight of shape {tuple(loaded_weight.shape)} across tp_size={tp_size}."
+            )
+
+        shard_size = loaded_weight.shape[0] // tp_size
+        start_idx = get_tensor_model_parallel_rank() * shard_size
+        shard = loaded_weight.narrow(0, start_idx, shard_size)
+        if param.shape != shard.shape:
+            raise ValueError(f"RMSNorm shard shape mismatch: param={tuple(param.shape)}, shard={tuple(shard.shape)}.")
+        param.data.copy_(shard)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        tp_size = get_tensor_model_parallel_world_size()
+        x_float = x.float()
+        local_sum_sq = x_float.pow(2).sum(dim=-1, keepdim=True)
+        local_count = x.shape[-1]
+
+        if tp_size > 1:
+            global_sum_sq = local_sum_sq.clone()
+            torch.distributed.all_reduce(global_sum_sq, group=get_tp_group().device_group)
+            global_count = local_count * tp_size
+        else:
+            global_sum_sq = local_sum_sq
+            global_count = local_count
+
+        mean_sq = global_sum_sq / global_count
+        return (x_float * torch.rsqrt(mean_sq + self.eps)).type_as(x) * self.weight
+
+
+# ── Projections ─────────────────────────────────────────────────────
+
+
+class MLPProj(nn.Module):
+    """CLIP feature projection for i2v.
+    Uses ColumnParallelLinear + RowParallelLinear (Qwen3_VisionMLP pattern).
+    """
+
+    def __init__(self, in_dim: int, out_dim: int) -> None:
+        super().__init__()
+        self.norm1 = nn.LayerNorm(in_dim)
+        self.fc1 = ColumnParallelLinear(
+            in_dim,
+            in_dim,
+            bias=True,
+            return_bias=False,
+        )
+        self.act = nn.GELU()
+        self.fc2 = RowParallelLinear(
+            in_dim,
+            out_dim,
+            bias=True,
+            return_bias=False,
+        )
+        self.norm2 = nn.LayerNorm(out_dim)
+
+    def forward(self, image_embeds: torch.Tensor) -> torch.Tensor:
+        x = self.norm1(image_embeds)
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.fc2(x)
+        x = self.norm2(x)
+        return x
+
+
+# ── Cross-Attention ─────────────────────────────────────────────────
+# T2V and I2V cross-attention variants
+
+
+class WanT2VCrossAttention(nn.Module):
+    """Text-to-video cross-attention.
+    Uses vllm-omni Attention for FlashAttn backend.
+    """
+
+    def __init__(self, dim: int, num_heads: int, window_size=(-1, -1), qk_norm: bool = True, eps: float = 1e-6) -> None:
+        super().__init__()
+        if dim % num_heads != 0:
+            raise ValueError(f"dim={dim} must be divisible by num_heads={num_heads}.")
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        tp_size = get_tensor_model_parallel_world_size()
+        if num_heads % tp_size != 0:
+            raise ValueError(f"num_heads={num_heads} must be divisible by tp_size={tp_size}.")
+        self.tp_num_heads = num_heads // tp_size
+        self.tp_inner_dim = self.tp_num_heads * self.head_dim
+        self.q = ColumnParallelLinear(dim, dim, bias=True, gather_output=False, return_bias=False)
+        self.k = ColumnParallelLinear(dim, dim, bias=True, gather_output=False, return_bias=False)
+        self.v = ColumnParallelLinear(dim, dim, bias=True, gather_output=False, return_bias=False)
+        self.o = RowParallelLinear(dim, dim, bias=True, input_is_parallel=True, return_bias=False)
+        self.norm_q = DistributedRMSNorm(self.tp_inner_dim, eps=eps) if qk_norm else nn.Identity()
+        self.norm_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps) if qk_norm else nn.Identity()
+        self.attn = Attention(
+            self.tp_num_heads,
+            self.head_dim,
+            causal=False,
+            softmax_scale=self.head_dim**-0.5,
+            skip_sequence_parallel=True,
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        context: torch.Tensor,
+        context_lens: torch.Tensor | None = None,
+        crossattn_cache: dict | None = None,
+    ) -> torch.Tensor:
+        del context_lens
+        n, d = self.tp_num_heads, self.head_dim
+        q = self.norm_q(self.q(x)).unflatten(2, (n, d))
+        if crossattn_cache is not None:
+            if not crossattn_cache["is_init"]:
+                crossattn_cache["is_init"] = True
+                k = self.norm_k(self.k(context)).unflatten(2, (n, d))
+                v = self.v(context).unflatten(2, (n, d))
+                crossattn_cache["k"] = k
+                crossattn_cache["v"] = v
+            else:
+                k = crossattn_cache["k"]
+                v = crossattn_cache["v"]
+        else:
+            k = self.norm_k(self.k(context)).unflatten(2, (n, d))
+            v = self.v(context).unflatten(2, (n, d))
+        x = self.attn(q, k, v)
+        x = x.flatten(2)
+        x = self.o(x)
+        return x
+
+
+class WanI2VCrossAttention(nn.Module):
+    """Image-to-video cross-attention (splits first 257 image tokens).
+    Uses vllm-omni Attention for FlashAttn backend.
+    """
+
+    def __init__(self, dim: int, num_heads: int, window_size=(-1, -1), qk_norm: bool = True, eps: float = 1e-6) -> None:
+        super().__init__()
+        if dim % num_heads != 0:
+            raise ValueError(f"dim={dim} must be divisible by num_heads={num_heads}.")
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        tp_size = get_tensor_model_parallel_world_size()
+        if num_heads % tp_size != 0:
+            raise ValueError(f"num_heads={num_heads} must be divisible by tp_size={tp_size}.")
+        self.tp_num_heads = num_heads // tp_size
+        self.tp_inner_dim = self.tp_num_heads * self.head_dim
+        self.q = ColumnParallelLinear(dim, dim, bias=True, gather_output=False, return_bias=False)
+        self.k = ColumnParallelLinear(dim, dim, bias=True, gather_output=False, return_bias=False)
+        self.v = ColumnParallelLinear(dim, dim, bias=True, gather_output=False, return_bias=False)
+        self.o = RowParallelLinear(dim, dim, bias=True, input_is_parallel=True, return_bias=False)
+        self.norm_q = DistributedRMSNorm(self.tp_inner_dim, eps=eps) if qk_norm else nn.Identity()
+        self.norm_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps) if qk_norm else nn.Identity()
+        self.k_img = ColumnParallelLinear(dim, dim, bias=True, gather_output=False, return_bias=False)
+        self.v_img = ColumnParallelLinear(dim, dim, bias=True, gather_output=False, return_bias=False)
+        self.norm_k_img = DistributedRMSNorm(self.tp_inner_dim, eps=eps) if qk_norm else nn.Identity()
+        self.attn = Attention(
+            self.tp_num_heads,
+            self.head_dim,
+            causal=False,
+            softmax_scale=self.head_dim**-0.5,
+            skip_sequence_parallel=True,
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        context: torch.Tensor,
+        context_lens: torch.Tensor | None = None,
+        crossattn_cache: dict | None = None,
+    ) -> torch.Tensor:
+        del context_lens
+        context_img = context[:, :257]
+        context = context[:, 257:]
+        n, d = self.tp_num_heads, self.head_dim
+        q = self.norm_q(self.q(x)).unflatten(2, (n, d))
+        if crossattn_cache is not None:
+            if not crossattn_cache["is_init"]:
+                crossattn_cache["is_init"] = True
+                k = self.norm_k(self.k(context)).unflatten(2, (n, d))
+                v = self.v(context).unflatten(2, (n, d))
+                crossattn_cache["k"] = k
+                crossattn_cache["v"] = v
+            else:
+                k = crossattn_cache["k"]
+                v = crossattn_cache["v"]
+        else:
+            k = self.norm_k(self.k(context)).unflatten(2, (n, d))
+            v = self.v(context).unflatten(2, (n, d))
+        x = self.attn(q, k, v)
+        k_img = self.norm_k_img(self.k_img(context_img)).unflatten(2, (n, d))
+        v_img = self.v_img(context_img).unflatten(2, (n, d))
+        img_x = self.attn(q, k_img, v_img)
+        x = x.flatten(2)
+        img_x = img_x.flatten(2)
+        x = x + img_x
+        x = self.o(x)
+        return x
+
+
+WAN_CROSSATTENTION_CLASSES = {
+    "t2v_cross_attn": WanT2VCrossAttention,
+    "i2v_cross_attn": WanI2VCrossAttention,
+}
+
+
+# ── Self-Attention with causal masking + KV cache ───────────────────
+
+
+class CausalWanSelfAttention(nn.Module):
+    """Causal self-attention with KV cache + action/state tokens."""
+
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        frame_seqlen: int,
+        local_attn_size: int = -1,
+        sink_size: int = 0,
+        num_frame_per_block: int = 1,
+        qk_norm: bool = True,
+        eps: float = 1e-6,
+        num_action_per_block: int = 32,
+        num_state_per_block: int = 1,
+    ) -> None:
+        super().__init__()
+        if dim % num_heads != 0:
+            raise ValueError(f"dim={dim} must be divisible by num_heads={num_heads}.")
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        tp_size = get_tensor_model_parallel_world_size()
+        if num_heads % tp_size != 0:
+            raise ValueError(f"num_heads={num_heads} must be divisible by tp_size={tp_size}.")
+        self.tp_num_heads = num_heads // tp_size
+        self.tp_inner_dim = self.tp_num_heads * self.head_dim
+        self.local_attn_size = local_attn_size
+        self.num_frame_per_block = num_frame_per_block
+        self.frame_seqlen = frame_seqlen
+        self.num_action_per_block = num_action_per_block
+        self.num_state_per_block = num_state_per_block
+        self.max_attention_size = 21 * frame_seqlen if local_attn_size == -1 else local_attn_size * frame_seqlen
+        self.q = ColumnParallelLinear(dim, dim, bias=True, gather_output=False, return_bias=False)
+        self.k = ColumnParallelLinear(dim, dim, bias=True, gather_output=False, return_bias=False)
+        self.v = ColumnParallelLinear(dim, dim, bias=True, gather_output=False, return_bias=False)
+        self.o = RowParallelLinear(dim, dim, bias=True, input_is_parallel=True, return_bias=False)
+        self.norm_q = DistributedRMSNorm(self.tp_inner_dim, eps=eps) if qk_norm else nn.Identity()
+        self.norm_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps) if qk_norm else nn.Identity()
+        self.attn = Attention(
+            self.tp_num_heads,
+            self.head_dim,
+            causal=False,
+            softmax_scale=self.head_dim**-0.5,
+            skip_sequence_parallel=True,
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        freqs: torch.Tensor,
+        freqs_action: torch.Tensor,
+        freqs_state: torch.Tensor,
+        action_register_length: int | None,
+        kv_cache: torch.Tensor | None = None,
+        current_start_frame: int = 0,
+        is_tf: bool = True,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        """Inference-only forward (KV cache path)."""
+        n, d = self.tp_num_heads, self.head_dim
+
+        q = self.norm_q(self.q(x)).unflatten(2, (n, d))
+        k = self.norm_k(self.k(x)).unflatten(2, (n, d))
+        v = self.v(x).unflatten(2, (n, d))
+
+        updated_kv_cache: torch.Tensor | None = None
+
+        if kv_cache is None:
+            raise RuntimeError("Inference only: kv_cache is required.")
+
+        action_state_index = max(0, (current_start_frame - 1) // self.num_frame_per_block)
+
+        roped_query = causal_rope_action_apply(
+            q,
+            freqs,
+            freqs_action,
+            freqs_state,
+            action_register_length,
+            self.num_action_per_block,
+            self.num_state_per_block,
+            action_state_index,
+        ).type_as(v)
+        roped_key = causal_rope_action_apply(
+            k,
+            freqs,
+            freqs_action,
+            freqs_state,
+            action_register_length,
+            self.num_action_per_block,
+            self.num_state_per_block,
+            action_state_index,
+        ).type_as(v)
+
+        roped_action_query = None
+        roped_action_key = None
+        action_v = None
+
+        if action_register_length is not None:
+            roped_action_query = roped_query[:, -action_register_length:]
+            roped_query = roped_query[:, :-action_register_length]
+            roped_action_key = roped_key[:, -action_register_length:]
+            roped_key = roped_key[:, :-action_register_length]
+            action_v = v[:, -action_register_length:]
+            v = v[:, :-action_register_length]
+
+        updated_k = kv_cache[0]
+        updated_v = kv_cache[1]
+        new_k = torch.cat([updated_k, roped_key], dim=1)
+        new_v = torch.cat([updated_v, v], dim=1)
+        new_k = new_k[:, -self.max_attention_size :]
+        new_v = new_v[:, -self.max_attention_size :]
+
+        if action_register_length is not None:
+            q_cat = torch.cat([roped_query, roped_action_query], dim=1)
+            k_cat = torch.cat([new_k, roped_action_key], dim=1)
+            v_cat = torch.cat([new_v, action_v], dim=1)
+        else:
+            q_cat = roped_query
+            k_cat = new_k
+            v_cat = new_v
+
+        x = self.attn(q_cat, k_cat, v_cat)
+        updated_kv_cache = torch.stack([new_k, new_v], dim=0)
+
+        x = x.flatten(2)
+        x = self.o(x)
+        return x, updated_kv_cache
+
+
+# ── Attention Block ─────────────────────────────────────────────────
+
+
+class CausalWanAttentionBlock(nn.Module):
+    """Transformer block: self-attn + cross-attn + FFN with 6-param modulation."""
+
+    def __init__(
+        self,
+        cross_attn_type: str,
+        dim: int,
+        ffn_dim: int,
+        num_heads: int,
+        frame_seqlen: int,
+        local_attn_size: int = -1,
+        sink_size: int = 0,
+        num_frame_per_block: int = 1,
+        qk_norm: bool = True,
+        cross_attn_norm: bool = False,
+        eps: float = 1e-6,
+        num_action_per_block: int = 32,
+        num_state_per_block: int = 1,
+    ) -> None:
+        super().__init__()
+        self.norm1 = WanLayerNorm(dim, eps)
+        self.self_attn = CausalWanSelfAttention(
+            dim=dim,
+            num_heads=num_heads,
+            frame_seqlen=frame_seqlen,
+            local_attn_size=local_attn_size,
+            sink_size=sink_size,
+            num_frame_per_block=num_frame_per_block,
+            qk_norm=qk_norm,
+            eps=eps,
+            num_action_per_block=num_action_per_block,
+            num_state_per_block=num_state_per_block,
+        )
+        self.norm3 = WanLayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity()
+        self.cross_attn = WAN_CROSSATTENTION_CLASSES[cross_attn_type](dim, num_heads, (-1, -1), qk_norm, eps)
+        self.norm2 = WanLayerNorm(dim, eps)
+        self.ffn = nn.Sequential(
+            ColumnParallelLinear(dim, ffn_dim, bias=True, gather_output=False, return_bias=False),
+            nn.GELU(approximate="tanh"),
+            RowParallelLinear(ffn_dim, dim, bias=True, input_is_parallel=True, return_bias=False),
+        )
+        self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        e: torch.Tensor,
+        freqs: torch.Tensor,
+        freqs_action: torch.Tensor,
+        freqs_state: torch.Tensor,
+        context: torch.Tensor,
+        action_register_length: int | None = None,
+        kv_cache: torch.Tensor | None = None,
+        crossattn_cache: dict | None = None,
+        current_start_frame: int = 0,
+        is_tf: bool = True,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        e = (self.modulation.unsqueeze(1) + e).chunk(6, dim=2)
+
+        y, updated_kv_cache = self.self_attn(
+            x=(self.norm1(x) * (1 + e[1].squeeze(2)) + e[0].squeeze(2)),
+            freqs=freqs,
+            freqs_action=freqs_action,
+            freqs_state=freqs_state,
+            action_register_length=action_register_length,
+            kv_cache=kv_cache,
+            is_tf=is_tf,
+            current_start_frame=current_start_frame,
+        )
+        x = x + (y * e[2].squeeze(2))
+
+        x = x + self.cross_attn(self.norm3(x), context, crossattn_cache=crossattn_cache)
+        y = self.ffn(self.norm2(x) * (1 + e[4].squeeze(2)) + e[3].squeeze(2))
+        x = x + (y * e[5].squeeze(2))
+        return x, updated_kv_cache
+
+
+# ── Output Head ─────────────────────────────────────────────────────
+
+
+class CausalHead(nn.Module):
+    """Output norm + linear with 2-param modulation.
+    Runs once per step (not TP-critical), uses nn.Linear.
+    """
+
+    def __init__(self, dim: int, out_dim: int, patch_size: tuple, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.dim = dim
+        self.out_dim = out_dim
+        self.patch_size = patch_size
+        out_channels = math.prod(patch_size) * out_dim
+        self.norm = WanLayerNorm(dim, eps)
+        self.head = nn.Linear(dim, out_channels)
+        self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5)
+
+    def forward(self, x: torch.Tensor, e: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: [B, L1, C]
+            e: [B, F, 1, C]     (time embedding, unsqueezed)
+        """
+        e = (self.modulation.unsqueeze(1) + e).chunk(2, dim=2)
+        x = self.head(self.norm(x) * (1 + e[1].squeeze(2)) + e[0].squeeze(2))
+        return x
+
+
+# ── Main Model ──────────────────────────────────────────────────────
+
+
+class CausalWanModel(nn.Module):
+    """Causal video diffusion transformer for DreamZero.
+
+    Architecture (14B): 40 layers, dim=5120, heads=40, ffn=13824
+    """
+
+    def __init__(
+        self,
+        model_type: str = "t2v",
+        patch_size: tuple[int, int, int] = (1, 2, 2),
+        frame_seqlen: int = 220,
+        text_len: int = 512,
+        in_dim: int = 16,
+        dim: int = 2048,
+        ffn_dim: int = 8192,
+        freq_dim: int = 256,
+        text_dim: int = 4096,
+        out_dim: int = 16,
+        num_heads: int = 16,
+        num_layers: int = 32,
+        max_chunk_size: int = -1,
+        sink_size: int = 0,
+        qk_norm: bool = True,
+        cross_attn_norm: bool = True,
+        eps: float = 1e-6,
+        num_frame_per_block: int = 1,
+        action_dim: int = 32,
+        num_registers: int = 8,
+        max_state_dim: int = 64,
+        max_num_embodiments: int = 32,
+        hidden_size: int = 1024,
+        diffusion_model_pretrained_path: str | None = None,
+        num_action_per_block: int = 32,
+        num_state_per_block: int = 1,
+    ) -> None:
+        super().__init__()
+        if model_type not in ["t2v", "i2v", "ti2v"]:
+            raise ValueError(f"Unsupported model_type={model_type!r}; expected one of ['t2v', 'i2v', 'ti2v'].")
+        self.model_type = model_type
+        self.patch_size = patch_size
+        self.frame_seqlen = frame_seqlen
+        self.text_len = text_len
+        self.dim = dim
+        self.freq_dim = freq_dim
+        self.out_dim = out_dim
+        self.num_heads = num_heads
+        self.num_layers = num_layers
+        self.local_attn_size = max_chunk_size * num_frame_per_block + 1 if max_chunk_size != -1 else -1
+        self.num_frame_per_block = num_frame_per_block
+        self.action_dim = action_dim
+        self.num_action_per_block = num_action_per_block
+        self.num_state_per_block = num_state_per_block
+
+        max_num_embodiments_local = 1
+        self.state_encoder = CategorySpecificMLP(
+            num_categories=max_num_embodiments_local,
+            input_dim=max_state_dim,
+            hidden_dim=hidden_size,
+            output_dim=dim,
+        )
+        self.action_encoder = MultiEmbodimentActionEncoder(
+            action_dim=action_dim,
+            hidden_size=dim,
+            num_embodiments=max_num_embodiments_local,
+        )
+        self.action_decoder = CategorySpecificMLP(
+            num_categories=max_num_embodiments_local,
+            input_dim=dim,
+            hidden_dim=hidden_size,
+            output_dim=action_dim,
+        )
+
+        # Disable the Conv3d GEMM rewrite for patch embedding.
+        self.patch_embedding = Conv3dLayer(
+            in_dim,
+            dim,
+            kernel_size=patch_size,
+            stride=patch_size,
+        )
+        self.patch_embedding.enable_linear = False
+        self.text_embedding = nn.Sequential(
+            nn.Linear(text_dim, dim),
+            nn.GELU(approximate="tanh"),
+            nn.Linear(dim, dim),
+        )
+        self.time_embedding = nn.Sequential(
+            nn.Linear(freq_dim, dim),
+            nn.SiLU(),
+            nn.Linear(dim, dim),
+        )
+        self.time_projection = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(dim, dim * 6),
+        )
+
+        cross_attn_type = "t2v_cross_attn" if model_type == "t2v" else "i2v_cross_attn"
+        self.blocks = nn.ModuleList(
+            [
+                CausalWanAttentionBlock(
+                    cross_attn_type,
+                    dim,
+                    ffn_dim,
+                    num_heads,
+                    frame_seqlen,
+                    self.local_attn_size,
+                    sink_size,
+                    num_frame_per_block,
+                    qk_norm,
+                    cross_attn_norm,
+                    eps,
+                    num_action_per_block,
+                    num_state_per_block,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+
+        self.head = CausalHead(dim, out_dim, patch_size, eps)
+
+        if dim % num_heads != 0:
+            raise ValueError(f"dim={dim} must be divisible by num_heads={num_heads}.")
+        if (dim // num_heads) % 2 != 0:
+            raise ValueError(f"dim // num_heads must be even, got {dim // num_heads}.")
+        d = dim // num_heads
+        self.freqs_action = rope_params(1024 * 10, d)
+        self.freqs_state = rope_params(1024, d)
+        self.freqs = [
+            rope_params(1024, d - 4 * (d // 6)),
+            rope_params(1024, 2 * (d // 6)),
+            rope_params(1024, 2 * (d // 6)),
+        ]
+
+        if model_type == "i2v":
+            self.img_emb = MLPProj(1280, dim)
+
+        self.init_weights()
+
+    def init_weights(self) -> None:
+        """Initialize parameters."""
+
+        def _init_linear_like(module: nn.Module) -> None:
+            if isinstance(module, nn.Linear):
+                nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.zeros_(module.bias)
+                return
+
+            if isinstance(module, (ColumnParallelLinear, RowParallelLinear)):
+                fan_in = module.input_size
+                fan_out = module.output_size
+                bound = math.sqrt(6.0 / float(fan_in + fan_out))
+                nn.init.uniform_(module.weight, -bound, bound)
+                if module.bias is not None:
+                    nn.init.zeros_(module.bias)
+
+        for module in self.modules():
+            _init_linear_like(module)
+
+        nn.init.xavier_uniform_(self.patch_embedding.weight.flatten(1))
+        if self.patch_embedding.bias is not None:
+            fan_in = self.patch_embedding.in_channels * math.prod(self.patch_embedding.kernel_size)
+            bound = 1 / math.sqrt(fan_in)
+            nn.init.uniform_(self.patch_embedding.bias, -bound, bound)
+
+        for module in self.text_embedding.modules():
+            if isinstance(module, nn.Linear):
+                nn.init.normal_(module.weight, std=0.02)
+
+        for module in self.time_embedding.modules():
+            if isinstance(module, nn.Linear):
+                nn.init.normal_(module.weight, std=0.02)
+
+        nn.init.zeros_(self.head.head.weight)
+
+    def _create_freqs(self, grid_size: torch.Tensor, start_frame: int) -> torch.Tensor:
+        """Create 3D RoPE frequency tensor."""
+        device = self.patch_embedding.weight.device
+        if any(freq.device != device for freq in self.freqs):
+            self.freqs = [freq.to(device) for freq in self.freqs]
+        if self.freqs_action.device != device:
+            self.freqs_action = self.freqs_action.to(device)
+        if self.freqs_state.device != device:
+            self.freqs_state = self.freqs_state.to(device)
+
+        f, h, w = grid_size.tolist()
+        freqs = torch.cat(
+            [
+                self.freqs[0][start_frame : start_frame + f].view(f, 1, 1, -1).expand(f, h, w, -1),
+                self.freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),
+                self.freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1),
+            ],
+            dim=-1,
+        ).reshape(f * h * w, 1, -1)
+        return freqs
+
+    def unpatchify(self, x: torch.Tensor, grid_size: torch.Tensor) -> torch.Tensor:
+        """Reconstruct video from patch embeddings."""
+        B = x.shape[0]
+        c = self.out_dim
+        grid_size = grid_size.tolist()
+        expected_seq_len = math.prod(grid_size)
+        if x.shape[1] != expected_seq_len:
+            raise ValueError(f"x sequence length must equal product(grid_size)={expected_seq_len}, got {x.shape[1]}.")
+        x = x.view(B, *grid_size, *self.patch_size, c)
+        x = torch.einsum("bfhwpqrc->bcfphqwr", x)
+        x = x.reshape(B, c, *[i * j for i, j in zip(grid_size, self.patch_size)])
+        return x
+
+    def _forward_blocks(
+        self,
+        x: torch.Tensor,
+        seq_len: int,
+        freqs: torch.Tensor,
+        timestep: torch.Tensor,
+        context: torch.Tensor,
+        clip_feature: torch.Tensor | None,
+        embodiment_id: torch.Tensor | None,
+        action: torch.Tensor | None,
+        timestep_action: torch.Tensor | None,
+        state: torch.Tensor | None,
+        kv_cache: list[torch.Tensor],
+        current_start_frame: int,
+    ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor]]:
+        x = x.flatten(start_dim=2).transpose(1, 2)
+        B = x.shape[0]
+        F_t = timestep.shape[1]
+
+        if action is not None:
+            # Current DreamZero checkpoints have one local action/state adapter.
+            # Global embodiment IDs are used by transforms and normalization.
+            adapter_category_id = torch.zeros(B, dtype=torch.long, device=x.device)
+            action_features = self.action_encoder(action, timestep_action, adapter_category_id)
+            state_features = self.state_encoder(state, adapter_category_id)
+            action_register = torch.cat([action_features, state_features], dim=1)
+            action_length = action_features.shape[1]
+            action_register_length = action_register.shape[1]
+            x = torch.cat([x, action_register], dim=1)
+        else:
+            action_length = 0
+            action_register_length = None
+
+        timestep = timestep.unsqueeze(-1).expand(B, F_t, seq_len // F_t).reshape(B, -1)
+        if action is not None:
+            if timestep_action is None or state is None:
+                raise RuntimeError("timestep_action and state are required when action is provided.")
+            state_features_t = self.state_encoder(state, adapter_category_id)
+            stride = timestep_action.shape[1] // state_features_t.shape[1]
+            timestep_state = timestep_action[:, ::stride]
+            timestep = torch.cat([timestep, timestep_action, timestep_state], dim=1)
+
+        e = self.time_embedding(sinusoidal_embedding_1d(self.freq_dim, timestep.flatten()).type_as(x))
+        e = e.unflatten(dim=0, sizes=(B, -1))
+        e0 = self.time_projection(e)
+        e0 = e0.unflatten(dim=2, sizes=(6, self.dim))
+
+        context = self.text_embedding(context)
+        if clip_feature is not None:
+            clip_embedding = self.img_emb(clip_feature)
+            context = torch.cat([clip_embedding, context], dim=1)
+
+        updated_kv_caches: list[torch.Tensor] = []
+        for block_index, block in enumerate(self.blocks):
+            x, updated_kv_cache = block(
+                x=x,
+                e=e0,
+                freqs=freqs,
+                freqs_action=self.freqs_action,
+                freqs_state=self.freqs_state,
+                context=context,
+                action_register_length=action_register_length,
+                kv_cache=kv_cache[block_index] if kv_cache else None,
+                current_start_frame=current_start_frame,
+            )
+            updated_kv_caches.append(updated_kv_cache)
+
+        if action is not None:
+            action_noise_pred = x[:, seq_len : seq_len + action_length]
+            action_noise_pred = self.action_decoder(action_noise_pred, adapter_category_id)
+        else:
+            action_noise_pred = None
+
+        x_video = x[:, :seq_len]
+        e_video = e[:, :seq_len]
+        x_video = self.head(x_video, e_video.unsqueeze(2))
+
+        return x_video, action_noise_pred, updated_kv_caches
+
+    def _forward_inference(
+        self,
+        x: torch.Tensor,
+        timestep: torch.Tensor,
+        context: torch.Tensor,
+        seq_len: int,
+        kv_cache: list[torch.Tensor],
+        crossattn_cache: list[torch.Tensor],
+        current_start_frame: int,
+        y: torch.Tensor | None = None,
+        clip_feature: torch.Tensor | None = None,
+        action: torch.Tensor | None = None,
+        timestep_action: torch.Tensor | None = None,
+        state: torch.Tensor | None = None,
+        embodiment_id: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor]]:
+        if self.model_type == "i2v":
+            if clip_feature is None or y is None:
+                raise RuntimeError("clip_feature and y are required for i2v inference.")
+        if context.shape[1] != self.text_len:
+            raise ValueError(f"context length must be {self.text_len}, got {context.shape[1]}.")
+
+        if y is not None:
+            x = torch.cat([x, y.to(dtype=x.dtype)], dim=1)
+
+        x = self.patch_embedding(x)
+        grid_size = torch.tensor(x.shape[2:], dtype=torch.long)
+        freqs = self._create_freqs(grid_size, current_start_frame)
+
+        x_video, action_noise_pred, updated_kv_caches = self._forward_blocks(
+            x=x,
+            seq_len=seq_len,
+            freqs=freqs,
+            timestep=timestep,
+            context=context,
+            clip_feature=clip_feature,
+            embodiment_id=embodiment_id,
+            action=action,
+            timestep_action=timestep_action,
+            state=state,
+            kv_cache=kv_cache,
+            current_start_frame=current_start_frame,
+        )
+
+        x_video = x_video.clone()
+        if action_noise_pred is not None:
+            action_noise_pred = action_noise_pred.clone()
+
+        video_noise_pred = self.unpatchify(x_video, grid_size)
+        return video_noise_pred, action_noise_pred, updated_kv_caches
+
+    def forward(self, *args: Any, **kwargs: Any):
+        """Inference only. Requires kv_cache."""
+        return self._forward_inference(*args, **kwargs)

vllm_omni/diffusion/models/dreamzero/image_encoder.py

@@ -0,0 +1,243 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""DreamZero image encoder.
+
+Only the visual tower used by DreamZero I2V inference is ported here. The
+checkpoint keys under `action_head.image_encoder.*` load via simple prefix
+stripping.
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.transforms as T
+
+
+class DreamZeroLayerNorm(nn.LayerNorm):
+    """LayerNorm that preserves the input dtype."""
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return super().forward(x).type_as(x)
+
+
+class DreamZeroVisionSelfAttention(nn.Module):
+    """Self-attention for the vision tower."""
+
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        proj_dropout: float = 0.0,
+    ) -> None:
+        super().__init__()
+        if dim % num_heads != 0:
+            raise ValueError(f"dim={dim} must be divisible by num_heads={num_heads}.")
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.proj_dropout = proj_dropout
+
+        self.to_qkv = nn.Linear(dim, dim * 3)
+        self.proj = nn.Linear(dim, dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        batch_size, seq_len, _ = x.shape
+        q, k, v = self.to_qkv(x).chunk(3, dim=-1)
+        q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
+        k = k.view(batch_size, seq_len, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
+        v = v.view(batch_size, seq_len, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
+        x = F.scaled_dot_product_attention(q, k, v)
+        x = x.permute(0, 2, 1, 3).reshape(batch_size, seq_len, self.dim)
+        x = self.proj(x)
+        return F.dropout(x, self.proj_dropout, self.training)
+
+
+class DreamZeroVisionAttentionBlock(nn.Module):
+    """Attention block for the vision tower."""
+
+    def __init__(
+        self,
+        dim: int,
+        mlp_ratio: float,
+        num_heads: int,
+        post_norm: bool = False,
+        activation: str = "gelu",
+        proj_dropout: float = 0.0,
+        norm_eps: float = 1e-5,
+    ) -> None:
+        super().__init__()
+        if activation != "gelu":
+            raise ValueError(f"DreamZero image encoder uses GELU; got activation={activation!r}.")
+        self.post_norm = post_norm
+        hidden_dim = int(dim * mlp_ratio)
+
+        self.norm1 = DreamZeroLayerNorm(dim, eps=norm_eps)
+        self.attn = DreamZeroVisionSelfAttention(
+            dim,
+            num_heads,
+            proj_dropout=proj_dropout,
+        )
+        self.norm2 = DreamZeroLayerNorm(dim, eps=norm_eps)
+        self.mlp = nn.Sequential(
+            nn.Linear(dim, hidden_dim),
+            nn.GELU(),
+            nn.Linear(hidden_dim, dim),
+            nn.Dropout(proj_dropout),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.post_norm:
+            x = x + self.norm1(self.attn(x))
+            x = x + self.norm2(self.mlp(x))
+        else:
+            x = x + self.attn(self.norm1(x))
+            x = x + self.mlp(self.norm2(x))
+        return x
+
+
+class DreamZeroVisionTransformer(nn.Module):
+    """Vision transformer used by the image encoder."""
+
+    def __init__(
+        self,
+        image_size: int = 224,
+        patch_size: int = 14,
+        dim: int = 1280,
+        mlp_ratio: float = 4.0,
+        out_dim: int = 1024,
+        num_heads: int = 16,
+        num_layers: int = 32,
+        pool_type: str = "token",
+        pre_norm: bool = True,
+        post_norm: bool = False,
+        activation: str = "gelu",
+        proj_dropout: float = 0.0,
+        embedding_dropout: float = 0.0,
+        norm_eps: float = 1e-5,
+    ) -> None:
+        super().__init__()
+        if pool_type != "token":
+            raise ValueError(f"DreamZero image encoder only supports pool_type='token', got {pool_type!r}.")
+        self.image_size = image_size
+        self.patch_size = patch_size
+        self.num_patches = (image_size // patch_size) ** 2
+        self.dim = dim
+        self.num_heads = num_heads
+        self.num_layers = num_layers
+        self.pool_type = pool_type
+
+        gain = 1.0 / math.sqrt(dim)
+        self.patch_embedding = nn.Conv2d(
+            3,
+            dim,
+            kernel_size=patch_size,
+            stride=patch_size,
+            bias=not pre_norm,
+        )
+        self.cls_embedding = nn.Parameter(gain * torch.randn(1, 1, dim))
+        self.pos_embedding = nn.Parameter(
+            gain * torch.randn(1, self.num_patches + 1, dim),
+        )
+        self.dropout = nn.Dropout(embedding_dropout)
+        self.pre_norm = DreamZeroLayerNorm(dim, eps=norm_eps) if pre_norm else None
+        self.transformer = nn.Sequential(
+            *[
+                DreamZeroVisionAttentionBlock(
+                    dim=dim,
+                    mlp_ratio=mlp_ratio,
+                    num_heads=num_heads,
+                    post_norm=post_norm,
+                    activation=activation,
+                    proj_dropout=proj_dropout,
+                    norm_eps=norm_eps,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+        self.post_norm = DreamZeroLayerNorm(dim, eps=norm_eps)
+        self.head = nn.Parameter(gain * torch.randn(dim, out_dim))
+
+    def forward(self, x: torch.Tensor, use_31_block: bool = False) -> torch.Tensor:
+        batch_size = x.shape[0]
+        x = self.patch_embedding(x).flatten(2).permute(0, 2, 1)
+        x = torch.cat(
+            [
+                self.cls_embedding.expand(batch_size, -1, -1).to(dtype=x.dtype, device=x.device),
+                x,
+            ],
+            dim=1,
+        )
+        x = self.dropout(x + self.pos_embedding.to(dtype=x.dtype, device=x.device))
+        if self.pre_norm is not None:
+            x = self.pre_norm(x)
+
+        if use_31_block:
+            return self.transformer[:-1](x)
+        return self.transformer(x)
+
+
+class _DreamZeroCLIPContainer(nn.Module):
+    """Container matching checkpoint names under `model.visual.*`."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.log_scale = nn.Parameter(torch.ones(()))
+        self.visual = DreamZeroVisionTransformer(
+            image_size=224,
+            patch_size=14,
+            dim=1280,
+            mlp_ratio=4.0,
+            out_dim=1024,
+            num_heads=16,
+            num_layers=32,
+            pool_type="token",
+            pre_norm=True,
+            post_norm=False,
+            activation="gelu",
+            proj_dropout=0.0,
+            embedding_dropout=0.0,
+            norm_eps=1e-5,
+        )
+
+
+class DreamZeroImageEncoder(nn.Module):
+    """Image encoder wrapper."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.model = _DreamZeroCLIPContainer()
+        # returns a composed transform whose last stage is CLIP normalization.
+        self.transforms = T.Compose(
+            [
+                T.Normalize(
+                    mean=[0.48145466, 0.4578275, 0.40821073],
+                    std=[0.26862954, 0.26130258, 0.27577711],
+                ),
+            ]
+        )
+
+    def encode_image(self, videos: torch.Tensor) -> torch.Tensor:
+        """Encode images for I2V conditioning."""
+        size = (self.model.visual.image_size,) * 2
+        videos = torch.cat(
+            [
+                F.interpolate(
+                    frame_batch,
+                    size=size,
+                    mode="bicubic",
+                    align_corners=False,
+                )
+                for frame_batch in videos
+            ]
+        )
+        videos = self.transforms.transforms[-1](videos.mul_(0.5).add_(0.5))
+
+        param_dtype = next(iter(self.model.visual.parameters())).dtype
+        videos = videos.to(dtype=param_dtype)
+        out = self.model.visual(videos, use_31_block=True)
+        return out.clone()

vllm_omni/diffusion/models/dreamzero/state_dreamzero.py

@@ -0,0 +1,161 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""DreamZero pipeline persistent state."""
+
+from __future__ import annotations
+
+import logging
+from collections import deque
+
+import numpy as np
+import torch
+
+logger = logging.getLogger(__name__)
+
+# Number of frames per chunk for subsequent calls (first call uses 1)
+FRAMES_PER_CHUNK = 4
+
+
+class DreamZeroState:
+    """Pipeline persistent state across forward() calls.
+
+    Lifecycle:
+        - Created once in DreamZeroPipeline.__init__()
+        - Mutated every forward() call (frame append, KV cache grow)
+        - reset() on new session / language change / local_attn_size exceeded
+    """
+
+    def __init__(self) -> None:
+        self.reset()
+
+    # ------------------------------------------------------------------
+    # Frame accumulation (single stitched buffer)
+    # Transform outputs stitched single frame per call.
+    # We accumulate here to build multi-frame video for AR inference.
+    # ------------------------------------------------------------------
+
+    def accumulate_frames(self, stitched: np.ndarray) -> np.ndarray:
+        """Accumulate stitched frames and return multi-frame video.
+
+        Args:
+            stitched: (H, W, C) single frame or (T, H, W, C) multi-frame,
+                      already stitched by transform.
+
+        Returns:
+            (T, H, W, C) ndarray. T=1 for first call, T=FRAMES_PER_CHUNK(4) after.
+        """
+        if stitched.ndim == 3:
+            self.stitched_buffer.append(stitched)
+        elif stitched.ndim == 4:
+            self.stitched_buffer.extend(list(stitched))
+        else:
+            raise ValueError(f"Expected 3D or 4D stitched, got {stitched.ndim}D")
+
+        num_frames = 1 if self.call_count == 0 else FRAMES_PER_CHUNK
+
+        buffer_frames = list(self.stitched_buffer)
+        if len(buffer_frames) >= num_frames:
+            frames = buffer_frames[-num_frames:]
+        else:
+            frames = buffer_frames
+            while len(frames) < num_frames:
+                frames.insert(0, buffer_frames[0])
+
+        self.call_count += 1
+        return np.stack(frames, axis=0)  # (T, H, W, C)
+
+    # ------------------------------------------------------------------
+    # Reset / should_reset
+    # ------------------------------------------------------------------
+
+    def reset(self) -> None:
+        """Clear all state."""
+        self.stitched_buffer: deque[np.ndarray] = deque(maxlen=FRAMES_PER_CHUNK)
+        self.call_count: int = 0
+
+        # KV cache once robot-policy diffusion supports that integration.
+        self.kv_cache: list[torch.Tensor] | None = None
+        self.kv_cache_neg: list[torch.Tensor] | None = None
+        self.crossattn_cache: list[dict[str, bool | torch.Tensor | None]] | None = None
+        self.crossattn_cache_neg: list[dict[str, bool | torch.Tensor | None]] | None = None
+        self.current_start_frame: int = 0
+
+        self.clip_feas: torch.Tensor | None = None
+        self.ys: torch.Tensor | None = None
+        self.language: torch.Tensor | None = None
+
+    def should_reset(self, text_tokens: torch.Tensor | None, num_video_frames: int, local_attn_size: int) -> bool:
+        """Determine if state should be reset before this forward()."""
+        if self.language is None:
+            logger.info("language is None, resetting")
+            return True
+
+        if text_tokens is not None and not torch.equal(self.language, text_tokens):
+            logger.info("language changed, resetting")
+            return True
+
+        # NOTE: after accumulate_frames, num_video_frames is the accumulated T
+        # (1 for first call, 4 for subsequent). Only reset on true single-frame
+        # which happens when the stitched_buffer was cleared externally.
+        if num_video_frames == 1 and self.call_count > 1:
+            logger.info("single frame input after first call, resetting")
+            return True
+
+        if local_attn_size != -1 and self.current_start_frame >= local_attn_size:
+            logger.info(
+                "current_start_frame %d >= local_attn_size %d, resetting", self.current_start_frame, local_attn_size
+            )
+            return True
+
+        return False
+
+    # ------------------------------------------------------------------
+    # KV cache management
+    # ------------------------------------------------------------------
+
+    def create_kv_caches(
+        self,
+        batch_size: int,
+        dtype: torch.dtype,
+        device: torch.device,
+        num_layers: int,
+        num_heads: int,
+        head_dim: int,
+    ) -> None:
+        """Initialize empty KV caches and cross-attention caches."""
+        self.kv_cache = [
+            torch.zeros(2, batch_size, 0, num_heads, head_dim, dtype=dtype, device=device) for _ in range(num_layers)
+        ]
+        self.kv_cache_neg = [
+            torch.zeros(2, batch_size, 0, num_heads, head_dim, dtype=dtype, device=device) for _ in range(num_layers)
+        ]
+
+        self.crossattn_cache = [{"is_init": False, "k": None, "v": None} for _ in range(num_layers)]
+        self.crossattn_cache_neg = [{"is_init": False, "k": None, "v": None} for _ in range(num_layers)]
+
+    def update_kv_cache(
+        self,
+        layer_index: int,
+        updated_kv: torch.Tensor,
+        is_negative: bool = False,
+    ) -> None:
+        """Update a single layer's KV cache after prefill."""
+        cache = self.kv_cache_neg if is_negative else self.kv_cache
+        if cache is None:
+            raise RuntimeError("KV caches not initialized, call create_kv_caches first.")
+        cache[layer_index] = updated_kv.clone()
+
+    def get_kv_caches(self, is_negative: bool = False) -> list[torch.Tensor]:
+        """Get KV caches for the specified branch."""
+        cache = self.kv_cache_neg if is_negative else self.kv_cache
+        if cache is None:
+            raise RuntimeError("KV caches not initialized.")
+        return cache
+
+    def get_crossattn_caches(self, is_negative: bool = False) -> list[dict[str, bool | torch.Tensor | None]]:
+        """Get cross-attention caches for the specified branch."""
+        cache = self.crossattn_cache_neg if is_negative else self.crossattn_cache
+        if cache is None:
+            raise RuntimeError("Cross-attn caches not initialized.")
+        return cache

vllm_omni/diffusion/models/dreamzero/transform/__init__.py

@@ -0,0 +1,31 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import importlib
+
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.models.dreamzero.transform.base import TRANSFORMS
+
+logger = init_logger(__name__)
+
+DEFAULT_EMBODIMENT = "roboarena"
+_BUILTIN_TRANSFORM_MODULES = (
+    "vllm_omni.diffusion.models.dreamzero.transform.droid",
+    "vllm_omni.diffusion.models.dreamzero.transform.roboarena",
+)
+
+
+def ensure_transforms_loaded() -> None:
+    """Import DreamZero transform modules and verify registration."""
+    for module_name in _BUILTIN_TRANSFORM_MODULES:
+        try:
+            importlib.import_module(module_name)
+        except Exception as exc:
+            logger.exception("Failed to import DreamZero transform module %s", module_name)
+            raise RuntimeError(f"Failed to import DreamZero transform module '{module_name}'.") from exc
+
+    if DEFAULT_EMBODIMENT not in TRANSFORMS:
+        raise RuntimeError(f"Built-in DreamZero transform '{DEFAULT_EMBODIMENT}' is not registered after import.")

vllm_omni/diffusion/models/dreamzero/transform/base.py

@@ -0,0 +1,126 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Base transform interface for DreamZero robot policy serving.
+
+Transforms handle dataset-specific concerns ONLY:
+  - Observation key mapping
+  - Multi-view stitching (embodiment-specific layout)
+  - Language template wrapping (embodiment-specific)
+  - Raw state extraction (dataset-specific keys)
+  - Output action slicing (to actual action_dim)
+
+Model-specific concerns belong in the pipeline:
+  - Tokenization (pipeline owns tokenizer)
+  - State padding (pipeline knows MAX_STATE_DIM)
+  - Negative prompt (pipeline owns the string)
+  - Noise generation, encoding, decoding
+
+Flow:
+  raw obs (dataset format)
+    → DreamZeroPipeline selects transform by embodiment
+      → unified dict (stitched video, templated prompt str, raw state)
+        → tokenize, pad, encode, denoise
+          → transform_action_output()
+            → ndarray (N, action_dim)
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+
+
+class RobotPolicyTransform:
+    """Base class for dataset-specific observation transforms.
+
+    Subclasses MUST define:
+        IMAGE_KEY_MAP: dict — dataset obs keys → unified keys
+        EMBODIMENT_NAME: str — embodiment identity (pipeline maps to numeric ID)
+        ACTION_DIM: int — actual action dimensions (for output slicing)
+
+    Subclasses MUST override:
+        _stitch_views()      — multi-view → single stitched image
+        _language_template() — prompt → embodiment-aware template
+        _extract_raw_state() — obs → raw state ndarray
+    """
+
+    IMAGE_KEY_MAP: dict[str, str]
+    EMBODIMENT_NAME: str
+    ACTION_DIM: int
+
+    def transform_input(self, obs: dict) -> dict:
+        """Dataset-specific transform: key map → stitch → template → state."""
+        # 1. Map image keys → unified keys
+        images: dict[str, np.ndarray] = {}
+        for src_key, dst_key in self.IMAGE_KEY_MAP.items():
+            if src_key in obs:
+                images[dst_key] = np.asarray(obs[src_key])
+
+        # 2. Multi-view stitching
+        stitched = self._stitch_views(images)
+
+        # 3. Language template (string only, pipeline tokenizes)
+        prompt = obs.get("prompt", "")
+        templated_prompt = self._language_template(prompt)
+
+        # 4. Raw state extraction (pipeline pads)
+        raw_state = self._extract_raw_state(obs)
+
+        # 5. Build unified output
+        unified: dict[str, Any] = {
+            "images": stitched,  # ndarray (T, H_out, W_out, 3)
+            "prompt": templated_prompt,  # str (templated, not tokenized)
+            "state": raw_state,  # ndarray (state_dim,) — pipeline pads
+            "embodiment_name": self.EMBODIMENT_NAME,
+        }
+        if "session_id" in obs:
+            unified["session_id"] = obs["session_id"]
+        return unified
+
+    def transform_action_output(self, actions: Any) -> np.ndarray:
+        """Adapt model action output to this transform's action dimensions."""
+        actions = np.asarray(actions, dtype=np.float32)
+        # Handle any remaining batch dims: squeeze to 2D (horizon, dim)
+        while actions.ndim > 2:
+            actions = actions[0]
+        # Slice padded dim to actual ACTION_DIM
+        if actions.ndim == 2 and actions.shape[-1] > self.ACTION_DIM:
+            actions = actions[:, : self.ACTION_DIM]
+        return actions
+
+    # ------------------------------------------------------------------
+    # Subclass MUST override
+    # ------------------------------------------------------------------
+
+    def _stitch_views(self, images: dict[str, np.ndarray]) -> np.ndarray:
+        """Stitch camera views into single image.
+        Input: unified key → ndarray (H,W,3) or (T,H,W,3).
+        Output: ndarray (T, H_out, W_out, 3).
+        """
+        raise NotImplementedError
+
+    def _language_template(self, prompt: str) -> str:
+        """Wrap prompt in embodiment-specific template string."""
+        raise NotImplementedError
+
+    def _extract_raw_state(self, obs: dict) -> np.ndarray:
+        """Extract raw state vector from obs.
+        Returns: ndarray (state_dim,) float64. Pipeline handles padding.
+        """
+        raise NotImplementedError
+
+
+# Transform registry — keyed by embodiment/dataset name
+TRANSFORMS: dict[str, RobotPolicyTransform] = {}
+
+
+def register_transform(name: str, transform: RobotPolicyTransform) -> None:
+    TRANSFORMS[name] = transform
+
+
+def get_transform(name: str) -> RobotPolicyTransform:
+    if name not in TRANSFORMS:
+        raise KeyError(f"Unknown transform '{name}'. Available: {list(TRANSFORMS.keys())}")
+    return TRANSFORMS[name]

vllm_omni/diffusion/models/dreamzero/transform/droid.py

@@ -0,0 +1,138 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""DROID dataset transform.
+
+DROID uses 1-indexed exterior cameras, 3 views total (OXE_DROID embodiment).
+Stitching layout (same as RoboArena — both are OXE_DROID):
+    ┌─────────────────────────┐
+    │   wrist (2x width)      │  ← pixel-repeat along width
+    ├────────────┬────────────┤
+    │  left ext  │ right ext  │
+    └────────────┴────────────┘
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import torch
+import torchvision.transforms.v2 as T
+
+from vllm_omni.diffusion.models.dreamzero.transform.base import (
+    RobotPolicyTransform,
+    register_transform,
+)
+
+
+class DroidTransform(RobotPolicyTransform):
+    """Transform for DROID dataset (OXE_DROID embodiment).
+
+    DROID observation keys (1-indexed exterior cameras):
+        observation/exterior_image_1_left  → left exterior
+        observation/exterior_image_2_left  → right exterior
+        observation/wrist_image_left       → wrist
+    """
+
+    IMAGE_KEY_MAP = {
+        "observation/exterior_image_1_left": "images/exterior_0",
+        "observation/exterior_image_2_left": "images/exterior_1",
+        "observation/wrist_image_left": "images/wrist",
+    }
+    EMBODIMENT_NAME = "oxe_droid"
+    ACTION_DIM = 8  # 7 joint + 1 gripper
+    _VIDEO_CROP_SCALE = 0.95
+    _VIDEO_RESIZE_HW = (176, 320)
+
+    @classmethod
+    def _preprocess_view(cls, arr: np.ndarray) -> np.ndarray:
+        """Apply per-view crop and resize before stitching."""
+        frames = torch.from_numpy(arr).to(torch.float32).permute(0, 3, 1, 2) / 255.0
+        crop_h = int(arr.shape[1] * cls._VIDEO_CROP_SCALE)
+        crop_w = int(arr.shape[2] * cls._VIDEO_CROP_SCALE)
+        frames = T.CenterCrop((crop_h, crop_w))(frames)
+        frames = T.Resize(
+            cls._VIDEO_RESIZE_HW,
+            interpolation=T.InterpolationMode.BILINEAR,
+            antialias=True,
+        )(frames)
+        return (frames.permute(0, 2, 3, 1) * 255.0).to(torch.uint8).cpu().numpy()
+
+    def _stitch_views(self, images: dict[str, np.ndarray]) -> np.ndarray:
+        """OXE_DROID 2x2 stitching: wrist top (2x wide), exteriors bottom."""
+        left_ext = images.get("images/exterior_0")
+        right_ext = images.get("images/exterior_1")
+        wrist = images.get("images/wrist")
+
+        # Ensure 4D: (T, H, W, C)
+        def ensure_4d(arr: np.ndarray | None) -> np.ndarray | None:
+            if arr is None:
+                return None
+            return arr if arr.ndim == 4 else arr[np.newaxis]
+
+        left_ext = ensure_4d(left_ext)
+        right_ext = ensure_4d(right_ext)
+        wrist = ensure_4d(wrist)
+
+        # Determine shape from first available view.
+        ref = next((v for v in [wrist, left_ext, right_ext] if v is not None), None)
+        if ref is None:
+            return np.zeros((1, 352, 640, 3), dtype=np.uint8)
+
+        # Apply per-view crop + resize before stitching.
+        def maybe_preprocess(arr: np.ndarray | None) -> np.ndarray | None:
+            if arr is None:
+                return None
+            return self._preprocess_view(arr)
+
+        left_ext = maybe_preprocess(left_ext)
+        right_ext = maybe_preprocess(right_ext)
+        wrist = maybe_preprocess(wrist)
+        ref = next((v for v in [wrist, left_ext, right_ext] if v is not None), None)
+        if ref is None:
+            raise RuntimeError("Expected at least one DROID camera view after preprocessing.")
+        t, h, w, c = ref.shape
+
+        out = np.zeros((t, 2 * h, 2 * w, c), dtype=ref.dtype)  # (T, 2H, 2W, C)
+
+        # Top row: wrist repeated 2x along width.
+        if wrist is not None:
+            wrist_wide = np.repeat(wrist, 2, axis=2)  # (T, H, 2W, C)
+            out[:, :h, :] = wrist_wide
+
+        # Bottom row: left exterior | right exterior.
+        if left_ext is not None:
+            out[:, h:, :w] = left_ext
+        if right_ext is not None:
+            out[:, h:, w:] = right_ext
+
+        return out
+
+    def _language_template(self, prompt: str) -> str:
+        """Expand the language prompt for the OXE_DROID multi-view format."""
+        prompt = (prompt or "Perform the default behavior.").strip()
+        prompt_lower = prompt.lower()
+        return (
+            "A multi-view video shows that a robot "
+            + prompt_lower
+            + " The video is split into three views: The top view shows the "
+            + "camera view from the robot's wrist, the bottom-left view shows "
+            + "the camera view from the left exterior camera, and the "
+            + "bottom-right view shows the camera view from the right exterior "
+            + "camera. During training, one of the two bottom exterior views "
+            + "may be a black screen (dropped view). The robot "
+            + prompt_lower
+        )
+
+    def _extract_raw_state(self, obs: dict) -> np.ndarray:
+        """OXE_DROID state: 7 joint + 1 gripper = 8 dims."""
+        parts = []
+        if "observation/joint_position" in obs:
+            parts.append(np.asarray(obs["observation/joint_position"], dtype=np.float64).flatten())
+        if "observation/gripper_position" in obs:
+            parts.append(np.asarray(obs["observation/gripper_position"], dtype=np.float64).flatten())
+        if parts:
+            return np.concatenate(parts)
+        return np.zeros(8, dtype=np.float64)
+
+
+register_transform("droid", DroidTransform())

vllm_omni/diffusion/models/dreamzero/transform/roboarena.py

@@ -0,0 +1,41 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""RoboArena dataset transform.
+
+RoboArena uses 0-indexed exterior cameras, 3 views total (OXE_DROID embodiment).
+Same stitching layout as DROID — both map to OXE_DROID in DreamZero.
+
+"""
+
+from __future__ import annotations
+
+from vllm_omni.diffusion.models.dreamzero.transform.base import (
+    register_transform,
+)
+from vllm_omni.diffusion.models.dreamzero.transform.droid import (
+    DroidTransform,
+)
+
+
+class RoboArenaTransform(DroidTransform):
+    """Transform for RoboArena dataset.
+
+    Same embodiment as DROID (OXE_DROID), same stitching and template.
+    Only difference: 0-indexed exterior camera keys.
+
+    RoboArena observation keys (0-indexed):
+        observation/exterior_image_0_left  → left exterior
+        observation/exterior_image_1_left  → right exterior
+        observation/wrist_image_left       → wrist
+
+    """
+
+    IMAGE_KEY_MAP = {
+        "observation/exterior_image_0_left": "images/exterior_0",
+        "observation/exterior_image_1_left": "images/exterior_1",
+        "observation/wrist_image_left": "images/wrist",
+    }
+
+
+register_transform("roboarena", RoboArenaTransform())

vllm_omni/diffusion/models/dreamzero/utils.py

@@ -0,0 +1,29 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""DreamZero model constants shared by the pipeline."""
+
+DEFAULT_NUM_INFERENCE_STEPS = 16
+DEFAULT_CFG_SCALE = 5.0
+DEFAULT_SIGMA_SHIFT = 5.0
+DEFAULT_SEED = 1140
+
+DEFAULT_NEGATIVE_PROMPT = (
+    "Vibrant colors, overexposed, static, blurry details, text, subtitles, "
+    "style, artwork, painting, image, still, grayscale, dull, worst quality, "
+    "low quality, JPEG artifacts, ugly, mutilated, extra fingers, bad hands, "
+    "bad face, deformed, disfigured, mutated limbs, fused fingers, stagnant "
+    "image, cluttered background, three legs, many people in the background, "
+    "walking backwards."
+)
+
+DEFAULT_EMBODIMENT_NAME_TO_ID = {
+    "oxe_droid": 17,
+    "agibot": 26,
+    "gr1_unified": 24,
+    "xdof": 22,
+    "yam": 32,
+    "mecka_hands": 27,
+    "lapa": 27,
+    "dream": 31,
+}

vllm_omni/diffusion/models/dreamzero/video_export_worker.py

@@ -0,0 +1,21 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import torch
+
+
+class DreamZeroVideoExportWorkerExtension:
+    """DreamZero worker RPCs used by offline example video export."""
+
+    def decode_video_latents_to_uint8(self, video_latents: torch.Tensor) -> torch.Tensor:
+        if self.model_runner is None or self.model_runner.pipeline is None:
+            raise RuntimeError("DreamZero pipeline is not initialized on this worker.")
+
+        with torch.inference_mode():
+            decoded = self.model_runner.pipeline.decode_video_latents(video_latents)
+            decoded = decoded.squeeze(0).permute(1, 2, 3, 0).contiguous()
+            decoded = decoded.clamp(-1, 1) * 0.5 + 0.5
+            decoded = (decoded * 255.0).round().to(torch.uint8).cpu()
+        return decoded

vllm_omni/diffusion/registry.py

@@ -281,6 +281,11 @@ _DIFFUSION_MODELS = {
         "pipeline_hidream_image",
         "HiDreamImagePipeline",
     ),
+    "DreamZeroPipeline": (
+        "dreamzero",
+        "pipeline_dreamzero",
+        "DreamZeroPipeline",
+    ),
 }
 
 

vllm_omni/diffusion/utils/hf_utils.py

@@ -1,4 +1,5 @@
 import os
+from collections.abc import Mapping
 from functools import lru_cache
 
 from vllm.logger import init_logger
@@ -27,6 +28,31 @@ def _looks_like_bagel(model_name: str) -> bool:
         return False
 
 
+def _looks_like_dreamzero(model_name: str) -> bool:
+    """Best-effort detection for DreamZero-style VLA diffusion checkpoints."""
+    try:
+        cfg = get_hf_file_to_dict("config.json", model_name)
+        if cfg.get("model_type") != "vla":
+            return False
+        action_head_cfg = cfg.get("action_head_cfg") or {}
+        if not isinstance(action_head_cfg, Mapping):
+            return False
+        action_head_config = action_head_cfg.get("config") or {}
+        if not isinstance(action_head_config, Mapping):
+            return False
+        diffusion_model_cfg = action_head_config.get("diffusion_model_cfg") or {}
+        if not isinstance(diffusion_model_cfg, Mapping):
+            return False
+        return (
+            action_head_cfg.get("_target_")
+            == "groot.vla.model.dreamzero.action_head.wan_flow_matching_action_tf.WANPolicyHead"
+            and diffusion_model_cfg.get("_target_")
+            == ("groot.vla.model.dreamzero.modules.wan_video_dit_action_casual_chunk.CausalWanModel")
+        )
+    except Exception:
+        return False
+
+
 @lru_cache
 def is_diffusion_model(model_name: str) -> bool:
     """Check if a model is a diffusion model.
@@ -72,6 +98,7 @@ def is_diffusion_model(model_name: str) -> bool:
     except Exception as e:
         logger.debug("Failed to load diffusers config via DiffusionPipeline: %s", e)
 
-        # Bagel is not a diffusers pipeline (no model_index.json), but is still a
-        # diffusion-style model in vllm-omni. Detect it via config.json.
-    return _looks_like_bagel(model_name)
+        # Bagel and DreamZero are not diffusers pipelines (no model_index.json),
+        # but are still diffusion-style models in vllm-omni. Detect them via
+        # config.json.
+    return _looks_like_bagel(model_name) or _looks_like_dreamzero(model_name)

vllm_omni/entrypoints/cli/benchmark/serve.py

@@ -131,6 +131,35 @@ def add_seed_tts_cli_args(parser: argparse.ArgumentParser) -> None:
     )
 
 
+_OMNI_BENCH_DATASET_CHOICES = (
+    "daily-omni",
+    "seed-tts",
+    "seed-tts-text",
+    "seed-tts-design",
+    "ttsd",
+    "sound-effect",
+)
+
+
+def _extend_omni_dataset_name_choices(parser: argparse.ArgumentParser) -> None:
+    """Append omni benchmark dataset names to --dataset-name choices.
+
+    TrackingArgumentParser keeps a shadow parser for explicit-arg tracking; both
+    parsers must list the same choices or parse_args rejects valid omni values.
+    """
+    parsers = [parser]
+    shadow = getattr(parser, "_shadow", None)
+    if shadow is not None:
+        parsers.append(shadow)
+
+    for p in parsers:
+        for action in p._actions:
+            if action.dest == "dataset_name" and action.choices is not None:
+                extra = [c for c in _OMNI_BENCH_DATASET_CHOICES if c not in action.choices]
+                if extra:
+                    action.choices = list(action.choices) + extra
+
+
 class OmniBenchmarkServingSubcommand(OmniBenchmarkSubcommandBase):
     """The `serve` subcommand for vllm bench."""
 
@@ -145,22 +174,7 @@ class OmniBenchmarkServingSubcommand(OmniBenchmarkSubcommandBase):
         add_daily_omni_cli_args(parser)
         add_seed_tts_cli_args(parser)
 
-        for action in parser._actions:
-            if action.dest == "dataset_name" and action.choices is not None:
-                extra = [
-                    c
-                    for c in (
-                        "daily-omni",
-                        "seed-tts",
-                        "seed-tts-text",
-                        "seed-tts-design",
-                        "ttsd",
-                        "sound-effect",
-                    )
-                    if c not in action.choices
-                ]
-                if extra:
-                    action.choices = list(action.choices) + extra
+        _extend_omni_dataset_name_choices(parser)
 
         # Update help messages for omni-specific features
         for action in parser._actions:

vllm_omni/entrypoints/openai/api_server.py

@@ -2627,6 +2627,8 @@ async def _parse_video_form(
     flow_shift: float | None = Form(default=None),
     true_cfg_scale: float | None = Form(default=None),
     seed: int | None = Form(default=None),
+    generate_sound: bool | None = Form(default=None),
+    sound_duration: float | None = Form(default=None, gt=0.0),
     negative_prompt: str | None = Form(default=None),
     enable_frame_interpolation: bool | None = Form(default=None),
     frame_interpolation_exp: int | None = Form(default=None, ge=1),
@@ -2667,6 +2669,8 @@ async def _parse_video_form(
         "flow_shift": flow_shift,
         "true_cfg_scale": true_cfg_scale,
         "seed": seed,
+        "generate_sound": generate_sound,
+        "sound_duration": sound_duration,
         "negative_prompt": negative_prompt,
         "enable_frame_interpolation": enable_frame_interpolation,
         "frame_interpolation_exp": frame_interpolation_exp,

vllm_omni/entrypoints/openai/protocol/videos.py

@@ -149,6 +149,15 @@ class VideoGenerationRequest(BaseModel):
         description="True CFG scale (model-specific parameter, may be ignored if not supported)",
     )
     seed: int | None = Field(default=None, description="Random seed for reproducibility")
+    generate_sound: bool = Field(
+        default=False,
+        description="Request model-generated audio for video models that support sound generation.",
+    )
+    sound_duration: float | None = Field(
+        default=None,
+        gt=0.0,
+        description="Duration in seconds for model-generated audio. Defaults to the generated video duration.",
+    )
 
     # vllm-omni extensions for post-generation frame interpolation.
     enable_frame_interpolation: bool = Field(

vllm_omni/entrypoints/openai/serving_video.py

@@ -148,6 +148,10 @@ class OmniOpenAIServingVideo:
         )
         if "flow_shift" in provided_fields and request.flow_shift is not None:
             gen_params.extra_args["flow_shift"] = request.flow_shift
+        if "generate_sound" in provided_fields:
+            gen_params.extra_args["generate_sound"] = request.generate_sound
+        if "sound_duration" in provided_fields and request.sound_duration is not None:
+            gen_params.extra_args["sound_duration"] = request.sound_duration
 
         # Apply model-specific extra parameters
         if request.extra_params is not None:

vllm_omni/entrypoints/utils.py

@@ -13,6 +13,7 @@ from vllm.transformers_utils.repo_utils import file_or_path_exists
 
 from vllm_omni.config.stage_config import StageConfigFactory
 from vllm_omni.config.yaml_util import create_config, load_yaml_config, merge_configs
+from vllm_omni.diffusion.utils.hf_utils import _looks_like_dreamzero
 from vllm_omni.entrypoints.stage_utils import _to_dict
 from vllm_omni.inputs.data import OmniSamplingParams
 from vllm_omni.platforms import current_omni_platform
@@ -287,6 +288,9 @@ def resolve_model_config_path(model: str) -> str:
                 )
 
     default_config_path = current_omni_platform.get_default_stage_config_path()
+    if model_type == "vla" and _looks_like_dreamzero(model):
+        model_type = "dreamzero"
+
     if model_type in _DIFFUSERS_CLASS_TO_CONFIG:
         normalized_model_type = _DIFFUSERS_CLASS_TO_CONFIG[model_type]
     else:

vllm_omni/model_executor/models/dreamzero/__init__.py

@@ -0,0 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm_omni.model_executor.models.dreamzero.pipeline import DREAMZERO_PIPELINE
+
+__all__ = ["DREAMZERO_PIPELINE"]

vllm_omni/model_executor/models/dreamzero/pipeline.py

@@ -0,0 +1,25 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""DreamZero single-stage diffusion topology."""
+
+from vllm_omni.config.stage_config import (
+    PipelineConfig,
+    StageExecutionType,
+    StagePipelineConfig,
+)
+
+DREAMZERO_PIPELINE = PipelineConfig(
+    model_type="dreamzero",
+    model_arch="DreamZeroPipeline",
+    stages=(
+        StagePipelineConfig(
+            stage_id=0,
+            model_stage="diffusion",
+            execution_type=StageExecutionType.DIFFUSION,
+            input_sources=(),
+            final_output=True,
+            final_output_type="image",
+            model_arch="DreamZeroPipeline",
+        ),
+    ),
+)

vllm_omni/model_executor/models/dynin_omni/dynin_omni_common.py

@@ -144,7 +144,6 @@ _DYNIN_CONFIG_CANDIDATE_RELPATHS = (
     "configs/dynin_omni.yaml",
     "models/configs/dynin_omni.yaml",
     "vllm_omni/model_executor/models/dynin_omni/configs/dynin_omni.yaml",
-    "vllm_omni/model_executor/stage_configs/dynin_omni.yaml",
     "dynin_omni.yaml",
 )
 
@@ -482,7 +481,6 @@ def _resolve_config_path(vllm_config: VllmConfig, runtime_info: dict[str, Any])
     for bundled in (
         module_root / "configs" / "dynin_omni.yaml",
         module_root / "models" / "configs" / "dynin_omni.yaml",
-        module_root.parent / "stage_configs" / "dynin_omni.yaml",
     ):
         if bundled.exists():
             return str(bundled)

vllm_omni/model_executor/models/dynin_omni/pipeline.py

@@ -0,0 +1,64 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Dynin-Omni pipeline topology (frozen).
+
+Stage 0: token2text  — multimodal understanding / text generation (comprehension)
+Stage 1: token2image — stage-0 tokens → image latents
+Stage 2: token2audio — stage-1 tokens → audio latents
+
+All three stages run on the generation worker (``LLM_GENERATION``); the
+inter-stage hand-off uses the worker-connector full-payload data plane
+(``*_full_payload`` producers + ``token2text_to_token2image`` /
+``token2image_to_token2audio`` consumers). Deploy knobs (devices, GPU memory,
+batched tokens, connectors) live in ``vllm_omni/deploy/dynin_omni*.yaml``.
+"""
+
+from vllm_omni.config.stage_config import (
+    PipelineConfig,
+    StageExecutionType,
+    StagePipelineConfig,
+)
+
+_PROC = "vllm_omni.model_executor.stage_input_processors.dynin_omni"
+
+DYNIN_OMNI_PIPELINE = PipelineConfig(
+    model_type="dynin_omni",
+    model_arch="DyninOmniForConditionalGeneration",
+    # Arch-fallback safety net: route by hf_config.architectures when the
+    # auto-detected model_type does not match the registry key exactly.
+    hf_architectures=("DyninOmniForConditionalGeneration",),
+    stages=(
+        StagePipelineConfig(
+            stage_id=0,
+            model_stage="token2text",
+            execution_type=StageExecutionType.LLM_GENERATION,
+            input_sources=(),
+            final_output=True,
+            final_output_type="text",
+            owns_tokenizer=True,
+            engine_output_type="latent",
+            custom_process_next_stage_input_func=f"{_PROC}.token2text_to_token2image_full_payload",
+        ),
+        StagePipelineConfig(
+            stage_id=1,
+            model_stage="token2image",
+            execution_type=StageExecutionType.LLM_GENERATION,
+            input_sources=(0,),
+            final_output=True,
+            final_output_type="image",
+            engine_output_type="latent",
+            custom_process_input_func=f"{_PROC}.token2text_to_token2image",
+            custom_process_next_stage_input_func=f"{_PROC}.token2image_to_token2audio_full_payload",
+        ),
+        StagePipelineConfig(
+            stage_id=2,
+            model_stage="token2audio",
+            execution_type=StageExecutionType.LLM_GENERATION,
+            input_sources=(1,),
+            final_output=True,
+            final_output_type="audio",
+            engine_output_type="latent",
+            custom_process_input_func=f"{_PROC}.token2image_to_token2audio",
+        ),
+    ),
+)

vllm_omni/model_executor/stage_configs/dynin_omni.yaml

@@ -1,76 +0,0 @@
-stage_args:
-  - stage_id: 0
-    stage_type: llm
-    runtime:
-      devices: "0"
-      max_batch_size: 1
-    engine_args:
-      model_stage: token2text
-      custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.dynin_omni.token2text_to_token2image_full_payload
-      model_arch: DyninOmniForConditionalGeneration
-      worker_type: generation
-      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
-      engine_output_type: latent
-      trust_remote_code: true
-      gpu_memory_utilization: 0.5
-      enforce_eager: true
-      enable_prefix_caching: false
-      async_scheduling: false
-      max_num_batched_tokens: 32768
-    is_comprehension: true
-    final_output: true
-    final_output_type: text
-
-  - stage_id: 1
-    stage_type: llm
-    runtime:
-      devices: "0"
-      max_batch_size: 1
-    engine_args:
-      model_stage: token2image
-      model_arch: DyninOmniForConditionalGeneration
-      worker_type: generation
-      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
-      engine_output_type: latent
-      trust_remote_code: true
-      gpu_memory_utilization: 0.1
-      enforce_eager: true
-      enable_prefix_caching: false
-      async_scheduling: false
-      max_num_batched_tokens: 32768
-    engine_input_source: [0]
-    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.dynin_omni.token2text_to_token2image
-    final_output: true
-    final_output_type: image
-
-  - stage_id: 2
-    stage_type: llm
-    runtime:
-      devices: "0"
-      max_batch_size: 1
-    engine_args:
-      model_stage: token2audio
-      model_arch: DyninOmniForConditionalGeneration
-      worker_type: generation
-      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
-      engine_output_type: latent
-      trust_remote_code: true
-      gpu_memory_utilization: 0.1
-      enforce_eager: true
-      enable_prefix_caching: false
-      async_scheduling: false
-      max_num_batched_tokens: 32768
-    engine_input_source: [1]
-    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.dynin_omni.token2image_to_token2audio
-    final_output: true
-    final_output_type: audio
-
-# Top-level runtime config (concise): default windows and stage edges
-runtime:
-  enabled: true
-
-  edges:
-    - from: 0
-      to: 1
-    - from: 1
-      to: 2

vllm_omni/model_executor/stage_configs/dynin_omni_multiconnector.yaml

@@ -1,110 +0,0 @@
-stage_args:
-  - stage_id: 0
-    stage_type: llm
-    runtime:
-      devices: "0"
-      max_batch_size: 1
-    engine_args:
-      model_stage: token2text
-      custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.dynin_omni.token2text_to_token2image_full_payload
-      model_arch: DyninOmniForConditionalGeneration
-      worker_type: generation
-      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
-      engine_output_type: latent
-      trust_remote_code: false
-      enforce_eager: true
-      enable_prefix_caching: false
-      async_scheduling: false
-      max_num_batched_tokens: 32768
-    output_connectors:
-      to_stage_1: mooncake_connector
-    final_output: true
-    final_output_type: text
-
-  - stage_id: 1
-    stage_type: llm
-    runtime:
-      devices: "0"
-      max_batch_size: 1
-    engine_args:
-      model_stage: token2image
-      model_arch: DyninOmniForConditionalGeneration
-      worker_type: generation
-      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
-      engine_output_type: latent
-      trust_remote_code: false
-      enforce_eager: true
-      enable_prefix_caching: false
-      async_scheduling: false
-      max_num_batched_tokens: 32768
-    engine_input_source: [0]
-    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.dynin_omni.token2text_to_token2image
-    final_output: true
-    final_output_type: image
-    input_connectors:
-      from_stage_0: mooncake_connector
-    output_connectors:
-      to_stage_2: mooncake_connector
-
-  - stage_id: 2
-    stage_type: llm
-    runtime:
-      devices: "0"
-      max_batch_size: 1
-    engine_args:
-      model_stage: token2audio
-      model_arch: DyninOmniForConditionalGeneration
-      worker_type: generation
-      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
-      engine_output_type: latent
-      trust_remote_code: false
-      enforce_eager: true
-      enable_prefix_caching: false
-      async_scheduling: false
-      max_num_batched_tokens: 32768
-    engine_input_source: [1]
-    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.dynin_omni.token2image_to_token2audio
-    final_output: true
-    final_output_type: audio
-    input_connectors:
-      from_stage_1: mooncake_connector
-
-# Top-level runtime config (concise): default windows and stage edges
-runtime:
-  enabled: true
-  ####
-  # same as Qwen2.5_omni version
-  # Distributed connectors configuration (optional)
-  # More connectors will be supported in the future.
-  connectors:
-    # Mooncake connector for cross-node/intra-node communication
-    mooncake_connector:
-      name: MooncakeConnector
-      extra:
-        host: "127.0.0.1"
-        metadata_server: "http://10.90.67.86:8080/metadata"
-        master: "10.90.67.86:50051"
-        segment: 512000000    # 512MB
-        localbuf: 64000000     # 64MB
-        proto: "tcp"
-
-    # Yuanrong connector for cross-node/intra-node communication
-    yuanrong_connector:
-      name: YuanrongConnector
-      extra:
-        host: "127.0.0.1"
-        port: "35000"
-
-    # SharedMemory connector for intra-node communication
-    # Alternative SHM connector with different threshold
-    shared_memory_connector:
-      name: SharedMemoryConnector
-      extra:
-        shm_threshold_bytes: 65536 # 64KB threshold
-    ####
-
-  edges:
-    - from: 0
-      to: 1
-    - from: 1
-      to: 2