From 78934753ff7e092b24ada0e34062cf3d121efd25 Mon Sep 17 00:00:00 2001
From: Nick Cao <ncao@redhat.com>
Date: Tue, 15 Sep 2026 12:36:34 -0400
Subject: [PATCH] [Tests] Test image edit input limits in tiny model framework
 (#7434)

Signed-off-by: Nick Cao <ncao@redhat.com>
Co-authored-by: Codex <noreply@openai.com>
---
 tests/model_tests/diffusion/task_runners.py   | 50 +++++++++++++++++++
 .../diffusion/test_common_online.py           | 17 +++++++
 2 files changed, 67 insertions(+)

diff --git a/tests/model_tests/diffusion/task_runners.py b/tests/model_tests/diffusion/task_runners.py
index 010d73777..991c74603 100644
--- a/tests/model_tests/diffusion/task_runners.py
+++ b/tests/model_tests/diffusion/task_runners.py
@@ -15,6 +15,8 @@ import numpy as np
 import pytest
 from PIL import Image
 
+from benchmarks.accuracy.common import pil_to_png_bytes
+from tests.helpers.assertions import assert_images_generations_response
 from tests.helpers.runtime import DiffusionResponse, OmniServer, OnlineOmniClient, dummy_messages_from_mix_data
 from tests.model_tests.diffusion.config_types import DiffusionTasks
 from vllm_omni.entrypoints.omni import Omni
@@ -204,6 +206,25 @@ def _run_online_i2i(
     return client.send_diffusion_request(request_config)
 
 
+def _build_online_image_edits_request(model: str, num_images: int) -> dict[str, Any]:
+    """Build a multipart ``/v1/images/edits`` request with repeated image fields."""
+    files = [
+        ("image", (f"image_{index}.png", pil_to_png_bytes(INPUT_IMAGE), "image/png")) for index in range(num_images)
+    ]
+    return {
+        "data": {
+            "model": model,
+            "prompt": PROMPT,
+            "size": f"{WIDTH}x{HEIGHT}",
+            "n": 1,
+            "response_format": "b64_json",
+            "num_inference_steps": 2,
+            "seed": 42,
+        },
+        "files": files,
+    }
+
+
 ### Offline task runners
 def run_and_validate_text_to_image_request(omni: Omni):
     """Run and validate a text to image request."""
@@ -324,6 +345,35 @@ def run_and_validate_online_image_to_image_request(server: OmniServer, client: O
     _validate_images(_get_online_images(_run_online_i2i(server, client)))
 
 
+def run_and_validate_online_image_edits(
+    server: OmniServer,
+    client: OnlineOmniClient,
+    num_images: int,
+    max_multimodal_image_inputs: int,
+):
+    """Run and validate one request through ``/v1/images/edits``."""
+    request = _build_online_image_edits_request(server.model, num_images)
+    if num_images <= max_multimodal_image_inputs:
+        response = client.send_images_edits_http_request(request)[0]
+        assert response.success, response.error_message
+        assert isinstance(response.json_body, dict)
+        assert_images_generations_response(
+            response.json_body,
+            {"json": {"n": 1, "size": f"{WIDTH}x{HEIGHT}"}},
+        )
+    else:
+        err_message = (
+            "Only a single image is supported by this model."
+            if max_multimodal_image_inputs == 1
+            else f"At most {max_multimodal_image_inputs} images are supported"
+        )
+        client.send_images_edits_http_request(
+            request,
+            err_code=400,
+            err_message=err_message,
+        )
+
+
 def run_and_validate_online_determinism(server: OmniServer, client: OnlineOmniClient, task_type: DiffusionTasks):
     """Checks for determinism through the server, dispatching by task type.
 
diff --git a/tests/model_tests/diffusion/test_common_online.py b/tests/model_tests/diffusion/test_common_online.py
index 65089e3ef..ed7f0e2fa 100644
--- a/tests/model_tests/diffusion/test_common_online.py
+++ b/tests/model_tests/diffusion/test_common_online.py
@@ -20,12 +20,14 @@ from tests.model_tests.diffusion.config_types import (
 from tests.model_tests.diffusion.model_settings import DIFFUSION_TEST_SETTINGS
 from tests.model_tests.diffusion.task_runners import (
     run_and_validate_online_determinism,
+    run_and_validate_online_image_edits,
     run_and_validate_online_image_to_image_request,
     run_and_validate_online_image_to_video_request,
     run_and_validate_online_multi_output,
     run_and_validate_online_text_to_image_request,
     run_and_validate_online_text_to_video_request,
 )
+from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata
 
 # NOTE : Hardware marks are added dynamically based on test requirements
 pytestmark = [pytest.mark.diffusion, pytest.mark.xdist]
@@ -84,6 +86,21 @@ def test_online_on_supported_tasks(
                     run_and_validate_online_text_to_image_request(server, client)
                 elif task_type == DiffusionTasks.IMAGE_TO_IMAGE:
                     run_and_validate_online_image_to_image_request(server, client)
+                    max_multimodal_image_inputs = (
+                        get_diffusion_model_metadata(model_name).max_multimodal_image_inputs or 1
+                    )
+                    image_counts = [1]
+                    if max_multimodal_image_inputs != 1:
+                        image_counts.append(max_multimodal_image_inputs)
+                    image_counts.append(max_multimodal_image_inputs + 1)
+                    for num_images in image_counts:
+                        with subtests.test(api="/v1/images/edits", num_images=num_images):
+                            run_and_validate_online_image_edits(
+                                server,
+                                client,
+                                num_images=num_images,
+                                max_multimodal_image_inputs=max_multimodal_image_inputs,
+                            )
                 elif task_type == DiffusionTasks.TEXT_TO_VIDEO:
                     run_and_validate_online_text_to_video_request(server, client)
                 elif task_type == DiffusionTasks.IMAGE_TO_VIDEO: