Address video-audio example review feedback

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

docs/user_guide/examples/offline_inference/x_to_video_audio.md

@@ -12,9 +12,10 @@ another model-specific runner.
 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.
+full by design; the runner does not apply a time window. Prepare a trimmed input
+file explicitly when only part of a recording should condition the model.
+DreamID-Omni's audio conditioning sequence grows with the complete input
+duration, so long recordings increase denoising memory and step time.
 
 Model-specific request translation (for example, DreamID-Omni's two negative
 prompts and its 16 kHz input requirement) is declared in `model_extras`. Model

examples/offline_inference/x_to_video_audio/x_to_video_audio.md

@@ -9,9 +9,10 @@ through `vllm_omni.model_extras` without adding another model-specific runner.
 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.
+full by design; the runner does not apply a time window. Prepare a trimmed input
+file explicitly when only part of a recording should condition the model.
+DreamID-Omni's audio conditioning sequence grows with the complete input
+duration, so long recordings increase denoising memory and step time.
 
 Model-specific request translation (for example, DreamID-Omni's two negative
 prompts and its 16 kHz input requirement) is declared in `model_extras`. Model

examples/offline_inference/x_to_video_audio/x_to_video_audio.py

@@ -4,12 +4,15 @@
 import argparse
 import json
 import math
+import re
 import time
 from collections.abc import Mapping
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Any
 
+import numpy as np
+import torch
 from PIL import Image
 from vllm.multimodal.media.audio import load_audio
 
@@ -25,6 +28,7 @@ from vllm_omni.model_extras import (
     get_extra_body_params,
     get_input_audio_sample_rate,
     get_model_class_name,
+    get_output_tensor_range,
 )
 
 IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".gif"}
@@ -39,6 +43,7 @@ class XToVideoAudioOutput:
     audio: Any | None
     fps: float
     audio_sample_rate: int | None
+    output_tensor_range: str = "negative_one_to_one"
 
 
 def parse_json_object(value: str) -> dict[str, Any]:
@@ -62,6 +67,24 @@ def build_x_to_video_audio_prompt(
     return result
 
 
+def _clean_official_prompt(prompt: str) -> str:
+    """Remove metadata tags used by the official DreamID 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 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.")
@@ -185,6 +208,7 @@ def extract_x_to_video_audio_output(
     *,
     fps: float | None = None,
     audio_sample_rate: int | None = None,
+    output_tensor_range: str = "negative_one_to_one",
 ) -> XToVideoAudioOutput:
     """Normalize engine output without depending on a model's tensor layout."""
     if isinstance(outputs, list):
@@ -228,13 +252,32 @@ def extract_x_to_video_audio_output(
     elif resolved_sample_rate is not None:
         resolved_sample_rate = int(resolved_sample_rate)
 
-    return XToVideoAudioOutput(video, audio, resolved_fps, resolved_sample_rate)
+    return XToVideoAudioOutput(video, audio, resolved_fps, resolved_sample_rate, output_tensor_range)
+
+
+def _normalize_output_tensor_range(video: Any, source_range: str) -> Any:
+    """Normalize floating-point output according to the pipeline contract."""
+    if isinstance(video, torch.Tensor):
+        if not video.is_floating_point():
+            return video
+        video = video.detach().cpu().float().numpy()
+    if isinstance(video, np.ndarray):
+        if not np.issubdtype(video.dtype, np.floating):
+            return video
+        if source_range == "negative_one_to_one":
+            return np.clip(video, -1.0, 1.0) * 0.5 + 0.5
+        if source_range == "zero_to_one":
+            return np.clip(video, 0.0, 1.0)
+        raise ValueError(f"Unsupported floating-point tensor range: {source_range!r}")
+    if isinstance(video, list):
+        return [_normalize_output_tensor_range(frame, source_range) for frame in video]
+    return video
 
 
 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,
+        _normalize_output_tensor_range(output.video, output.output_tensor_range),
         output.fps,
         output.audio,
         output.audio_sample_rate,
@@ -271,6 +314,8 @@ def main() -> None:
     if args.prompt_file:
         with open(args.prompt_file) as f:
             text_prompt = json.load(f)
+            if isinstance(text_prompt, str):
+                text_prompt = _clean_official_prompt(text_prompt)
     if not isinstance(text_prompt, str):
         raise ValueError("Prompt content must be a JSON string or --prompt text.")
 
@@ -313,7 +358,7 @@ def main() -> None:
         omni_kwargs["quantization"] = args.quantization
     omni = Omni(**omni_kwargs)
     try:
-        model_class_name = get_model_class_name(omni) or args.model_class_name
+        model_class_name = args.model_class_name or get_model_class_name(omni)
         input_sample_rate = get_input_audio_sample_rate(model_class_name)
         images, audios = load_image_and_audio(
             args.image_path,
@@ -361,6 +406,7 @@ def main() -> None:
             outputs,
             fps=args.fps,
             audio_sample_rate=args.audio_sample_rate,
+            output_tensor_range=get_output_tensor_range(model_class_name),
         )
 
         output_path = Path(args.output)

tests/examples/offline_inference/test_x_to_video_audio.py

@@ -3,10 +3,11 @@
 
 from __future__ import annotations
 
+import io
 from pathlib import Path
 from types import SimpleNamespace
-from typing import Any
 
+import av
 import numpy as np
 import pytest
 from PIL import Image
@@ -39,6 +40,15 @@ def test_canonical_prompt_omits_empty_media() -> None:
     }
 
 
+def test_official_prompt_cleanup_stays_in_prompt_file_handling() -> None:
+    prompt = (
+        "[SPEAKER_TIMESTAMPS_START]metadata[SPEAKER_TIMESTAMPS_END]\n\n"
+        "A person waves.\n[AUDIO_DESCRIPTION_START]noise[AUDIO_DESCRIPTION_END]"
+    )
+
+    assert x_to_video_audio._clean_official_prompt(prompt) == "A person waves."
+
+
 def test_media_loader_preserves_complete_audio(
     monkeypatch: pytest.MonkeyPatch,
     tmp_path: Path,
@@ -115,19 +125,20 @@ def test_output_contract_requires_model_neutral_metadata() -> None:
         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 test_output_encoder_normalizes_declared_range_and_channel_first_layout() -> None:
+    video = np.zeros((3, 2, 4, 6), dtype=np.float32)
+    output = x_to_video_audio.XToVideoAudioOutput(
+        video=video,
+        audio=None,
+        fps=8,
+        audio_sample_rate=None,
+        output_tensor_range="negative_one_to_one",
+    )
 
-    def fake_encode(*args: Any) -> bytes:
-        received.append(args)
-        return b"encoded"
+    payload = x_to_video_audio.encode_x_to_video_audio_output(output)
 
-    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)]
+    with av.open(io.BytesIO(payload)) as container:
+        decoded = [frame.to_ndarray(format="rgb24") for frame in container.decode(video=0)]
+    assert len(decoded) == 2
+    assert decoded[0].shape == (4, 6, 3)
+    np.testing.assert_allclose(decoded, 128, atol=2)

tests/model_extras/test_model_extras.py

@@ -149,10 +149,7 @@ 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]"
-        ),
+        "prompt": "A person waves.",
         "modalities": ["video"],
         "multi_modal_data": {"image": [image], "audio": [(waveform, 16000)]},
     }

vllm_omni/model_extras/dreamid_omni.py

@@ -3,7 +3,6 @@
 
 from __future__ import annotations
 
-import re
 from collections.abc import Mapping
 from typing import Any
 
@@ -11,31 +10,12 @@ 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: