diff --git a/tests/model_extras/test_mammothmoda2.py b/tests/model_extras/test_mammothmoda2.py new file mode 100644 index 000000000..c4797bc0e --- /dev/null +++ b/tests/model_extras/test_mammothmoda2.py @@ -0,0 +1,72 @@ +# 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}, + ) diff --git a/vllm_omni/model_extras/mammothmodal2_preview.py b/vllm_omni/model_extras/mammothmodal2_preview.py index 96f3fb939..23ad7e5d9 100644 --- a/vllm_omni/model_extras/mammothmodal2_preview.py +++ b/vllm_omni/model_extras/mammothmodal2_preview.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 @@ -37,9 +37,16 @@ def build_text_to_image_prompt( MammothModa2 t2i uses classifier-free guidance via ``text_guidance_scale`` and has no explicit negative-prompt path, so ``negative_prompt`` is accepted for signature compatibility but not injected. + + Image dimensions must be positive multiples of 16. Omitted dimensions + default to 1024. """ - h = height or 1024 - w = width or 1024 + h = 1024 if height is None else height + w = 1024 if width is None else width + for name, size in (("height", h), ("width", w)): + if size <= 0 or size % _PATCH_SIZE != 0: + raise ValueError(f"MammothModa2 {name} must be a positive multiple of {_PATCH_SIZE}, got {size}") + ar_height = h // _PATCH_SIZE ar_width = w // _PATCH_SIZE @@ -97,5 +104,5 @@ def build_x_to_text_prompt( ) -MAMMOTHMODA2_PREVIEW_EXTRA_OUTPUT_PARAMS = frozenset() +MAMMOTHMODA2_PREVIEW_EXTRA_OUTPUT_PARAMS: frozenset[str] = frozenset() MAMMOTHMODA2_PREVIEW_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES = True