Update app.py

app.py

@@ -1,45 +1,54 @@
 import os
-import gradio as gr
 import requests
+import gradio as gr
 
-# Load API key from secret
+# Get the API key from Hugging Face Secrets
 API_KEY = os.getenv("GOOGLE_API_KEY")
 
 def image_generation(prompt):
     if not API_KEY:
         return "❌ API key not found. Please add GOOGLE_API_KEY in your Hugging Face secrets."
 
-    url = "https://generativelanguage.googleapis.com/v1beta/models/imagegeneration:generate?key=" + API_KEY
+    # Google Imagen 3.0 endpoint for image generation
+    url = f"https://generativelanguage.googleapis.com/v1beta/models/imagen-3.0-generate:generate?key={API_KEY}"
 
     headers = {"Content-Type": "application/json"}
     data = {"prompt": {"text": prompt}}
 
+    # Send POST request
     res = requests.post(url, headers=headers, json=data)
 
-    # Debugging info
+    # Check for HTTP errors
     if res.status_code != 200:
-        return f"❌ Error {res.status_code}: {res.text}"
+        return f"❌ Error {res.status_code}:\n{res.text}"
 
     try:
         result = res.json()
-    except Exception as e:
-        return f"❌ Failed to parse JSON response: {str(e)}\n\nRaw response:\n{res.text}"
-
-    # Extract image
-    try:
-        image_base64 = result["candidates"][0]["image"]["base64"]
-        return gr.Image(value=image_base64)
-    except Exception as e:
-        return f"⚠️ Could not extract image: {str(e)}\n\nFull JSON:\n{result}"
-
-# Gradio UI
-demo = gr.Interface(
+    except Exception:
+        return f"❌ Invalid JSON response:\n{res.text}"
+
+    # If API returned base64 image data
+    if "images" in result and len(result["images"]) > 0:
+        image_base64 = result["images"][0].get("image_base64", None)
+        if image_base64:
+            import base64
+            from io import BytesIO
+            from PIL import Image
+            image_bytes = base64.b64decode(image_base64)
+            image = Image.open(BytesIO(image_bytes))
+            return image
+
+    # Otherwise show the raw response for debugging
+    return f"⚠️ Unexpected response:\n{result}"
+
+# Gradio interface
+iface = gr.Interface(
     fn=image_generation,
-    inputs=gr.Textbox(label="Enter your prompt"),
-    outputs="image",
+    inputs=gr.Textbox(label="Enter your prompt", placeholder="e.g. A robot painting a sunset over the ocean"),
+    outputs=gr.Image(label="Generated Image"),
     title="Google AI Studio Image Generator",
-    description="Generate images using Google AI Studio API."
+    description="Generate images using Google AI Studio (Imagen 3.0 API). Educational use only."
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    iface.launch()