diff --git a/docs/serving/image_generation_api.md b/docs/serving/image_generation_api.md index 747cef995..4702a19c1 100644 --- a/docs/serving/image_generation_api.md +++ b/docs/serving/image_generation_api.md @@ -32,6 +32,20 @@ curl -X POST http://localhost:8000/v1/images/generations \ }' | jq -r '.data[0].b64_json' | base64 -d > dragon.png ``` +**Using curl save to file:** + +```bash +curl -o dragon.png -X POST http://localhost:8000/v1/images/generations \ + -H "Content-Type: application/json" \ + -d '{ + "prompt": "a dragon laying over the spine of the Green Mountains of Vermont", + "size": "1024x1024", + "seed": 42, + "response_format":"file" + }' +``` + + **Using Python:** ```python @@ -57,6 +71,36 @@ img = Image.open(io.BytesIO(img_bytes)) img.save("cat.png") ``` +**Using Python save to file:** + +```python +import requests +import base64 +from PIL import Image +import io +import re + +response = requests.post( + "http://localhost:8000/v1/images/generations", + json={ + "prompt": "a black and white cat wearing a princess tiara", + "size": "1024x1024", + "num_inference_steps": 50, + "seed": 42, + "response_format":"file" + } +) + +# save to file +content_disposition = response.headers.get("Content-Disposition", "") +match = re.search(r'filename="?(.+)"?', content_disposition) +filename = match.group(1) if match else "save.png" +with open(filename, "wb") as f: + for chunk in response.iter_content(8192): + f.write(chunk) +print("saved:", filename) +``` + **Using OpenAI SDK:** ```python diff --git a/tests/entrypoints/openai_api/test_image_server.py b/tests/entrypoints/openai_api/test_image_server.py index bf3d8bd61..31a82102a 100644 --- a/tests/entrypoints/openai_api/test_image_server.py +++ b/tests/entrypoints/openai_api/test_image_server.py @@ -1301,6 +1301,58 @@ 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): + """Test response_format=file with n>1 returns ZIP archive""" + response = test_client.post( + "/v1/images/generations", + json={ + "prompt": "a dog", + "n": 3, + "response_format": "file", + }, + ) + + assert response.status_code == 200 + assert response.headers["content-type"] == "application/zip" + 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 + 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) + + # Verify each file is a valid PNG + for filename in files: + img_bytes = zf.read(filename) + img = Image.open(io.BytesIO(img_bytes)) + assert img.format == "PNG" + + +def test_image_file_response_format_single(test_client): + """Test response_format=file with n=1 returns a single image file.""" + response = test_client.post( + "/v1/images/generations", + json={ + "prompt": "a dog", + "n": 1, + "response_format": "file", + }, + ) + + assert response.status_code == 200 + assert response.headers["content-type"] == "image/png" + assert "attachment" in response.headers.get("content-disposition", "") + assert ".png" in response.headers.get("content-disposition", "") + + img = Image.open(io.BytesIO(response.content)) + assert img.format == "PNG" + + def make_test_image_bytes(size=(64, 64)) -> bytes: img = Image.new( "RGB", diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 405803d9c..11186a0e7 100755 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -110,6 +110,7 @@ from vllm_omni.entrypoints.openai.protocol.images import ( ImageData, ImageGenerationRequest, ImageGenerationResponse, + ResponseFormat, ) from vllm_omni.entrypoints.openai.protocol.videos import ( SecondStr, @@ -1706,6 +1707,7 @@ async def show_available_models(raw_request: Request) -> JSONResponse: @router.post( "/v1/images/generations", dependencies=[Depends(validate_json_request)], + response_model=None, responses={ HTTPStatus.OK.value: {"model": ImageGenerationResponse}, HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse}, @@ -1714,7 +1716,9 @@ async def show_available_models(raw_request: Request) -> JSONResponse: }, ) @with_cancellation -async def generate_images(request: ImageGenerationRequest, raw_request: Request): +async def generate_images( + request: ImageGenerationRequest, raw_request: Request +) -> ImageGenerationResponse | StreamingResponse: """Generate images from text prompts using diffusion models. OpenAI DALL-E compatible endpoint for text-to-image generation. @@ -1906,7 +1910,10 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request) } if request.size: response_kwargs["size"] = size_str - return ImageGenerationResponse(**response_kwargs) + response = ImageGenerationResponse(**response_kwargs) + if request.response_format != ResponseFormat.FILE: + return response + return response.stream_response() except (EngineGenerateError, EngineDeadError) as exc: return _create_engine_error_json_response(raw_request, exc) diff --git a/vllm_omni/entrypoints/openai/protocol/images.py b/vllm_omni/entrypoints/openai/protocol/images.py index 83338624a..78f63d387 100644 --- a/vllm_omni/entrypoints/openai/protocol/images.py +++ b/vllm_omni/entrypoints/openai/protocol/images.py @@ -7,9 +7,16 @@ This module provides Pydantic models that follow the OpenAI DALL-E API specifica for text-to-image generation, with vllm-omni specific extensions. """ +import base64 +import io +import uuid +import zipfile from enum import Enum +from http import HTTPStatus from typing import Any, Literal +from fastapi import HTTPException +from fastapi.responses import StreamingResponse from pydantic import BaseModel, Field, field_validator from vllm_omni.entrypoints.openai.image_api_utils import validate_layered_layers @@ -20,6 +27,7 @@ class ResponseFormat(str, Enum): B64_JSON = "b64_json" URL = "url" # Not implemented in PoC + FILE = "file" # file response class ImageGenerationRequest(BaseModel): @@ -73,9 +81,9 @@ class ImageGenerationRequest(BaseModel): @field_validator("response_format") @classmethod def validate_response_format(cls, v): - """Validate response format - only b64_json is supported.""" - if v is not None and v != ResponseFormat.B64_JSON: - raise ValueError(f"Only 'b64_json' response format is supported, got: {v}") + """Validate response format - only b64_json and file are supported.""" + if v is not None and v not in (ResponseFormat.B64_JSON, ResponseFormat.FILE): + raise ValueError(f"Only 'b64_json' or 'file' response format is supported, got: {v}") return v @field_validator("layers") @@ -183,6 +191,40 @@ class ImageGenerationResponse(BaseModel): "Only present for image editing (IT2I) with CoT-enabled models.", ) + def stream_response(self) -> StreamingResponse: + if not self.data or not self.data[0].b64_json: + raise HTTPException( + status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, + detail="No image data available for file response.", + ) + if len(self.data) == 1: + image_bytes = base64.b64decode(self.data[0].b64_json) + filename = f"image_{uuid.uuid4().hex[:8]}.png" + return StreamingResponse( + io.BytesIO(image_bytes), + media_type="image/png", + headers={ + "Content-Disposition": f'attachment; filename="{filename}"', + "Content-Length": str(len(image_bytes)), + }, + ) + else: + zip_buffer = io.BytesIO() + 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)) + zip_bytes = zip_buffer.getvalue() + filename = f"images_{uuid.uuid4().hex[:8]}.zip" + return StreamingResponse( + io.BytesIO(zip_bytes), + media_type="application/zip", + headers={ + "Content-Disposition": f'attachment; filename="{filename}"', + "Content-Length": str(len(zip_bytes)), + }, + ) + class ImageEditARDeltaChunk(BaseModel): """Streaming chunk carrying a text delta from the image-edit AR stage."""