diff --git a/examples/offline_inference/image_to_video/image_to_video.py b/examples/offline_inference/image_to_video/image_to_video.py index 1a7e40b12..65e7389f0 100644 --- a/examples/offline_inference/image_to_video/image_to_video.py +++ b/examples/offline_inference/image_to_video/image_to_video.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Image-to-Video generation example using Wan2.2 I2V/TI2V models, LTX2/LTX-2.3, @@ -111,6 +111,13 @@ def build_image_to_video_prompt( return result +def _validate_video_output_type(output_type: str) -> None: + if output_type not in {"image", "video"}: + raise ValueError( + f"Unexpected output type '{output_type}', expected 'video' or legacy 'image' for video generation." + ) + + def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description=( @@ -665,10 +672,7 @@ def main(): frames = frames[0] if frames else None if isinstance(frames, OmniRequestOutput): - if frames.final_output_type != "image": - raise ValueError( - f"Unexpected output type '{frames.final_output_type}', expected 'image' for video generation." - ) + _validate_video_output_type(frames.final_output_type) if frames.multimodal_output and "audio" in frames.multimodal_output: audio = frames.multimodal_output["audio"] audio_sample_rate = frames.multimodal_output.get("audio_sample_rate", audio_sample_rate) diff --git a/examples/offline_inference/text_to_video/text_to_video.py b/examples/offline_inference/text_to_video/text_to_video.py index ce46fba49..02383c6e4 100644 --- a/examples/offline_inference/text_to_video/text_to_video.py +++ b/examples/offline_inference/text_to_video/text_to_video.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import argparse import json @@ -174,6 +174,13 @@ def build_text_to_video_prompt(prompt: str, negative_prompt: str | None) -> dict return result +def _validate_video_output_type(output_type: str) -> None: + if output_type not in {"image", "video"}: + raise ValueError( + f"Unexpected output type '{output_type}', expected 'video' or legacy 'image' for video generation." + ) + + def _normalize_float_tensor(tensor: torch.Tensor, source_range: str) -> torch.Tensor: """Normalize decoded frames according to the pipeline's output contract.""" if not tensor.is_floating_point(): @@ -620,10 +627,7 @@ def main(): frames = frames[0] if frames else None if isinstance(frames, OmniRequestOutput): - if frames.final_output_type != "image": - raise ValueError( - f"Unexpected output type '{frames.final_output_type}', expected 'image' for video generation." - ) + _validate_video_output_type(frames.final_output_type) if frames.multimodal_output and "audio" in frames.multimodal_output: audio = frames.multimodal_output["audio"] audio_sample_rate = frames.multimodal_output.get("audio_sample_rate", audio_sample_rate) diff --git a/tests/examples/offline_inference/test_image_task_prompts.py b/tests/examples/offline_inference/test_image_task_prompts.py index 421082ed6..054756e33 100644 --- a/tests/examples/offline_inference/test_image_task_prompts.py +++ b/tests/examples/offline_inference/test_image_task_prompts.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations @@ -65,3 +65,38 @@ def test_image_to_video_builds_canonical_prompt() -> None: "multi_modal_data": {"image": image}, "negative_prompt": "flicker", } + + +@pytest.mark.parametrize( + ("relative_path", "module_name"), + [ + ("offline_inference/text_to_video/text_to_video.py", "text_to_video_output_example"), + ("offline_inference/image_to_video/image_to_video.py", "image_to_video_output_example"), + ], +) +@pytest.mark.parametrize("output_type", ["video", "image"]) +def test_video_examples_accept_semantic_and_legacy_output_types( + relative_path: str, + module_name: str, + output_type: str, +) -> None: + mod = _load_example_module(relative_path, module_name) + + mod._validate_video_output_type(output_type) + + +@pytest.mark.parametrize( + ("relative_path", "module_name"), + [ + ("offline_inference/text_to_video/text_to_video.py", "text_to_video_invalid_output_example"), + ("offline_inference/image_to_video/image_to_video.py", "image_to_video_invalid_output_example"), + ], +) +def test_video_examples_reject_unrelated_output_types( + relative_path: str, + module_name: str, +) -> None: + mod = _load_example_module(relative_path, module_name) + + with pytest.raises(ValueError, match="Unexpected output type 'audio'"): + mod._validate_video_output_type("audio")