From 0fd2620d5669c9aa2cfbbc6d394fb1412120bfc6 Mon Sep 17 00:00:00 2001 From: Jiantong Men <40687627+menjiantong@users.noreply.github.com> Date: Thu, 25 Jun 2026 13:14:02 +0800 Subject: [PATCH] [BugFix] Fix MammothModa2 Text-to-Image Bug (#4299) Signed-off-by: menjiantong Co-authored-by: menjiantong Co-authored-by: Didan Deng <33117903+wtomin@users.noreply.github.com> --- .../test_mammoth_moda2_config.py | 91 +++++++++++++++++++ .../pipeline_mammothmoda2_dit.py | 14 ++- .../stage_input_processors/mammoth_moda2.py | 6 +- .../configs/mammoth_moda2.py | 10 +- 4 files changed, 116 insertions(+), 5 deletions(-) create mode 100644 tests/unit/mammoth_moda2/test_mammoth_moda2_config.py diff --git a/tests/unit/mammoth_moda2/test_mammoth_moda2_config.py b/tests/unit/mammoth_moda2/test_mammoth_moda2_config.py new file mode 100644 index 000000000..77b68ca3f --- /dev/null +++ b/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") diff --git a/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py b/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py index e230fab33..39fbb92b4 100644 --- a/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py +++ b/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) diff --git a/vllm_omni/model_executor/stage_input_processors/mammoth_moda2.py b/vllm_omni/model_executor/stage_input_processors/mammoth_moda2.py index 7c821c243..5ab4b6f48 100644 --- a/vllm_omni/model_executor/stage_input_processors/mammoth_moda2.py +++ b/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"] diff --git a/vllm_omni/transformers_utils/configs/mammoth_moda2.py b/vllm_omni/transformers_utils/configs/mammoth_moda2.py index 93cddd735..1c6352264 100644 --- a/vllm_omni/transformers_utils/configs/mammoth_moda2.py +++ b/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: