[Feat][HunyuanImage3] Stream AR text for IT2I image edits (#3723)

Signed-off-by: TaffyOfficial <2324465096@qq.com>
Co-authored-by: TaffyOfficial <2324465096@qq.com>

docs/serving/image_edit_api.md

@@ -13,7 +13,7 @@ For example...
 ```bash
 # Qwen-Image
 vllm serve Qwen/Qwen-Image-Edit-2511 --omni --port 8000
-
+```
 
 ### Generate Images
 
@@ -92,7 +92,8 @@ Content-Type: multipart/form-data
 | `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).
+| `background` | string or null | "auto" | Allows to set transparency for the background of the generated image(s). |
+| `stream` | boolean | false | Return Server-Sent Events instead of a single JSON response. Currently supported only by multi-stage image edit pipelines that produce AR text before the final image, such as HunyuanImage3 IT2I AR+DiT. |
 
 #### vllm-omni Extension Parameters
 
@@ -109,6 +110,8 @@ Content-Type: multipart/form-data
 
 ### Response Format
 
+When `stream=false` or omitted, the endpoint returns the standard image edit response:
+
 ```json
 {
   "created": 1701234567,
@@ -120,10 +123,56 @@ Content-Type: multipart/form-data
     }
   ],
   "output_format": null,
-  "size": null,
+  "size": null
 }
 ```
 
+### Streaming Response Format
+
+Set `stream=true` when you want to receive the HunyuanImage3 IT2I AR
+recaption text before the final edited image is ready. Streaming is only
+available for multi-stage image edit pipelines where stage 0 produces AR text
+and a later diffusion stage produces the image. Single-stage diffusion image
+edit pipelines reject `stream=true` with HTTP 400.
+
+The response uses Server-Sent Events with `Content-Type: text/event-stream`.
+Each event is sent as one `data:` line. The stream order is:
+
+1. One or more AR text delta chunks.
+2. One final image chunk.
+3. `data: [DONE]`.
+
+AR text chunks expose only the generated text delta and the AR completion
+`index`. They do not expose token ids, ratio tokens, KV cache data, or internal
+prompt token ids.
+
+Example AR delta event:
+
+```text
+data: {"object":"image.edit.chunk","type":"ar_delta","delta":"A close-up product photo...","index":0,"created":1701234567,"model":"tencent/HunyuanImage-3.0-Instruct"}
+```
+
+Example final image event:
+
+```text
+data: {"object":"image.edit.chunk","type":"image","data":[{"b64_json":"<base64-encoded PNG>","url":null,"revised_prompt":null}],"output_format":"png","size":"1024x1024","created":1701234567,"model":"tencent/HunyuanImage-3.0-Instruct"}
+```
+
+Terminal event:
+
+```text
+data: [DONE]
+```
+
+If the engine fails after the SSE response has started, the stream emits an
+error chunk followed by `[DONE]`:
+
+```text
+data: {"object":"error","created":1701234567,"model":"tencent/HunyuanImage-3.0-Instruct","error":{"message":"<error message>","type":"server_error","code":500}}
+
+data: [DONE]
+```
+
 ## Examples
 
 ### Multiple Images input
@@ -134,12 +183,28 @@ curl -s -D >(grep -i x-request-id >&2) \
   -X POST "http://localhost:8000/v1/images/edits" \
   -F "model=xxx" \
   -F "image=@xx.png" \
-  -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"
 ```
 
+### Streaming HunyuanImage3 IT2I AR Text
+
+```bash
+curl -N -X POST "http://localhost:8000/v1/images/edits" \
+  -F "model=tencent/HunyuanImage-3.0-Instruct" \
+  -F "image=@./input.png" \
+  -F "prompt=Turn the product into a clean studio advertisement." \
+  -F "size=1024x1024" \
+  -F "output_format=png" \
+  -F "stream=true"
+```
+
+Use the `ar_delta` events for progressive display of the AR-generated
+recaption. Decode the `b64_json` field from the final `image` event to get the
+edited image.
+
 
 ## Parameter Handling
 

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
+import json
 from argparse import Namespace
 from types import SimpleNamespace
 
@@ -19,6 +20,7 @@ from PIL import Image
 from pytest_mock import MockerFixture
 from vllm import SamplingParams
 from vllm.entrypoints.openai.models.protocol import BaseModelPath
+from vllm.sampling_params import RequestOutputKind
 
 from vllm_omni.entrypoints.async_omni import AsyncOmni
 from vllm_omni.entrypoints.openai.api_server import _DiffusionServingModels, router
@@ -120,6 +122,27 @@ class MockGenerationResult:
         self.peak_memory_mb = 0.0
 
 
+class MockStageResult:
+    """Mock multi-stage output for streaming image edit tests."""
+
+    def __init__(self, *, stage_id, final_output_type, text="", texts=None, images=None):
+        self.stage_id = stage_id
+        self.final_output_type = final_output_type
+        self.images = images or []
+        if texts is not None:
+            outputs = [SimpleNamespace(text=item, index=index) for index, item in enumerate(texts)]
+        elif text:
+            outputs = [SimpleNamespace(text=text, index=0)]
+        else:
+            outputs = []
+        self.request_output = SimpleNamespace(
+            outputs=outputs,
+            images=self.images,
+        )
+        self.stage_durations = {}
+        self.peak_memory_mb = 0.0
+
+
 class FakeAsyncOmni:
     """Fake AsyncOmni that yields a single diffusion output."""
 
@@ -391,6 +414,72 @@ def async_omni_stage_configs_only_client():
     return TestClient(app)
 
 
+@pytest.fixture
+def streaming_image_edit_client():
+    """Create a multi-stage client whose engine yields AR text before image output."""
+    from fastapi import FastAPI
+
+    from vllm_omni.entrypoints.async_omni import AsyncOmni
+    from vllm_omni.entrypoints.openai.api_server import router
+    from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
+
+    class FakeAsyncOmniClass(AsyncOmni):
+        def __init__(self):
+            stage_configs = [
+                SimpleNamespace(stage_type="llm", is_comprehension=True),
+                SimpleNamespace(stage_type="diffusion", is_comprehension=False),
+            ]
+            default_sampling_params_list = [
+                SamplingParams(temperature=0.1),
+                OmniDiffusionSamplingParams(),
+            ]
+            self.engine = SimpleNamespace(
+                stage_configs=stage_configs,
+                default_sampling_params_list=default_sampling_params_list,
+            )
+            self.default_sampling_params_list = default_sampling_params_list
+            self.captured_sampling_params_list = None
+            self.captured_prompt = None
+            self.od_config = SimpleNamespace(supports_multimodal_inputs=True)
+
+        async def generate(self, prompt, request_id, sampling_params=None, sampling_params_list=None):
+            self.captured_prompt = prompt
+            self.captured_sampling_params_list = sampling_params_list or [sampling_params]
+            assert self.captured_sampling_params_list[0].output_kind == RequestOutputKind.DELTA
+            yield MockStageResult(stage_id=0, final_output_type="text", text="recap")
+            yield MockStageResult(stage_id=0, final_output_type="text", text=" done")
+            yield MockStageResult(
+                stage_id=1,
+                final_output_type="image",
+                images=[Image.new("RGB", (32, 24), color="purple")],
+            )
+
+        def __class_getitem__(cls, item):
+            return cls
+
+        def get_diffusion_od_config(self):
+            return self.od_config
+
+    app = FastAPI()
+    app.include_router(router)
+
+    engine = FakeAsyncOmniClass()
+    chat_handler = object.__new__(OmniOpenAIServingChat)
+    chat_handler.engine_client = engine
+    chat_handler._diffusion_engine = None
+    app.state.openai_serving_chat = chat_handler
+    app.state.engine_client = engine
+    app.state.stage_configs = [
+        SimpleNamespace(stage_type="llm"),
+        SimpleNamespace(stage_type="diffusion"),
+    ]
+    app.state.args = Namespace(
+        default_sampling_params='{"1": {"num_inference_steps":4, "guidance_scale":7.5, "generator_device":"cpu"}}',
+        max_generated_image_size=1024 * 1792,
+    )
+    return TestClient(app)
+
+
 def test_health_endpoint(test_client):
     """Test health check endpoint for diffusion mode"""
     response = test_client.get("/health")
