[Bugfix] Fix image num check error (#5838)
Signed-off-by: dengyunyang <584797741@qq.com>
This commit is contained in:
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)
|
||||
|
||||
@@ -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()
|
||||
Reference in new issue
Block a user