Update app.py

app.py

@@ -1,28 +1,39 @@
 import gradio as gr
 import requests
-from PIL import Image
 import io
+from PIL import Image
 import os
+import base64
 
 API_URL = "https://api-inference.huggingface.co/models/CompVis/stable-diffusion-v1-4"
 headers = {"Authorization": f"Bearer {os.getenv('HF_API_KEY')}"}
 
 def generate_image(prompt):
-    response = requests.post(API_URL, headers=headers, json={"inputs": prompt})
-    if response.status_code == 200:
-        image_bytes = response.content
-        image = Image.open(io.BytesIO(image_bytes))
+    payload = {"inputs": prompt}
+    response = requests.post(API_URL, headers=headers, json=payload)
+
+    # If error or waiting for model to load
+    if response.status_code != 200:
+        try:
+            return f"Error: {response.json()}"
+        except:
+            return f"Error: {response.status_code} - {response.text}"
+
+    # If it's actually an image, convert and return
+    try:
+        image = Image.open(io.BytesIO(response.content))
         return image
-    else:
-        return f"Error: {response.status_code} - {response.text}"
+    except Exception as e:
+        return f"Could not load image: {str(e)}"
 
 iface = gr.Interface(
     fn=generate_image,
     inputs=gr.Textbox(label="Enter your prompt"),
     outputs=gr.Image(type="pil", label="Generated Image"),
     title="Text-to-Image Generator",
-    theme="default",
+    theme="default"
 )
 
 iface.launch()
 
+