@@ -680,6 +769,154 @@ def test_image_edits_async_omni_stage_configs_only(async_omni_stage_configs_only
     assert len(captured) == 2
 
 
+def _parse_sse_payloads(body: str):
+    payloads = []
+    for line in body.splitlines():
+        if not line.startswith("data: "):
+            continue
+        data = line[len("data: ") :]
+        payloads.append(data if data == "[DONE]" else json.loads(data))
+    return payloads
+
+
+def test_image_edits_streaming_returns_ar_delta_then_image(streaming_image_edit_client):
+    img_bytes = make_test_image_bytes((16, 16))
+    response = streaming_image_edit_client.post(
+        "/v1/images/edits",
+        files=[("image", img_bytes)],
+        data={
+            "prompt": "edit me",
+            "size": "auto",
+            "stream": "true",
+            "output_format": "png",
+        },
+    )
+
+    assert response.status_code == 200
+    assert response.headers["content-type"].startswith("text/event-stream")
+
+    payloads = _parse_sse_payloads(response.text)
+    assert [p["type"] if isinstance(p, dict) and "type" in p else p for p in payloads] == [
+        "ar_delta",
+        "ar_delta",
+        "image",
+        "[DONE]",
+    ]
+    assert payloads[0]["delta"] == "recap"
+    assert payloads[0]["index"] == 0
+    assert payloads[1]["delta"] == " done"
+    assert payloads[2]["output_format"] == "png"
+    assert payloads[2]["size"] == "16x16"
+
+    image_payload = payloads[2]["data"][0]
+    img = Image.open(io.BytesIO(base64.b64decode(image_payload["b64_json"])))
+    assert img.size == (32, 24)
+
+
+def test_image_edits_streaming_ar_delta_chunks_include_index(streaming_image_edit_client):
+    async def generate_multi_output_delta(prompt, request_id, sampling_params=None, sampling_params_list=None):
+        yield MockStageResult(stage_id=0, final_output_type="text", texts=["first", "second"])
+        yield MockStageResult(
+            stage_id=1,
+            final_output_type="image",
+            images=[Image.new("RGB", (32, 24), color="purple")],
+        )
+
+    streaming_image_edit_client.app.state.engine_client.generate = generate_multi_output_delta
+    img_bytes = make_test_image_bytes((16, 16))
+    response = streaming_image_edit_client.post(
+        "/v1/images/edits",
+        files=[("image", img_bytes)],
+        data={
+            "prompt": "edit me",
+            "stream": "true",
+        },
+    )
+
+    assert response.status_code == 200
+    payloads = _parse_sse_payloads(response.text)
+    assert [(payloads[0]["index"], payloads[0]["delta"]), (payloads[1]["index"], payloads[1]["delta"])] == [
+        (0, "first"),
+        (1, "second"),
+    ]
+
+
+def test_image_edits_streaming_errors_without_final_image(streaming_image_edit_client):
+    async def generate_without_image(prompt, request_id, sampling_params=None, sampling_params_list=None):
+        yield MockStageResult(stage_id=0, final_output_type="text", text="recap")
+
+    streaming_image_edit_client.app.state.engine_client.generate = generate_without_image
+    img_bytes = make_test_image_bytes((16, 16))
+    response = streaming_image_edit_client.post(
+        "/v1/images/edits",
+        files=[("image", img_bytes)],
+        data={
+            "prompt": "edit me",
+            "stream": "true",
+        },
+    )
+
+    assert response.status_code == 200
+    payloads = _parse_sse_payloads(response.text)
+    assert payloads[0]["type"] == "ar_delta"
+    assert payloads[1]["object"] == "error"
+    assert "without a final image" in payloads[1]["error"]["message"]
+    assert payloads[2] == "[DONE]"
+
+
+def test_image_edits_streaming_errors_on_empty_final_image(streaming_image_edit_client):
+    async def generate_empty_image(prompt, request_id, sampling_params=None, sampling_params_list=None):
+        yield MockStageResult(stage_id=0, final_output_type="text", text="recap")
+        yield MockStageResult(stage_id=1, final_output_type="image", images=[])
+
+    streaming_image_edit_client.app.state.engine_client.generate = generate_empty_image
+    img_bytes = make_test_image_bytes((16, 16))
+    response = streaming_image_edit_client.post(
+        "/v1/images/edits",
+        files=[("image", img_bytes)],
+        data={
+            "prompt": "edit me",
+            "stream": "true",
+        },
+    )
+
+    assert response.status_code == 200
+    payloads = _parse_sse_payloads(response.text)
+    assert payloads[0]["type"] == "ar_delta"
+    assert payloads[1]["object"] == "error"
+    assert "empty final image" in payloads[1]["error"]["message"]
+    assert payloads[2] == "[DONE]"
+
+
+def test_image_edits_streaming_rejects_single_stage(test_client):
+    img_bytes = make_test_image_bytes((16, 16))
+    response = test_client.post(
+        "/v1/images/edits",
+        files=[("image", img_bytes)],
+        data={
+            "prompt": "edit me",
+            "stream": "true",
+        },
+    )
+
+    assert response.status_code == 400
+    assert "multi-stage" in response.json()["detail"]
+
+
+def test_image_edits_streaming_rejects_single_stage_before_loading_url(test_client):
+    response = test_client.post(
+        "/v1/images/edits",
+        data={
+            "prompt": "edit me",
+            "url": "https://example.invalid/not-fetched.png",
+            "stream": "true",
+        },
+    )
+
+    assert response.status_code == 400
+    assert "multi-stage" in response.json()["detail"]
+
+
 def test_generate_images_max_size_rejected(async_omni_test_client):
     """Test that a size exceeding max_generated_image_size returns 400."""
     response = async_omni_test_client.post(

vllm_omni/entrypoints/openai/api_server.py

@@ -92,6 +92,7 @@ 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,
 )
@@ -1644,7 +1645,7 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request)
 
         # 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)
