Update app.py

app.py

@@ -34,15 +34,14 @@
 # if __name__ == "__main__":
 #     demo.launch()
 
-
-
 import os
 import gradio as gr
-from google import genai
+import google.genai as genai   # ✅ correct import
 from PIL import Image
 from io import BytesIO
+import base64
 
-# Load API key from environment
+# Load API key
 api_key = os.getenv("Gemini_api_key")
 
 # Initialize Gemini client
@@ -52,48 +51,52 @@ client = genai.Client(api_key=api_key)
 def generate_or_edit_image(prompt, uploaded_img):
     contents = [prompt]
 
-    # If user uploads an image, include it in the request
+    # If user uploads an image, include it
     if uploaded_img is not None:
-        # Validate image format
-        if uploaded_img.format not in ["JPEG", "JPG", "PNG"]:
-            return "❌ Please upload a JPG, JPEG, or PNG image."
-
-        # Convert uploaded image to PNG bytes (safe for Gemini)
         buffered = BytesIO()
         uploaded_img.save(buffered, format="PNG")
         image_bytes = buffered.getvalue()
-
         contents.append({
             "mime_type": "image/png",
             "data": image_bytes
         })
 
-    # Call Gemini image model
-    response = client.models.generate_content(
-        model="gemini-2.5-flash-image-preview",
-        contents=contents,
-    )
+    try:
+        response = client.models.generate_content(
+            model="gemini-2.5-flash-image-preview",
+            contents=contents,
+        )
+
+        # 🔎 Debug: print structure
+        print("Gemini Response:", response)
+
+        for part in response.candidates[0].content.parts:
+            if part.inline_data is not None:
+                # Sometimes Gemini returns base64 string
+                data = part.inline_data.data
+                if isinstance(data, str):
+                    image_bytes = base64.b64decode(data)
+                else:
+                    image_bytes = data
+
+                image = Image.open(BytesIO(image_bytes))
+                return image
 
-    # Extract image from response
-    for part in response.candidates[0].content.parts:
-        if part.inline_data is not None:
-            image = Image.open(BytesIO(part.inline_data.data))
-            return image
-    return None
+        return "⚠️ No image generated. Try a different prompt."
 
-# Gradio interface
+    except Exception as e:
+        return f"❌ Error: {e}"
+
+# Gradio UI
 demo = gr.Interface(
     fn=generate_or_edit_image,
     inputs=[
         gr.Textbox(label="Enter your prompt"),
-        gr.Image(
-            label="Upload an image (optional)",
-            type="pil"   # just keep type=pil, no file_types in your version
-        )
+        gr.Image(label="Upload an image (optional)", type="pil")
     ],
     outputs=gr.Image(label="Generated/Edited Image"),
     title="Image Generator & Editor",
-    description="Enter a text prompt to generate a new image, or upload an image (JPG, JPEG, PNG) to edit it."
+    description="Enter a prompt to generate a new image, or upload a JPG/PNG to edit."
 )
 
 if __name__ == "__main__":