← back to Exo
feat: initial image generation support (#1095)
307f454b961ac0227a6004909af966ffc1feccfa · 2026-01-21 18:21:58 +0000 · ciaranbor
## Motivation
Enable distributed image generation across exo clusters
## Changes
- Added OpenAI-compatible /v1/images/generations and /v1/images/edits
API endpoints - Added /bench/images/generations and /bench/images/edits
endpoints that return generation statistics (timing, throughput metrics)
- Implemented PipeFusion distributed inference for diffusion models,
enabling patch-based parallelism across nodes
- Added model adapters for Flux (schnell, dev) and Qwen image models
## Why It Works
https://arxiv.org/abs/2405.14430
## Test Plan
### Manual Testing
- Generate images using /v1/images/generations endpoint with single and
multi-node clusters
- Test image editing via /v1/images/edits with source images
- Verify streaming partial images appear progressively in the dashboard
- Use /bench/images/generations to measure generation performance
- Test both Flux and Qwen model families
---------
Co-authored-by: Sami Khan <smsak99@gmail.com>
Files touched
M dashboard/package-lock.jsonM dashboard/src/lib/components/ChatForm.svelteM dashboard/src/lib/components/ChatMessages.svelteA dashboard/src/lib/components/ImageParamsPanel.svelteM dashboard/src/lib/components/index.tsM dashboard/src/lib/stores/app.svelte.tsM dashboard/src/routes/+page.svelteM docs/api.mdM src/exo/master/api.pyA src/exo/master/image_store.pyM src/exo/master/main.pyM src/exo/master/tests/test_api_error_handling.pyM src/exo/master/tests/test_master.pyM src/exo/master/tests/test_placement.pyM src/exo/master/tests/test_placement_utils.pyM src/exo/shared/apply.pyM src/exo/shared/constants.pyM src/exo/shared/models/model_cards.pyM src/exo/shared/tests/conftest.pyM src/exo/shared/types/api.pyM src/exo/shared/types/chunks.pyM src/exo/shared/types/commands.pyM src/exo/shared/types/events.pyM src/exo/shared/types/tasks.pyM src/exo/shared/types/worker/runner_response.pyM src/exo/worker/download/download_utils.pyM src/exo/worker/download/huggingface_utils.pyM src/exo/worker/download/shard_downloader.pyA src/exo/worker/engines/image/__init__.pyA src/exo/worker/engines/image/config.pyA src/exo/worker/engines/image/distributed_model.pyA src/exo/worker/engines/image/generate.pyA src/exo/worker/engines/image/models/__init__.pyA src/exo/worker/engines/image/models/base.pyA src/exo/worker/engines/image/models/flux/__init__.pyA src/exo/worker/engines/image/models/flux/adapter.pyA src/exo/worker/engines/image/models/flux/config.pyA src/exo/worker/engines/image/models/flux/wrappers.pyA src/exo/worker/engines/image/models/qwen/__init__.pyA src/exo/worker/engines/image/models/qwen/adapter.pyA src/exo/worker/engines/image/models/qwen/config.pyA src/exo/worker/engines/image/models/qwen/edit_adapter.pyA src/exo/worker/engines/image/models/qwen/wrappers.pyA src/exo/worker/engines/image/pipeline/__init__.pyA src/exo/worker/engines/image/pipeline/block_wrapper.pyA src/exo/worker/engines/image/pipeline/kv_cache.pyA src/exo/worker/engines/image/pipeline/runner.pyM src/exo/worker/engines/mlx/generator/generate.pyM src/exo/worker/main.pyM src/exo/worker/plan.pyM src/exo/worker/runner/runner.pyM src/exo/worker/tests/unittests/conftest.pyM src/exo/worker/tests/unittests/test_mlx/conftest.py
Diff
commit 307f454b961ac0227a6004909af966ffc1feccfa
Author: ciaranbor <81697641+ciaranbor@users.noreply.github.com>
Date: Wed Jan 21 18:21:58 2026 +0000
feat: initial image generation support (#1095)
## Motivation
Enable distributed image generation across exo clusters
## Changes
- Added OpenAI-compatible /v1/images/generations and /v1/images/edits
API endpoints - Added /bench/images/generations and /bench/images/edits
endpoints that return generation statistics (timing, throughput metrics)
- Implemented PipeFusion distributed inference for diffusion models,
enabling patch-based parallelism across nodes
- Added model adapters for Flux (schnell, dev) and Qwen image models
## Why It Works
https://arxiv.org/abs/2405.14430
## Test Plan
### Manual Testing
- Generate images using /v1/images/generations endpoint with single and
multi-node clusters
- Test image editing via /v1/images/edits with source images
- Verify streaming partial images appear progressively in the dashboard
- Use /bench/images/generations to measure generation performance
- Test both Flux and Qwen model families
---------
Co-authored-by: Sami Khan <smsak99@gmail.com>
---
dashboard/package-lock.json | 9 +
dashboard/src/lib/components/ChatForm.svelte | 485 +++++++---
dashboard/src/lib/components/ChatMessages.svelte | 83 +-
.../src/lib/components/ImageParamsPanel.svelte | 502 +++++++++++
dashboard/src/lib/components/index.ts | 1 +
dashboard/src/lib/stores/app.svelte.ts | 465 +++++++++-
dashboard/src/routes/+page.svelte | 54 +-
docs/api.md | 77 +-
src/exo/master/api.py | 539 +++++++++++-
src/exo/master/image_store.py | 79 ++
src/exo/master/main.py | 104 ++-
src/exo/master/tests/test_api_error_handling.py | 37 +-
src/exo/master/tests/test_master.py | 4 +-
src/exo/master/tests/test_placement.py | 6 +-
src/exo/master/tests/test_placement_utils.py | 4 +-
src/exo/shared/apply.py | 5 +-
src/exo/shared/constants.py | 4 +
src/exo/shared/models/model_cards.py | 217 ++++-
src/exo/shared/tests/conftest.py | 3 +-
src/exo/shared/types/api.py | 129 ++-
src/exo/shared/types/chunks.py | 38 +-
src/exo/shared/types/commands.py | 24 +-
src/exo/shared/types/events.py | 8 +-
src/exo/shared/types/tasks.py | 24 +-
src/exo/shared/types/worker/runner_response.py | 32 +-
src/exo/worker/download/download_utils.py | 36 +-
src/exo/worker/download/huggingface_utils.py | 68 +-
src/exo/worker/download/shard_downloader.py | 3 +-
src/exo/worker/engines/image/__init__.py | 12 +
src/exo/worker/engines/image/config.py | 50 ++
src/exo/worker/engines/image/distributed_model.py | 166 ++++
src/exo/worker/engines/image/generate.py | 176 ++++
src/exo/worker/engines/image/models/__init__.py | 86 ++
src/exo/worker/engines/image/models/base.py | 295 +++++++
.../worker/engines/image/models/flux/__init__.py | 11 +
.../worker/engines/image/models/flux/adapter.py | 215 +++++
src/exo/worker/engines/image/models/flux/config.py | 34 +
.../worker/engines/image/models/flux/wrappers.py | 279 ++++++
.../worker/engines/image/models/qwen/__init__.py | 13 +
.../worker/engines/image/models/qwen/adapter.py | 292 +++++++
src/exo/worker/engines/image/models/qwen/config.py | 29 +
.../engines/image/models/qwen/edit_adapter.py | 434 +++++++++
.../worker/engines/image/models/qwen/wrappers.py | 200 +++++
src/exo/worker/engines/image/pipeline/__init__.py | 15 +
.../worker/engines/image/pipeline/block_wrapper.py | 303 +++++++
src/exo/worker/engines/image/pipeline/kv_cache.py | 72 ++
src/exo/worker/engines/image/pipeline/runner.py | 972 +++++++++++++++++++++
src/exo/worker/engines/mlx/generator/generate.py | 2 +-
src/exo/worker/main.py | 62 +-
src/exo/worker/plan.py | 20 +-
src/exo/worker/runner/runner.py | 273 +++++-
src/exo/worker/tests/unittests/conftest.py | 3 +-
.../worker/tests/unittests/test_mlx/conftest.py | 4 +-
53 files changed, 6856 insertions(+), 202 deletions(-)
diff --git a/dashboard/package-lock.json b/dashboard/package-lock.json
index 1768cb96..fddf99de 100644
--- a/dashboard/package-lock.json
+++ b/dashboard/package-lock.json
@@ -863,6 +863,7 @@
"integrity": "sha512-oH8tXw7EZnie8FdOWYrF7Yn4IKrqTFHhXvl8YxXxbKwTMcD/5NNCryUSEXRk2ZR4ojnub0P8rNrsVGHXWqIDtA==",
"dev": true,
"license": "MIT",
+ "peer": true,
"dependencies": {
"@standard-schema/spec": "^1.0.0",
"@sveltejs/acorn-typescript": "^1.0.5",
@@ -902,6 +903,7 @@
"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",
@@ -1518,6 +1520,7 @@
"integrity": "sha512-LCCV0HdSZZZb34qifBsyWlUmok6W7ouER+oQIGBScS8EsZsQbrtFTUrDX4hOl+CS6p7cnNC4td+qrSVGSCTUfQ==",
"dev": true,
"license": "MIT",
+ "peer": true,
"dependencies": {
"undici-types": "~6.21.0"
}
@@ -1527,6 +1530,7 @@
"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"
},
@@ -1939,6 +1943,7 @@
"integrity": "sha512-fmTRWbNMmsmWq6xJV8D19U/gw/bwrHfNXxrIN+HfZgnzqTHp9jOmKMhsTUjXOJnZOdZY9Q28y4yebKzqDKlxlQ==",
"dev": true,
"license": "ISC",
+ "peer": true,
"engines": {
"node": ">=12"
}
@@ -2646,6 +2651,7 @@
"integrity": "sha512-5gTmgEY/sqK6gFXLIsQNH19lWb4ebPDLA4SdLP7dsWkIXHWlG66oPuVvXSGFPppYZz8ZDZq0dYYrbHfBCVUb1Q==",
"dev": true,
"license": "MIT",
+ "peer": true,
"engines": {
"node": ">=12"
},
@@ -2833,6 +2839,7 @@
"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",
@@ -2977,6 +2984,7 @@
"integrity": "sha512-jl1vZzPDinLr9eUt3J/t7V6FgNEw9QjvBPdysz9KfQDD41fQrC2Y4vKQdiaUpFT4bXlb1RHhLpp8wtm6M5TgSw==",
"dev": true,
"license": "Apache-2.0",
+ "peer": true,
"bin": {
"tsc": "bin/tsc",
"tsserver": "bin/tsserver"
@@ -2998,6 +3006,7 @@
"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/components/ChatForm.svelte b/dashboard/src/lib/components/ChatForm.svelte
index 269c01f2..d1648480 100644
--- a/dashboard/src/lib/components/ChatForm.svelte
+++ b/dashboard/src/lib/components/ChatForm.svelte
@@ -1,8 +1,22 @@
<script lang="ts">
- import { isLoading, sendMessage, selectedChatModel, setSelectedChatModel, instances, ttftMs, tps, totalTokens } from '$lib/stores/app.svelte';
- import ChatAttachments from './ChatAttachments.svelte';
- import type { ChatUploadedFile } from '$lib/types/files';
- import { processUploadedFiles, getAcceptString } from '$lib/types/files';
+ import {
+ isLoading,
+ sendMessage,
+ generateImage,
+ editImage,
+ editingImage,
+ clearEditingImage,
+ selectedChatModel,
+ setSelectedChatModel,
+ instances,
+ ttftMs,
+ tps,
+ totalTokens,
+ } from "$lib/stores/app.svelte";
+ import ChatAttachments from "./ChatAttachments.svelte";
+ import ImageParamsPanel from "./ImageParamsPanel.svelte";
+ import type { ChatUploadedFile } from "$lib/types/files";
+ import { processUploadedFiles, getAcceptString } from "$lib/types/files";
interface Props {
class?: string;
@@ -10,17 +24,19 @@
showHelperText?: boolean;
autofocus?: boolean;
showModelSelector?: boolean;
+ modelTasks?: Record<string, string[]>;
}
- let {
- class: className = '',
- placeholder = 'Ask anything',
+ let {
+ class: className = "",
+ placeholder = "Ask anything",
showHelperText = false,
autofocus = true,
- showModelSelector = false
+ showModelSelector = false,
+ modelTasks = {},
}: Props = $props();
- let message = $state('');
+ let message = $state("");
let textareaRef: HTMLTextAreaElement | undefined = $state();
let fileInputRef: HTMLInputElement | undefined = $state();
let uploadedFiles = $state<ChatUploadedFile[]>([]);
@@ -31,30 +47,82 @@
const currentTtft = $derived(ttftMs());
const currentTps = $derived(tps());
const currentTokens = $derived(totalTokens());
-
+ const currentEditingImage = $derived(editingImage());
+ const isEditMode = $derived(currentEditingImage !== null);
+
// Custom dropdown state
let isModelDropdownOpen = $state(false);
let dropdownButtonRef: HTMLButtonElement | undefined = $state();
let dropdownPosition = $derived(() => {
- if (!dropdownButtonRef || !isModelDropdownOpen) return { top: 0, left: 0, width: 0 };
+ if (!dropdownButtonRef || !isModelDropdownOpen)
+ return { top: 0, left: 0, width: 0 };
const rect = dropdownButtonRef.getBoundingClientRect();
return {
top: rect.top,
left: rect.left,
- width: rect.width
+ width: rect.width,
};
});
// Accept all supported file types
- const acceptString = getAcceptString(['image', 'text', 'pdf']);
+ const acceptString = getAcceptString(["image", "text", "pdf"]);
+
+ function modelSupportsImageGeneration(modelId: string): boolean {
+ const tasks = modelTasks[modelId] || [];
+ return tasks.includes("TextToImage") || tasks.includes("ImageToImage");
+ }
+
+ function modelSupportsTextToImage(modelId: string): boolean {
+ const tasks = modelTasks[modelId] || [];
+ return tasks.includes("TextToImage");
+ }
+
+ function modelSupportsOnlyImageEditing(modelId: string): boolean {
+ const tasks = modelTasks[modelId] || [];
+ return tasks.includes("ImageToImage") && !tasks.includes("TextToImage");
+ }
+
+ function modelSupportsImageEditing(modelId: string): boolean {
+ const tasks = modelTasks[modelId] || [];
+ return tasks.includes("ImageToImage");
+ }
+
+ const isImageModel = $derived(() => {
+ if (!currentModel) return false;
+ return modelSupportsTextToImage(currentModel);
+ });
+
+ const isEditOnlyWithoutImage = $derived(
+ currentModel !== null &&
+ modelSupportsOnlyImageEditing(currentModel) &&
+ !isEditMode &&
+ uploadedFiles.length === 0,
+ );
+
+ // Show edit mode when: explicit edit mode OR (model supports ImageToImage AND files attached)
+ const shouldShowEditMode = $derived(
+ isEditMode ||
+ (currentModel &&
+ modelSupportsImageEditing(currentModel) &&
+ uploadedFiles.length > 0),
+ );
// Extract available models from running instances
const availableModels = $derived(() => {
- const models: Array<{id: string, label: string}> = [];
+ const models: Array<{ id: string; label: string; isImageModel: boolean }> =
+ [];
for (const [, instance] of Object.entries(instanceData)) {
const modelId = getInstanceModelId(instance);
- if (modelId && modelId !== 'Unknown' && !models.some(m => m.id === modelId)) {
- models.push({ id: modelId, label: modelId.split('/').pop() || modelId });
+ if (
+ modelId &&
+ modelId !== "Unknown" &&
+ !models.some((m) => m.id === modelId)
+ ) {
+ models.push({
+ id: modelId,
+ label: modelId.split("/").pop() || modelId,
+ isImageModel: modelSupportsImageGeneration(modelId),
+ });
}
}
return models;
@@ -66,18 +134,18 @@
// Auto-select the first available model if none is selected, if current selection is stale, or if a new model is added
$effect(() => {
const models = availableModels();
- const currentModelIds = new Set(models.map(m => m.id));
+ const currentModelIds = new Set(models.map((m) => m.id));
if (models.length > 0) {
// Find newly added models (in current but not in previous)
- const newModels = models.filter(m => !previousModelIds.has(m.id));
+ const newModels = models.filter((m) => !previousModelIds.has(m.id));
// If no model selected, select the first available
if (!currentModel) {
setSelectedChatModel(models[0].id);
}
// If current model is stale (no longer has a running instance), reset to first available
- else if (!models.some(m => m.id === currentModel)) {
+ else if (!models.some((m) => m.id === currentModel)) {
setSelectedChatModel(models[0].id);
}
// If a new model was just added, select it
@@ -87,7 +155,7 @@
} else {
// No instances running - clear the selected model
if (currentModel) {
- setSelectedChatModel('');
+ setSelectedChatModel("");
}
}
@@ -96,13 +164,15 @@
});
function getInstanceModelId(instanceWrapped: unknown): string {
- if (!instanceWrapped || typeof instanceWrapped !== 'object') return '';
+ if (!instanceWrapped || typeof instanceWrapped !== "object") return "";
const keys = Object.keys(instanceWrapped as Record<string, unknown>);
if (keys.length === 1) {
- const instance = (instanceWrapped as Record<string, unknown>)[keys[0]] as { shardAssignments?: { modelId?: string } };
- return instance?.shardAssignments?.modelId || '';
+ const instance = (instanceWrapped as Record<string, unknown>)[
+ keys[0]
+ ] as { shardAssignments?: { modelId?: string } };
+ return instance?.shardAssignments?.modelId || "";
}
- return '';
+ return "";
}
async function handleFiles(files: File[]) {
@@ -115,33 +185,35 @@
const input = event.target as HTMLInputElement;
if (input.files && input.files.length > 0) {
handleFiles(Array.from(input.files));
- input.value = ''; // Reset for next selection
+ input.value = ""; // Reset for next selection
}
}
function handleFileRemove(fileId: string) {
- uploadedFiles = uploadedFiles.filter(f => f.id !== fileId);
+ uploadedFiles = uploadedFiles.filter((f) => f.id !== fileId);
}
function handlePaste(event: ClipboardEvent) {
if (!event.clipboardData) return;
-
+
const files = Array.from(event.clipboardData.items)
- .filter(item => item.kind === 'file')
- .map(item => item.getAsFile())
+ .filter((item) => item.kind === "file")
+ .map((item) => item.getAsFile())
.filter((file): file is File => file !== null);
-
+
if (files.length > 0) {
event.preventDefault();
handleFiles(files);
return;
}
-
+
// Handle long text paste as file
- const text = event.clipboardData.getData('text/plain');
+ const text = event.clipboardData.getData("text/plain");
if (text.length > 2500) {
event.preventDefault();
- const textFile = new File([text], 'pasted-text.txt', { type: 'text/plain' });
+ const textFile = new File([text], "pasted-text.txt", {
+ type: "text/plain",
+ });
handleFiles([textFile]);
}
}
@@ -159,7 +231,7 @@
function handleDrop(event: DragEvent) {
event.preventDefault();
isDragOver = false;
-
+
if (event.dataTransfer?.files) {
handleFiles(Array.from(event.dataTransfer.files));
}
@@ -170,8 +242,8 @@
if (event.isComposing || event.keyCode === 229) {
return;
}
-
- if (event.key === 'Enter' && !event.shiftKey) {
+
+ if (event.key === "Enter" && !event.shiftKey) {
event.preventDefault();
handleSubmit();
}
@@ -179,29 +251,50 @@
function handleSubmit() {
if ((!message.trim() && uploadedFiles.length === 0) || loading) return;
-
+
const content = message.trim();
const files = [...uploadedFiles];
-
- message = '';
+
+ message = "";
uploadedFiles = [];
resetTextareaHeight();
-
- sendMessage(content, files);
-
+
+ // Use image editing if in edit mode
+ if (isEditMode && currentEditingImage && content) {
+ editImage(content, currentEditingImage.imageDataUrl);
+ }
+ // If user attached an image with an ImageToImage model, use edit endpoint
+ else if (
+ currentModel &&
+ modelSupportsImageEditing(currentModel) &&
+ files.length > 0 &&
+ content
+ ) {
+ // Use the first attached image for editing
+ const imageFile = files[0];
+ if (imageFile.preview) {
+ editImage(content, imageFile.preview);
+ }
+ } else if (isImageModel() && content) {
+ // Use image generation for text-to-image models
+ generateImage(content);
+ } else {
+ sendMessage(content, files);
+ }
+
// Refocus the textarea after sending
setTimeout(() => textareaRef?.focus(), 10);
}
function handleInput() {
if (!textareaRef) return;
- textareaRef.style.height = 'auto';
- textareaRef.style.height = Math.min(textareaRef.scrollHeight, 150) + 'px';
+ textareaRef.style.height = "auto";
+ textareaRef.style.height = Math.min(textareaRef.scrollHeight, 150) + "px";
}
function resetTextareaHeight() {
if (textareaRef) {
- textareaRef.style.height = 'auto';
+ textareaRef.style.height = "auto";
}
}
@@ -211,13 +304,13 @@
// Track previous loading state to detect when loading completes
let wasLoading = $state(false);
-
+
$effect(() => {
if (autofocus && textareaRef) {
setTimeout(() => textareaRef?.focus(), 10);
}
});
-
+
// Refocus after loading completes (AI response finished)
$effect(() => {
if (wasLoading && !loading && textareaRef) {
@@ -226,7 +319,9 @@
wasLoading = loading;
});
- const canSend = $derived(message.trim().length > 0 || uploadedFiles.length > 0);
+ const canSend = $derived(
+ message.trim().length > 0 || uploadedFiles.length > 0,
+ );
</script>
<!-- Hidden file input -->
@@ -239,69 +334,132 @@
onchange={handleFileInputChange}
/>
-<form
- onsubmit={(e) => { e.preventDefault(); handleSubmit(); }}
+<form
+ onsubmit={(e) => {
+ e.preventDefault();
+ handleSubmit();
+ }}
class="w-full {className}"
ondragover={handleDragOver}
ondragleave={handleDragLeave}
ondrop={handleDrop}
>
- <div
- class="relative command-panel rounded overflow-hidden transition-all duration-200 {isDragOver ? 'ring-2 ring-exo-yellow ring-opacity-50' : ''}"
+ <div
+ class="relative command-panel rounded overflow-hidden transition-all duration-200 {isDragOver
+ ? 'ring-2 ring-exo-yellow ring-opacity-50'
+ : ''}"
>
<!-- Top accent line -->
- <div class="absolute top-0 left-0 right-0 h-px bg-gradient-to-r from-transparent via-exo-yellow/50 to-transparent"></div>
-
+ <div
+ class="absolute top-0 left-0 right-0 h-px bg-gradient-to-r from-transparent via-exo-yellow/50 to-transparent"
+ ></div>
+
<!-- Drag overlay -->
{#if isDragOver}
- <div class="absolute inset-0 bg-exo-dark-gray/80 z-10 flex items-center justify-center">
+ <div
+ class="absolute inset-0 bg-exo-dark-gray/80 z-10 flex items-center justify-center"
+ >
<div class="text-exo-yellow text-sm font-mono tracking-wider uppercase">
DROP FILES HERE
</div>
</div>
{/if}
-
+
+ <!-- Edit mode banner -->
+ {#if isEditMode && currentEditingImage}
+ <div
+ class="flex items-center gap-3 px-3 py-2 bg-exo-yellow/10 border-b border-exo-yellow/30"
+ >
+ <img
+ src={currentEditingImage.imageDataUrl}
+ alt="Source for editing"
+ class="w-10 h-10 object-cover rounded border border-exo-yellow/30"
+ />
+ <div class="flex-1">
+ <span
+ class="text-xs font-mono tracking-wider uppercase text-exo-yellow"
+ >EDITING IMAGE</span
+ >
+ </div>
+ <button
+ type="button"
+ onclick={() => clearEditingImage()}
+ class="px-2 py-1 text-xs font-mono tracking-wider uppercase bg-exo-medium-gray/30 text-exo-light-gray border border-exo-medium-gray/50 rounded hover:bg-exo-medium-gray/50 hover:text-exo-yellow transition-colors cursor-pointer"
+ >
+ CANCEL
+ </button>
+ </div>
+ {/if}
+
<!-- Model selector (when enabled) -->
{#if showModelSelector && availableModels().length > 0}
- <div class="flex items-center justify-between gap-2 px-3 py-2 border-b border-exo-medium-gray/30">
+ <div
+ class="flex items-center justify-between gap-2 px-3 py-2 border-b border-exo-medium-gray/30"
+ >
<div class="flex items-center gap-2 flex-1">
- <span class="text-xs text-exo-light-gray uppercase tracking-wider flex-shrink-0">MODEL:</span>
+ <span
+ class="text-xs text-exo-light-gray uppercase tracking-wider flex-shrink-0"
+ >MODEL:</span
+ >
<!-- Custom dropdown -->
<div class="relative flex-1 max-w-xs">
<button
bind:this={dropdownButtonRef}
type="button"
- onclick={() => isModelDropdownOpen = !isModelDropdownOpen}
- class="w-full bg-exo-medium-gray/50 border border-exo-yellow/30 rounded pl-3 pr-8 py-1.5 text-xs font-mono text-left tracking-wide cursor-pointer transition-all duration-200 hover:border-exo-yellow/50 focus:outline-none focus:border-exo-yellow/70 {isModelDropdownOpen ? 'border-exo-yellow/70' : ''}"
+ onclick={() => (isModelDropdownOpen = !isModelDropdownOpen)}
+ class="w-full bg-exo-medium-gray/50 border border-exo-yellow/30 rounded pl-3 pr-8 py-1.5 text-xs font-mono text-left tracking-wide cursor-pointer transition-all duration-200 hover:border-exo-yellow/50 focus:outline-none focus:border-exo-yellow/70 {isModelDropdownOpen
+ ? 'border-exo-yellow/70'
+ : ''}"
>
- {#if availableModels().find(m => m.id === currentModel)}
- <span class="text-exo-yellow truncate">{availableModels().find(m => m.id === currentModel)?.label}</span>
+ {#if availableModels().find((m) => m.id === currentModel)}
+ <span class="text-exo-yellow truncate"
+ >{availableModels().find((m) => m.id === currentModel)
+ ?.label}</span
+ >
{:else if availableModels().length > 0}
- <span class="text-exo-yellow truncate">{availableModels()[0].label}</span>
+ <span class="text-exo-yellow truncate"
+ >{availableModels()[0].label}</span
+ >
{:else}
<span class="text-exo-light-gray/50">— SELECT MODEL —</span>
{/if}
</button>
- <div class="absolute right-2 top-1/2 -translate-y-1/2 pointer-events-none transition-transform duration-200 {isModelDropdownOpen ? 'rotate-180' : ''}">
- <svg class="w-3 h-3 text-exo-yellow/60" fill="none" viewBox="0 0 24 24" stroke="currentColor">
- <path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M19 9l-7 7-7-7" />
+ <div
+ class="absolute right-2 top-1/2 -translate-y-1/2 pointer-events-none transition-transform duration-200 {isModelDropdownOpen
+ ? 'rotate-180'
+ : ''}"
+ >
+ <svg
+ class="w-3 h-3 text-exo-yellow/60"
+ fill="none"
+ viewBox="0 0 24 24"
+ stroke="currentColor"
+ >
+ <path
+ stroke-linecap="round"
+ stroke-linejoin="round"
+ stroke-width="2"
+ d="M19 9l-7 7-7-7"
+ />
</svg>
</div>
</div>
-
+
{#if isModelDropdownOpen}
<!-- Backdrop to close dropdown -->
- <button
+ <button
type="button"
- class="fixed inset-0 z-[9998] cursor-default"
- onclick={() => isModelDropdownOpen = false}
+ class="fixed inset-0 z-[9998] cursor-default"
+ onclick={() => (isModelDropdownOpen = false)}
aria-label="Close dropdown"
></button>
-
+
<!-- Dropdown Panel - fixed positioning to escape overflow:hidden -->
- <div
+ <div
class="fixed bg-exo-dark-gray border border-exo-yellow/30 rounded shadow-lg shadow-black/50 z-[9999] max-h-48 overflow-y-auto"
- style="bottom: calc(100vh - {dropdownPosition().top}px + 4px); left: {dropdownPosition().left}px; width: {dropdownPosition().width}px;"
+ style="bottom: calc(100vh - {dropdownPosition()
+ .top}px + 4px); left: {dropdownPosition()
+ .left}px; width: {dropdownPosition().width}px;"
>
<div class="py-1">
{#each availableModels() as model}
@@ -311,20 +469,48 @@
setSelectedChatModel(model.id);
isModelDropdownOpen = false;
}}
- class="w-full px-3 py-2 text-left text-xs font-mono tracking-wide transition-colors duration-100 flex items-center gap-2 {
- currentModel === model.id
- ? 'bg-transparent text-exo-yellow'
- : 'text-exo-light-gray hover:text-exo-yellow'
- }"
+ class="w-full px-3 py-2 text-left text-xs font-mono tracking-wide transition-colors duration-100 flex items-center gap-2 {currentModel ===
+ model.id
+ ? 'bg-transparent text-exo-yellow'
+ : 'text-exo-light-gray hover:text-exo-yellow'}"
>
{#if currentModel === model.id}
- <svg class="w-3 h-3 flex-shrink-0" fill="currentColor" viewBox="0 0 20 20">
- <path fill-rule="evenodd" d="M16.707 5.293a1 1 0 010 1.414l-8 8a1 1 0 01-1.414 0l-4-4a1 1 0 011.414-1.414L8 12.586l7.293-7.293a1 1 0 011.414 0z" clip-rule="evenodd" />
+ <svg
+ class="w-3 h-3 flex-shrink-0"
+ fill="currentColor"
+ viewBox="0 0 20 20"
+ >
+ <path
+ fill-rule="evenodd"
+ d="M16.707 5.293a1 1 0 010 1.414l-8 8a1 1 0 01-1.414 0l-4-4a1 1 0 011.414-1.414L8 12.586l7.293-7.293a1 1 0 011.414 0z"
+ clip-rule="evenodd"
+ />
</svg>
{:else}
<span class="w-3"></span>
{/if}
- <span class="truncate">{model.label}</span>
+ {#if model.isImageModel}
+ <svg
+ class="w-3.5 h-3.5 flex-shrink-0 text-exo-yellow"
+ fill="none"
+ viewBox="0 0 24 24"
+ stroke="currentColor"
+ stroke-width="2"
+ aria-label="Image generation model"
+ >
+ <rect
+ x="3"
+ y="3"
+ width="18"
+ height="18"
+ rx="2"
+ ry="2"
+ />
+ <circle cx="8.5" cy="8.5" r="1.5" />
+ <polyline points="21 15 16 10 5 21" />
+ </svg>
+ {/if}
+ <span class="truncate flex-1">{model.label}</span>
</button>
{/each}
</div>
@@ -336,30 +522,37 @@
<div class="flex items-center gap-4 text-xs font-mono flex-shrink-0">
{#if currentTtft !== null}
<span class="text-exo-light-gray">
- <span class="text-white/70">TTFT</span> <span class="text-exo-yellow">{currentTtft.toFixed(1)}ms</span>
+ <span class="text-white/70">TTFT</span>
+ <span class="text-exo-yellow">{currentTtft.toFixed(1)}ms</span>
</span>
{/if}
{#if currentTps !== null}
<span class="text-exo-light-gray">
- <span class="text-white/70">TPS</span> <span class="text-exo-yellow">{currentTps.toFixed(1)}</span> <span class="text-white/60">tok/s</span>
- <span class="text-white/50">({(1000 / currentTps).toFixed(1)} ms/tok)</span>
+ <span class="text-white/70">TPS</span>
+ <span class="text-exo-yellow">{currentTps.toFixed(1)}</span>
+ <span class="text-white/60">tok/s</span>
+ <span class="text-white/50"
+ >({(1000 / currentTps).toFixed(1)} ms/tok)</span
+ >
</span>
{/if}
</div>
{/if}
</div>
{/if}
-
+
+ <!-- Image params panel (shown for image models or edit mode) -->
+ {#if showModelSelector && (isImageModel() || isEditMode)}
+ <ImageParamsPanel {isEditMode} />
+ {/if}
+
<!-- Attached files preview -->
{#if uploadedFiles.length > 0}
<div class="px-3 pt-3">
- <ChatAttachments
- files={uploadedFiles}
- onRemove={handleFileRemove}
- />
+ <ChatAttachments files={uploadedFiles} onRemove={handleFileRemove} />
</div>
{/if}
-
+
<!-- Input area -->
<div class="flex items-start gap-2 sm:gap-3 py-3 px-3 sm:px-4">
<!-- Attach file button -->
@@ -370,58 +563,130 @@
class="flex items-center justify-center w-7 h-7 rounded text-exo-light-gray hover:text-exo-yellow transition-all disabled:opacity-50 disabled:cursor-not-allowed flex-shrink-0 cursor-pointer"
title="Attach file"
>
- <svg class="w-4 h-4" fill="none" viewBox="0 0 24 24" stroke="currentColor">
- <path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M15.172 7l-6.586 6.586a2 2 0 102.828 2.828l6.414-6.586a4 4 0 00-5.656-5.656l-6.415 6.585a6 6 0 108.486 8.486L20.5 13" />
+ <svg
+ class="w-4 h-4"
+ fill="none"
+ viewBox="0 0 24 24"
+ stroke="currentColor"
+ >
+ <path
+ stroke-linecap="round"
+ stroke-linejoin="round"
+ stroke-width="2"
+ d="M15.172 7l-6.586 6.586a2 2 0 102.828 2.828l6.414-6.586a4 4 0 00-5.656-5.656l-6.415 6.585a6 6 0 108.486 8.486L20.5 13"
+ />
</svg>
</button>
-
+
<!-- Terminal prompt -->
- <span class="text-exo-yellow text-sm font-bold flex-shrink-0 leading-7">▶</span>
-
+ <span class="text-exo-yellow text-sm font-bold flex-shrink-0 leading-7"
+ >▶</span
+ >
+
<textarea
bind:this={textareaRef}
bind:value={message}
onkeydown={handleKeydown}
oninput={handleInput}
onpaste={handlePaste}
- {placeholder}
+ placeholder={isEditOnlyWithoutImage
+ ? "Attach an image to edit..."
+ : isEditMode
+ ? "Describe how to edit this image..."
+ : isImageModel()
+ ? "Describe the image you want to generate..."
+ : placeholder}
disabled={loading}
rows={1}
class="flex-1 resize-none bg-transparent text-foreground placeholder:text-exo-light-gray/60 placeholder:text-sm placeholder:tracking-[0.15em] placeholder:leading-7 focus:outline-none focus:ring-0 focus:border-none disabled:opacity-50 text-sm leading-7 font-mono"
style="min-height: 28px; max-height: 150px;"
></textarea>
-
+
<button
type="submit"
- disabled={!canSend || loading}
+ disabled={!canSend || loading || isEditOnlyWithoutImage}
class="px-2.5 sm:px-4 py-1.5 sm:py-2 rounded text-xs sm:text-xs tracking-[0.1em] sm:tracking-[0.15em] uppercase font-medium transition-all duration-200 whitespace-nowrap
- {!canSend || loading
- ? 'bg-exo-medium-gray/50 text-exo-light-gray cursor-not-allowed'
- : 'bg-exo-yellow text-exo-black hover:bg-exo-yellow-darker hover:shadow-[0_0_20px_rgba(255,215,0,0.3)]'}"
- aria-label="Send message"
+ {!canSend || loading || isEditOnlyWithoutImage
+ ? 'bg-exo-medium-gray/50 text-exo-light-gray cursor-not-allowed'
+ : 'bg-exo-yellow text-exo-black hover:bg-exo-yellow-darker hover:shadow-[0_0_20px_rgba(255,215,0,0.3)]'}"
+ aria-label={shouldShowEditMode
+ ? "Edit image"
+ : isImageModel()
+ ? "Generate image"
+ : "Send message"}
>
{#if loading}
<span class="inline-flex items-center gap-1 sm:gap-2">
- <span class="w-2.5 h-2.5 sm:w-3 sm:h-3 border-2 border-current border-t-transparent rounded-full animate-spin"></span>
- <span class="hidden sm:inline">PROCESSING</span>
+ <span
+ class="w-2.5 h-2.5 sm:w-3 sm:h-3 border-2 border-current border-t-transparent rounded-full animate-spin"
+ ></span>
+ <span class="hidden sm:inline"
+ >{shouldShowEditMode
+ ? "EDITING"
+ : isImageModel()
+ ? "GENERATING"
+ : "PROCESSING"}</span
+ >
<span class="sm:hidden">...</span>
</span>
+ {:else if shouldShowEditMode}
+ <span class="inline-flex items-center gap-1.5">
+ <svg
+ class="w-3.5 h-3.5"
+ fill="none"
+ viewBox="0 0 24 24"
+ stroke="currentColor"
+ stroke-width="2"
+ >
+ <path
+ stroke-linecap="round"
+ stroke-linejoin="round"
+ d="M11 5H6a2 2 0 00-2 2v11a2 2 0 002 2h11a2 2 0 002-2v-5m-1.414-9.414a2 2 0 112.828 2.828L11.828 15H9v-2.828l8.586-8.586z"
+ />
+ </svg>
+ <span>EDIT</span>
+ </span>
+ {:else if isImageModel()}
+ <span class="inline-flex items-center gap-1.5">
+ <svg
+ class="w-3.5 h-3.5"
+ fill="none"
+ viewBox="0 0 24 24"
+ stroke="currentColor"
+ stroke-width="2"
+ >
+ <rect x="3" y="3" width="18" height="18" rx="2" ry="2" />
+ <circle cx="8.5" cy="8.5" r="1.5" />
+ <polyline points="21 15 16 10 5 21" />
+ </svg>
+ <span>GENERATE</span>
+ </span>
{:else}
SEND
{/if}
</button>
</div>
-
+
<!-- Bottom accent line -->
- <div class="absolute bottom-0 left-0 right-0 h-px bg-gradient-to-r from-transparent via-exo-yellow/30 to-transparent"></div>
+ <div
+ class="absolute bottom-0 left-0 right-0 h-px bg-gradient-to-r from-transparent via-exo-yellow/30 to-transparent"
+ ></div>
</div>
-
+
{#if showHelperText}
- <p class="mt-2 sm:mt-3 text-center text-xs sm:text-xs text-exo-light-gray tracking-[0.1em] sm:tracking-[0.15em] uppercase">
- <kbd class="px-1 sm:px-1.5 py-0.5 rounded bg-exo-medium-gray/30 text-exo-light-gray border border-exo-medium-gray/50">ENTER</kbd>
+ <p
+ class="mt-2 sm:mt-3 text-center text-xs sm:text-xs text-exo-light-gray tracking-[0.1em] sm:tracking-[0.15em] uppercase"
+ >
+ <kbd
+ class="px-1 sm:px-1.5 py-0.5 rounded bg-exo-medium-gray/30 text-exo-light-gray border border-exo-medium-gray/50"
+ >ENTER</kbd
+ >
<span class="mx-0.5 sm:mx-1">TO SEND</span>
<span class="text-exo-medium-gray mx-1 sm:mx-2">|</span>
- <kbd class="px-1 sm:px-1.5 py-0.5 rounded bg-exo-medium-gray/30 text-exo-light-gray border border-exo-medium-gray/50">SHIFT+ENTER</kbd>
+ <kbd
+ class="px-1 sm:px-1.5 py-0.5 rounded bg-exo-medium-gray/30 text-exo-light-gray border border-exo-medium-gray/50"
+ >SHIFT+ENTER</kbd
+ >
<span class="mx-0.5 sm:mx-1">NEW LINE</span>
<span class="text-exo-medium-gray mx-1 sm:mx-2">|</span>
<span class="text-exo-light-gray">DRAG & DROP OR PASTE FILES</span>
diff --git a/dashboard/src/lib/components/ChatMessages.svelte b/dashboard/src/lib/components/ChatMessages.svelte
index 2bbd09f8..90e03128 100644
--- a/dashboard/src/lib/components/ChatMessages.svelte
+++ b/dashboard/src/lib/components/ChatMessages.svelte
@@ -1,12 +1,14 @@
<script lang="ts">
- import {
- messages,
- currentResponse,
+ import {
+ messages,
+ currentResponse,
isLoading,
deleteMessage,
editAndRegenerate,
- regenerateLastResponse
+ regenerateLastResponse,
+ setEditingImage
} from '$lib/stores/app.svelte';
+ import type { Message } from '$lib/stores/app.svelte';
import type { MessageAttachment } from '$lib/stores/app.svelte';
import MarkdownContent from './MarkdownContent.svelte';
@@ -365,10 +367,77 @@ function isThinkingExpanded(messageId: string): boolean {
{/if}
</div>
{/if}
+
+ <!-- Generated Images -->
+ {#if message.attachments?.some(a => a.type === 'generated-image')}
+ <div class="mb-3">
+ {#each message.attachments.filter(a => a.type === 'generated-image') as attachment}
+ <div class="relative group/img inline-block">
+ <img
+ src={attachment.preview}
+ alt=""
+ class="max-w-full max-h-[512px] rounded-lg border border-exo-yellow/20 shadow-lg shadow-black/20"
+ />
+ <!-- Button overlay -->
+ <div class="absolute top-2 right-2 flex gap-1 opacity-0 group-hover/img:opacity-100 transition-opacity">
+ <!-- Edit button -->
+ <button
+ type="button"
+ class="p-2 rounded-lg bg-exo-dark-gray/80 border border-exo-yellow/30 text-exo-yellow hover:bg-exo-dark-gray hover:border-exo-yellow/50 cursor-pointer"
+ onclick={() => {
+ if (attachment.preview) {
+ setEditingImage(attachment.preview, message);
+ }
+ }}
+ title="Edit image"
+ >
+ <svg class="w-4 h-4" fill="none" viewBox="0 0 24 24" stroke="currentColor" stroke-width="2">
+ <path stroke-linecap="round" stroke-linejoin="round" d="M11 5H6a2 2 0 00-2 2v11a2 2 0 002 2h11a2 2 0 002-2v-5m-1.414-9.414a2 2 0 112.828 2.828L11.828 15H9v-2.828l8.586-8.586z" />
+ </svg>
+ </button>
+ <!-- Download button -->
+ <button
+ type="button"
+ class="p-2 rounded-lg bg-exo-dark-gray/80 border border-exo-yellow/30 text-exo-yellow hover:bg-exo-dark-gray hover:border-exo-yellow/50 cursor-pointer"
+ onclick={() => {
+ if (attachment.preview) {
+ const link = document.createElement('a');
+ link.href = attachment.preview;
+ const ext = attachment.name?.split('.').pop() || 'png';
+ link.download = `generated-image-${Date.now()}.${ext}`;
+ link.click();
+ }
+ }}
+ title="Download image"
+ >
+ <svg class="w-4 h-4" fill="none" viewBox="0 0 24 24" stroke="currentColor" stroke-width="2">
+ <path stroke-linecap="round" stroke-linejoin="round" d="M4 16v1a3 3 0 003 3h10a3 3 0 003-3v-1m-4-4l-4 4m0 0l-4-4m4 4V4" />
+ </svg>
+ </button>
+ </div>
+ </div>
+ {/each}
+ </div>
+ {/if}
+
<div class="text-xs text-foreground">
- <MarkdownContent content={message.content || (loading ? response : '')} />
- {#if loading && !message.content}
- <span class="inline-block w-2 h-4 bg-exo-yellow/70 ml-1 cursor-blink"></span>
+ {#if message.content === 'Generating image...' || message.content === 'Editing image...' || message.content?.startsWith('Generating...') || message.content?.startsWith('Editing...')}
+ <div class="flex items-center gap-3 text-exo-yellow">
+ <div class="relative">
+ <div class="w-8 h-8 border-2 border-exo-yellow/30 border-t-exo-yellow rounded-full animate-spin"></div>
+ <svg class="absolute inset-0 w-8 h-8 p-1.5 text-exo-yellow/60" fill="none" viewBox="0 0 24 24" stroke="currentColor" stroke-width="2">
+ <rect x="3" y="3" width="18" height="18" rx="2" ry="2"/>
+ <circle cx="8.5" cy="8.5" r="1.5"/>
+ <polyline points="21 15 16 10 5 21"/>
+ </svg>
+ </div>
+ <span class="font-mono tracking-wider uppercase text-sm">{message.content}</span>
+ </div>
+ {:else if message.content || (loading && !message.attachments?.some(a => a.type === 'generated-image'))}
+ <MarkdownContent content={message.content || (loading ? response : '')} />
+ {#if loading && !message.content}
+ <span class="inline-block w-2 h-4 bg-exo-yellow/70 ml-1 cursor-blink"></span>
+ {/if}
{/if}
</div>
</div>
diff --git a/dashboard/src/lib/components/ImageParamsPanel.svelte b/dashboard/src/lib/components/ImageParamsPanel.svelte
new file mode 100644
index 00000000..39742b22
--- /dev/null
+++ b/dashboard/src/lib/components/ImageParamsPanel.svelte
@@ -0,0 +1,502 @@
+<script lang="ts">
+ import {
+ imageGenerationParams,
+ setImageGenerationParams,
+ resetImageGenerationParams,
+ type ImageGenerationParams,
+ } from "$lib/stores/app.svelte";
+
+ interface Props {
+ isEditMode?: boolean;
+ }
+
+ let { isEditMode = false }: Props = $props();
+
+ let showAdvanced = $state(false);
+
+ // Custom dropdown state
+ let isSizeDropdownOpen = $state(false);
+ let isQualityDropdownOpen = $state(false);
+ let sizeButtonRef: HTMLButtonElement | undefined = $state();
+ let qualityButtonRef: HTMLButtonElement | undefined = $state();
+
+ const sizeDropdownPosition = $derived(() => {
+ if (!sizeButtonRef || !isSizeDropdownOpen) return { top: 0, left: 0, width: 0 };
+ const rect = sizeButtonRef.getBoundingClientRect();
+ return { top: rect.top, left: rect.left, width: rect.width };
+ });
+
+ const qualityDropdownPosition = $derived(() => {
+ if (!qualityButtonRef || !isQualityDropdownOpen) return { top: 0, left: 0, width: 0 };
+ const rect = qualityButtonRef.getBoundingClientRect();
+ return { top: rect.top, left: rect.left, width: rect.width };
+ });
+
+ const params = $derived(imageGenerationParams());
+
+ const inputFidelityOptions: ImageGenerationParams["inputFidelity"][] = [
+ "low",
+ "high",
+ ];
+
+ const outputFormatOptions: ImageGenerationParams["outputFormat"][] = [
+ "png",
+ "jpeg",
+ ];
+
+ function handleInputFidelityChange(value: ImageGenerationParams["inputFidelity"]) {
+ setImageGenerationParams({ inputFidelity: value });
+ }
+
+ function handleOutputFormatChange(value: ImageGenerationParams["outputFormat"]) {
+ setImageGenerationParams({ outputFormat: value });
+ }
+
+ const sizeOptions: ImageGenerationParams["size"][] = [
+ "512x512",
+ "768x768",
+ "1024x1024",
+ "1024x768",
+ "768x1024",
+ ];
+
+ const qualityOptions: ImageGenerationParams["quality"][] = [
+ "low",
+ "medium",
+ "high",
+ ];
+
+ function selectSize(value: ImageGenerationParams["size"]) {
+ setImageGenerationParams({ size: value });
+ isSizeDropdownOpen = false;
+ }
+
+ function selectQuality(value: ImageGenerationParams["quality"]) {
+ setImageGenerationParams({ quality: value });
+ isQualityDropdownOpen = false;
+ }
+
+ function handleSeedChange(event: Event) {
+ const input = event.target as HTMLInputElement;
+ const value = input.value.trim();
+ if (value === "") {
+ setImageGenerationParams({ seed: null });
+ } else {
+ const num = parseInt(value, 10);
+ if (!isNaN(num) && num >= 0) {
+ setImageGenerationParams({ seed: num });
+ }
+ }
+ }
+
+ function handleStepsChange(event: Event) {
+ const value = parseInt((event.target as HTMLInputElement).value, 10);
+ setImageGenerationParams({ numInferenceSteps: value });
+ }
+
+ function handleGuidanceChange(event: Event) {
+ const value = parseFloat((event.target as HTMLInputElement).value);
+ setImageGenerationParams({ guidance: value });
+ }
+
+ function handleNegativePromptChange(event: Event) {
+ const value = (event.target as HTMLTextAreaElement).value;
+ setImageGenerationParams({ negativePrompt: value || null });
+ }
+
+ function clearSteps() {
+ setImageGenerationParams({ numInferenceSteps: null });
+ }
+
+ function clearGuidance() {
+ setImageGenerationParams({ guidance: null });
+ }
+
+ function handleReset() {
+ resetImageGenerationParams();
+ showAdvanced = false;
+ }
+
+ const hasAdvancedParams = $derived(
+ params.seed !== null ||
+ params.numInferenceSteps !== null ||
+ params.guidance !== null ||
+ (params.negativePrompt !== null && params.negativePrompt.trim() !== ""),
+ );
+</script>
+
+<div class="border-b border-exo-medium-gray/30 px-3 py-2">
+ <!-- Basic params row -->
+ <div class="flex items-center gap-3 flex-wrap">
+ <!-- Size -->
+ <div class="flex items-center gap-1.5">
+ <span class="text-xs text-exo-light-gray uppercase tracking-wider"
+ >SIZE:</span
+ >
+ <div class="relative">
+ <button
+ bind:this={sizeButtonRef}
+ type="button"
+ onclick={() => isSizeDropdownOpen = !isSizeDropdownOpen}
+ class="bg-exo-medium-gray/50 border border-exo-yellow/30 rounded pl-2 pr-6 py-1 text-xs font-mono text-exo-yellow cursor-pointer transition-all duration-200 hover:border-exo-yellow/50 focus:outline-none focus:border-exo-yellow/70 {isSizeDropdownOpen ? 'border-exo-yellow/70' : ''}"
+ >
+ {params.size}
+ </button>
+ <div class="absolute right-1.5 top-1/2 -translate-y-1/2 pointer-events-none transition-transform duration-200 {isSizeDropdownOpen ? 'rotate-180' : ''}">
+ <svg class="w-3 h-3 text-exo-yellow/60" fill="none" viewBox="0 0 24 24" stroke="currentColor">
+ <path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M19 9l-7 7-7-7" />
+ </svg>
+ </div>
+ </div>
+
+ {#if isSizeDropdownOpen}
+ <!-- Backdrop to close dropdown -->
+ <button
+ type="button"
+ class="fixed inset-0 z-[9998] cursor-default"
+ onclick={() => isSizeDropdownOpen = false}
+ aria-label="Close dropdown"
+ ></button>
+
+ <!-- Dropdown Panel - fixed positioning to escape overflow:hidden -->
+ <div
+ class="fixed bg-exo-dark-gray border border-exo-yellow/30 rounded shadow-lg shadow-black/50 z-[9999] max-h-48 overflow-y-auto min-w-max"
+ style="bottom: calc(100vh - {sizeDropdownPosition().top}px + 4px); left: {sizeDropdownPosition().left}px;"
+ >
+ <div class="py-1">
+ {#each sizeOptions as size}
+ <button
+ type="button"
+ onclick={() => selectSize(size)}
+ class="w-full px-3 py-1.5 text-left text-xs font-mono tracking-wide transition-colors duration-100 flex items-center gap-2 {
+ params.size === size
+ ? 'bg-transparent text-exo-yellow'
+ : 'text-exo-light-gray hover:text-exo-yellow'
+ }"
+ >
+ {#if params.size === size}
+ <svg class="w-3 h-3 flex-shrink-0" fill="currentColor" viewBox="0 0 20 20">
+ <path fill-rule="evenodd" d="M16.707 5.293a1 1 0 010 1.414l-8 8a1 1 0 01-1.414 0l-4-4a1 1 0 011.414-1.414L8 12.586l7.293-7.293a1 1 0 011.414 0z" clip-rule="evenodd" />
+ </svg>
+ {:else}
+ <span class="w-3"></span>
+ {/if}
+ <span>{size}</span>
+ </button>
+ {/each}
+ </div>
+ </div>
+ {/if}
+ </div>
+
+ <!-- Quality -->
+ <div class="flex items-center gap-1.5">
+ <span class="text-xs text-exo-light-gray uppercase tracking-wider"
+ >QUALITY:</span
+ >
+ <div class="relative">
+ <button
+ bind:this={qualityButtonRef}
+ type="button"
+ onclick={() => isQualityDropdownOpen = !isQualityDropdownOpen}
+ class="bg-exo-medium-gray/50 border border-exo-yellow/30 rounded pl-2 pr-6 py-1 text-xs font-mono text-exo-yellow cursor-pointer transition-all duration-200 hover:border-exo-yellow/50 focus:outline-none focus:border-exo-yellow/70 {isQualityDropdownOpen ? 'border-exo-yellow/70' : ''}"
+ >
+ {params.quality.toUpperCase()}
+ </button>
+ <div class="absolute right-1.5 top-1/2 -translate-y-1/2 pointer-events-none transition-transform duration-200 {isQualityDropdownOpen ? 'rotate-180' : ''}">
+ <svg class="w-3 h-3 text-exo-yellow/60" fill="none" viewBox="0 0 24 24" stroke="currentColor">
+ <path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M19 9l-7 7-7-7" />
+ </svg>
+ </div>
+ </div>
+
+ {#if isQualityDropdownOpen}
+ <!-- Backdrop to close dropdown -->
+ <button
+ type="button"
+ class="fixed inset-0 z-[9998] cursor-default"
+ onclick={() => isQualityDropdownOpen = false}
+ aria-label="Close dropdown"
+ ></button>
+
+ <!-- Dropdown Panel - fixed positioning to escape overflow:hidden -->
+ <div
+ class="fixed bg-exo-dark-gray border border-exo-yellow/30 rounded shadow-lg shadow-black/50 z-[9999] max-h-48 overflow-y-auto min-w-max"
+ style="bottom: calc(100vh - {qualityDropdownPosition().top}px + 4px); left: {qualityDropdownPosition().left}px;"
+ >
+ <div class="py-1">
+ {#each qualityOptions as quality}
+ <button
+ type="button"
+ onclick={() => selectQuality(quality)}
+ class="w-full px-3 py-1.5 text-left text-xs font-mono tracking-wide transition-colors duration-100 flex items-center gap-2 {
+ params.quality === quality
+ ? 'bg-transparent text-exo-yellow'
+ : 'text-exo-light-gray hover:text-exo-yellow'
+ }"
+ >
+ {#if params.quality === quality}
+ <svg class="w-3 h-3 flex-shrink-0" fill="currentColor" viewBox="0 0 20 20">
+ <path fill-rule="evenodd" d="M16.707 5.293a1 1 0 010 1.414l-8 8a1 1 0 01-1.414 0l-4-4a1 1 0 011.414-1.414L8 12.586l7.293-7.293a1 1 0 011.414 0z" clip-rule="evenodd" />
+ </svg>
+ {:else}
+ <span class="w-3"></span>
+ {/if}
+ <span>{quality.toUpperCase()}</span>
+ </button>
+ {/each}
+ </div>
+ </div>
+ {/if}
+ </div>
+
+ <!-- Format -->
+ <div class="flex items-center gap-1.5">
+ <span class="text-xs text-exo-light-gray uppercase tracking-wider"
+ >FORMAT:</span
+ >
+ <div class="flex rounded overflow-hidden border border-exo-yellow/30">
+ {#each outputFormatOptions as format}
+ <button
+ type="button"
+ onclick={() => handleOutputFormatChange(format)}
+ class="px-2 py-1 text-xs font-mono uppercase transition-all duration-200 cursor-pointer {
+ params.outputFormat === format
+ ? 'bg-exo-yellow text-exo-black'
+ : 'bg-exo-medium-gray/50 text-exo-light-gray hover:text-exo-yellow'
+ }"
+ >
+ {format}
+ </button>
+ {/each}
+ </div>
+ </div>
+
+ <!-- Input Fidelity (edit mode only) -->
+ {#if isEditMode}
+ <div class="flex items-center gap-1.5">
+ <span class="text-xs text-exo-light-gray uppercase tracking-wider"
+ >FIDELITY:</span
+ >
+ <div class="flex rounded overflow-hidden border border-exo-yellow/30">
+ {#each inputFidelityOptions as fidelity}
+ <button
+ type="button"
+ onclick={() => handleInputFidelityChange(fidelity)}
+ class="px-2 py-1 text-xs font-mono uppercase transition-all duration-200 cursor-pointer {
+ params.inputFidelity === fidelity
+ ? 'bg-exo-yellow text-exo-black'
+ : 'bg-exo-medium-gray/50 text-exo-light-gray hover:text-exo-yellow'
+ }"
+ title={fidelity === 'low' ? 'More creative variation' : 'Closer to original'}
+ >
+ {fidelity}
+ </button>
+ {/each}
+ </div>
+ </div>
+ {/if}
+
+ <!-- Spacer -->
+ <div class="flex-1"></div>
+
+ <!-- Advanced toggle -->
+ <button
+ type="button"
+ onclick={() => (showAdvanced = !showAdvanced)}
+ class="flex items-center gap-1 text-xs font-mono tracking-wider uppercase transition-colors duration-200 {showAdvanced ||
+ hasAdvancedParams
+ ? 'text-exo-yellow'
+ : 'text-exo-light-gray hover:text-exo-yellow'}"
+ >
+ <span>ADVANCED</span>
+ <svg
+ class="w-3 h-3 transition-transform duration-200 {showAdvanced
+ ? 'rotate-180'
+ : ''}"
+ fill="none"
+ viewBox="0 0 24 24"
+ stroke="currentColor"
+ >
+ <path
+ stroke-linecap="round"
+ stroke-linejoin="round"
+ stroke-width="2"
+ d="M19 9l-7 7-7-7"
+ />
+ </svg>
+ {#if hasAdvancedParams && !showAdvanced}
+ <span class="w-1.5 h-1.5 rounded-full bg-exo-yellow"></span>
+ {/if}
+ </button>
+ </div>
+
+ <!-- Advanced params section -->
+ {#if showAdvanced}
+ <div class="mt-3 pt-3 border-t border-exo-medium-gray/20 space-y-3">
+ <!-- Row 1: Seed and Steps -->
+ <div class="flex items-center gap-4 flex-wrap">
+ <!-- Seed -->
+ <div class="flex items-center gap-1.5">
+ <span class="text-xs text-exo-light-gray uppercase tracking-wider"
+ >SEED:</span
+ >
+ <input
+ type="number"
+ min="0"
+ value={params.seed ?? ""}
+ oninput={handleSeedChange}
+ placeholder="Random"
+ class="w-24 bg-exo-medium-gray/50 border border-exo-yellow/30 rounded px-2 py-1 text-xs font-mono text-exo-yellow placeholder:text-exo-light-gray/50 transition-all duration-200 hover:border-exo-yellow/50 focus:outline-none focus:border-exo-yellow/70"
+ />
+ </div>
+
+ <!-- Steps Slider -->
+ <div class="flex items-center gap-1.5 flex-1 min-w-[200px]">
+ <span
+ class="text-xs text-exo-light-gray uppercase tracking-wider whitespace-nowrap"
+ >STEPS:</span
+ >
+ <div class="flex items-center gap-2 flex-1">
+ <input
+ type="range"
+ min="1"
+ max="100"
+ value={params.numInferenceSteps ?? 50}
+ oninput={handleStepsChange}
+ class="flex-1 h-1 bg-exo-medium-gray/50 rounded appearance-none cursor-pointer accent-exo-yellow"
+ />
+ <span class="text-xs font-mono text-exo-yellow w-8 text-right">
+ {params.numInferenceSteps ?? "--"}
+ </span>
+ {#if params.numInferenceSteps !== null}
+ <button
+ type="button"
+ onclick={clearSteps}
+ class="text-exo-light-gray hover:text-exo-yellow transition-colors"
+ title="Clear"
+ >
+ <svg
+ class="w-3 h-3"
+ fill="none"
+ viewBox="0 0 24 24"
+ stroke="currentColor"
+ >
+ <path
+ stroke-linecap="round"
+ stroke-linejoin="round"
+ stroke-width="2"
+ d="M6 18L18 6M6 6l12 12"
+ />
+ </svg>
+ </button>
+ {/if}
+ </div>
+ </div>
+ </div>
+
+ <!-- Row 2: Guidance -->
+ <div class="flex items-center gap-1.5">
+ <span
+ class="text-xs text-exo-light-gray uppercase tracking-wider whitespace-nowrap"
+ >GUIDANCE:</span
+ >
+ <div class="flex items-center gap-2 flex-1 max-w-xs">
+ <input
+ type="range"
+ min="1"
+ max="20"
+ step="0.5"
+ value={params.guidance ?? 7.5}
+ oninput={handleGuidanceChange}
+ class="flex-1 h-1 bg-exo-medium-gray/50 rounded appearance-none cursor-pointer accent-exo-yellow"
+ />
+ <span class="text-xs font-mono text-exo-yellow w-8 text-right">
+ {params.guidance !== null ? params.guidance.toFixed(1) : "--"}
+ </span>
+ {#if params.guidance !== null}
+ <button
+ type="button"
+ onclick={clearGuidance}
+ class="text-exo-light-gray hover:text-exo-yellow transition-colors"
+ title="Clear"
+ >
+ <svg
+ class="w-3 h-3"
+ fill="none"
+ viewBox="0 0 24 24"
+ stroke="currentColor"
+ >
+ <path
+ stroke-linecap="round"
+ stroke-linejoin="round"
+ stroke-width="2"
+ d="M6 18L18 6M6 6l12 12"
+ />
+ </svg>
+ </button>
+ {/if}
+ </div>
+ </div>
+
+ <!-- Row 3: Negative Prompt -->
+ <div class="flex flex-col gap-1.5">
+ <span class="text-xs text-exo-light-gray uppercase tracking-wider"
+ >NEGATIVE PROMPT:</span
+ >
+ <textarea
+ value={params.negativePrompt ?? ""}
+ oninput={handleNegativePromptChange}
+ placeholder="Things to avoid in the image..."
+ rows={2}
+ class="w-full bg-exo-medium-gray/50 border border-exo-yellow/30 rounded px-2 py-1.5 text-xs font-mono text-exo-yellow placeholder:text-exo-light-gray/50 resize-none transition-all duration-200 hover:border-exo-yellow/50 focus:outline-none focus:border-exo-yellow/70"
+ ></textarea>
+ </div>
+
+ <!-- Reset Button -->
+ <div class="flex justify-end pt-1">
+ <button
+ type="button"
+ onclick={handleReset}
+ class="text-xs font-mono tracking-wider uppercase text-exo-light-gray hover:text-exo-yellow transition-colors duration-200"
+ >
+ RESET TO DEFAULTS
+ </button>
+ </div>
+ </div>
+ {/if}
+</div>
+
+<style>
+ /* Custom range slider styling */
+ input[type="range"]::-webkit-slider-thumb {
+ -webkit-appearance: none;
+ appearance: none;
+ width: 12px;
+ height: 12px;
+ border-radius: 50%;
+ background: #ffd700;
+ cursor: pointer;
+ border: none;
+ }
+
+ input[type="range"]::-moz-range-thumb {
+ width: 12px;
+ height: 12px;
+ border-radius: 50%;
+ background: #ffd700;
+ cursor: pointer;
+ border: none;
+ }
+
+ /* Hide number input spinners */
+ input[type="number"]::-webkit-inner-spin-button,
+ input[type="number"]::-webkit-outer-spin-button {
+ -webkit-appearance: none;
+ margin: 0;
+ }
+
+ input[type="number"] {
+ -moz-appearance: textfield;
+ }
+</style>
diff --git a/dashboard/src/lib/components/index.ts b/dashboard/src/lib/components/index.ts
index 5b35654e..dc8a7d76 100644
--- a/dashboard/src/lib/components/index.ts
+++ b/dashboard/src/lib/components/index.ts
@@ -5,3 +5,4 @@ export { default as ChatAttachments } from "./ChatAttachments.svelte";
export { default as ChatSidebar } from "./ChatSidebar.svelte";
export { default as ModelCard } from "./ModelCard.svelte";
export { default as MarkdownContent } from "./MarkdownContent.svelte";
+export { default as ImageParamsPanel } from "./ImageParamsPanel.svelte";
diff --git a/dashboard/src/lib/stores/app.svelte.ts b/dashboard/src/lib/stores/app.svelte.ts
index 23d32741..b67c103e 100644
--- a/dashboard/src/lib/stores/app.svelte.ts
+++ b/dashboard/src/lib/stores/app.svelte.ts
@@ -193,7 +193,7 @@ interface RawStateResponse {
}
export interface MessageAttachment {
- type: "image" | "text" | "file";
+ type: "image" | "text" | "file" | "generated-image";
name: string;
content?: string;
preview?: string;
@@ -223,6 +223,39 @@ export interface Conversation {
}
const STORAGE_KEY = "exo-conversations";
+const IMAGE_PARAMS_STORAGE_KEY = "exo-image-generation-params";
+
+// Image generation params interface matching backend API
+export interface ImageGenerationParams {
+ // Basic params
+ size: "512x512" | "768x768" | "1024x1024" | "1024x768" | "768x1024";
+ quality: "low" | "medium" | "high";
+ outputFormat: "png" | "jpeg";
+ // Advanced params
+ seed: number | null;
+ numInferenceSteps: number | null;
+ guidance: number | null;
+ negativePrompt: string | null;
+ // Edit mode params
+ inputFidelity: "low" | "high";
+}
+
+// Image being edited
+export interface EditingImage {
+ imageDataUrl: string;
+ sourceMessage: Message;
+}
+
+const DEFAULT_IMAGE_PARAMS: ImageGenerationParams = {
+ size: "1024x1024",
+ quality: "medium",
+ outputFormat: "png",
+ seed: null,
+ numInferenceSteps: null,
+ guidance: null,
+ negativePrompt: null,
+ inputFidelity: "low",
+};
interface GranularNodeState {
nodeIdentities?: Record<string, RawNodeIdentity>;
@@ -395,6 +428,14 @@ class AppStore {
topologyOnlyMode = $state(false);
chatSidebarVisible = $state(true); // Shown by default
+ // Image generation params
+ imageGenerationParams = $state<ImageGenerationParams>({
+ ...DEFAULT_IMAGE_PARAMS,
+ });
+
+ // Image editing state
+ editingImage = $state<EditingImage | null>(null);
+
private fetchInterval: ReturnType<typeof setInterval> | null = null;
private previewsInterval: ReturnType<typeof setInterval> | null = null;
private lastConversationPersistTs = 0;
@@ -406,6 +447,7 @@ class AppStore {
this.loadDebugModeFromStorage();
this.loadTopologyOnlyModeFromStorage();
this.loadChatSidebarVisibleFromStorage();
+ this.loadImageGenerationParamsFromStorage();
}
}
@@ -507,6 +549,57 @@ class AppStore {
}
}
+ private loadImageGenerationParamsFromStorage() {
+ try {
+ const stored = localStorage.getItem(IMAGE_PARAMS_STORAGE_KEY);
+ if (stored) {
+ const parsed = JSON.parse(stored) as Partial<ImageGenerationParams>;
+ this.imageGenerationParams = {
+ ...DEFAULT_IMAGE_PARAMS,
+ ...parsed,
+ };
+ }
+ } catch (error) {
+ console.error("Failed to load image generation params:", error);
+ }
+ }
+
+ private saveImageGenerationParamsToStorage() {
+ try {
+ localStorage.setItem(
+ IMAGE_PARAMS_STORAGE_KEY,
+ JSON.stringify(this.imageGenerationParams),
+ );
+ } catch (error) {
+ console.error("Failed to save image generation params:", error);
+ }
+ }
+
+ getImageGenerationParams(): ImageGenerationParams {
+ return this.imageGenerationParams;
+ }
+
+ setImageGenerationParams(params: Partial<ImageGenerationParams>) {
+ this.imageGenerationParams = {
+ ...this.imageGenerationParams,
+ ...params,
+ };
+ this.saveImageGenerationParamsToStorage();
+ }
+
+ resetImageGenerationParams() {
+ this.imageGenerationParams = { ...DEFAULT_IMAGE_PARAMS };
+ this.saveImageGenerationParamsToStorage();
+ }
+
+ setEditingImage(imageDataUrl: string, sourceMessage: Message) {
+ this.editingImage = { imageDataUrl, sourceMessage };
+ }
+
+ clearEditingImage() {
+ this.editingImage = null;
+ }
+
/**
* Create a new conversation
*/
@@ -1577,6 +1670,357 @@ class AppStore {
}
}
+ /**
+ * Generate an image using the image generation API
+ */
+ async generateImage(prompt: string, modelId?: string): Promise<void> {
+ if (!prompt.trim() || this.isLoading) return;
+
+ if (!this.hasStartedChat) {
+ this.startChat();
+ }
+
+ this.isLoading = true;
+ this.currentResponse = "";
+
+ // Add user message
+ const userMessage: Message = {
+ id: generateUUID(),
+ role: "user",
+ content: prompt,
+ timestamp: Date.now(),
+ };
+ this.messages.push(userMessage);
+
+ // Create placeholder for assistant message with generating state
+ const assistantMessage = this.addMessage("assistant", "");
+ this.messages[this.messages.length - 1].content = "Generating image...";
+ this.updateActiveConversation();
+
+ try {
+ // Determine the model to use
+ let model = modelId || this.selectedChatModel;
+ if (!model) {
+ throw new Error(
+ "No model selected. Please select an image generation model.",
+ );
+ }
+
+ // Build request body using image generation params
+ const params = this.imageGenerationParams;
+ const hasAdvancedParams =
+ params.seed !== null ||
+ params.numInferenceSteps !== null ||
+ params.guidance !== null ||
+ (params.negativePrompt !== null && params.negativePrompt.trim() !== "");
+
+ const requestBody: Record<string, unknown> = {
+ model,
+ prompt,
+ quality: params.quality,
+ size: params.size,
+ output_format: params.outputFormat,
+ response_format: "b64_json",
+ stream: true,
+ partial_images: 3,
+ };
+
+ if (hasAdvancedParams) {
+ requestBody.advanced_params = {
+ ...(params.seed !== null && { seed: params.seed }),
+ ...(params.numInferenceSteps !== null && {
+ num_inference_steps: params.numInferenceSteps,
+ }),
+ ...(params.guidance !== null && { guidance: params.guidance }),
+ ...(params.negativePrompt !== null &&
+ params.negativePrompt.trim() !== "" && {
+ negative_prompt: params.negativePrompt,
+ }),
+ };
+ }
+
+ const response = await fetch("/v1/images/generations", {
+ method: "POST",
+ headers: {
+ "Content-Type": "application/json",
+ },
+ body: JSON.stringify(requestBody),
+ });
+
+ if (!response.ok) {
+ const errorText = await response.text();
+ throw new Error(`API error: ${response.status} - ${errorText}`);
+ }
+
+ const reader = response.body?.getReader();
+ if (!reader) {
+ throw new Error("No response body");
+ }
+
+ const decoder = new TextDecoder();
+ let buffer = "";
+ const idx = this.messages.findIndex((m) => m.id === assistantMessage.id);
+
+ while (true) {
+ const { done, value } = await reader.read();
+ if (done) break;
+
+ buffer += decoder.decode(value, { stream: true });
+
+ // Process complete lines
+ const lines = buffer.split("\n");
+ buffer = lines.pop() || ""; // Keep incomplete line in buffer
+
+ for (const line of lines) {
+ const trimmed = line.trim();
+ if (!trimmed) continue;
+
+ if (trimmed.startsWith("data: ")) {
+ const data = trimmed.slice(6);
+ if (data === "[DONE]") continue;
+
+ try {
+ const parsed = JSON.parse(data);
+ const imageData = parsed.data?.b64_json;
+
+ if (imageData && idx !== -1) {
+ 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.messages[idx].content =
+ `Generating... ${partialNum}/${totalPartials}`;
+ this.messages[idx].attachments = [
+ {
+ type: "generated-image",
+ name: `generated-image.${format}`,
+ preview: `data:${mimeType};base64,${imageData}`,
+ mimeType,
+ },
+ ];
+ } else if (parsed.type === "final") {
+ // Final image
+ this.messages[idx].content = "";
+ this.messages[idx].attachments = [
+ {
+ type: "generated-image",
+ name: `generated-image.${format}`,
+ preview: `data:${mimeType};base64,${imageData}`,
+ mimeType,
+ },
+ ];
+ }
+ }
+ } catch {
+ // Ignore parse errors for incomplete JSON
+ }
+ }
+ }
+ }
+ } catch (error) {
+ console.error("Error generating image:", error);
+ const idx = this.messages.findIndex((m) => m.id === assistantMessage.id);
+ if (idx !== -1) {
+ this.messages[idx].content =
+ `Error: ${error instanceof Error ? error.message : "Failed to generate image"}`;
+ }
+ } finally {
+ this.isLoading = false;
+ this.updateActiveConversation();
+ }
+ }
+
+ /**
+ * Edit an image using the image edit API
+ */
+ async editImage(
+ prompt: string,
+ imageDataUrl: string,
+ modelId?: string,
+ ): Promise<void> {
+ if (!prompt.trim() || !imageDataUrl || this.isLoading) return;
+
+ if (!this.hasStartedChat) {
+ this.startChat();
+ }
+
+ this.isLoading = true;
+ this.currentResponse = "";
+
+ // Add user message with the edit prompt
+ const userMessage: Message = {
+ id: generateUUID(),
+ role: "user",
+ content: prompt,
+ timestamp: Date.now(),
+ };
+ this.messages.push(userMessage);
+
+ // Create placeholder for assistant message with generating state
+ const assistantMessage = this.addMessage("assistant", "");
+ this.messages[this.messages.length - 1].content = "Editing image...";
+ this.updateActiveConversation();
+
+ // Clear editing state
+ this.editingImage = null;
+
+ try {
+ // Determine the model to use
+ let model = modelId || this.selectedChatModel;
+ if (!model) {
+ throw new Error(
+ "No model selected. Please select an image generation model.",
+ );
+ }
+
+ // Convert base64 data URL to blob
+ const response = await fetch(imageDataUrl);
+ const imageBlob = await response.blob();
+
+ // Build FormData request
+ const formData = new FormData();
+ formData.append("model", model);
+ formData.append("prompt", prompt);
+ formData.append("image", imageBlob, "image.png");
+
+ // Add params from image generation params
+ const params = this.imageGenerationParams;
+ formData.append("quality", params.quality);
+ formData.append("size", params.size);
+ formData.append("output_format", params.outputFormat);
+ formData.append("response_format", "b64_json");
+ formData.append("stream", "1"); // Use "1" instead of "true" for reliable FastAPI boolean parsing
+ formData.append("partial_images", "3");
+ formData.append("input_fidelity", params.inputFidelity);
+
+ // Advanced params
+ if (params.seed !== null) {
+ formData.append(
+ "advanced_params",
+ JSON.stringify({
+ seed: params.seed,
+ ...(params.numInferenceSteps !== null && {
+ num_inference_steps: params.numInferenceSteps,
+ }),
+ ...(params.guidance !== null && { guidance: params.guidance }),
+ ...(params.negativePrompt !== null &&
+ params.negativePrompt.trim() !== "" && {
+ negative_prompt: params.negativePrompt,
+ }),
+ }),
+ );
+ } else if (
+ params.numInferenceSteps !== null ||
+ params.guidance !== null ||
+ (params.negativePrompt !== null && params.negativePrompt.trim() !== "")
+ ) {
+ formData.append(
+ "advanced_params",
+ JSON.stringify({
+ ...(params.numInferenceSteps !== null && {
+ num_inference_steps: params.numInferenceSteps,
+ }),
+ ...(params.guidance !== null && { guidance: params.guidance }),
+ ...(params.negativePrompt !== null &&
+ params.negativePrompt.trim() !== "" && {
+ negative_prompt: params.negativePrompt,
+ }),
+ }),
+ );
+ }
+
+ const apiResponse = await fetch("/v1/images/edits", {
+ method: "POST",
+ body: formData,
+ });
+
+ if (!apiResponse.ok) {
+ const errorText = await apiResponse.text();
+ throw new Error(`API error: ${apiResponse.status} - ${errorText}`);
+ }
+
+ const reader = apiResponse.body?.getReader();
+ if (!reader) {
+ throw new Error("No response body");
+ }
+
+ const decoder = new TextDecoder();
+ let buffer = "";
+ const idx = this.messages.findIndex((m) => m.id === assistantMessage.id);
+
+ while (true) {
+ const { done, value } = await reader.read();
+ if (done) break;
+
+ buffer += decoder.decode(value, { stream: true });
+
+ // Process complete lines
+ const lines = buffer.split("\n");
+ buffer = lines.pop() || ""; // Keep incomplete line in buffer
+
+ for (const line of lines) {
+ const trimmed = line.trim();
+ if (!trimmed) continue;
+
+ if (trimmed.startsWith("data: ")) {
+ const data = trimmed.slice(6);
+ if (data === "[DONE]") continue;
+
+ try {
+ const parsed = JSON.parse(data);
+ const imageData = parsed.data?.b64_json;
+
+ if (imageData && idx !== -1) {
+ 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.messages[idx].content =
+ `Editing... ${partialNum}/${totalPartials}`;
+ this.messages[idx].attachments = [
+ {
+ type: "generated-image",
+ name: `edited-image.${format}`,
+ preview: `data:${mimeType};base64,${imageData}`,
+ mimeType,
+ },
+ ];
+ } else if (parsed.type === "final") {
+ // Final image
+ this.messages[idx].content = "";
+ this.messages[idx].attachments = [
+ {
+ type: "generated-image",
+ name: `edited-image.${format}`,
+ preview: `data:${mimeType};base64,${imageData}`,
+ mimeType,
+ },
+ ];
+ }
+ }
+ } catch {
+ // Ignore parse errors for incomplete JSON
+ }
+ }
+ }
+ }
+ } catch (error) {
+ console.error("Error editing image:", error);
+ const idx = this.messages.findIndex((m) => m.id === assistantMessage.id);
+ if (idx !== -1) {
+ this.messages[idx].content =
+ `Error: ${error instanceof Error ? error.message : "Failed to edit image"}`;
+ }
+ } finally {
+ this.isLoading = false;
+ this.updateActiveConversation();
+ }
+ }
+
/**
* Clear current chat and go back to welcome state
*/
@@ -1638,6 +2082,17 @@ export const sendMessage = (
preview?: string;
}[],
) => appStore.sendMessage(content, files);
+export const generateImage = (prompt: string, modelId?: string) =>
+ appStore.generateImage(prompt, modelId);
+export const editImage = (
+ prompt: string,
+ imageDataUrl: string,
+ modelId?: string,
+) => appStore.editImage(prompt, imageDataUrl, modelId);
+export const editingImage = () => appStore.editingImage;
+export const setEditingImage = (imageDataUrl: string, sourceMessage: Message) =>
+ appStore.setEditingImage(imageDataUrl, sourceMessage);
+export const clearEditingImage = () => appStore.clearEditingImage();
export const clearChat = () => appStore.clearChat();
export const setSelectedChatModel = (modelId: string) =>
appStore.setSelectedModel(modelId);
@@ -1678,3 +2133,11 @@ export const toggleChatSidebarVisible = () =>
export const setChatSidebarVisible = (visible: boolean) =>
appStore.setChatSidebarVisible(visible);
export const refreshState = () => appStore.fetchState();
+
+// Image generation params
+export const imageGenerationParams = () => appStore.getImageGenerationParams();
+export const setImageGenerationParams = (
+ params: Partial<ImageGenerationParams>,
+) => appStore.setImageGenerationParams(params);
+export const resetImageGenerationParams = () =>
+ appStore.resetImageGenerationParams();
diff --git a/dashboard/src/routes/+page.svelte b/dashboard/src/routes/+page.svelte
index c6dc8d14..4be37a3d 100644
--- a/dashboard/src/routes/+page.svelte
+++ b/dashboard/src/routes/+page.svelte
@@ -47,7 +47,30 @@ const sidebarVisible = $derived(chatSidebarVisible());
let mounted = $state(false);
// Instance launch state
- let models = $state<Array<{id: string, name?: string, storage_size_megabytes?: number}>>([]);
+ let models = $state<Array<{id: string, name?: string, storage_size_megabytes?: number, tasks?: string[], hugging_face_id?: string}>>([]);
+
+ // Model tasks lookup for ChatForm - maps both short IDs and full HuggingFace IDs
+ const modelTasks = $derived(() => {
+ const tasks: Record<string, string[]> = {};
+ for (const model of models) {
+ if (model.tasks && model.tasks.length > 0) {
+ // Map by short ID
+ tasks[model.id] = model.tasks;
+ // Also map by hugging_face_id from the API response
+ if (model.hugging_face_id) {
+ tasks[model.hugging_face_id] = model.tasks;
+ }
+ }
+ }
+ return tasks;
+ });
+
+ // Helper to check if a model supports image generation
+ function modelSupportsImageGeneration(modelId: string): boolean {
+ const model = models.find(m => m.id === modelId || m.hugging_face_id === modelId);
+ if (!model?.tasks) return false;
+ return model.tasks.includes('TextToImage') || model.tasks.includes('ImageToImage');
+ }
let selectedSharding = $state<'Pipeline' | 'Tensor'>('Pipeline');
type InstanceMeta = 'MlxRing' | 'MlxIbv' | 'MlxJaccl';
@@ -1270,6 +1293,7 @@ function toggleInstanceDownloadDetails(nodeId: string): void {
placeholder="Ask anything"
showHelperText={false}
showModelSelector={true}
+ modelTasks={modelTasks()}
/>
</div>
</div>
@@ -1491,8 +1515,18 @@ function toggleInstanceDownloadDetails(nodeId: string): void {
{@const foundModel = models.find(m => m.id === selectedModelId)}
{#if foundModel}
{@const sizeGB = getModelSizeGB(foundModel)}
- <span class="flex items-center justify-between gap-2 w-full pr-4">
- <span class="text-exo-light-gray truncate">{foundModel.name || foundModel.id}</span>
+ {@const isImageModel = modelSupportsImageGeneration(foundModel.id)}
+ <span class="flex items-center justify-between gap-2 w-full pr-4">
+ <span class="flex items-center gap-2 text-exo-light-gray truncate">
+ {#if isImageModel}
+ <svg class="w-4 h-4 flex-shrink-0 text-exo-yellow" fill="none" viewBox="0 0 24 24" stroke="currentColor" stroke-width="2">
+ <rect x="3" y="3" width="18" height="18" rx="2" ry="2"/>
+ <circle cx="8.5" cy="8.5" r="1.5"/>
+ <polyline points="21 15 16 10 5 21"/>
+ </svg>
+ {/if}
+ <span class="truncate">{foundModel.name || foundModel.id}</span>
+ </span>
<span class="text-white/50 text-xs flex-shrink-0">{sizeGB >= 1 ? sizeGB.toFixed(0) : sizeGB.toFixed(1)}GB</span>
</span>
{:else}
@@ -1537,6 +1571,7 @@ function toggleInstanceDownloadDetails(nodeId: string): void {
) as model}
{@const sizeGB = getModelSizeGB(model)}
{@const modelCanFit = hasEnoughMemory(model)}
+ {@const isImageModel = modelSupportsImageGeneration(model.id)}
<button
type="button"
onclick={() => {
@@ -1556,7 +1591,16 @@ function toggleInstanceDownloadDetails(nodeId: string): void {
: 'text-white/30 cursor-default'
}"
>
- <span class="truncate">{model.name || model.id}</span>
+ <span class="flex items-center gap-2 truncate flex-1">
+ {#if isImageModel}
+ <svg class="w-4 h-4 flex-shrink-0 text-exo-yellow" fill="none" viewBox="0 0 24 24" stroke="currentColor" stroke-width="2" aria-label="Image generation model">
+ <rect x="3" y="3" width="18" height="18" rx="2" ry="2"/>
+ <circle cx="8.5" cy="8.5" r="1.5"/>
+ <polyline points="21 15 16 10 5 21"/>
+ </svg>
+ {/if}
+ <span class="truncate">{model.name || model.id}</span>
+ </span>
<span class="flex-shrink-0 text-xs {modelCanFit ? 'text-white/50' : 'text-red-400/60'}">
{sizeGB >= 1 ? sizeGB.toFixed(0) : sizeGB.toFixed(1)}GB
</span>
@@ -1753,7 +1797,7 @@ function toggleInstanceDownloadDetails(nodeId: string): void {
<div class="flex-shrink-0 px-8 pb-6 pt-4 bg-gradient-to-t from-exo-black via-exo-black to-transparent">
<div class="max-w-7xl mx-auto">
- <ChatForm placeholder="Ask anything" showModelSelector={true} />
+ <ChatForm placeholder="Ask anything" showModelSelector={true} modelTasks={modelTasks()} />
</div>
</div>
</div>
diff --git a/docs/api.md b/docs/api.md
index 46daa244..9c4ec096 100644
--- a/docs/api.md
+++ b/docs/api.md
@@ -1,6 +1,6 @@
# EXO API – Technical Reference
-This document describes the REST API exposed by the **EXO ** service, as implemented in:
+This document describes the REST API exposed by the **EXO** service, as implemented in:
`src/exo/master/api.py`
@@ -183,7 +183,70 @@ Same schema as `/v1/chat/completions`.
**Response:**
Chat completion plus benchmarking metrics.
-## 5. Complete Endpoint Summary
+## 5. Image Generation & Editing
+
+### Image Generation
+
+**POST** `/v1/images/generations`
+
+Executes an image generation request using an OpenAI-compatible schema with additional advanced_params.
+
+**Request body (example):**
+
+```json
+{
+ "prompt": "a robot playing chess",
+ "model": "flux-dev",
+ "stream": false,
+}
+```
+
+**Advanced Parameters (`advanced_params`):**
+
+| Parameter | Type | Constraints | Description |
+|-----------|------|-------------|-------------|
+| `seed` | int | >= 0 | Random seed for reproducible generation |
+| `num_inference_steps` | int | 1-100 | Number of denoising steps |
+| `guidance` | float | 1.0-20.0 | Classifier-free guidance scale |
+| `negative_prompt` | string | - | Text describing what to avoid in the image |
+
+**Response:**
+OpenAI-compatible image generation response.
+
+### Benchmarked Image Generation
+
+**POST** `/bench/images/generations`
+
+Same as `/v1/images/generations`, but also returns generation statistics.
+
+**Request body:**
+Same schema as `/v1/images/generations`.
+
+**Response:**
+Image generation plus benchmarking metrics.
+
+### Image Editing
+
+**POST** `/v1/images/edits`
+
+Executes an image editing request using an OpenAI-compatible schema with additional advanced_params (same as `/v1/images/generations`).
+
+**Response:**
+Same format as `/v1/images/generations`.
+
+### Benchmarked Image Editing
+
+**POST** `/bench/images/edits`
+
+Same as `/v1/images/edits`, but also returns generation statistics.
+
+**Request:**
+Same schema as `/v1/images/edits`.
+
+**Response:**
+Same format as `/bench/images/generations`, including `generation_stats`.
+
+## 6. Complete Endpoint Summary
```
GET /node_id
@@ -203,10 +266,16 @@ GET /v1/models
POST /v1/chat/completions
POST /bench/chat/completions
+
+POST /v1/images/generations
+POST /bench/images/generations
+POST /v1/images/edits
+POST /bench/images/edits
```
-## 6. Notes
+## 7. Notes
-* The `/v1/chat/completions` endpoint is compatible with the OpenAI API format, so existing OpenAI clients can be pointed to EXO by changing the base URL.
+* The `/v1/chat/completions` endpoint is compatible with the OpenAI Chat API format, so existing OpenAI clients can be pointed to EXO by changing the base URL.
+* The `/v1/images/generations` and `/v1/images/edits` endpoints are compatible with the OpenAI Images API format.
* The instance placement endpoints allow you to plan and preview cluster allocations before actually creating instances.
* The `/events` and `/state` endpoints are primarily intended for operational visibility and debugging.
diff --git a/src/exo/master/api.py b/src/exo/master/api.py
index 5402802f..9ea3fd5a 100644
--- a/src/exo/master/api.py
+++ b/src/exo/master/api.py
@@ -1,22 +1,26 @@
+import base64
+import json
import time
from collections.abc import AsyncGenerator
from http import HTTPStatus
-from typing import cast
+from typing import Literal, cast
import anyio
from anyio import BrokenResourceError, create_task_group
from anyio.abc import TaskGroup
-from fastapi import FastAPI, HTTPException, Request
+from fastapi import FastAPI, File, Form, HTTPException, Request, UploadFile
from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import JSONResponse, StreamingResponse
+from fastapi.responses import FileResponse, JSONResponse, StreamingResponse
from fastapi.staticfiles import StaticFiles
from hypercorn.asyncio import serve # pyright: ignore[reportUnknownVariableType]
from hypercorn.config import Config
from hypercorn.typing import ASGIFramework
from loguru import logger
+from exo.master.image_store import ImageStore
from exo.master.placement import place_instance as get_instance_placements
from exo.shared.apply import apply
+from exo.shared.constants import EXO_IMAGE_CACHE_DIR, EXO_MAX_CHUNK_SIZE
from exo.shared.election import ElectionMessage
from exo.shared.logging import InterceptLogger
from exo.shared.models.model_cards import (
@@ -27,6 +31,8 @@ from exo.shared.models.model_cards import (
from exo.shared.types.api import (
BenchChatCompletionResponse,
BenchChatCompletionTaskParams,
+ BenchImageGenerationResponse,
+ BenchImageGenerationTaskParams,
ChatCompletionChoice,
ChatCompletionMessage,
ChatCompletionResponse,
@@ -37,6 +43,13 @@ from exo.shared.types.api import (
ErrorResponse,
FinishReason,
GenerationStats,
+ ImageData,
+ ImageEditsInternalParams,
+ ImageGenerationResponse,
+ ImageGenerationStats,
+ ImageGenerationTaskParams,
+ ImageListItem,
+ ImageListResponse,
ModelList,
ModelListModel,
PlaceInstanceParams,
@@ -44,17 +57,20 @@ from exo.shared.types.api import (
PlacementPreviewResponse,
StreamingChoiceResponse,
)
-from exo.shared.types.chunks import TokenChunk
+from exo.shared.types.chunks import ImageChunk, InputImageChunk, TokenChunk
from exo.shared.types.commands import (
ChatCompletion,
Command,
CreateInstance,
DeleteInstance,
ForwarderCommand,
+ ImageEdits,
+ ImageGeneration,
PlaceInstance,
+ SendInputChunk,
TaskFinished,
)
-from exo.shared.types.common import CommandId, NodeId, SessionId
+from exo.shared.types.common import CommandId, Id, NodeId, SessionId
from exo.shared.types.events import (
ChunkGenerated,
Event,
@@ -72,6 +88,10 @@ from exo.utils.dashboard_path import find_dashboard
from exo.utils.event_buffer import OrderedBuffer
+def _format_to_content_type(image_format: Literal["png", "jpeg", "webp"] | None) -> str:
+ return f"image/{image_format or 'png'}"
+
+
def chunk_to_response(
chunk: TokenChunk, command_id: CommandId
) -> ChatCompletionResponse:
@@ -93,8 +113,12 @@ async def resolve_model_card(model_id: ModelId) -> ModelCard:
if model_id in MODEL_CARDS:
model_card = MODEL_CARDS[model_id]
return model_card
- else:
- return await ModelCard.from_hf(model_id)
+
+ for card in MODEL_CARDS.values():
+ if card.model_id == ModelId(model_id):
+ return card
+
+ return await ModelCard.from_hf(model_id)
class API:
@@ -139,6 +163,8 @@ class API:
)
self._chat_completion_queues: dict[CommandId, Sender[TokenChunk]] = {}
+ self._image_generation_queues: dict[CommandId, Sender[ImageChunk]] = {}
+ self._image_store = ImageStore(EXO_IMAGE_CACHE_DIR)
self._tg: TaskGroup | None = None
def reset(self, new_session_id: SessionId, result_clock: int):
@@ -147,6 +173,7 @@ class API:
self.session_id = new_session_id
self.event_buffer = OrderedBuffer[Event]()
self._chat_completion_queues = {}
+ self._image_generation_queues = {}
self.unpause(result_clock)
def unpause(self, result_clock: int):
@@ -194,6 +221,14 @@ class API:
self.chat_completions
)
self.app.post("/bench/chat/completions")(self.bench_chat_completions)
+ self.app.post("/v1/images/generations", response_model=None)(
+ self.image_generations
+ )
+ self.app.post("/bench/images/generations")(self.bench_image_generations)
+ self.app.post("/v1/images/edits", response_model=None)(self.image_edits)
+ self.app.post("/bench/images/edits")(self.bench_image_edits)
+ self.app.get("/images")(self.list_images)
+ self.app.get("/images/{image_id}")(self.get_image)
self.app.get("/state")(lambda: self.state)
self.app.get("/events")(lambda: self._event_log)
@@ -601,6 +636,458 @@ class API:
response = await self._collect_chat_completion_with_stats(command.command_id)
return response
+ async def _validate_image_model(self, model: str) -> ModelId:
+ """Validate model exists and return resolved model ID.
+
+ Raises HTTPException 404 if no instance is found for the model.
+ """
+ model_card = await resolve_model_card(ModelId(model))
+ resolved_model = model_card.model_id
+ if not any(
+ instance.shard_assignments.model_id == resolved_model
+ for instance in self.state.instances.values()
+ ):
+ await self._trigger_notify_user_to_download_model(resolved_model)
+ raise HTTPException(
+ status_code=404, detail=f"No instance found for model {resolved_model}"
+ )
+ return resolved_model
+
+ async def get_image(self, image_id: str) -> FileResponse:
+ stored = self._image_store.get(Id(image_id))
+ if stored is None:
+ raise HTTPException(status_code=404, detail="Image not found or expired")
+ return FileResponse(path=stored.file_path, media_type=stored.content_type)
+
+ async def list_images(self, request: Request) -> ImageListResponse:
+ """List all stored images."""
+ stored_images = self._image_store.list_images()
+ return ImageListResponse(
+ data=[
+ ImageListItem(
+ image_id=img.image_id,
+ url=self._build_image_url(request, img.image_id),
+ content_type=img.content_type,
+ expires_at=img.expires_at,
+ )
+ for img in stored_images
+ ]
+ )
+
+ def _build_image_url(self, request: Request, image_id: Id) -> str:
+ host = request.headers.get("host", f"localhost:{self.port}")
+ scheme = "https" if request.url.scheme == "https" else "http"
+ return f"{scheme}://{host}/v1/images/{image_id}"
+
+ async def image_generations(
+ self, request: Request, payload: ImageGenerationTaskParams
+ ) -> 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.
+ """
+ payload.model = await self._validate_image_model(payload.model)
+
+ command = ImageGeneration(
+ request_params=payload,
+ )
+ await self._send(command)
+
+ # Check if streaming is requested
+ if payload.stream and payload.partial_images and payload.partial_images > 0:
+ return StreamingResponse(
+ self._generate_image_stream(
+ request=request,
+ 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(
+ request=request,
+ 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,
+ request: Request,
+ 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:
+ if chunk.finish_reason == "error":
+ error_response = ErrorResponse(
+ error=ErrorInfo(
+ message=chunk.error_message or "Internal server error",
+ type="InternalServerError",
+ code=500,
+ )
+ )
+ yield f"data: {error_response.model_dump_json()}\n\n"
+ yield "data: [DONE]\n\n"
+ return
+
+ 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 (always use b64_json for partials)
+ event_data = {
+ "type": "partial",
+ "partial_index": partial_idx,
+ "total_partials": total_partials,
+ "format": str(chunk.format),
+ "data": {
+ "b64_json": full_data
+ if response_format == "b64_json"
+ else None,
+ },
+ }
+ yield f"data: {json.dumps(event_data)}\n\n"
+ else:
+ # Final image
+ if response_format == "url":
+ image_bytes = base64.b64decode(full_data)
+ content_type = _format_to_content_type(chunk.format)
+ stored = self._image_store.store(
+ image_bytes, content_type
+ )
+ url = self._build_image_url(request, stored.image_id)
+ event_data = {
+ "type": "final",
+ "image_index": chunk.image_index,
+ "format": str(chunk.format),
+ "data": {"url": url},
+ }
+ else:
+ event_data = {
+ "type": "final",
+ "image_index": chunk.image_index,
+ "format": str(chunk.format),
+ "data": {"b64_json": full_data},
+ }
+ 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_chunks(
+ self,
+ request: Request | None,
+ command_id: CommandId,
+ num_images: int,
+ response_format: str,
+ capture_stats: bool = False,
+ ) -> tuple[list[ImageData], ImageGenerationStats | None]:
+ """Collect image chunks and optionally capture stats."""
+ # 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] = {}
+ image_formats: dict[int, Literal["png", "jpeg", "webp"] | None] = {}
+ images_complete = 0
+ stats: ImageGenerationStats | None = None
+
+ try:
+ self._image_generation_queues[command_id], recv = channel[ImageChunk]()
+
+ while images_complete < num_images:
+ with recv as chunks:
+ async for chunk in chunks:
+ if chunk.finish_reason == "error":
+ raise HTTPException(
+ status_code=500,
+ detail=chunk.error_message or "Internal server error",
+ )
+
+ 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
+ image_formats[chunk.image_index] = chunk.format
+
+ image_chunks[chunk.image_index][chunk.chunk_index] = chunk.data
+
+ if capture_stats and chunk.stats is not None:
+ stats = chunk.stats
+
+ if (
+ len(image_chunks[chunk.image_index])
+ == image_total_chunks[chunk.image_index]
+ ):
+ images_complete += 1
+
+ if images_complete >= num_images:
+ break
+
+ images: list[ImageData] = []
+ for image_idx in range(num_images):
+ chunks_dict = image_chunks[image_idx]
+ full_data = "".join(chunks_dict[i] for i in range(len(chunks_dict)))
+ if response_format == "url" and request is not None:
+ image_bytes = base64.b64decode(full_data)
+ content_type = _format_to_content_type(image_formats.get(image_idx))
+ stored = self._image_store.store(image_bytes, content_type)
+ url = self._build_image_url(request, stored.image_id)
+ images.append(ImageData(b64_json=None, url=url))
+ else:
+ images.append(
+ ImageData(
+ b64_json=full_data
+ if response_format == "b64_json"
+ else None,
+ url=None,
+ )
+ )
+
+ return (images, stats if capture_stats else None)
+ 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,
+ request: Request,
+ command_id: CommandId,
+ num_images: int,
+ response_format: str,
+ ) -> ImageGenerationResponse:
+ """Collect all image chunks (non-streaming) and return a single response."""
+ images, _ = await self._collect_image_chunks(
+ request, command_id, num_images, response_format, capture_stats=False
+ )
+ return ImageGenerationResponse(data=images)
+
+ async def _collect_image_generation_with_stats(
+ self,
+ request: Request | None,
+ command_id: CommandId,
+ num_images: int,
+ response_format: str,
+ ) -> BenchImageGenerationResponse:
+ images, stats = await self._collect_image_chunks(
+ request, command_id, num_images, response_format, capture_stats=True
+ )
+ return BenchImageGenerationResponse(data=images, generation_stats=stats)
+
+ async def bench_image_generations(
+ self, request: Request, payload: BenchImageGenerationTaskParams
+ ) -> BenchImageGenerationResponse:
+ payload.model = await self._validate_image_model(payload.model)
+
+ payload.stream = False
+ payload.partial_images = 0
+
+ command = ImageGeneration(
+ request_params=payload,
+ )
+ await self._send(command)
+
+ return await self._collect_image_generation_with_stats(
+ request=request,
+ command_id=command.command_id,
+ num_images=payload.n or 1,
+ response_format=payload.response_format or "b64_json",
+ )
+
+ async def _send_image_edits_command(
+ self,
+ image: UploadFile,
+ prompt: str,
+ model: str,
+ n: int,
+ size: str,
+ response_format: Literal["url", "b64_json"],
+ input_fidelity: Literal["low", "high"],
+ stream: bool,
+ partial_images: int,
+ bench: bool,
+ ) -> ImageEdits:
+ """Prepare and send an image edits command with chunked image upload."""
+ resolved_model = await self._validate_image_model(model)
+
+ image_content = await image.read()
+ image_data = base64.b64encode(image_content).decode("utf-8")
+
+ image_strength = 0.7 if input_fidelity == "high" else 0.3
+
+ data_chunks = [
+ image_data[i : i + EXO_MAX_CHUNK_SIZE]
+ for i in range(0, len(image_data), EXO_MAX_CHUNK_SIZE)
+ ]
+ total_chunks = len(data_chunks)
+
+ command = ImageEdits(
+ request_params=ImageEditsInternalParams(
+ image_data="",
+ total_input_chunks=total_chunks,
+ prompt=prompt,
+ model=resolved_model,
+ n=n,
+ size=size,
+ response_format=response_format,
+ image_strength=image_strength,
+ stream=stream,
+ partial_images=partial_images,
+ bench=bench,
+ ),
+ )
+
+ logger.info(
+ f"Sending input image: {len(image_data)} bytes in {total_chunks} chunks"
+ )
+ for chunk_index, chunk_data in enumerate(data_chunks):
+ await self._send(
+ SendInputChunk(
+ chunk=InputImageChunk(
+ idx=chunk_index,
+ model=resolved_model,
+ command_id=command.command_id,
+ data=chunk_data,
+ chunk_index=chunk_index,
+ total_chunks=total_chunks,
+ )
+ )
+ )
+
+ await self._send(command)
+ return command
+
+ async def image_edits(
+ self,
+ request: Request,
+ image: UploadFile = File(...), # noqa: B008
+ prompt: str = Form(...),
+ model: str = Form(...),
+ n: int = Form(1),
+ size: str = Form("1024x1024"),
+ response_format: Literal["url", "b64_json"] = Form("b64_json"),
+ input_fidelity: Literal["low", "high"] = Form("low"),
+ stream: str = Form("false"),
+ partial_images: str = Form("0"),
+ ) -> ImageGenerationResponse | StreamingResponse:
+ """Handle image editing requests (img2img)."""
+ # Parse string form values to proper types
+ stream_bool = stream.lower() in ("true", "1", "yes")
+ partial_images_int = int(partial_images) if partial_images.isdigit() else 0
+
+ command = await self._send_image_edits_command(
+ image=image,
+ prompt=prompt,
+ model=model,
+ n=n,
+ size=size,
+ response_format=response_format,
+ input_fidelity=input_fidelity,
+ stream=stream_bool,
+ partial_images=partial_images_int,
+ bench=False,
+ )
+
+ if stream_bool and partial_images_int > 0:
+ return StreamingResponse(
+ self._generate_image_stream(
+ request=request,
+ command_id=command.command_id,
+ num_images=n,
+ response_format=response_format,
+ ),
+ media_type="text/event-stream",
+ )
+
+ return await self._collect_image_generation(
+ request=request,
+ command_id=command.command_id,
+ num_images=n,
+ response_format=response_format,
+ )
+
+ async def bench_image_edits(
+ self,
+ request: Request,
+ image: UploadFile = File(...), # noqa: B008
+ prompt: str = Form(...),
+ model: str = Form(...),
+ n: int = Form(1),
+ size: str = Form("1024x1024"),
+ response_format: Literal["url", "b64_json"] = Form("b64_json"),
+ input_fidelity: Literal["low", "high"] = Form("low"),
+ ) -> BenchImageGenerationResponse:
+ """Handle benchmark image editing requests with generation stats."""
+ command = await self._send_image_edits_command(
+ image=image,
+ prompt=prompt,
+ model=model,
+ n=n,
+ size=size,
+ response_format=response_format,
+ input_fidelity=input_fidelity,
+ stream=False,
+ partial_images=0,
+ bench=True,
+ )
+
+ return await self._collect_image_generation_with_stats(
+ request=request,
+ command_id=command.command_id,
+ num_images=n,
+ response_format=response_format,
+ )
+
def _calculate_total_available_memory(self) -> Memory:
"""Calculate total available memory across all nodes in bytes."""
total_available = Memory()
@@ -622,6 +1109,7 @@ class API:
tags=[],
storage_size_megabytes=int(card.storage_size.in_mb),
supports_tensor=card.supports_tensor,
+ tasks=[task.value for task in card.tasks],
)
for card in MODEL_CARDS.values()
]
@@ -640,6 +1128,7 @@ class API:
logger.info("Starting API")
tg.start_soon(self._apply_state)
tg.start_soon(self._pause_on_new_election)
+ tg.start_soon(self._cleanup_expired_images)
print_startup_banner(self.port)
await serve(
cast(ASGIFramework, self.app),
@@ -660,13 +1149,26 @@ class API:
self._event_log.append(event)
self.state = apply(self.state, IndexedEvent(event=event, idx=idx))
if isinstance(event, ChunkGenerated):
- assert isinstance(event.chunk, TokenChunk)
- queue = self._chat_completion_queues.get(event.command_id)
- if queue is not None:
- try:
- await queue.send(event.chunk)
- except BrokenResourceError:
- self._chat_completion_queues.pop(event.command_id, None)
+ if event.command_id in self._chat_completion_queues:
+ assert isinstance(event.chunk, TokenChunk)
+ queue = self._chat_completion_queues.get(event.command_id)
+ if queue is not None:
+ try:
+ await queue.send(event.chunk)
+ except BrokenResourceError:
+ self._chat_completion_queues.pop(
+ event.command_id, None
+ )
+ elif event.command_id in self._image_generation_queues:
+ assert isinstance(event.chunk, ImageChunk)
+ queue = self._image_generation_queues.get(event.command_id)
+ if queue is not None:
+ try:
+ await queue.send(event.chunk)
+ except BrokenResourceError:
+ self._image_generation_queues.pop(
+ event.command_id, None
+ )
async def _pause_on_new_election(self):
with self.election_receiver as ems:
@@ -674,6 +1176,15 @@ class API:
if message.clock > self.last_completed_election:
self.paused = True
+ async def _cleanup_expired_images(self):
+ """Periodically clean up expired images from the store."""
+ cleanup_interval_seconds = 300 # 5 minutes
+ while True:
+ await anyio.sleep(cleanup_interval_seconds)
+ removed = self._image_store.cleanup_expired()
+ if removed > 0:
+ logger.debug(f"Cleaned up {removed} expired images")
+
async def _send(self, command: Command):
while self.paused:
await self.paused_ev.wait()
diff --git a/src/exo/master/image_store.py b/src/exo/master/image_store.py
new file mode 100644
index 00000000..c9c3550c
--- /dev/null
+++ b/src/exo/master/image_store.py
@@ -0,0 +1,79 @@
+import time
+from pathlib import Path
+
+from pydantic import BaseModel
+
+from exo.shared.types.common import Id
+
+
+class StoredImage(BaseModel, frozen=True):
+ image_id: Id
+ file_path: Path
+ content_type: str
+ expires_at: float
+
+
+class ImageStore:
+ def __init__(self, storage_dir: Path, default_expiry_seconds: int = 3600) -> None:
+ self._storage_dir = storage_dir
+ self._default_expiry_seconds = default_expiry_seconds
+ self._images: dict[Id, StoredImage] = {}
+ self._storage_dir.mkdir(parents=True, exist_ok=True)
+
+ def store(self, image_bytes: bytes, content_type: str) -> StoredImage:
+ image_id = Id()
+ extension = _content_type_to_extension(content_type)
+ file_path = self._storage_dir / f"{image_id}{extension}"
+ file_path.write_bytes(image_bytes)
+
+ stored = StoredImage(
+ image_id=image_id,
+ file_path=file_path,
+ content_type=content_type,
+ expires_at=time.time() + self._default_expiry_seconds,
+ )
+ self._images[image_id] = stored
+ return stored
+
+ def get(self, image_id: Id) -> StoredImage | None:
+ stored = self._images.get(image_id)
+ if stored is None:
+ return None
+
+ if time.time() > stored.expires_at:
+ self._remove(image_id)
+ return None
+
+ return stored
+
+ def list_images(self) -> list[StoredImage]:
+ now = time.time()
+ return [stored for stored in self._images.values() if now <= stored.expires_at]
+
+ def cleanup_expired(self) -> int:
+ now = time.time()
+ expired_ids = [
+ image_id
+ for image_id, stored in self._images.items()
+ if now > stored.expires_at
+ ]
+
+ for image_id in expired_ids:
+ self._remove(image_id)
+
+ return len(expired_ids)
+
+ def _remove(self, image_id: Id) -> None:
+ stored = self._images.pop(image_id, None)
+ if stored is not None and stored.file_path.exists():
+ stored.file_path.unlink()
+
+
+def _content_type_to_extension(
+ content_type: str,
+) -> str:
+ ext = f"{content_type.split('/')[1]}"
+ if ext == "jpeg":
+ ext = "jpg"
+
+ return f".{ext}"
diff --git a/src/exo/master/main.py b/src/exo/master/main.py
index e1bb4997..407f68d4 100644
--- a/src/exo/master/main.py
+++ b/src/exo/master/main.py
@@ -16,8 +16,11 @@ from exo.shared.types.commands import (
CreateInstance,
DeleteInstance,
ForwarderCommand,
+ ImageEdits,
+ ImageGeneration,
PlaceInstance,
RequestEventLog,
+ SendInputChunk,
TaskFinished,
TestCommand,
)
@@ -26,6 +29,7 @@ from exo.shared.types.events import (
Event,
ForwarderEvent,
IndexedEvent,
+ InputChunkReceived,
InstanceDeleted,
NodeGatheredInfo,
NodeTimedOut,
@@ -36,6 +40,12 @@ from exo.shared.types.state import State
from exo.shared.types.tasks import (
ChatCompletion as ChatCompletionTask,
)
+from exo.shared.types.tasks import (
+ ImageEdits as ImageEditsTask,
+)
+from exo.shared.types.tasks import (
+ ImageGeneration as ImageGenerationTask,
+)
from exo.shared.types.tasks import (
TaskId,
TaskStatus,
@@ -100,13 +110,14 @@ class Master:
async for forwarder_command in commands:
try:
logger.info(f"Executing command: {forwarder_command.command}")
+
generated_events: list[Event] = []
command = forwarder_command.command
+ instance_task_counts: dict[InstanceId, int] = {}
match command:
case TestCommand():
pass
case ChatCompletion():
- instance_task_counts: dict[InstanceId, int] = {}
for instance in self.state.instances.values():
if (
instance.shard_assignments.model_id
@@ -147,6 +158,90 @@ class Master:
)
)
+ self.command_task_mapping[command.command_id] = task_id
+ case ImageGeneration():
+ for instance in self.state.instances.values():
+ if (
+ instance.shard_assignments.model_id
+ == command.request_params.model
+ ):
+ task_count = sum(
+ 1
+ for task in self.state.tasks.values()
+ if task.instance_id == instance.instance_id
+ )
+ instance_task_counts[instance.instance_id] = (
+ task_count
+ )
+
+ if not instance_task_counts:
+ raise ValueError(
+ f"No instance found for model {command.request_params.model}"
+ )
+
+ available_instance_ids = sorted(
+ instance_task_counts.keys(),
+ key=lambda instance_id: instance_task_counts[
+ instance_id
+ ],
+ )
+
+ task_id = TaskId()
+ generated_events.append(
+ TaskCreated(
+ task_id=task_id,
+ task=ImageGenerationTask(
+ task_id=task_id,
+ command_id=command.command_id,
+ instance_id=available_instance_ids[0],
+ task_status=TaskStatus.Pending,
+ task_params=command.request_params,
+ ),
+ )
+ )
+
+ self.command_task_mapping[command.command_id] = task_id
+ case ImageEdits():
+ for instance in self.state.instances.values():
+ if (
+ instance.shard_assignments.model_id
+ == command.request_params.model
+ ):
+ task_count = sum(
+ 1
+ for task in self.state.tasks.values()
+ if task.instance_id == instance.instance_id
+ )
+ instance_task_counts[instance.instance_id] = (
+ task_count
+ )
+
+ if not instance_task_counts:
+ raise ValueError(
+ f"No instance found for model {command.request_params.model}"
+ )
+
+ available_instance_ids = sorted(
+ instance_task_counts.keys(),
+ key=lambda instance_id: instance_task_counts[
+ instance_id
+ ],
+ )
+
+ task_id = TaskId()
+ generated_events.append(
+ TaskCreated(
+ task_id=task_id,
+ task=ImageEditsTask(
+ task_id=task_id,
+ command_id=command.command_id,
+ instance_id=available_instance_ids[0],
+ task_status=TaskStatus.Pending,
+ task_params=command.request_params,
+ ),
+ )
+ )
+
self.command_task_mapping[command.command_id] = task_id
case DeleteInstance():
placement = delete_instance(command, self.state.instances)
@@ -176,6 +271,13 @@ class Master:
self.state.instances, placement
)
generated_events.extend(transition_events)
+ case SendInputChunk(chunk=chunk):
+ generated_events.append(
+ InputChunkReceived(
+ command_id=chunk.command_id,
+ chunk=chunk,
+ )
+ )
case TaskFinished():
generated_events.append(
TaskDeleted(
diff --git a/src/exo/master/tests/test_api_error_handling.py b/src/exo/master/tests/test_api_error_handling.py
index 76001043..85a3f688 100644
--- a/src/exo/master/tests/test_api_error_handling.py
+++ b/src/exo/master/tests/test_api_error_handling.py
@@ -5,7 +5,7 @@ from fastapi import FastAPI, HTTPException
from fastapi.testclient import TestClient
from exo.shared.types.api import ErrorInfo, ErrorResponse, FinishReason
-from exo.shared.types.chunks import TokenChunk
+from exo.shared.types.chunks import ImageChunk, TokenChunk
from exo.worker.tests.constants import MODEL_A_ID
@@ -105,3 +105,38 @@ def test_normal_finish_reasons_still_work() -> None:
finish_reason=reason, # type: ignore[arg-type]
)
assert chunk.finish_reason == reason
+
+
+def test_image_chunk_with_error_fields() -> None:
+ chunk = ImageChunk(
+ idx=0,
+ model=MODEL_A_ID,
+ data="",
+ chunk_index=0,
+ total_chunks=1,
+ image_index=0,
+ finish_reason="error",
+ error_message="Image generation failed",
+ )
+
+ assert chunk.finish_reason == "error"
+ assert chunk.error_message == "Image generation failed"
+ assert chunk.data == ""
+ assert chunk.chunk_index == 0
+ assert chunk.total_chunks == 1
+ assert chunk.image_index == 0
+
+
+def test_image_chunk_without_error() -> None:
+ chunk = ImageChunk(
+ idx=0,
+ model=MODEL_A_ID,
+ data="base64encodeddata",
+ chunk_index=0,
+ total_chunks=1,
+ image_index=0,
+ )
+
+ assert chunk.finish_reason is None
+ assert chunk.error_message is None
+ assert chunk.data == "base64encodeddata"
diff --git a/src/exo/master/tests/test_master.py b/src/exo/master/tests/test_master.py
index 0c1cbea8..4796c264 100644
--- a/src/exo/master/tests/test_master.py
+++ b/src/exo/master/tests/test_master.py
@@ -7,7 +7,7 @@ from loguru import logger
from exo.master.main import Master
from exo.routing.router import get_node_id_keypair
-from exo.shared.models.model_cards import ModelCard, ModelId
+from exo.shared.models.model_cards import ModelCard, ModelId, ModelTask
from exo.shared.types.api import ChatCompletionMessage, ChatCompletionTaskParams
from exo.shared.types.commands import (
ChatCompletion,
@@ -115,6 +115,7 @@ async def test_master():
storage_size=Memory.from_bytes(678948),
hidden_size=7168,
supports_tensor=True,
+ tasks=[ModelTask.TextGeneration],
),
sharding=Sharding.Pipeline,
instance_meta=InstanceMeta.MlxRing,
@@ -172,6 +173,7 @@ async def test_master():
storage_size=Memory.from_bytes(678948),
hidden_size=7168,
supports_tensor=True,
+ tasks=[ModelTask.TextGeneration],
),
device_rank=0,
world_size=1,
diff --git a/src/exo/master/tests/test_placement.py b/src/exo/master/tests/test_placement.py
index ea2501e9..baa1a0f8 100644
--- a/src/exo/master/tests/test_placement.py
+++ b/src/exo/master/tests/test_placement.py
@@ -10,7 +10,7 @@ from exo.master.tests.conftest import (
create_rdma_connection,
create_socket_connection,
)
-from exo.shared.models.model_cards import ModelCard, ModelId
+from exo.shared.models.model_cards import ModelCard, ModelId, ModelTask
from exo.shared.topology import Topology
from exo.shared.types.commands import PlaceInstance
from exo.shared.types.common import CommandId, NodeId
@@ -50,6 +50,7 @@ def model_card() -> ModelCard:
n_layers=10,
hidden_size=30,
supports_tensor=True,
+ tasks=[ModelTask.TextGeneration],
)
@@ -169,6 +170,7 @@ def test_get_instance_placements_one_node_exact_fit() -> None:
n_layers=10,
hidden_size=1000,
supports_tensor=True,
+ tasks=[ModelTask.TextGeneration],
),
)
placements = place_instance(cic, topology, {}, node_memory, node_network)
@@ -195,6 +197,7 @@ def test_get_instance_placements_one_node_fits_with_extra_memory() -> None:
n_layers=10,
hidden_size=1000,
supports_tensor=True,
+ tasks=[ModelTask.TextGeneration],
),
)
placements = place_instance(cic, topology, {}, node_memory, node_network)
@@ -221,6 +224,7 @@ def test_get_instance_placements_one_node_not_fit() -> None:
n_layers=10,
hidden_size=1000,
supports_tensor=True,
+ tasks=[ModelTask.TextGeneration],
),
)
diff --git a/src/exo/master/tests/test_placement_utils.py b/src/exo/master/tests/test_placement_utils.py
index d26cf0c6..871e1b84 100644
--- a/src/exo/master/tests/test_placement_utils.py
+++ b/src/exo/master/tests/test_placement_utils.py
@@ -12,7 +12,7 @@ from exo.master.tests.conftest import (
create_node_memory,
create_socket_connection,
)
-from exo.shared.models.model_cards import ModelCard, ModelId
+from exo.shared.models.model_cards import ModelCard, ModelId, ModelTask
from exo.shared.topology import Topology
from exo.shared.types.common import Host, NodeId
from exo.shared.types.memory import Memory
@@ -238,6 +238,7 @@ def test_get_shard_assignments(
storage_size=Memory.from_kb(1000),
hidden_size=1000,
supports_tensor=True,
+ tasks=[ModelTask.TextGeneration],
)
cycles = topology.get_cycles()
@@ -517,6 +518,7 @@ def test_get_shard_assignments_insufficient_memory_raises():
storage_size=Memory.from_kb(1000),
hidden_size=1000,
supports_tensor=True,
+ tasks=[ModelTask.TextGeneration],
)
cycles = topology.get_cycles()
selected_cycle = cycles[0]
diff --git a/src/exo/shared/apply.py b/src/exo/shared/apply.py
index b31281af..5c1f6d03 100644
--- a/src/exo/shared/apply.py
+++ b/src/exo/shared/apply.py
@@ -9,6 +9,7 @@ from exo.shared.types.events import (
ChunkGenerated,
Event,
IndexedEvent,
+ InputChunkReceived,
InstanceCreated,
InstanceDeleted,
NodeDownloadProgress,
@@ -52,8 +53,8 @@ def event_apply(event: Event, state: State) -> State:
"""Apply an event to state."""
match event:
case (
- TestEvent() | ChunkGenerated() | TaskAcknowledged()
- ): # TaskAcknowledged should never be sent by a worker but i dont mind if it just gets ignored
+ TestEvent() | ChunkGenerated() | TaskAcknowledged() | InputChunkReceived()
+ ): # Pass-through events that don't modify state
return state
case InstanceCreated():
return apply_instance_created(event, state)
diff --git a/src/exo/shared/constants.py b/src/exo/shared/constants.py
index 7112895e..242d520d 100644
--- a/src/exo/shared/constants.py
+++ b/src/exo/shared/constants.py
@@ -45,3 +45,7 @@ LIBP2P_LOCAL_EVENTS_TOPIC = "worker_events"
LIBP2P_GLOBAL_EVENTS_TOPIC = "global_events"
LIBP2P_ELECTION_MESSAGES_TOPIC = "election_message"
LIBP2P_COMMANDS_TOPIC = "commands"
+
+EXO_MAX_CHUNK_SIZE = 512 * 1024
+
+EXO_IMAGE_CACHE_DIR = EXO_CACHE_HOME / "images"
diff --git a/src/exo/shared/models/model_cards.py b/src/exo/shared/models/model_cards.py
index 8f593316..c58f321a 100644
--- a/src/exo/shared/models/model_cards.py
+++ b/src/exo/shared/models/model_cards.py
@@ -1,3 +1,4 @@
+from enum import Enum
from typing import Annotated
import aiofiles
@@ -6,7 +7,7 @@ import tomlkit
from anyio import Path, open_file
from huggingface_hub import model_info
from loguru import logger
-from pydantic import BaseModel, Field, PositiveInt
+from pydantic import BaseModel, Field, PositiveInt, field_validator
from exo.shared.types.common import ModelId
from exo.shared.types.memory import Memory
@@ -15,12 +16,34 @@ from exo.utils.pydantic_ext import CamelCaseModel
_card_cache: dict[str, "ModelCard"] = {}
+class ModelTask(str, Enum):
+ TextGeneration = "TextGeneration"
+ TextToImage = "TextToImage"
+ ImageToImage = "ImageToImage"
+
+
+class ComponentInfo(CamelCaseModel):
+ component_name: str
+ component_path: str
+ storage_size: Memory
+ n_layers: PositiveInt | None
+ can_shard: bool
+ safetensors_index_filename: str | None
+
+
class ModelCard(CamelCaseModel):
model_id: ModelId
storage_size: Memory
n_layers: PositiveInt
hidden_size: PositiveInt
supports_tensor: bool
+ tasks: list[ModelTask]
+ components: list[ComponentInfo] | None = None
+
+ @field_validator("tasks", mode="before")
+ @classmethod
+ def _validate_tasks(cls, v: list[str | ModelTask]) -> list[ModelTask]:
+ return [item if isinstance(item, ModelTask) else ModelTask(item) for item in v]
async def save(self, path: Path) -> None:
async with await open_file(path, "w") as f:
@@ -55,6 +78,7 @@ class ModelCard(CamelCaseModel):
n_layers=num_layers,
hidden_size=config_data.hidden_size or 0,
supports_tensor=config_data.supports_tensor,
+ tasks=[ModelTask.TextGeneration],
)
_card_cache[model_id] = mc
return mc
@@ -68,6 +92,7 @@ MODEL_CARDS: dict[str, ModelCard] = {
n_layers=61,
hidden_size=7168,
supports_tensor=True,
+ tasks=[ModelTask.TextGeneration],
),
"deepseek-v3.1-8bit": ModelCard(
model_id=ModelId("mlx-community/DeepSeek-V3.1-8bit"),
@@ -75,6 +100,7 @@ MODEL_CARDS: dict[str, ModelCard] = {
n_layers=61,
hidden_size=7168,
supports_tensor=True,
+ tasks=[ModelTask.TextGeneration],
),
# kimi k2
"kimi-k2-instruct-4bit": ModelCard(
@@ -83,6 +109,7 @@ MODEL_CARDS: dict[str, ModelCard] = {
n_layers=61,
hidden_size=7168,
supports_tensor=True,
+ tasks=[ModelTask.TextGeneration],
),
"kimi-k2-thinking": ModelCard(
model_id=ModelId("mlx-community/Kimi-K2-Thinking"),
@@ -90,6 +117,7 @@ MODEL_CARDS: dict[str, ModelCard] = {
n_layers=61,
hidden_size=7168,
supports_tensor=True,
+ tasks=[ModelTask.TextGeneration],
),
# llama-3.1
"llama-3.1-8b": ModelCard(
@@ -98,6 +126,7 @@ MODEL_CARDS: dict[str, ModelCard] = {
n_layers=32,
hidden_size=4096,
supports_tensor=True,
+ tasks=[ModelTask.TextGeneration],
),
"llama-3.1-8b-8bit": ModelCard(
model_id=ModelId("mlx-community/Meta-Llama-3.1-8B-Instruct-8bit"),
@@ -105,6 +134,7 @@ MODEL_CARDS: dict[str, ModelCard] = {
n_layers=32,
hidden_size=4096,
supports_tensor=True,
+ tasks=[ModelTask.TextGeneration],
),
"llama-3.1-8b-bf16": ModelCard(
model_id=ModelId("mlx-community/Meta-Llama-3.1-8B-Instruct-bf16"),
@@ -112,6 +142,7 @@ MODEL_CARDS: dict[str, ModelCard] = {
n_layers=32,
hidden_size=4096,
supports_tensor=True,
+ tasks=[ModelTask.TextGeneration],
),
"llama-3.1-70b": ModelCard(
model_id=ModelId("mlx-community/Meta-Llama-3.1-70B-Instruct-4bit"),
@@ -119,6 +150,7 @@ MODEL_CARDS: dict[str, ModelCard] = {
n_layers=80,
hidden_size=8192,
supports_tensor=True,
+ tasks=[ModelTask.TextGeneration],
),
# llama-3.2
"llama-3.2-1b": ModelCard(
@@ -127,6 +159,7 @@ MODEL_CARDS: dict[str, ModelCard] = {
n_layers=16,
hidden_size=2048,
supports_tensor=True,
+ tasks=[ModelTask.TextGeneration],
),
"llama-3.2-3b": ModelCard(
model_id=ModelId("mlx-community/Llama-3.2-3B-Instruct-4bit"),
@@ -134,6 +167,7 @@ MODEL_CARDS: dict[str, ModelCard] = {
n_layers=28,
hidden_size=3072,
supports_tensor=True,
+ tasks=[ModelTask.TextGeneration],
),
"llama-3.2-3b-8bit": ModelCard(
model_id=ModelId("mlx-community/Llama-3.2-3B-Instruct-8bit"),
@@ -141,6 +175,7 @@ MODEL_CARDS: dict[str, ModelCard] = {
n_layers=28,
hidden_size=3072,
supports_tensor=True,
+ tasks=[ModelTask.TextGeneration],
),
# llama-3.3
"llama-3.3-70b": ModelCard(
@@ -149,6 +184,7 @@ MODEL_CARDS: dict[str, ModelCard] = {
n_layers=80,
hidden_size=8192,
supports_tensor=True,
+ tasks=[ModelTask.TextGeneration],
),
"llama-3.3-70b-8bit": ModelCard(
model_id=ModelId("mlx-community/Llama-3.3-70B-Instruct-8bit"),
@@ -156,6 +192,7 @@ MODEL_CARDS: dict[str, ModelCard] = {
n_layers=80,
hidden_size=8192,
supports_tensor=True,
+ tasks=[ModelTask.TextGeneration],
),
"llama-3.3-70b-fp16": ModelCard(
model_id=ModelId("mlx-community/llama-3.3-70b-instruct-fp16"),
@@ -163,6 +200,7 @@ MODEL_CARDS: dict[str, ModelCard] = {
n_layers=80,
hidden_size=8192,
supports_tensor=True,
+ tasks=[ModelTask.TextGeneration],
),
# qwen3
"qwen3-0.6b": ModelCard(
@@ -171,6 +209,7 @@ MODEL_CARDS: dict[str, ModelCard] = {
n_layers=28,
hidden_size=1024,
supports_tensor=False,
+ tasks=[ModelTask.TextGeneration],
),
"qwen3-0.6b-8bit": ModelCard(
model_id=ModelId("mlx-community/Qwen3-0.6B-8bit"),
@@ -178,6 +217,7 @@ MODEL_CARDS: dict[str, ModelCard] = {
n_layers=28,
hidden_size=1024,
supports_tensor=False,
+ tasks=[ModelTask.TextGeneration],
),
"qwen3-30b": ModelCard(
model_id=ModelId("mlx-community/Qwen3-30B-A3B-4bit"),
@@ -185,6 +225,7 @@ MODEL_CARDS: dict[str, ModelCard] = {
n_layers=48,
hidden_size=2048,
supports_tensor=True,
+ tasks=[ModelTask.TextGeneration],
),
"qwen3-30b-8bit": ModelCard(
model_id=ModelId("mlx-community/Qwen3-30B-A3B-8bit"),
@@ -192,6 +233,7 @@ MODEL_CARDS: dict[str, ModelCard] = {
n_layers=48,
hidden_size=2048,
supports_tensor=True,
+ tasks=[ModelTask.TextGeneration],
),
"qwen3-80b-a3B-4bit": ModelCard(
model_id=ModelId("mlx-community/Qwen3-Next-80B-A3B-Instruct-4bit"),
@@ -199,6 +241,7 @@ MODEL_CARDS: dict[str, ModelCard] = {
n_layers=48,
hidden_size=2048,
supports_tensor=True,
+ tasks=[ModelTask.TextGeneration],
),
"qwen3-80b-a3B-8bit": ModelCard(
model_id=ModelId("mlx-community/Qwen3-Next-80B-A3B-Instruct-8bit"),
@@ -206,6 +249,7 @@ MODEL_CARDS: dict[str, ModelCard] = {
n_layers=48,
hidden_size=2048,
supports_tensor=True,
+ tasks=[ModelTask.TextGeneration],
),
"qwen3-80b-a3B-thinking-4bit": ModelCard(
model_id=ModelId("mlx-community/Qwen3-Next-80B-A3B-Thinking-4bit"),
@@ -213,6 +257,7 @@ MODEL_CARDS: dict[str, ModelCard] = {
n_layers=48,
hidden_size=2048,
supports_tensor=True,
+ tasks=[ModelTask.TextGeneration],
),
"qwen3-80b-a3B-thinking-8bit": ModelCard(
model_id=ModelId("mlx-community/Qwen3-Next-80B-A3B-Thinking-8bit"),
@@ -220,6 +265,7 @@ MODEL_CARDS: dict[str, ModelCard] = {
n_layers=48,
hidden_size=2048,
supports_tensor=True,
+ tasks=[ModelTask.TextGeneration],
),
"qwen3-235b-a22b-4bit": ModelCard(
model_id=ModelId("mlx-community/Qwen3-235B-A22B-Instruct-2507-4bit"),
@@ -227,6 +273,7 @@ MODEL_CARDS: dict[str, ModelCard] = {
n_layers=94,
hidden_size=4096,
supports_tensor=True,
+ tasks=[ModelTask.TextGeneration],
),
"qwen3-235b-a22b-8bit": ModelCard(
model_id=ModelId("mlx-community/Qwen3-235B-A22B-Instruct-2507-8bit"),
@@ -234,6 +281,7 @@ MODEL_CARDS: dict[str, ModelCard] = {
n_layers=94,
hidden_size=4096,
supports_tensor=True,
+ tasks=[ModelTask.TextGeneration],
),
"qwen3-coder-480b-a35b-4bit": ModelCard(
model_id=ModelId("mlx-community/Qwen3-Coder-480B-A35B-Instruct-4bit"),
@@ -241,6 +289,7 @@ MODEL_CARDS: dict[str, ModelCard] = {
n_layers=62,
hidden_size=6144,
supports_tensor=True,
+ tasks=[ModelTask.TextGeneration],
),
"qwen3-coder-480b-a35b-8bit": ModelCard(
model_id=ModelId("mlx-community/Qwen3-Coder-480B-A35B-Instruct-8bit"),
@@ -248,6 +297,7 @@ MODEL_CARDS: dict[str, ModelCard] = {
n_layers=62,
hidden_size=6144,
supports_tensor=True,
+ tasks=[ModelTask.TextGeneration],
),
# gpt-oss
"gpt-oss-120b-MXFP4-Q8": ModelCard(
@@ -256,6 +306,7 @@ MODEL_CARDS: dict[str, ModelCard] = {
n_layers=36,
hidden_size=2880,
supports_tensor=True,
+ tasks=[ModelTask.TextGeneration],
),
"gpt-oss-20b-MXFP4-Q8": ModelCard(
model_id=ModelId("mlx-community/gpt-oss-20b-MXFP4-Q8"),
@@ -263,6 +314,7 @@ MODEL_CARDS: dict[str, ModelCard] = {
n_layers=24,
hidden_size=2880,
supports_tensor=True,
+ tasks=[ModelTask.TextGeneration],
),
# glm 4.5
"glm-4.5-air-8bit": ModelCard(
@@ -272,6 +324,7 @@ MODEL_CARDS: dict[str, ModelCard] = {
n_layers=46,
hidden_size=4096,
supports_tensor=False,
+ tasks=[ModelTask.TextGeneration],
),
"glm-4.5-air-bf16": ModelCard(
model_id=ModelId("mlx-community/GLM-4.5-Air-bf16"),
@@ -279,6 +332,7 @@ MODEL_CARDS: dict[str, ModelCard] = {
n_layers=46,
hidden_size=4096,
supports_tensor=True,
+ tasks=[ModelTask.TextGeneration],
),
# glm 4.7
"glm-4.7-4bit": ModelCard(
@@ -287,6 +341,7 @@ MODEL_CARDS: dict[str, ModelCard] = {
n_layers=91,
hidden_size=5120,
supports_tensor=True,
+ tasks=[ModelTask.TextGeneration],
),
"glm-4.7-6bit": ModelCard(
model_id=ModelId("mlx-community/GLM-4.7-6bit"),
@@ -294,6 +349,7 @@ MODEL_CARDS: dict[str, ModelCard] = {
n_layers=91,
hidden_size=5120,
supports_tensor=True,
+ tasks=[ModelTask.TextGeneration],
),
"glm-4.7-8bit-gs32": ModelCard(
model_id=ModelId("mlx-community/GLM-4.7-8bit-gs32"),
@@ -301,6 +357,7 @@ MODEL_CARDS: dict[str, ModelCard] = {
n_layers=91,
hidden_size=5120,
supports_tensor=True,
+ tasks=[ModelTask.TextGeneration],
),
# glm 4.7 flash
"glm-4.7-flash-4bit": ModelCard(
@@ -309,6 +366,7 @@ MODEL_CARDS: dict[str, ModelCard] = {
n_layers=47,
hidden_size=2048,
supports_tensor=True,
+ tasks=[ModelTask.TextGeneration],
),
"glm-4.7-flash-5bit": ModelCard(
model_id=ModelId("mlx-community/GLM-4.7-Flash-5bit"),
@@ -316,6 +374,7 @@ MODEL_CARDS: dict[str, ModelCard] = {
n_layers=47,
hidden_size=2048,
supports_tensor=True,
+ tasks=[ModelTask.TextGeneration],
),
"glm-4.7-flash-6bit": ModelCard(
model_id=ModelId("mlx-community/GLM-4.7-Flash-6bit"),
@@ -323,6 +382,7 @@ MODEL_CARDS: dict[str, ModelCard] = {
n_layers=47,
hidden_size=2048,
supports_tensor=True,
+ tasks=[ModelTask.TextGeneration],
),
"glm-4.7-flash-8bit": ModelCard(
model_id=ModelId("mlx-community/GLM-4.7-Flash-8bit"),
@@ -330,6 +390,7 @@ MODEL_CARDS: dict[str, ModelCard] = {
n_layers=47,
hidden_size=2048,
supports_tensor=True,
+ tasks=[ModelTask.TextGeneration],
),
# minimax-m2
"minimax-m2.1-8bit": ModelCard(
@@ -338,6 +399,7 @@ MODEL_CARDS: dict[str, ModelCard] = {
n_layers=61,
hidden_size=3072,
supports_tensor=True,
+ tasks=[ModelTask.TextGeneration],
),
"minimax-m2.1-3bit": ModelCard(
model_id=ModelId("mlx-community/MiniMax-M2.1-3bit"),
@@ -345,6 +407,159 @@ MODEL_CARDS: dict[str, ModelCard] = {
n_layers=61,
hidden_size=3072,
supports_tensor=True,
+ tasks=[ModelTask.TextGeneration],
+ ),
+ "flux1-schnell": ModelCard(
+ model_id=ModelId("black-forest-labs/FLUX.1-schnell"),
+ storage_size=Memory.from_bytes(23782357120 + 9524621312),
+ n_layers=57,
+ hidden_size=1,
+ supports_tensor=False,
+ tasks=[ModelTask.TextToImage],
+ components=[
+ ComponentInfo(
+ component_name="text_encoder",
+ component_path="text_encoder/",
+ storage_size=Memory.from_kb(0),
+ n_layers=12,
+ can_shard=False,
+ safetensors_index_filename=None, # Single file
+ ),
+ ComponentInfo(
+ component_name="text_encoder_2",
+ component_path="text_encoder_2/",
+ storage_size=Memory.from_bytes(9524621312),
+ n_layers=24,
+ can_shard=False,
+ safetensors_index_filename="model.safetensors.index.json",
+ ),
+ ComponentInfo(
+ component_name="transformer",
+ component_path="transformer/",
+ storage_size=Memory.from_bytes(23782357120),
+ n_layers=57, # 19 transformer_blocks + 38 single_transformer_blocks
+ can_shard=True,
+ safetensors_index_filename="diffusion_pytorch_model.safetensors.index.json",
+ ),
+ ComponentInfo(
+ component_name="vae",
+ component_path="vae/",
+ storage_size=Memory.from_kb(0),
+ n_layers=None,
+ can_shard=False,
+ safetensors_index_filename=None,
+ ),
+ ],
+ ),
+ "flux1-dev": ModelCard(
+ model_id=ModelId("black-forest-labs/FLUX.1-dev"),
+ storage_size=Memory.from_bytes(23782357120 + 9524621312),
+ n_layers=57,
+ hidden_size=1,
+ supports_tensor=False,
+ tasks=[ModelTask.TextToImage, ModelTask.ImageToImage],
+ components=[
+ ComponentInfo(
+ component_name="text_encoder",
+ component_path="text_encoder/",
+ storage_size=Memory.from_kb(0),
+ n_layers=12,
+ can_shard=False,
+ safetensors_index_filename=None, # Single file
+ ),
+ ComponentInfo(
+ component_name="text_encoder_2",
+ component_path="text_encoder_2/",
+ storage_size=Memory.from_bytes(9524621312),
+ n_layers=24,
+ can_shard=False,
+ safetensors_index_filename="model.safetensors.index.json",
+ ),
+ ComponentInfo(
+ component_name="transformer",
+ component_path="transformer/",
+ storage_size=Memory.from_bytes(23802816640),
+ n_layers=57, # 19 transformer_blocks + 38 single_transformer_blocks
+ can_shard=True,
+ safetensors_index_filename="diffusion_pytorch_model.safetensors.index.json",
+ ),
+ ComponentInfo(
+ component_name="vae",
+ component_path="vae/",
+ storage_size=Memory.from_kb(0),
+ n_layers=None,
+ can_shard=False,
+ safetensors_index_filename=None,
+ ),
+ ],
+ ),
+ "qwen-image": ModelCard(
+ model_id=ModelId("Qwen/Qwen-Image"),
+ storage_size=Memory.from_bytes(16584333312 + 40860802176),
+ n_layers=60, # Qwen has 60 transformer blocks (all joint-style)
+ hidden_size=1,
+ supports_tensor=False,
+ tasks=[ModelTask.TextToImage, ModelTask.ImageToImage],
+ components=[
+ ComponentInfo(
+ component_name="text_encoder",
+ component_path="text_encoder/",
+ storage_size=Memory.from_kb(16584333312),
+ n_layers=12,
+ can_shard=False,
+ safetensors_index_filename=None, # Single file
+ ),
+ ComponentInfo(
+ component_name="transformer",
+ component_path="transformer/",
+ storage_size=Memory.from_bytes(40860802176),
+ n_layers=60,
+ can_shard=True,
+ safetensors_index_filename="diffusion_pytorch_model.safetensors.index.json",
+ ),
+ ComponentInfo(
+ component_name="vae",
+ component_path="vae/",
+ storage_size=Memory.from_kb(0),
+ n_layers=None,
+ can_shard=False,
+ safetensors_index_filename=None,
+ ),
+ ],
+ ),
+ "qwen-image-edit-2509": ModelCard(
+ model_id=ModelId("Qwen/Qwen-Image-Edit-2509"),
+ storage_size=Memory.from_bytes(16584333312 + 40860802176),
+ n_layers=60, # Qwen has 60 transformer blocks (all joint-style)
+ hidden_size=1,
+ supports_tensor=False,
+ tasks=[ModelTask.ImageToImage],
+ components=[
+ ComponentInfo(
+ component_name="text_encoder",
+ component_path="text_encoder/",
+ storage_size=Memory.from_kb(16584333312),
+ n_layers=12,
+ can_shard=False,
+ safetensors_index_filename=None, # Single file
+ ),
+ ComponentInfo(
+ component_name="transformer",
+ component_path="transformer/",
+ storage_size=Memory.from_bytes(40860802176),
+ n_layers=60,
+ can_shard=True,
+ safetensors_index_filename="diffusion_pytorch_model.safetensors.index.json",
+ ),
+ ComponentInfo(
+ component_name="vae",
+ component_path="vae/",
+ storage_size=Memory.from_kb(0),
+ n_layers=None,
+ can_shard=False,
+ safetensors_index_filename=None,
+ ),
+ ],
),
}
diff --git a/src/exo/shared/tests/conftest.py b/src/exo/shared/tests/conftest.py
index 7ed28ad9..0622cb0b 100644
--- a/src/exo/shared/tests/conftest.py
+++ b/src/exo/shared/tests/conftest.py
@@ -7,7 +7,7 @@ import pytest
from _pytest.logging import LogCaptureFixture
from loguru import logger
-from exo.shared.models.model_cards import ModelCard, ModelId
+from exo.shared.models.model_cards import ModelCard, ModelId, ModelTask
from exo.shared.types.memory import Memory
from exo.shared.types.worker.shards import PipelineShardMetadata, ShardMetadata
@@ -37,6 +37,7 @@ def get_pipeline_shard_metadata(
n_layers=32,
hidden_size=1000,
supports_tensor=True,
+ tasks=[ModelTask.TextGeneration],
),
device_rank=device_rank,
world_size=world_size,
diff --git a/src/exo/shared/types/api.py b/src/exo/shared/types/api.py
index 479dd210..c43f95f4 100644
--- a/src/exo/shared/types/api.py
+++ b/src/exo/shared/types/api.py
@@ -1,6 +1,8 @@
import time
-from typing import Any, Literal
+from collections.abc import Generator
+from typing import Annotated, Any, Literal
+from fastapi import UploadFile
from pydantic import BaseModel, Field, field_validator
from pydantic_core import PydanticUseDefault
@@ -39,6 +41,7 @@ class ModelListModel(BaseModel):
tags: list[str] = Field(default=[])
storage_size_megabytes: int = Field(default=0)
supports_tensor: bool = Field(default=False)
+ tasks: list[str] = Field(default=[])
class ModelList(BaseModel):
@@ -137,6 +140,19 @@ class GenerationStats(BaseModel):
peak_memory_usage: Memory
+class ImageGenerationStats(BaseModel):
+ seconds_per_step: float
+ total_generation_time: float
+
+ num_inference_steps: int
+ num_images: int
+
+ image_width: int
+ image_height: int
+
+ peak_memory_usage: Memory
+
+
class BenchChatCompletionResponse(ChatCompletionResponse):
generation_stats: GenerationStats | None = None
@@ -213,3 +229,114 @@ class DeleteInstanceResponse(BaseModel):
message: str
command_id: CommandId
instance_id: InstanceId
+
+
+class AdvancedImageParams(BaseModel):
+ seed: Annotated[int, Field(ge=0)] | None = None
+ num_inference_steps: Annotated[int, Field(ge=1, le=100)] | None = None
+ guidance: Annotated[float, Field(ge=1.0, le=20.0)] | None = None
+ negative_prompt: str | None = None
+
+
+class ImageGenerationTaskParams(BaseModel):
+ prompt: str
+ background: str | None = None
+ model: str
+ moderation: str | None = None
+ n: int | None = 1
+ output_compression: int | None = None
+ output_format: Literal["png", "jpeg", "webp"] = "png"
+ partial_images: int | None = 0
+ quality: Literal["high", "medium", "low"] | None = "medium"
+ response_format: Literal["url", "b64_json"] | None = "b64_json"
+ size: str | None = "1024x1024"
+ stream: bool | None = False
+ style: str | None = "vivid"
+ user: str | None = None
+ advanced_params: AdvancedImageParams | None = None
+ # Internal flag for benchmark mode - set by API, preserved through serialization
+ bench: bool = False
+
+
+class BenchImageGenerationTaskParams(ImageGenerationTaskParams):
+ bench: bool = True
+
+
+class ImageEditsTaskParams(BaseModel):
+ image: UploadFile
+ prompt: str
+ background: str | None = None
+ input_fidelity: float | None = None
+ mask: UploadFile | None = None
+ model: str
+ n: int | None = 1
+ output_compression: int | None = None
+ output_format: Literal["png", "jpeg", "webp"] = "png"
+ partial_images: int | None = 0
+ quality: Literal["high", "medium", "low"] | None = "medium"
+ response_format: Literal["url", "b64_json"] | None = "b64_json"
+ size: str | None = "1024x1024"
+ stream: bool | None = False
+ user: str | None = None
+ advanced_params: AdvancedImageParams | None = None
+ # Internal flag for benchmark mode - set by API, preserved through serialization
+ bench: bool = False
+
+
+class ImageEditsInternalParams(BaseModel):
+ """Serializable version of ImageEditsTaskParams for distributed task execution."""
+
+ image_data: str = "" # Base64-encoded image (empty when using chunked transfer)
+ total_input_chunks: int = 0
+ prompt: str
+ model: str
+ n: int | None = 1
+ quality: Literal["high", "medium", "low"] | None = "medium"
+ output_format: Literal["png", "jpeg", "webp"] = "png"
+ response_format: Literal["url", "b64_json"] | None = "b64_json"
+ size: str | None = "1024x1024"
+ image_strength: float | None = 0.7
+ stream: bool = False
+ partial_images: int | None = 0
+ advanced_params: AdvancedImageParams | None = None
+ bench: bool = False
+
+ def __repr_args__(self) -> Generator[tuple[str, Any], None, None]:
+ for name, value in super().__repr_args__(): # pyright: ignore[reportAny]
+ if name == "image_data":
+ yield name, f"<{len(self.image_data)} chars>"
+ elif name is not None:
+ yield name, value
+
+
+class ImageData(BaseModel):
+ b64_json: str | None = None
+ url: str | None = None
+ revised_prompt: str | None = None
+
+ def __repr_args__(self) -> Generator[tuple[str, Any], None, None]:
+ for name, value in super().__repr_args__(): # pyright: ignore[reportAny]
+ if name == "b64_json" and self.b64_json is not None:
+ yield name, f"<{len(self.b64_json)} chars>"
+ elif name is not None:
+ yield name, value
+
+
+class ImageGenerationResponse(BaseModel):
+ created: int = Field(default_factory=lambda: int(time.time()))
+ data: list[ImageData]
+
+
+class BenchImageGenerationResponse(ImageGenerationResponse):
+ generation_stats: ImageGenerationStats | None = None
+
+
+class ImageListItem(BaseModel, frozen=True):
+ image_id: str
+ url: str
+ content_type: str
+ expires_at: float
+
+
+class ImageListResponse(BaseModel, frozen=True):
+ data: list[ImageListItem]
diff --git a/src/exo/shared/types/chunks.py b/src/exo/shared/types/chunks.py
index 5e5cfed1..175051cd 100644
--- a/src/exo/shared/types/chunks.py
+++ b/src/exo/shared/types/chunks.py
@@ -1,10 +1,13 @@
+from collections.abc import Generator
from enum import Enum
+from typing import Any, Literal
from exo.shared.models.model_cards import ModelId
-from exo.shared.types.api import GenerationStats
+from exo.shared.types.api import GenerationStats, ImageGenerationStats
from exo.utils.pydantic_ext import TaggedModel
from .api import FinishReason
+from .common import CommandId
class ChunkType(str, Enum):
@@ -26,7 +29,38 @@ class TokenChunk(BaseChunk):
class ImageChunk(BaseChunk):
- data: bytes
+ data: str
+ chunk_index: int
+ total_chunks: int
+ image_index: int
+ is_partial: bool = False
+ partial_index: int | None = None
+ total_partials: int | None = None
+ stats: ImageGenerationStats | None = None
+ format: Literal["png", "jpeg", "webp"] | None = None
+ finish_reason: FinishReason | None = None
+ error_message: str | None = None
+
+ def __repr_args__(self) -> Generator[tuple[str, Any], None, None]:
+ for name, value in super().__repr_args__(): # pyright: ignore[reportAny]
+ if name == "data" and hasattr(value, "__len__"): # pyright: ignore[reportAny]
+ yield name, f"<{len(self.data)} chars>"
+ elif name is not None:
+ yield name, value
+
+
+class InputImageChunk(BaseChunk):
+ command_id: CommandId
+ data: str
+ chunk_index: int
+ total_chunks: int
+
+ def __repr_args__(self) -> Generator[tuple[str, Any], None, None]:
+ for name, value in super().__repr_args__(): # pyright: ignore[reportAny]
+ if name == "data" and hasattr(value, "__len__"): # pyright: ignore[reportAny]
+ yield name, f"<{len(self.data)} chars>"
+ elif name is not None:
+ yield name, value
GenerationChunk = TokenChunk | ImageChunk
diff --git a/src/exo/shared/types/commands.py b/src/exo/shared/types/commands.py
index c264a592..a0ddb968 100644
--- a/src/exo/shared/types/commands.py
+++ b/src/exo/shared/types/commands.py
@@ -1,7 +1,12 @@
from pydantic import Field
from exo.shared.models.model_cards import ModelCard
-from exo.shared.types.api import ChatCompletionTaskParams
+from exo.shared.types.api import (
+ ChatCompletionTaskParams,
+ ImageEditsInternalParams,
+ ImageGenerationTaskParams,
+)
+from exo.shared.types.chunks import InputImageChunk
from exo.shared.types.common import CommandId, NodeId
from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
from exo.shared.types.worker.shards import Sharding
@@ -20,6 +25,14 @@ class ChatCompletion(BaseCommand):
request_params: ChatCompletionTaskParams
+class ImageGeneration(BaseCommand):
+ request_params: ImageGenerationTaskParams
+
+
+class ImageEdits(BaseCommand):
+ request_params: ImageEditsInternalParams
+
+
class PlaceInstance(BaseCommand):
model_card: ModelCard
sharding: Sharding
@@ -39,6 +52,12 @@ class TaskFinished(BaseCommand):
finished_command_id: CommandId
+class SendInputChunk(BaseCommand):
+ """Command to send an input image chunk (converted to event by master)."""
+
+ chunk: InputImageChunk
+
+
class RequestEventLog(BaseCommand):
since_idx: int
@@ -47,10 +66,13 @@ Command = (
TestCommand
| RequestEventLog
| ChatCompletion
+ | ImageGeneration
+ | ImageEdits
| PlaceInstance
| CreateInstance
| DeleteInstance
| TaskFinished
+ | SendInputChunk
)
diff --git a/src/exo/shared/types/events.py b/src/exo/shared/types/events.py
index 779ce546..647e510f 100644
--- a/src/exo/shared/types/events.py
+++ b/src/exo/shared/types/events.py
@@ -3,7 +3,7 @@ from datetime import datetime
from pydantic import Field
from exo.shared.topology import Connection
-from exo.shared.types.chunks import GenerationChunk
+from exo.shared.types.chunks import GenerationChunk, InputImageChunk
from exo.shared.types.common import CommandId, Id, NodeId, SessionId
from exo.shared.types.tasks import Task, TaskId, TaskStatus
from exo.shared.types.worker.downloads import DownloadProgress
@@ -96,6 +96,11 @@ class ChunkGenerated(BaseEvent):
chunk: GenerationChunk
+class InputChunkReceived(BaseEvent):
+ command_id: CommandId
+ chunk: InputImageChunk
+
+
class TopologyEdgeCreated(BaseEvent):
conn: Connection
@@ -119,6 +124,7 @@ Event = (
| NodeGatheredInfo
| NodeDownloadProgress
| ChunkGenerated
+ | InputChunkReceived
| TopologyEdgeCreated
| TopologyEdgeDeleted
)
diff --git a/src/exo/shared/types/tasks.py b/src/exo/shared/types/tasks.py
index 04de0b95..2c58087e 100644
--- a/src/exo/shared/types/tasks.py
+++ b/src/exo/shared/types/tasks.py
@@ -2,7 +2,11 @@ from enum import Enum
from pydantic import Field
-from exo.shared.types.api import ChatCompletionTaskParams
+from exo.shared.types.api import (
+ ChatCompletionTaskParams,
+ ImageEditsInternalParams,
+ ImageGenerationTaskParams,
+)
from exo.shared.types.common import CommandId, Id
from exo.shared.types.worker.instances import BoundInstance, InstanceId
from exo.shared.types.worker.runners import RunnerId
@@ -56,6 +60,22 @@ class ChatCompletion(BaseTask): # emitted by Master
error_message: str | None = Field(default=None)
+class ImageGeneration(BaseTask): # emitted by Master
+ command_id: CommandId
+ task_params: ImageGenerationTaskParams
+
+ error_type: str | None = Field(default=None)
+ error_message: str | None = Field(default=None)
+
+
+class ImageEdits(BaseTask): # emitted by Master
+ command_id: CommandId
+ task_params: ImageEditsInternalParams
+
+ error_type: str | None = Field(default=None)
+ error_message: str | None = Field(default=None)
+
+
class Shutdown(BaseTask): # emitted by Worker
runner_id: RunnerId
@@ -67,5 +87,7 @@ Task = (
| LoadModel
| StartWarmup
| ChatCompletion
+ | ImageGeneration
+ | ImageEdits
| Shutdown
)
diff --git a/src/exo/shared/types/worker/runner_response.py b/src/exo/shared/types/worker/runner_response.py
index 1be5abca..9d36d1cb 100644
--- a/src/exo/shared/types/worker/runner_response.py
+++ b/src/exo/shared/types/worker/runner_response.py
@@ -1,4 +1,7 @@
-from exo.shared.types.api import FinishReason, GenerationStats
+from collections.abc import Generator
+from typing import Any, Literal
+
+from exo.shared.types.api import FinishReason, GenerationStats, ImageGenerationStats
from exo.utils.pydantic_ext import TaggedModel
@@ -18,5 +21,32 @@ class GenerationResponse(BaseRunnerResponse):
stats: GenerationStats | None = None
+class ImageGenerationResponse(BaseRunnerResponse):
+ image_data: bytes
+ format: Literal["png", "jpeg", "webp"] = "png"
+ stats: ImageGenerationStats | None = None
+
+ def __repr_args__(self) -> Generator[tuple[str, Any], None, None]:
+ for name, value in super().__repr_args__(): # pyright: ignore[reportAny]
+ if name == "image_data":
+ yield name, f"<{len(self.image_data)} bytes>"
+ elif name is not 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__(): # pyright: ignore[reportAny]
+ if name == "image_data":
+ yield name, f"<{len(self.image_data)} bytes>"
+ elif name is not None:
+ yield name, value
+
+
class FinishedResponse(BaseRunnerResponse):
pass
diff --git a/src/exo/worker/download/download_utils.py b/src/exo/worker/download/download_utils.py
index 6791267c..c64add0f 100644
--- a/src/exo/worker/download/download_utils.py
+++ b/src/exo/worker/download/download_utils.py
@@ -15,6 +15,9 @@ import aiofiles
import aiofiles.os as aios
import aiohttp
import certifi
+from huggingface_hub import (
+ snapshot_download, # pyright: ignore[reportUnknownVariableType]
+)
from loguru import logger
from pydantic import (
DirectoryPath,
@@ -399,12 +402,33 @@ def calculate_repo_progress(
async def get_weight_map(model_id: ModelId, revision: str = "main") -> dict[str, str]:
target_dir = (await ensure_models_dir()) / model_id.normalize()
await aios.makedirs(target_dir, exist_ok=True)
- index_file = await download_file_with_retry(
- model_id, revision, "model.safetensors.index.json", target_dir
+
+ index_files_dir = snapshot_download(
+ repo_id=model_id,
+ local_dir=target_dir,
+ allow_patterns="*.safetensors.index.json",
)
- async with aiofiles.open(index_file, "r") as f:
- index_data = ModelSafetensorsIndex.model_validate_json(await f.read())
- return index_data.weight_map
+
+ index_files = list(Path(index_files_dir).glob("**/*.safetensors.index.json"))
+
+ weight_map: dict[str, str] = {}
+
+ for index_file in index_files:
+ relative_dir = index_file.parent.relative_to(index_files_dir)
+
+ async with aiofiles.open(index_file, "r") as f:
+ index_data = ModelSafetensorsIndex.model_validate_json(await f.read())
+
+ if relative_dir != Path("."):
+ prefixed_weight_map = {
+ f"{relative_dir}/{key}": str(relative_dir / value)
+ for key, value in index_data.weight_map.items()
+ }
+ weight_map = weight_map | prefixed_weight_map
+ else:
+ weight_map = weight_map | index_data.weight_map
+
+ return weight_map
async def resolve_allow_patterns(shard: ShardMetadata) -> list[str]:
@@ -455,8 +479,6 @@ async def download_shard(
logger.info(f"Downloading {shard.model_card.model_id=} with {allow_patterns=}")
all_start_time = time.time()
- # TODO: currently not recursive. Some models might require subdirectories - thus this will need to be changed.
- # Update: <- This does not seem to be the case. Yay?
file_list = await fetch_file_list_with_cache(
shard.model_card.model_id, revision, recursive=True
)
diff --git a/src/exo/worker/download/huggingface_utils.py b/src/exo/worker/download/huggingface_utils.py
index 8a1b179e..abe12107 100644
--- a/src/exo/worker/download/huggingface_utils.py
+++ b/src/exo/worker/download/huggingface_utils.py
@@ -100,26 +100,68 @@ def get_allow_patterns(weight_map: dict[str, str], shard: ShardMetadata) -> list
"*.py",
"tokenizer.model",
"tiktoken.model",
+ "*/spiece.model",
"*.tiktoken",
"*.txt",
"*.jinja",
]
)
shard_specific_patterns: set[str] = set()
- if weight_map:
- for tensor_name, filename in weight_map.items():
- layer_num = extract_layer_num(tensor_name)
- if (
- layer_num is not None
- and shard.start_layer <= layer_num <= shard.end_layer
- ):
- shard_specific_patterns.add(filename)
- layer_independent_files = set(
- [v for k, v in weight_map.items() if extract_layer_num(k) is None]
+
+ if shard.model_card.components is not None:
+ shardable_component = next(
+ (c for c in shard.model_card.components if c.can_shard), None
)
- shard_specific_patterns.update(layer_independent_files)
- logger.debug(f"get_allow_patterns {shard=} {layer_independent_files=}")
+
+ if weight_map and shardable_component:
+ for tensor_name, filename in weight_map.items():
+ # Strip component prefix from tensor name (added by weight map namespacing)
+ # E.g., "transformer/blocks.0.weight" -> "blocks.0.weight"
+ if "/" in tensor_name:
+ _, tensor_name_no_prefix = tensor_name.split("/", 1)
+ else:
+ tensor_name_no_prefix = tensor_name
+
+ # Determine which component this file belongs to from filename
+ component_path = Path(filename).parts[0] if "/" in filename else None
+
+ if component_path == shardable_component.component_path.rstrip("/"):
+ layer_num = extract_layer_num(tensor_name_no_prefix)
+ if (
+ layer_num is not None
+ and shard.start_layer <= layer_num < shard.end_layer
+ ):
+ shard_specific_patterns.add(filename)
+
+ if shard.is_first_layer or shard.is_last_layer:
+ shard_specific_patterns.add(filename)
+ else:
+ shard_specific_patterns.add(filename)
+
+ else:
+ shard_specific_patterns = set(["*.safetensors"])
+
+ # TODO(ciaran): temporary - Include all files from non-shardable components that have no index file
+ for component in shard.model_card.components:
+ if not component.can_shard and component.safetensors_index_filename is None:
+ component_pattern = f"{component.component_path.rstrip('/')}/*"
+ shard_specific_patterns.add(component_pattern)
else:
- shard_specific_patterns = set(["*.safetensors"])
+ if weight_map:
+ for tensor_name, filename in weight_map.items():
+ layer_num = extract_layer_num(tensor_name)
+ if (
+ layer_num is not None
+ and shard.start_layer <= layer_num < shard.end_layer
+ ):
+ shard_specific_patterns.add(filename)
+ layer_independent_files = set(
+ [v for k, v in weight_map.items() if extract_layer_num(k) is None]
+ )
+ shard_specific_patterns.update(layer_independent_files)
+ logger.debug(f"get_allow_patterns {shard=} {layer_independent_files=}")
+ else:
+ shard_specific_patterns = set(["*.safetensors"])
+
logger.info(f"get_allow_patterns {shard=} {shard_specific_patterns=}")
return list(default_patterns | shard_specific_patterns)
diff --git a/src/exo/worker/download/shard_downloader.py b/src/exo/worker/download/shard_downloader.py
index d964196a..deb0aaee 100644
--- a/src/exo/worker/download/shard_downloader.py
+++ b/src/exo/worker/download/shard_downloader.py
@@ -5,7 +5,7 @@ from datetime import timedelta
from pathlib import Path
from typing import AsyncIterator, Callable
-from exo.shared.models.model_cards import ModelCard, ModelId
+from exo.shared.models.model_cards import ModelCard, ModelId, ModelTask
from exo.shared.types.memory import Memory
from exo.shared.types.worker.shards import (
PipelineShardMetadata,
@@ -92,6 +92,7 @@ NOOP_DOWNLOAD_PROGRESS = RepoDownloadProgress(
n_layers=1,
hidden_size=1,
supports_tensor=False,
+ tasks=[ModelTask.TextGeneration],
),
device_rank=0,
world_size=1,
diff --git a/src/exo/worker/engines/image/__init__.py b/src/exo/worker/engines/image/__init__.py
new file mode 100644
index 00000000..c83b4702
--- /dev/null
+++ b/src/exo/worker/engines/image/__init__.py
@@ -0,0 +1,12 @@
+from exo.worker.engines.image.distributed_model import (
+ DistributedImageModel,
+ initialize_image_model,
+)
+from exo.worker.engines.image.generate import generate_image, warmup_image_generator
+
+__all__ = [
+ "DistributedImageModel",
+ "generate_image",
+ "initialize_image_model",
+ "warmup_image_generator",
+]
diff --git a/src/exo/worker/engines/image/config.py b/src/exo/worker/engines/image/config.py
new file mode 100644
index 00000000..69c46ac8
--- /dev/null
+++ b/src/exo/worker/engines/image/config.py
@@ -0,0 +1,50 @@
+from enum import Enum
+from math import ceil
+
+from pydantic import BaseModel
+
+
+class BlockType(Enum):
+ JOINT = "joint" # Separate image/text streams
+ SINGLE = "single" # Concatenated streams
+
+
+class TransformerBlockConfig(BaseModel):
+ model_config = {"frozen": True}
+
+ block_type: BlockType
+ count: int
+ has_separate_text_output: bool # True for joint blocks that output text separately
+
+
+class ImageModelConfig(BaseModel):
+ model_family: str
+
+ block_configs: tuple[TransformerBlockConfig, ...]
+
+ default_steps: dict[str, int] # {"low": X, "medium": Y, "high": Z}
+ num_sync_steps_factor: float # Fraction of steps for sync phase
+
+ guidance_scale: float | None = None # None or <= 1.0 disables CFG
+
+ @property
+ def total_blocks(self) -> int:
+ return sum(bc.count for bc in self.block_configs)
+
+ @property
+ def joint_block_count(self) -> int:
+ return sum(
+ bc.count for bc in self.block_configs if bc.block_type == BlockType.JOINT
+ )
+
+ @property
+ def single_block_count(self) -> int:
+ return sum(
+ bc.count for bc in self.block_configs if bc.block_type == BlockType.SINGLE
+ )
+
+ def get_steps_for_quality(self, quality: str) -> int:
+ return self.default_steps[quality]
+
+ def get_num_sync_steps(self, steps: int) -> int:
+ return ceil(steps * self.num_sync_steps_factor)
diff --git a/src/exo/worker/engines/image/distributed_model.py b/src/exo/worker/engines/image/distributed_model.py
new file mode 100644
index 00000000..9c127f36
--- /dev/null
+++ b/src/exo/worker/engines/image/distributed_model.py
@@ -0,0 +1,166 @@
+from collections.abc import Generator
+from pathlib import Path
+from typing import Any, Literal, Optional
+
+import mlx.core as mx
+from mflux.models.common.config.config import Config
+from PIL import Image
+
+from exo.shared.types.api import AdvancedImageParams
+from exo.shared.types.worker.instances import BoundInstance
+from exo.shared.types.worker.shards import PipelineShardMetadata
+from exo.worker.download.download_utils import build_model_path
+from exo.worker.engines.image.config import ImageModelConfig
+from exo.worker.engines.image.models import (
+ create_adapter_for_model,
+ get_config_for_model,
+)
+from exo.worker.engines.image.models.base import ModelAdapter
+from exo.worker.engines.image.pipeline import DiffusionRunner
+from exo.worker.engines.mlx.utils_mlx import mlx_distributed_init, mx_barrier
+from exo.worker.runner.bootstrap import logger
+
+
+class DistributedImageModel:
+ _config: ImageModelConfig
+ _adapter: ModelAdapter[Any, Any]
+ _runner: DiffusionRunner
+
+ def __init__(
+ self,
+ model_id: str,
+ local_path: Path,
+ shard_metadata: PipelineShardMetadata,
+ group: Optional[mx.distributed.Group] = None,
+ quantize: int | None = None,
+ ):
+ config = get_config_for_model(model_id)
+ adapter = create_adapter_for_model(config, model_id, local_path, quantize)
+
+ if group is not None:
+ adapter.slice_transformer_blocks(
+ start_layer=shard_metadata.start_layer,
+ end_layer=shard_metadata.end_layer,
+ )
+
+ runner = DiffusionRunner(
+ config=config,
+ adapter=adapter,
+ group=group,
+ shard_metadata=shard_metadata,
+ )
+
+ if group is not None:
+ logger.info("Initialized distributed diffusion runner")
+
+ mx.eval(adapter.model.parameters()) # pyright: ignore[reportAny]
+
+ # TODO(ciaran): Do we need this?
+ mx.eval(adapter.model) # pyright: ignore[reportAny]
+
+ mx_barrier(group)
+ logger.info(f"Transformer sharded for rank {group.rank()}")
+ else:
+ logger.info("Single-node initialization")
+
+ self._config = config
+ self._adapter = adapter
+ self._runner = runner
+
+ @classmethod
+ def from_bound_instance(
+ cls, bound_instance: BoundInstance
+ ) -> "DistributedImageModel":
+ model_id = bound_instance.bound_shard.model_card.model_id
+ model_path = build_model_path(model_id)
+
+ shard_metadata = bound_instance.bound_shard
+ if not isinstance(shard_metadata, PipelineShardMetadata):
+ raise ValueError("Expected PipelineShardMetadata for image generation")
+
+ is_distributed = (
+ len(bound_instance.instance.shard_assignments.node_to_runner) > 1
+ )
+
+ if is_distributed:
+ logger.info("Starting distributed init for image model")
+ group = mlx_distributed_init(bound_instance)
+ else:
+ group = None
+
+ return cls(
+ model_id=model_id,
+ local_path=model_path,
+ shard_metadata=shard_metadata,
+ group=group,
+ )
+
+ def get_steps_for_quality(self, quality: Literal["low", "medium", "high"]) -> int:
+ """Get the number of inference steps for a quality level."""
+ return self._config.get_steps_for_quality(quality)
+
+ def generate(
+ self,
+ prompt: str,
+ height: int,
+ width: int,
+ quality: Literal["low", "medium", "high"] = "medium",
+ seed: int = 2,
+ image_path: Path | None = None,
+ partial_images: int = 0,
+ advanced_params: AdvancedImageParams | None = None,
+ ) -> Generator[Image.Image | tuple[Image.Image, int, int], None, None]:
+ if (
+ advanced_params is not None
+ and advanced_params.num_inference_steps is not None
+ ):
+ steps = advanced_params.num_inference_steps
+ else:
+ steps = self._config.get_steps_for_quality(quality)
+
+ guidance_override: float | None = None
+ if advanced_params is not None and advanced_params.guidance is not None:
+ guidance_override = advanced_params.guidance
+
+ negative_prompt: str | None = None
+ if advanced_params is not None and advanced_params.negative_prompt is not None:
+ negative_prompt = advanced_params.negative_prompt
+
+ # For edit mode: compute dimensions from input image
+ # This also stores image_paths in the adapter for encode_prompt()
+ if image_path is not None:
+ computed_dims = self._adapter.set_image_dimensions(image_path)
+ if computed_dims is not None:
+ # Override user-provided dimensions with computed ones
+ width, height = computed_dims
+
+ config = Config(
+ num_inference_steps=steps,
+ height=height,
+ width=width,
+ image_path=image_path,
+ model_config=self._adapter.model.model_config, # pyright: ignore[reportAny]
+ )
+
+ num_sync_steps = self._config.get_num_sync_steps(steps)
+
+ for result in self._runner.generate_image(
+ runtime_config=config,
+ prompt=prompt,
+ seed=seed,
+ partial_images=partial_images,
+ guidance_override=guidance_override,
+ negative_prompt=negative_prompt,
+ num_sync_steps=num_sync_steps,
+ ):
+ if isinstance(result, tuple):
+ # Partial image: (GeneratedImage, partial_index, total_partials)
+ image, partial_idx, total_partials = result
+ yield (image, partial_idx, total_partials)
+ else:
+ logger.info("generated image")
+ yield result
+
+
+def initialize_image_model(bound_instance: BoundInstance) -> DistributedImageModel:
+ return DistributedImageModel.from_bound_instance(bound_instance)
diff --git a/src/exo/worker/engines/image/generate.py b/src/exo/worker/engines/image/generate.py
new file mode 100644
index 00000000..8bd34749
--- /dev/null
+++ b/src/exo/worker/engines/image/generate.py
@@ -0,0 +1,176 @@
+import base64
+import io
+import random
+import tempfile
+import time
+from pathlib import Path
+from typing import Generator, Literal
+
+import mlx.core as mx
+from PIL import Image
+
+from exo.shared.types.api import (
+ AdvancedImageParams,
+ ImageEditsInternalParams,
+ ImageGenerationStats,
+ ImageGenerationTaskParams,
+)
+from exo.shared.types.memory import Memory
+from exo.shared.types.worker.runner_response import (
+ ImageGenerationResponse,
+ PartialImageResponse,
+)
+from exo.worker.engines.image.distributed_model import DistributedImageModel
+
+
+def parse_size(size_str: str | None) -> tuple[int, int]:
+ """Parse size parameter like '1024x1024' to (width, height) tuple."""
+ if not size_str:
+ return (1024, 1024)
+
+ try:
+ parts = size_str.split("x")
+ if len(parts) == 2:
+ width, height = int(parts[0]), int(parts[1])
+ if width > 0 and height > 0:
+ return (width, height)
+ except (ValueError, AttributeError):
+ pass
+
+ raise ValueError(
+ f"Invalid size format: '{size_str}'. Expected 'WIDTHxHEIGHT' (e.g., '1024x1024')"
+ )
+
+
+def warmup_image_generator(model: DistributedImageModel) -> Image.Image | None:
+ """Warmup the image generator with a small image."""
+ with tempfile.TemporaryDirectory() as tmpdir:
+ # Create a small dummy image for warmup (needed for edit models)
+ dummy_image = Image.new("RGB", (256, 256), color=(128, 128, 128))
+ dummy_path = Path(tmpdir) / "warmup.png"
+ dummy_image.save(dummy_path)
+
+ warmup_params = AdvancedImageParams(num_inference_steps=2)
+
+ for result in model.generate(
+ prompt="Warmup",
+ height=256,
+ width=256,
+ quality="low",
+ image_path=dummy_path,
+ advanced_params=warmup_params,
+ ):
+ if not isinstance(result, tuple):
+ return result
+ return None
+
+
+def generate_image(
+ model: DistributedImageModel,
+ task: ImageGenerationTaskParams | ImageEditsInternalParams,
+) -> Generator[ImageGenerationResponse | PartialImageResponse, None, None]:
+ """Generate image(s), optionally yielding partial results.
+
+ When partial_images > 0 or stream=True, yields PartialImageResponse for
+ intermediate images, then ImageGenerationResponse for the final image.
+
+ Yields:
+ PartialImageResponse for intermediate images (if partial_images > 0)
+ ImageGenerationResponse for the final complete image
+ """
+ width, height = parse_size(task.size)
+ quality: Literal["low", "medium", "high"] = task.quality or "medium"
+
+ advanced_params = task.advanced_params
+ if advanced_params is not None and advanced_params.seed is not None:
+ seed = advanced_params.seed
+ else:
+ seed = random.randint(0, 2**32 - 1)
+
+ is_bench = getattr(task, "bench", False)
+
+ generation_start_time: float = 0.0
+
+ if is_bench:
+ mx.reset_peak_memory()
+ generation_start_time = time.perf_counter()
+
+ partial_images = task.partial_images or (3 if task.stream else 0)
+
+ image_path: Path | None = None
+
+ with tempfile.TemporaryDirectory() as tmpdir:
+ if isinstance(task, ImageEditsInternalParams):
+ # Decode base64 image data and save to temp file
+ image_path = Path(tmpdir) / "input.png"
+ image_path.write_bytes(base64.b64decode(task.image_data))
+
+ # Iterate over generator results
+ for result in model.generate(
+ prompt=task.prompt,
+ height=height,
+ width=width,
+ quality=quality,
+ seed=seed,
+ image_path=image_path,
+ partial_images=partial_images,
+ advanced_params=advanced_params,
+ ):
+ if isinstance(result, tuple):
+ # Partial image: (Image, partial_index, total_partials)
+ image, partial_idx, total_partials = result
+ buffer = io.BytesIO()
+ image_format = task.output_format.upper()
+ if image_format == "JPG":
+ image_format = "JPEG"
+ if image_format == "JPEG" and image.mode == "RGBA":
+ image = image.convert("RGB")
+ image.save(buffer, format=image_format)
+
+ yield PartialImageResponse(
+ image_data=buffer.getvalue(),
+ format=task.output_format,
+ partial_index=partial_idx,
+ total_partials=total_partials,
+ )
+ else:
+ image = result
+
+ stats: ImageGenerationStats | None = None
+ if is_bench:
+ generation_end_time = time.perf_counter()
+ total_generation_time = generation_end_time - generation_start_time
+
+ num_inference_steps = model.get_steps_for_quality(quality)
+
+ seconds_per_step = (
+ total_generation_time / num_inference_steps
+ if num_inference_steps > 0
+ else 0.0
+ )
+
+ peak_memory_gb = mx.get_peak_memory() / (1024**3)
+
+ stats = ImageGenerationStats(
+ seconds_per_step=seconds_per_step,
+ total_generation_time=total_generation_time,
+ num_inference_steps=num_inference_steps,
+ num_images=task.n or 1,
+ image_width=width,
+ image_height=height,
+ peak_memory_usage=Memory.from_gb(peak_memory_gb),
+ )
+
+ buffer = io.BytesIO()
+ image_format = task.output_format.upper()
+ if image_format == "JPG":
+ image_format = "JPEG"
+ if image_format == "JPEG" and image.mode == "RGBA":
+ image = image.convert("RGB")
+ image.save(buffer, format=image_format)
+
+ yield ImageGenerationResponse(
+ image_data=buffer.getvalue(),
+ format=task.output_format,
+ stats=stats,
+ )
diff --git a/src/exo/worker/engines/image/models/__init__.py b/src/exo/worker/engines/image/models/__init__.py
new file mode 100644
index 00000000..b205af60
--- /dev/null
+++ b/src/exo/worker/engines/image/models/__init__.py
@@ -0,0 +1,86 @@
+from pathlib import Path
+from typing import Any, Callable
+
+from exo.worker.engines.image.config import ImageModelConfig
+from exo.worker.engines.image.models.base import ModelAdapter
+from exo.worker.engines.image.models.flux import (
+ FLUX_DEV_CONFIG,
+ FLUX_SCHNELL_CONFIG,
+ FluxModelAdapter,
+)
+from exo.worker.engines.image.models.qwen import (
+ QWEN_IMAGE_CONFIG,
+ QWEN_IMAGE_EDIT_CONFIG,
+ QwenEditModelAdapter,
+ QwenModelAdapter,
+)
+
+__all__: list[str] = []
+
+# Type alias for adapter factory functions
+# Factory takes (config, model_id, local_path, quantize) and returns a ModelAdapter
+AdapterFactory = Callable[
+ [ImageModelConfig, str, Path, int | None], ModelAdapter[Any, Any]
+]
+
+# Registry maps model_family string to adapter factory
+_ADAPTER_REGISTRY: dict[str, AdapterFactory] = {
+ "flux": FluxModelAdapter,
+ "qwen-edit": QwenEditModelAdapter,
+ "qwen": QwenModelAdapter,
+}
+
+# Config registry: maps model ID patterns to configs
+_CONFIG_REGISTRY: dict[str, ImageModelConfig] = {
+ "flux.1-schnell": FLUX_SCHNELL_CONFIG,
+ "flux.1-dev": FLUX_DEV_CONFIG,
+ "qwen-image-edit": QWEN_IMAGE_EDIT_CONFIG, # Must come before "qwen-image" for pattern matching
+ "qwen-image": QWEN_IMAGE_CONFIG,
+}
+
+
+def get_config_for_model(model_id: str) -> ImageModelConfig:
+ """Get configuration for a model ID.
+
+ Args:
+ model_id: The model identifier (e.g., "black-forest-labs/FLUX.1-schnell")
+
+ Returns:
+ The model configuration
+
+ Raises:
+ ValueError: If no configuration found for model ID
+ """
+ model_id_lower = model_id.lower()
+
+ for pattern, config in _CONFIG_REGISTRY.items():
+ if pattern in model_id_lower:
+ return config
+
+ raise ValueError(f"No configuration found for model: {model_id}")
+
+
+def create_adapter_for_model(
+ config: ImageModelConfig,
+ model_id: str,
+ local_path: Path,
+ quantize: int | None = None,
+) -> ModelAdapter[Any, Any]:
+ """Create a model adapter for the given configuration.
+
+ Args:
+ config: The model configuration
+ model_id: The model identifier
+ local_path: Path to the model weights
+ quantize: Optional quantization bits
+
+ Returns:
+ A ModelAdapter instance
+
+ Raises:
+ ValueError: If no adapter found for model family
+ """
+ factory = _ADAPTER_REGISTRY.get(config.model_family)
+ if factory is None:
+ raise ValueError(f"No adapter found for model family: {config.model_family}")
+ return factory(config, model_id, local_path, quantize)
diff --git a/src/exo/worker/engines/image/models/base.py b/src/exo/worker/engines/image/models/base.py
new file mode 100644
index 00000000..90439823
--- /dev/null
+++ b/src/exo/worker/engines/image/models/base.py
@@ -0,0 +1,295 @@
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Generic, TypeVar
+
+import mlx.core as mx
+from mflux.models.common.config.config import Config
+from mflux.models.common.latent_creator.latent_creator import Img2Img, LatentCreator
+from mflux.utils.image_util import ImageUtil
+from PIL import Image
+
+from exo.worker.engines.image.config import ImageModelConfig
+
+if TYPE_CHECKING:
+ from exo.worker.engines.image.pipeline.block_wrapper import (
+ JointBlockWrapper,
+ SingleBlockWrapper,
+ )
+
+ModelT = TypeVar("ModelT")
+TransformerT = TypeVar("TransformerT")
+
+RotaryEmbeddings = mx.array | tuple[mx.array, mx.array]
+
+
+class PromptData(ABC):
+ @property
+ @abstractmethod
+ def prompt_embeds(self) -> mx.array: ...
+
+ @property
+ @abstractmethod
+ def pooled_prompt_embeds(self) -> mx.array: ...
+
+ @property
+ @abstractmethod
+ def negative_prompt_embeds(self) -> mx.array | None: ...
+
+ @property
+ @abstractmethod
+ def negative_pooled_prompt_embeds(self) -> mx.array | None: ...
+
+ @abstractmethod
+ def get_encoder_hidden_states_mask(
+ self, positive: bool = True
+ ) -> mx.array | None: ...
+
+ @property
+ @abstractmethod
+ def cond_image_grid(
+ self,
+ ) -> tuple[int, int, int] | list[tuple[int, int, int]] | None:
+ """Conditioning image grid dimensions for edit mode.
+
+ Returns:
+ Grid dimensions (edit) or None (standard generation).
+ """
+ ...
+
+ @property
+ @abstractmethod
+ def conditioning_latents(self) -> mx.array | None:
+ """Conditioning latents for edit mode.
+
+ Returns:
+ Conditioning latents array for image editing, None for standard generation.
+ """
+ ...
+
+ @abstractmethod
+ def get_batched_cfg_data(
+ self,
+ ) -> tuple[mx.array, mx.array, mx.array | None, mx.array | None] | None:
+ """Get embeddings for CFG with batch_size=2.
+
+ Combines positive and negative embeddings into batched tensors for
+ a single forward pass. Pads shorter sequences to max length. Attention
+ mask is used to mask padding.
+
+ Returns:
+ None if model doesn't support CFG, otherwise tuple of:
+ - batched_embeds: [2, max_seq, hidden] (positive then negative)
+ - batched_mask: [2, max_seq] attention mask
+ - batched_pooled: [2, hidden] pooled embeddings or None
+ - conditioning_latents: [2, latent_seq, latent_dim] or None
+ TODO(ciaran): type this
+ """
+ ...
+
+
+class ModelAdapter(ABC, Generic[ModelT, TransformerT]):
+ _config: ImageModelConfig
+ _model: ModelT
+ _transformer: TransformerT
+
+ @property
+ def config(self) -> ImageModelConfig:
+ return self._config
+
+ @property
+ def model(self) -> ModelT:
+ return self._model
+
+ @property
+ def transformer(self) -> TransformerT:
+ return self._transformer
+
+ @property
+ @abstractmethod
+ def hidden_dim(self) -> int: ...
+
+ @property
+ @abstractmethod
+ def needs_cfg(self) -> bool:
+ """Whether this model uses classifier-free guidance."""
+ ...
+
+ @abstractmethod
+ def _get_latent_creator(self) -> type: ...
+
+ @abstractmethod
+ def get_joint_block_wrappers(
+ self,
+ text_seq_len: int,
+ encoder_hidden_states_mask: mx.array | None = None,
+ ) -> list["JointBlockWrapper[Any]"]:
+ """Create wrapped joint transformer blocks with pipefusion support.
+
+ Args:
+ text_seq_len: Number of text tokens (constant for generation)
+ encoder_hidden_states_mask: Attention mask for text (Qwen only)
+
+ Returns:
+ List of wrapped joint blocks ready for pipefusion
+ """
+ ...
+
+ @abstractmethod
+ def get_single_block_wrappers(
+ self,
+ text_seq_len: int,
+ ) -> list["SingleBlockWrapper[Any]"]:
+ """Create wrapped single transformer blocks with pipefusion support.
+
+ Args:
+ text_seq_len: Number of text tokens (constant for generation)
+
+ Returns:
+ List of wrapped single blocks ready for pipefusion
+ """
+ ...
+
+ @abstractmethod
+ def slice_transformer_blocks(
+ self,
+ start_layer: int,
+ end_layer: int,
+ ):
+ """Remove transformer blocks outside the assigned range.
+
+ This should be called BEFORE mx.eval() to avoid loading unused weights
+ in distributed mode.
+
+ Args:
+ start_layer: First layer index (inclusive) assigned to this node
+ end_layer: Last layer index (exclusive) assigned to this node
+ """
+ ...
+
+ def set_image_dimensions(self, image_path: Path) -> tuple[int, int] | None:
+ """Default implementation: no dimension computation needed.
+
+ Override in edit adapters to compute dimensions from input image.
+ TODO(ciaran): this is a hack
+
+ Returns:
+ None (use user-specified dimensions)
+ """
+ return None
+
+ def create_latents(self, seed: int, runtime_config: Config) -> mx.array:
+ """Create initial latents. Uses model-specific latent creator."""
+ model: Any = self.model
+ return LatentCreator.create_for_txt2img_or_img2img(
+ seed=seed,
+ height=runtime_config.height,
+ width=runtime_config.width,
+ img2img=Img2Img(
+ vae=model.vae, # pyright: ignore[reportAny]
+ latent_creator=self._get_latent_creator(),
+ sigmas=runtime_config.scheduler.sigmas, # pyright: ignore[reportAny]
+ init_time_step=runtime_config.init_time_step,
+ image_path=runtime_config.image_path,
+ ),
+ )
+
+ def decode_latents(
+ self,
+ latents: mx.array,
+ runtime_config: Config,
+ seed: int,
+ prompt: str,
+ ) -> Image.Image:
+ model: Any = self.model # Allow attribute access on model
+ latents = self._get_latent_creator().unpack_latents( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
+ latents=latents,
+ height=runtime_config.height,
+ width=runtime_config.width,
+ )
+ decoded = model.vae.decode(latents) # pyright: ignore[reportAny]
+ # TODO(ciaran):
+ # from mflux.models.common.vae.vae_util import VAEUtil
+ # VAEUtil.decode(vae=model.vae, latents=latents, tiling_config=self.tiling_config)
+ generated_image = ImageUtil.to_image(
+ decoded_latents=decoded, # pyright: ignore[reportAny]
+ config=runtime_config,
+ seed=seed,
+ prompt=prompt,
+ quantization=model.bits, # pyright: ignore[reportAny]
+ lora_paths=model.lora_paths, # pyright: ignore[reportAny]
+ lora_scales=model.lora_scales, # pyright: ignore[reportAny]
+ image_path=runtime_config.image_path,
+ image_strength=runtime_config.image_strength,
+ generation_time=0,
+ )
+ return generated_image.image
+
+ @abstractmethod
+ def encode_prompt(
+ self, prompt: str, negative_prompt: str | None = None
+ ) -> "PromptData": ...
+
+ @abstractmethod
+ def compute_embeddings(
+ self,
+ hidden_states: mx.array,
+ prompt_embeds: mx.array,
+ ) -> tuple[mx.array, mx.array]: ...
+
+ @abstractmethod
+ def compute_text_embeddings(
+ self,
+ t: int,
+ runtime_config: Config,
+ pooled_prompt_embeds: mx.array | None = None,
+ hidden_states: mx.array | None = None,
+ ) -> mx.array: ...
+
+ @abstractmethod
+ def compute_rotary_embeddings(
+ self,
+ prompt_embeds: mx.array,
+ runtime_config: Config,
+ encoder_hidden_states_mask: mx.array | None = None,
+ cond_image_grid: tuple[int, int, int]
+ | list[tuple[int, int, int]]
+ | None = None,
+ kontext_image_ids: mx.array | None = None,
+ ) -> RotaryEmbeddings: ...
+
+ def merge_streams(
+ self,
+ hidden_states: mx.array,
+ encoder_hidden_states: mx.array,
+ ) -> mx.array:
+ return mx.concatenate([encoder_hidden_states, hidden_states], axis=1)
+
+ @abstractmethod
+ def apply_guidance(
+ self,
+ noise_positive: mx.array,
+ noise_negative: mx.array,
+ guidance_scale: float,
+ ) -> mx.array:
+ """Apply classifier-free guidance to combine positive/negative predictions.
+
+ Only called when needs_cfg is True.
+
+ Args:
+ noise_positive: Noise prediction from positive prompt
+ noise_negative: Noise prediction from negative prompt
+ guidance_scale: Guidance strength
+
+ Returns:
+ Guided noise prediction
+ """
+ ...
+
+ def final_projection(
+ self,
+ hidden_states: mx.array,
+ text_embeddings: mx.array,
+ ) -> mx.array:
+ transformer: Any = self.transformer
+ hidden_states = transformer.norm_out(hidden_states, text_embeddings) # pyright: ignore[reportAny]
+ return transformer.proj_out(hidden_states) # pyright: ignore[reportAny]
diff --git a/src/exo/worker/engines/image/models/flux/__init__.py b/src/exo/worker/engines/image/models/flux/__init__.py
new file mode 100644
index 00000000..3adc2626
--- /dev/null
+++ b/src/exo/worker/engines/image/models/flux/__init__.py
@@ -0,0 +1,11 @@
+from exo.worker.engines.image.models.flux.adapter import FluxModelAdapter
+from exo.worker.engines.image.models.flux.config import (
+ FLUX_DEV_CONFIG,
+ FLUX_SCHNELL_CONFIG,
+)
+
+__all__ = [
+ "FluxModelAdapter",
+ "FLUX_DEV_CONFIG",
+ "FLUX_SCHNELL_CONFIG",
+]
diff --git a/src/exo/worker/engines/image/models/flux/adapter.py b/src/exo/worker/engines/image/models/flux/adapter.py
new file mode 100644
index 00000000..be9b43f8
--- /dev/null
+++ b/src/exo/worker/engines/image/models/flux/adapter.py
@@ -0,0 +1,215 @@
+from pathlib import Path
+from typing import Any
+
+import mlx.core as mx
+from mflux.models.common.config.config import Config
+from mflux.models.common.config.model_config import ModelConfig
+from mflux.models.flux.latent_creator.flux_latent_creator import FluxLatentCreator
+from mflux.models.flux.model.flux_text_encoder.prompt_encoder import PromptEncoder
+from mflux.models.flux.model.flux_transformer.transformer import Transformer
+from mflux.models.flux.variants.txt2img.flux import Flux1
+
+from exo.worker.engines.image.config import ImageModelConfig
+from exo.worker.engines.image.models.base import (
+ ModelAdapter,
+ PromptData,
+ RotaryEmbeddings,
+)
+from exo.worker.engines.image.models.flux.wrappers import (
+ FluxJointBlockWrapper,
+ FluxSingleBlockWrapper,
+)
+from exo.worker.engines.image.pipeline.block_wrapper import (
+ JointBlockWrapper,
+ SingleBlockWrapper,
+)
+
+
+class FluxPromptData(PromptData):
+ def __init__(self, prompt_embeds: mx.array, pooled_prompt_embeds: mx.array):
+ self._prompt_embeds = prompt_embeds
+ self._pooled_prompt_embeds = pooled_prompt_embeds
+
+ @property
+ def prompt_embeds(self) -> mx.array:
+ return self._prompt_embeds
+
+ @property
+ def pooled_prompt_embeds(self) -> mx.array:
+ return self._pooled_prompt_embeds
+
+ @property
+ def negative_prompt_embeds(self) -> mx.array | None:
+ return None
+
+ @property
+ def negative_pooled_prompt_embeds(self) -> mx.array | None:
+ return None
+
+ def get_encoder_hidden_states_mask(self, positive: bool = True) -> mx.array | None:
+ return None
+
+ @property
+ def cond_image_grid(
+ self,
+ ) -> tuple[int, int, int] | list[tuple[int, int, int]] | None:
+ return None
+
+ @property
+ def conditioning_latents(self) -> mx.array | None:
+ return None
+
+ def get_batched_cfg_data(
+ self,
+ ) -> tuple[mx.array, mx.array, mx.array | None, mx.array | None] | None:
+ return None
+
+
+class FluxModelAdapter(ModelAdapter[Flux1, Transformer]):
+ def __init__(
+ self,
+ config: ImageModelConfig,
+ model_id: str,
+ local_path: Path,
+ quantize: int | None = None,
+ ):
+ self._config = config
+ self._model = Flux1(
+ model_config=ModelConfig.from_name(model_name=model_id, base_model=None),
+ model_path=str(local_path),
+ quantize=quantize,
+ )
+ self._transformer = self._model.transformer
+
+ @property
+ def hidden_dim(self) -> int:
+ return self._transformer.x_embedder.weight.shape[0] # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
+
+ @property
+ def needs_cfg(self) -> bool:
+ return False
+
+ def _get_latent_creator(self) -> type:
+ return FluxLatentCreator
+
+ def get_joint_block_wrappers(
+ self,
+ text_seq_len: int,
+ encoder_hidden_states_mask: mx.array | None = None,
+ ) -> list[JointBlockWrapper[Any]]:
+ """Create wrapped joint blocks for Flux."""
+ return [
+ FluxJointBlockWrapper(block, text_seq_len)
+ for block in self._transformer.transformer_blocks
+ ]
+
+ def get_single_block_wrappers(
+ self,
+ text_seq_len: int,
+ ) -> list[SingleBlockWrapper[Any]]:
+ """Create wrapped single blocks for Flux."""
+ return [
+ FluxSingleBlockWrapper(block, text_seq_len)
+ for block in self._transformer.single_transformer_blocks
+ ]
+
+ def slice_transformer_blocks(
+ self,
+ start_layer: int,
+ end_layer: int,
+ ):
+ all_joint = list(self._transformer.transformer_blocks)
+ all_single = list(self._transformer.single_transformer_blocks)
+ total_joint_blocks = len(all_joint)
+ if end_layer <= total_joint_blocks:
+ # All assigned are joint blocks
+ joint_start, joint_end = start_layer, end_layer
+ single_start, single_end = 0, 0
+ elif start_layer >= total_joint_blocks:
+ # All assigned are single blocks
+ joint_start, joint_end = 0, 0
+ single_start = start_layer - total_joint_blocks
+ single_end = end_layer - total_joint_blocks
+ else:
+ # Spans both joint and single
+ joint_start, joint_end = start_layer, total_joint_blocks
+ single_start = 0
+ single_end = end_layer - total_joint_blocks
+
+ self._transformer.transformer_blocks = all_joint[joint_start:joint_end]
+
+ self._transformer.single_transformer_blocks = all_single[
+ single_start:single_end
+ ]
+
+ def encode_prompt(
+ self, prompt: str, negative_prompt: str | None = None
+ ) -> FluxPromptData:
+ del negative_prompt
+
+ assert isinstance(self.model.prompt_cache, dict)
+ assert isinstance(self.model.tokenizers, dict)
+
+ prompt_embeds, pooled_prompt_embeds = PromptEncoder.encode_prompt(
+ prompt=prompt,
+ prompt_cache=self.model.prompt_cache,
+ t5_tokenizer=self.model.tokenizers["t5"], # pyright: ignore[reportAny]
+ clip_tokenizer=self.model.tokenizers["clip"], # pyright: ignore[reportAny]
+ t5_text_encoder=self.model.t5_text_encoder,
+ clip_text_encoder=self.model.clip_text_encoder,
+ )
+ return FluxPromptData(
+ prompt_embeds=prompt_embeds,
+ pooled_prompt_embeds=pooled_prompt_embeds,
+ )
+
+ def compute_embeddings(
+ self,
+ hidden_states: mx.array,
+ prompt_embeds: mx.array,
+ ) -> tuple[mx.array, mx.array]:
+ embedded_hidden = self._transformer.x_embedder(hidden_states)
+ embedded_encoder = self._transformer.context_embedder(prompt_embeds)
+ return embedded_hidden, embedded_encoder
+
+ def compute_text_embeddings(
+ self,
+ t: int,
+ runtime_config: Config,
+ pooled_prompt_embeds: mx.array | None = None,
+ hidden_states: mx.array | None = None, # Ignored by Flux
+ ) -> mx.array:
+ if pooled_prompt_embeds is None:
+ raise ValueError(
+ "pooled_prompt_embeds is required for Flux text embeddings"
+ )
+
+ # hidden_states is ignored - Flux uses pooled_prompt_embeds instead
+ return Transformer.compute_text_embeddings(
+ t, pooled_prompt_embeds, self._transformer.time_text_embed, runtime_config
+ )
+
+ def compute_rotary_embeddings(
+ self,
+ prompt_embeds: mx.array,
+ runtime_config: Config,
+ encoder_hidden_states_mask: mx.array | None = None,
+ cond_image_grid: tuple[int, int, int]
+ | list[tuple[int, int, int]]
+ | None = None,
+ kontext_image_ids: mx.array | None = None,
+ ) -> RotaryEmbeddings:
+ return Transformer.compute_rotary_embeddings(
+ prompt_embeds,
+ self._transformer.pos_embed,
+ runtime_config,
+ kontext_image_ids,
+ )
+
+ def apply_guidance(
+ self,
+ noise_positive: mx.array,
+ noise_negative: mx.array,
+ guidance_scale: float,
+ ) -> mx.array:
+ raise NotImplementedError("Flux does not use classifier-free guidance")
diff --git a/src/exo/worker/engines/image/models/flux/config.py b/src/exo/worker/engines/image/models/flux/config.py
new file mode 100644
index 00000000..0cf85d74
--- /dev/null
+++ b/src/exo/worker/engines/image/models/flux/config.py
@@ -0,0 +1,34 @@
+from exo.worker.engines.image.config import (
+ BlockType,
+ ImageModelConfig,
+ TransformerBlockConfig,
+)
+
+FLUX_SCHNELL_CONFIG = ImageModelConfig(
+ model_family="flux",
+ block_configs=(
+ TransformerBlockConfig(
+ block_type=BlockType.JOINT, count=19, has_separate_text_output=True
+ ),
+ TransformerBlockConfig(
+ block_type=BlockType.SINGLE, count=38, has_separate_text_output=False
+ ),
+ ),
+ default_steps={"low": 1, "medium": 2, "high": 4},
+ num_sync_steps_factor=0.5, # 1 sync step for medium (2 steps)
+)
+
+
+FLUX_DEV_CONFIG = ImageModelConfig(
+ model_family="flux",
+ block_configs=(
+ TransformerBlockConfig(
+ block_type=BlockType.JOINT, count=19, has_separate_text_output=True
+ ),
+ TransformerBlockConfig(
+ block_type=BlockType.SINGLE, count=38, has_separate_text_output=False
+ ),
+ ),
+ default_steps={"low": 10, "medium": 25, "high": 50},
+ num_sync_steps_factor=0.125, # ~3 sync steps for medium (25 steps)
+)
diff --git a/src/exo/worker/engines/image/models/flux/wrappers.py b/src/exo/worker/engines/image/models/flux/wrappers.py
new file mode 100644
index 00000000..ef0a6336
--- /dev/null
+++ b/src/exo/worker/engines/image/models/flux/wrappers.py
@@ -0,0 +1,279 @@
+from typing import final
+
+import mlx.core as mx
+from mflux.models.flux.model.flux_transformer.common.attention_utils import (
+ AttentionUtils,
+)
+from mflux.models.flux.model.flux_transformer.joint_transformer_block import (
+ JointTransformerBlock,
+)
+from mflux.models.flux.model.flux_transformer.single_transformer_block import (
+ SingleTransformerBlock,
+)
+from pydantic import BaseModel, ConfigDict
+
+from exo.worker.engines.image.models.base import RotaryEmbeddings
+from exo.worker.engines.image.pipeline.block_wrapper import (
+ JointBlockWrapper,
+ SingleBlockWrapper,
+)
+
+
+@final
+class FluxModulationParams(BaseModel):
+ model_config = ConfigDict(frozen=True, strict=True, arbitrary_types_allowed=True)
+
+ gate_msa: mx.array
+ shift_mlp: mx.array
+ scale_mlp: mx.array
+ gate_mlp: mx.array
+
+
+@final
+class FluxNormGateState(BaseModel):
+ model_config = ConfigDict(frozen=True, strict=True, arbitrary_types_allowed=True)
+
+ norm_hidden: mx.array
+ gate: mx.array
+
+
+class FluxJointBlockWrapper(JointBlockWrapper[JointTransformerBlock]):
+ def __init__(self, block: JointTransformerBlock, text_seq_len: int):
+ super().__init__(block, text_seq_len)
+ self._num_heads = block.attn.num_heads
+ self._head_dim = block.attn.head_dimension
+
+ # Intermediate state stored between _compute_qkv and _apply_output
+ self._hidden_mod: FluxModulationParams | None = None
+ self._context_mod: FluxModulationParams | None = None
+
+ def _compute_qkv(
+ self,
+ hidden_states: mx.array,
+ encoder_hidden_states: mx.array,
+ text_embeddings: mx.array,
+ rotary_embeddings: RotaryEmbeddings,
+ patch_mode: bool = False,
+ ) -> tuple[mx.array, mx.array, mx.array]:
+ assert isinstance(rotary_embeddings, mx.array)
+
+ attn = self.block.attn
+
+ (
+ norm_hidden,
+ gate_msa,
+ shift_mlp,
+ scale_mlp,
+ gate_mlp,
+ ) = self.block.norm1(
+ hidden_states=hidden_states,
+ text_embeddings=text_embeddings,
+ )
+ self._hidden_mod = FluxModulationParams(
+ gate_msa=gate_msa,
+ shift_mlp=shift_mlp,
+ scale_mlp=scale_mlp,
+ gate_mlp=gate_mlp,
+ )
+
+ (
+ norm_encoder,
+ c_gate_msa,
+ c_shift_mlp,
+ c_scale_mlp,
+ c_gate_mlp,
+ ) = self.block.norm1_context(
+ hidden_states=encoder_hidden_states,
+ text_embeddings=text_embeddings,
+ )
+ self._context_mod = FluxModulationParams(
+ gate_msa=c_gate_msa,
+ shift_mlp=c_shift_mlp,
+ scale_mlp=c_scale_mlp,
+ gate_mlp=c_gate_mlp,
+ )
+
+ img_query, img_key, img_value = AttentionUtils.process_qkv(
+ hidden_states=norm_hidden,
+ to_q=attn.to_q,
+ to_k=attn.to_k,
+ to_v=attn.to_v,
+ norm_q=attn.norm_q,
+ norm_k=attn.norm_k,
+ num_heads=self._num_heads,
+ head_dim=self._head_dim,
+ )
+
+ txt_query, txt_key, txt_value = AttentionUtils.process_qkv(
+ hidden_states=norm_encoder,
+ to_q=attn.add_q_proj,
+ to_k=attn.add_k_proj,
+ to_v=attn.add_v_proj,
+ norm_q=attn.norm_added_q,
+ norm_k=attn.norm_added_k,
+ num_heads=self._num_heads,
+ head_dim=self._head_dim,
+ )
+
+ query = mx.concatenate([txt_query, img_query], axis=2)
+ key = mx.concatenate([txt_key, img_key], axis=2)
+ value = mx.concatenate([txt_value, img_value], axis=2)
+
+ if patch_mode:
+ text_rope = rotary_embeddings[:, :, : self._text_seq_len, ...]
+ patch_img_rope = rotary_embeddings[
+ :,
+ :,
+ self._text_seq_len + self._patch_start : self._text_seq_len
+ + self._patch_end,
+ ...,
+ ]
+ rope = mx.concatenate([text_rope, patch_img_rope], axis=2)
+ else:
+ rope = rotary_embeddings
+
+ query, key = AttentionUtils.apply_rope(xq=query, xk=key, freqs_cis=rope)
+
+ return query, key, value
+
+ def _compute_attention(
+ self, query: mx.array, key: mx.array, value: mx.array
+ ) -> mx.array:
+ batch_size = query.shape[0]
+ return AttentionUtils.compute_attention(
+ query=query,
+ key=key,
+ value=value,
+ batch_size=batch_size,
+ num_heads=self._num_heads,
+ head_dim=self._head_dim,
+ )
+
+ def _apply_output(
+ self,
+ attn_out: mx.array,
+ hidden_states: mx.array,
+ encoder_hidden_states: mx.array,
+ text_embeddings: mx.array,
+ ) -> tuple[mx.array, mx.array]:
+ attn = self.block.attn
+
+ context_attn_output = attn_out[:, : self._text_seq_len, :]
+ hidden_attn_output = attn_out[:, self._text_seq_len :, :]
+
+ hidden_attn_output = attn.to_out[0](hidden_attn_output) # pyright: ignore[reportAny]
+ context_attn_output = attn.to_add_out(context_attn_output)
+
+ assert self._hidden_mod is not None
+ assert self._context_mod is not None
+
+ hidden_states = JointTransformerBlock.apply_norm_and_feed_forward(
+ hidden_states=hidden_states,
+ attn_output=hidden_attn_output, # pyright: ignore[reportAny]
+ gate_mlp=self._hidden_mod.gate_mlp,
+ gate_msa=self._hidden_mod.gate_msa,
+ scale_mlp=self._hidden_mod.scale_mlp,
+ shift_mlp=self._hidden_mod.shift_mlp,
+ norm_layer=self.block.norm2,
+ ff_layer=self.block.ff,
+ )
+ encoder_hidden_states = JointTransformerBlock.apply_norm_and_feed_forward(
+ hidden_states=encoder_hidden_states,
+ attn_output=context_attn_output,
+ gate_mlp=self._context_mod.gate_mlp,
+ gate_msa=self._context_mod.gate_msa,
+ scale_mlp=self._context_mod.scale_mlp,
+ shift_mlp=self._context_mod.shift_mlp,
+ norm_layer=self.block.norm2_context,
+ ff_layer=self.block.ff_context,
+ )
+
+ return encoder_hidden_states, hidden_states
+
+
+class FluxSingleBlockWrapper(SingleBlockWrapper[SingleTransformerBlock]):
+ """Flux-specific single block wrapper with pipefusion support."""
+
+ def __init__(self, block: SingleTransformerBlock, text_seq_len: int):
+ super().__init__(block, text_seq_len)
+ self._num_heads = block.attn.num_heads
+ self._head_dim = block.attn.head_dimension
+
+ # Intermediate state stored between _compute_qkv and _apply_output
+ self._norm_state: FluxNormGateState | None = None
+
+ def _compute_qkv(
+ self,
+ hidden_states: mx.array,
+ text_embeddings: mx.array,
+ rotary_embeddings: RotaryEmbeddings,
+ patch_mode: bool = False,
+ ) -> tuple[mx.array, mx.array, mx.array]:
+ assert isinstance(rotary_embeddings, mx.array)
+
+ attn = self.block.attn
+
+ norm_hidden, gate = self.block.norm(
+ hidden_states=hidden_states,
+ text_embeddings=text_embeddings,
+ )
+ self._norm_state = FluxNormGateState(norm_hidden=norm_hidden, gate=gate)
+
+ query, key, value = AttentionUtils.process_qkv(
+ hidden_states=norm_hidden,
+ to_q=attn.to_q,
+ to_k=attn.to_k,
+ to_v=attn.to_v,
+ norm_q=attn.norm_q,
+ norm_k=attn.norm_k,
+ num_heads=self._num_heads,
+ head_dim=self._head_dim,
+ )
+
+ if patch_mode:
+ text_rope = rotary_embeddings[:, :, : self._text_seq_len, ...]
+ patch_img_rope = rotary_embeddings[
+ :,
+ :,
+ self._text_seq_len + self._patch_start : self._text_seq_len
+ + self._patch_end,
+ ...,
+ ]
+ rope = mx.concatenate([text_rope, patch_img_rope], axis=2)
+ else:
+ rope = rotary_embeddings
+
+ query, key = AttentionUtils.apply_rope(xq=query, xk=key, freqs_cis=rope)
+
+ return query, key, value
+
+ def _compute_attention(
+ self, query: mx.array, key: mx.array, value: mx.array
+ ) -> mx.array:
+ batch_size = query.shape[0]
+ return AttentionUtils.compute_attention(
+ query=query,
+ key=key,
+ value=value,
+ batch_size=batch_size,
+ num_heads=self._num_heads,
+ head_dim=self._head_dim,
+ )
+
+ def _apply_output(
+ self,
+ attn_out: mx.array,
+ hidden_states: mx.array,
+ text_embeddings: mx.array,
+ ) -> mx.array:
+ residual = hidden_states
+
+ assert self._norm_state is not None
+
+ output = self.block._apply_feed_forward_and_projection(
+ norm_hidden_states=self._norm_state.norm_hidden,
+ attn_output=attn_out,
+ gate=self._norm_state.gate,
+ )
+
+ return residual + output
diff --git a/src/exo/worker/engines/image/models/qwen/__init__.py b/src/exo/worker/engines/image/models/qwen/__init__.py
new file mode 100644
index 00000000..1eed0ca8
--- /dev/null
+++ b/src/exo/worker/engines/image/models/qwen/__init__.py
@@ -0,0 +1,13 @@
+from exo.worker.engines.image.models.qwen.adapter import QwenModelAdapter
+from exo.worker.engines.image.models.qwen.config import (
+ QWEN_IMAGE_CONFIG,
+ QWEN_IMAGE_EDIT_CONFIG,
+)
+from exo.worker.engines.image.models.qwen.edit_adapter import QwenEditModelAdapter
+
+__all__ = [
+ "QwenModelAdapter",
+ "QwenEditModelAdapter",
+ "QWEN_IMAGE_CONFIG",
+ "QWEN_IMAGE_EDIT_CONFIG",
+]
diff --git a/src/exo/worker/engines/image/models/qwen/adapter.py b/src/exo/worker/engines/image/models/qwen/adapter.py
new file mode 100644
index 00000000..d9f009ec
--- /dev/null
+++ b/src/exo/worker/engines/image/models/qwen/adapter.py
@@ -0,0 +1,292 @@
+from pathlib import Path
+from typing import Any
+
+import mlx.core as mx
+from mflux.models.common.config import ModelConfig
+from mflux.models.common.config.config import Config
+from mflux.models.qwen.latent_creator.qwen_latent_creator import QwenLatentCreator
+from mflux.models.qwen.model.qwen_text_encoder.qwen_prompt_encoder import (
+ QwenPromptEncoder,
+)
+from mflux.models.qwen.model.qwen_transformer.qwen_transformer import QwenTransformer
+from mflux.models.qwen.variants.txt2img.qwen_image import QwenImage
+
+from exo.worker.engines.image.config import ImageModelConfig
+from exo.worker.engines.image.models.base import (
+ ModelAdapter,
+ PromptData,
+ RotaryEmbeddings,
+)
+from exo.worker.engines.image.models.qwen.wrappers import QwenJointBlockWrapper
+from exo.worker.engines.image.pipeline.block_wrapper import (
+ JointBlockWrapper,
+ SingleBlockWrapper,
+)
+
+
+class QwenPromptData(PromptData):
+ def __init__(
+ self,
+ prompt_embeds: mx.array,
+ prompt_mask: mx.array,
+ negative_prompt_embeds: mx.array,
+ negative_prompt_mask: mx.array,
+ ):
+ self._prompt_embeds = prompt_embeds
+ self._prompt_mask = prompt_mask
+ self._negative_prompt_embeds = negative_prompt_embeds
+ self._negative_prompt_mask = negative_prompt_mask
+
+ @property
+ def prompt_embeds(self) -> mx.array:
+ return self._prompt_embeds
+
+ @property
+ def pooled_prompt_embeds(self) -> mx.array:
+ return self._prompt_embeds
+
+ @property
+ def negative_prompt_embeds(self) -> mx.array:
+ return self._negative_prompt_embeds
+
+ @property
+ def negative_pooled_prompt_embeds(self) -> mx.array:
+ return self._negative_prompt_embeds
+
+ def get_encoder_hidden_states_mask(self, positive: bool = True) -> mx.array:
+ if positive:
+ return self._prompt_mask
+ else:
+ return self._negative_prompt_mask
+
+ @property
+ def cond_image_grid(
+ self,
+ ) -> tuple[int, int, int] | list[tuple[int, int, int]] | None:
+ return None
+
+ @property
+ def conditioning_latents(self) -> mx.array | None:
+ return None
+
+ def get_batched_cfg_data(
+ self,
+ ) -> tuple[mx.array, mx.array, mx.array | None, mx.array | None] | None:
+ """Batch positive and negative embeddings for CFG with batch_size=2.
+
+ Pads shorter sequence to max length using zeros for embeddings
+ and zeros (masked) for attention mask.
+
+ Returns:
+ Tuple of (batched_embeds, batched_mask, None, conditioning_latents)
+ - batched_embeds: [2, max_seq, hidden]
+ - batched_mask: [2, max_seq]
+ - None for pooled (Qwen doesn't use it)
+ - conditioning_latents: [2, latent_seq, latent_dim] or None
+ """
+ pos_embeds = self._prompt_embeds
+ neg_embeds = self._negative_prompt_embeds
+ pos_mask = self._prompt_mask
+ neg_mask = self._negative_prompt_mask
+
+ pos_seq_len = pos_embeds.shape[1]
+ neg_seq_len = neg_embeds.shape[1]
+ max_seq_len = max(pos_seq_len, neg_seq_len)
+ hidden_dim = pos_embeds.shape[2]
+
+ if pos_seq_len < max_seq_len:
+ pad_len = max_seq_len - pos_seq_len
+ pos_embeds = mx.concatenate(
+ [
+ pos_embeds,
+ mx.zeros((1, pad_len, hidden_dim), dtype=pos_embeds.dtype),
+ ],
+ axis=1,
+ )
+ pos_mask = mx.concatenate(
+ [pos_mask, mx.zeros((1, pad_len), dtype=pos_mask.dtype)],
+ axis=1,
+ )
+
+ elif neg_seq_len < max_seq_len:
+ pad_len = max_seq_len - neg_seq_len
+ neg_embeds = mx.concatenate(
+ [
+ neg_embeds,
+ mx.zeros((1, pad_len, hidden_dim), dtype=neg_embeds.dtype),
+ ],
+ axis=1,
+ )
+ neg_mask = mx.concatenate(
+ [neg_mask, mx.zeros((1, pad_len), dtype=neg_mask.dtype)],
+ axis=1,
+ )
+
+ batched_embeds = mx.concatenate([pos_embeds, neg_embeds], axis=0)
+ batched_mask = mx.concatenate([pos_mask, neg_mask], axis=0)
+
+ # TODO(ciaran): currently None but maybe we will deduplicate with edit
+ # adapter
+ cond_latents = self.conditioning_latents
+ if cond_latents is not None:
+ cond_latents = mx.concatenate([cond_latents, cond_latents], axis=0)
+
+ return batched_embeds, batched_mask, None, cond_latents
+
+
+class QwenModelAdapter(ModelAdapter[QwenImage, QwenTransformer]):
+ """Adapter for Qwen-Image model.
+
+ Key differences from Flux:
+ - Single text encoder (vs dual T5+CLIP)
+ - 60 joint-style blocks, no single blocks
+ - 3D RoPE returning ((img_cos, img_sin), (txt_cos, txt_sin))
+ - Norm-preserving CFG with negative prompts
+ - Uses attention mask for variable-length text
+ """
+
+ def __init__(
+ self,
+ config: ImageModelConfig,
+ model_id: str,
+ local_path: Path,
+ quantize: int | None = None,
+ ):
+ self._config = config
+ self._model = QwenImage(
+ model_config=ModelConfig.from_name(model_name=model_id, base_model=None),
+ model_path=str(local_path),
+ quantize=quantize,
+ )
+ self._transformer = self._model.transformer
+
+ @property
+ def hidden_dim(self) -> int:
+ return self._transformer.inner_dim
+
+ @property
+ def needs_cfg(self) -> bool:
+ gs = self._config.guidance_scale
+ return gs is not None and gs > 1.0
+
+ def _get_latent_creator(self) -> type:
+ return QwenLatentCreator
+
+ def get_joint_block_wrappers(
+ self,
+ text_seq_len: int,
+ encoder_hidden_states_mask: mx.array | None = None,
+ ) -> list[JointBlockWrapper[Any]]:
+ """Create wrapped joint blocks for Qwen."""
+ return [
+ QwenJointBlockWrapper(block, text_seq_len, encoder_hidden_states_mask)
+ for block in self._transformer.transformer_blocks
+ ]
+
+ def get_single_block_wrappers(
+ self,
+ text_seq_len: int,
+ ) -> list[SingleBlockWrapper[Any]]:
+ return []
+
+ def slice_transformer_blocks(
+ self,
+ start_layer: int,
+ end_layer: int,
+ ):
+ self._transformer.transformer_blocks = self._transformer.transformer_blocks[
+ start_layer:end_layer
+ ]
+
+ def encode_prompt(
+ self, prompt: str, negative_prompt: str | None = None
+ ) -> QwenPromptData:
+ assert isinstance(self.model.prompt_cache, dict)
+ assert isinstance(self.model.tokenizers, dict)
+
+ if negative_prompt is None or negative_prompt == "":
+ negative_prompt = " "
+
+ prompt_embeds, prompt_mask, neg_embeds, neg_mask = (
+ QwenPromptEncoder.encode_prompt(
+ prompt=prompt,
+ negative_prompt=negative_prompt,
+ prompt_cache=self.model.prompt_cache,
+ qwen_tokenizer=self.model.tokenizers["qwen"], # pyright: ignore[reportAny]
+ qwen_text_encoder=self.model.text_encoder,
+ )
+ )
+
+ return QwenPromptData(
+ prompt_embeds=prompt_embeds,
+ prompt_mask=prompt_mask,
+ negative_prompt_embeds=neg_embeds,
+ negative_prompt_mask=neg_mask,
+ )
+
+ def compute_embeddings(
+ self,
+ hidden_states: mx.array,
+ prompt_embeds: mx.array,
+ ) -> tuple[mx.array, mx.array]:
+ embedded_hidden = self._transformer.img_in(hidden_states)
+ encoder_hidden_states = self._transformer.txt_norm(prompt_embeds)
+ embedded_encoder = self._transformer.txt_in(encoder_hidden_states)
+ return embedded_hidden, embedded_encoder
+
+ def compute_text_embeddings(
+ self,
+ t: int,
+ runtime_config: Config,
+ pooled_prompt_embeds: mx.array | None = None,
+ hidden_states: mx.array | None = None,
+ ) -> mx.array:
+ # Use hidden_states if provided, otherwise fall back to pooled_prompt_embeds
+ # (which for Qwen is the same as prompt_embeds)
+ ref_tensor = (
+ hidden_states if hidden_states is not None else pooled_prompt_embeds
+ )
+ if ref_tensor is None:
+ raise ValueError(
+ "Either hidden_states or pooled_prompt_embeds is required "
+ "for Qwen text embeddings"
+ )
+
+ timestep = QwenTransformer._compute_timestep(t, runtime_config) # noqa: SLF001
+ batch_size = ref_tensor.shape[0]
+ timestep = mx.broadcast_to(timestep, (batch_size,)).astype(mx.float32)
+ return self._transformer.time_text_embed(timestep, ref_tensor) # pyright: ignore[reportAny]
+
+ def compute_rotary_embeddings(
+ self,
+ prompt_embeds: mx.array,
+ runtime_config: Config,
+ encoder_hidden_states_mask: mx.array | None = None,
+ cond_image_grid: tuple[int, int, int]
+ | list[tuple[int, int, int]]
+ | None = None,
+ kontext_image_ids: mx.array | None = None,
+ ) -> RotaryEmbeddings:
+ if encoder_hidden_states_mask is None:
+ raise ValueError(
+ "encoder_hidden_states_mask is required for Qwen RoPE computation"
+ )
+
+ return QwenTransformer._compute_rotary_embeddings(
+ encoder_hidden_states_mask=encoder_hidden_states_mask,
+ pos_embed=self._transformer.pos_embed, # pyright: ignore[reportAny]
+ config=runtime_config,
+ cond_image_grid=cond_image_grid,
+ )
+
+ def apply_guidance(
+ self,
+ noise_positive: mx.array,
+ noise_negative: mx.array,
+ guidance_scale: float,
+ ) -> mx.array:
+ return self._model.compute_guided_noise(
+ noise=noise_positive,
+ noise_negative=noise_negative,
+ guidance=guidance_scale,
+ )
diff --git a/src/exo/worker/engines/image/models/qwen/config.py b/src/exo/worker/engines/image/models/qwen/config.py
new file mode 100644
index 00000000..d5da1bac
--- /dev/null
+++ b/src/exo/worker/engines/image/models/qwen/config.py
@@ -0,0 +1,29 @@
+from exo.worker.engines.image.config import (
+ BlockType,
+ ImageModelConfig,
+ TransformerBlockConfig,
+)
+
+QWEN_IMAGE_CONFIG = ImageModelConfig(
+ model_family="qwen",
+ block_configs=(
+ TransformerBlockConfig(
+ block_type=BlockType.JOINT, count=60, has_separate_text_output=True
+ ),
+ ),
+ default_steps={"low": 10, "medium": 25, "high": 50},
+ num_sync_steps_factor=0.125, # ~3 sync steps for medium (30 steps)
+ guidance_scale=3.5, # Set to None or < 1.0 to disable CFG
+)
+
+QWEN_IMAGE_EDIT_CONFIG = ImageModelConfig(
+ model_family="qwen-edit",
+ block_configs=(
+ TransformerBlockConfig(
+ block_type=BlockType.JOINT, count=60, has_separate_text_output=True
+ ),
+ ),
+ default_steps={"low": 10, "medium": 25, "high": 50},
+ num_sync_steps_factor=0.125,
+ guidance_scale=3.5,
+)
diff --git a/src/exo/worker/engines/image/models/qwen/edit_adapter.py b/src/exo/worker/engines/image/models/qwen/edit_adapter.py
new file mode 100644
index 00000000..e327eb0c
--- /dev/null
+++ b/src/exo/worker/engines/image/models/qwen/edit_adapter.py
@@ -0,0 +1,434 @@
+import math
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import mlx.core as mx
+from mflux.models.common.config.config import Config
+from mflux.models.qwen.latent_creator.qwen_latent_creator import QwenLatentCreator
+from mflux.models.qwen.model.qwen_transformer.qwen_transformer import QwenTransformer
+from mflux.models.qwen.variants.edit.qwen_edit_util import QwenEditUtil
+from mflux.models.qwen.variants.edit.qwen_image_edit import QwenImageEdit
+
+from exo.worker.engines.image.config import ImageModelConfig
+from exo.worker.engines.image.models.base import (
+ ModelAdapter,
+ PromptData,
+ RotaryEmbeddings,
+)
+from exo.worker.engines.image.models.qwen.wrappers import QwenJointBlockWrapper
+from exo.worker.engines.image.pipeline.block_wrapper import (
+ JointBlockWrapper,
+ SingleBlockWrapper,
+)
+
+
+@dataclass(frozen=True)
+class EditImageDimensions:
+ vl_width: int
+ vl_height: int
+ vae_width: int
+ vae_height: int
+ image_paths: list[str]
+
+
+class QwenEditPromptData(PromptData):
+ def __init__(
+ self,
+ prompt_embeds: mx.array,
+ prompt_mask: mx.array,
+ negative_prompt_embeds: mx.array,
+ negative_prompt_mask: mx.array,
+ conditioning_latents: mx.array,
+ qwen_image_ids: mx.array,
+ cond_image_grid: tuple[int, int, int] | list[tuple[int, int, int]],
+ ):
+ self._prompt_embeds = prompt_embeds
+ self._prompt_mask = prompt_mask
+ self._negative_prompt_embeds = negative_prompt_embeds
+ self._negative_prompt_mask = negative_prompt_mask
+ self._conditioning_latents = conditioning_latents
+ self._qwen_image_ids = qwen_image_ids
+ self._cond_image_grid = cond_image_grid
+
+ @property
+ def prompt_embeds(self) -> mx.array:
+ return self._prompt_embeds
+
+ @property
+ def pooled_prompt_embeds(self) -> mx.array:
+ return self._prompt_embeds
+
+ @property
+ def negative_prompt_embeds(self) -> mx.array:
+ return self._negative_prompt_embeds
+
+ @property
+ def negative_pooled_prompt_embeds(self) -> mx.array:
+ return self._negative_prompt_embeds
+
+ def get_encoder_hidden_states_mask(self, positive: bool = True) -> mx.array:
+ if positive:
+ return self._prompt_mask
+ else:
+ return self._negative_prompt_mask
+
+ @property
+ def cond_image_grid(self) -> tuple[int, int, int] | list[tuple[int, int, int]]:
+ return self._cond_image_grid
+
+ @property
+ def conditioning_latents(self) -> mx.array:
+ return self._conditioning_latents
+
+ @property
+ def qwen_image_ids(self) -> mx.array:
+ return self._qwen_image_ids
+
+ @property
+ def is_edit_mode(self) -> bool:
+ return True
+
+ def get_batched_cfg_data(
+ self,
+ ) -> tuple[mx.array, mx.array, mx.array | None, mx.array | None] | None:
+ """Batch positive and negative embeddings for CFG with batch_size=2.
+
+ Pads shorter sequence to max length using zeros for embeddings
+ and zeros (masked) for attention mask. Duplicates conditioning
+ latents for both positive and negative passes.
+
+ Returns:
+ Tuple of (batched_embeds, batched_mask, None, batched_cond_latents)
+ - batched_embeds: [2, max_seq, hidden]
+ - batched_mask: [2, max_seq]
+ - None for pooled (Qwen doesn't use it)
+ - batched_cond_latents: [2, latent_seq, latent_dim]
+ TODO(ciaran): type this
+ """
+ pos_embeds = self._prompt_embeds
+ neg_embeds = self._negative_prompt_embeds
+ pos_mask = self._prompt_mask
+ neg_mask = self._negative_prompt_mask
+
+ pos_seq_len = pos_embeds.shape[1]
+ neg_seq_len = neg_embeds.shape[1]
+ max_seq_len = max(pos_seq_len, neg_seq_len)
+ hidden_dim = pos_embeds.shape[2]
+
+ if pos_seq_len < max_seq_len:
+ pad_len = max_seq_len - pos_seq_len
+ pos_embeds = mx.concatenate(
+ [
+ pos_embeds,
+ mx.zeros((1, pad_len, hidden_dim), dtype=pos_embeds.dtype),
+ ],
+ axis=1,
+ )
+ pos_mask = mx.concatenate(
+ [pos_mask, mx.zeros((1, pad_len), dtype=pos_mask.dtype)],
+ axis=1,
+ )
+
+ if neg_seq_len < max_seq_len:
+ pad_len = max_seq_len - neg_seq_len
+ neg_embeds = mx.concatenate(
+ [
+ neg_embeds,
+ mx.zeros((1, pad_len, hidden_dim), dtype=neg_embeds.dtype),
+ ],
+ axis=1,
+ )
+ neg_mask = mx.concatenate(
+ [neg_mask, mx.zeros((1, pad_len), dtype=neg_mask.dtype)],
+ axis=1,
+ )
+
+ batched_embeds = mx.concatenate([pos_embeds, neg_embeds], axis=0)
+ batched_mask = mx.concatenate([pos_mask, neg_mask], axis=0)
+
+ batched_cond_latents = mx.concatenate(
+ [self._conditioning_latents, self._conditioning_latents], axis=0
+ )
+
+ return batched_embeds, batched_mask, None, batched_cond_latents
+
+
+class QwenEditModelAdapter(ModelAdapter[QwenImageEdit, QwenTransformer]):
+ """Adapter for Qwen-Image-Edit model.
+
+ Key differences from standard QwenModelAdapter:
+ - Uses QwenImageEdit model with vision-language components
+ - Encodes prompts WITH input images via VL tokenizer/encoder
+ - Creates conditioning latents from input images
+ - Supports image editing with concatenated latents during diffusion
+ """
+
+ def __init__(
+ self,
+ config: ImageModelConfig,
+ model_id: str,
+ local_path: Path,
+ quantize: int | None = None,
+ ):
+ self._config = config
+ self._model = QwenImageEdit(
+ quantize=quantize,
+ model_path=str(local_path),
+ )
+ self._transformer = self._model.transformer
+
+ self._edit_dimensions: EditImageDimensions | None = None
+
+ @property
+ def config(self) -> ImageModelConfig:
+ return self._config
+
+ @property
+ def model(self) -> QwenImageEdit:
+ return self._model
+
+ @property
+ def transformer(self) -> QwenTransformer:
+ return self._transformer
+
+ @property
+ def hidden_dim(self) -> int:
+ return self._transformer.inner_dim
+
+ @property
+ def needs_cfg(self) -> bool:
+ gs = self._config.guidance_scale
+ return gs is not None and gs > 1.0
+
+ def _get_latent_creator(self) -> type[QwenLatentCreator]:
+ return QwenLatentCreator
+
+ def get_joint_block_wrappers(
+ self,
+ text_seq_len: int,
+ encoder_hidden_states_mask: mx.array | None = None,
+ ) -> list[JointBlockWrapper[Any]]:
+ """Create wrapped joint blocks for Qwen Edit."""
+ return [
+ QwenJointBlockWrapper(block, text_seq_len, encoder_hidden_states_mask)
+ for block in self._transformer.transformer_blocks
+ ]
+
+ def get_single_block_wrappers(
+ self,
+ text_seq_len: int,
+ ) -> list[SingleBlockWrapper[Any]]:
+ """Qwen has no single blocks."""
+ return []
+
+ def slice_transformer_blocks(
+ self,
+ start_layer: int,
+ end_layer: int,
+ ):
+ self._transformer.transformer_blocks = self._transformer.transformer_blocks[
+ start_layer:end_layer
+ ]
+
+ def set_image_dimensions(self, image_path: Path) -> tuple[int, int]:
+ """Compute and store dimensions from input image.
+
+ Also stores image_paths for use in encode_prompt().
+
+ Returns:
+ (output_width, output_height) for runtime config
+ """
+ vl_w, vl_h, vae_w, vae_h, out_w, out_h = self._compute_dimensions_from_image(
+ image_path
+ )
+ self._edit_dimensions = EditImageDimensions(
+ vl_width=vl_w,
+ vl_height=vl_h,
+ vae_width=vae_w,
+ vae_height=vae_h,
+ image_paths=[str(image_path)],
+ )
+ return out_w, out_h
+
+ def create_latents(self, seed: int, runtime_config: Config) -> mx.array:
+ """Create initial noise latents (pure noise for edit mode)."""
+ return QwenLatentCreator.create_noise(
+ seed=seed,
+ height=runtime_config.height,
+ width=runtime_config.width,
+ )
+
+ def encode_prompt(
+ self, prompt: str, negative_prompt: str | None = None
+ ) -> QwenEditPromptData:
+ dims = self._edit_dimensions
+ if dims is None:
+ raise RuntimeError(
+ "set_image_dimensions() must be called before encode_prompt() "
+ "for QwenEditModelAdapter"
+ )
+
+ if negative_prompt is None or negative_prompt == "":
+ negative_prompt = " "
+
+ # TODO(ciaran): config is untyped and unused, unsure if Config or RuntimeConfig is intended
+ (
+ prompt_embeds,
+ prompt_mask,
+ negative_prompt_embeds,
+ negative_prompt_mask,
+ ) = self._model._encode_prompts_with_images(
+ prompt,
+ negative_prompt,
+ dims.image_paths,
+ self._config, # pyright: ignore[reportArgumentType]
+ dims.vl_width,
+ dims.vl_height,
+ )
+
+ (
+ conditioning_latents,
+ qwen_image_ids,
+ cond_h_patches,
+ cond_w_patches,
+ num_images,
+ ) = QwenEditUtil.create_image_conditioning_latents( # pyright: ignore[reportUnknownMemberType]
+ vae=self._model.vae,
+ height=dims.vae_height,
+ width=dims.vae_width,
+ image_paths=dims.image_paths,
+ vl_width=dims.vl_width,
+ vl_height=dims.vl_height,
+ )
+
+ if num_images > 1:
+ cond_image_grid: tuple[int, int, int] | list[tuple[int, int, int]] = [
+ (1, cond_h_patches, cond_w_patches) for _ in range(num_images)
+ ]
+ else:
+ cond_image_grid = (1, cond_h_patches, cond_w_patches)
+
+ return QwenEditPromptData(
+ prompt_embeds=prompt_embeds,
+ prompt_mask=prompt_mask,
+ negative_prompt_embeds=negative_prompt_embeds,
+ negative_prompt_mask=negative_prompt_mask,
+ conditioning_latents=conditioning_latents,
+ qwen_image_ids=qwen_image_ids,
+ cond_image_grid=cond_image_grid,
+ )
+
+ def compute_embeddings(
+ self,
+ hidden_states: mx.array,
+ prompt_embeds: mx.array,
+ ) -> tuple[mx.array, mx.array]:
+ embedded_hidden = self._transformer.img_in(hidden_states)
+ encoder_hidden_states = self._transformer.txt_norm(prompt_embeds)
+ embedded_encoder = self._transformer.txt_in(encoder_hidden_states)
+ return embedded_hidden, embedded_encoder
+
+ def compute_text_embeddings(
+ self,
+ t: int,
+ runtime_config: Config,
+ pooled_prompt_embeds: mx.array | None = None,
+ hidden_states: mx.array | None = None,
+ ) -> mx.array:
+ ref_tensor = (
+ hidden_states if hidden_states is not None else pooled_prompt_embeds
+ )
+ if ref_tensor is None:
+ raise ValueError(
+ "Either hidden_states or pooled_prompt_embeds is required "
+ "for Qwen text embeddings"
+ )
+
+ timestep = QwenTransformer._compute_timestep(t, runtime_config) # noqa: SLF001
+ batch_size = ref_tensor.shape[0]
+ timestep = mx.broadcast_to(timestep, (batch_size,)).astype(mx.float32)
+ return self._transformer.time_text_embed(timestep, ref_tensor) # pyright: ignore[reportAny]
+
+ def compute_rotary_embeddings(
+ self,
+ prompt_embeds: mx.array,
+ runtime_config: Config,
+ encoder_hidden_states_mask: mx.array | None = None,
+ cond_image_grid: tuple[int, int, int]
+ | list[tuple[int, int, int]]
+ | None = None,
+ kontext_image_ids: mx.array | None = None,
+ ) -> RotaryEmbeddings:
+ if encoder_hidden_states_mask is None:
+ raise ValueError(
+ "encoder_hidden_states_mask is required for Qwen RoPE computation"
+ )
+
+ return QwenTransformer._compute_rotary_embeddings(
+ encoder_hidden_states_mask=encoder_hidden_states_mask,
+ pos_embed=self._transformer.pos_embed, # pyright: ignore[reportAny]
+ config=runtime_config,
+ cond_image_grid=cond_image_grid,
+ )
+
+ def apply_guidance(
+ self,
+ noise_positive: mx.array,
+ noise_negative: mx.array,
+ guidance_scale: float,
+ ) -> mx.array:
+ from mflux.models.qwen.variants.txt2img.qwen_image import QwenImage
+
+ return QwenImage.compute_guided_noise(
+ noise=noise_positive,
+ noise_negative=noise_negative,
+ guidance=guidance_scale,
+ )
+
+ def _compute_dimensions_from_image(
+ self, image_path: Path
+ ) -> tuple[int, int, int, int, int, int]:
+ from mflux.utils.image_util import ImageUtil
+
+ pil_image = ImageUtil.load_image(str(image_path)).convert("RGB")
+ image_size = pil_image.size
+
+ # Vision-language dimensions (384x384 target area)
+ condition_image_size = 384 * 384
+ condition_ratio = image_size[0] / image_size[1]
+ vl_width = math.sqrt(condition_image_size * condition_ratio)
+ vl_height = vl_width / condition_ratio
+ vl_width = round(vl_width / 32) * 32
+ vl_height = round(vl_height / 32) * 32
+
+ # VAE dimensions (1024x1024 target area)
+ vae_image_size = 1024 * 1024
+ vae_ratio = image_size[0] / image_size[1]
+ vae_width = math.sqrt(vae_image_size * vae_ratio)
+ vae_height = vae_width / vae_ratio
+ vae_width = round(vae_width / 32) * 32
+ vae_height = round(vae_height / 32) * 32
+
+ # Output dimensions from input image aspect ratio
+ target_area = 1024 * 1024
+ ratio = image_size[0] / image_size[1]
+ output_width = math.sqrt(target_area * ratio)
+ output_height = output_width / ratio
+ output_width = round(output_width / 32) * 32
+ output_height = round(output_height / 32) * 32
+
+ # Ensure multiple of 16 for VAE
+ vae_scale_factor = 8
+ multiple_of = vae_scale_factor * 2
+ output_width = output_width // multiple_of * multiple_of
+ output_height = output_height // multiple_of * multiple_of
+
+ return (
+ int(vl_width),
+ int(vl_height),
+ int(vae_width),
+ int(vae_height),
+ int(output_width),
+ int(output_height),
+ )
diff --git a/src/exo/worker/engines/image/models/qwen/wrappers.py b/src/exo/worker/engines/image/models/qwen/wrappers.py
new file mode 100644
index 00000000..106f8fec
--- /dev/null
+++ b/src/exo/worker/engines/image/models/qwen/wrappers.py
@@ -0,0 +1,200 @@
+from typing import final
+
+import mlx.core as mx
+from mflux.models.qwen.model.qwen_transformer.qwen_attention import QwenAttention
+from mflux.models.qwen.model.qwen_transformer.qwen_transformer_block import (
+ QwenTransformerBlock,
+)
+from pydantic import BaseModel, ConfigDict
+
+from exo.worker.engines.image.models.base import RotaryEmbeddings
+from exo.worker.engines.image.pipeline.block_wrapper import JointBlockWrapper
+
+
+@final
+class QwenStreamModulation(BaseModel):
+ model_config = ConfigDict(frozen=True, strict=True, arbitrary_types_allowed=True)
+
+ mod1: mx.array
+ mod2: mx.array
+ gate1: mx.array
+
+
+class QwenJointBlockWrapper(JointBlockWrapper[QwenTransformerBlock]):
+ def __init__(
+ self,
+ block: QwenTransformerBlock,
+ text_seq_len: int,
+ encoder_hidden_states_mask: mx.array | None = None,
+ ):
+ super().__init__(block, text_seq_len)
+ self._encoder_hidden_states_mask = encoder_hidden_states_mask
+
+ self._num_heads = block.attn.num_heads
+ self._head_dim = block.attn.head_dim
+
+ # Intermediate state stored between _compute_qkv and _apply_output
+ self._img_mod: QwenStreamModulation | None = None
+ self._txt_mod: QwenStreamModulation | None = None
+
+ def set_encoder_mask(self, mask: mx.array | None) -> None:
+ """Set the encoder hidden states mask for attention."""
+ self._encoder_hidden_states_mask = mask
+
+ def _compute_qkv(
+ self,
+ hidden_states: mx.array,
+ encoder_hidden_states: mx.array,
+ text_embeddings: mx.array,
+ rotary_embeddings: RotaryEmbeddings,
+ patch_mode: bool = False,
+ ) -> tuple[mx.array, mx.array, mx.array]:
+ assert isinstance(rotary_embeddings, tuple)
+
+ batch_size = hidden_states.shape[0]
+ img_seq_len = hidden_states.shape[1]
+ attn = self.block.attn
+
+ img_mod_params = self.block.img_mod_linear(
+ self.block.img_mod_silu(text_embeddings) # pyright: ignore[reportUnknownArgumentType]
+ )
+ txt_mod_params = self.block.txt_mod_linear(
+ self.block.txt_mod_silu(text_embeddings) # pyright: ignore[reportUnknownArgumentType]
+ )
+
+ img_mod1, img_mod2 = mx.split(img_mod_params, 2, axis=-1)
+ txt_mod1, txt_mod2 = mx.split(txt_mod_params, 2, axis=-1)
+
+ img_normed = self.block.img_norm1(hidden_states)
+ img_modulated, img_gate1 = QwenTransformerBlock._modulate(img_normed, img_mod1)
+ self._img_mod = QwenStreamModulation(
+ mod1=img_mod1, mod2=img_mod2, gate1=img_gate1
+ )
+
+ txt_normed = self.block.txt_norm1(encoder_hidden_states)
+ txt_modulated, txt_gate1 = QwenTransformerBlock._modulate(txt_normed, txt_mod1)
+ self._txt_mod = QwenStreamModulation(
+ mod1=txt_mod1, mod2=txt_mod2, gate1=txt_gate1
+ )
+
+ img_query = attn.to_q(img_modulated)
+ img_key = attn.to_k(img_modulated)
+ img_value = attn.to_v(img_modulated)
+
+ txt_query = attn.add_q_proj(txt_modulated)
+ txt_key = attn.add_k_proj(txt_modulated)
+ txt_value = attn.add_v_proj(txt_modulated)
+
+ img_query = mx.reshape(
+ img_query, (batch_size, img_seq_len, self._num_heads, self._head_dim)
+ )
+ img_key = mx.reshape(
+ img_key, (batch_size, img_seq_len, self._num_heads, self._head_dim)
+ )
+ img_value = mx.reshape(
+ img_value, (batch_size, img_seq_len, self._num_heads, self._head_dim)
+ )
+
+ txt_query = mx.reshape(
+ txt_query,
+ (batch_size, self._text_seq_len, self._num_heads, self._head_dim),
+ )
+ txt_key = mx.reshape(
+ txt_key, (batch_size, self._text_seq_len, self._num_heads, self._head_dim)
+ )
+ txt_value = mx.reshape(
+ txt_value, (batch_size, self._text_seq_len, self._num_heads, self._head_dim)
+ )
+
+ img_query = attn.norm_q(img_query)
+ img_key = attn.norm_k(img_key)
+ txt_query = attn.norm_added_q(txt_query)
+ txt_key = attn.norm_added_k(txt_key)
+
+ (img_cos, img_sin), (txt_cos, txt_sin) = rotary_embeddings
+
+ if patch_mode:
+ # Slice image RoPE for patch, keep full text RoPE
+ img_cos = img_cos[self._patch_start : self._patch_end]
+ img_sin = img_sin[self._patch_start : self._patch_end]
+
+ img_query = QwenAttention._apply_rope_qwen(img_query, img_cos, img_sin)
+ img_key = QwenAttention._apply_rope_qwen(img_key, img_cos, img_sin)
+ txt_query = QwenAttention._apply_rope_qwen(txt_query, txt_cos, txt_sin)
+ txt_key = QwenAttention._apply_rope_qwen(txt_key, txt_cos, txt_sin)
+
+ img_query = mx.transpose(img_query, (0, 2, 1, 3))
+ img_key = mx.transpose(img_key, (0, 2, 1, 3))
+ img_value = mx.transpose(img_value, (0, 2, 1, 3))
+
+ txt_query = mx.transpose(txt_query, (0, 2, 1, 3))
+ txt_key = mx.transpose(txt_key, (0, 2, 1, 3))
+ txt_value = mx.transpose(txt_value, (0, 2, 1, 3))
+
+ query = mx.concatenate([txt_query, img_query], axis=2)
+ key = mx.concatenate([txt_key, img_key], axis=2)
+ value = mx.concatenate([txt_value, img_value], axis=2)
+
+ return query, key, value
+
+ def _compute_attention(
+ self, query: mx.array, key: mx.array, value: mx.array
+ ) -> mx.array:
+ attn = self.block.attn
+
+ mask = QwenAttention._convert_mask_for_qwen(
+ mask=self._encoder_hidden_states_mask,
+ joint_seq_len=key.shape[2],
+ txt_seq_len=self._text_seq_len,
+ )
+
+ query_bshd = mx.transpose(query, (0, 2, 1, 3))
+ key_bshd = mx.transpose(key, (0, 2, 1, 3))
+ value_bshd = mx.transpose(value, (0, 2, 1, 3))
+
+ return attn._compute_attention_qwen(
+ query=query_bshd,
+ key=key_bshd,
+ value=value_bshd,
+ mask=mask,
+ block_idx=None,
+ )
+
+ def _apply_output(
+ self,
+ attn_out: mx.array,
+ hidden_states: mx.array,
+ encoder_hidden_states: mx.array,
+ text_embeddings: mx.array,
+ ) -> tuple[mx.array, mx.array]:
+ attn = self.block.attn
+
+ assert self._img_mod is not None
+ assert self._txt_mod is not None
+
+ txt_attn_output = attn_out[:, : self._text_seq_len, :]
+ img_attn_output = attn_out[:, self._text_seq_len :, :]
+
+ img_attn_output = attn.attn_to_out[0](img_attn_output) # pyright: ignore[reportAny]
+ txt_attn_output = attn.to_add_out(txt_attn_output)
+
+ hidden_states = hidden_states + self._img_mod.gate1 * img_attn_output # pyright: ignore[reportAny]
+ encoder_hidden_states = (
+ encoder_hidden_states + self._txt_mod.gate1 * txt_attn_output
+ )
+
+ img_normed2 = self.block.img_norm2(hidden_states)
+ img_modulated2, img_gate2 = QwenTransformerBlock._modulate(
+ img_normed2, self._img_mod.mod2
+ )
+ img_mlp_output = self.block.img_ff(img_modulated2) # pyright: ignore[reportAny]
+ hidden_states = hidden_states + img_gate2 * img_mlp_output # pyright: ignore[reportAny]
+
+ txt_normed2 = self.block.txt_norm2(encoder_hidden_states)
+ txt_modulated2, txt_gate2 = QwenTransformerBlock._modulate(
+ txt_normed2, self._txt_mod.mod2
+ )
+ txt_mlp_output = self.block.txt_ff(txt_modulated2) # pyright: ignore[reportAny]
+ encoder_hidden_states = encoder_hidden_states + txt_gate2 * txt_mlp_output # pyright: ignore[reportAny]
+
+ return encoder_hidden_states, hidden_states
diff --git a/src/exo/worker/engines/image/pipeline/__init__.py b/src/exo/worker/engines/image/pipeline/__init__.py
new file mode 100644
index 00000000..5b90d0cb
--- /dev/null
+++ b/src/exo/worker/engines/image/pipeline/__init__.py
@@ -0,0 +1,15 @@
+from exo.worker.engines.image.pipeline.block_wrapper import (
+ BlockWrapperMode,
+ JointBlockWrapper,
+ SingleBlockWrapper,
+)
+from exo.worker.engines.image.pipeline.kv_cache import ImagePatchKVCache
+from exo.worker.engines.image.pipeline.runner import DiffusionRunner
+
+__all__ = [
+ "BlockWrapperMode",
+ "DiffusionRunner",
+ "ImagePatchKVCache",
+ "JointBlockWrapper",
+ "SingleBlockWrapper",
+]
diff --git a/src/exo/worker/engines/image/pipeline/block_wrapper.py b/src/exo/worker/engines/image/pipeline/block_wrapper.py
new file mode 100644
index 00000000..527eec8e
--- /dev/null
+++ b/src/exo/worker/engines/image/pipeline/block_wrapper.py
@@ -0,0 +1,303 @@
+from abc import ABC, abstractmethod
+from enum import Enum
+from typing import Generic, Self, TypeVar
+
+import mlx.core as mx
+
+from exo.worker.engines.image.models.base import RotaryEmbeddings
+from exo.worker.engines.image.pipeline.kv_cache import ImagePatchKVCache
+
+BlockT = TypeVar("BlockT")
+
+
+class BlockWrapperMode(Enum):
+ CACHING = "caching" # Sync mode: compute full attention, populate cache
+ PATCHED = "patched" # Async mode: compute patch attention, use cached KV
+
+
+class BlockWrapperMixin:
+ """Common cache management logic for block wrappers.
+
+ Including:
+ - KV cache creation and management
+ - Mode
+ - Patch range setting
+ """
+
+ _text_seq_len: int
+ _kv_cache: ImagePatchKVCache | None
+ _mode: BlockWrapperMode
+ _patch_start: int
+ _patch_end: int
+
+ def _init_cache_state(self, text_seq_len: int) -> None:
+ self._text_seq_len = text_seq_len
+ self._kv_cache = None
+ self._mode = BlockWrapperMode.CACHING
+ self._patch_start = 0
+ self._patch_end = 0
+
+ def set_patch(
+ self,
+ mode: BlockWrapperMode,
+ patch_start: int = 0,
+ patch_end: int = 0,
+ ) -> Self:
+ """Set mode and patch range.
+
+ Args:
+ mode: CACHING (full attention) or PATCHED (use cached KV)
+ patch_start: Start token index within image (for PATCHED mode)
+ patch_end: End token index within image (for PATCHED mode)
+
+ Returns:
+ Self for method chaining
+ """
+ self._mode = mode
+ self._patch_start = patch_start
+ self._patch_end = patch_end
+ return self
+
+ def set_text_seq_len(self, text_seq_len: int) -> None:
+ self._text_seq_len = text_seq_len
+
+ def _get_active_cache(self) -> ImagePatchKVCache | None:
+ return self._kv_cache
+
+ def _ensure_cache(self, img_key: mx.array) -> None:
+ if self._kv_cache is None:
+ batch, num_heads, img_seq_len, head_dim = img_key.shape
+ self._kv_cache = ImagePatchKVCache(
+ batch_size=batch,
+ num_heads=num_heads,
+ image_seq_len=img_seq_len,
+ head_dim=head_dim,
+ )
+
+ def _cache_full_image_kv(self, img_key: mx.array, img_value: mx.array) -> None:
+ self._ensure_cache(img_key)
+ cache = self._get_active_cache()
+ assert cache is not None
+ cache.update_image_patch(0, img_key.shape[2], img_key, img_value)
+
+ def _cache_patch_kv(self, img_key: mx.array, img_value: mx.array) -> None:
+ cache = self._get_active_cache()
+ assert cache is not None
+ cache.update_image_patch(self._patch_start, self._patch_end, img_key, img_value)
+
+ def _get_full_kv(
+ self, text_key: mx.array, text_value: mx.array
+ ) -> tuple[mx.array, mx.array]:
+ cache = self._get_active_cache()
+ assert cache is not None
+ return cache.get_full_kv(text_key, text_value)
+
+ def reset_cache(self) -> None:
+ self._kv_cache = None
+
+
+class JointBlockWrapper(BlockWrapperMixin, ABC, Generic[BlockT]):
+ """Base class for joint transformer block wrappers with pipefusion support.
+
+ The wrapper:
+ - Owns its KV cache (created lazily on first CACHING forward)
+ - Controls the forward pass flow (CACHING vs PATCHED mode)
+ - Handles patch slicing and cache operations
+ """
+
+ block: BlockT
+
+ def __init__(self, block: BlockT, text_seq_len: int):
+ self.block = block
+ self._init_cache_state(text_seq_len)
+
+ def set_encoder_mask(self, mask: mx.array | None) -> None: # noqa: B027
+ """Set the encoder hidden states mask for attention.
+
+ Override in subclasses that use attention masks
+ Default is a no-op for models that don't use masks
+ """
+ del mask # Unused in base class
+
+ def __call__(
+ self,
+ hidden_states: mx.array,
+ encoder_hidden_states: mx.array,
+ text_embeddings: mx.array,
+ rotary_embeddings: RotaryEmbeddings,
+ ) -> tuple[mx.array, mx.array]:
+ if self._mode == BlockWrapperMode.CACHING:
+ return self._forward_caching(
+ hidden_states, encoder_hidden_states, text_embeddings, rotary_embeddings
+ )
+ return self._forward_patched(
+ hidden_states, encoder_hidden_states, text_embeddings, rotary_embeddings
+ )
+
+ def _forward_caching(
+ self,
+ hidden_states: mx.array,
+ encoder_hidden_states: mx.array,
+ text_embeddings: mx.array,
+ rotary_embeddings: RotaryEmbeddings,
+ ) -> tuple[mx.array, mx.array]:
+ """CACHING mode: Full attention, store image K/V in cache."""
+ query, key, value = self._compute_qkv(
+ hidden_states, encoder_hidden_states, text_embeddings, rotary_embeddings
+ )
+
+ img_key = key[:, :, self._text_seq_len :, :]
+ img_value = value[:, :, self._text_seq_len :, :]
+ self._cache_full_image_kv(img_key, img_value)
+
+ attn_out = self._compute_attention(query, key, value)
+
+ return self._apply_output(
+ attn_out, hidden_states, encoder_hidden_states, text_embeddings
+ )
+
+ def _forward_patched(
+ self,
+ hidden_states: mx.array,
+ encoder_hidden_states: mx.array,
+ text_embeddings: mx.array,
+ rotary_embeddings: RotaryEmbeddings,
+ ) -> tuple[mx.array, mx.array]:
+ # hidden_states is already the patch (provided by runner)
+ patch_hidden = hidden_states
+
+ query, key, value = self._compute_qkv(
+ patch_hidden,
+ encoder_hidden_states,
+ text_embeddings,
+ rotary_embeddings,
+ patch_mode=True,
+ )
+
+ text_key = key[:, :, : self._text_seq_len, :]
+ text_value = value[:, :, : self._text_seq_len, :]
+ img_key = key[:, :, self._text_seq_len :, :]
+ img_value = value[:, :, self._text_seq_len :, :]
+
+ self._cache_patch_kv(img_key, img_value)
+ full_key, full_value = self._get_full_kv(text_key, text_value)
+
+ attn_out = self._compute_attention(query, full_key, full_value)
+
+ return self._apply_output(
+ attn_out, patch_hidden, encoder_hidden_states, text_embeddings
+ )
+
+ @abstractmethod
+ def _compute_qkv(
+ self,
+ hidden_states: mx.array,
+ encoder_hidden_states: mx.array,
+ text_embeddings: mx.array,
+ rotary_embeddings: RotaryEmbeddings,
+ patch_mode: bool = False,
+ ) -> tuple[mx.array, mx.array, mx.array]: ...
+
+ @abstractmethod
+ def _compute_attention(
+ self, query: mx.array, key: mx.array, value: mx.array
+ ) -> mx.array: ...
+
+ @abstractmethod
+ def _apply_output(
+ self,
+ attn_out: mx.array,
+ hidden_states: mx.array,
+ encoder_hidden_states: mx.array,
+ text_embeddings: mx.array,
+ ) -> tuple[mx.array, mx.array]: ...
+
+
+class SingleBlockWrapper(BlockWrapperMixin, ABC, Generic[BlockT]):
+ """Base class for single-stream transformer block wrappers.
+
+ Similar to JointBlockWrapper but for blocks that operate on a single
+ concatenated [text, image] stream rather than separate streams.
+ """
+
+ block: BlockT
+
+ def __init__(self, block: BlockT, text_seq_len: int):
+ self.block = block
+ self._init_cache_state(text_seq_len)
+
+ def __call__(
+ self,
+ hidden_states: mx.array,
+ text_embeddings: mx.array,
+ rotary_embeddings: RotaryEmbeddings,
+ ) -> mx.array:
+ if self._mode == BlockWrapperMode.CACHING:
+ return self._forward_caching(
+ hidden_states, text_embeddings, rotary_embeddings
+ )
+ return self._forward_patched(hidden_states, text_embeddings, rotary_embeddings)
+
+ def _forward_caching(
+ self,
+ hidden_states: mx.array,
+ text_embeddings: mx.array,
+ rotary_embeddings: RotaryEmbeddings,
+ ) -> mx.array:
+ """CACHING mode: Full attention, store image K/V in cache."""
+ query, key, value = self._compute_qkv(
+ hidden_states, text_embeddings, rotary_embeddings
+ )
+
+ img_key = key[:, :, self._text_seq_len :, :]
+ img_value = value[:, :, self._text_seq_len :, :]
+ self._cache_full_image_kv(img_key, img_value)
+
+ attn_out = self._compute_attention(query, key, value)
+
+ return self._apply_output(attn_out, hidden_states, text_embeddings)
+
+ def _forward_patched(
+ self,
+ hidden_states: mx.array,
+ text_embeddings: mx.array,
+ rotary_embeddings: RotaryEmbeddings,
+ ) -> mx.array:
+ """PATCHED mode: Compute patch Q/K/V, use cached image K/V for attention."""
+ query, key, value = self._compute_qkv(
+ hidden_states, text_embeddings, rotary_embeddings, patch_mode=True
+ )
+
+ text_key = key[:, :, : self._text_seq_len, :]
+ text_value = value[:, :, : self._text_seq_len, :]
+ img_key = key[:, :, self._text_seq_len :, :]
+ img_value = value[:, :, self._text_seq_len :, :]
+
+ self._cache_patch_kv(img_key, img_value)
+ full_key, full_value = self._get_full_kv(text_key, text_value)
+
+ attn_out = self._compute_attention(query, full_key, full_value)
+
+ return self._apply_output(attn_out, hidden_states, text_embeddings)
+
+ @abstractmethod
+ def _compute_qkv(
+ self,
+ hidden_states: mx.array,
+ text_embeddings: mx.array,
+ rotary_embeddings: RotaryEmbeddings,
+ patch_mode: bool = False,
+ ) -> tuple[mx.array, mx.array, mx.array]: ...
+
+ @abstractmethod
+ def _compute_attention(
+ self, query: mx.array, key: mx.array, value: mx.array
+ ) -> mx.array: ...
+
+ @abstractmethod
+ def _apply_output(
+ self,
+ attn_out: mx.array,
+ hidden_states: mx.array,
+ text_embeddings: mx.array,
+ ) -> mx.array: ...
diff --git a/src/exo/worker/engines/image/pipeline/kv_cache.py b/src/exo/worker/engines/image/pipeline/kv_cache.py
new file mode 100644
index 00000000..22c6f0d4
--- /dev/null
+++ b/src/exo/worker/engines/image/pipeline/kv_cache.py
@@ -0,0 +1,72 @@
+import mlx.core as mx
+
+
+class ImagePatchKVCache:
+ """KV cache that stores only IMAGE K/V with patch-level updates.
+
+ Only caches image K/V since:
+ - Text K/V is always computed fresh (same for all patches)
+ - Only image portion needs stale/fresh cache management across patches
+ """
+
+ def __init__(
+ self,
+ batch_size: int,
+ num_heads: int,
+ image_seq_len: int,
+ head_dim: int,
+ dtype: mx.Dtype = mx.float32,
+ ):
+ self.batch_size = batch_size
+ self.num_heads = num_heads
+ self.image_seq_len = image_seq_len
+ self.head_dim = head_dim
+ self._dtype = dtype
+
+ self.key_cache = mx.zeros(
+ (batch_size, num_heads, image_seq_len, head_dim), dtype=dtype
+ )
+ self.value_cache = mx.zeros(
+ (batch_size, num_heads, image_seq_len, head_dim), dtype=dtype
+ )
+
+ def update_image_patch(
+ self, patch_start: int, patch_end: int, key: mx.array, value: mx.array
+ ) -> None:
+ """Update cache with fresh K/V for an image patch slice.
+
+ Args:
+ patch_start: Start token index within image portion (0-indexed)
+ patch_end: End token index within image portion
+ key: Fresh key tensor [batch, heads, patch_seq_len, head_dim]
+ value: Fresh value tensor [batch, heads, patch_seq_len, head_dim]
+ """
+ self.key_cache[:, :, patch_start:patch_end, :] = key
+ self.value_cache[:, :, patch_start:patch_end, :] = value
+
+ def get_full_kv(
+ self, text_key: mx.array, text_value: mx.array
+ ) -> tuple[mx.array, mx.array]:
+ """Return full K/V by concatenating fresh text K/V with cached image K/V.
+
+ Args:
+ text_key: Fresh text key tensor [batch, heads, text_seq_len, head_dim]
+ text_value: Fresh text value tensor [batch, heads, text_seq_len, head_dim]
+
+ Returns:
+ Tuple of (full_key, full_value) with shape [batch, heads, text+image, head_dim]
+ """
+ full_key = mx.concatenate([text_key, self.key_cache], axis=2)
+ full_value = mx.concatenate([text_value, self.value_cache], axis=2)
+ return full_key, full_value
+
+ def reset(self) -> None:
+ """Reset cache to zeros."""
+ self.key_cache = mx.zeros(
+ (self.batch_size, self.num_heads, self.image_seq_len, self.head_dim),
+ dtype=self._dtype,
+ )
+ self.value_cache = mx.zeros(
+ (self.batch_size, self.num_heads, self.image_seq_len, self.head_dim),
+ dtype=self._dtype,
+ )
diff --git a/src/exo/worker/engines/image/pipeline/runner.py b/src/exo/worker/engines/image/pipeline/runner.py
new file mode 100644
index 00000000..cd5d92d5
--- /dev/null
+++ b/src/exo/worker/engines/image/pipeline/runner.py
@@ -0,0 +1,972 @@
+from math import ceil
+from typing import Any, Optional
+
+import mlx.core as mx
+from mflux.models.common.config.config import Config
+from mflux.utils.exceptions import StopImageGenerationException
+from tqdm import tqdm
+
+from exo.shared.types.worker.shards import PipelineShardMetadata
+from exo.worker.engines.image.config import ImageModelConfig
+from exo.worker.engines.image.models.base import (
+ ModelAdapter,
+ PromptData,
+ RotaryEmbeddings,
+)
+from exo.worker.engines.image.pipeline.block_wrapper import (
+ BlockWrapperMode,
+ JointBlockWrapper,
+ SingleBlockWrapper,
+)
+
+
+def calculate_patch_heights(
+ latent_height: int, num_patches: int
+) -> tuple[list[int], int]:
+ patch_height = ceil(latent_height / num_patches)
+
+ actual_num_patches = ceil(latent_height / patch_height)
+ patch_heights = [patch_height] * (actual_num_patches - 1)
+
+ last_height = latent_height - patch_height * (actual_num_patches - 1)
+ patch_heights.append(last_height)
+
+ return patch_heights, actual_num_patches
+
+
+def calculate_token_indices(
+ patch_heights: list[int], latent_width: int
+) -> list[tuple[int, int]]:
+ tokens_per_row = latent_width
+
+ token_ranges: list[tuple[int, int]] = []
+ cumulative_height = 0
+
+ for h in patch_heights:
+ start_token = tokens_per_row * cumulative_height
+ end_token = tokens_per_row * (cumulative_height + h)
+
+ token_ranges.append((start_token, end_token))
+ cumulative_height += h
+
+ return token_ranges
+
+
+class DiffusionRunner:
+ """Orchestrates the diffusion loop for image generation.
+
+ In distributed mode, it implements PipeFusion with:
+ - Sync pipeline for initial timesteps (full image, all devices in lockstep)
+ - Async pipeline for later timesteps (patches processed independently)
+ """
+
+ def __init__(
+ self,
+ config: ImageModelConfig,
+ adapter: ModelAdapter[Any, Any],
+ group: Optional[mx.distributed.Group],
+ shard_metadata: PipelineShardMetadata,
+ num_patches: Optional[int] = None,
+ ):
+ self.config = config
+ self.adapter = adapter
+ self.group = group
+
+ if group is None:
+ self.rank = 0
+ self.world_size = 1
+ self.next_rank = 0
+ self.prev_rank = 0
+ self.start_layer = 0
+ self.end_layer = config.total_blocks
+ else:
+ self.rank = shard_metadata.device_rank
+ self.world_size = shard_metadata.world_size
+ self.next_rank = (self.rank + 1) % self.world_size
+ self.prev_rank = (self.rank - 1 + self.world_size) % self.world_size
+ self.start_layer = shard_metadata.start_layer
+ self.end_layer = shard_metadata.end_layer
+
+ self.num_patches = num_patches if num_patches else max(1, self.world_size)
+
+ self.total_joint = config.joint_block_count
+ self.total_single = config.single_block_count
+ self.total_layers = config.total_blocks
+
+ self._guidance_override: float | None = None
+
+ self._compute_assigned_blocks()
+
+ def _compute_assigned_blocks(self) -> None:
+ """Determine which joint/single blocks this stage owns."""
+ start = self.start_layer
+ end = self.end_layer
+
+ if end <= self.total_joint:
+ self.joint_start = start
+ self.joint_end = end
+ self.single_start = 0
+ self.single_end = 0
+ elif start >= self.total_joint:
+ self.joint_start = 0
+ self.joint_end = 0
+ self.single_start = start - self.total_joint
+ self.single_end = end - self.total_joint
+ else:
+ self.joint_start = start
+ self.joint_end = self.total_joint
+ self.single_start = 0
+ self.single_end = end - self.total_joint
+
+ self.has_joint_blocks = self.joint_end > self.joint_start
+ self.has_single_blocks = self.single_end > self.single_start
+
+ self.owns_concat_stage = self.has_joint_blocks and (
+ self.has_single_blocks or self.end_layer == self.total_joint
+ )
+
+ # Wrappers created lazily on first forward (need text_seq_len)
+ self.joint_block_wrappers: list[JointBlockWrapper[Any]] | None = None
+ self.single_block_wrappers: list[SingleBlockWrapper[Any]] | None = None
+ self._wrappers_initialized = False
+ self._current_text_seq_len: int | None = None
+
+ @property
+ def is_first_stage(self) -> bool:
+ return self.rank == 0
+
+ @property
+ def is_last_stage(self) -> bool:
+ return self.rank == self.world_size - 1
+
+ @property
+ def is_distributed(self) -> bool:
+ return self.group is not None
+
+ def _get_effective_guidance_scale(self) -> float | None:
+ if self._guidance_override is not None:
+ return self._guidance_override
+ return self.config.guidance_scale
+
+ def _ensure_wrappers(
+ self,
+ text_seq_len: int,
+ encoder_hidden_states_mask: mx.array | None = None,
+ ) -> None:
+ """Lazily create block wrappers on first forward pass.
+
+ Wrappers need text_seq_len which is only known after prompt encoding.
+ Re-initializes if text_seq_len changes (e.g., warmup vs real generation).
+ """
+ if self._wrappers_initialized and self._current_text_seq_len == text_seq_len:
+ return
+
+ self.joint_block_wrappers = self.adapter.get_joint_block_wrappers(
+ text_seq_len=text_seq_len,
+ encoder_hidden_states_mask=encoder_hidden_states_mask,
+ )
+ self.single_block_wrappers = self.adapter.get_single_block_wrappers(
+ text_seq_len=text_seq_len,
+ )
+ self._wrappers_initialized = True
+ self._current_text_seq_len = text_seq_len
+
+ def _reset_all_caches(self) -> None:
+ """Reset KV caches on all wrappers for a new generation."""
+ if self.joint_block_wrappers:
+ for wrapper in self.joint_block_wrappers:
+ wrapper.reset_cache()
+ if self.single_block_wrappers:
+ for wrapper in self.single_block_wrappers:
+ wrapper.reset_cache()
+
+ def _set_text_seq_len(self, text_seq_len: int) -> None:
+ if self.joint_block_wrappers:
+ for wrapper in self.joint_block_wrappers:
+ wrapper.set_text_seq_len(text_seq_len)
+ if self.single_block_wrappers:
+ for wrapper in self.single_block_wrappers:
+ wrapper.set_text_seq_len(text_seq_len)
+
+ def _calculate_capture_steps(
+ self,
+ partial_images: int,
+ init_time_step: int,
+ num_inference_steps: int,
+ ) -> set[int]:
+ """Calculate which timesteps should produce partial images.
+
+ Places the first partial after step 1 for fast initial feedback,
+ then evenly spaces remaining partials with equal gaps between them
+ and from the last partial to the final image.
+
+ Args:
+ partial_images: Number of partial images to capture
+ init_time_step: Starting timestep (for img2img this may not be 0)
+ num_inference_steps: Total inference steps
+
+ Returns:
+ Set of timestep indices to capture
+ """
+ if partial_images <= 0:
+ return set()
+
+ total_steps = num_inference_steps - init_time_step
+ if total_steps <= 1:
+ return set()
+
+ if partial_images >= total_steps - 1:
+ return set(range(init_time_step, num_inference_steps - 1))
+
+ capture_steps: set[int] = set()
+
+ first_capture = init_time_step + 1
+ capture_steps.add(first_capture)
+
+ if partial_images == 1:
+ return capture_steps
+
+ final_step = num_inference_steps - 1
+ remaining_range = final_step - first_capture
+
+ for i in range(1, partial_images):
+ step_idx = first_capture + int(i * remaining_range / partial_images)
+ capture_steps.add(step_idx)
+
+ return capture_steps
+
+ def generate_image(
+ self,
+ runtime_config: Config,
+ prompt: str,
+ seed: int,
+ partial_images: int = 0,
+ guidance_override: float | None = None,
+ negative_prompt: str | None = None,
+ num_sync_steps: int = 1,
+ ):
+ """Primary entry point for image generation.
+
+ Orchestrates the full generation flow:
+ 1. Create runtime config
+ 2. Create initial latents
+ 3. Encode prompt
+ 4. Run diffusion loop (yielding partials if requested)
+ 5. Decode to image
+
+ Args:
+ settings: Generation config (steps, height, width)
+ prompt: Text prompt
+ seed: Random seed
+ partial_images: Number of intermediate images to yield (0 for none)
+ guidance_override: Optional override for guidance scale (CFG)
+
+ Yields:
+ Partial images as (GeneratedImage, partial_index, total_partials) tuples
+ Final GeneratedImage
+ """
+ self._guidance_override = guidance_override
+ latents = self.adapter.create_latents(seed, runtime_config)
+ prompt_data = self.adapter.encode_prompt(prompt, negative_prompt)
+
+ capture_steps = self._calculate_capture_steps(
+ partial_images=partial_images,
+ init_time_step=runtime_config.init_time_step,
+ num_inference_steps=runtime_config.num_inference_steps,
+ )
+
+ diffusion_gen = self._run_diffusion_loop(
+ latents=latents,
+ prompt_data=prompt_data,
+ runtime_config=runtime_config,
+ seed=seed,
+ prompt=prompt,
+ capture_steps=capture_steps,
+ num_sync_steps=num_sync_steps,
+ )
+
+ partial_index = 0
+ total_partials = len(capture_steps)
+
+ if capture_steps:
+ try:
+ while True:
+ partial_latents, _step = next(diffusion_gen)
+ if self.is_last_stage:
+ partial_image = self.adapter.decode_latents(
+ partial_latents, runtime_config, seed, prompt
+ )
+ yield (partial_image, partial_index, total_partials)
+ partial_index += 1
+ except StopIteration as e:
+ latents = e.value # pyright: ignore[reportAny]
+ else:
+ try:
+ while True:
+ next(diffusion_gen)
+ except StopIteration as e:
+ latents = e.value # pyright: ignore[reportAny]
+
+ if self.is_last_stage:
+ yield self.adapter.decode_latents(latents, runtime_config, seed, prompt) # pyright: ignore[reportAny]
+
+ def _run_diffusion_loop(
+ self,
+ latents: mx.array,
+ prompt_data: PromptData,
+ runtime_config: Config,
+ seed: int,
+ prompt: str,
+ num_sync_steps: int,
+ capture_steps: set[int] | None = None,
+ ):
+ if capture_steps is None:
+ capture_steps = set()
+
+ self._reset_all_caches()
+
+ time_steps = tqdm(range(runtime_config.num_inference_steps))
+
+ ctx = self.adapter.model.callbacks.start( # pyright: ignore[reportAny]
+ seed=seed, prompt=prompt, config=runtime_config
+ )
+
+ ctx.before_loop( # pyright: ignore[reportAny]
+ latents=latents,
+ )
+
+ for t in time_steps:
+ try:
+ latents = self._diffusion_step(
+ t=t,
+ config=runtime_config,
+ latents=latents,
+ prompt_data=prompt_data,
+ num_sync_steps=num_sync_steps,
+ )
+
+ ctx.in_loop( # pyright: ignore[reportAny]
+ t=t,
+ latents=latents,
+ )
+
+ mx.eval(latents)
+
+ if t in capture_steps and self.is_last_stage:
+ yield (latents, t)
+
+ except KeyboardInterrupt: # noqa: PERF203
+ ctx.interruption(t=t, latents=latents) # pyright: ignore[reportAny]
+ raise StopImageGenerationException(
+ f"Stopping image generation at step {t + 1}/{len(time_steps)}"
+ ) from None
+
+ ctx.after_loop(latents=latents) # pyright: ignore[reportAny]
+
+ return latents
+
+ def _forward_pass(
+ self,
+ latents: mx.array,
+ prompt_embeds: mx.array,
+ pooled_prompt_embeds: mx.array,
+ t: int,
+ config: Config,
+ encoder_hidden_states_mask: mx.array | None = None,
+ cond_image_grid: tuple[int, int, int]
+ | list[tuple[int, int, int]]
+ | None = None,
+ conditioning_latents: mx.array | None = None,
+ ) -> mx.array:
+ """Run a single forward pass through the transformer.
+ Args:
+ latents: Input latents (already scaled by caller)
+ prompt_embeds: Text embeddings
+ pooled_prompt_embeds: Pooled text embeddings (Flux) or placeholder (Qwen)
+ t: Current timestep
+ config: Runtime configuration
+ encoder_hidden_states_mask: Attention mask for text (Qwen)
+ cond_image_grid: Conditioning image grid dimensions (Qwen edit)
+ conditioning_latents: Conditioning latents for edit mode
+
+ Returns:
+ Noise prediction tensor
+ """
+ text_seq_len = prompt_embeds.shape[1]
+
+ self._ensure_wrappers(text_seq_len, encoder_hidden_states_mask)
+
+ if self.joint_block_wrappers and encoder_hidden_states_mask is not None:
+ for wrapper in self.joint_block_wrappers:
+ wrapper.set_encoder_mask(encoder_hidden_states_mask)
+
+ scaled_latents = config.scheduler.scale_model_input(latents, t) # pyright: ignore[reportAny]
+
+ # For edit mode: concatenate with conditioning latents
+ original_latent_tokens: int = scaled_latents.shape[1] # pyright: ignore[reportAny]
+ if conditioning_latents is not None:
+ scaled_latents = mx.concatenate(
+ [scaled_latents, conditioning_latents], axis=1
+ )
+
+ hidden_states, encoder_hidden_states = self.adapter.compute_embeddings(
+ scaled_latents, prompt_embeds
+ )
+ text_embeddings = self.adapter.compute_text_embeddings(
+ t, config, pooled_prompt_embeds, hidden_states=hidden_states
+ )
+ rotary_embeddings = self.adapter.compute_rotary_embeddings(
+ prompt_embeds,
+ config,
+ encoder_hidden_states_mask=encoder_hidden_states_mask,
+ cond_image_grid=cond_image_grid,
+ )
+
+ assert self.joint_block_wrappers is not None
+ for wrapper in self.joint_block_wrappers:
+ encoder_hidden_states, hidden_states = wrapper(
+ hidden_states=hidden_states,
+ encoder_hidden_states=encoder_hidden_states,
+ text_embeddings=text_embeddings,
+ rotary_embeddings=rotary_embeddings,
+ )
+
+ if self.joint_block_wrappers:
+ hidden_states = self.adapter.merge_streams(
+ hidden_states, encoder_hidden_states
+ )
+
+ assert self.single_block_wrappers is not None
+ for wrapper in self.single_block_wrappers:
+ hidden_states = wrapper(
+ hidden_states=hidden_states,
+ text_embeddings=text_embeddings,
+ rotary_embeddings=rotary_embeddings,
+ )
+
+ # Extract image portion
+ hidden_states = hidden_states[:, text_seq_len:, ...]
+
+ # For edit mode: extract only the generated portion (exclude conditioning latents)
+ if conditioning_latents is not None:
+ hidden_states = hidden_states[:, :original_latent_tokens, ...]
+
+ return self.adapter.final_projection(hidden_states, text_embeddings)
+
+ def _diffusion_step(
+ self,
+ t: int,
+ config: Config,
+ latents: mx.array,
+ prompt_data: PromptData,
+ num_sync_steps: int,
+ ) -> mx.array:
+ if self.group is None:
+ return self._single_node_step(t, config, latents, prompt_data)
+ elif t < config.init_time_step + num_sync_steps:
+ return self._sync_pipeline_step(
+ t,
+ config,
+ latents,
+ prompt_data,
+ )
+ else:
+ return self._async_pipeline_step(
+ t,
+ config,
+ latents,
+ prompt_data,
+ is_first_async_step=t == config.init_time_step + num_sync_steps,
+ )
+
+ def _single_node_step(
+ self,
+ t: int,
+ config: Config,
+ latents: mx.array,
+ prompt_data: PromptData,
+ ) -> mx.array:
+ cond_image_grid = prompt_data.cond_image_grid
+ needs_cfg = self.adapter.needs_cfg
+
+ if needs_cfg:
+ batched_data = prompt_data.get_batched_cfg_data()
+ assert batched_data is not None, "CFG model must provide batched data"
+ prompt_embeds, encoder_mask, batched_pooled, cond_latents = batched_data
+ pooled_embeds = (
+ batched_pooled if batched_pooled is not None else prompt_embeds
+ )
+ step_latents = mx.concatenate([latents, latents], axis=0)
+ else:
+ prompt_embeds = prompt_data.prompt_embeds
+ pooled_embeds = prompt_data.pooled_prompt_embeds
+ encoder_mask = prompt_data.get_encoder_hidden_states_mask(positive=True)
+ cond_latents = prompt_data.conditioning_latents
+ step_latents = latents
+
+ noise = self._forward_pass(
+ step_latents,
+ prompt_embeds,
+ pooled_embeds,
+ t=t,
+ config=config,
+ encoder_hidden_states_mask=encoder_mask,
+ cond_image_grid=cond_image_grid,
+ conditioning_latents=cond_latents,
+ )
+
+ if needs_cfg:
+ noise_pos, noise_neg = mx.split(noise, 2, axis=0)
+ guidance_scale = self._get_effective_guidance_scale()
+ assert guidance_scale is not None
+ noise = self.adapter.apply_guidance(
+ noise_pos, noise_neg, guidance_scale=guidance_scale
+ )
+
+ return config.scheduler.step(noise=noise, timestep=t, latents=latents) # pyright: ignore[reportAny]
+
+ def _create_patches(
+ self,
+ latents: mx.array,
+ config: Config,
+ ) -> tuple[list[mx.array], list[tuple[int, int]]]:
+ latent_height = config.height // 16
+ latent_width = config.width // 16
+
+ patch_heights, _ = calculate_patch_heights(latent_height, self.num_patches)
+ token_indices = calculate_token_indices(patch_heights, latent_width)
+
+ patch_latents = [latents[:, start:end, :] for start, end in token_indices]
+
+ return patch_latents, token_indices
+
+ def _run_sync_pass(
+ self,
+ t: int,
+ config: Config,
+ scaled_hidden_states: mx.array,
+ prompt_embeds: mx.array,
+ pooled_prompt_embeds: mx.array,
+ encoder_hidden_states_mask: mx.array | None,
+ cond_image_grid: tuple[int, int, int] | list[tuple[int, int, int]] | None,
+ kontext_image_ids: mx.array | None,
+ num_img_tokens: int,
+ original_latent_tokens: int,
+ conditioning_latents: mx.array | None,
+ ) -> mx.array | None:
+ hidden_states = scaled_hidden_states
+ batch_size = hidden_states.shape[0]
+ text_seq_len = prompt_embeds.shape[1]
+ hidden_dim = self.adapter.hidden_dim
+ dtype = scaled_hidden_states.dtype
+
+ self._set_text_seq_len(text_seq_len)
+
+ if self.joint_block_wrappers:
+ for wrapper in self.joint_block_wrappers:
+ wrapper.set_encoder_mask(encoder_hidden_states_mask)
+
+ encoder_hidden_states: mx.array | None = None
+ if self.is_first_stage:
+ hidden_states, encoder_hidden_states = self.adapter.compute_embeddings(
+ hidden_states, prompt_embeds
+ )
+
+ text_embeddings = self.adapter.compute_text_embeddings(
+ t, config, pooled_prompt_embeds
+ )
+ image_rotary_embeddings = self.adapter.compute_rotary_embeddings(
+ prompt_embeds,
+ config,
+ encoder_hidden_states_mask=encoder_hidden_states_mask,
+ cond_image_grid=cond_image_grid,
+ kontext_image_ids=kontext_image_ids,
+ )
+
+ if self.has_joint_blocks:
+ if not self.is_first_stage:
+ hidden_states = mx.distributed.recv(
+ (batch_size, num_img_tokens, hidden_dim),
+ dtype,
+ self.prev_rank,
+ group=self.group,
+ )
+ encoder_hidden_states = mx.distributed.recv(
+ (batch_size, text_seq_len, hidden_dim),
+ dtype,
+ self.prev_rank,
+ group=self.group,
+ )
+ mx.eval(hidden_states, encoder_hidden_states)
+
+ assert self.joint_block_wrappers is not None
+ assert encoder_hidden_states is not None
+ for wrapper in self.joint_block_wrappers:
+ wrapper.set_patch(BlockWrapperMode.CACHING)
+ encoder_hidden_states, hidden_states = wrapper(
+ hidden_states=hidden_states,
+ encoder_hidden_states=encoder_hidden_states,
+ text_embeddings=text_embeddings,
+ rotary_embeddings=image_rotary_embeddings,
+ )
+
+ if self.owns_concat_stage:
+ assert encoder_hidden_states is not None
+ concatenated = self.adapter.merge_streams(
+ hidden_states, encoder_hidden_states
+ )
+
+ if self.has_single_blocks or self.is_last_stage:
+ hidden_states = concatenated
+ else:
+ concatenated = mx.distributed.send(
+ concatenated, self.next_rank, group=self.group
+ )
+ mx.async_eval(concatenated)
+
+ elif self.has_joint_blocks and not self.is_last_stage:
+ assert encoder_hidden_states is not None
+ hidden_states = mx.distributed.send(
+ hidden_states, self.next_rank, group=self.group
+ )
+ encoder_hidden_states = mx.distributed.send(
+ encoder_hidden_states, self.next_rank, group=self.group
+ )
+ mx.async_eval(hidden_states, encoder_hidden_states)
+
+ if self.has_single_blocks:
+ if not self.owns_concat_stage and not self.is_first_stage:
+ hidden_states = mx.distributed.recv(
+ (batch_size, text_seq_len + num_img_tokens, hidden_dim),
+ dtype,
+ self.prev_rank,
+ group=self.group,
+ )
+ mx.eval(hidden_states)
+
+ assert self.single_block_wrappers is not None
+ for wrapper in self.single_block_wrappers:
+ wrapper.set_patch(BlockWrapperMode.CACHING)
+ hidden_states = wrapper(
+ hidden_states=hidden_states,
+ text_embeddings=text_embeddings,
+ rotary_embeddings=image_rotary_embeddings,
+ )
+
+ if not self.is_last_stage:
+ hidden_states = mx.distributed.send(
+ hidden_states, self.next_rank, group=self.group
+ )
+ mx.async_eval(hidden_states)
+
+ hidden_states = hidden_states[:, text_seq_len:, ...]
+
+ if conditioning_latents is not None:
+ hidden_states = hidden_states[:, :original_latent_tokens, ...]
+
+ if self.is_last_stage:
+ return self.adapter.final_projection(hidden_states, text_embeddings)
+
+ return None
+
+ def _sync_pipeline_step(
+ self,
+ t: int,
+ config: Config,
+ hidden_states: mx.array,
+ prompt_data: PromptData,
+ kontext_image_ids: mx.array | None = None,
+ ) -> mx.array:
+ prev_latents = hidden_states
+ needs_cfg = self.adapter.needs_cfg
+ cond_image_grid = prompt_data.cond_image_grid
+
+ scaled_hidden_states = config.scheduler.scale_model_input(hidden_states, t) # pyright: ignore[reportAny]
+ original_latent_tokens: int = scaled_hidden_states.shape[1] # pyright: ignore[reportAny]
+
+ if needs_cfg:
+ batched_data = prompt_data.get_batched_cfg_data()
+ assert batched_data is not None, "CFG model must provide batched data"
+ prompt_embeds, encoder_mask, batched_pooled, cond_latents = batched_data
+ pooled_embeds = (
+ batched_pooled if batched_pooled is not None else prompt_embeds
+ )
+ step_latents = mx.concatenate(
+ [scaled_hidden_states, scaled_hidden_states], axis=0
+ )
+ else:
+ prompt_embeds = prompt_data.prompt_embeds
+ pooled_embeds = prompt_data.pooled_prompt_embeds
+ encoder_mask = prompt_data.get_encoder_hidden_states_mask(positive=True)
+ cond_latents = prompt_data.conditioning_latents
+ step_latents = scaled_hidden_states # pyright: ignore[reportAny]
+
+ if cond_latents is not None:
+ num_img_tokens: int = original_latent_tokens + cond_latents.shape[1]
+ else:
+ num_img_tokens = original_latent_tokens
+
+ if self.is_first_stage and cond_latents is not None:
+ step_latents = mx.concatenate([step_latents, cond_latents], axis=1)
+
+ text_seq_len = prompt_embeds.shape[1]
+ self._ensure_wrappers(text_seq_len, encoder_mask)
+
+ noise = self._run_sync_pass(
+ t,
+ config,
+ step_latents,
+ prompt_embeds,
+ pooled_embeds,
+ encoder_mask,
+ cond_image_grid,
+ kontext_image_ids,
+ num_img_tokens,
+ original_latent_tokens,
+ cond_latents,
+ )
+
+ if self.is_last_stage:
+ assert noise is not None
+ if needs_cfg:
+ noise_pos, noise_neg = mx.split(noise, 2, axis=0)
+ guidance_scale = self._get_effective_guidance_scale()
+ assert guidance_scale is not None
+ noise = self.adapter.apply_guidance(
+ noise_pos, noise_neg, guidance_scale
+ )
+
+ hidden_states = config.scheduler.step( # pyright: ignore[reportAny]
+ noise=noise, timestep=t, latents=prev_latents
+ )
+
+ if not self.is_first_stage:
+ hidden_states = mx.distributed.send(hidden_states, 0, group=self.group)
+ mx.async_eval(hidden_states)
+
+ elif self.is_first_stage:
+ hidden_states = mx.distributed.recv_like(
+ prev_latents, src=self.world_size - 1, group=self.group
+ )
+ mx.eval(hidden_states)
+
+ else:
+ hidden_states = prev_latents
+
+ return hidden_states
+
+ def _async_pipeline_step(
+ self,
+ t: int,
+ config: Config,
+ latents: mx.array,
+ prompt_data: PromptData,
+ is_first_async_step: bool,
+ kontext_image_ids: mx.array | None = None,
+ ) -> mx.array:
+ patch_latents, token_indices = self._create_patches(latents, config)
+ needs_cfg = self.adapter.needs_cfg
+ cond_image_grid = prompt_data.cond_image_grid
+
+ if needs_cfg:
+ batched_data = prompt_data.get_batched_cfg_data()
+ assert batched_data is not None, "CFG model must provide batched data"
+ prompt_embeds, encoder_mask, batched_pooled, _ = batched_data
+ pooled_embeds = (
+ batched_pooled if batched_pooled is not None else prompt_embeds
+ )
+ else:
+ prompt_embeds = prompt_data.prompt_embeds
+ pooled_embeds = prompt_data.pooled_prompt_embeds
+ encoder_mask = prompt_data.get_encoder_hidden_states_mask(positive=True)
+
+ text_seq_len = prompt_embeds.shape[1]
+ self._ensure_wrappers(text_seq_len, encoder_mask)
+ self._set_text_seq_len(text_seq_len)
+
+ if self.joint_block_wrappers:
+ for wrapper in self.joint_block_wrappers:
+ wrapper.set_encoder_mask(encoder_mask)
+
+ text_embeddings = self.adapter.compute_text_embeddings(t, config, pooled_embeds)
+ image_rotary_embeddings = self.adapter.compute_rotary_embeddings(
+ prompt_embeds,
+ config,
+ encoder_hidden_states_mask=encoder_mask,
+ cond_image_grid=cond_image_grid,
+ kontext_image_ids=kontext_image_ids,
+ )
+
+ prev_patch_latents = [p for p in patch_latents]
+ encoder_hidden_states: mx.array | None = None
+
+ for patch_idx in range(len(patch_latents)):
+ patch = patch_latents[patch_idx]
+
+ if (
+ self.is_first_stage
+ and not self.is_last_stage
+ and not is_first_async_step
+ ):
+ patch = mx.distributed.recv_like(
+ patch, src=self.prev_rank, group=self.group
+ )
+ mx.eval(patch)
+
+ step_patch = mx.concatenate([patch, patch], axis=0) if needs_cfg else patch
+
+ noise, encoder_hidden_states = self._run_single_patch_pass(
+ patch=step_patch,
+ patch_idx=patch_idx,
+ token_indices=token_indices[patch_idx],
+ prompt_embeds=prompt_embeds,
+ text_embeddings=text_embeddings,
+ image_rotary_embeddings=image_rotary_embeddings,
+ encoder_hidden_states=encoder_hidden_states,
+ )
+
+ if self.is_last_stage:
+ assert noise is not None
+ if needs_cfg:
+ noise_pos, noise_neg = mx.split(noise, 2, axis=0)
+ guidance_scale = self._get_effective_guidance_scale()
+ assert guidance_scale is not None
+ noise = self.adapter.apply_guidance(
+ noise_pos, noise_neg, guidance_scale
+ )
+
+ patch_latents[patch_idx] = config.scheduler.step( # pyright: ignore[reportAny]
+ noise=noise,
+ timestep=t,
+ latents=prev_patch_latents[patch_idx],
+ )
+
+ if not self.is_first_stage and t != config.num_inference_steps - 1:
+ patch_latents[patch_idx] = mx.distributed.send(
+ patch_latents[patch_idx], self.next_rank, group=self.group
+ )
+ mx.async_eval(patch_latents[patch_idx])
+
+ return mx.concatenate(patch_latents, axis=1)
+
+ def _run_single_patch_pass(
+ self,
+ patch: mx.array,
+ patch_idx: int,
+ token_indices: tuple[int, int],
+ prompt_embeds: mx.array,
+ text_embeddings: mx.array,
+ image_rotary_embeddings: RotaryEmbeddings,
+ encoder_hidden_states: mx.array | None,
+ ) -> tuple[mx.array | None, mx.array | None]:
+ """Process a single patch through the forward pipeline.
+
+ Handles stage-to-stage communication (stage i -> stage i+1).
+ Ring communication (last stage -> first stage) is handled by the caller.
+
+ Args:
+ patch: The patch latents to process
+ patch_idx: Index of this patch (0-indexed)
+ token_indices: (start_token, end_token) for this patch
+ prompt_embeds: Text embeddings (for compute_embeddings on first stage)
+ text_embeddings: Precomputed text embeddings
+ image_rotary_embeddings: Precomputed rotary embeddings
+ encoder_hidden_states: Encoder hidden states (passed between patches)
+
+ Returns:
+ (noise_prediction, encoder_hidden_states) - noise is None for non-last stages
+ """
+ start_token, end_token = token_indices
+ batch_size = patch.shape[0]
+ text_seq_len = prompt_embeds.shape[1]
+ hidden_dim = self.adapter.hidden_dim
+
+ if self.has_joint_blocks:
+ if not self.is_first_stage:
+ patch_len = patch.shape[1]
+ patch = mx.distributed.recv(
+ (batch_size, patch_len, hidden_dim),
+ patch.dtype,
+ self.prev_rank,
+ group=self.group,
+ )
+ mx.eval(patch)
+
+ if patch_idx == 0:
+ encoder_hidden_states = mx.distributed.recv(
+ (batch_size, text_seq_len, hidden_dim),
+ patch.dtype,
+ self.prev_rank,
+ group=self.group,
+ )
+ mx.eval(encoder_hidden_states)
+
+ if self.is_first_stage:
+ patch, encoder_hidden_states = self.adapter.compute_embeddings(
+ patch, prompt_embeds
+ )
+
+ assert self.joint_block_wrappers is not None
+ assert encoder_hidden_states is not None
+ for wrapper in self.joint_block_wrappers:
+ wrapper.set_patch(BlockWrapperMode.PATCHED, start_token, end_token)
+ encoder_hidden_states, patch = wrapper(
+ hidden_states=patch,
+ encoder_hidden_states=encoder_hidden_states,
+ text_embeddings=text_embeddings,
+ rotary_embeddings=image_rotary_embeddings,
+ )
+
+ if self.owns_concat_stage:
+ assert encoder_hidden_states is not None
+ patch_concat = self.adapter.merge_streams(patch, encoder_hidden_states)
+
+ if self.has_single_blocks or self.is_last_stage:
+ patch = patch_concat
+ else:
+ patch_concat = mx.distributed.send(
+ patch_concat, self.next_rank, group=self.group
+ )
+ mx.async_eval(patch_concat)
+
+ elif self.has_joint_blocks and not self.is_last_stage:
+ patch = mx.distributed.send(patch, self.next_rank, group=self.group)
+ mx.async_eval(patch)
+
+ if patch_idx == 0:
+ assert encoder_hidden_states is not None
+ encoder_hidden_states = mx.distributed.send(
+ encoder_hidden_states, self.next_rank, group=self.group
+ )
+ mx.async_eval(encoder_hidden_states)
+
+ if self.has_single_blocks:
+ if not self.owns_concat_stage and not self.is_first_stage:
+ patch_len = patch.shape[1]
+ patch = mx.distributed.recv(
+ (batch_size, text_seq_len + patch_len, hidden_dim),
+ patch.dtype,
+ self.prev_rank,
+ group=self.group,
+ )
+ mx.eval(patch)
+
+ assert self.single_block_wrappers is not None
+ for wrapper in self.single_block_wrappers:
+ wrapper.set_patch(BlockWrapperMode.PATCHED, start_token, end_token)
+ patch = wrapper(
+ hidden_states=patch,
+ text_embeddings=text_embeddings,
+ rotary_embeddings=image_rotary_embeddings,
+ )
+
+ if not self.is_last_stage:
+ patch = mx.distributed.send(patch, self.next_rank, group=self.group)
+ mx.async_eval(patch)
+
+ noise: mx.array | None = None
+ if self.is_last_stage:
+ patch_img_only = patch[:, text_seq_len:, :]
+ noise = self.adapter.final_projection(patch_img_only, text_embeddings)
+
+ return noise, encoder_hidden_states
diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py
index 2482cc0c..65dba4cb 100644
--- a/src/exo/worker/engines/mlx/generator/generate.py
+++ b/src/exo/worker/engines/mlx/generator/generate.py
@@ -1,7 +1,7 @@
from typing import Any, Callable, Generator, cast, get_args
import mlx.core as mx
-from mlx_lm import stream_generate
+from mlx_lm.generate import stream_generate
from mlx_lm.models.cache import KVCache
from mlx_lm.sample_utils import make_sampler
from mlx_lm.tokenizer_utils import TokenizerWrapper
diff --git a/src/exo/worker/main.py b/src/exo/worker/main.py
index 7f331a15..153a9d67 100644
--- a/src/exo/worker/main.py
+++ b/src/exo/worker/main.py
@@ -9,13 +9,15 @@ from loguru import logger
from exo.routing.connection_message import ConnectionMessage, ConnectionMessageType
from exo.shared.apply import apply
from exo.shared.models.model_cards import ModelId
+from exo.shared.types.api import ImageEditsInternalParams
from exo.shared.types.commands import ForwarderCommand, RequestEventLog
-from exo.shared.types.common import NodeId, SessionId
+from exo.shared.types.common import CommandId, NodeId, SessionId
from exo.shared.types.events import (
Event,
EventId,
ForwarderEvent,
IndexedEvent,
+ InputChunkReceived,
NodeDownloadProgress,
NodeGatheredInfo,
TaskCreated,
@@ -28,6 +30,7 @@ from exo.shared.types.state import State
from exo.shared.types.tasks import (
CreateRunner,
DownloadModel,
+ ImageEdits,
Shutdown,
Task,
TaskStatus,
@@ -93,6 +96,10 @@ class Worker:
self.event_sender, self.event_receiver = channel[Event]()
+ # Buffer for input image chunks (for image editing)
+ self.input_chunk_buffer: dict[CommandId, dict[int, str]] = {}
+ self.input_chunk_counts: dict[CommandId, int] = {}
+
async def run(self):
logger.info("Starting Worker")
@@ -157,6 +164,17 @@ class Worker:
for idx, event in indexed_events:
self.state = apply(self.state, IndexedEvent(idx=idx, event=event))
+ # Buffer input image chunks for image editing
+ if isinstance(event, InputChunkReceived):
+ cmd_id = event.command_id
+ if cmd_id not in self.input_chunk_buffer:
+ self.input_chunk_buffer[cmd_id] = {}
+ self.input_chunk_counts[cmd_id] = event.chunk.total_chunks
+
+ self.input_chunk_buffer[cmd_id][event.chunk.chunk_index] = (
+ event.chunk.data
+ )
+
async def plan_step(self):
while True:
await anyio.sleep(0.1)
@@ -169,6 +187,8 @@ class Worker:
self.state.instances,
self.state.runners,
self.state.tasks,
+ self.input_chunk_buffer,
+ self.input_chunk_counts,
)
if task is None:
continue
@@ -232,6 +252,46 @@ class Worker:
task_id=task.task_id, task_status=TaskStatus.TimedOut
)
)
+ case ImageEdits() if task.task_params.total_input_chunks > 0:
+ # Assemble image from chunks and inject into task
+ cmd_id = task.command_id
+ chunks = self.input_chunk_buffer.get(cmd_id, {})
+ assembled = "".join(chunks[i] for i in range(len(chunks)))
+ logger.info(
+ f"Assembled input image from {len(chunks)} chunks, "
+ f"total size: {len(assembled)} bytes"
+ )
+ # Create modified task with assembled image data
+ modified_task = ImageEdits(
+ task_id=task.task_id,
+ command_id=task.command_id,
+ instance_id=task.instance_id,
+ task_status=task.task_status,
+ task_params=ImageEditsInternalParams(
+ image_data=assembled,
+ total_input_chunks=task.task_params.total_input_chunks,
+ prompt=task.task_params.prompt,
+ model=task.task_params.model,
+ n=task.task_params.n,
+ quality=task.task_params.quality,
+ output_format=task.task_params.output_format,
+ response_format=task.task_params.response_format,
+ size=task.task_params.size,
+ image_strength=task.task_params.image_strength,
+ bench=task.task_params.bench,
+ stream=task.task_params.stream,
+ partial_images=task.task_params.partial_images,
+ advanced_params=task.task_params.advanced_params,
+ ),
+ )
+ # Cleanup buffers
+ if cmd_id in self.input_chunk_buffer:
+ del self.input_chunk_buffer[cmd_id]
+ if cmd_id in self.input_chunk_counts:
+ del self.input_chunk_counts[cmd_id]
+ await self.runners[self._task_to_runner_id(task)].start_task(
+ modified_task
+ )
case task:
await self.runners[self._task_to_runner_id(task)].start_task(task)
diff --git a/src/exo/worker/plan.py b/src/exo/worker/plan.py
index 54ca1106..9f8f7d52 100644
--- a/src/exo/worker/plan.py
+++ b/src/exo/worker/plan.py
@@ -3,12 +3,14 @@
from collections.abc import Mapping, Sequence
from exo.shared.models.model_cards import ModelId
-from exo.shared.types.common import NodeId
+from exo.shared.types.common import CommandId, NodeId
from exo.shared.types.tasks import (
ChatCompletion,
ConnectToGroup,
CreateRunner,
DownloadModel,
+ ImageEdits,
+ ImageGeneration,
LoadModel,
Shutdown,
StartWarmup,
@@ -49,6 +51,8 @@ def plan(
instances: Mapping[InstanceId, Instance],
all_runners: Mapping[RunnerId, RunnerStatus], # all global
tasks: Mapping[TaskId, Task],
+ input_chunk_buffer: Mapping[CommandId, dict[int, str]] | None = None,
+ input_chunk_counts: Mapping[CommandId, int] | None = None,
) -> Task | None:
# Python short circuiting OR logic should evaluate these sequentially.
return (
@@ -58,7 +62,7 @@ def plan(
or _init_distributed_backend(runners, all_runners)
or _load_model(runners, all_runners, global_download_status)
or _ready_to_warmup(runners, all_runners)
- or _pending_tasks(runners, tasks, all_runners)
+ or _pending_tasks(runners, tasks, all_runners, input_chunk_buffer)
)
@@ -262,14 +266,24 @@ def _pending_tasks(
runners: Mapping[RunnerId, RunnerSupervisor],
tasks: Mapping[TaskId, Task],
all_runners: Mapping[RunnerId, RunnerStatus],
+ input_chunk_buffer: Mapping[CommandId, dict[int, str]] | None = None,
) -> Task | None:
for task in tasks.values():
# for now, just forward chat completions
- if not isinstance(task, ChatCompletion):
+ # TODO(ciaran): do this better!
+ if not isinstance(task, (ChatCompletion, ImageGeneration, ImageEdits)):
continue
if task.task_status not in (TaskStatus.Pending, TaskStatus.Running):
continue
+ # For ImageEdits tasks, verify all input chunks have been received
+ if isinstance(task, ImageEdits) and task.task_params.total_input_chunks > 0:
+ cmd_id = task.command_id
+ expected = task.task_params.total_input_chunks
+ received = len((input_chunk_buffer or {}).get(cmd_id, {}))
+ if received < expected:
+ continue # Wait for all chunks to arrive
+
for runner in runners.values():
if task.instance_id != runner.bound_instance.instance.instance_id:
continue
diff --git a/src/exo/worker/runner/runner.py b/src/exo/worker/runner/runner.py
index dead1191..9628b8d3 100644
--- a/src/exo/worker/runner/runner.py
+++ b/src/exo/worker/runner/runner.py
@@ -1,6 +1,8 @@
+import base64
import time
from collections.abc import Generator
from functools import cache
+from typing import Literal
import mlx.core as mx
from mlx_lm.models.gpt_oss import Model as GptOssModel
@@ -12,8 +14,11 @@ from openai_harmony import ( # pyright: ignore[reportMissingTypeStubs]
load_harmony_encoding,
)
-from exo.shared.types.api import ChatCompletionMessageText
-from exo.shared.types.chunks import TokenChunk
+from exo.shared.constants import EXO_MAX_CHUNK_SIZE
+from exo.shared.models.model_cards import ModelId, ModelTask
+from exo.shared.types.api import ChatCompletionMessageText, ImageGenerationStats
+from exo.shared.types.chunks import ImageChunk, TokenChunk
+from exo.shared.types.common import CommandId
from exo.shared.types.events import (
ChunkGenerated,
Event,
@@ -24,6 +29,8 @@ from exo.shared.types.events import (
from exo.shared.types.tasks import (
ChatCompletion,
ConnectToGroup,
+ ImageEdits,
+ ImageGeneration,
LoadModel,
Shutdown,
StartWarmup,
@@ -33,6 +40,8 @@ from exo.shared.types.tasks import (
from exo.shared.types.worker.instances import BoundInstance
from exo.shared.types.worker.runner_response import (
GenerationResponse,
+ ImageGenerationResponse,
+ PartialImageResponse,
)
from exo.shared.types.worker.runners import (
RunnerConnected,
@@ -48,7 +57,15 @@ from exo.shared.types.worker.runners import (
RunnerStatus,
RunnerWarmingUp,
)
+from exo.shared.types.worker.shards import ShardMetadata
from exo.utils.channels import MpReceiver, MpSender
+from exo.worker.engines.image import (
+ DistributedImageModel,
+ generate_image,
+ initialize_image_model,
+ warmup_image_generator,
+)
+from exo.worker.engines.mlx import Model
from exo.worker.engines.mlx.generator.generate import mlx_generate, warmup_inference
from exo.worker.engines.mlx.utils_mlx import (
apply_chat_template,
@@ -79,7 +96,7 @@ def main(
setup_start_time = time.time()
- model = None
+ model: Model | DistributedImageModel | None = None
tokenizer = None
group = None
@@ -133,15 +150,25 @@ def main(
)
time.sleep(0.5)
- model, tokenizer = load_mlx_items(
- bound_instance, group, on_timeout=on_model_load_timeout
- )
+ if ModelTask.TextGeneration in shard_metadata.model_card.tasks:
+ model, tokenizer = load_mlx_items(
+ bound_instance, group, on_timeout=on_model_load_timeout
+ )
+ elif (
+ ModelTask.TextToImage in shard_metadata.model_card.tasks
+ or ModelTask.ImageToImage in shard_metadata.model_card.tasks
+ ):
+ model = initialize_image_model(bound_instance)
+ else:
+ raise ValueError(
+ f"Unknown model task(s): {shard_metadata.model_card.tasks}"
+ )
current_status = RunnerLoaded()
logger.info("runner loaded")
case StartWarmup() if isinstance(current_status, RunnerLoaded):
assert model
- assert tokenizer
+
current_status = RunnerWarmingUp()
logger.info("runner warming up")
event_sender.send(
@@ -151,15 +178,30 @@ def main(
)
logger.info(f"warming up inference for instance: {instance}")
- toks = warmup_inference(
- model=model,
- tokenizer=tokenizer,
- # kv_prefix_cache=kv_prefix_cache, # supply for warmup-time prefix caching
- )
- logger.info(f"warmed up by generating {toks} tokens")
- logger.info(
- f"runner initialized in {time.time() - setup_start_time} seconds"
- )
+ if ModelTask.TextGeneration in shard_metadata.model_card.tasks:
+ assert not isinstance(model, DistributedImageModel)
+ assert tokenizer
+
+ toks = warmup_inference(
+ model=model,
+ tokenizer=tokenizer,
+ # kv_prefix_cache=kv_prefix_cache, # supply for warmup-time prefix caching
+ )
+ logger.info(f"warmed up by generating {toks} tokens")
+ logger.info(
+ f"runner initialized in {time.time() - setup_start_time} seconds"
+ )
+ elif (
+ ModelTask.TextToImage in shard_metadata.model_card.tasks
+ or ModelTask.ImageToImage in shard_metadata.model_card.tasks
+ ):
+ assert isinstance(model, DistributedImageModel)
+ image = warmup_image_generator(model=model)
+ if image is not None:
+ logger.info(f"warmed up by generating {image.size} image")
+ else:
+ logger.info("warmup completed (non-primary node)")
+
current_status = RunnerReady()
logger.info("runner ready")
case ChatCompletion(task_params=task_params, command_id=command_id) if (
@@ -173,7 +215,7 @@ def main(
runner_id=runner_id, runner_status=current_status
)
)
- assert model
+ assert model and not isinstance(model, DistributedImageModel)
assert tokenizer
assert task_params.messages[0].content is not None
@@ -240,6 +282,134 @@ def main(
)
raise
+ current_status = RunnerReady()
+ logger.info("runner ready")
+ case ImageGeneration(
+ task_params=task_params, command_id=command_id
+ ) if isinstance(current_status, RunnerReady):
+ assert isinstance(model, DistributedImageModel)
+ logger.info(f"received image generation request: {str(task)[:500]}")
+ current_status = RunnerRunning()
+ logger.info("runner running")
+ event_sender.send(
+ RunnerStatusUpdated(
+ runner_id=runner_id, runner_status=current_status
+ )
+ )
+
+ try:
+ # Generate images using the image generation backend
+ # Track image_index for final images only
+ image_index = 0
+ for response in generate_image(model=model, task=task_params):
+ if (
+ shard_metadata.device_rank
+ == shard_metadata.world_size - 1
+ ):
+ match response:
+ case PartialImageResponse():
+ logger.info(
+ f"sending partial ImageChunk {response.partial_index}/{response.total_partials}"
+ )
+ _process_image_response(
+ response,
+ command_id,
+ shard_metadata,
+ event_sender,
+ image_index,
+ )
+ case ImageGenerationResponse():
+ logger.info("sending final ImageChunk")
+ _process_image_response(
+ response,
+ command_id,
+ shard_metadata,
+ event_sender,
+ image_index,
+ )
+ image_index += 1
+ except Exception as e:
+ if shard_metadata.device_rank == shard_metadata.world_size - 1:
+ event_sender.send(
+ ChunkGenerated(
+ command_id=command_id,
+ chunk=ImageChunk(
+ idx=0,
+ model=shard_metadata.model_card.model_id,
+ data="",
+ chunk_index=0,
+ total_chunks=1,
+ image_index=0,
+ finish_reason="error",
+ error_message=str(e),
+ ),
+ )
+ )
+ raise
+
+ current_status = RunnerReady()
+ logger.info("runner ready")
+ case ImageEdits(task_params=task_params, command_id=command_id) if (
+ isinstance(current_status, RunnerReady)
+ ):
+ assert isinstance(model, DistributedImageModel)
+ logger.info(f"received image edits request: {str(task)[:500]}")
+ current_status = RunnerRunning()
+ logger.info("runner running")
+ event_sender.send(
+ RunnerStatusUpdated(
+ runner_id=runner_id, runner_status=current_status
+ )
+ )
+
+ try:
+ image_index = 0
+ for response in generate_image(model=model, task=task_params):
+ if (
+ shard_metadata.device_rank
+ == shard_metadata.world_size - 1
+ ):
+ match response:
+ case PartialImageResponse():
+ logger.info(
+ f"sending partial ImageChunk {response.partial_index}/{response.total_partials}"
+ )
+ _process_image_response(
+ response,
+ command_id,
+ shard_metadata,
+ event_sender,
+ image_index,
+ )
+ case ImageGenerationResponse():
+ logger.info("sending final ImageChunk")
+ _process_image_response(
+ response,
+ command_id,
+ shard_metadata,
+ event_sender,
+ image_index,
+ )
+ image_index += 1
+ except Exception as e:
+ if shard_metadata.device_rank == shard_metadata.world_size - 1:
+ event_sender.send(
+ ChunkGenerated(
+ command_id=command_id,
+ chunk=ImageChunk(
+ idx=0,
+ model=shard_metadata.model_card.model_id,
+ data="",
+ chunk_index=0,
+ total_chunks=1,
+ image_index=0,
+ finish_reason="error",
+ error_message=str(e),
+ ),
+ )
+ )
+ raise
+
current_status = RunnerReady()
logger.info("runner ready")
case Shutdown():
@@ -329,6 +499,75 @@ def parse_thinking_models(
yield response
+def _send_image_chunk(
+ encoded_data: str,
+ command_id: CommandId,
+ model_id: ModelId,
+ event_sender: MpSender[Event],
+ image_index: int,
+ is_partial: bool,
+ partial_index: int | None = None,
+ total_partials: int | None = None,
+ stats: ImageGenerationStats | None = None,
+ image_format: Literal["png", "jpeg", "webp"] | None = None,
+) -> None:
+ """Send base64-encoded image data as chunks via events."""
+ data_chunks = [
+ encoded_data[i : i + EXO_MAX_CHUNK_SIZE]
+ for i in range(0, len(encoded_data), EXO_MAX_CHUNK_SIZE)
+ ]
+ total_chunks = len(data_chunks)
+ for chunk_index, chunk_data in enumerate(data_chunks):
+ # Only include stats on the last chunk of the final image
+ chunk_stats = (
+ stats if chunk_index == total_chunks - 1 and not is_partial else None
+ )
+ event_sender.send(
+ ChunkGenerated(
+ command_id=command_id,
+ chunk=ImageChunk(
+ idx=chunk_index,
+ model=model_id,
+ data=chunk_data,
+ chunk_index=chunk_index,
+ total_chunks=total_chunks,
+ image_index=image_index,
+ is_partial=is_partial,
+ partial_index=partial_index,
+ total_partials=total_partials,
+ stats=chunk_stats,
+ format=image_format,
+ ),
+ )
+ )
+
+
+def _process_image_response(
+ response: ImageGenerationResponse | PartialImageResponse,
+ command_id: CommandId,
+ shard_metadata: ShardMetadata,
+ event_sender: MpSender[Event],
+ image_index: int,
+) -> None:
+ """Process a single image response and send chunks."""
+ encoded_data = base64.b64encode(response.image_data).decode("utf-8")
+ is_partial = isinstance(response, PartialImageResponse)
+ # Extract stats from final ImageGenerationResponse if available
+ stats = response.stats if isinstance(response, ImageGenerationResponse) else None
+ _send_image_chunk(
+ encoded_data=encoded_data,
+ command_id=command_id,
+ model_id=shard_metadata.model_card.model_id,
+ event_sender=event_sender,
+ image_index=response.partial_index if is_partial else image_index,
+ is_partial=is_partial,
+ partial_index=response.partial_index if is_partial else None,
+ total_partials=response.total_partials if is_partial else None,
+ stats=stats,
+ image_format=response.format,
+ )
+
+
EXO_RUNNER_MUST_FAIL = "EXO RUNNER MUST FAIL"
EXO_RUNNER_MUST_OOM = "EXO RUNNER MUST OOM"
EXO_RUNNER_MUST_TIMEOUT = "EXO RUNNER MUST TIMEOUT"
diff --git a/src/exo/worker/tests/unittests/conftest.py b/src/exo/worker/tests/unittests/conftest.py
index 67b19f3c..ad0a4b97 100644
--- a/src/exo/worker/tests/unittests/conftest.py
+++ b/src/exo/worker/tests/unittests/conftest.py
@@ -1,6 +1,6 @@
from dataclasses import dataclass, field
-from exo.shared.models.model_cards import ModelCard, ModelId
+from exo.shared.models.model_cards import ModelCard, ModelId, ModelTask
from exo.shared.types.common import NodeId
from exo.shared.types.memory import Memory
from exo.shared.types.tasks import BaseTask, TaskId
@@ -38,6 +38,7 @@ def get_pipeline_shard_metadata(
n_layers=32,
hidden_size=2048,
supports_tensor=False,
+ tasks=[ModelTask.TextGeneration],
),
device_rank=device_rank,
world_size=world_size,
diff --git a/src/exo/worker/tests/unittests/test_mlx/conftest.py b/src/exo/worker/tests/unittests/test_mlx/conftest.py
index c267ba6b..87ecce7d 100644
--- a/src/exo/worker/tests/unittests/test_mlx/conftest.py
+++ b/src/exo/worker/tests/unittests/test_mlx/conftest.py
@@ -11,7 +11,7 @@ import mlx.core as mx
import mlx.nn as nn
from exo.shared.constants import EXO_MODELS_DIR
-from exo.shared.models.model_cards import ModelCard
+from exo.shared.models.model_cards import ModelCard, ModelTask
from exo.shared.types.api import ChatCompletionMessage
from exo.shared.types.common import ModelId
from exo.shared.types.memory import Memory
@@ -88,6 +88,7 @@ def run_gpt_oss_pipeline_device(
n_layers=24,
hidden_size=2880,
supports_tensor=False,
+ tasks=[ModelTask.TextGeneration],
),
device_rank=rank,
world_size=world_size,
@@ -157,6 +158,7 @@ def run_gpt_oss_tensor_parallel_device(
n_layers=24,
hidden_size=2880,
supports_tensor=True,
+ tasks=[ModelTask.TextGeneration],
),
device_rank=rank,
world_size=world_size,
← a31b6ee0 Import download utils once all modules are loaded (#1238)
·
back to Exo
·
Fix tests broken in recent commits (#1239) a354aaa3 →