[Fix] Restore MammothModa2 support

Signed-off-by: hsliuustc0106 <hsliuustc0106@users.noreply.github.com>

tests/model_extras/test_model_extras.py

@@ -31,6 +31,7 @@ from vllm_omni.model_extras import (
     [
         ({"model_type": "bagel"}, "bagel"),
         ({"architectures": ["HunyuanImage3ForConditionalGeneration"]}, "hunyuan_image3"),
+        ({"model_type": "mammothmoda2"}, "mammoth_moda2"),
         ({"model_type": "qwen2_vl"}, "generic"),
     ],
 )
@@ -50,6 +51,22 @@ def test_bagel_x_to_text_prompt_builder() -> None:
     assert stop_token_ids is None
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_mammoth_x_to_text_prompt_builder() -> None:
+    prompt, stop_token_ids = build_x_to_text_prompt("mammoth_moda2", "unused", "Hello.", has_image=False)
+    assert prompt == {
+        "prompt": (
+            "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
+            "<|im_start|>user\nHello.<|im_end|>\n"
+            "<|im_start|>assistant\n"
+        ),
+        "modalities": ["text"],
+        "additional_information": {"omni_task": ["chat"]},
+    }
+    assert stop_token_ids is None
+
+
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_generic_x_to_text_prompt_builder() -> None:
@@ -555,3 +572,64 @@ def test_declared_extra_args_apply_to_existing_sampling_params() -> None:
         "cfg_text_scale": 4.0,
         "think": False,
     }
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_mammothmoda2_extra_registry_declares_request_and_response_params() -> None:
+    for model_class_name in (
+        "MammothModa2DiTPipeline",
+        "MammothModa2ForConditionalGeneration",
+        "Mammothmoda2Model",
+    ):
+        assert get_extra_body_params(model_class_name) == frozenset(
+            {
+                "text_guidance_scale",
+                "cfg_range",
+                "num_inference_steps",
+            }
+        )
+        assert get_extra_output_params(model_class_name) == frozenset()
+        assert should_init_extra_args_for_non_diffusion_stages(model_class_name) is True
+
+    wrapper_prompt = build_text_to_image_prompt(
+        "MammothModa2ForConditionalGeneration",
+        {"prompt": "a cat", "modalities": ["image"]},
+        height=256,
+        width=256,
+    )
+    assert wrapper_prompt["additional_information"]["omni_task"] == ["t2i"]
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_mammothmoda2_text_to_image_prompt_builder() -> None:
+    # Image dims are converted to the AR grid (width/16 x height/16); the negative
+    # prompt is ignored (MammothModa2 t2i uses CFG, not an explicit negative path).
+    assert build_text_to_image_prompt(
+        "MammothModa2DiTPipeline",
+        {
+            "prompt": "a cat",
+            "modalities": ["image"],
+            "negative_prompt": "blurry",
+        },
+        height=512,
+        width=768,
+    ) == {
+        "prompt": (
+            "<|im_start|>system\nYou are a helpful image generator.<|im_end|>\n"
+            "<|im_start|>user\na cat<|im_end|>\n"
+            "<|im_start|>assistant\n"
+            "<|image start|>48*32<|image token|>"
+        ),
+        "additional_information": {
+            "omni_task": ["t2i"],
+            "ar_width": [48],
+            "ar_height": [32],
+            "image_height": [512],
+            "image_width": [768],
+            "eol_token_id": [152064],
+            "visual_token_start_id": [152072],
+            "visual_token_end_id": [168456],
+        },
+    }