[Bugfix] Limit Qwen-Image-Edit-2511 input image count (#2840)

Signed-off-by: david6666666 <530634352@qq.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>

tests/diffusion/models/qwen_image/test_qwen_image_edit_plus.py

@@ -0,0 +1,38 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import json
+from pathlib import Path
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+from PIL import Image
+
+from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image_edit_plus import (
+    get_qwen_image_edit_plus_pre_process_func,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def test_qwen_image_edit_plus_rejects_too_many_input_images(tmp_path: Path):
+    vae_dir = tmp_path / "vae"
+    vae_dir.mkdir()
+    # Keep the mock config intentionally minimal: this test only needs the
+    # fields touched during pre-process initialization.
+    (vae_dir / "config.json").write_text(json.dumps({"z_dim": 16}))
+
+    pre_process = get_qwen_image_edit_plus_pre_process_func(SimpleNamespace(model=str(tmp_path)))
+    image = Image.fromarray(np.zeros((32, 32, 3), dtype=np.uint8))
+    request = SimpleNamespace(
+        prompts=[
+            {
+                "prompt": "combine",
+                "multi_modal_data": {"image": [image, image, image, image, image]},
+            }
+        ],
+        sampling_params=SimpleNamespace(height=None, width=None),
+    )
+
+    with pytest.raises(ValueError, match=r"At most 4 images are supported by this model"):
+        pre_process(request)

tests/entrypoints/openai_api/test_image_server.py

@@ -782,6 +782,63 @@ def test_image_edit_rejects_multiple_images_when_model_does_not_support_them(asy
     assert engine.captured_prompt is None
 
 
+def test_image_edit_rejects_too_many_images_for_qwen_image_edit_2511(async_omni_test_client):
+    engine = async_omni_test_client.app.state.engine_client
+    engine.get_diffusion_od_config = lambda: SimpleNamespace(
+        supports_multimodal_inputs=True,
+        max_multimodal_image_inputs=4,
+    )
+
+    response = async_omni_test_client.post(
+        "/v1/images/edits",
+        files=[
+            ("image", make_test_image_bytes((16, 16))),
+            ("image", make_test_image_bytes((16, 16))),
+            ("image", make_test_image_bytes((16, 16))),
+            ("image", make_test_image_bytes((16, 16))),
+            ("image", make_test_image_bytes((16, 16))),
+        ],
+        data={"prompt": "hello world."},
+    )
+
+    assert response.status_code == 400
+    assert response.json()["detail"] == "Received 5 input images. At most 4 images are supported by this model."
+    assert engine.captured_prompt is None
+
+
+def test_image_edit_rejects_too_many_images_for_qwen_image_edit_2511_before_loading(
+    async_omni_test_client, monkeypatch: pytest.MonkeyPatch
+):
+    import vllm_omni.entrypoints.openai.api_server as api_server_module
+
+    engine = async_omni_test_client.app.state.engine_client
+    engine.get_diffusion_od_config = lambda: SimpleNamespace(
+        supports_multimodal_inputs=True,
+        max_multimodal_image_inputs=4,
+    )
+
+    def _fail_load(*args, **kwargs):
+        raise AssertionError("_load_input_images should not run for over-limit requests")
+
+    monkeypatch.setattr(api_server_module, "_load_input_images", _fail_load)
+
+    response = async_omni_test_client.post(
+        "/v1/images/edits",
+        files=[
+            ("image", make_test_image_bytes((16, 16))),
+            ("image", make_test_image_bytes((16, 16))),
+            ("image", make_test_image_bytes((16, 16))),
+            ("image", make_test_image_bytes((16, 16))),
+            ("image", make_test_image_bytes((16, 16))),
+        ],
+        data={"prompt": "hello world."},
+    )
+
+    assert response.status_code == 400
+    assert response.json()["detail"] == "Received 5 input images. At most 4 images are supported by this model."
+    assert engine.captured_prompt is None
+
+
 def test_image_edit_parameter_pass(async_omni_test_client):
     img_bytes_1 = make_test_image_bytes((16, 16))
 

tests/test_diffusion_config_propagation.py

@@ -15,6 +15,7 @@ from vllm_omni.diffusion.data import (
     DiffusionParallelConfig,
     OmniDiffusionConfig,
 )
+from vllm_omni.diffusion.model_metadata import QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -109,3 +110,12 @@ class TestCreateDefaultDiffusion:
         ea = stages[0]["engine_args"]
         assert ea["enforce_eager"] is True
         assert ea["lora_path"] == "/tmp/lora"
+
+
+def test_qwen_image_edit_plus_sets_generic_multimodal_limit():
+    od_config = OmniDiffusionConfig(model="Qwen/Qwen-Image-Edit-2511", model_class_name="QwenImageEditPlusPipeline")
+
+    od_config.update_multimodal_support()
+
+    assert od_config.supports_multimodal_inputs is True
+    assert od_config.max_multimodal_image_inputs == QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES