From 6e8c279a68c5dda298d418c369d10af2a2a8035c Mon Sep 17 00:00:00 2001 From: SYLAR <125541396+lishunyang12@users.noreply.github.com> Date: Sat, 1 Aug 2026 23:09:27 +0800 Subject: [PATCH] [Bugfix] Match image file metadata to output format (#5619) Signed-off-by: lishunyang12 Co-authored-by: lishunyang12 --- .../openai_api/test_image_server.py | 77 +++++++++++++++--- vllm_omni/entrypoints/openai/api_server.py | 78 +++++++++++-------- .../entrypoints/openai/image_api_utils.py | 3 +- .../entrypoints/openai/protocol/images.py | 19 ++++- 4 files changed, 129 insertions(+), 48 deletions(-) diff --git a/tests/entrypoints/openai_api/test_image_server.py b/tests/entrypoints/openai_api/test_image_server.py index 34ae5ba20..2d6c0299e 100644 --- a/tests/entrypoints/openai_api/test_image_server.py +++ b/tests/entrypoints/openai_api/test_image_server.py @@ -1301,7 +1301,11 @@ def test_generate_images_rejects_model_mismatch(test_client): assert "model mismatch" in response.json()["detail"].lower() -def test_image_file_response_format_multiple(test_client): +@pytest.mark.parametrize( + ("output_format", "expected_pil_format"), + [("png", "PNG"), ("jpeg", "JPEG"), ("webp", "WEBP")], +) +def test_image_file_response_format_multiple(test_client, output_format, expected_pil_format): """Test response_format=file with n>1 returns ZIP archive""" response = test_client.post( "/v1/images/generations", @@ -1309,6 +1313,7 @@ def test_image_file_response_format_multiple(test_client): "prompt": "a dog", "n": 3, "response_format": "file", + "output_format": output_format, }, ) @@ -1317,23 +1322,32 @@ def test_image_file_response_format_multiple(test_client): assert "attachment" in response.headers.get("content-disposition", "") assert ".zip" in response.headers.get("content-disposition", "") - # Verify it's a valid ZIP with 3 PNG files + # Verify it's a valid ZIP with 3 correctly labelled image files import zipfile zip_buffer = io.BytesIO(response.content) with zipfile.ZipFile(zip_buffer, "r") as zf: files = zf.namelist() assert len(files) == 3 - assert all(f.endswith(".png") for f in files) + assert all(f.endswith(f".{output_format}") for f in files) - # Verify each file is a valid PNG + # Verify each file's bytes match its advertised format for filename in files: img_bytes = zf.read(filename) img = Image.open(io.BytesIO(img_bytes)) - assert img.format == "PNG" + assert img.format == expected_pil_format -def test_image_file_response_format_single(test_client): +@pytest.mark.parametrize( + ("output_format", "expected_media_type", "expected_pil_format"), + [ + ("png", "image/png", "PNG"), + ("jpg", "image/jpeg", "JPEG"), + ("jpeg", "image/jpeg", "JPEG"), + ("webp", "image/webp", "WEBP"), + ], +) +def test_image_file_response_format_single(test_client, output_format, expected_media_type, expected_pil_format): """Test response_format=file with n=1 returns a single image file.""" response = test_client.post( "/v1/images/generations", @@ -1341,16 +1355,61 @@ def test_image_file_response_format_single(test_client): "prompt": "a dog", "n": 1, "response_format": "file", + "output_format": output_format, }, ) assert response.status_code == 200 - assert response.headers["content-type"] == "image/png" + assert response.headers["content-type"] == expected_media_type assert "attachment" in response.headers.get("content-disposition", "") - assert ".png" in response.headers.get("content-disposition", "") + assert f".{output_format}" in response.headers.get("content-disposition", "") img = Image.open(io.BytesIO(response.content)) - assert img.format == "PNG" + assert img.format == expected_pil_format + + +@pytest.mark.parametrize( + ("output_format", "image_count", "expected_media_type", "expected_pil_format"), + [ + ("jpeg", 1, "image/jpeg", "JPEG"), + ("webp", 2, "application/zip", "WEBP"), + ], +) +def test_multistage_image_file_response_uses_requested_format( + async_omni_test_client, + output_format, + image_count, + expected_media_type, + expected_pil_format, +): + engine = async_omni_test_client.app.state.engine_client + engine._images = [Image.new("RGB", (16, 16), color="green") for _ in range(image_count)] + + response = async_omni_test_client.post( + "/v1/images/generations", + json={ + "prompt": "a dog", + "n": image_count, + "response_format": "file", + "output_format": output_format, + }, + ) + + assert response.status_code == 200 + assert response.headers["content-type"] == expected_media_type + if image_count == 1: + assert f".{output_format}" in response.headers["content-disposition"] + assert Image.open(io.BytesIO(response.content)).format == expected_pil_format + return + + import zipfile + + with zipfile.ZipFile(io.BytesIO(response.content), "r") as archive: + filenames = archive.namelist() + assert len(filenames) == image_count + assert all(filename.endswith(f".{output_format}") for filename in filenames) + for filename in filenames: + assert Image.open(io.BytesIO(archive.read(filename))).format == expected_pil_format def make_test_image_bytes(size=(64, 64)) -> bytes: diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 6d088f65e..2e92c864b 100755 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -98,7 +98,6 @@ from vllm_omni.entrypoints.openai.duplex_capability import should_enable_duplex_ from vllm_omni.entrypoints.openai.errors import InvalidInputReferenceError from vllm_omni.entrypoints.openai.image_api_utils import ( SUPPORTED_LAYERED_RESOLUTIONS, - encode_image_base64, encode_image_base64_with_compression, parse_size, validate_layered_layers, @@ -1731,6 +1730,39 @@ async def show_available_models(raw_request: Request) -> JSONResponse: # Image generation API endpoints +def _build_image_generation_response( + *, + images: list[Image.Image], + request: ImageGenerationRequest, + stage_durations: Any, + peak_memory_mb: Any, +) -> ImageGenerationResponse | StreamingResponse: + """Encode generated images and apply the requested response format.""" + output_format = _choose_output_format(request.output_format or "png", None) + image_data = [ + ImageData( + b64_json=encode_image_base64_with_compression(image, format=output_format), + revised_prompt=None, + ) + for image in images + ] + response_kwargs: dict[str, Any] = { + "created": int(time.time()), + "data": image_data, + "output_format": output_format, + "metrics": { + "stage_durations": stage_durations or None, + "peak_memory_mb": float(peak_memory_mb) if peak_memory_mb else None, + }, + } + if request.size is not None: + response_kwargs["size"] = request.size + response = ImageGenerationResponse(**response_kwargs) + if request.response_format == ResponseFormat.FILE: + return response.stream_response() + return response + + @router.post( "/v1/images/generations", dependencies=[Depends(validate_json_request)], @@ -1832,15 +1864,11 @@ async def generate_images( content=generation_result.model_dump(), ) flat_images, stage_durations, peak_memory_mb, _ = generation_result - image_data = [ImageData(b64_json=encode_image_base64(img), revised_prompt=None) for img in flat_images] - - return ImageGenerationResponse( - created=int(time.time()), - data=image_data, - metrics={ - "stage_durations": stage_durations or None, - "peak_memory_mb": float(peak_memory_mb) if peak_memory_mb else None, - }, + return _build_image_generation_response( + images=flat_images, + request=request, + stage_durations=stage_durations, + peak_memory_mb=peak_memory_mb, ) # Build params - pass through user values directly @@ -1928,32 +1956,14 @@ async def generate_images( logger.debug(f"Successfully generated {len(images)} image(s)") - # Determine output format (default to png) - output_format = _choose_output_format(request.output_format or "png", None) - - # Encode images to base64 with the specified format - image_data = [ - ImageData(b64_json=encode_image_base64_with_compression(img, format=output_format), revised_prompt=None) - for img in images - ] - stage_durations = getattr(result, "stage_durations", None) peak_memory_mb = getattr(result, "peak_memory_mb", None) - response_kwargs = { - "created": int(time.time()), - "data": image_data, - "output_format": output_format, - "metrics": { - "stage_durations": stage_durations or None, - "peak_memory_mb": float(peak_memory_mb) if peak_memory_mb else None, - }, - } - if request.size: - response_kwargs["size"] = size_str - response = ImageGenerationResponse(**response_kwargs) - if request.response_format != ResponseFormat.FILE: - return response - return response.stream_response() + return _build_image_generation_response( + images=images, + request=request, + stage_durations=stage_durations, + peak_memory_mb=peak_memory_mb, + ) except (EngineGenerateError, EngineDeadError) as exc: return _create_engine_error_json_response(raw_request, exc) diff --git a/vllm_omni/entrypoints/openai/image_api_utils.py b/vllm_omni/entrypoints/openai/image_api_utils.py index 5bd79e936..ed759a080 100644 --- a/vllm_omni/entrypoints/openai/image_api_utils.py +++ b/vllm_omni/entrypoints/openai/image_api_utils.py @@ -101,13 +101,14 @@ def encode_image_base64_with_compression( """ buffer = io.BytesIO() image = prepare_image_for_output_format(image, format) + pil_format = "jpeg" if format == "jpg" else format save_kwargs = {} if format in ("jpg", "jpeg", "webp"): save_kwargs["quality"] = output_compression elif format == "png": save_kwargs["compress_level"] = max(0, min(9, 9 - output_compression // 11)) - image.save(buffer, format=format, **save_kwargs) + image.save(buffer, format=pil_format, **save_kwargs) buffer.seek(0) return base64.b64encode(buffer.read()).decode("utf-8") diff --git a/vllm_omni/entrypoints/openai/protocol/images.py b/vllm_omni/entrypoints/openai/protocol/images.py index e83c7b5c0..0a20ba1c8 100644 --- a/vllm_omni/entrypoints/openai/protocol/images.py +++ b/vllm_omni/entrypoints/openai/protocol/images.py @@ -21,6 +21,13 @@ from pydantic import BaseModel, Field, field_validator from vllm_omni.entrypoints.openai.image_api_utils import validate_layered_layers +_IMAGE_FILE_METADATA = { + "jpg": ("jpg", "image/jpeg"), + "jpeg": ("jpeg", "image/jpeg"), + "png": ("png", "image/png"), + "webp": ("webp", "image/webp"), +} + class ResponseFormat(str, Enum): """Image response format""" @@ -183,7 +190,7 @@ class ImageGenerationResponse(BaseModel): created: int = Field(..., description="Unix timestamp of when the generation completed") data: list[ImageData] = Field(..., description="Array of generated images") - output_format: str = Field(None, description="The output format of the image generation") + output_format: str | None = Field(None, description="The output format of the image generation") size: str = Field(None, description="The size of the image generated") cot_output: str | None = Field( None, @@ -201,12 +208,16 @@ class ImageGenerationResponse(BaseModel): status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail="No image data available for file response.", ) + extension, media_type = _IMAGE_FILE_METADATA.get( + (self.output_format or "png").lower(), + _IMAGE_FILE_METADATA["png"], + ) if len(self.data) == 1: image_bytes = base64.b64decode(self.data[0].b64_json) - filename = f"image_{uuid.uuid4().hex[:8]}.png" + filename = f"image_{uuid.uuid4().hex[:8]}.{extension}" return StreamingResponse( io.BytesIO(image_bytes), - media_type="image/png", + media_type=media_type, headers={ "Content-Disposition": f'attachment; filename="{filename}"', "Content-Length": str(len(image_bytes)), @@ -217,7 +228,7 @@ class ImageGenerationResponse(BaseModel): with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zf: for idx, item in enumerate(self.data): if item.b64_json: - zf.writestr(f"image_{idx}.png", base64.b64decode(item.b64_json)) + zf.writestr(f"image_{idx}.{extension}", base64.b64decode(item.b64_json)) zip_bytes = zip_buffer.getvalue() filename = f"images_{uuid.uuid4().hex[:8]}.zip" return StreamingResponse(