[Fix] Restore MammothModa2 support

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

vllm_omni/model_executor/models/mammoth_moda2/__init__.py

@@ -0,0 +1,8 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# NOTE: Do not import model classes in this file. Importing any
+# submodule in this package triggers __init__.py execution, and
+# both the model registry and pipeline registry import submodules
+# directly — heavy imports here would be loaded as a side effect
+# even though nothing depends on these re-exports.

vllm_omni/model_executor/models/mammoth_moda2/pipeline.py

@@ -0,0 +1,68 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""MammothModa2 pipeline topology (frozen).
+
+Stage 0: AR  — multimodal understanding + latent generation
+Stage 1: DiT — latent → image
+
+For text/image understanding tasks (text output only), use MAMMOTH_MODA2_AR_PIPELINE.
+"""
+
+from vllm_omni.config.stage_config import (
+    PipelineConfig,
+    StageExecutionType,
+    StagePipelineConfig,
+)
+
+_PROC = "vllm_omni.model_executor.stage_input_processors.mammoth_moda2"
+
+MAMMOTH_MODA2_PIPELINE = PipelineConfig(
+    model_type="mammoth_moda2",
+    default_deploy_config_name="mammoth_moda2.yaml",
+    model_arch="MammothModa2ForConditionalGeneration",
+    hf_architectures=("Mammothmoda2Model", "MammothModa2ForConditionalGeneration"),
+    stages=(
+        StagePipelineConfig(
+            stage_id=0,
+            model_stage="ar",
+            execution_type=StageExecutionType.LLM_AR,
+            input_sources=(),
+            final_output=False,
+            owns_tokenizer=True,
+            requires_multimodal_data=True,
+            engine_output_type="latent",
+        ),
+        StagePipelineConfig(
+            stage_id=1,
+            model_stage="dit",
+            execution_type=StageExecutionType.LLM_GENERATION,
+            input_sources=(0,),
+            final_output=True,
+            final_output_type="image",
+            owns_tokenizer=False,
+            requires_multimodal_data=False,
+            engine_output_type="image",
+            custom_process_input_func=f"{_PROC}.ar2dit",
+        ),
+    ),
+)
+
+# Single-stage AR variant for understanding / summarization tasks
+MAMMOTH_MODA2_AR_PIPELINE = PipelineConfig(
+    model_type="mammoth_moda2_ar",
+    default_deploy_config_name="mammoth_moda2_ar.yaml",
+    model_arch="MammothModa2ForConditionalGeneration",
+    stages=(
+        StagePipelineConfig(
+            stage_id=0,
+            model_stage="ar",
+            execution_type=StageExecutionType.LLM_AR,
+            input_sources=(),
+            final_output=True,
+            final_output_type="text",
+            owns_tokenizer=True,
+            requires_multimodal_data=True,
+            engine_output_type="text",
+        ),
+    ),
+)

vllm_omni/model_executor/models/mammoth_moda2/pipeline_mammothmoda2_dit.py

@@ -0,0 +1,15 @@
+"""
+Compatibility shim.
+
+The MammothModa2 DiT implementation lives under `vllm_omni.diffusion` to align
+with other ARDiT structured models. We keep this module path so existing
+OmniModelRegistry entries (and downstream code) keep working.
+"""
+
+from vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit import (  # noqa: F401
+    MammothModa2DiTPipeline,
+)
+
+__all__ = [
+    "MammothModa2DiTPipeline",
+]

vllm_omni/model_executor/models/registry.py

@@ -114,6 +114,41 @@ _OMNI_MODELS = {
         "omnivoice",
         "OmniVoiceModel",
     ),
+    "MammothModa2Qwen2ForCausalLM": (
+        "mammoth_moda2",
+        "mammoth_moda2",
+        "MammothModa2Qwen2ForCausalLM",
+    ),
+    "MammothModa2ARForConditionalGeneration": (
+        "mammoth_moda2",
+        "mammoth_moda2",
+        "MammothModa2ARForConditionalGeneration",
+    ),
+    "MammothModa2Qwen3ARForConditionalGeneration": (
+        "mammoth_moda2",
+        "mammoth_moda2",
+        "MammothModa2Qwen3ARForConditionalGeneration",
+    ),
+    "MammothModa2Qwen3ForCausalLM": (
+        "mammoth_moda2",
+        "mammoth_moda2",
+        "MammothModa2Qwen3ForCausalLM",
+    ),
+    "MammothModa2DiTPipeline": (
+        "mammoth_moda2",
+        "pipeline_mammothmoda2_dit",
+        "MammothModa2DiTPipeline",
+    ),
+    "MammothModa2ForConditionalGeneration": (
+        "mammoth_moda2",
+        "mammoth_moda2",
+        "MammothModa2ForConditionalGeneration",
+    ),
+    "Mammothmoda2Model": (
+        "mammoth_moda2",
+        "mammoth_moda2",
+        "MammothModa2ForConditionalGeneration",
+    ),
     "Qwen3TTSForConditionalGeneration": (
         "qwen3_tts",
         "qwen3_tts_talker",