Harden the shared video-audio example

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>

docs/user_guide/examples/offline_inference/x_to_video_audio.md

@@ -3,7 +3,25 @@
 Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/x_to_video_audio>.
 
 
-The `DreamID-Omni` pipeline generates short videos from text, image and video.
+The shared runner generates video with optional audio from text, image, and audio
+inputs. It can be extended through `vllm_omni.model_extras` without adding
+another model-specific runner.
+
+## Shared request and output contract
+
+The example builds one canonical request envelope: `prompt` is text,
+`modalities` is `["video"]`, and optional `multi_modal_data` contains RGB PIL
+images plus audio `(waveform, sample_rate)` tuples. Audio files are loaded in
+full; the runner does not silently select or crop a time range. Prepare a
+trimmed input file explicitly when only part of a recording should condition
+the model.
+
+Model-specific request translation (for example, DreamID-Omni's two negative
+prompts and its 16 kHz input requirement) is declared in `model_extras`. Model
+outputs must provide `fps` and, when audio is present, `audio_sample_rate`.
+`--fps` and `--audio-sample-rate` are available as explicit metadata overrides.
+The shared encoder normalizes supported video tensor layouts, so this runner
+does not branch on a model name or assume a model-specific frame layout.
 
 ## Local CLI Usage
 ### Download the Model locally
@@ -91,7 +109,8 @@ python x_to_video_audio.py \
 Key arguments:
 - `--prompt`: text description (string).
 - `--model`: path to the model local directory.
-- `--height/--width`: output resolution (defaults 704 * 1024).
+- `--model-class-name`: optional diffusion pipeline class override.
+- `--height/--width`: output resolution (defaults 704 * 1280).
 - `--image-path`: path to the input image list.
 - `--audio-path`: path to the input audio list, indicate the timbre of the output video.
 - `--cfg-parallel-size`: number of parallel cfg parallel (defaults 1).
@@ -101,6 +120,9 @@ Key arguments:
 - `--num-inference-steps`: number of denoising steps (defaults 45).
 - `--video-negative-prompt`: negative prompt for video generation.
 - `--audio-negative-prompt`: negative prompt for audio generation.
+- `--extra-body`: JSON object containing parameters declared by the selected model.
+- `--fps`: explicit output FPS override when a model does not return that metadata.
+- `--audio-sample-rate`: explicit output audio sample-rate override.
 - `--enable-cpu-offload`: enable CPU offload (defaults False).
 
 ## Example materials

examples/offline_inference/x_to_video_audio/x_to_video_audio.md

@@ -1,6 +1,24 @@
 # X-To-Video-Audio
 
-The `DreamID-Omni` pipeline generates short videos from text, image and video.
+The shared runner generates video with optional audio from text, image, and audio
+inputs. It currently documents DreamID-Omni and MagiHuman, and can be extended
+through `vllm_omni.model_extras` without adding another model-specific runner.
+
+## Shared request and output contract
+
+The example builds one canonical request envelope: `prompt` is text,
+`modalities` is `["video"]`, and optional `multi_modal_data` contains RGB PIL
+images plus audio `(waveform, sample_rate)` tuples. Audio files are loaded in
+full; the runner does not silently select or crop a time range. Prepare a
+trimmed input file explicitly when only part of a recording should condition
+the model.
+
+Model-specific request translation (for example, DreamID-Omni's two negative
+prompts and its 16 kHz input requirement) is declared in `model_extras`. Model
+outputs must provide `fps` and, when audio is present, `audio_sample_rate`.
+`--fps` and `--audio-sample-rate` are available as explicit metadata overrides.
+The shared encoder normalizes supported video tensor layouts, so this runner
+does not branch on a model name or assume a model-specific frame layout.
 
 ## Local CLI Usage
 ### Download the Model locally
@@ -147,7 +165,8 @@ in `vllm_omni/model_extras/magi_human.py`, routed via `extra_args`):
 Key arguments:
 - `--prompt`: text description (string).
 - `--model`: path to the model local directory.
-- `--height/--width`: output resolution (defaults 704 * 1024).
+- `--model-class-name`: optional diffusion pipeline class override.
+- `--height/--width`: output resolution (defaults 704 * 1280).
 - `--image-path`: path to the input image list.
 - `--audio-path`: path to the input audio list, indicate the timbre of the output video.
 - `--cfg-parallel-size`: number of parallel cfg parallel (defaults 1).
