Update app.py

app.py

@@ -1,108 +1,45 @@
+import os
 import gradio as gr
 import requests
-import base64
-from io import BytesIO
-from PIL import Image
-import os
 
-# 🔐 Use environment variable for your key (set it in Hugging Face "Secrets")
+# Load API key from secret
 API_KEY = os.getenv("GOOGLE_API_KEY")
 
-# ---------- Helper Functions ---------- #
-
-def text_generation(prompt):
-    url = "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:generateContent"
-    headers = {"Content-Type": "application/json", "x-goog-api-key": API_KEY}
-    data = {"contents": [{"parts": [{"text": prompt}]}]}
-    res = requests.post(url, headers=headers, json=data)
-    output = res.json()
-    try:
-        return output["candidates"][0]["content"]["parts"][0]["text"]
-    except Exception:
-        return str(output)
-
-
 def image_generation(prompt):
-    url = "https://generativelanguage.googleapis.com/v1beta/models/imagen-3.0-generate:generateImage"
-    headers = {"Content-Type": "application/json", "x-goog-api-key": API_KEY}
-    data = {"prompt": {"text": prompt}}
-    res = requests.post(url, headers=headers, json=data)
-    result = res.json()
-    try:
-        img_b64 = result["images"][0]["data"]
-        image_bytes = base64.b64decode(img_b64)
-        return Image.open(BytesIO(image_bytes))
-    except Exception:
-        return str(result)
+    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
 
-def video_generation(prompt):
-    url = "https://generativelanguage.googleapis.com/v1beta/models/veo-2.0-generate-001:generateVideo"
-    headers = {"Content-Type": "application/json", "x-goog-api-key": API_KEY}
+    headers = {"Content-Type": "application/json"}
     data = {"prompt": {"text": prompt}}
+
     res = requests.post(url, headers=headers, json=data)
-    result = res.json()
-    try:
-        video_url = result["videos"][0].get("url", None)
-        if video_url:
-            return video_url
-        else:
-            return "Video generated (but no URL returned): " + str(result)
-    except Exception:
-        return str(result)
 
+    # Debugging info
+    if res.status_code != 200:
+        return f"❌ Error {res.status_code}: {res.text}"
 
-def text_to_speech(text):
-    url = "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-tts:generateContent"
-    headers = {"Content-Type": "application/json", "x-goog-api-key": API_KEY}
-    data = {"contents": [{"parts": [{"text": text}]}]}
-    res = requests.post(url, headers=headers, json=data)
-    result = res.json()
     try:
-        audio_b64 = result["candidates"][0]["content"]["parts"][0]["inline_data"]["data"]
-        audio_bytes = base64.b64decode(audio_b64)
-        return (audio_bytes, "audio/wav")
-    except Exception:
-        return str(result)
-
-# ---------- Gradio Interfaces ---------- #
+        result = res.json()
+    except Exception as e:
+        return f"❌ Failed to parse JSON response: {str(e)}\n\nRaw response:\n{res.text}"
 
-text_tab = gr.Interface(
-    fn=text_generation,
-    inputs=gr.Textbox(label="Enter your text prompt"),
-    outputs=gr.Textbox(label="Generated Text"),
-    title="🧠 Text Generation (Gemini 2.5 Flash)"
-)
+    # 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}"
 
-image_tab = gr.Interface(
+# Gradio UI
+demo = gr.Interface(
     fn=image_generation,
-    inputs=gr.Textbox(label="Enter image prompt"),
-    outputs=gr.Image(type="pil"),
-    title="🖼️ Image Generation (Imagen 3)"
-)
-
-video_tab = gr.Interface(
-    fn=video_generation,
-    inputs=gr.Textbox(label="Describe your video idea"),
-    outputs=gr.Textbox(label="Generated Video URL / Response"),
-    title="🎥 Video Generation (Veo 2)"
+    inputs=gr.Textbox(label="Enter your prompt"),
+    outputs="image",
+    title="Google AI Studio Image Generator",
+    description="Generate images using Google AI Studio API."
 )
 
-tts_tab = gr.Interface(
-    fn=text_to_speech,
-    inputs=gr.Textbox(label="Enter text for speech"),
-    outputs=gr.Audio(label="Generated Speech"),
-    title="🗣️ Text-to-Speech (Gemini TTS)"
-)
-
-# ---------- Combine All Tabs ---------- #
-
-demo = gr.TabbedInterface(
-    [text_tab, image_tab, video_tab, tts_tab],
-    tab_names=["Text", "Image", "Video", "Speech"]
-)
-
-# ---------- Launch ---------- #
-
 if __name__ == "__main__":
     demo.launch()