73 lines
2.5 KiB
Python
73 lines
2.5 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
|
|
from vllm_omni.model_extras import build_text_to_image_prompt
|
|
|
|
pytestmark = [
|
|
pytest.mark.core_model,
|
|
pytest.mark.cpu,
|
|
pytest.mark.parametrize(
|
|
"model_class_name",
|
|
["MammothModa2DiTPipeline", "MammothModa2ForConditionalGeneration", "Mammothmoda2Model"],
|
|
),
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("height", "width", "expected_height", "expected_width"),
|
|
[
|
|
(None, None, 1024, 1024),
|
|
(None, 512, 1024, 512),
|
|
(512, None, 512, 1024),
|
|
(16, 16, 16, 16),
|
|
(16, 32, 16, 32),
|
|
(32, 16, 32, 16),
|
|
(256, 256, 256, 256),
|
|
(512, 768, 512, 768),
|
|
(768, 512, 768, 512),
|
|
],
|
|
)
|
|
def test_text_to_image_dimensions(
|
|
model_class_name: str, height: int | None, width: int | None, expected_height: int, expected_width: int
|
|
) -> None:
|
|
result = build_text_to_image_prompt(
|
|
model_class_name,
|
|
{"prompt": "a cat", "modalities": ["image"]},
|
|
height=height,
|
|
width=width,
|
|
)
|
|
|
|
info = result["additional_information"]
|
|
assert info["image_height"] == [expected_height]
|
|
assert info["image_width"] == [expected_width]
|
|
assert info["ar_height"] == [expected_height // 16]
|
|
assert info["ar_width"] == [expected_width // 16]
|
|
assert result["prompt"].endswith(f"<|image start|>{expected_width // 16}*{expected_height // 16}<|image token|>")
|
|
|
|
|
|
@pytest.mark.parametrize("dimension", ["height", "width"])
|
|
@pytest.mark.parametrize("value", [-16, 0, 1, 15, 17, 513])
|
|
def test_text_to_image_rejects_invalid_dimensions(model_class_name: str, dimension: str, value: int) -> None:
|
|
dimensions = {"height": 512, "width": 512}
|
|
dimensions[dimension] = value
|
|
|
|
with pytest.raises(ValueError, match=rf"MammothModa2 {dimension} must be a positive multiple of 16, got {value}"):
|
|
build_text_to_image_prompt(model_class_name, {"prompt": "a cat", "modalities": ["image"]}, **dimensions)
|
|
|
|
|
|
@pytest.mark.parametrize("dimension", ["height", "width"])
|
|
@pytest.mark.parametrize("value", [0, 513])
|
|
def test_text_to_image_rejects_invalid_dimension_when_other_is_omitted(
|
|
model_class_name: str, dimension: str, value: int
|
|
) -> None:
|
|
with pytest.raises(ValueError, match=rf"MammothModa2 {dimension} must be a positive multiple of 16, got {value}"):
|
|
build_text_to_image_prompt(
|
|
model_class_name,
|
|
{"prompt": "a cat", "modalities": ["image"]},
|
|
**{dimension: value},
|
|
)
|