From 78e1a0ae5ce6bec260043a085df2d35d42ce1d03 Mon Sep 17 00:00:00 2001 From: Qi Jia Date: Fri, 11 Sep 2026 22:18:07 +0800 Subject: [PATCH] [Bugfix][Model] Fix FLUX.2 Klein multi-image edit metadata (#7430) Signed-off-by: QI JIA Co-authored-by: QI JIA Co-authored-by: Cursor --- .../diffusion/test_diffusion_config_propagation.py | 13 +++++++++++++ vllm_omni/diffusion/model_metadata.py | 6 ++++++ 2 files changed, 19 insertions(+) diff --git a/tests/diffusion/test_diffusion_config_propagation.py b/tests/diffusion/test_diffusion_config_propagation.py index 394f7c04f..23ff796e2 100644 --- a/tests/diffusion/test_diffusion_config_propagation.py +++ b/tests/diffusion/test_diffusion_config_propagation.py @@ -17,6 +17,7 @@ from vllm_omni.diffusion.data import ( ) from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode from vllm_omni.diffusion.model_metadata import ( + FLUX2_KLEIN_MAX_INPUT_IMAGES, HUNYUAN_IMAGE3_MAX_INPUT_IMAGES, QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES, ) @@ -196,6 +197,18 @@ def test_qwen_image_edit_plus_sets_generic_multimodal_limit(): assert od_config.max_multimodal_image_inputs == QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES +def test_flux2_klein_sets_generic_multimodal_limit(): + od_config = OmniDiffusionConfig( + model="black-forest-labs/FLUX.2-klein-9B", + model_class_name="Flux2KleinPipeline", + ) + + od_config.update_multimodal_support() + + assert od_config.supports_multimodal_inputs is True + assert od_config.max_multimodal_image_inputs == FLUX2_KLEIN_MAX_INPUT_IMAGES + + def test_task_type_roundtrip(): od = _roundtrip_diffusion_config(model="x", task_type="model-defined-task") assert od.task_type == "model-defined-task" diff --git a/vllm_omni/diffusion/model_metadata.py b/vllm_omni/diffusion/model_metadata.py index c91d2ead5..2ee1b1f4b 100644 --- a/vllm_omni/diffusion/model_metadata.py +++ b/vllm_omni/diffusion/model_metadata.py @@ -19,6 +19,8 @@ class DiffusionModelMetadata: final_output_type: str | None = None +# FLUX.2 Klein supports up to four reference images. +FLUX2_KLEIN_MAX_INPUT_IMAGES = 4 QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES = 4 # Upstream HunyuanImage-3.0 "Multi-Image Fusion" caps reference images at 3. HUNYUAN_IMAGE3_MAX_INPUT_IMAGES = 3 @@ -27,6 +29,10 @@ BOOGU_IMAGE_MAX_INPUT_IMAGES = 1 _DIFFUSION_MODEL_METADATA: dict[str, DiffusionModelMetadata] = { + "Flux2KleinPipeline": DiffusionModelMetadata( + supports_multimodal_inputs=True, + max_multimodal_image_inputs=FLUX2_KLEIN_MAX_INPUT_IMAGES, + ), "QwenImageEditPlusPipeline": DiffusionModelMetadata( supports_multimodal_inputs=True, max_multimodal_image_inputs=QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES,