From f11411340978a80fbec06d6d7ef389a6d84e697e Mon Sep 17 00:00:00 2001
From: dengyunyang <584797741@qq.com>
Date: Fri, 7 Aug 2026 16:50:24 +0800
Subject: [PATCH] [Bugfix] Fix image num check error (#5838)

Signed-off-by: dengyunyang <584797741@qq.com>
---
 .../test_diffusion_config_propagation.py      | 29 ++++++++++++++++++-
 vllm_omni/diffusion/model_metadata.py         | 19 +++++++++++-
 2 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/tests/diffusion/test_diffusion_config_propagation.py b/tests/diffusion/test_diffusion_config_propagation.py
index 205b2e491..91d366723 100644
--- a/tests/diffusion/test_diffusion_config_propagation.py
+++ b/tests/diffusion/test_diffusion_config_propagation.py
@@ -15,7 +15,10 @@ from vllm_omni.diffusion.data import (
     DiffusionParallelConfig,
     OmniDiffusionConfig,
 )
-from vllm_omni.diffusion.model_metadata import QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES
+from vllm_omni.diffusion.model_metadata import (
+    HUNYUAN_IMAGE3_MAX_INPUT_IMAGES,
+    QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES,
+)
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -136,6 +139,30 @@ def test_task_type_roundtrip():
     assert od.task_type == "model-defined-task"
 
 
+def test_architecture_name_resolves_via_pipeline_class_fallback():
+    # Direct hit: the architecture name itself is the metadata table key.
+    qwen_od_config = OmniDiffusionConfig(
+        model="Qwen/Qwen-Image-Edit-2511",
+        model_class_name="QwenImageEditPlusPipeline",
+    )
+    qwen_od_config.update_multimodal_support()
+
+    assert qwen_od_config.supports_multimodal_inputs is True
+    assert qwen_od_config.max_multimodal_image_inputs == QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES
+
+    # Fallback: the HF architecture name differs from the internal pipeline class
+    # name, so ``get_diffusion_model_metadata`` must consult ``_DIFFUSION_MODELS``
+    # and resolve the metadata via the pipeline class name.
+    hunyuan_od_config = OmniDiffusionConfig(
+        model="tencent/HunyuanImage-3.0-Instruct",
+        model_class_name="HunyuanImage3ForCausalMM",
+    )
+    hunyuan_od_config.update_multimodal_support()
+
+    assert hunyuan_od_config.supports_multimodal_inputs is True
+    assert hunyuan_od_config.max_multimodal_image_inputs == HUNYUAN_IMAGE3_MAX_INPUT_IMAGES
+
+
 def test_additional_config_roundtrip():
     additional_config = {"torchair_graph_config": {"enabled": True}}
     od = _roundtrip_diffusion_config(model="x", additional_config=additional_config)
diff --git a/vllm_omni/diffusion/model_metadata.py b/vllm_omni/diffusion/model_metadata.py
index 6184478b8..ecc53cdc6 100644
--- a/vllm_omni/diffusion/model_metadata.py
+++ b/vllm_omni/diffusion/model_metadata.py
@@ -64,4 +64,21 @@ def get_diffusion_model_metadata(model_class_name: str | None) -> DiffusionModel
     # pipelines do not accidentally inherit limits meant for other models.
     if model_class_name is None:
         return DiffusionModelMetadata()
-    return _DIFFUSION_MODEL_METADATA.get(model_class_name, DiffusionModelMetadata())
+    metadata = _DIFFUSION_MODEL_METADATA.get(model_class_name)
+    if metadata is not None:
+        return metadata
+    # Some checkpoints report the HF architecture name diff from internal pipeline class name
+    # (e.g. HunyuanImage3ForCausalMM, WanVACEPipeline, OmniVoice ...).
+    from vllm_omni.diffusion.registry import _DIFFUSION_MODELS
+
+    entry = _DIFFUSION_MODELS.get(model_class_name)
+    if entry is not None:
+        # Unpack instead of indexing so a future change to the registry tuple
+        # shape fails loudly instead of silently reading the wrong element.
+        _, _, pipeline_cls_name = entry
+        # Note: the registry ``cls_name`` and the metadata keys are two separate
+        # key spaces. Aliases whose pipeline class has no metadata entry (e.g.
+        # Wan22VACEPipeline, OmniVoicePipeline) still fall back to the defaults;
+        # that is not a regression, it just means no capability override.
+        return _DIFFUSION_MODEL_METADATA.get(pipeline_cls_name, DiffusionModelMetadata())
+    return DiffusionModelMetadata()