Compare commits

...
Author SHA1 Message Date
Hongsheng Liu f1bc6c8569 Address video-audio example review feedback
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-17 16:45:55 +08:00
Hongsheng Liu 6bdbaeb038 Harden the shared video-audio example
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-17 13:48:27 +08:00
9 changed files with 567 additions and 108 deletions

No files matched your search

@@ -3,7 +3,26 @@
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/x_to_video_audio>. 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 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
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 ## Local CLI Usage
### Download the Model locally ### Download the Model locally
@@ -91,7 +110,8 @@ python x_to_video_audio.py \
Key arguments: Key arguments:
- `--prompt`: text description (string). - `--prompt`: text description (string).
- `--model`: path to the model local directory. - `--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. - `--image-path`: path to the input image list.
- `--audio-path`: path to the input audio list, indicate the timbre of the output video. - `--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). - `--cfg-parallel-size`: number of parallel cfg parallel (defaults 1).
@@ -101,6 +121,9 @@ Key arguments:
- `--num-inference-steps`: number of denoising steps (defaults 45). - `--num-inference-steps`: number of denoising steps (defaults 45).
- `--video-negative-prompt`: negative prompt for video generation. - `--video-negative-prompt`: negative prompt for video generation.
- `--audio-negative-prompt`: negative prompt for audio 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). - `--enable-cpu-offload`: enable CPU offload (defaults False).
## Example materials ## Example materials
@@ -1,6 +1,25 @@
# X-To-Video-Audio # 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 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
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 ## Local CLI Usage
### Download the Model locally ### Download the Model locally
@@ -147,7 +166,8 @@ in `vllm_omni/model_extras/magi_human.py`, routed via `extra_args`):
Key arguments: Key arguments:
- `--prompt`: text description (string). - `--prompt`: text description (string).
- `--model`: path to the model local directory. - `--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. - `--image-path`: path to the input image list.
- `--audio-path`: path to the input audio list, indicate the timbre of the output video. - `--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). - `--cfg-parallel-size`: number of parallel cfg parallel (defaults 1).
@@ -157,6 +177,9 @@ Key arguments:
- `--num-inference-steps`: number of denoising steps (defaults 45). - `--num-inference-steps`: number of denoising steps (defaults 45).
- `--video-negative-prompt`: negative prompt for video generation. - `--video-negative-prompt`: negative prompt for video generation.
- `--audio-negative-prompt`: negative prompt for audio 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). - `--enable-cpu-offload`: enable CPU offload (defaults False).
- `--cache-backend`: enable `cache_dit` for acceleration. - `--cache-backend`: enable `cache_dit` for acceleration.
- `--quantization`: online (dynamic) quantization method — `fp8` or `int8`. - `--quantization`: online (dynamic) quantization method — `fp8` or `int8`.
@@ -3,32 +3,97 @@
import argparse import argparse
import json import json
import math
import re import re
import time import time
from collections.abc import Mapping
from dataclasses import dataclass
from pathlib import Path
from typing import Any from typing import Any
import numpy as np import numpy as np
import torch
from PIL import Image from PIL import Image
from vllm.multimodal.media.audio import load_audio from vllm.multimodal.media.audio import load_audio
from vllm_omni.diffusion.data import DiffusionParallelConfig 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.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.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,
get_output_tensor_range,
)
IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".gif"} IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".gif"}
AUDIO_EXTENSIONS = {".wav", ".mp3", ".flac", ".m4a", ".aac", ".ogg"} 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
output_tensor_range: str = "negative_one_to_one"
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 _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: def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Offline inference for X -> video+audio models.") 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", required=True, help="Model ckpt root directory.")
parser.add_argument( parser.add_argument(
"--model-type", "--model-type",
default="dreamid-omni", default=None,
choices=["dreamid-omni", "magi-human"], help="Optional legacy model type override (for example, magi-human).",
help="Model type.",
) )
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("--prompt", default=None, help="Text prompt.")
parser.add_argument("--image-path", type=str, nargs="+", help="list of image-path") parser.add_argument("--image-path", type=str, nargs="+", help="list of image-path")
@@ -36,9 +101,9 @@ def parse_args() -> argparse.Namespace:
parser.add_argument("--prompt-file", type=str, default=None, help="Text prompt in json format.") parser.add_argument("--prompt-file", type=str, default=None, help="Text prompt in json format.")
parser.add_argument( parser.add_argument(
"--extra-body", "--extra-body",
type=str, type=parse_json_object,
default=None, 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}\'.', 'merged into sampling extra_args. Example: \'{"image_path": "/path/to/img.jpg", "seconds": 5}\'.',
) )
@@ -71,13 +136,20 @@ def parse_args() -> argparse.Namespace:
default="robotic, muffled, echo, distorted", default="robotic, muffled, echo, distorted",
help="Negative prompt for audio.", 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("--output", default="dreamid_output.mp4", help="Output video path.")
parser.add_argument( parser.add_argument(
"--quantization", "--quantization",
type=str, type=str,
default=None, default=None,
choices=["fp8", "int8"], 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( parser.add_argument(
"--enable-cpu-offload", "--enable-cpu-offload",
@@ -94,7 +166,7 @@ def parse_args() -> argparse.Namespace:
parser.add_argument( parser.add_argument(
"--use-hsdp", "--use-hsdp",
action="store_true", action="store_true",
help="Enable HSDP for DreamID-Omni fused transformer blocks.", help="Enable HSDP for supported transformer blocks.",
) )
parser.add_argument( parser.add_argument(
"--hsdp-shard-size", "--hsdp-shard-size",
@@ -111,20 +183,105 @@ def parse_args() -> argparse.Namespace:
return parser.parse_args() return parser.parse_args()
def load_image_and_audio(image_paths, audio_paths): def load_image_and_audio(
image = [] image_paths: list[str] | None,
audio = [] 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: with Image.open(path) as img:
img = img.convert("RGB") images.append(img.convert("RGB"))
image.append(img)
for path in audio_paths: for path in audio_paths or []:
audio_array, sr = load_audio(path, sr=16000) waveform, sample_rate = load_audio(path, sr=audio_sample_rate)
audio_array = audio_array[int(sr * 1) : int(sr * 3)] audios.append((waveform, sample_rate))
audio.append(audio_array) return images, audios
return image, audio
def extract_x_to_video_audio_output(
outputs: Any,
*,
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):
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, 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(
_normalize_output_tensor_range(output.video, output.output_tensor_range),
output.fps,
output.audio,
output.audio_sample_rate,
)
def _extract_peak_memory_mb(result: Any) -> float: def _extract_peak_memory_mb(result: Any) -> float:
@@ -157,45 +314,10 @@ def main() -> None:
if args.prompt_file: if args.prompt_file:
with open(args.prompt_file) as f: with open(args.prompt_file) as f:
text_prompt = json.load(f) text_prompt = json.load(f)
text_prompt = re.sub( if isinstance(text_prompt, str):
r"\[SPEAKER_TIMESTAMPS_START\].*?\[SPEAKER_TIMESTAMPS_END\]", "", text_prompt, flags=re.DOTALL text_prompt = _clean_official_prompt(text_prompt)
).strip() if not isinstance(text_prompt, str):
text_prompt = re.sub( raise ValueError("Prompt content must be a JSON string or --prompt text.")
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,
},
)
parallel_config = DiffusionParallelConfig( parallel_config = DiffusionParallelConfig(
cfg_parallel_size=args.cfg_parallel_size, cfg_parallel_size=args.cfg_parallel_size,
@@ -220,64 +342,83 @@ def main() -> None:
"scm_steps_policy": "dynamic", "scm_steps_policy": "dynamic",
} }
omni_kwargs = dict( omni_kwargs: dict[str, Any] = dict(
model=args.model, model=args.model,
parallel_config=parallel_config, parallel_config=parallel_config,
model_type=args.model_type,
enable_cpu_offload=args.enable_cpu_offload, enable_cpu_offload=args.enable_cpu_offload,
enable_layerwise_offload=args.enable_layerwise_offload, enable_layerwise_offload=args.enable_layerwise_offload,
cache_backend=args.cache_backend, cache_backend=args.cache_backend,
cache_config=cache_config, 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: if args.quantization is not None:
omni_kwargs["quantization"] = args.quantization omni_kwargs["quantization"] = args.quantization
omni = Omni(**omni_kwargs) omni = Omni(**omni_kwargs)
start = time.perf_counter() try:
outputs = omni.generate(prompt, sampling_params) model_class_name = args.model_class_name or get_model_class_name(omni)
elapsed = time.perf_counter() - start 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: canonical_prompt = build_x_to_video_audio_prompt(text_prompt, media_inputs)
raise RuntimeError("No output returned from the model.") prompt = build_model_x_to_video_audio_prompt(
result = outputs[0] model_class_name,
if not result.images: canonical_prompt,
raise RuntimeError("No video frames found in model output.") {
generated_video = result.images[0] "video_negative_prompt": args.video_negative_prompt,
mm = result.multimodal_output or {} "audio_negative_prompt": args.audio_negative_prompt,
generated_audio = mm.get("audio") },
)
if args.model_type == "magi-human": sampling_params = OmniDiffusionSamplingParams(
# MagiHuman returns frames already as (F, H, W, C) uint8, ready to mux. height=args.height,
fps = float(mm.get("fps", 25)) width=args.width,
sample_rate = int(mm.get("audio_sample_rate", 24000)) num_inference_steps=args.num_inference_steps,
frames = generated_video seed=args.seed,
else: extra_args={},
fps = float(mm.get("fps", 24)) )
sample_rate = int(mm.get("audio_sample_rate", 16000)) declared_extra_body_params = get_extra_body_params(model_class_name)
# DreamID-Omni returns video as (C, F, H, W) float32 in [-1, 1]. extra_body = dict(args.extra_body or {})
# mux_video_audio_bytes expects (F, H, W, C) uint8. legacy_extra_body = {"solver_name": args.solver_name, "shift": args.shift}
if not isinstance(generated_video, np.ndarray) or generated_video.ndim != 4: for key, value in legacy_extra_body.items():
raise RuntimeError(f"Unexpected video shape: {getattr(generated_video, 'shape', None)}") if key in declared_extra_body_params:
frames = generated_video.transpose(1, 2, 3, 0) extra_body.setdefault(key, value)
frames = (np.clip((frames + 1.0) / 2.0, 0.0, 1.0) * 255.0).round().astype(np.uint8) 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 start = time.perf_counter()
if generated_audio is not None: outputs = omni.generate(prompt, sampling_params)
audio_np = np.squeeze(np.asarray(generated_audio)).astype(np.float32) elapsed = time.perf_counter() - start
output = extract_x_to_video_audio_output(
outputs,
fps=args.fps,
audio_sample_rate=args.audio_sample_rate,
output_tensor_range=get_output_tensor_range(model_class_name),
)
output_path = args.output output_path = Path(args.output)
video_bytes = mux_video_audio_bytes( output_path.parent.mkdir(parents=True, exist_ok=True)
frames, output_path.write_bytes(encode_x_to_video_audio_output(output))
audio_np, print(f"Saved generated video to {output_path}")
fps=float(fps), print(f"Total time: {elapsed:.2f}s")
audio_sample_rate=sample_rate, peak_mb = _extract_peak_memory_mb(outputs)
) if peak_mb:
with open(output_path, "wb") as f: print(f"Worker peak GPU memory (reserved): {peak_mb:.2f} MiB ({peak_mb / 1024:.2f} GiB)")
f.write(video_bytes) finally:
print(f"Saved generated video to {output_path}") omni.close()
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)")
if __name__ == "__main__": if __name__ == "__main__":
@@ -0,0 +1,144 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from __future__ import annotations
import io
from pathlib import Path
from types import SimpleNamespace
import av
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_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,
) -> 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_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",
)
payload = x_to_video_audio.encode_x_to_video_audio_output(output)
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)
+46
View File
@@ -15,8 +15,10 @@ from vllm_omni.model_extras import (
build_image_to_video_prompt, build_image_to_video_prompt,
build_text_to_image_prompt, build_text_to_image_prompt,
build_x_to_text_prompt, build_x_to_text_prompt,
build_x_to_video_audio_prompt,
get_extra_body_params, get_extra_body_params,
get_extra_output_params, get_extra_output_params,
get_input_audio_sample_rate,
get_output_tensor_range, get_output_tensor_range,
get_x_to_text_model_family, get_x_to_text_model_family,
should_init_extra_args_for_non_diffusion_stages, should_init_extra_args_for_non_diffusion_stages,
@@ -141,6 +143,50 @@ def test_lingbot_extra_registry_declares_request_params() -> None:
assert should_init_extra_args_for_non_diffusion_stages("LingBotVideoPipeline") is False 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": "A person waves.",
"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.core_model
@pytest.mark.cpu @pytest.mark.cpu
@pytest.mark.parametrize("pipeline_name", ["Cosmos3OmniDiffusersPipeline", "Cosmos3OmniPipeline"]) @pytest.mark.parametrize("pipeline_name", ["Cosmos3OmniDiffusersPipeline", "Cosmos3OmniPipeline"])
@@ -509,9 +509,9 @@ def _coerce_video_to_uint8_frames(video: Any) -> np.ndarray:
return frames_u8 return frames_u8
def _encode_video_bytes( def encode_video_bytes(
video: Any, video: Any,
fps: int, fps: int | float,
audio: Any | None = None, audio: Any | None = None,
audio_sample_rate: int | None = None, audio_sample_rate: int | None = None,
video_codec_options: dict[str, str] | 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: class FragmentedMP4VideoEncoder:
"""Normalize video chunks and append them to one fragmented MP4 stream.""" """Normalize video chunks and append them to one fragmented MP4 stream."""
+4
View File
@@ -6,8 +6,10 @@ from vllm_omni.model_extras.registry import (
build_image_to_video_prompt, build_image_to_video_prompt,
build_text_to_image_prompt, build_text_to_image_prompt,
build_x_to_text_prompt, build_x_to_text_prompt,
build_x_to_video_audio_prompt,
get_extra_body_params, get_extra_body_params,
get_extra_output_params, get_extra_output_params,
get_input_audio_sample_rate,
get_model_class_name, get_model_class_name,
get_output_tensor_range, get_output_tensor_range,
get_transformer_config_subfolder, get_transformer_config_subfolder,
@@ -21,8 +23,10 @@ __all__ = [
"build_image_to_video_prompt", "build_image_to_video_prompt",
"build_text_to_image_prompt", "build_text_to_image_prompt",
"build_x_to_text_prompt", "build_x_to_text_prompt",
"build_x_to_video_audio_prompt",
"get_extra_body_params", "get_extra_body_params",
"get_extra_output_params", "get_extra_output_params",
"get_input_audio_sample_rate",
"get_model_class_name", "get_model_class_name",
"get_output_tensor_range", "get_output_tensor_range",
"get_transformer_config_subfolder", "get_transformer_config_subfolder",
+44
View File
@@ -0,0 +1,44 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from __future__ import annotations
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 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)
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
+30
View File
@@ -28,6 +28,13 @@ from vllm_omni.model_extras.cosmos3 import (
COSMOS3_EXTRA_BODY_PARAMS, COSMOS3_EXTRA_BODY_PARAMS,
COSMOS3_EXTRA_OUTPUT_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 ( from vllm_omni.model_extras.helios import (
HELIOS_EXTRA_BODY_PARAMS, HELIOS_EXTRA_BODY_PARAMS,
HELIOS_EXTRA_OUTPUT_PARAMS, HELIOS_EXTRA_OUTPUT_PARAMS,
@@ -98,6 +105,7 @@ ImageToVideoPromptBuilder = Callable[
dict[str, Any], dict[str, Any],
] ]
XToTextPromptBuilder = Callable[[str, str, bool], tuple[dict[str, Any], list[int] | None]] 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"] 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_body_params": COSMOS3_EXTRA_BODY_PARAMS,
"extra_output_params": COSMOS3_EXTRA_OUTPUT_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": { "MagiHumanPipeline": {
"extra_body_params": MAGI_HUMAN_EXTRA_BODY_PARAMS, "extra_body_params": MAGI_HUMAN_EXTRA_BODY_PARAMS,
"extra_output_params": MAGI_HUMAN_EXTRA_OUTPUT_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() 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: def get_output_tensor_range(model_class_name: str | None) -> OutputTensorRange:
"""Return the declared range for floating-point tensor outputs. """Return the declared range for floating-point tensor outputs.
@@ -401,3 +420,14 @@ def build_image_to_video_prompt(
width=width, width=width,
num_frames=num_frames, 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