diff --git a/docs/user_guide/examples/offline_inference/x_to_video_audio.md b/docs/user_guide/examples/offline_inference/x_to_video_audio.md index d5cd0ea46..122e10363 100644 --- a/docs/user_guide/examples/offline_inference/x_to_video_audio.md +++ b/docs/user_guide/examples/offline_inference/x_to_video_audio.md @@ -3,7 +3,25 @@ Source . -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 diff --git a/examples/offline_inference/x_to_video_audio/x_to_video_audio.md b/examples/offline_inference/x_to_video_audio/x_to_video_audio.md index 8305c785f..b625956e0 100644 --- a/examples/offline_inference/x_to_video_audio/x_to_video_audio.md +++ b/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`. diff --git a/examples/offline_inference/x_to_video_audio/x_to_video_audio.py b/examples/offline_inference/x_to_video_audio/x_to_video_audio.py index 9512ac65c..2485b3fdf 100644 --- a/examples/offline_inference/x_to_video_audio/x_to_video_audio.py +++ b/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__": diff --git a/tests/examples/offline_inference/test_x_to_video_audio.py b/tests/examples/offline_inference/test_x_to_video_audio.py new file mode 100644 index 000000000..1b73dab1e --- /dev/null +++ b/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)] diff --git a/tests/model_extras/test_model_extras.py b/tests/model_extras/test_model_extras.py index 573f0bd79..1ae838681 100644 --- a/tests/model_extras/test_model_extras.py +++ b/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"]) diff --git a/vllm_omni/entrypoints/openai/video_api_utils.py b/vllm_omni/entrypoints/openai/video_api_utils.py index 5a4f9f826..5cba72d16 100644 --- a/vllm_omni/entrypoints/openai/video_api_utils.py +++ b/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.""" diff --git a/vllm_omni/model_extras/__init__.py b/vllm_omni/model_extras/__init__.py index f10ca9806..0c83d923d 100644 --- a/vllm_omni/model_extras/__init__.py +++ b/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", diff --git a/vllm_omni/model_extras/dreamid_omni.py b/vllm_omni/model_extras/dreamid_omni.py new file mode 100644 index 000000000..7a653aab5 --- /dev/null +++ b/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 diff --git a/vllm_omni/model_extras/registry.py b/vllm_omni/model_extras/registry.py index c394b8999..fbe36c7a2 100644 --- a/vllm_omni/model_extras/registry.py +++ b/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