@@ -157,6 +176,9 @@ Key arguments:
 - `--num-inference-steps`: number of denoising steps (defaults 45).
 - `--video-negative-prompt`: negative prompt for video generation.
 - `--audio-negative-prompt`: negative prompt for audio generation.
+- `--extra-body`: JSON object containing parameters declared by the selected model.
+- `--fps`: explicit output FPS override when a model does not return that metadata.
+- `--audio-sample-rate`: explicit output audio sample-rate override.
 - `--enable-cpu-offload`: enable CPU offload (defaults False).
 - `--cache-backend`: enable `cache_dit` for acceleration.
 - `--quantization`: online (dynamic) quantization method — `fp8` or `int8`.

examples/offline_inference/x_to_video_audio/x_to_video_audio.py

@@ -3,32 +3,74 @@
 
 import argparse
 import json
-import re
+import math
 import time
+from collections.abc import Mapping
+from dataclasses import dataclass
+from pathlib import Path
 from typing import Any
 
-import numpy as np
 from PIL import Image
 from vllm.multimodal.media.audio import load_audio
 
 from vllm_omni.diffusion.data import DiffusionParallelConfig
-from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
+from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
 from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.entrypoints.openai.video_api_utils import encode_video_bytes
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.model_extras import (
+    build_x_to_video_audio_prompt as build_model_x_to_video_audio_prompt,
+)
+from vllm_omni.model_extras import (
+    get_extra_body_params,
+    get_input_audio_sample_rate,
+    get_model_class_name,
+)
 
 IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".gif"}
 AUDIO_EXTENSIONS = {".wav", ".mp3", ".flac", ".m4a", ".aac", ".ogg"}
 
 
+@dataclass(frozen=True)
+class XToVideoAudioOutput:
+    """Canonical generated video+audio payload and its required metadata."""
+
+    video: Any
+    audio: Any | None
+    fps: float
+    audio_sample_rate: int | None
+
+
+def parse_json_object(value: str) -> dict[str, Any]:
+    try:
+        parsed = json.loads(value)
+    except json.JSONDecodeError as exc:
+        raise argparse.ArgumentTypeError(f"--extra-body must be valid JSON: {exc}") from exc
+    if not isinstance(parsed, dict):
+        raise argparse.ArgumentTypeError("--extra-body must be a JSON object")
+    return parsed
+
+
+def build_x_to_video_audio_prompt(
+    prompt: str,
+    media_inputs: dict[str, Any] | None = None,
+) -> dict[str, Any]:
+    """Build the canonical request envelope for image/audio-to-video models."""
+    result: dict[str, Any] = {"prompt": prompt, "modalities": ["video"]}
+    if media_inputs:
+        result["multi_modal_data"] = media_inputs
+    return result
+
+
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(description="Offline inference for X -> video+audio models.")
     parser.add_argument("--model", required=True, help="Model ckpt root directory.")
     parser.add_argument(
         "--model-type",
-        default="dreamid-omni",
-        choices=["dreamid-omni", "magi-human"],
-        help="Model type.",
+        default=None,
+        help="Optional legacy model type override (for example, magi-human).",
     )
+    parser.add_argument("--model-class-name", default=None, help="Optional diffusion pipeline class override.")
     parser.add_argument("--prompt", default=None, help="Text prompt.")
 
     parser.add_argument("--image-path", type=str, nargs="+", help="list of image-path")
@@ -36,9 +78,9 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--prompt-file", type=str, default=None, help="Text prompt in json format.")
     parser.add_argument(
         "--extra-body",
-        type=str,
+        type=parse_json_object,
         default=None,
-        help="[magi-human] JSON dict of model-specific extra params (declared in vllm_omni/model_extras/), "
+        help="JSON dict of model-specific extra params (declared in vllm_omni/model_extras/), "
         'merged into sampling extra_args. Example: \'{"image_path": "/path/to/img.jpg", "seconds": 5}\'.',
     )
 
@@ -71,13 +113,20 @@ def parse_args() -> argparse.Namespace:
         default="robotic, muffled, echo, distorted",
         help="Negative prompt for audio.",
     )
