[BugFix] Fix MammothModa2 Text-to-Image Bug (#4299)

Signed-off-by: menjiantong <menjiantong@foxmail.com>
Co-authored-by: menjiantong <menjiantong@foxmail.com>
Co-authored-by: Didan Deng <33117903+wtomin@users.noreply.github.com>

tests/unit/mammoth_moda2/test_mammoth_moda2_config.py

@@ -0,0 +1,91 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Unit tests for MammothModa2 config classes.
+
+These tests verify the correct behavior of the Mammothmoda2Config class,
+particularly the get_text_config() method which is called during parent
+class initialization (PretrainedConfig.__init__ -> validation -> get_text_config).
+
+Regression test for: AttributeError: 'Mammothmoda2Config' object has no attribute 'llm_config'
+"""
+
+import pytest
+
+from vllm_omni.transformers_utils.configs.mammoth_moda2 import (
+    Mammothmoda2Config,
+    Mammothmoda2Qwen2_5_VLConfig,
+    Mammothmoda2Qwen2_5_VLTextConfig,
+)
+
+
+@pytest.mark.cpu
+class TestMammothmoda2Config:
+    """Tests for Mammothmoda2Config class."""
+
+    def test_get_text_config_returns_nested_text_config(self):
+        """
+        Verify that get_text_config() returns the nested text_config from llm_config.
+
+        This is a regression test for the AttributeError that occurred when
+        get_text_config() was called during parent class initialization before
+        llm_config was set.
+        """
+        config = Mammothmoda2Config(
+            llm_config={
+                "model_type": "mammothmoda2_qwen2_5_vl",
+            }
+        )
+
+        assert config.llm_config is not None
+        assert isinstance(config.llm_config, Mammothmoda2Qwen2_5_VLConfig)
+
+        text_config = config.get_text_config()
+        assert text_config is not None
+        assert isinstance(text_config, Mammothmoda2Qwen2_5_VLTextConfig)
+
+        # The text_config should be the same object as llm_config.text_config
+        assert text_config is config.llm_config.text_config
+
+    def test_get_text_config_returns_none_when_llm_config_is_none(self):
+        """Verify that get_text_config() returns None when llm_config is None."""
+        config = Mammothmoda2Config(llm_config=None)
+
+        assert config.llm_config is None
+        assert config.get_text_config() is None
+
+    def test_init_sets_llm_config_before_parent_init(self):
+        """
+        Verify that llm_config is set before parent __init__ is called.
+
+        This is critical because the parent PretrainedConfig.__init__ calls
+        validation which in turn calls get_text_config(). If llm_config is not
+        set before super().__init__(), an AttributeError will be raised.
+        """
+        # This should not raise an AttributeError
+        config = Mammothmoda2Config(
+            llm_config={
+                "model_type": "mammothmoda2_qwen2_5_vl",
+            }
+        )
+        assert config.llm_config is not None
+
+
+@pytest.mark.cpu
+class TestMammothmoda2Qwen2_5_VLConfig:
+    """Tests for Mammothmoda2Qwen2_5_VLConfig class."""
+
+    def test_has_text_config_attribute(self):
+        """Verify that Mammothmoda2Qwen2_5_VLConfig has text_config attribute."""
+        config = Mammothmoda2Qwen2_5_VLConfig()
+        assert hasattr(config, "text_config")
+        assert isinstance(config.text_config, Mammothmoda2Qwen2_5_VLTextConfig)
+
+    def test_text_config_has_expected_attributes(self):
+        """Verify that the nested text_config has expected attributes for token validation."""
+        config = Mammothmoda2Qwen2_5_VLConfig()
+        text_config = config.text_config
+
+        assert hasattr(text_config, "vocab_size")
+        assert hasattr(text_config, "image_token_id")
+        assert hasattr(text_config, "video_token_id")

vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py

@@ -59,9 +59,19 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
         self.gen_vae = AutoencoderKL.from_config(self.config.gen_vae_config)
         self.gen_transformer = Transformer2DModel.from_config(self.config.gen_dit_config)
 
-        llm_hidden_size = int(getattr(self.config.llm_config, "hidden_size", 0) or 0)
+        # llm_config is a Mammothmoda2Qwen2_5_VLConfig which has nested text_config
+        llm_hidden_size = int(
+            getattr(
+                getattr(self.config.llm_config, "text_config", None),
+                "hidden_size",
+                0,
+            )
+            or 0
+        )
         if llm_hidden_size <= 0:
-            raise ValueError("Failed to infer llm hidden_size from Mammothmoda2Config.llm_config.hidden_size")
+            raise ValueError(
+                "Failed to infer llm hidden_size from Mammothmoda2Config.llm_config.text_config.hidden_size"
+            )
         self._reinit_caption_embedder(llm_hidden_size)
 
         # Optional: image condition refiner (Q-Former)

vllm_omni/model_executor/stage_input_processors/mammoth_moda2.py

@@ -11,12 +11,16 @@ from vllm_omni.inputs.data import OmniTokensPrompt
 
 def ar2dit(
     source_outputs: list[Any],
-    prompts: OmniTokensPrompt | TextPrompt | None = None,
+    prompts: OmniTokensPrompt | TextPrompt | list[OmniTokensPrompt | TextPrompt] | None = None,
     _requires_multimodal_data: bool = False,
 ) -> list[OmniTokensPrompt]:
     """Convert AR stage outputs to DiT stage inputs."""
     ar_outputs = source_outputs
 
+    # Normalize prompts to list
+    if not isinstance(prompts, list):
+        prompts = [prompts] if prompts is not None else [{}]
+
     dit_inputs: list[OmniTokensPrompt] = []
     for ar_output, prompt in zip(ar_outputs, prompts):
         addi_info = prompt["additional_information"]

vllm_omni/transformers_utils/configs/mammoth_moda2.py

@@ -231,7 +231,8 @@ class Mammothmoda2Config(PretrainedConfig):
         architectures: list[str] | None = None,
         **kwargs,
     ) -> None:
-        super().__init__(**kwargs)
+        # Must set llm_config before super().__init__() because the parent's
+        # validation calls get_text_config() which accesses self.llm_config
         self.llm_config = AutoConfig.for_model(**llm_config) if llm_config is not None else None
         self.gen_vae_config = gen_vae_config
         self.gen_dit_config = gen_dit_config
@@ -245,8 +246,13 @@ class Mammothmoda2Config(PretrainedConfig):
         self.tokenizer_class = "MammothUTokenizer"
         self.architectures = ["Mammothmoda2Model"]
 
+        super().__init__(**kwargs)
+
     def get_text_config(self, decoder: bool = False) -> PretrainedConfig:  # noqa: ARG002
-        return self.llm_config
+        if self.llm_config is None:
+            return None
+        # llm_config is a Mammothmoda2Qwen2_5_VLConfig which has nested text_config
+        return self.llm_config.text_config
 
     def _require_llm_config(self) -> PretrainedConfig:
         if self.llm_config is None: