From 24d99ada4ae02989e30429514edbb4e4a135ff80 Mon Sep 17 00:00:00 2001
From: Sami Khan <smsak99@gmail.com>
Date: Wed, 7 Jan 2026 03:44:50 +0500
Subject: [PATCH] image gen in dashboard

---
 dashboard/src/lib/components/ChatForm.svelte  | 61 ++++++++++---
 .../src/lib/components/ChatMessages.svelte    | 54 +++++++++++-
 dashboard/src/lib/stores/app.svelte.ts        | 87 ++++++++++++++++++-
 dashboard/src/routes/+page.svelte             | 52 ++++++++++-
 src/exo/master/api.py                         |  1 +
 src/exo/shared/types/api.py                   |  1 +
 6 files changed, 238 insertions(+), 18 deletions(-)

diff --git a/dashboard/src/lib/components/ChatForm.svelte b/dashboard/src/lib/components/ChatForm.svelte
index 3f622602..2b19bdb0 100644
--- a/dashboard/src/lib/components/ChatForm.svelte
+++ b/dashboard/src/lib/components/ChatForm.svelte
@@ -1,5 +1,5 @@
 <script lang="ts">
-	import { isLoading, sendMessage, selectedChatModel, setSelectedChatModel, instances, ttftMs, tps, totalTokens } from '$lib/stores/app.svelte';
+	import { isLoading, sendMessage, generateImage, 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';
@@ -10,6 +10,7 @@
 		showHelperText?: boolean;
 		autofocus?: boolean;
 		showModelSelector?: boolean;
+		modelTasks?: Record<string, string[]>;
 	}
 
 	let { 
@@ -17,7 +18,8 @@
 		placeholder = 'Ask anything',
 		showHelperText = false,
 		autofocus = true,
-		showModelSelector = false
+		showModelSelector = false,
+		modelTasks = {}
 	}: Props = $props();
 
 	let message = $state('');
@@ -48,13 +50,29 @@
 	// Accept all supported file types
 	const acceptString = getAcceptString(['image', 'text', 'pdf']);
 
+	// Check if a model supports image generation
+	function modelSupportsImageGeneration(modelId: string): boolean {
+		const tasks = modelTasks[modelId] || [];
+		return tasks.includes('TextToImage') || tasks.includes('ImageToImage');
+	}
+
+	// Check if the currently selected model supports image generation
+	const isImageModel = $derived(() => {
+		if (!currentModel) return false;
+		return modelSupportsImageGeneration(currentModel);
+	});
+
 	// 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 });
+				models.push({ 
+					id: modelId, 
+					label: modelId.split('/').pop() || modelId,
+					isImageModel: modelSupportsImageGeneration(modelId)
+				});
 			}
 		}
 		return models;
@@ -160,7 +178,12 @@
 		uploadedFiles = [];
 		resetTextareaHeight();
 		
-		sendMessage(content, files);
+		// Use image generation for image models
+		if (isImageModel() && content) {
+			generateImage(content);
+		} else {
+			sendMessage(content, files);
+		}
 		
 		// Refocus the textarea after sending
 		setTimeout(() => textareaRef?.focus(), 10);
@@ -297,7 +320,14 @@
 										{:else}
 											<span class="w-3"></span>
 										{/if}
-										<span class="truncate">{model.label}</span>
+										<span class="truncate flex-1">{model.label}</span>
+										{#if model.isImageModel}
+											<svg class="w-3.5 h-3.5 flex-shrink-0 text-purple-400" 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}
 									</button>
 								{/each}
 							</div>
@@ -357,7 +387,7 @@
 				onkeydown={handleKeydown}
 				oninput={handleInput}
 				onpaste={handlePaste}
-				{placeholder}
+				placeholder={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"
@@ -370,15 +400,26 @@
 				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"
+						: isImageModel() 
+							? 'bg-gradient-to-r from-purple-500 to-pink-500 text-white hover:from-purple-600 hover:to-pink-600 hover:shadow-[0_0_20px_rgba(168,85,247,0.4)]'
+							: 'bg-exo-yellow text-exo-black hover:bg-exo-yellow-darker hover:shadow-[0_0_20px_rgba(255,215,0,0.3)]'}"
+				aria-label={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="hidden sm:inline">{isImageModel() ? 'GENERATING' : 'PROCESSING'}</span>
 						<span class="sm:hidden">...</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}
