Update app.py

app.py

@@ -1,89 +1,164 @@
-import gradio as gr
-from google import genai
+import os
+import tempfile
 from PIL import Image
-from io import BytesIO
-import base64
+import gradio as gr
+import google.generativeai as genai
+
+# Hardcode the Gemini API key
+GEMINI_API_KEY = "AIzaSyDL5Rilo7ptJpUOZdY6wy8PJYUcVcnDADs"
 
-# Configure the Gemini API client with the hardcoded API key
-GOOGLE_API_KEY = "AIzaSyDL5Rilo7ptJpUOZdY6wy8PJYUcVcnDADs"
-client = genai.Client(api_key=GOOGLE_API_KEY)
-GEMINI_MODEL_NAME = 'gemini-2.5-flash-image-preview'
+def configure_api_key():
+    """Configure the Gemini API key."""
+    if not GEMINI_API_KEY:
+        raise gr.Error("Gemini API key is not set.")
+    genai.configure(api_key=GEMINI_API_KEY)
 
-def process_image(image, prompt):
+def generate(text, images, model="gemini-2.5-flash"):
+    """Generate content using the Gemini model."""
+    configure_api_key()
+    
+    # Convert images to Gemini-compatible format
+    contents = []
+    for img in images:
+        if img.mode == "RGBA":
+            img = img.convert("RGB")
+        contents.append(img)
+    contents.append(text)
+    
     try:
-        # Prepare the content for the Gemini API
-        contents = []
-        if image:
-            # Convert Gradio image (PIL Image) to base64
-            buffered = BytesIO()
-            image.save(buffered, format="PNG")
-            img_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
-            contents.append({
-                "parts": [
-                    {"text": prompt},
-                    {
-                        "inline_data": {
-                            "mime_type": "image/png",
-                            "data": img_base64
-                        }
-                    }
-                ]
-            })
-        else:
-            # Text-to-image generation
-            contents.append({"parts": [{"text": prompt}]})
+        response = genai.GenerativeModel(model).generate_content(contents)
+        
+        text_response = ""
+        image_path = None
+        
+        for part in response.candidates[0].content.parts:
+            if hasattr(part, 'text') and part.text:
+                text_response += part.text + "\n"
+            elif hasattr(part, 'inline_data') and part.inline_data:
+                # Save generated image to a temporary file
+                with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
+                    image_path = tmp.name
+                    generated_image = Image.open(BytesIO(part.inline_data.data))
+                    generated_image.save(image_path)
+                    print(f"Generated image saved to: {image_path} with prompt: {text}")
+        
+        return image_path, text_response
+    except Exception as e:
+        raise gr.Error(f"Error generating content: {str(e)}")
 
-        # Call the Gemini API
-        response = client.models.generate_content(
-            model=GEMINI_MODEL_NAME,
-            contents=contents
-        )
+def load_uploaded_images(uploaded_files):
+    """Load and display uploaded images immediately."""
+    uploaded_images = []
+    if uploaded_files:
+        for file in uploaded_files:
+            if file.name.lower().endswith(('.png', '.jpg', '.jpeg', '.webp')):
+                img = Image.open(file.name)
+                uploaded_images.append(img)
+    return uploaded_images
 
-        # Process the response
-        for candidate in response.candidates:
-            for part in candidate.content.parts:
-                if hasattr(part, 'inline_data') and part.inline_data:
-                    # Decode the generated image
-                    img_data = base64.b64decode(part.inline_data.data)
-                    return Image.open(BytesIO(img_data))
-                elif part.text:
-                    return f"Text response: {part.text}"
+def process_image_and_prompt(uploaded_files, prompt):
+    """Process uploaded images and prompt to generate output."""
+    try:
+        if not uploaded_files:
+            raise gr.Error("Please upload at least one image.")
+        if not prompt:
+            raise gr.Error("Please provide a prompt.")
         
-        return "No image or text returned by the model."
-    
+        # Load images
+        images = load_uploaded_images(uploaded_files)
+        
+        # Generate content
+        image_path, text_response = generate(
+            text=prompt,
+            images=images
+        )
+        
+        # Prepare outputs
+        output_images = [Image.open(image_path)] if image_path else None
+        return images, output_images, text_response
     except Exception as e:
-        return f"Error: {str(e)}"
+        raise gr.Error(f"Error: {str(e)}")
 
-# Create the Gradio interface
-css = '''
-.grid-container img {object-fit: contain}
-.grid-container {display: grid; grid-template-columns: 1fr}
-'''
+# Gradio interface
+with gr.Blocks(css="style.css") as demo:
+    gr.HTML("""
+    <div class="header-container">
+      <div>
+          <img src="https://www.gstatic.com/lamda/images/gemini_favicon_f069958c85030456e93de685481c559f160ea06b.png" alt="Gemini logo">
+      </div>
+      <div>
+          <h1>Gemini for Image Editing</h1>
+          <p>Powered by <a href="https://gradio.app/">Gradio</a>⚡️ | 
+          <a href="https://huggingface.co/spaces">Duplicate this Space</a> |
+          <a href="https://aistudio.google.com/apikey">Learn about Gemini API</a></p>
+      </div>
+    </div>
+    """)
+    
+    with gr.Accordion("⚠️ API Configuration ⚠️", open=False):
+        gr.Markdown("""
+        - **Note:** The Gemini API is pre-configured for this Space.
+        - **Issue:** ❗ Sometimes the model returns text instead of an image.
+        ### 🔧 Steps to Address:
+        1. **🛠️ Duplicate the Space**  
+           - Create a copy on Hugging Face Spaces for modifications.
+        2. **🔑 API Key Info**  
+           - The API key is already set up for this demo. For custom deployments, get your own key from Google AI Studio.
+        """)
 
-with gr.Blocks(theme=gr.themes.Citrus(), css=css) as demo:
-    gr.HTML('''
-    <img src='https://huggingface.co/spaces/multimodalart/nano-banana/resolve/main/nano_banana_pros_light.png' style='margin: 0 auto; max-width: 500px' />
-    <h3 style='text-align:center'>Nano Banana: Gemini 2.5 Flash Image Preview</h3>
-    ''')
+    with gr.Accordion("📌 Usage Instructions", open=False):
+        gr.Markdown("""
+        ### 📌 Usage  
+        - Upload an image and enter a prompt to generate outputs.
+        - If text is returned instead of an image, it will appear in the text output.
+        - Supported formats: PNG, JPG, JPEG, WEBP
+        - ❌ **Do not use NSFW images!**
+        """)
 
     with gr.Row():
-        with gr.Column(scale=1):
-            image_input = gr.Image(type="pil", label="Upload Image (Optional)", file_types=["image"])
+        with gr.Column():
+            image_input = gr.File(
+                file_types=["image"],
+                file_count="multiple",
+                label="Upload Images"
+            )
             prompt_input = gr.Textbox(
-                label="Prompt",
-                placeholder="e.g., 'Add a nano-banana to the image in a fancy restaurant setting' or 'Generate a cat eating a nano-banana'"
+                lines=2,
+                placeholder="Enter prompt here...",
+                label="Prompt"
+            )
+            submit_btn = gr.Button("Generate")
+        
+        with gr.Column():
+            uploaded_gallery = gr.Gallery(label="Uploaded Images")
+            output_gallery = gr.Gallery(label="Generated Outputs")
+            output_text = gr.Textbox(
+                label="Gemini Output",
+                placeholder="Text response will appear here if no image is generated."
             )
-            generate_button = gr.Button("Generate", variant="primary")
-
-        with gr.Column(scale=1):
-            output_image = gr.Image(label="Generated Image", type="pil")
 
-    # Event handler
-    generate_button.click(
-        fn=process_image,
+    submit_btn.click(
+        fn=process_image_and_prompt,
         inputs=[image_input, prompt_input],
-        outputs=[output_image]
+        outputs=[uploaded_gallery, output_gallery, output_text]
+    )
+    
+    image_input.upload(
+        fn=load_uploaded_images,
+        inputs=[image_input],
+        outputs=[uploaded_gallery]
+    )
+
+    gr.Markdown("## Try these examples")
+    examples = [
+        ["data/1.webp", "change text to 'HUGGINGFACE'"],
+        ["data/2.webp", "remove the spoon from hand only"],
+        ["data/3.webp", "change text to 'AI POWERED'"],
+        ["data/1.jpg", "add futuristic style to background"],
+    ]
+    gr.Examples(
+        examples=examples,
+        inputs=[image_input, prompt_input]
     )
 
-if __name__ == "__main__":
-    demo.launch()
+demo.launch()