[FEATURE] /v1/images/edit interface (#1101)

Signed-off-by: dengyunyang <584797741@qq.com>

docs/.nav.yml

@@ -7,6 +7,7 @@ nav:
   - Serving:
     - OpenAI-Compatible API:
       - Image Generation: serving/image_generation_api.md
+      - Image Edit: serving/image_edit_api.md
   - Examples:
     - examples/README.md
     - Offline Inference:

docs/serving/image_edit_api.md

@@ -0,0 +1,205 @@
+# Image Edit API
+
+vLLM-Omni provides an OpenAI DALL-E compatible API for image edit using diffusion models.
+
+Each server instance runs a single model (specified at startup via `vllm serve <model> --omni`).
+
+## Quick Start
+
+### Start the Server
+
+For example...
+
+```bash
+# Qwen-Image
+vllm serve Qwen/Qwen-Image-Edit-2511 --omni --port 8000
+
+
+### Generate Images
+
+**Using curl:**
+
+```bash
+curl -s -D >(grep -i x-request-id >&2) \
+  -o >(jq -r '.data[0].b64_json' | base64 --decode > gift-basket.png) \
+  -X POST "http://localhost:8000/v1/images/edits" \
+  -F "model=xxx" \
+  -F "image=@./xx.png" \
+  -F "prompt='this bear is wearing sportwear. holding a basketball, and bending one leg.'" \
+  -F "size=1024x1024" \
+  -F "output_format=png"
+```
+
+
+**Using OpenAI SDK:**
+
+```python
+import base64
+from openai import OpenAI
+from pathlib import Path
+client = OpenAI(
+    api_key="None",
+    base_url="http://localhost:8000/v1"
+)
+
+input_image_url = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/omni-assets/qwen-bear.png"
+
+result = client.images.edit(
+    image=[],
+    model="Qwen-Image-Edit-2511",
+    prompt="Change the bears in the two input images into walking together.",
+    size='512x512',
+    stream=False,
+    output_format='jpeg',
+    # url格式
+    extra_body={
+        "url": [input_image_url1,input_image_url],
+        "num_inference_steps": 50,
+        "guidance_scale": 1,
+        "seed": 777,
+    }
+)
+
+image_base64 = result.data[0].b64_json
+image_bytes = base64.b64decode(image_base64)
+
+# Save the image to a file
+with open("edit_out_http.jpeg", "wb") as f:
+    f.write(image_bytes)
+```
+
+## API Reference
+
+### Endpoint
+
+```
+POST /v1/images/edits
+Content-Type: multipart/form-data
+```
+
+### Request Parameters
+
+#### OpenAI Standard Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `prompt` | string | **required** | A text description of the desired image |
+| `model` | string | server's model | Model to use (optional, should match server if specified) |
+| `image` | string or array | **required** | The image(s) to edit. |
+| `n` | integer | 1 | Number of images to generate (1-10) |
+| `size` | string | "auto" | Image dimensions in WxH format (e.g., "1024x1024", "512x512"), when set to auto, it decide size from first input image. |
+| `response_format` | string | "b64_json" | Response format (only "b64_json" supported) |
+| `user` | string | null | User identifier for tracking |
+| `output_format` | string | "png" | The format in which the generated images are returned. Must be one of "png", "jpg", "jpeg", "webp". |
+| `output_compression` | integer | 100 | The compression level (0-100%) for the generated images. |
+| `background` | string or null | "auto" | Allows to set transparency for the background of the generated image(s).
+
+#### vllm-omni Extension Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `url` | string or array | None | The image(s) to edit. |
+| `negative_prompt` | string | null | Text describing what to avoid in the image |
+| `num_inference_steps` | integer | model defaults | Number of diffusion steps |
+| `guidance_scale` | float | model defaults | Classifier-free guidance scale (typically 0.0-20.0) |
+| `true_cfg_scale` | float | model defaults | True CFG scale (model-specific parameter, may be ignored if not supported) |
+| `seed` | integer | null | Random seed for reproducibility |
+
+### Response Format
+
+```json
+{
+  "created": 1701234567,
+  "data": [
+    {
+      "b64_json": "<base64-encoded PNG>",
+      "url": null,
+      "revised_prompt": null
+    }
+  ],
+  "output_format": null,
+  "size": null,
+}
+```
+
+## Examples
+
+### Multiple Images input
+
+```bash
+curl -s -D >(grep -i x-request-id >&2) \
+  -o >(jq -r '.data[0].b64_json' | base64 --decode > gift-basket.png) \
+  -X POST "http://localhost:8000/v1/images/edits" \
+  -F "model=xxx" \
+  -F "image=@xx.png" \
+  -F "image=@xx.png"
+  -F "prompt='this bear is wearing sportwear. holding a basketball, and bending one leg.'" \
+  -F "size=1024x1024" \
+  -F "output_format=png"
+```
+
+
+## Parameter Handling
+
+The API passes parameters directly to the diffusion pipeline without model-specific transformation:
+
+- **Default values**: When parameters are not specified, the underlying model uses its own defaults
+- **Pass-through design**: User-provided values are forwarded directly to the diffusion engine
+- **Minimal validation**: Only basic type checking and range validation at the API level
+
+### Parameter Compatibility
+
+The API passes parameters directly to the diffusion pipeline without model-specific validation.
+
+- Unsupported parameters may be silently ignored by the model
+- Incompatible values will result in errors from the underlying pipeline
+- Recommended values vary by model - consult model documentation
+
+**Best Practice:** Start with the model's recommended parameters, then adjust based on your needs.
+
+## Error Responses
+
+### 400 Bad Request
+
+Invalid parameters (e.g., model mismatch):
+
+```json
+{
+  "detail": "Invalid size format: '1024x'. Expected format: 'WIDTHxHEIGHT' (e.g., '1024x1024')."
+}
+```
+
+### 422 Unprocessable Entity
+
+Validation errors (missing required fields):
+
+```json
+{
+  "detail": "Field 'image' or 'url' is required"
+}
+```
+
+## Troubleshooting
+
+### Server Not Running
+
+```bash
+# Check if server is responding
+curl -X http://localhost:8000/v1/images/edit \
+  -F "prompt='test'"
+```
+
+### Out of Memory
+
+If you encounter OOM errors:
+1. Reduce image size: `"size": "512x512"`
+2. Reduce inference steps: `"num_inference_steps": 25`
+
+## Development
+
+Enable debug logging to see prompts and generation details:
+
+```bash
+vllm serve Qwen/Qwen-Image-Edit-2511 --omni \
+  --uvicorn-log-level debug
+```

tests/entrypoints/openai_api/test_image_server.py

@@ -9,6 +9,7 @@ OpenAI-compatible async text-to-image generation API endpoints in api_server.py.
 
 import base64
 import io
+from argparse import Namespace
 from unittest.mock import AsyncMock, Mock
 
 import pytest
@@ -115,9 +116,11 @@ class FakeAsyncOmni:
         self.stage_list = ["llm", "diffusion"]
         self.default_sampling_params_list = [SamplingParams(temperature=0.1), OmniDiffusionSamplingParams()]
         self.captured_sampling_params_list = None
+        self.captured_prompt = None
 
     async def generate(self, prompt, request_id, sampling_params_list):
         self.captured_sampling_params_list = sampling_params_list
+        self.captured_prompt = prompt
         images = [Image.new("RGB", (64, 64), color="green")]
         yield MockGenerationResult(images)
 
@@ -133,6 +136,8 @@ def mock_async_diffusion():
         # Return n PIL images wrapped in result object
         print("!!!!!!!!!!!!!!!!!!!!! kwargs", kwargs)
         n = kwargs["sampling_params_list"][0].num_outputs_per_prompt
+        mock.captured_sampling_params_list = kwargs["sampling_params_list"]
+        mock.captured_prompt = kwargs["prompt"]
         images = [Image.new("RGB", (64, 64), color="blue") for _ in range(n)]
         return MockGenerationResult(images)
 
@@ -155,6 +160,10 @@ def test_client(mock_async_diffusion):
     app.state.diffusion_engine = mock_async_diffusion  # Also set for health endpoint
     app.state.stage_configs = [{"stage_type": "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
+    )
 
     return TestClient(app)
 
@@ -171,7 +180,10 @@ def async_omni_test_client():
 
     app.state.engine_client = FakeAsyncOmni()
     app.state.stage_configs = [{"stage_type": "llm"}, {"stage_type": "diffusion"}]
-
+    app.state.args = Namespace(
+        default_sampling_params='{"1": {"num_inference_steps":4, "guidance_scale":7.5}}',
+        max_generated_image_size=4096,  # 64*64
+    )
     return TestClient(app)
 
 
@@ -526,3 +538,279 @@ def test_model_field_omitted_works(test_client):
         },
     )
     assert response.status_code == 200
+
+
+def make_test_image_bytes(size=(64, 64)) -> bytes:
+    img = Image.new(
+        "RGB",
+        size,
+    )
+    buf = io.BytesIO()
+    img.save(buf, format="PNG")
+    return buf.getvalue()
+
+
+def test_image_edit_images_processing(async_omni_test_client):
+    img_bytes_1 = make_test_image_bytes((16, 16))
+    img_bytes_2 = make_test_image_bytes((32, 32))
+
+    # uploadfile with image key
+    response = async_omni_test_client.post(
+        "/v1/images/edits",
+        files=[
+            ("image", img_bytes_1),
+            ("image", img_bytes_2),
+        ],
+        data={"prompt": "hello world."},
+    )
+    assert response.status_code == 200
+    engine = async_omni_test_client.app.state.engine_client
+    captured_prompt = engine.captured_prompt
+    processed_images = captured_prompt["multi_modal_data"]["image"]
+    assert len(processed_images) == 2
+    assert isinstance(processed_images[0], Image.Image)
+    assert isinstance(processed_images[1], Image.Image)
+    assert processed_images[0].size == (16, 16)
+    assert processed_images[1].size == (32, 32)
+
+    # uploadfile with image[] key
+    response = async_omni_test_client.post(
+        "/v1/images/edits",
+        files=[
+            ("image[]", img_bytes_2),
+            ("image[]", img_bytes_1),
+        ],
+        data={"prompt": "hello world."},
+    )
+
+    assert response.status_code == 200
+    engine = async_omni_test_client.app.state.engine_client
+    captured_prompt = engine.captured_prompt
+    processed_images = captured_prompt["multi_modal_data"]["image"]
+    assert len(processed_images) == 2
+    assert isinstance(processed_images[0], Image.Image)
+    assert isinstance(processed_images[1], Image.Image)
+    assert processed_images[0].size == (32, 32)
+    assert processed_images[1].size == (16, 16)
+
+    # base64 url
+    buf1 = io.BytesIO()
+    img1 = Image.new("RGB", (16, 16))
+    img1.save(buf1, format="PNG")
+    b64_1 = "data:image/png;base64," + base64.b64encode(buf1.getvalue()).decode()
+
+    buf2 = io.BytesIO()
+    img2 = Image.new("RGB", (24, 24))
+    img2.save(buf2, format="PNG")
+    b64_2 = "data:image/png;base64," + base64.b64encode(buf2.getvalue()).decode()
+
+    response = async_omni_test_client.post(
+        "/v1/images/edits",
+        data={
+            "prompt": "hello from base64",
+            "url": [b64_1, b64_2],
+        },
+    )
+    assert response.status_code == 200
+    processed_images = engine.captured_prompt["multi_modal_data"]["image"]
+    assert len(processed_images) == 2
+    assert isinstance(processed_images[0], Image.Image)
+    assert isinstance(processed_images[1], Image.Image)
+    assert processed_images[0].size == (16, 16)
+    assert processed_images[1].size == (24, 24)
+
+
+def test_image_edit_parameter_pass(async_omni_test_client):
+    img_bytes_1 = make_test_image_bytes((16, 16))
+
+    # uploadfile with image key
+    response = async_omni_test_client.post(
+        "/v1/images/edits",
+        files=[("image", img_bytes_1)],
+        data={
+            "prompt": "hello world.",
+            "size": "16x24",
+            "output_format": "jpeg",
+            "num_inference_steps": 20,
+            "guidance_scale": 8.0,
+            "seed": 1234,
+            "negative_prompt": "negative",
+            "n": 2,
+        },
+    )
+    assert response.status_code == 200
+    engine = async_omni_test_client.app.state.engine_client
+    captured_prompt = engine.captured_prompt
+    captured_sampling_params = engine.captured_sampling_params_list[-1]
+
+    assert captured_prompt["prompt"] == "hello world."
+    assert captured_prompt["negative_prompt"] == "negative"
+    assert captured_sampling_params.num_inference_steps == 20
+    assert captured_sampling_params.guidance_scale == 8.0
+    assert captured_sampling_params.seed == 1234
+    assert captured_sampling_params.num_outputs_per_prompt == 2
+    assert captured_sampling_params.width == 16
+    assert captured_sampling_params.height == 24
+
+    data = response.json()
+    # All images should be valid
+    for img_data in data["data"]:
+        assert "b64_json" in img_data
+        img_bytes = base64.b64decode(img_data["b64_json"])
+        img = Image.open(io.BytesIO(img_bytes))
+        assert img.format.lower() == "jpeg"
+        assert data["output_format"] == "jpeg"
+        assert data["size"] == "16x24"
+
+
+def test_image_edit_parameter_default(async_omni_test_client):
+    img_bytes_1 = make_test_image_bytes((24, 16))
+
+    # uploadfile with image key
+    response = async_omni_test_client.post(
+        "/v1/images/edits",
+        files=[("image", img_bytes_1)],
+        data={
+            "prompt": "hello world.",
+            "size": "auto",
+        },
+    )
+    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.width == 24
+    assert captured_sampling_params.height == 16
+    assert captured_sampling_params.num_outputs_per_prompt == 1
+    assert captured_sampling_params.num_inference_steps == 4
+    assert captured_sampling_params.guidance_scale == 7.5
+
+    response = async_omni_test_client.post(
+        "/v1/images/edits",
+        files=[("image", img_bytes_1)],
+        data={
+            "prompt": "hello world.",
+            "size": "96x96",
+        },
+    )
+    assert response.status_code == 400
+
+
+def test_image_edit_parameter_default_single_stage(test_client):
+    img_bytes_1 = make_test_image_bytes((24, 16))
+
+    # uploadfile with image key
+    response = test_client.post(
+        "/v1/images/edits",
+        files=[("image", img_bytes_1)],
+        data={
+            "prompt": "hello world.",
+        },
+    )
+    assert response.status_code == 200
+    engine = test_client.app.state.engine_client
+    captured_sampling_params = engine.captured_sampling_params_list[0]
+
+    assert captured_sampling_params.width == 24
+    assert captured_sampling_params.height == 16
+    assert captured_sampling_params.num_outputs_per_prompt == 1
+    assert captured_sampling_params.num_inference_steps == 4
+    assert captured_sampling_params.guidance_scale == 7.5
+
+    response = test_client.post(
+        "/v1/images/edits",
+        files=[("image", img_bytes_1)],
+        data={
+            "prompt": "hello world.",
+            "size": "96x96",
+        },
+    )
+    assert response.status_code == 400
+
+
+def test_image_edit_compression_jpeg(test_client):
+    img_bytes_1 = make_test_image_bytes((16, 16))
+    # uploadfile with image key
+    response = test_client.post(
+        "/v1/images/edits",
+        files=[("image", img_bytes_1)],
+        data={"prompt": "hello world.", "output_format": "jpeg", "output_compression": 100},
+    )
+    assert response.status_code == 200
+    data = response.json()
+    img_bytes_100 = base64.b64decode(data["data"][0]["b64_json"])
+    img = Image.open(io.BytesIO(img_bytes_100))
+    assert img.format.lower() == "jpeg"
+
+    response = test_client.post(
+        "/v1/images/edits",
+        files=[("image", img_bytes_1)],
+        data={
+            "prompt": "hello world.",
+            "output_format": "jpeg",
+            "output_compression": 50,
+        },
+    )
+    assert response.status_code == 200
+    data = response.json()
+    img_bytes_50 = base64.b64decode(data["data"][0]["b64_json"])
+
+    response = test_client.post(
+        "/v1/images/edits",
+        files=[("image", img_bytes_1)],
+        data={
+            "prompt": "hello world.",
+            "output_format": "jpeg",
+            "output_compression": 10,
+        },
+    )
+    assert response.status_code == 200
+    data = response.json()
+    img_bytes_10 = base64.b64decode(data["data"][0]["b64_json"])
+
+    assert len(img_bytes_10) < len(img_bytes_50)
+    assert len(img_bytes_50) < len(img_bytes_100)
+
+
+def test_image_edit_compression_png(async_omni_test_client):
+    img_bytes_1 = make_test_image_bytes((16, 16))
+    # uploadfile with image key
+    response = async_omni_test_client.post(
+        "/v1/images/edits",
+        files=[("image", img_bytes_1)],
+        data={"prompt": "hello world.", "output_format": "PNG", "output_compression": 100},
+    )
+    assert response.status_code == 200
+    data = response.json()
+    img_bytes_100 = base64.b64decode(data["data"][0]["b64_json"])
+    img = Image.open(io.BytesIO(img_bytes_100))
+    assert img.format.lower() == "png"
+
+    response = async_omni_test_client.post(
+        "/v1/images/edits",
+        files=[("image", img_bytes_1)],
+        data={
+            "prompt": "hello world.",
+            "output_format": "PNG",
+            "output_compression": 50,
+        },
+    )
+    assert response.status_code == 200
+    data = response.json()
+    img_bytes_50 = base64.b64decode(data["data"][0]["b64_json"])
+
+    response = async_omni_test_client.post(
+        "/v1/images/edits",
+        files=[("image", img_bytes_1)],
+        data={
+            "prompt": "hello world.",
+            "output_format": "PNG",
+            "output_compression": 10,
+        },
+    )
+    assert response.status_code == 200
+    data = response.json()
+    img_bytes_10 = base64.b64decode(data["data"][0]["b64_json"])
+
+    assert len(img_bytes_10) < len(img_bytes_50)
+    assert len(img_bytes_50) < len(img_bytes_100)

vllm_omni/entrypoints/cli/serve.py

@@ -232,6 +232,22 @@ class OmniServeCommand(CLISubcommand):
             help="Number of devices for CFG parallel computation for diffusion models. "
             "Equivalent to setting DiffusionParallelConfig.cfg_parallel_size.",
         )
+
+        # Default sampling parameters
+        omni_config_group.add_argument(
+            "--default-sampling-params",
+            type=str,
+            help="Json str for Default sampling parameters, \n"
+            'Structure: {"<stage_id>": {<sampling_param>: value, ...}, ...}\n'
+            'e.g., \'{"0": {"num_inference_steps":50, "guidance_scale":1}}\'. '
+            "Currently only supports diffusion models.",
+        )
+        # Diffusion model mixed precision
+        omni_config_group.add_argument(
+            "--max-generated-image-size",
+            type=float,
+            help="The max size of generate image (height * width).",
+        )
         return serve_parser
 
 

vllm_omni/entrypoints/openai/api_server.py

@@ -1,5 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import base64
+import io
+import json
 import multiprocessing
 import multiprocessing.forkserver as forkserver
 import os
@@ -12,11 +15,13 @@ from argparse import Namespace
 from collections.abc import AsyncIterator
 from contextlib import asynccontextmanager
 from http import HTTPStatus
-from typing import Any, cast
+from typing import Annotated, Any, cast
 
+import httpx
 import vllm.envs as envs
-from fastapi import Depends, HTTPException, Request
+from fastapi import Depends, File, Form, HTTPException, Request, UploadFile
 from fastapi.responses import JSONResponse, StreamingResponse
+from PIL import Image
 from starlette.datastructures import State
 from starlette.routing import Route
 from vllm import SamplingParams
@@ -890,6 +895,251 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request)
     Raises:
         HTTPException: For validation errors, missing engine, or generation failures
     """
+    # Get engine client (AsyncOmni) from app state
+    engine_client, model_name, stage_types = _get_engine_and_model(raw_request)
+
+    # Validate model field (warn if mismatch, don't error)
+    if request.model is not None and request.model != model_name:
+        logger.warning(
+            f"Model mismatch: request specifies '{request.model}' but "
+            f"server is running '{model_name}'. Using server model."
+        )
+
+    try:
+        # Build params - pass through user values directly
+        prompt: OmniTextPrompt = {"prompt": request.prompt}
+        if request.negative_prompt is not None:
+            prompt["negative_prompt"] = request.negative_prompt
+        gen_params = OmniDiffusionSamplingParams(num_outputs_per_prompt=request.n)
+
+        # Parse per-request LoRA (compatible with chat's extra_body.lora shape).
+        lora_request, lora_scale = _parse_lora_request(request.lora)
+        _update_if_not_none(gen_params, "lora_request", lora_request)
+        _update_if_not_none(gen_params, "lora_scale", lora_scale)
+
+        # Parse and add size if provided
+        width, height = None, None
+        if request.size:
+            width, height = parse_size(request.size)
+            size_str = f"{width}x{height}"
+        else:
+            size_str = "model default"
+        _update_if_not_none(gen_params, "width", width)
+        _update_if_not_none(gen_params, "height", height)
+
+        # 3.3 Add optional parameters ONLY if provided
+        _update_if_not_none(gen_params, "num_inference_steps", request.num_inference_steps)
+        _update_if_not_none(gen_params, "guidance_scale", request.guidance_scale)
+        _update_if_not_none(gen_params, "true_cfg_scale", request.true_cfg_scale)
+        # If seed is not provided, generate a random one to ensure
+        # a proper generator is initialized in the backend.
+        # This fixes issues where using the default global generator
+        # might produce blurry images in some environments.
+        _update_if_not_none(gen_params, "seed", random.randint(0, 2**32 - 1) if request.seed is None else request.seed)
+
+        request_id = f"img_gen_{uuid.uuid4().hex}"
+
+        logger.info(f"Generating {request.n} image(s) {size_str}")
+
+        # Generate images using AsyncOmni (multi-stage mode)
+        result = await _generate_with_async_omni(
+            engine_client=engine_client,
+            gen_params=gen_params,
+            stage_types=stage_types,
+            prompt=prompt,
+            request_id=request_id,
+        )
+
+        if result is None:
+            raise HTTPException(
+                status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value,
+                detail="No output generated from multi-stage pipeline.",
+            )
+
+        # Extract images from result
+        images = _extract_images_from_result(result)
+
+        logger.info(f"Successfully generated {len(images)} image(s)")
+
+        # Encode images to base64
+        image_data = [ImageData(b64_json=encode_image_base64(img), revised_prompt=None) for img in images]
+
+        return ImageGenerationResponse(
+            created=int(time.time()),
+            data=image_data,
+        )
+
+    except HTTPException:
+        # Re-raise HTTPExceptions as-is
+        raise
+    except ValueError as e:
+        logger.error(f"Validation error: {e}")
+        raise HTTPException(status_code=HTTPStatus.BAD_REQUEST.value, detail=str(e))
+    except Exception as e:
+        logger.exception(f"Image generation failed: {e}")
+        raise HTTPException(
+            status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=f"Image generation failed: {str(e)}"
+        )
+
+
+@router.post(
+    "/v1/images/edits",
+    responses={
+        HTTPStatus.OK.value: {"model": ImageGenerationResponse},
+        HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
+        HTTPStatus.SERVICE_UNAVAILABLE.value: {"model": ErrorResponse},
+        HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
+    },
+)
+async def edit_images(
+    raw_request: Request,
+    image: list[UploadFile] | None = File(None),
+    image_array: list[UploadFile] | None = File(None, alias="image[]"),
+    url: list[str] | None = Form(None),
+    url_array: list[str] | None = Form(None, alias="url[]"),
+    prompt: str = Form(...),
+    model: str = Form(None),
+    n: int = Form(1),
+    size: str = Form("auto"),
+    response_format: str = Form("b64_json"),
+    output_format: str | None = Form("png"),
+    background: str | None = Form("auto"),
+    output_compression: Annotated[int, Form(ge=0, le=100)] = 100,
+    user: str | None = Form(None),  # unused now
+    # vllm-omni extensions for diffusion control
+    negative_prompt: str | None = Form(None),
+    num_inference_steps: int | None = Form(None),
+    guidance_scale: float | None = Form(None),
+    true_cfg_scale: float | None = Form(None),
+    seed: int | None = Form(None),
+    # vllm-omni extension for per-request LoRA.
+    lora: str | None = Form(None),  # Json string
+) -> ImageGenerationResponse:
+    """
+    OpenAI-compatible image edit endpoint.
+    """
+    # 1. get engine and model
+    engine_client, model_name, stage_types = _get_engine_and_model(raw_request)
+    if model is not None and model != model_name:
+        logger.warning(
+            f"Model mismatch: request specifies '{model}' but server is running '{model_name}'. Using server model."
+        )
+    # 2. get output format & compression
+    output_format = _choose_output_format(output_format, background)
+    if response_format != "b64_json":
+        raise HTTPException(
+            status_code=HTTPStatus.BAD_REQUEST.value,
+            detail="Only response_format 'b64_json' is supported now.",
+        )
+    try:
+        # 2. Build prompt & images params
+        prompt: OmniTextPrompt = {"prompt": prompt}
+        if negative_prompt is not None:
+            prompt["negative_prompt"] = negative_prompt
+        input_images_list = []
+        images = image or image_array
+        urls = url or url_array
+        if images:
+            input_images_list.extend(images)
+        if urls:
+            input_images_list.extend(urls)
+        if not input_images_list:
+            raise HTTPException(status_code=422, detail="Field 'image' or 'url' is required")
+        pil_images = await _load_input_images(input_images_list)
+        prompt["multi_modal_data"] = {}
+        prompt["multi_modal_data"]["image"] = pil_images
+
+        # 3 Build sample params
+        gen_params = OmniDiffusionSamplingParams()
+        # 3.0 Init with system default values
+        app_state_args = getattr(raw_request.app.state, "args", None)
+        default_sample_param = getattr(app_state_args, "default_sampling_params", None)
+        # Currently only have one diffusion stage
+        diffusion_stage_id = [i for i, t in enumerate(stage_types) if t == "diffusion"][0]
+        apply_stage_default_sampling_params(
+            default_sample_param,
+            gen_params,
+            str(diffusion_stage_id),
+        )
+        _update_if_not_none(gen_params, "num_outputs_per_prompt", n)
+        # 3.1 Parse per-request LoRA (compatible with chat's extra_body.lora shape).
+        lora_dict = _get_lora_from_json_str(lora)
+        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
+        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
+        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}"
+        _update_if_not_none(gen_params, "width", width)
+        _update_if_not_none(gen_params, "height", height)
+
+        # 3.3 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)
+        # If seed is not provided, generate a random one to ensure
+        # a proper generator is initialized in the backend.
+        # This fixes issues where using the default global generator
+        # might produce blurry images in some environments.
+        _update_if_not_none(gen_params, "seed", seed or random.randint(0, 2**32 - 1))
+
+        # 4. Generate images using AsyncOmni (multi-stage mode)
+        request_id = f"img_edit_{int(time.time())}"
+        logger.info(f"Generating {n} image(s) {size_str}")
+        result = await _generate_with_async_omni(
+            engine_client=engine_client,
+            gen_params=gen_params,
+            stage_types=stage_types,
+            prompt=prompt,
+            request_id=request_id,
+        )
+
+        # 5. Extract images from result
+        images = _extract_images_from_result(result)
+        logger.info(f"Successfully generated {len(images)} image(s)")
+
+        # Encode images to base64
+        image_data = [
+            ImageData(
+                b64_json=_encode_image_base64_with_compression(
+                    img, format=output_format, output_compression=output_compression
+                ),
+                revised_prompt=None,
+            )
+            for img in images
+        ]
+
+        return ImageGenerationResponse(
+            created=int(time.time()),
+            data=image_data,
+            output_format=output_format,
+            size=size_str,
+        )
+
+    except HTTPException:
+        # Re-raise HTTPExceptions as-is
+        raise
+    except ValueError as e:
+        logger.error(f"Validation error: {e}")
+        raise HTTPException(status_code=HTTPStatus.BAD_REQUEST.value, detail=str(e))
+    except Exception as e:
+        logger.exception(f"Image edit failed: {e}")
+        raise HTTPException(status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=f"Image edit failed: {str(e)}")
+
+
+def _get_engine_and_model(raw_request: Request):
     # Get engine client (AsyncOmni) from app state
     engine_client: EngineClient | AsyncOmni | None = getattr(raw_request.app.state, "engine_client", None)
     if engine_client is None or not hasattr(engine_client, "stage_list"):
@@ -942,157 +1192,229 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request)
     else:
         model_name = "unknown"
 
-    # Validate model field (warn if mismatch, don't error)
-    if request.model is not None and request.model != model_name:
-        logger.warning(
-            f"Model mismatch: request specifies '{request.model}' but "
-            f"server is running '{model_name}'. Using server model."
-        )
+    return engine_client, model_name, stage_types
 
+
+def _get_lora_from_json_str(lora_body):
+    if lora_body is None:
+        return None
     try:
-        # Build params - pass through user values directly
-        prompt: OmniTextPrompt = {"prompt": request.prompt}
-        gen_params = OmniDiffusionSamplingParams(num_outputs_per_prompt=request.n)
+        lora_dict = json.loads(lora_body)
+    except json.JSONDecodeError:
+        raise HTTPException(status_code=400, detail="Invalid LoRA JSON string")
 
-        # Parse per-request LoRA (compatible with chat's extra_body.lora shape).
-        if request.lora is not None:
-            if not isinstance(request.lora, dict):
-                raise HTTPException(
-                    status_code=HTTPStatus.BAD_REQUEST.value,
-                    detail="Invalid lora field: expected an object.",
-                )
-            lora_body = request.lora
-            lora_name = lora_body.get("name") or lora_body.get("lora_name") or lora_body.get("adapter")
-            lora_path = (
-                lora_body.get("local_path")
-                or lora_body.get("path")
-                or lora_body.get("lora_path")
-                or lora_body.get("lora_local_path")
-            )
-            lora_scale = lora_body.get("scale")
-            if lora_scale is None:
-                lora_scale = lora_body.get("lora_scale")
-            lora_int_id = lora_body.get("int_id")
-            if lora_int_id is None:
-                lora_int_id = lora_body.get("lora_int_id")
-            if lora_int_id is None and lora_path:
-                lora_int_id = stable_lora_int_id(str(lora_path))
+    if not isinstance(lora_dict, dict):
+        raise HTTPException(status_code=400, detail="LoRA must be a JSON object")
 
-            if not lora_name or not lora_path:
-                raise HTTPException(
-                    status_code=HTTPStatus.BAD_REQUEST.value,
-                    detail="Invalid lora object: both name and path are required.",
-                )
+    return lora_dict
 
-            gen_params.lora_request = LoRARequest(str(lora_name), int(lora_int_id), str(lora_path))
-            if lora_scale is not None:
-                gen_params.lora_scale = float(lora_scale)
 
-        # Parse and add size if provided
-        if request.size:
-            width, height = parse_size(request.size)
-            gen_params.height = height
-            gen_params.width = width
-            size_str = f"{width}x{height}"
-        else:
-            size_str = "model default"
-
-        # Add optional parameters ONLY if provided
-        if request.num_inference_steps is not None:
-            gen_params.num_inference_steps = request.num_inference_steps
-        if request.negative_prompt is not None:
-            prompt["negative_prompt"] = request.negative_prompt
-        if request.guidance_scale is not None:
-            gen_params.guidance_scale = request.guidance_scale
-        if request.true_cfg_scale is not None:
-            gen_params.true_cfg_scale = request.true_cfg_scale
-        if request.seed is not None:
-            gen_params.seed = request.seed
-        else:
-            # If seed is not provided, generate a random one to ensure
-            # a proper generator is initialized in the backend.
-            # This fixes issues where using the default global generator
-            # might produce blurry images in some environments.
-            gen_params.seed = random.randint(0, 2**32 - 1)
-
-        request_id = f"img_gen_{uuid.uuid4().hex}"
-
-        logger.info(f"Generating {request.n} image(s) {size_str}")
-
-        # Generate images using AsyncOmni (multi-stage mode)
-        engine_client = cast(AsyncOmni, engine_client)
-        result = None
-        stage_list = getattr(engine_client, "stage_list", None)
-        if isinstance(stage_list, list):
-            default_params_list: list[OmniSamplingParams] | None = getattr(
-                engine_client, "default_sampling_params_list", None
-            )
-            if not isinstance(default_params_list, list):
-                default_params_list = [
-                    OmniDiffusionSamplingParams() if st == "diffusion" else SamplingParams() for st in stage_types
-                ]
-            else:
-                default_params_list = list(default_params_list)
-            if len(default_params_list) != len(stage_types):
-                default_params_list = (
-                    default_params_list
-                    + [OmniDiffusionSamplingParams() if st == "diffusion" else SamplingParams() for st in stage_types]
-                )[: len(stage_types)]
-
-            sampling_params_list: list[OmniSamplingParams] = []
-            for idx, stage_type in enumerate(stage_types):
-                if stage_type == "diffusion":
-                    sampling_params_list.append(gen_params)
-                else:
-                    base_params = default_params_list[idx]
-                    sampling_params_list.append(base_params)
-
-            async for output in engine_client.generate(
-                prompt=prompt,
-                request_id=request_id,
-                sampling_params_list=sampling_params_list,
-            ):
-                result = output
-        else:
-            result = await engine_client.generate(
-                prompt=prompt, request_id=request_id, sampling_params_list=[gen_params]
-            )
-
-        if result is None:
+def _parse_lora_request(lora_body: dict[str, Any]):
+    if lora_body is not None:
+        if not isinstance(lora_body, dict):
             raise HTTPException(
-                status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value,
-                detail="No output generated from multi-stage pipeline.",
+                status_code=HTTPStatus.BAD_REQUEST.value,
+                detail="Invalid lora field: expected an object.",
+            )
+        lora_name = lora_body.get("name") or lora_body.get("lora_name") or lora_body.get("adapter")
+        lora_path = (
+            lora_body.get("local_path")
+            or lora_body.get("path")
+            or lora_body.get("lora_path")
+            or lora_body.get("lora_local_path")
+        )
+        lora_scale = lora_body.get("scale")
+        if lora_scale is None:
+            lora_scale = lora_body.get("lora_scale")
+        lora_int_id = lora_body.get("int_id")
+        if lora_int_id is None:
+            lora_int_id = lora_body.get("lora_int_id")
+        if lora_int_id is None and lora_path:
+            lora_int_id = stable_lora_int_id(str(lora_path))
+
+        if not lora_name or not lora_path:
+            raise HTTPException(
+                status_code=HTTPStatus.BAD_REQUEST.value,
+                detail="Invalid lora object: both name and path are required.",
             )
 
-        # Extract images from result
-        images = []
-        if hasattr(result, "images") and result.images:
-            images = result.images
-        elif hasattr(result, "request_output"):
-            request_output = result.request_output
-            if isinstance(request_output, dict) and request_output.get("images"):
-                images = request_output["images"]
-            elif hasattr(request_output, "images") and request_output.images:
-                images = request_output.images
+        return LoRARequest(str(lora_name), int(lora_int_id), str(lora_path)), lora_scale
+    return None, None
 
-        logger.info(f"Successfully generated {len(images)} image(s)")
 
-        # Encode images to base64
-        image_data = [ImageData(b64_json=encode_image_base64(img), revised_prompt=None) for img in images]
+async def _generate_with_async_omni(
+    engine_client: AsyncOmni | Any,
+    gen_params: Any,
+    stage_types: list[str],
+    **kwargs,
+):
+    engine_client = cast(AsyncOmni, engine_client)
+    result = None
+    stage_list = getattr(engine_client, "stage_list", None)
+    if isinstance(stage_list, list):
+        default_params_list: list[OmniSamplingParams] | None = getattr(
+            engine_client, "default_sampling_params_list", None
+        )
+        if not isinstance(default_params_list, list):
+            default_params_list = [
+                OmniDiffusionSamplingParams() if st == "diffusion" else SamplingParams() for st in stage_types
+            ]
+        else:
+            default_params_list = list(default_params_list)
+        if len(default_params_list) != len(stage_types):
+            default_params_list = (
+                default_params_list
+                + [OmniDiffusionSamplingParams() if st == "diffusion" else SamplingParams() for st in stage_types]
+            )[: len(stage_types)]
 
-        return ImageGenerationResponse(
-            created=int(time.time()),
-            data=image_data,
+        sampling_params_list: list[OmniSamplingParams] = []
+        for idx, stage_type in enumerate(stage_types):
+            if stage_type == "diffusion":
+                sampling_params_list.append(gen_params)
+            else:
+                base_params = default_params_list[idx]
+                sampling_params_list.append(base_params)
+
+        async for output in engine_client.generate(
+            sampling_params_list=sampling_params_list,
+            **kwargs,
+        ):
+            result = output
+    else:
+        result = await engine_client.generate(
+            sampling_params_list=[gen_params],
+            **kwargs,
         )
 
-    except HTTPException:
-        # Re-raise HTTPExceptions as-is
-        raise
-    except ValueError as e:
-        logger.error(f"Validation error: {e}")
-        raise HTTPException(status_code=HTTPStatus.BAD_REQUEST.value, detail=str(e))
-    except Exception as e:
-        logger.exception(f"Image generation failed: {e}")
+    if result is None:
         raise HTTPException(
-            status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=f"Image generation failed: {str(e)}"
+            status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value,
+            detail="No output generated from multi-stage pipeline.",
         )
+    return result
+
+
+def _update_if_not_none(object: any, key: str, val: any) -> None:
+    if val is not None:
+        setattr(object, key, val)
+
+
+def _extract_images_from_result(result: Any) -> list[Any]:
+    images = []
+    if hasattr(result, "images") and result.images:
+        images = result.images
+    elif hasattr(result, "request_output"):
+        request_output = result.request_output
+        if isinstance(request_output, dict) and request_output.get("images"):
+            images = request_output["images"]
+        elif hasattr(request_output, "images") and request_output.images:
+            images = request_output.images
+    return images
+
+
+async def _load_input_images(
+    inputs: list[str],
+) -> list[Image.Image]:
+    """
+    convert to PIL.Image.Image list
+    """
+    if isinstance(inputs, str):
+        inputs = [inputs]
+
+    images: list[Image.Image] = []
+
+    for inp in inputs:
+        # 1. URL + base64
+        if isinstance(inp, str) and inp.startswith("data:image"):
+            try:
+                _, b64_data = inp.split(",", 1)
+                image_bytes = base64.b64decode(b64_data)
+                img = Image.open(io.BytesIO(image_bytes))
+                images.append(img)
+            except Exception as e:
+                raise ValueError(f"Invalid base64 image: {e}")
+
+        # 2. URL
+        elif isinstance(inp, str) and inp.startswith("http"):
+            async with httpx.AsyncClient(timeout=60) as client:
+                try:
+                    resp = await client.get(inp)
+                    resp.raise_for_status()
+                    img = Image.open(io.BytesIO(resp.content))
+                    images.append(img)
+                except Exception as e:
+                    raise ValueError(f"Failed to download image from URL {inp}: {e}")
+
+        # 3. UploadFile
+        elif hasattr(inp, "file"):
+            try:
+                img_data = await inp.read()
+                img = Image.open(io.BytesIO(img_data))
+                images.append(img)
+            except Exception as e:
+                raise ValueError(f"Failed to open uploaded file: {e}")
+        else:
+            raise ValueError(f"Unsupported input: {inp}")
+
+    if not images:
+        raise ValueError("No valid input images found")
+
+    return images
+
+
+def _choose_output_format(output_format: str | None, background: str | None) -> str:
+    # Normalize and choose extension
+    fmt = (output_format or "").lower()
+    if fmt in {"jpg", "png", "webp", "jpeg"}:
+        return fmt
+    # If transparency requested, prefer png
+    if (background or "auto").lower() == "transparent":
+        return "png"
+    # Default
+    return "jpeg"
+
+
+def _encode_image_base64_with_compression(
+    image: Image.Image, format: str = "png", output_compression: int = 100
+) -> str:
+    """Encode PIL Image to base64 PNG string.
+
+    Args:
+        image: PIL Image object
+        format: Output image format (e.g., "PNG", "JPEG", "WEBP")
+        output_compression: Compression level (0-100%), 100 for best quality
+    Returns:
+        Base64-encoded image as string
+    """
+    buffer = io.BytesIO()
+    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))  # Map 0-100 to 9-0
+
+    image.save(buffer, format=format, **save_kwargs)
+    buffer.seek(0)
+    return base64.b64encode(buffer.read()).decode("utf-8")
+
+
+def apply_stage_default_sampling_params(
+    default_params_json: str | None,
+    sampling_params: any,
+    stage_key: str,
+) -> None:
+    """
+    Update a stage's sampling parameters with vLLM-Omni defaults.
+
+    Args:
+        default_params_json: JSON string of stage-keyed default parameters
+        sampling_params: The sampling parameters object to update
+        stage_key: The stage ID/key in the pipeline
+    """
+    if default_params_json is not None:
+        default_params_dict = json.loads(default_params_json)
+        if stage_key in default_params_dict:
+            stage_defaults = default_params_dict[stage_key]
+            for param_name, param_value in stage_defaults.items():
+                if hasattr(sampling_params, param_name):
+                    setattr(sampling_params, param_name, param_value)

vllm_omni/entrypoints/openai/protocol/images.py

@@ -123,3 +123,5 @@ 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")
+    size: str = Field(None, description="The size of the image generated")