[Model] Migrate MammothModa2 DiT to the shared diffusion runtime (#7134)

Signed-off-by: levius <2114377220@qq.com>
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
Co-authored-by: Yueqian Lin <linyueqian@outlook.com>

tests/config/test_mammoth_moda2_shared_runtime.py

@@ -0,0 +1,60 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+
+from tests.helpers.stage_config import get_deploy_config_path
+from vllm_omni.config.endpoint_policy import OmniServingCapability
+from vllm_omni.config.stage_config import StageExecutionType, load_deploy_config
+from vllm_omni.diffusion.registry import _DIFFUSION_MODELS
+from vllm_omni.model_executor.models.mammoth_moda2.pipeline import (
+    MAMMOTH_MODA2_AR_PIPELINE,
+    MAMMOTH_MODA2_PIPELINE,
+)
+from vllm_omni.model_executor.models.registry import _OMNI_MODELS
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_mammothmoda2_generation_stage_uses_shared_diffusion_runtime() -> None:
+    ar_stage, dit_stage = MAMMOTH_MODA2_PIPELINE.stages
+    assert ar_stage.execution_type is StageExecutionType.LLM_AR
+    assert dit_stage.execution_type is StageExecutionType.DIFFUSION
+    assert dit_stage.model_arch == "MammothModa2DiTPipeline"
+    assert dit_stage.custom_process_input_func.endswith(".ar2diffusion")
+    assert dit_stage.omni_kv_config == {"need_recv_cache": False}
+    assert dit_stage.input_sources == (0,)
+    assert dit_stage.final_output is True
+    assert dit_stage.final_output_type == "image"
+
+
+def test_mammothmoda2_dit_is_registered_only_with_diffusion_runtime() -> None:
+    assert _DIFFUSION_MODELS["MammothModa2DiTPipeline"] == (
+        "mammoth_moda2",
+        "pipeline_mammothmoda2_dit",
+        "MammothModa2DiTPipeline",
+    )
+    assert "MammothModa2DiTPipeline" not in _OMNI_MODELS
+
+
+def test_mammothmoda2_deploy_does_not_shadow_request_sampling_fields() -> None:
+    deploy = load_deploy_config(get_deploy_config_path("mammoth_moda2.yaml"))
+
+    assert deploy.stages[1].default_sampling_params is None
+
+
+def test_mammothmoda2_rejects_unsupported_image_edits() -> None:
+    restrictions = {
+        restriction.capability: restriction.reason for restriction in MAMMOTH_MODA2_PIPELINE.endpoint_restrictions
+    }
+
+    assert set(restrictions) == {OmniServingCapability.IMAGE_EDITS}
+    assert "img2img" in restrictions[OmniServingCapability.IMAGE_EDITS]
+    assert MAMMOTH_MODA2_AR_PIPELINE.endpoint_restrictions == ()
+
+
+def test_mammothmoda2_ar_only_topology_is_unchanged() -> None:
+    assert len(MAMMOTH_MODA2_AR_PIPELINE.stages) == 1
+    stage = MAMMOTH_MODA2_AR_PIPELINE.stages[0]
+    assert stage.execution_type is StageExecutionType.LLM_AR
+    assert stage.final_output_type == "text"