[Feature] add a new response_format is file type when image generate (#1673)

Signed-off-by: rongfu.leng <lenronfu@gmail.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>

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

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",

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)

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."""