← back to Exo
Ciaran/image non streaming (#1328)
a31942ce122c9ce4050c8dc7b8027623520abca3 · 2026-01-29 17:24:32 +0000 · ciaranbor
## Motivation
The dashboard UI attempted to parse all image generation responses as
SSE streams, even when streaming was disabled. This broke non-streaming
image generation.
## Changes
- Parse JSON responses directly when not streaming, use SSE parser only
when stream=true AND partialImages > 0
- explicitly disable partial images when not streaming
## Why It Works
Both API and dashboard now use the same condition (stream &&
partialImages > 0) to determine response format, ensuring correct
parsing.
## Test Plan
### Manual Testing
Non-streamed image generation results appear in the UI. Streamed image
generation still works
Files touched
M dashboard/package-lock.jsonM dashboard/src/lib/stores/app.svelte.tsM src/exo/worker/engines/image/generate.pyM src/exo/worker/engines/image/pipeline/runner.py
Diff
commit a31942ce122c9ce4050c8dc7b8027623520abca3
Author: ciaranbor <81697641+ciaranbor@users.noreply.github.com>
Date: Thu Jan 29 17:24:32 2026 +0000
Ciaran/image non streaming (#1328)
## Motivation
The dashboard UI attempted to parse all image generation responses as
SSE streams, even when streaming was disabled. This broke non-streaming
image generation.
## Changes
- Parse JSON responses directly when not streaming, use SSE parser only
when stream=true AND partialImages > 0
- explicitly disable partial images when not streaming
## Why It Works
Both API and dashboard now use the same condition (stream &&
partialImages > 0) to determine response format, ensuring correct
parsing.
## Test Plan
### Manual Testing
Non-streamed image generation results appear in the UI. Streamed image
generation still works
---
dashboard/package-lock.json | 10 -
dashboard/src/lib/stores/app.svelte.ts | 378 ++++++++++++++----------
src/exo/worker/engines/image/generate.py | 4 +-
src/exo/worker/engines/image/pipeline/runner.py | 1 +
4 files changed, 224 insertions(+), 169 deletions(-)
diff --git a/dashboard/package-lock.json b/dashboard/package-lock.json
index e28a6b60..345c73d2 100644
--- a/dashboard/package-lock.json
+++ b/dashboard/package-lock.json
@@ -865,7 +865,6 @@
"integrity": "sha512-oH8tXw7EZnie8FdOWYrF7Yn4IKrqTFHhXvl8YxXxbKwTMcD/5NNCryUSEXRk2ZR4ojnub0P8rNrsVGHXWqIDtA==",
"dev": true,
"license": "MIT",
- "peer": true,
"dependencies": {
"@standard-schema/spec": "^1.0.0",
"@sveltejs/acorn-typescript": "^1.0.5",
@@ -905,7 +904,6 @@
"integrity": "sha512-Y1Cs7hhTc+a5E9Va/xwKlAJoariQyHY+5zBgCZg4PFWNYQ1nMN9sjK1zhw1gK69DuqVP++sht/1GZg1aRwmAXQ==",
"dev": true,
"license": "MIT",
- "peer": true,
"dependencies": {
"@sveltejs/vite-plugin-svelte-inspector": "^4.0.1",
"debug": "^4.4.1",
@@ -1522,7 +1520,6 @@
"integrity": "sha512-LCCV0HdSZZZb34qifBsyWlUmok6W7ouER+oQIGBScS8EsZsQbrtFTUrDX4hOl+CS6p7cnNC4td+qrSVGSCTUfQ==",
"dev": true,
"license": "MIT",
- "peer": true,
"dependencies": {
"undici-types": "~6.21.0"
}
@@ -1532,7 +1529,6 @@
"resolved": "https://registry.npmjs.org/acorn/-/acorn-8.15.0.tgz",
"integrity": "sha512-NZyJarBfL7nWwIq+FDL6Zp/yHEhePMNnnJ0y3qfieCrmNvYct8uvtiV41UvlSe6apAfk0fY1FbWx+NwfmpvtTg==",
"license": "MIT",
- "peer": true,
"bin": {
"acorn": "bin/acorn"
},
@@ -1945,7 +1941,6 @@
"integrity": "sha512-fmTRWbNMmsmWq6xJV8D19U/gw/bwrHfNXxrIN+HfZgnzqTHp9jOmKMhsTUjXOJnZOdZY9Q28y4yebKzqDKlxlQ==",
"dev": true,
"license": "ISC",
- "peer": true,
"engines": {
"node": ">=12"
}
@@ -2653,7 +2648,6 @@
"integrity": "sha512-5gTmgEY/sqK6gFXLIsQNH19lWb4ebPDLA4SdLP7dsWkIXHWlG66oPuVvXSGFPppYZz8ZDZq0dYYrbHfBCVUb1Q==",
"dev": true,
"license": "MIT",
- "peer": true,
"engines": {
"node": ">=12"
},
@@ -2696,7 +2690,6 @@
"integrity": "sha512-UOnG6LftzbdaHZcKoPFtOcCKztrQ57WkHDeRD9t/PTQtmT0NHSeWWepj6pS0z/N7+08BHFDQVUrfmfMRcZwbMg==",
"dev": true,
"license": "MIT",
- "peer": true,
"bin": {
"prettier": "bin/prettier.cjs"
},
@@ -2869,7 +2862,6 @@
"resolved": "https://registry.npmjs.org/svelte/-/svelte-5.45.3.tgz",
"integrity": "sha512-ngKXNhNvwPzF43QqEhDOue7TQTrG09em1sd4HBxVF0Wr2gopAmdEWan+rgbdgK4fhBtSOTJO8bYU4chUG7VXZQ==",
"license": "MIT",
- "peer": true,
"dependencies": {
"@jridgewell/remapping": "^2.3.4",
"@jridgewell/sourcemap-codec": "^1.5.0",
@@ -3014,7 +3006,6 @@
"integrity": "sha512-jl1vZzPDinLr9eUt3J/t7V6FgNEw9QjvBPdysz9KfQDD41fQrC2Y4vKQdiaUpFT4bXlb1RHhLpp8wtm6M5TgSw==",
"dev": true,
"license": "Apache-2.0",
- "peer": true,
"bin": {
"tsc": "bin/tsc",
"tsserver": "bin/tsserver"
@@ -3036,7 +3027,6 @@
"integrity": "sha512-+Oxm7q9hDoLMyJOYfUYBuHQo+dkAloi33apOPP56pzj+vsdJDzr+j1NISE5pyaAuKL4A3UD34qd0lx5+kfKp2g==",
"dev": true,
"license": "MIT",
- "peer": true,
"dependencies": {
"esbuild": "^0.25.0",
"fdir": "^6.4.4",
diff --git a/dashboard/src/lib/stores/app.svelte.ts b/dashboard/src/lib/stores/app.svelte.ts
index 9e3198ae..9fa371fb 100644
--- a/dashboard/src/lib/stores/app.svelte.ts
+++ b/dashboard/src/lib/stores/app.svelte.ts
@@ -173,6 +173,11 @@ export interface PlacementPreviewResponse {
previews: PlacementPreview[];
}
+interface ImageApiResponse {
+ created: number;
+ data: Array<{ b64_json?: string; url?: string }>;
+}
+
interface RawStateResponse {
topology?: RawTopology;
instances?: Record<
@@ -2095,107 +2100,137 @@ class AppStore {
throw new Error(`API error: ${response.status} - ${errorText}`);
}
- const reader = response.body?.getReader();
- if (!reader) {
- throw new Error("No response body");
- }
+ // Streaming requires both stream=true AND partialImages > 0
+ const isStreaming = params.stream && params.partialImages > 0;
- interface ImageGenerationChunk {
- data?: { b64_json?: string };
- format?: string;
- type?: "partial" | "final";
- image_index?: number;
- partial_index?: number;
- total_partials?: number;
- }
+ if (!isStreaming) {
+ // Non-streaming: parse JSON response directly
+ const jsonResponse = (await response.json()) as ImageApiResponse;
+ const format = params.outputFormat || "png";
+ const mimeType = `image/${format}`;
- const numImages = params.numImages;
+ const attachments: MessageAttachment[] = jsonResponse.data
+ .filter((img) => img.b64_json)
+ .map((img, index) => ({
+ type: "generated-image" as const,
+ name: `generated-image-${index + 1}.${format}`,
+ preview: `data:${mimeType};base64,${img.b64_json}`,
+ mimeType,
+ }));
- await this.parseSSEStream<ImageGenerationChunk>(
- reader,
- targetConversationId,
- (parsed) => {
- const imageData = parsed.data?.b64_json;
-
- if (imageData) {
- const format = parsed.format || "png";
- const mimeType = `image/${format}`;
- const imageIndex = parsed.image_index ?? 0;
-
- if (parsed.type === "partial") {
- // Update with partial image and progress
- const partialNum = (parsed.partial_index ?? 0) + 1;
- const totalPartials = parsed.total_partials ?? 3;
- const progressText =
- numImages > 1
- ? `Generating image ${imageIndex + 1}/${numImages}... ${partialNum}/${totalPartials}`
- : `Generating... ${partialNum}/${totalPartials}`;
-
- const partialAttachment: MessageAttachment = {
- type: "generated-image",
- name: `generated-image.${format}`,
- preview: `data:${mimeType};base64,${imageData}`,
- mimeType,
- };
-
- this.updateConversationMessage(
- targetConversationId,
- assistantMessage.id,
- (msg) => {
- msg.content = progressText;
- if (imageIndex === 0) {
- // First image - safe to replace attachments with partial preview
- msg.attachments = [partialAttachment];
- } else {
- // Subsequent images - keep existing finals, show partial at current position
- const existingAttachments = msg.attachments || [];
- // Keep only the completed final images (up to current imageIndex)
- const finals = existingAttachments.slice(0, imageIndex);
- msg.attachments = [...finals, partialAttachment];
- }
- },
- );
- } else if (parsed.type === "final") {
- // Final image - replace partial at this position
- const newAttachment: MessageAttachment = {
- type: "generated-image",
- name: `generated-image-${imageIndex + 1}.${format}`,
- preview: `data:${mimeType};base64,${imageData}`,
- mimeType,
- };
-
- this.updateConversationMessage(
- targetConversationId,
- assistantMessage.id,
- (msg) => {
- if (imageIndex === 0) {
- // First final image - replace any partial preview
- msg.attachments = [newAttachment];
- } else {
- // Subsequent images - keep previous finals, replace partial at current position
- const existingAttachments = msg.attachments || [];
- // Slice keeps indices 0 to imageIndex-1 (the previous final images)
- const previousFinals = existingAttachments.slice(
- 0,
- imageIndex,
- );
- msg.attachments = [...previousFinals, newAttachment];
- }
-
- // Update progress message for multiple images
- if (numImages > 1 && imageIndex < numImages - 1) {
- msg.content = `Generating image ${imageIndex + 2}/${numImages}...`;
- } else {
- msg.content = "";
- }
- },
- );
- }
+ this.updateConversationMessage(
+ targetConversationId,
+ assistantMessage.id,
+ (msg) => {
+ msg.content = "";
+ msg.attachments = attachments;
+ },
+ );
+ this.syncActiveMessagesIfNeeded(targetConversationId);
+ } else {
+ // Streaming mode: use SSE parser
+ const reader = response.body?.getReader();
+ if (!reader) {
+ throw new Error("No response body");
+ }
- this.syncActiveMessagesIfNeeded(targetConversationId);
- }
- },
- );
+ interface ImageGenerationChunk {
+ data?: { b64_json?: string };
+ format?: string;
+ type?: "partial" | "final";
+ image_index?: number;
+ partial_index?: number;
+ total_partials?: number;
+ }
+
+ const numImages = params.numImages;
+
+ await this.parseSSEStream<ImageGenerationChunk>(
+ reader,
+ targetConversationId,
+ (parsed) => {
+ const imageData = parsed.data?.b64_json;
+
+ if (imageData) {
+ const format = parsed.format || "png";
+ const mimeType = `image/${format}`;
+ const imageIndex = parsed.image_index ?? 0;
+
+ if (parsed.type === "partial") {
+ // Update with partial image and progress
+ const partialNum = (parsed.partial_index ?? 0) + 1;
+ const totalPartials = parsed.total_partials ?? 3;
+ const progressText =
+ numImages > 1
+ ? `Generating image ${imageIndex + 1}/${numImages}... ${partialNum}/${totalPartials}`
+ : `Generating... ${partialNum}/${totalPartials}`;
+
+ const partialAttachment: MessageAttachment = {
+ type: "generated-image",
+ name: `generated-image.${format}`,
+ preview: `data:${mimeType};base64,${imageData}`,
+ mimeType,
+ };
+
+ this.updateConversationMessage(
+ targetConversationId,
+ assistantMessage.id,
+ (msg) => {
+ msg.content = progressText;
+ if (imageIndex === 0) {
+ // First image - safe to replace attachments with partial preview
+ msg.attachments = [partialAttachment];
+ } else {
+ // Subsequent images - keep existing finals, show partial at current position
+ const existingAttachments = msg.attachments || [];
+ // Keep only the completed final images (up to current imageIndex)
+ const finals = existingAttachments.slice(0, imageIndex);
+ msg.attachments = [...finals, partialAttachment];
+ }
+ },
+ );
+ } else if (parsed.type === "final") {
+ // Final image - replace partial at this position
+ const newAttachment: MessageAttachment = {
+ type: "generated-image",
+ name: `generated-image-${imageIndex + 1}.${format}`,
+ preview: `data:${mimeType};base64,${imageData}`,
+ mimeType,
+ };
+
+ this.updateConversationMessage(
+ targetConversationId,
+ assistantMessage.id,
+ (msg) => {
+ if (imageIndex === 0) {
+ // First final image - replace any partial preview
+ msg.attachments = [newAttachment];
+ } else {
+ // Subsequent images - keep previous finals, replace partial at current position
+ const existingAttachments = msg.attachments || [];
+ // Slice keeps indices 0 to imageIndex-1 (the previous final images)
+ const previousFinals = existingAttachments.slice(
+ 0,
+ imageIndex,
+ );
+ msg.attachments = [...previousFinals, newAttachment];
+ }
+
+ // Update progress message for multiple images
+ if (numImages > 1 && imageIndex < numImages - 1) {
+ msg.content = `Generating image ${imageIndex + 2}/${numImages}...`;
+ } else {
+ msg.content = "";
+ }
+ },
+ );
+ }
+
+ this.syncActiveMessagesIfNeeded(targetConversationId);
+ }
+ },
+ );
+ }
} catch (error) {
console.error("Error generating image:", error);
this.handleStreamingError(
@@ -2343,69 +2378,98 @@ class AppStore {
throw new Error(`API error: ${apiResponse.status} - ${errorText}`);
}
- const reader = apiResponse.body?.getReader();
- if (!reader) {
- throw new Error("No response body");
- }
+ // Streaming requires both stream=true AND partialImages > 0
+ const isStreaming = params.stream && params.partialImages > 0;
+
+ if (!isStreaming) {
+ // Non-streaming: parse JSON response directly
+ const jsonResponse = (await apiResponse.json()) as ImageApiResponse;
+ const format = params.outputFormat || "png";
+ const mimeType = `image/${format}`;
+ const attachments: MessageAttachment[] = jsonResponse.data
+ .filter((img) => img.b64_json)
+ .map((img) => ({
+ type: "generated-image" as const,
+ name: `edited-image.${format}`,
+ preview: `data:${mimeType};base64,${img.b64_json}`,
+ mimeType,
+ }));
- interface ImageEditChunk {
- data?: { b64_json?: string };
- format?: string;
- type?: "partial" | "final";
- partial_index?: number;
- total_partials?: number;
- }
+ this.updateConversationMessage(
+ targetConversationId,
+ assistantMessage.id,
+ (msg) => {
+ msg.content = "";
+ msg.attachments = attachments;
+ },
+ );
+ this.syncActiveMessagesIfNeeded(targetConversationId);
+ } else {
+ // Streaming mode: use SSE parser
+ const reader = apiResponse.body?.getReader();
+ if (!reader) {
+ throw new Error("No response body");
+ }
- await this.parseSSEStream<ImageEditChunk>(
- reader,
- targetConversationId,
- (parsed) => {
- const imageData = parsed.data?.b64_json;
-
- if (imageData) {
- const format = parsed.format || "png";
- const mimeType = `image/${format}`;
- if (parsed.type === "partial") {
- // Update with partial image and progress
- const partialNum = (parsed.partial_index ?? 0) + 1;
- const totalPartials = parsed.total_partials ?? 3;
- this.updateConversationMessage(
- targetConversationId,
- assistantMessage.id,
- (msg) => {
- msg.content = `Editing... ${partialNum}/${totalPartials}`;
- msg.attachments = [
- {
- type: "generated-image",
- name: `edited-image.${format}`,
- preview: `data:${mimeType};base64,${imageData}`,
- mimeType,
- },
- ];
- },
- );
- } else if (parsed.type === "final") {
- // Final image
- this.updateConversationMessage(
- targetConversationId,
- assistantMessage.id,
- (msg) => {
- msg.content = "";
- msg.attachments = [
- {
- type: "generated-image",
- name: `edited-image.${format}`,
- preview: `data:${mimeType};base64,${imageData}`,
- mimeType,
- },
- ];
- },
- );
+ interface ImageEditChunk {
+ data?: { b64_json?: string };
+ format?: string;
+ type?: "partial" | "final";
+ partial_index?: number;
+ total_partials?: number;
+ }
+
+ await this.parseSSEStream<ImageEditChunk>(
+ reader,
+ targetConversationId,
+ (parsed) => {
+ const imageData = parsed.data?.b64_json;
+
+ if (imageData) {
+ const format = parsed.format || "png";
+ const mimeType = `image/${format}`;
+ if (parsed.type === "partial") {
+ // Update with partial image and progress
+ const partialNum = (parsed.partial_index ?? 0) + 1;
+ const totalPartials = parsed.total_partials ?? 3;
+ this.updateConversationMessage(
+ targetConversationId,
+ assistantMessage.id,
+ (msg) => {
+ msg.content = `Editing... ${partialNum}/${totalPartials}`;
+ msg.attachments = [
+ {
+ type: "generated-image",
+ name: `edited-image.${format}`,
+ preview: `data:${mimeType};base64,${imageData}`,
+ mimeType,
+ },
+ ];
+ },
+ );
+ } else if (parsed.type === "final") {
+ // Final image
+ this.updateConversationMessage(
+ targetConversationId,
+ assistantMessage.id,
+ (msg) => {
+ msg.content = "";
+ msg.attachments = [
+ {
+ type: "generated-image",
+ name: `edited-image.${format}`,
+ preview: `data:${mimeType};base64,${imageData}`,
+ mimeType,
+ },
+ ];
+ },
+ );
+ }
+ this.syncActiveMessagesIfNeeded(targetConversationId);
}
- this.syncActiveMessagesIfNeeded(targetConversationId);
- }
- },
- );
+ },
+ );
+ }
} catch (error) {
console.error("Error editing image:", error);
this.handleStreamingError(
diff --git a/src/exo/worker/engines/image/generate.py b/src/exo/worker/engines/image/generate.py
index 66f21a1e..014b3b2f 100644
--- a/src/exo/worker/engines/image/generate.py
+++ b/src/exo/worker/engines/image/generate.py
@@ -98,8 +98,8 @@ def generate_image(
partial_images = (
task.partial_images
- if task.partial_images is not None
- else (3 if task.stream else 0)
+ if task.partial_images is not None and task.stream is not None and task.stream
+ else 0
)
image_path: Path | None = None
diff --git a/src/exo/worker/engines/image/pipeline/runner.py b/src/exo/worker/engines/image/pipeline/runner.py
index cd5d92d5..698887b0 100644
--- a/src/exo/worker/engines/image/pipeline/runner.py
+++ b/src/exo/worker/engines/image/pipeline/runner.py
@@ -348,6 +348,7 @@ class DiffusionRunner:
ctx.in_loop( # pyright: ignore[reportAny]
t=t,
latents=latents,
+ time_steps=time_steps,
)
mx.eval(latents)
← 7cc313b2 Treat Swift/Xcode build warnings as errors (#1322)
·
back to Exo
·
Fix bench after recent updates (#1331) 9dabde7e →