[Bugfix] Fix image num check error (#5838)

Signed-off-by: dengyunyang <584797741@qq.com>
This commit is contained in:
dengyunyang authored and GitHub committed 2026-08-07 08:50:24 +00:00
1 parent ab459b36c2
commit f114113409
2 files changed
+46 -2

No files matched your search

@@ -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)
+18 -1
View File
@@ -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()