diff --git a/dashboard/src/lib/components/ChatMessages.svelte b/dashboard/src/lib/components/ChatMessages.svelte
index 2bbd09f8..31adb257 100644
--- a/dashboard/src/lib/components/ChatMessages.svelte
+++ b/dashboard/src/lib/components/ChatMessages.svelte
@@ -365,10 +365,58 @@ 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"
+												/>
+												<!-- Download button overlay -->
+												<button
+													type="button"
+													class="absolute top-2 right-2 p-2 rounded-lg bg-exo-dark-gray/80 border border-exo-yellow/30 text-exo-yellow opacity-0 group-hover/img:opacity-100 transition-opacity 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;
+															link.download = `generated-image-${Date.now()}.png`;
+															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>
+										{/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...'}
+										<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">Generating image...</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/stores/app.svelte.ts b/dashboard/src/lib/stores/app.svelte.ts
index e0bfbad1..d9ef9aff 100644
--- a/dashboard/src/lib/stores/app.svelte.ts
+++ b/dashboard/src/lib/stores/app.svelte.ts
@@ -163,7 +163,7 @@ interface RawStateResponse {
 }
 
 export interface MessageAttachment {
-	type: 'image' | 'text' | 'file';
+	type: 'image' | 'text' | 'file' | 'generated-image';
 	name: string;
 	content?: string;
 	preview?: string;
@@ -1413,6 +1413,90 @@ 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.');
+			}
+			
+			const response = await fetch('/v1/images/generations', {
+				method: 'POST',
+				headers: {
+					'Content-Type': 'application/json'
+				},
+				body: JSON.stringify({
+					model,
+					prompt,
+					quality: 'medium',
+					size: '1024x1024',
+					response_format: 'b64_json'
+				})
+			});
+
+			if (!response.ok) {
+				const errorText = await response.text();
+				throw new Error(`API error: ${response.status} - ${errorText}`);
+			}
+
+			const data = await response.json();
+			const imageData = data.data?.[0]?.b64_json;
+			
+			if (!imageData) {
+				throw new Error('No image data received from the API');
+			}
+			
+			// Update the assistant message with the generated image
+			const idx = this.messages.findIndex(m => m.id === assistantMessage.id);
+			if (idx !== -1) {
+				this.messages[idx].content = '';
+				this.messages[idx].attachments = [{
+					type: 'generated-image',
+					name: 'generated-image.png',
+					preview: `data:image/png;base64,${imageData}`,
+					mimeType: 'image/png'
+				}];
+			}
+			
+		} 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();
+		}
+	}
+
 	/**
 	 * Clear current chat and go back to welcome state
 	 */
@@ -1463,6 +1547,7 @@ export const chatSidebarVisible = () => appStore.getChatSidebarVisible();
 // Actions
 export const startChat = () => appStore.startChat();
 export const sendMessage = (content: string, files?: { id: string; name: string; type: string; textContent?: string; preview?: string }[]) => appStore.sendMessage(content, files);
+export const generateImage = (prompt: string, modelId?: string) => appStore.generateImage(prompt, modelId);
 export const clearChat = () => appStore.clearChat();
 export const setSelectedChatModel = (modelId: string) => appStore.setSelectedModel(modelId);
 export const selectPreviewModel = (modelId: string | null) => appStore.selectPreviewModel(modelId);
diff --git a/dashboard/src/routes/+page.svelte b/dashboard/src/routes/+page.svelte
index abc603ff..a2731313 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';
 	
@@ -1250,6 +1273,7 @@ function toggleInstanceDownloadDetails(nodeId: string): void {
 								placeholder="Ask anything" 
 								showHelperText={false}
 								showModelSelector={true}
+								modelTasks={modelTasks()}
 							/>
 						</div>
 					</div>
@@ -1471,8 +1495,18 @@ function toggleInstanceDownloadDetails(nodeId: string): void {
 									{@const foundModel = models.find(m => m.id === selectedModelId)}
 									{#if foundModel}
 										{@const sizeGB = getModelSizeGB(foundModel)}
+										{@const isImageModel = modelSupportsImageGeneration(foundModel.id)}
 										<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>
+											<span class="flex items-center gap-2 text-exo-light-gray truncate">
+												<span class="truncate">{foundModel.name || foundModel.id}</span>
+												{#if isImageModel}
+													<svg class="w-4 h-4 flex-shrink-0 text-purple-400" 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>
 											<span class="text-white/50 text-xs flex-shrink-0">{sizeGB >= 1 ? sizeGB.toFixed(0) : sizeGB.toFixed(1)}GB</span>
 										</span>
 									{:else}
@@ -1517,6 +1551,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={() => {
@@ -1536,7 +1571,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">
+													<span class="truncate">{model.name || model.id}</span>
+													{#if isImageModel}
+														<svg class="w-4 h-4 flex-shrink-0 text-purple-400" 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>
 												<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>
@@ -1733,7 +1777,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/src/exo/master/api.py b/src/exo/master/api.py
index 571496db..e31d573f 100644
--- a/src/exo/master/api.py
+++ b/src/exo/master/api.py
@@ -649,6 +649,7 @@ class API:
                     tags=card.tags,
                     storage_size_megabytes=int(card.metadata.storage_size.in_mb),
                     supports_tensor=card.metadata.supports_tensor,
+                    tasks=[task.value for task in card.tasks],
                 )
                 for card in MODEL_CARDS.values()
             ]
diff --git a/src/exo/shared/types/api.py b/src/exo/shared/types/api.py
index 201dfc43..c34cbb5f 100644
--- a/src/exo/shared/types/api.py
+++ b/src/exo/shared/types/api.py
@@ -28,6 +28,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):