103 lines
3.1 KiB
Python
103 lines
3.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
from __future__ import annotations
|
|
|
|
import functools
|
|
import importlib.util
|
|
|
|
import pytest
|
|
from PIL import Image
|
|
|
|
from tests.examples.helpers import EXAMPLES
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
@functools.cache
|
|
def _load_example_module(relative_path: str, module_name: str):
|
|
path = EXAMPLES / relative_path
|
|
spec = importlib.util.spec_from_file_location(module_name, path)
|
|
assert spec is not None and spec.loader is not None
|
|
module = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("negative_prompt", "expected_negative_prompt"),
|
|
[(None, None), ("", ""), ("blurry", "blurry")],
|
|
)
|
|
def test_text_to_image_builds_canonical_prompt(
|
|
negative_prompt: str | None,
|
|
expected_negative_prompt: str | None,
|
|
) -> None:
|
|
mod = _load_example_module(
|
|
"offline_inference/text_to_image/text_to_image.py",
|
|
"text_to_image_example",
|
|
)
|
|
result = mod.build_text_to_image_prompt("a red fox", negative_prompt)
|
|
|
|
assert result["prompt"] == "a red fox"
|
|
assert result["modalities"] == ["image"]
|
|
if expected_negative_prompt is None:
|
|
assert "negative_prompt" not in result
|
|
else:
|
|
assert result["negative_prompt"] == expected_negative_prompt
|
|
|
|
|
|
def test_image_to_video_builds_canonical_prompt() -> None:
|
|
mod = _load_example_module(
|
|
"offline_inference/image_to_video/image_to_video.py",
|
|
"image_to_video_example",
|
|
)
|
|
image = Image.new("RGB", (32, 16), "red")
|
|
|
|
result = mod.build_image_to_video_prompt(
|
|
prompt="the fox turns toward the camera",
|
|
negative_prompt="flicker",
|
|
media_inputs={"image": image},
|
|
)
|
|
|
|
assert result == {
|
|
"prompt": "the fox turns toward the camera",
|
|
"modalities": ["video"],
|
|
"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")
|