Handle partial image streaming

Author: ciaranbor

src/exo/master/api.py

@@ -1,4 +1,5 @@
 import base64
+import json
 import time
 from collections.abc import AsyncGenerator
 from typing import Literal, cast
@@ -194,7 +195,9 @@ def _setup_routes(self) -> None:
         self.app.post("/v1/chat/completions", response_model=None)(
             self.chat_completions
         )
-        self.app.post("/v1/images/generations")(self.image_generations)
+        self.app.post("/v1/images/generations", response_model=None)(
+            self.image_generations
+        )
         self.app.post("/v1/images/edits")(self.image_edits)
         self.app.get("/state")(lambda: self.state)
         self.app.get("/events")(lambda: self._event_log)
@@ -551,8 +554,12 @@ async def chat_completions(
 
     async def image_generations(
         self, payload: ImageGenerationTaskParams
-    ) -> ImageGenerationResponse:
-        """Handle image generation requests."""
+    ) -> ImageGenerationResponse | StreamingResponse:
+        """Handle image generation requests.
+
+        When stream=True and partial_images > 0, returns a StreamingResponse
+        with SSE-formatted events for partial and final images.
+        """
         model_meta = await resolve_model_meta(payload.model)
         payload.model = model_meta.model_id
 
@@ -570,22 +577,128 @@ async def image_generations(
         )
         await self._send(command)
 
-        # Collect all image chunks (non-streaming)
-        num_images = payload.n or 1
+        # Check if streaming is requested
+        if payload.stream and payload.partial_images and payload.partial_images > 0:
+            return StreamingResponse(
+                self._generate_image_stream(
+                    command_id=command.command_id,
+                    num_images=payload.n or 1,
+                    response_format=payload.response_format or "b64_json",
+                ),
+                media_type="text/event-stream",
+            )
+
+        # Non-streaming: collect all image chunks
+        return await self._collect_image_generation(
+            command_id=command.command_id,
+            num_images=payload.n or 1,
+            response_format=payload.response_format or "b64_json",
+        )
+
+    async def _generate_image_stream(
+        self,
+        command_id: CommandId,
+        num_images: int,
+        response_format: str,
+    ) -> AsyncGenerator[str, None]:
+        """Generate SSE stream of partial and final images."""
+        # Track chunks: {(image_index, is_partial): {chunk_index: data}}
+        image_chunks: dict[tuple[int, bool], dict[int, str]] = {}
+        image_total_chunks: dict[tuple[int, bool], int] = {}
+        image_metadata: dict[tuple[int, bool], tuple[int | None, int | None]] = {}
+        images_complete = 0
+
+        try:
+            self._image_generation_queues[command_id], recv = channel[ImageChunk]()
+
+            with recv as chunks:
+                async for chunk in chunks:
+                    key = (chunk.image_index, chunk.is_partial)
+
+                    if key not in image_chunks:
+                        image_chunks[key] = {}
+                        image_total_chunks[key] = chunk.total_chunks
+                        image_metadata[key] = (
+                            chunk.partial_index,
+                            chunk.total_partials,
+                        )
+
+                    image_chunks[key][chunk.chunk_index] = chunk.data
+
+                    # Check if this image is complete
+                    if len(image_chunks[key]) == image_total_chunks[key]:
+                        full_data = "".join(
+                            image_chunks[key][i] for i in range(len(image_chunks[key]))
+                        )
+
+                        partial_idx, total_partials = image_metadata[key]
+
+                        if chunk.is_partial:
+                            # Yield partial image event
+                            event_data = {
+                                "type": "partial",
+                                "partial_index": partial_idx,
+                                "total_partials": total_partials,
+                                "data": {
+                                    "b64_json": full_data
+                                    if response_format == "b64_json"
+                                    else None,
+                                },
+                            }
+                            yield f"data: {json.dumps(event_data)}\n\n"
+                        else:
+                            # Final image
+                            event_data = {
+                                "type": "final",
+                                "image_index": chunk.image_index,
+                                "data": {
+                                    "b64_json": full_data
+                                    if response_format == "b64_json"
+                                    else None,
+                                },
+                            }
+                            yield f"data: {json.dumps(event_data)}\n\n"
+                            images_complete += 1
 
+                            if images_complete >= num_images:
+                                yield "data: [DONE]\n\n"
+                                break
+
+                        # Clean up completed image chunks
+                        del image_chunks[key]
+                        del image_total_chunks[key]
+                        del image_metadata[key]
+
+        except anyio.get_cancelled_exc_class():
+            raise
+        finally:
+            await self._send(TaskFinished(finished_command_id=command_id))
+            if command_id in self._image_generation_queues:
+                del self._image_generation_queues[command_id]
+
+    async def _collect_image_generation(
+        self,
+        command_id: CommandId,
+        num_images: int,
+        response_format: str,
+    ) -> ImageGenerationResponse:
+        """Collect all image chunks (non-streaming) and return a single response."""
         # Track chunks per image: {image_index: {chunk_index: data}}
+        # Only track non-partial (final) images
         image_chunks: dict[int, dict[int, str]] = {}
         image_total_chunks: dict[int, int] = {}
         images_complete = 0
 
         try:
-            self._image_generation_queues[command.command_id], recv = channel[
-                ImageChunk
-            ]()
+            self._image_generation_queues[command_id], recv = channel[ImageChunk]()
 
             while images_complete < num_images:
                 with recv as chunks:
                     async for chunk in chunks:
+                        # Skip partial images in non-streaming mode
+                        if chunk.is_partial:
+                            continue
+
                         if chunk.image_index not in image_chunks:
                             image_chunks[chunk.image_index] = {}
                             image_total_chunks[chunk.image_index] = chunk.total_chunks
@@ -609,26 +722,18 @@ async def image_generations(
                 full_data = "".join(chunks_dict[i] for i in range(len(chunks_dict)))
                 images.append(
                     ImageData(
-                        b64_json=full_data
-                        if payload.response_format == "b64_json"
-                        else None,
+                        b64_json=full_data if response_format == "b64_json" else None,
                         url=None,  # URL format not implemented yet
                     )
                 )
 
             return ImageGenerationResponse(data=images)
         except anyio.get_cancelled_exc_class():
-            # TODO(ciaran): TaskCancelled
-            """
-            self.command_sender.send_nowait(
-                ForwarderCommand(origin=self.node_id, command=command)
-            )
-            """
             raise
         finally:
-            # Send TaskFinished command
-            await self._send(TaskFinished(finished_command_id=command.command_id))
-            del self._image_generation_queues[command.command_id]
+            await self._send(TaskFinished(finished_command_id=command_id))
+            if command_id in self._image_generation_queues:
+                del self._image_generation_queues[command_id]
 
     async def image_edits(
         self,

src/exo/shared/types/chunks.py

@@ -30,6 +30,9 @@ class ImageChunk(BaseChunk):
     chunk_index: int
     total_chunks: int
     image_index: int
+    is_partial: bool = False
+    partial_index: int | None = None
+    total_partials: int | None = None
 
     def __repr_args__(self) -> Generator[tuple[str, Any], None, None]:
         for name, value in super().__repr_args__():

src/exo/shared/types/worker/runner_response.py

@@ -32,5 +32,19 @@ def __repr_args__(self) -> Generator[tuple[str, Any], None, None]:
                 yield name, value
 
 
+class PartialImageResponse(BaseRunnerResponse):
+    image_data: bytes
+    format: Literal["png", "jpeg", "webp"] = "png"
+    partial_index: int
+    total_partials: int
+
+    def __repr_args__(self) -> Generator[tuple[str, Any], None, None]:
+        for name, value in super().__repr_args__():
+            if name == "image_data":
+                yield name, f"<{len(self.image_data)} bytes>"
+            elif name is not None:
+                yield name, value
+
+
 class FinishedResponse(BaseRunnerResponse):
     pass

src/exo/worker/engines/image/base.py

@@ -1,5 +1,6 @@
+from collections.abc import Generator
 from pathlib import Path
-from typing import Literal, Optional, Protocol, runtime_checkable
+from typing import Literal, Protocol, runtime_checkable
 
 from PIL import Image
 
@@ -21,11 +22,18 @@ def generate(
         seed: int,
         image_path: Path | None = None,
         image_strength: float | None = None,
-    ) -> Optional[Image.Image]:
+        partial_images: int = 0,
+    ) -> Generator[Image.Image | tuple[Image.Image, int, int], None, None]:
         """Generate an image from a text prompt, or edit an existing image.
 
-        For distributed inference, only the first stage (rank 0) returns the image.
-        Other stages return None after participating in the pipeline.
+        For distributed inference, only the last stage returns images.
+        Other stages yield nothing after participating in the pipeline.
+
+        When partial_images > 0, yields intermediate images during diffusion
+        as tuples of (image, partial_index, total_partials), then yields
+        the final image.
+
+        When partial_images = 0 (default), only yields the final image.
 
         Args:
             prompt: Text description of the image to generate
@@ -35,8 +43,10 @@ def generate(
             seed: Random seed for reproducibility
             image_path: Optional path to input image for img2img
             image_strength: Optional strength for img2img (0.0-1.0, higher = more change)
+            partial_images: Number of intermediate images to yield (0 for none)
 
-        Returns:
-            Generated PIL Image (rank 0) or None (other ranks)
+        Yields:
+            Intermediate images as (Image, partial_index, total_partials) tuples
+            Final PIL Image (last stage) or nothing (other stages)
         """
         ...

src/exo/worker/engines/image/distributed_model.py

@@ -1,3 +1,4 @@
+from collections.abc import Generator
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Literal, Optional
 
@@ -186,7 +187,8 @@ def generate(
         seed: int = 2,
         image_path: Path | None = None,
         image_strength: float | None = None,
-    ) -> Optional[Image.Image]:
+        partial_images: int = 0,
+    ) -> Generator[Image.Image | tuple[Image.Image, int, int], None, None]:
         # Determine number of inference steps based on quality
         steps = self._config.get_steps_for_quality(quality)
 
@@ -197,20 +199,22 @@ def generate(
             image_path=image_path,
             image_strength=image_strength,
         )
-        image = self._generate_image(settings=config, prompt=prompt, seed=seed)
-        logger.info("generated image")
 
-        # Only final rank returns the actual image
-        if self.is_last_stage:
-            return image.image
-
-    def _generate_image(self, settings: Config, prompt: str, seed: int) -> Any:
-        """Generate image by delegating to the runner."""
-        return self._runner.generate_image(
-            settings=settings,
+        # Generate images via the runner
+        for result in self._runner.generate_image(
+            settings=config,
             prompt=prompt,
             seed=seed,
-        )
+            partial_images=partial_images,
+        ):
+            if isinstance(result, tuple):
+                # Partial image: (GeneratedImage, partial_index, total_partials)
+                generated_image, partial_idx, total_partials = result
+                yield (generated_image.image, partial_idx, total_partials)
+            else:
+                # Final image: GeneratedImage
+                logger.info("generated image")
+                yield result.image
 
 
 def initialize_image_model(bound_instance: BoundInstance) -> DistributedImageModel: