[Bugfix] Validate MammothModa2 text-to-image dimensions (#7482)

Signed-off-by: mudrobot <mudrobot@foxmail.com>

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},
+        )

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