From 3c48257f8d2aa001a6676fc6e39602a8e6c900be Mon Sep 17 00:00:00 2001 From: TaffyOfficial <2587297563@qq.com> Date: Wed, 27 May 2026 08:27:11 +0800 Subject: [PATCH] [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 | 73 +++++- .../openai_api/test_image_server.py | 237 ++++++++++++++++++ vllm_omni/entrypoints/openai/api_server.py | 63 ++--- .../entrypoints/openai/image_api_utils.py | 44 ++++ .../entrypoints/openai/protocol/__init__.py | 8 + .../entrypoints/openai/protocol/images.py | 37 ++- vllm_omni/entrypoints/openai/serving_chat.py | 156 +++++++++++- 7 files changed, 556 insertions(+), 62 deletions(-) diff --git a/docs/serving/image_edit_api.md b/docs/serving/image_edit_api.md index 79303e1a6..4fc13a62b 100644 --- a/docs/serving/image_edit_api.md +++ b/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":"","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":"","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 diff --git a/tests/entrypoints/openai_api/test_image_server.py b/tests/entrypoints/openai_api/test_image_server.py index 40adb7a91..2556ac71d 100644 --- a/tests/entrypoints/openai_api/test_image_server.py +++ b/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( diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 69a6dd603..d8d59285d 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/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, diff --git a/vllm_omni/entrypoints/openai/image_api_utils.py b/vllm_omni/entrypoints/openai/image_api_utils.py index 97847ad11..b958d786a 100644 --- a/vllm_omni/entrypoints/openai/image_api_utils.py +++ b/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: diff --git a/vllm_omni/entrypoints/openai/protocol/__init__.py b/vllm_omni/entrypoints/openai/protocol/__init__.py index c73203cc4..c68f6f598 100644 --- a/vllm_omni/entrypoints/openai/protocol/__init__.py +++ b/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", diff --git a/vllm_omni/entrypoints/openai/protocol/images.py b/vllm_omni/entrypoints/openai/protocol/images.py index dbf4c24b3..34fee760b 100644 --- a/vllm_omni/entrypoints/openai/protocol/images.py +++ b/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 diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index d3bf29d89..ec10606a1 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/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,