diff --git a/tests/entrypoints/openai_api/test_image_server.py b/tests/entrypoints/openai_api/test_image_server.py index 35fccce12..b5d22dc09 100644 --- a/tests/entrypoints/openai_api/test_image_server.py +++ b/tests/entrypoints/openai_api/test_image_server.py @@ -194,7 +194,7 @@ def async_omni_test_client(): ] 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=1048576, # 1024*1024 to support resolution tests ) return TestClient(app) @@ -754,6 +754,119 @@ def test_image_edit_parameter_pass(async_omni_test_client): assert data["size"] == "16x24" +def test_image_edit_layers_and_resolution(async_omni_test_client): + """Test layers and resolution parameters for layered models.""" + img_bytes = make_test_image_bytes((16, 16)) + + response = async_omni_test_client.post( + "/v1/images/edits", + files=[("image", img_bytes)], + data={ + "prompt": "decompose into layers", + "layers": 4, + "resolution": 1024, + }, + ) + assert response.status_code == 200 + engine = async_omni_test_client.app.state.engine_client + captured_sampling_params = engine.captured_sampling_params_list[-1] + assert captured_sampling_params.layers == 4 + assert captured_sampling_params.resolution == 1024 + + +def test_image_edit_resolution_auto_size(async_omni_test_client): + """Test that size='auto' with resolution lets pipeline calculate dimensions.""" + img_bytes = make_test_image_bytes((16, 16)) + + response = async_omni_test_client.post( + "/v1/images/edits", + files=[("image", img_bytes)], + data={ + "prompt": "test", + "size": "auto", + "resolution": 640, + }, + ) + assert response.status_code == 200 + engine = async_omni_test_client.app.state.engine_client + captured_sampling_params = engine.captured_sampling_params_list[-1] + # When resolution is set with size=auto, width/height should be None + # to let pipeline calculate based on resolution + assert captured_sampling_params.width is None + assert captured_sampling_params.height is None + assert captured_sampling_params.resolution == 640 + + +def test_image_edit_invalid_resolution(async_omni_test_client): + """Test that invalid resolution values are rejected with 400.""" + img_bytes = make_test_image_bytes((16, 16)) + + response = async_omni_test_client.post( + "/v1/images/edits", + files=[("image", img_bytes)], + data={ + "prompt": "test", + "resolution": 512, # Invalid, only 640 or 1024 are supported + }, + ) + assert response.status_code == 400 + detail = response.json()["detail"] + assert "Invalid resolution" in detail + assert "512" in detail + + +def test_image_edit_invalid_layers(async_omni_test_client): + """Test that invalid layers values (< 1) are rejected with 400.""" + img_bytes = make_test_image_bytes((16, 16)) + + # Test layers = 0 + response = async_omni_test_client.post( + "/v1/images/edits", + files=[("image", img_bytes)], + data={ + "prompt": "test", + "layers": 0, + }, + ) + assert response.status_code == 400 + detail = response.json()["detail"] + assert "Invalid layers" in detail + assert "layers must be >= 1" in detail + + # Test layers = -1 + response = async_omni_test_client.post( + "/v1/images/edits", + files=[("image", img_bytes)], + data={ + "prompt": "test", + "layers": -1, + }, + ) + assert response.status_code == 400 + detail = response.json()["detail"] + assert "Invalid layers" in detail + + +def test_image_edit_resolution_and_size_conflict(async_omni_test_client): + """Test that providing both resolution and explicit size raises 400.""" + img_bytes = make_test_image_bytes((16, 16)) + + response = async_omni_test_client.post( + "/v1/images/edits", + files=[("image", img_bytes)], + data={ + "prompt": "test", + "resolution": 1024, + "size": "512x512", # Conflict: both resolution and explicit size + }, + ) + assert response.status_code == 400 + detail = response.json()["detail"] + assert "Cannot specify both" in detail + assert "resolution" in detail + assert "size" in detail + + def test_image_edit_parameter_default(async_omni_test_client): img_bytes_1 = make_test_image_bytes((24, 16)) @@ -776,12 +889,13 @@ def test_image_edit_parameter_default(async_omni_test_client): assert captured_sampling_params.num_inference_steps == 4 assert captured_sampling_params.guidance_scale == 7.5 + # Test that a size exceeding max_generated_image_size returns 400 response = async_omni_test_client.post( "/v1/images/edits", files=[("image", img_bytes_1)], data={ "prompt": "hello world.", - "size": "96x96", + "size": "2048x2048", # 4,194,304 pixels > max_generated_image_size (1,048,576) }, ) assert response.status_code == 400 diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index d41d83a97..3e14bc2e5 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -115,6 +115,9 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniSamplingParam logger = init_logger(__name__) router = APIRouter() + +# Supported resolution buckets for layered models (e.g., Qwen-Image-Layered) +SUPPORTED_LAYERED_RESOLUTIONS = (640, 1024) profiler_router = APIRouter() @@ -1333,6 +1336,9 @@ async def edit_images( generator_device: str | None = Form(None), # vllm-omni extension for per-request LoRA. lora: str | None = Form(None), # Json string + # vllm-omni extension for layered models (e.g., Qwen-Image-Layered) + layers: int | None = Form(None), + resolution: int | None = Form(None), # See SUPPORTED_LAYERED_RESOLUTIONS ) -> ImageGenerationResponse: """ OpenAI-compatible image edit endpoint. @@ -1392,25 +1398,60 @@ async def edit_images( lora_request, lora_scale = _parse_lora_request(lora_dict) _update_if_not_none(gen_params, "lora_request", lora_request) _update_if_not_none(gen_params, "lora_scale", lora_scale) - # 3.2 Parse and add size if provided + # 3.2 Validate resolution if provided + if resolution is not None and resolution not in SUPPORTED_LAYERED_RESOLUTIONS: + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail=f"Invalid resolution {resolution}. Supported resolutions: {SUPPORTED_LAYERED_RESOLUTIONS}.", + ) + # 3.2.1 Validate layers if provided + if layers is not None and layers < 1: + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail=f"Invalid layers value {layers}. layers must be >= 1.", + ) + # 3.2.2 Check for conflicting size and resolution parameters + if resolution is not None and size.lower() != "auto": + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail="Cannot specify both 'resolution' and 'size'. " + "Use 'resolution' with size='auto', or use 'size' without 'resolution'.", + ) + + # 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": - width, height = pil_images[0].size # Use first image size + if resolution is None: + # No resolution specified, use input image size + width, height = pil_images[0].size + # else: let pipeline calculate dimensions based on resolution else: width, height = parse_size(size) - if max_generated_image_size is not None and (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.", - ) - size_str = f"{width}x{height}" + # 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.", + ) + + 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) _update_if_not_none(gen_params, "height", height) - # 3.3 Add optional parameters ONLY if provided + # 3.4 Add optional parameters ONLY if provided _update_if_not_none(gen_params, "num_inference_steps", num_inference_steps) _update_if_not_none(gen_params, "guidance_scale", guidance_scale) _update_if_not_none(gen_params, "true_cfg_scale", true_cfg_scale) @@ -1420,6 +1461,8 @@ async def edit_images( # might produce blurry images in some environments. _update_if_not_none(gen_params, "seed", seed if seed is not None else random.randint(0, 2**32 - 1)) _update_if_not_none(gen_params, "generator_device", generator_device) + _update_if_not_none(gen_params, "layers", layers) + _update_if_not_none(gen_params, "resolution", resolution) # 4. Generate images using AsyncOmni (multi-stage mode) request_id = f"img_edit-{random_uuid()}" @@ -1602,7 +1645,15 @@ def _extract_images_from_result(result: Any) -> list[Any]: images = request_output["images"] elif hasattr(request_output, "images") and request_output.images: images = request_output.images - return images + # Flatten nested lists (e.g., from layered models like Qwen-Image-Layered). + # Note: This only flattens one level deep. Deeper nesting is not supported. + flattened = [] + for img in images: + if isinstance(img, list): + flattened.extend(img) + else: + flattened.append(img) + return flattened async def _load_input_images(