From 589e6649472fcc08f40c6159280e304588cc80c4 Mon Sep 17 00:00:00 2001 From: WeiQing Chen <40507679+david6666666@users.noreply.github.com> Date: Thu, 9 Apr 2026 14:42:23 +0800 Subject: [PATCH] =?UTF-8?q?[cherry-pick=20release/v0.18.0post1][Bugfix]=20?= =?UTF-8?q?Enforce=20--max-generated-image-size=20on=20/v1/images/generati?= =?UTF-8?q?ons=E2=80=A6=20(#2625)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: David Chen <530634352@qq.com> --- .../openai_api/test_image_server.py | 19 +++++-- vllm_omni/entrypoints/openai/api_server.py | 51 ++++++++++++------- 2 files changed, 49 insertions(+), 21 deletions(-) diff --git a/tests/entrypoints/openai_api/test_image_server.py b/tests/entrypoints/openai_api/test_image_server.py index 40fd32cc3..8695311fb 100644 --- a/tests/entrypoints/openai_api/test_image_server.py +++ b/tests/entrypoints/openai_api/test_image_server.py @@ -172,7 +172,7 @@ def test_client(mock_async_diffusion): app.state.diffusion_model_name = "Qwen/Qwen-Image" # For models endpoint app.state.args = Namespace( default_sampling_params='{"0": {"num_inference_steps":4, "guidance_scale":7.5}}', - max_generated_image_size=4096, # 64*64 + max_generated_image_size=1024 * 1792, ) return TestClient(app) @@ -239,7 +239,7 @@ def async_omni_stage_configs_only_client(): # AsyncOmni exposes stage_configs on the engine instance. app.state.args = Namespace( default_sampling_params='{"1": {"num_inference_steps":4, "guidance_scale":7.5}}', - max_generated_image_size=4096, # 64*64 + max_generated_image_size=1024 * 1792, ) return TestClient(app) @@ -386,6 +386,18 @@ def test_image_edits_async_omni_stage_configs_only(async_omni_stage_configs_only assert len(captured) == 2 +def test_generate_images_max_size_rejected(async_omni_test_client): + """Test that a size exceeding max_generated_image_size returns 400.""" + response = async_omni_test_client.post( + "/v1/images/generations", + json={ + "prompt": "a cat", + "size": "2048x2048", # 4,194,304 pixels > max_generated_image_size (1,048,576) + }, + ) + assert response.status_code == 400 + + def test_generate_multiple_images(test_client): """Test generating multiple images""" response = test_client.post( @@ -976,12 +988,13 @@ def test_image_edit_parameter_default_single_stage(test_client): assert captured_sampling_params.num_inference_steps == 4 assert captured_sampling_params.guidance_scale == 7.5 + # Size exceeding max_generated_image_size (1024*1792) returns 400 response = test_client.post( "/v1/images/edits", files=[("image", img_bytes_1)], data={ "prompt": "hello world.", - "size": "96x96", + "size": "2048x2048", }, ) assert response.status_code == 400 diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 9f5b37981..ceb87f111 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -1299,6 +1299,10 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request) size_str = f"{width}x{height}" else: size_str = "model default" + + app_state_args = getattr(raw_request.app.state, "args", None) + _check_max_generated_image_size(app_state_args, width, height) + _update_if_not_none(gen_params, "width", width) _update_if_not_none(gen_params, "height", height) @@ -1484,7 +1488,6 @@ async def edit_images( ) # 3.3 Parse and add size if provided - max_generated_image_size = getattr(app_state_args, "max_generated_image_size", None) width, height = None, None if size.lower() == "auto": if resolution is None: @@ -1494,23 +1497,7 @@ async def edit_images( else: width, height = parse_size(size) - # Check max_generated_image_size - if max_generated_image_size is not None: - if width is not None and height is not None: - if width * height > max_generated_image_size: - raise HTTPException( - status_code=HTTPStatus.BAD_REQUEST.value, - detail=f"Requested image size {width}x{height} exceeds the maximum allowed " - f"size of {max_generated_image_size} pixels.", - ) - elif resolution is not None: - # When resolution is set, the output size is resolution * resolution - if resolution * resolution > max_generated_image_size: - raise HTTPException( - status_code=HTTPStatus.BAD_REQUEST.value, - detail=f"Requested resolution {resolution} (max {resolution}x{resolution} pixels) " - f"exceeds the maximum allowed size of {max_generated_image_size} pixels.", - ) + _check_max_generated_image_size(app_state_args, width, height, resolution) size_str = f"{width}x{height}" if width is not None and height is not None else "auto" _update_if_not_none(gen_params, "width", width) @@ -1709,6 +1696,34 @@ async def _generate_with_async_omni( return result +def _check_max_generated_image_size( + app_state_args: Any, + width: int | None, + height: int | None, + resolution: int | None = None, +) -> None: + """Raise 400 if the requested image size exceeds --max-generated-image-size.""" + max_generated_image_size = getattr(app_state_args, "max_generated_image_size", None) + # Check max_generated_image_size + if max_generated_image_size is None: + return + if width is not None and height is not None: + if width * height > max_generated_image_size: + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail=f"Requested image size {width}x{height} exceeds the maximum allowed " + f"size of {max_generated_image_size} pixels.", + ) + elif resolution is not None: + # When resolution is set, the output size is resolution * resolution + if resolution * resolution > max_generated_image_size: + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail=f"Requested resolution {resolution} (max {resolution}x{resolution} pixels) " + f"exceeds the maximum allowed size of {max_generated_image_size} pixels.", + ) + + def _update_if_not_none(object: Any, key: str, val: Any) -> None: if val is not None: setattr(object, key, val)