Update app.py

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 from diffusers import DDPMPipeline
 import torch
-from PIL import Image
+from PIL import Image, ImageDraw
 import os
 
 # --- CONFIGURATION ---
@@ -11,6 +11,7 @@ DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 print(f"🍋 Initializing Stable-Lime Protocol on {DEVICE}...")
 
 # --- LOAD MODEL ---
+pipeline = None
 try:
     # Load the pipeline directly from your Hub repo
     pipeline = DDPMPipeline.from_pretrained(MODEL_ID)
@@ -18,26 +19,40 @@ try:
     print("✅ Lime Status: ONLINE")
 except Exception as e:
     print(f"❌ CRITICAL FAILURE: Could not load the Lime. Error: {e}")
-    # Fallback to avoid crashing the space immediately, allows debugging
     pipeline = None
 
+def create_error_image(message):
+    """Generates a fallback image containing the error message to debug API issues."""
+    img = Image.new('RGB', (512, 512), color=(20, 0, 0))
+    d = ImageDraw.Draw(img)
+    try:
+        # Basic text drawing if font loading fails
+        d.text((20, 250), f"SYSTEM FAILURE:\n{message}", fill=(255, 50, 50))
+    except:
+        pass
+    return img
+
 def generate_lime():
     """
     The core inference function.
     Summons a lime from the latent void.
     """
     if pipeline is None:
-        return None
+        # Return a generated error image so the frontend has something to show
+        # instead of a broken link/null response
+        return create_error_image("MODEL_LOAD_FAIL")
     
     print("🍋 Generating new specimen...")
-    # Generate the image
-    # num_inference_steps=50 is a good balance for speed/quality on CPU Spaces
-    image = pipeline(num_inference_steps=100).images[0]
-    
-    return image
+    try:
+        # Reduced steps slightly for better responsiveness if on CPU
+        steps = 50 if DEVICE == "cpu" else 100
+        image = pipeline(num_inference_steps=steps).images[0]
+        return image
+    except Exception as e:
+        print(f"Generation Error: {e}")
+        return create_error_image("INFERENCE_ERR")
 
 # --- CUSTOM CSS ---
-# Giving it that dark, research-lab vibe from your screenshots
 custom_css = """
 body { background-color: #0d0d0d; color: #e0e0e0; font-family: 'Courier New', monospace; }
 .gradio-container { max-width: 700px !important; margin-top: 40px !important; }
@@ -53,9 +68,14 @@ with gr.Blocks(css=custom_css, title="Stable-Lime v1.1") as demo:
     
     with gr.Row():
         with gr.Column():
-            lime_output = gr.Image(label="Generated Artifact", type="pil", elem_id="lime-out")
+            # KEY CHANGE: type="filepath" ensures the API receives a string path
+            # instead of a complex object or base64 string.
+            lime_output = gr.Image(
+                label="Generated Artifact", 
+                type="filepath", 
+                elem_id="lime-out"
+            )
             
-            # Button is now directly below image, no status text
             generate_btn = gr.Button("INITIALIZE GENERATION", elem_classes="lime-btn")
     
     gr.HTML("<div class='footer'>Running on Unconditional U-Net Architecture | Powered by Limes</div>")