From c62b6817c407f57f0182c75dbf0e5152a7f0a689 Mon Sep 17 00:00:00 2001
From: koko210Serve <koko.yordanov@proton.me>
Date: Sat, 13 Dec 2025 00:36:35 +0200
Subject: [PATCH] Fix image generation UI: add image preview, serving endpoint,
 and proper error handling

- Fixed function name mismatch: generateImage() -> generateManualImage()
- Fixed status div ID mismatch in HTML
- Added /image/view/{filename} endpoint to serve generated images from ComfyUI output
- Implemented proper image preview with DOM element creation instead of innerHTML
- Added robust error handling with onload/onerror event handlers
- Added debug logging to image serving endpoint for troubleshooting
- Images now display directly in the Web UI after generation
---
 bot/api.py            | 43 +++++++++++++++++++++++++++++++++
 bot/static/index.html | 55 +++++++++++++++++++++++++++++++++++++++----
 2 files changed, 93 insertions(+), 5 deletions(-)

diff --git a/bot/api.py b/bot/api.py
index 3ad94f7..00b5937 100644
--- a/bot/api.py
+++ b/bot/api.py
@@ -990,6 +990,49 @@ async def test_image_detection(req: dict):
     except Exception as e:
         return {"status": "error", "message": f"Error: {e}"}
 
+@app.get("/image/view/{filename}")
+async def view_generated_image(filename: str):
+    """Serve generated images from ComfyUI output directory"""
+    try:
+        print(f"🖼️ Image view request for: {filename}")
+        
+        # Try multiple possible paths for ComfyUI output
+        possible_paths = [
+            f"/app/ComfyUI/output/{filename}",
+            f"/home/koko210Serve/ComfyUI/output/{filename}",
+            f"./ComfyUI/output/{filename}",
+        ]
+        
+        image_path = None
+        for path in possible_paths:
+            if os.path.exists(path):
+                image_path = path
+                print(f"✅ Found image at: {path}")
+                break
+            else:
+                print(f"❌ Not found at: {path}")
+        
+        if not image_path:
+            print(f"❌ Image not found anywhere: {filename}")
+            return {"status": "error", "message": f"Image not found: {filename}"}
+        
+        # Determine content type based on file extension
+        ext = filename.lower().split('.')[-1]
+        content_type = "image/png"
+        if ext == "jpg" or ext == "jpeg":
+            content_type = "image/jpeg"
+        elif ext == "gif":
+            content_type = "image/gif"
+        elif ext == "webp":
+            content_type = "image/webp"
+        
+        print(f"📤 Serving image: {image_path} as {content_type}")
+        return FileResponse(image_path, media_type=content_type)
+        
+    except Exception as e:
+        print(f"❌ Error serving image: {e}")
+        return {"status": "error", "message": f"Error serving image: {e}"}
+
 @app.post("/servers/{guild_id}/autonomous/tweet")
 async def trigger_autonomous_tweet_for_server(guild_id: int):
     """Trigger autonomous tweet sharing for a specific server"""
diff --git a/bot/static/index.html b/bot/static/index.html
index d1e8929..e43099b 100644
--- a/bot/static/index.html
+++ b/bot/static/index.html
@@ -991,8 +991,9 @@
             <label for="manual-image-prompt">Image Prompt:</label>
             <textarea id="manual-image-prompt" placeholder="Describe the image you want to generate..." rows="3" style="width: 100%; margin-top: 0.5rem;"></textarea>
           </div>
-          <button onclick="generateImage()" style="margin-right: 0.5rem;">🎨 Generate Image</button>
-          <div id="manual-generation-results" style="margin-top: 0.5rem; font-size: 0.9rem;"></div>
+          <button onclick="generateManualImage()" style="margin-right: 0.5rem;">🎨 Generate Image</button>
+          <div id="manual-image-status" style="margin-top: 0.5rem; font-size: 0.9rem;"></div>
+          <div id="manual-image-preview" style="margin-top: 1rem; text-align: center;"></div>
         </div>
         
         <!-- System Information -->
@@ -2562,6 +2563,7 @@ ${result.is_image_request ? `<br><strong>Extracted Prompt:</strong> "${result.ex
 async function generateManualImage() {
   const prompt = document.getElementById('manual-image-prompt').value.trim();
   const statusDiv = document.getElementById('manual-image-status');
+  const previewDiv = document.getElementById('manual-image-preview');
   
   if (!prompt) {
     statusDiv.innerHTML = '❌ Please enter an image prompt';
@@ -2570,6 +2572,9 @@ async function generateManualImage() {
   }
   
   try {
+    // Clear previous preview
+    previewDiv.innerHTML = '';
+    
     statusDiv.innerHTML = '🎨 Generating image... This may take a few minutes.';
     statusDiv.style.color = '#4CAF50';
     
@@ -2581,12 +2586,52 @@ async function generateManualImage() {
     
     const result = await response.json();
     
-    if (response.ok) {
-      statusDiv.innerHTML = `✅ Image generated successfully! Path: ${result.image_path}`;
+    if (response.ok && result.status === 'ok') {
+      statusDiv.innerHTML = `✅ Image generated successfully!`;
       statusDiv.style.color = '#4CAF50';
+      
+      // Display the generated image
+      if (result.image_path) {
+        const filename = result.image_path.split('/').pop();
+        const imageUrl = `/image/view/${encodeURIComponent(filename)}`;
+        
+        // Create image element with better error handling
+        const imgContainer = document.createElement('div');
+        imgContainer.style.cssText = 'background: #1e1e1e; padding: 1rem; border-radius: 8px; border: 1px solid #333;';
+        
+        const img = document.createElement('img');
+        img.src = imageUrl;
+        img.alt = 'Generated Image';
+        img.style.cssText = 'max-width: 100%; max-height: 600px; border-radius: 4px; display: block; margin: 0 auto;';
+        
+        img.onload = function() {
+          console.log('Image loaded successfully:', imageUrl);
+        };
+        
+        img.onerror = function() {
+          console.error('Failed to load image:', imageUrl);
+          imgContainer.innerHTML = `
+            <div style="color: #f44336; padding: 1rem; text-align: center;">
+              ❌ Failed to load image<br>
+              <span style="font-size: 0.85rem;">Path: ${result.image_path}</span><br>
+              <span style="font-size: 0.85rem;">URL: ${imageUrl}</span>
+            </div>
+          `;
+        };
+        
+        imgContainer.appendChild(img);
+        
+        const pathInfo = document.createElement('div');
+        pathInfo.style.cssText = 'margin-top: 0.5rem; color: #aaa; font-size: 0.85rem; text-align: center;';
+        pathInfo.innerHTML = `<strong>File:</strong> ${filename}`;
+        imgContainer.appendChild(pathInfo);
+        
+        previewDiv.appendChild(imgContainer);
+      }
+      
       document.getElementById('manual-image-prompt').value = '';
     } else {
-      statusDiv.innerHTML = `❌ Failed to generate image: ${result.message}`;
+      statusDiv.innerHTML = `❌ Failed to generate image: ${result.message || 'Unknown error'}`;
       statusDiv.style.color = 'red';
     }
   } catch (error) {