+    parser.add_argument("--fps", type=float, default=None, help="Override output FPS metadata.")
+    parser.add_argument(
+        "--audio-sample-rate",
+        type=int,
+        default=None,
+        help="Override output audio sample-rate metadata.",
+    )
     parser.add_argument("--output", default="dreamid_output.mp4", help="Output video path.")
     parser.add_argument(
         "--quantization",
         type=str,
         default=None,
         choices=["fp8", "int8"],
-        help=("Online (dynamic) quantization method for the DreamID-Omni transformer. "),
+        help="Online (dynamic) quantization method for the model transformer.",
     )
     parser.add_argument(
         "--enable-cpu-offload",
@@ -94,7 +143,7 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--use-hsdp",
         action="store_true",
-        help="Enable HSDP for DreamID-Omni fused transformer blocks.",
+        help="Enable HSDP for supported transformer blocks.",
     )
     parser.add_argument(
         "--hsdp-shard-size",
@@ -111,20 +160,85 @@ def parse_args() -> argparse.Namespace:
     return parser.parse_args()
 
 
-def load_image_and_audio(image_paths, audio_paths):
-    image = []
-    audio = []
+def load_image_and_audio(
+    image_paths: list[str] | None,
+    audio_paths: list[str] | None,
+    *,
+    audio_sample_rate: int | None = None,
+) -> tuple[list[Image.Image], list[tuple[Any, int]]]:
+    """Load complete media inputs into the canonical multimodal representation."""
+    images = []
+    audios = []
 
-    for path in image_paths:
+    for path in image_paths or []:
         with Image.open(path) as img:
-            img = img.convert("RGB")
-            image.append(img)
+            images.append(img.convert("RGB"))
 
-    for path in audio_paths:
-        audio_array, sr = load_audio(path, sr=16000)
-        audio_array = audio_array[int(sr * 1) : int(sr * 3)]
-        audio.append(audio_array)
-    return image, audio
+    for path in audio_paths or []:
+        waveform, sample_rate = load_audio(path, sr=audio_sample_rate)
+        audios.append((waveform, sample_rate))
+    return images, audios
+
+
+def extract_x_to_video_audio_output(
+    outputs: Any,
+    *,
+    fps: float | None = None,
+    audio_sample_rate: int | None = None,
+) -> XToVideoAudioOutput:
+    """Normalize engine output without depending on a model's tensor layout."""
+    if isinstance(outputs, list):
+        if not outputs:
+            raise RuntimeError("No output returned from the model.")
+        result = outputs[0]
+    else:
+        result = outputs
+    if result is None:
+        raise RuntimeError("No output returned from the model.")
+
+    multimodal_output = getattr(result, "multimodal_output", None) or {}
+    if not isinstance(multimodal_output, Mapping):
+        raise RuntimeError("Model multimodal output must be a mapping.")
+    images = getattr(result, "images", None)
+    video = images[0] if images else multimodal_output.get("video")
+    if video is None:
+        raise RuntimeError("No video payload found in model output.")
+
+    resolved_fps = fps if fps is not None else multimodal_output.get("fps")
+    try:
+        resolved_fps = float(resolved_fps)
+    except (TypeError, ValueError) as exc:
+        raise RuntimeError("Model output must declare a valid 'fps', or pass --fps.") from exc
+    if not math.isfinite(resolved_fps) or resolved_fps <= 0:
+        raise RuntimeError("Output FPS must be finite and positive.")
+
+    audio = multimodal_output.get("audio")
+    resolved_sample_rate = audio_sample_rate
+    if resolved_sample_rate is None:
+        resolved_sample_rate = multimodal_output.get("audio_sample_rate")
+    if audio is not None:
+        try:
+            resolved_sample_rate = int(resolved_sample_rate)
+        except (TypeError, ValueError) as exc:
+            raise RuntimeError(
+                "Model output with audio must declare a valid 'audio_sample_rate', or pass --audio-sample-rate."
+            ) from exc
+        if resolved_sample_rate <= 0:
+            raise RuntimeError("Output audio sample rate must be positive.")
+    elif resolved_sample_rate is not None:
+        resolved_sample_rate = int(resolved_sample_rate)
+
+    return XToVideoAudioOutput(video, audio, resolved_fps, resolved_sample_rate)
+
+
+def encode_x_to_video_audio_output(output: XToVideoAudioOutput) -> bytes:
+    """Encode any video layout supported by the shared video API utility."""
+    return encode_video_bytes(
+        output.video,
+        output.fps,
+        output.audio,
+        output.audio_sample_rate,
+    )
 
 
 def _extract_peak_memory_mb(result: Any) -> float:
@@ -157,45 +271,8 @@ def main() -> None:
     if args.prompt_file:
         with open(args.prompt_file) as f:
             text_prompt = json.load(f)
-            text_prompt = re.sub(
-                r"\[SPEAKER_TIMESTAMPS_START\].*?\[SPEAKER_TIMESTAMPS_END\]", "", text_prompt, flags=re.DOTALL
-            ).strip()
-            text_prompt = re.sub(
-                r"\[AUDIO_DESCRIPTION_START].*?\[AUDIO_DESCRIPTION_END]", "", text_prompt, flags=re.DOTALL
-            ).strip()
-            text_prompt = re.sub(r"\[[A-Z_]+\]", "", text_prompt)
-            text_prompt = re.sub(r"\n\s*\n", "\n", text_prompt).strip()
-
-    if args.model_type == "magi-human":
-        prompt = text_prompt
-        extra_args = json.loads(args.extra_body) if args.extra_body else {}
-        sampling_params = OmniDiffusionSamplingParams(
-            height=args.height,
-            width=args.width,
-            num_inference_steps=args.num_inference_steps,
-            seed=args.seed,
-            extra_args=extra_args,
-        )
-    else:
-        image, audio = load_image_and_audio(args.image_path, args.audio_path)
-
-        prompt = {
-            "prompt": text_prompt,
-            "video_negative_prompt": args.video_negative_prompt,
-            "audio_negative_prompt": args.audio_negative_prompt,
-            "multi_modal_data": {"image": image, "audio": audio},
-        }
-
-        sampling_params = OmniDiffusionSamplingParams(
-            height=args.height,
-            width=args.width,
-            num_inference_steps=args.num_inference_steps,
-            seed=args.seed,
-            extra_args={
-                "solver_name": args.solver_name,
-                "shift": args.shift,
-            },
-        )
+    if not isinstance(text_prompt, str):
+        raise ValueError("Prompt content must be a JSON string or --prompt text.")
 
     parallel_config = DiffusionParallelConfig(
         cfg_parallel_size=args.cfg_parallel_size,
@@ -220,64 +297,82 @@ def main() -> None:
             "scm_steps_policy": "dynamic",
         }
 
-    omni_kwargs = dict(
+    omni_kwargs: dict[str, Any] = dict(
         model=args.model,
         parallel_config=parallel_config,
-        model_type=args.model_type,
         enable_cpu_offload=args.enable_cpu_offload,
         enable_layerwise_offload=args.enable_layerwise_offload,
         cache_backend=args.cache_backend,
         cache_config=cache_config,
     )
+    if args.model_type is not None:
+        omni_kwargs["model_type"] = args.model_type
+    if args.model_class_name is not None:
+        omni_kwargs["model_class_name"] = args.model_class_name
     if args.quantization is not None:
         omni_kwargs["quantization"] = args.quantization
     omni = Omni(**omni_kwargs)
-    start = time.perf_counter()
-    outputs = omni.generate(prompt, sampling_params)
-    elapsed = time.perf_counter() - start
+    try:
+        model_class_name = get_model_class_name(omni) or args.model_class_name
+        input_sample_rate = get_input_audio_sample_rate(model_class_name)
+        images, audios = load_image_and_audio(
+            args.image_path,
+            args.audio_path,
+            audio_sample_rate=input_sample_rate,
+        )
+        media_inputs: dict[str, Any] = {}
+        if images:
+            media_inputs["image"] = images
+        if audios:
+            media_inputs["audio"] = audios
 
-    if not outputs:
-        raise RuntimeError("No output returned from the model.")
-    result = outputs[0]
-    if not result.images:
-        raise RuntimeError("No video frames found in model output.")
-    generated_video = result.images[0]
-    mm = result.multimodal_output or {}
-    generated_audio = mm.get("audio")
+        canonical_prompt = build_x_to_video_audio_prompt(text_prompt, media_inputs)
+        prompt = build_model_x_to_video_audio_prompt(
+            model_class_name,
+            canonical_prompt,
+            {
+                "video_negative_prompt": args.video_negative_prompt,
+                "audio_negative_prompt": args.audio_negative_prompt,
+            },
+        )
 
-    if args.model_type == "magi-human":
-        # MagiHuman returns frames already as (F, H, W, C) uint8, ready to mux.
-        fps = float(mm.get("fps", 25))
-        sample_rate = int(mm.get("audio_sample_rate", 24000))
-        frames = generated_video
-    else:
-        fps = float(mm.get("fps", 24))
-        sample_rate = int(mm.get("audio_sample_rate", 16000))
-        # DreamID-Omni returns video as (C, F, H, W) float32 in [-1, 1].
-        # mux_video_audio_bytes expects (F, H, W, C) uint8.
-        if not isinstance(generated_video, np.ndarray) or generated_video.ndim != 4:
-            raise RuntimeError(f"Unexpected video shape: {getattr(generated_video, 'shape', None)}")
-        frames = generated_video.transpose(1, 2, 3, 0)
-        frames = (np.clip((frames + 1.0) / 2.0, 0.0, 1.0) * 255.0).round().astype(np.uint8)
+        sampling_params = OmniDiffusionSamplingParams(
+            height=args.height,
+            width=args.width,
+            num_inference_steps=args.num_inference_steps,
+            seed=args.seed,
+            extra_args={},
+        )
+        declared_extra_body_params = get_extra_body_params(model_class_name)
+        extra_body = dict(args.extra_body or {})
+        legacy_extra_body = {"solver_name": args.solver_name, "shift": args.shift}
+        for key, value in legacy_extra_body.items():
+            if key in declared_extra_body_params:
+                extra_body.setdefault(key, value)
+        if declared_extra_body_params:
+            apply_declared_extra_args(sampling_params, declared_extra_body_params, extra_body)
+        elif extra_body:
+            sampling_params.extra_args.update({key: value for key, value in extra_body.items() if value is not None})
 
-    audio_np = None
-    if generated_audio is not None:
-        audio_np = np.squeeze(np.asarray(generated_audio)).astype(np.float32)
+        start = time.perf_counter()
+        outputs = omni.generate(prompt, sampling_params)
+        elapsed = time.perf_counter() - start
+        output = extract_x_to_video_audio_output(
+            outputs,
+            fps=args.fps,
+            audio_sample_rate=args.audio_sample_rate,
+        )
 
-    output_path = args.output
-    video_bytes = mux_video_audio_bytes(
-        frames,
-        audio_np,
-        fps=float(fps),
-        audio_sample_rate=sample_rate,
-    )
-    with open(output_path, "wb") as f:
-        f.write(video_bytes)
-    print(f"Saved generated video to {output_path}")
-    print(f"Total time: {elapsed:.2f}s")
-    peak_mb = _extract_peak_memory_mb(outputs)
-    if peak_mb:
-        print(f"Worker peak GPU memory (reserved): {peak_mb:.2f} MiB ({peak_mb / 1024:.2f} GiB)")
+        output_path = Path(args.output)
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+        output_path.write_bytes(encode_x_to_video_audio_output(output))
+        print(f"Saved generated video to {output_path}")
+        print(f"Total time: {elapsed:.2f}s")
+        peak_mb = _extract_peak_memory_mb(outputs)
+        if peak_mb:
+            print(f"Worker peak GPU memory (reserved): {peak_mb:.2f} MiB ({peak_mb / 1024:.2f} GiB)")
+    finally:
+        omni.close()
 
 
 if __name__ == "__main__":

tests/examples/offline_inference/test_x_to_video_audio.py

@@ -0,0 +1,133 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from pathlib import Path
+from types import SimpleNamespace
+from typing import Any
+
+import numpy as np
+import pytest
+from PIL import Image
+
+from examples.offline_inference.x_to_video_audio import x_to_video_audio
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_builds_canonical_image_audio_to_video_prompt() -> None:
+    image = Image.new("RGB", (16, 8), "red")
+    audio = (np.arange(8, dtype=np.float32), 24000)
+
+    prompt = x_to_video_audio.build_x_to_video_audio_prompt(
+        "A person speaks to the camera.",
+        {"image": [image], "audio": [audio]},
+    )
+
+    assert prompt == {
+        "prompt": "A person speaks to the camera.",
+        "modalities": ["video"],
+        "multi_modal_data": {"image": [image], "audio": [audio]},
+    }
+
+
+def test_canonical_prompt_omits_empty_media() -> None:
+    assert x_to_video_audio.build_x_to_video_audio_prompt("A landscape.") == {
+        "prompt": "A landscape.",
+        "modalities": ["video"],
+    }
+
+
+def test_media_loader_preserves_complete_audio(
+    monkeypatch: pytest.MonkeyPatch,
+    tmp_path: Path,
+) -> None:
+    image_path = tmp_path / "reference.png"
+    Image.new("RGBA", (4, 3), "blue").save(image_path)
+    complete_waveform = np.arange(80000, dtype=np.float32)
+    requested_rates: list[int | None] = []
+
+    def fake_load_audio(path: str, *, sr: int | None) -> tuple[np.ndarray, int]:
+        assert path == "reference.wav"
+        requested_rates.append(sr)
+        return complete_waveform, 16000
+
+    monkeypatch.setattr(x_to_video_audio, "load_audio", fake_load_audio)
+
+    images, audios = x_to_video_audio.load_image_and_audio(
+        [str(image_path)],
+        ["reference.wav"],
+        audio_sample_rate=16000,
+    )
+
+    assert images[0].mode == "RGB"
+    assert requested_rates == [16000]
+    assert audios[0][0] is complete_waveform
+    assert audios[0][0].shape == (80000,)
+    assert audios[0][1] == 16000
+
+
+@pytest.mark.parametrize(
+    "video",
+    [
+        np.zeros((3, 2, 4, 4), dtype=np.float32),
+        np.zeros((2, 4, 4, 3), dtype=np.uint8),
+    ],
+    ids=["channel_first_float", "frame_first_uint8"],
+)
+def test_output_contract_keeps_model_video_layout_opaque(video: np.ndarray) -> None:
+    audio = np.zeros(16, dtype=np.float32)
+    result = SimpleNamespace(
+        images=[video],
+        multimodal_output={"audio": audio, "fps": 25, "audio_sample_rate": 44100},
+    )
+
+    output = x_to_video_audio.extract_x_to_video_audio_output([result])
+
+    assert output.video is video
+    assert output.audio is audio
+    assert output.fps == 25.0
+    assert output.audio_sample_rate == 44100
+
+
+def test_output_contract_accepts_video_from_multimodal_output() -> None:
+    video = object()
+    result = SimpleNamespace(images=[], multimodal_output={"video": video, "fps": 12})
+
+    output = x_to_video_audio.extract_x_to_video_audio_output(result)
+
+    assert output == x_to_video_audio.XToVideoAudioOutput(
+        video=video,
+        audio=None,
+        fps=12.0,
+        audio_sample_rate=None,
+    )
+
+
+def test_output_contract_requires_model_neutral_metadata() -> None:
+    result = SimpleNamespace(images=[object()], multimodal_output={"audio": np.zeros(4)})
+
+    with pytest.raises(RuntimeError, match="must declare a valid 'fps'"):
+        x_to_video_audio.extract_x_to_video_audio_output(result)
+
+    with pytest.raises(RuntimeError, match="must declare a valid 'audio_sample_rate'"):
+        x_to_video_audio.extract_x_to_video_audio_output(result, fps=24)
+
+
+def test_output_encoder_receives_opaque_payload(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    video = object()
+    audio = object()
+    output = x_to_video_audio.XToVideoAudioOutput(video, audio, 23.976, 48000)
+    received: list[tuple[Any, ...]] = []
+
+    def fake_encode(*args: Any) -> bytes:
+        received.append(args)
+        return b"encoded"
+
+    monkeypatch.setattr(x_to_video_audio, "encode_video_bytes", fake_encode)
+
+    assert x_to_video_audio.encode_x_to_video_audio_output(output) == b"encoded"
+    assert received == [(video, 23.976, audio, 48000)]

tests/model_extras/test_model_extras.py

@@ -15,8 +15,10 @@ from vllm_omni.model_extras import (
     build_image_to_video_prompt,
     build_text_to_image_prompt,
     build_x_to_text_prompt,
+    build_x_to_video_audio_prompt,
     get_extra_body_params,
     get_extra_output_params,
+    get_input_audio_sample_rate,
     get_output_tensor_range,
     get_x_to_text_model_family,
     should_init_extra_args_for_non_diffusion_stages,
@@ -141,6 +143,53 @@ def test_lingbot_extra_registry_declares_request_params() -> None:
     assert should_init_extra_args_for_non_diffusion_stages("LingBotVideoPipeline") is False
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_dreamid_x_to_video_audio_adapter() -> None:
+    image = Image.new("RGB", (16, 16))
+    waveform = object()
+    canonical_prompt = {
+        "prompt": (
+            "[SPEAKER_TIMESTAMPS_START]metadata[SPEAKER_TIMESTAMPS_END]\n\n"
+            "A person waves.\n[AUDIO_DESCRIPTION_START]noise[AUDIO_DESCRIPTION_END]"
+        ),
+        "modalities": ["video"],
+        "multi_modal_data": {"image": [image], "audio": [(waveform, 16000)]},
+    }
+
+    result = build_x_to_video_audio_prompt(
+        "DreamIDOmniPipeline",
+        canonical_prompt,
+        {
+            "video_negative_prompt": "blur",
+            "audio_negative_prompt": "echo",
+        },
+    )
+
+    assert result == {
+        "prompt": "A person waves.",
+        "modalities": ["video"],
+        "video_negative_prompt": "blur",
+        "audio_negative_prompt": "echo",
+        "multi_modal_data": {"image": [image], "audio": [waveform]},
+    }
+    assert get_extra_body_params("DreamIDOmniPipeline") == frozenset({"solver_name", "shift"})
+    assert get_input_audio_sample_rate("DreamIDOmniPipeline") == 16000
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_generic_x_to_video_audio_adapter_is_model_neutral() -> None:
+    canonical_prompt = {
+        "prompt": "A person waves.",
+        "modalities": ["video"],
+        "multi_modal_data": {"audio": [(object(), 48000)]},
+    }
+
+    assert build_x_to_video_audio_prompt("FuturePipeline", canonical_prompt, {}) is canonical_prompt
+    assert get_input_audio_sample_rate("FuturePipeline") is None
+
+
 @pytest.mark.core_model
 @pytest.mark.cpu
 @pytest.mark.parametrize("pipeline_name", ["Cosmos3OmniDiffusersPipeline", "Cosmos3OmniPipeline"])

vllm_omni/entrypoints/openai/video_api_utils.py

@@ -509,9 +509,9 @@ def _coerce_video_to_uint8_frames(video: Any) -> np.ndarray:
     return frames_u8
 
 
-def _encode_video_bytes(
+def encode_video_bytes(
     video: Any,
-    fps: int,
+    fps: int | float,
     audio: Any | None = None,
     audio_sample_rate: int | None = None,
     video_codec_options: dict[str, str] | None = None,
@@ -530,6 +530,10 @@ def _encode_video_bytes(
     )
 
 
+# Backward-compatible private alias for existing internal callers.
+_encode_video_bytes = encode_video_bytes
+
+
 class FragmentedMP4VideoEncoder:
     """Normalize video chunks and append them to one fragmented MP4 stream."""
 

vllm_omni/model_extras/__init__.py

@@ -6,8 +6,10 @@ from vllm_omni.model_extras.registry import (
     build_image_to_video_prompt,
     build_text_to_image_prompt,
     build_x_to_text_prompt,
+    build_x_to_video_audio_prompt,
     get_extra_body_params,
     get_extra_output_params,
+    get_input_audio_sample_rate,
     get_model_class_name,
     get_output_tensor_range,
     get_transformer_config_subfolder,
@@ -21,8 +23,10 @@ __all__ = [
     "build_image_to_video_prompt",
     "build_text_to_image_prompt",
     "build_x_to_text_prompt",
+    "build_x_to_video_audio_prompt",
     "get_extra_body_params",
     "get_extra_output_params",
+    "get_input_audio_sample_rate",
     "get_model_class_name",
     "get_output_tensor_range",
     "get_transformer_config_subfolder",

vllm_omni/model_extras/dreamid_omni.py

@@ -0,0 +1,64 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import re
+from collections.abc import Mapping
+from typing import Any
+
+DREAMID_OMNI_EXTRA_BODY_PARAMS = frozenset({"solver_name", "shift"})
+DREAMID_OMNI_INPUT_AUDIO_SAMPLE_RATE = 16000
+
+
+def _clean_official_prompt(prompt: str) -> str:
+    """Remove metadata tags used by DreamID's official prompt fixtures."""
+    prompt = re.sub(
+        r"\[SPEAKER_TIMESTAMPS_START\].*?\[SPEAKER_TIMESTAMPS_END\]",
+        "",
+        prompt,
+        flags=re.DOTALL,
+    ).strip()
+    prompt = re.sub(
+        r"\[AUDIO_DESCRIPTION_START].*?\[AUDIO_DESCRIPTION_END]",
+        "",
+        prompt,
+        flags=re.DOTALL,
+    ).strip()
+    prompt = re.sub(r"\[[A-Z_]+\]", "", prompt)
+    return re.sub(r"\n\s*\n", "\n", prompt).strip()
+
+
+def build_x_to_video_audio_prompt(
+    prompt: dict[str, Any],
+    request_options: Mapping[str, Any],
+) -> dict[str, Any]:
+    """Translate the canonical X-to-video+audio envelope for DreamID-Omni."""
+    result = dict(prompt)
+    result["prompt"] = _clean_official_prompt(str(prompt["prompt"]))
+
+    for key in ("video_negative_prompt", "audio_negative_prompt"):
+        if request_options.get(key) is not None:
+            result[key] = request_options[key]
+
+    media_inputs = prompt.get("multi_modal_data")
+    if isinstance(media_inputs, Mapping) and "audio" in media_inputs:
+        adapted_media = dict(media_inputs)
+        audio_inputs = media_inputs["audio"]
+        if not isinstance(audio_inputs, list):
+            audio_inputs = [audio_inputs]
+        adapted_media["audio"] = [_extract_waveform(item) for item in audio_inputs]
+        result["multi_modal_data"] = adapted_media
+    return result
+
+
+def _extract_waveform(audio: Any) -> Any:
+    if not isinstance(audio, tuple) or len(audio) != 2:
+        raise TypeError("Canonical audio inputs must be (waveform, sample_rate) tuples.")
+    waveform, sample_rate = audio
+    if sample_rate != DREAMID_OMNI_INPUT_AUDIO_SAMPLE_RATE:
+        raise ValueError(
+            "DreamID-Omni audio inputs must be resampled to "
+            f"{DREAMID_OMNI_INPUT_AUDIO_SAMPLE_RATE} Hz; got {sample_rate}."
+        )
+    return waveform

vllm_omni/model_extras/registry.py

@@ -28,6 +28,13 @@ from vllm_omni.model_extras.cosmos3 import (
     COSMOS3_EXTRA_BODY_PARAMS,
     COSMOS3_EXTRA_OUTPUT_PARAMS,
 )
+from vllm_omni.model_extras.dreamid_omni import (
+    DREAMID_OMNI_EXTRA_BODY_PARAMS,
+    DREAMID_OMNI_INPUT_AUDIO_SAMPLE_RATE,
+)
+from vllm_omni.model_extras.dreamid_omni import (
+    build_x_to_video_audio_prompt as build_dreamid_x_to_video_audio_prompt,
+)
 from vllm_omni.model_extras.helios import (
     HELIOS_EXTRA_BODY_PARAMS,
     HELIOS_EXTRA_OUTPUT_PARAMS,
@@ -98,6 +105,7 @@ ImageToVideoPromptBuilder = Callable[
     dict[str, Any],
 ]
 XToTextPromptBuilder = Callable[[str, str, bool], tuple[dict[str, Any], list[int] | None]]
+XToVideoAudioPromptBuilder = Callable[[dict[str, Any], Mapping[str, Any]], dict[str, Any]]
 OutputTensorRange = Literal["negative_one_to_one", "zero_to_one"]
 
 
@@ -203,6 +211,11 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = {
         "extra_body_params": COSMOS3_EXTRA_BODY_PARAMS,
         "extra_output_params": COSMOS3_EXTRA_OUTPUT_PARAMS,
     },
+    "DreamIDOmniPipeline": {
+        "extra_body_params": DREAMID_OMNI_EXTRA_BODY_PARAMS,
+        "input_audio_sample_rate": DREAMID_OMNI_INPUT_AUDIO_SAMPLE_RATE,
+        "x_to_video_audio_prompt_builder": build_dreamid_x_to_video_audio_prompt,
+    },
     "MagiHumanPipeline": {
         "extra_body_params": MAGI_HUMAN_EXTRA_BODY_PARAMS,
         "extra_output_params": MAGI_HUMAN_EXTRA_OUTPUT_PARAMS,
@@ -295,6 +308,12 @@ def get_extra_output_params(model_class_name: str | None) -> frozenset[str]:
     return spec.get("extra_output_params", frozenset()) if spec is not None else frozenset()
 
 
+def get_input_audio_sample_rate(model_class_name: str | None) -> int | None:
+    """Return a pipeline's required input rate, or preserve the source rate."""
+    spec = _get_spec(model_class_name)
+    return spec.get("input_audio_sample_rate") if spec is not None else None
+
+
 def get_output_tensor_range(model_class_name: str | None) -> OutputTensorRange:
     """Return the declared range for floating-point tensor outputs.
 
@@ -401,3 +420,14 @@ def build_image_to_video_prompt(
         width=width,
         num_frames=num_frames,
     )
+
+
+def build_x_to_video_audio_prompt(
+    model_class_name: str | None,
+    prompt: dict[str, Any],
+    request_options: Mapping[str, Any],
+) -> dict[str, Any]:
+    """Apply a model adapter to an example-owned video+audio envelope."""
+    spec = _get_spec(model_class_name)
+    builder: XToVideoAudioPromptBuilder | None = spec.get("x_to_video_audio_prompt_builder") if spec else None
+    return builder(prompt, request_options) if builder is not None else prompt