+            ImageData(b64_json=encode_image_base64_with_compression(img, format=output_format), revised_prompt=None)
             for img in images
         ]
 
@@ -1694,6 +1695,7 @@ async def edit_images(
     output_format: str | None = Form("png"),
     background: str | None = Form("auto"),
     output_compression: Annotated[int, Form(ge=0, le=100)] = 100,
+    stream: bool = Form(False),
     user: str | None = Form(None),  # unused now
     # vllm-omni extensions for image editing
     mask_image: str | UploadFile | None = None,
@@ -1734,6 +1736,11 @@ async def edit_images(
             status_code=HTTPStatus.BAD_REQUEST.value,
             detail="Only response_format 'b64_json' is supported now.",
         )
+    if stream and len(stage_configs) <= 1:
+        raise HTTPException(
+            status_code=HTTPStatus.BAD_REQUEST.value,
+            detail="stream=true is only supported for multi-stage image editing pipelines.",
+        )
     try:
         # 2. Build prompt & images params
         cot_output = None
@@ -1941,7 +1948,17 @@ async def edit_images(
                 extra_body=extra_body,
                 reference_images=ref_b64_list,
                 request_id=request_id,
+                stream=stream,
+                model=model_name,
+                output_format=output_format,
+                output_compression=output_compression,
+                size=size_str,
             )
+            if stream and not isinstance(generation_result, ErrorResponse):
+                return StreamingResponse(
+                    content=generation_result,
+                    media_type="text/event-stream",
+                )
             if isinstance(generation_result, ErrorResponse):
                 raise HTTPException(
                     status_code=generation_result.error.code if generation_result.error else 400,
@@ -1964,7 +1981,7 @@ async def edit_images(
         # Encode images to base64
         image_data = [
             ImageData(
-                b64_json=_encode_image_base64_with_compression(
+                b64_json=encode_image_base64_with_compression(
                     img, format=output_format, output_compression=output_compression
                 ),
                 revised_prompt=None,
@@ -2289,48 +2306,6 @@ def _choose_output_format(output_format: str | None, background: str | None) ->
     return "jpeg"
 
 
-def _prepare_image_for_output_format(image: Image.Image, format: str) -> Image.Image:
-    fmt = format.lower()
-    if fmt not in {"jpg", "jpeg"}:
-        return image
-
-    if image.mode == "RGB":
-        return image
-
-    if image.mode in {"RGBA", "LA"} or (image.mode == "P" and "transparency" in image.info):
-        alpha_image = image.convert("RGBA")
-        flattened = Image.new("RGB", alpha_image.size, (255, 255, 255))
-        flattened.paste(alpha_image, mask=alpha_image.getchannel("A"))
-        return flattened
-
-    return image.convert("RGB")
-
-
-def _encode_image_base64_with_compression(
-    image: Image.Image, format: str = "png", output_compression: int = 100
-) -> str:
-    """Encode PIL Image to a base64 image 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()
-    image = _prepare_image_for_output_format(image, 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))  # 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,

vllm_omni/entrypoints/openai/image_api_utils.py

@@ -68,6 +68,50 @@ def encode_image_base64(image: PIL.Image.Image) -> str:
     return base64.b64encode(buffer.read()).decode("utf-8")
 
 
+def prepare_image_for_output_format(image: PIL.Image.Image, format: str) -> PIL.Image.Image:
+    fmt = format.lower()
+    if fmt not in {"jpg", "jpeg"}:
+        return image
+
+    if image.mode == "RGB":
+        return image
+
+    if image.mode in {"RGBA", "LA"} or (image.mode == "P" and "transparency" in image.info):
+        alpha_image = image.convert("RGBA")
+        flattened = PIL.Image.new("RGB", alpha_image.size, (255, 255, 255))
+        flattened.paste(alpha_image, mask=alpha_image.getchannel("A"))
+        return flattened
+
+    return image.convert("RGB")
+
+
+def encode_image_base64_with_compression(
+    image: PIL.Image.Image, format: str = "png", output_compression: int = 100
+) -> str:
+    """Encode a PIL image to a base64 image string.
+
+    Args:
+        image: PIL Image object.
+        format: Output image format, such as "png", "jpeg", or "webp".
+        output_compression: Compression level (0-100), where 100 keeps the
+            best quality for lossy formats and uses the lowest PNG compression.
+
+    Returns:
+        Base64-encoded image string.
+    """
+    buffer = io.BytesIO()
+    image = prepare_image_for_output_format(image, 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)
+    buffer.seek(0)
+    return base64.b64encode(buffer.read()).decode("utf-8")
+
+
 def validate_layered_layers(layers: int | None) -> int | None:
     """Validate the Qwen-Image-Layered ``layers`` parameter."""
     if layers is None:

vllm_omni/entrypoints/openai/protocol/__init__.py

@@ -4,6 +4,10 @@
 from vllm_omni.entrypoints.openai.protocol.chat_completion import OmniChatCompletionStreamResponse
 from vllm_omni.entrypoints.openai.protocol.images import (
     ImageData,
+    ImageEditARDeltaChunk,
+    ImageEditImageChunk,
+    ImageEditStreamError,
+    ImageEditStreamResponse,
     ImageGenerationRequest,
     ImageGenerationResponse,
     ResponseFormat,
@@ -16,6 +20,10 @@ from vllm_omni.entrypoints.openai.protocol.videos import (
 
 __all__ = [
     "ImageData",
+    "ImageEditARDeltaChunk",
+    "ImageEditImageChunk",
+    "ImageEditStreamError",
+    "ImageEditStreamResponse",
     "ImageGenerationRequest",
     "ImageGenerationResponse",
     "ResponseFormat",

vllm_omni/entrypoints/openai/protocol/images.py

@@ -8,7 +8,7 @@ for text-to-image generation, with vllm-omni specific extensions.
 """
 
 from enum import Enum
-from typing import Any
+from typing import Any, Literal
 
 from pydantic import BaseModel, Field, field_validator
 
@@ -175,3 +175,38 @@ class ImageGenerationResponse(BaseModel):
         description="Chain-of-thought text output from the AR stage. "
         "Only present for image editing (IT2I) with CoT-enabled models.",
     )
+
+
+class ImageEditARDeltaChunk(BaseModel):
+    """Streaming chunk carrying a text delta from the image-edit AR stage."""
+
+    object: Literal["image.edit.chunk"] = "image.edit.chunk"
+    type: Literal["ar_delta"] = "ar_delta"
+    delta: str = Field(..., description="Text delta generated by the AR stage")
+    index: int = Field(default=0, description="Completion index for the AR stage output")
+    created: int = Field(..., description="Unix timestamp of when the stream was created")
+    model: str = Field(..., description="Model used for the image edit request")
+
+
+class ImageEditImageChunk(BaseModel):
+    """Streaming chunk carrying the final image-edit result."""
+
+    object: Literal["image.edit.chunk"] = "image.edit.chunk"
+    type: Literal["image"] = "image"
+    data: list[ImageData] = Field(..., description="Array of generated images")
+    output_format: str = Field(..., description="The output format of the image generation")
+    size: str = Field(..., description="The generated image size")
+    created: int = Field(..., description="Unix timestamp of when the stream was created")
+    model: str = Field(..., description="Model used for the image edit request")
+
+
+class ImageEditStreamError(BaseModel):
+    """Streaming error chunk emitted before the terminal [DONE] event."""
+
+    object: Literal["error"] = "error"
+    created: int = Field(..., description="Unix timestamp of when the stream was created")
+    model: str = Field(..., description="Model used for the image edit request")
+    error: dict[str, Any] = Field(..., description="OpenAI-compatible streaming error payload")
+
+
+ImageEditStreamResponse = ImageEditARDeltaChunk | ImageEditImageChunk | ImageEditStreamError

vllm_omni/entrypoints/openai/serving_chat.py

@@ -86,9 +86,15 @@ from vllm.utils.collection_utils import as_list
 from vllm.v1.engine.exceptions import EngineDeadError
 
 from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin
-from vllm_omni.entrypoints.openai.image_api_utils import validate_layered_layers
+from vllm_omni.entrypoints.openai.image_api_utils import encode_image_base64_with_compression, validate_layered_layers
 from vllm_omni.entrypoints.openai.protocol import OmniChatCompletionStreamResponse
 from vllm_omni.entrypoints.openai.protocol.audio import AudioResponse, CreateAudio
+from vllm_omni.entrypoints.openai.protocol.images import (
+    ImageData,
+    ImageEditARDeltaChunk,
+    ImageEditImageChunk,
+    ImageEditStreamError,
+)
 from vllm_omni.entrypoints.openai.stage_params import (
     build_stage_sampling_params_list,
     get_default_sampling_params_list,
@@ -2414,17 +2420,13 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
 
         return engine_prompt, sampling_params_list
 
-    async def generate_diffusion_images(
+    def _prepare_diffusion_image_request(
         self,
         *,
         prompt: str,
         extra_body: dict[str, Any] | None = None,
         reference_images: list[str] | None = None,
-        request_id: str | None = None,
-    ) -> tuple[list[Image.Image], dict[str, Any], float] | ErrorResponse:
-        """Generate diffusion images and return raw images plus generation stats."""
-        if request_id is None:
-            request_id = f"chatcmpl-{uuid.uuid4().hex[:16]}"
+    ) -> tuple[Any, OmniTextPrompt, OmniDiffusionSamplingParams, list[Image.Image]] | ErrorResponse:
         if extra_body is None:
             extra_body = {}
         if reference_images is None:
@@ -2499,9 +2501,44 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                         status_code=400,
                     )
 
+        return engine, gen_prompt, gen_params, pil_images
+
+    async def generate_diffusion_images(
+        self,
+        *,
+        prompt: str,
+        extra_body: dict[str, Any] | None = None,
+        reference_images: list[str] | None = None,
+        request_id: str | None = None,
+        stream: bool = False,
+        model: str | None = None,
+        output_format: str = "png",
+        output_compression: int = 100,
+        size: str = "auto",
+    ) -> tuple[list[Image.Image], dict[str, Any], float, str | None] | ErrorResponse | AsyncIterator[str]:
+        """Generate diffusion images and return raw images plus generation stats."""
+        if request_id is None:
+            request_id = f"chatcmpl-{uuid.uuid4().hex[:16]}"
+
+        prepared = self._prepare_diffusion_image_request(
+            prompt=prompt,
+            extra_body=extra_body,
+            reference_images=reference_images,
+        )
+        if isinstance(prepared, ErrorResponse):
+            return prepared
+        engine, gen_prompt, gen_params, pil_images = prepared
+        if extra_body is None:
+            extra_body = {}
+
         if isinstance(engine, AsyncOmni):
             diffusion_engine = cast(AsyncOmni, engine)
             stage_configs = getattr(diffusion_engine, "stage_configs", None) or []
+            if stream and len(stage_configs) <= 1:
+                return self._create_error_response(
+                    "stream=true is only supported for multi-stage image editing pipelines.",
+                    status_code=400,
+                )
             if len(stage_configs) > 1:
                 # Pull tokenizer from the comprehension (AR) stage so we can
                 # build HF byte-for-byte prompt_token_ids in the helper. If
@@ -2526,15 +2563,31 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 engine_prompt = gen_prompt
                 sampling_params_list = [gen_params]
 
-            result = None
-            async for output in diffusion_engine.generate(
+            sampling_params_list = coerce_param_message_types(sampling_params_list, stream)
+            result_generator = diffusion_engine.generate(
                 prompt=engine_prompt,
                 sampling_params_list=sampling_params_list,
                 request_id=request_id,
-            ):
+            )
+            if stream:
+                return self._stream_diffusion_image_chunks(
+                    result_generator,
+                    model=model or self._diffusion_model_name,
+                    output_format=output_format,
+                    output_compression=output_compression,
+                    size=size,
+                )
+
+            result = None
+            async for output in result_generator:
                 result = output
             if result is None:
                 return self._create_error_response("No output generated from AsyncOmni", status_code=500)
+        elif stream:
+            return self._create_error_response(
+                "Streaming image edits require a multi-stage AsyncOmni engine.",
+                status_code=400,
+            )
         else:
             result = await engine.generate(
                 prompt=gen_prompt,
@@ -2557,14 +2610,91 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                     if isinstance(ar_text, str) and ar_text.strip():
                         cot_output = ar_text
 
+        return self._flatten_diffusion_images(images), stage_durations, peak_memory_mb, cot_output
+
+    async def _stream_diffusion_image_chunks(
+        self,
+        result_generator: AsyncIterator[OmniRequestOutput],
+        *,
+        model: str,
+        output_format: str,
+        output_compression: int,
+        size: str,
+    ) -> AsyncIterator[str]:
+        """Yield image edit SSE chunks from multi-stage diffusion outputs."""
+        created = int(time.time())
+        emitted_image = False
+        try:
+            async for output in result_generator:
+                final_output_type = getattr(output, "final_output_type", None)
+                stage_id = getattr(output, "stage_id", None)
+                if final_output_type == "text" and stage_id == 0:
+                    request_output = output.request_output
+                    for completion in request_output.outputs:
+                        text = completion.text or ""
+                        if not text:
+                            continue
+                        chunk = ImageEditARDeltaChunk(
+                            delta=text,
+                            index=completion.index,
+                            created=created,
+                            model=model,
+                        )
+                        yield f"data: {chunk.model_dump_json()}\n\n"
+                elif final_output_type == "image":
+                    images = self._flatten_diffusion_images(getattr(output.request_output, "images", []))
+                    if not images:
+                        raise RuntimeError("Streaming image edit produced an empty final image output.")
+                    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
+                    ]
+                    chunk = ImageEditImageChunk(
+                        data=image_data,
+                        output_format=output_format,
+                        size=size,
+                        created=created,
+                        model=model,
+                    )
+                    yield f"data: {chunk.model_dump_json()}\n\n"
+                    emitted_image = True
+            if not emitted_image:
+                raise RuntimeError("Streaming image edit completed without a final image output.")
+        except EngineDeadError as exc:
+            logger.error("EngineDeadError during streaming image edit: %s", exc)
+            data = self.create_streaming_error_response(exc)
+            yield f"data: {data}\n\n"
+            terminate_if_errored(engine=self.engine_client)
+        except Exception as exc:
+            logger.exception("Streaming image edit failed: %s", exc)
+            chunk = ImageEditStreamError(
+                created=created,
+                model=model,
+                error={
+                    "message": str(exc),
+                    "type": "server_error",
+                    "code": 500,
+                },
+            )
+            yield f"data: {chunk.model_dump_json()}\n\n"
+        yield "data: [DONE]\n\n"
+
+    @staticmethod
+    def _flatten_diffusion_images(images: Any) -> list[Image.Image]:
         flat_images: list[Image.Image] = []
-        for item in images:
+        for item in images or []:
             if isinstance(item, list):
                 flat_images.extend(item)
             else:
                 flat_images.append(item)
-
-        return flat_images, stage_durations, peak_memory_mb, cot_output
+        return flat_images
 
     async def _create_diffusion_chat_completion(
         self,