Update app.py

app.py

@@ -4,16 +4,16 @@ import mimetypes
 import tempfile
 import time
 import google.generativeai as genai
-from google.generativeai import types
 import gradio as gr
+from PIL import Image as PILImage
 
-def generate_image(api_key, prompt, file_name):
+def generate_image(api_key, prompt, input_image, file_name):
     # Validate inputs
     if not api_key:
         return None, "Please enter your Gemini API key"
     
-    if not prompt:
-        return None, "Please enter a prompt for image generation"
+    if not prompt and input_image is None:
+        return None, "Please enter a prompt and/or upload an input image"
     
     if not file_name:
         file_name = f"gemini_image_{int(time.time())}"
@@ -40,9 +40,27 @@ def generate_image(api_key, prompt, file_name):
         generated_image_path = None
         generation_text = ""
         
+        # Prepare content based on inputs
+        content = []
+        
+        # Add text prompt if provided
+        if prompt:
+            content.append(prompt)
+        
+        # Add image if provided
+        if input_image is not None:
+            # If the input is a file path (string), open the image
+            if isinstance(input_image, str):
+                img = PILImage.open(input_image)
+                content.append(img)
+            # If the input is already a numpy array from Gradio
+            else:
+                img = PILImage.fromarray(input_image)
+                content.append(img)
+        
         # Generate the content
         response = model.generate_content(
-            prompt,
+            content,
             generation_config=generation_config,
             stream=True
         )
@@ -67,7 +85,7 @@ def generate_image(api_key, prompt, file_name):
         if generated_image_path:
             return generated_image_path, generation_text
         else:
-            return None, "No image was generated. Try a different prompt."
+            return None, "No image was generated. Try a different prompt or input image."
             
     except Exception as e:
         return None, f"Error: {str(e)}"
@@ -89,6 +107,10 @@ def create_ui():
                     placeholder="Describe the image you want to generate",
                     lines=3
                 )
+                input_image = gr.Image(
+                    label="Input Image (Optional)",
+                    type="numpy"
+                )
                 file_name = gr.Textbox(
                     label="Output File Name (optional)", 
                     placeholder="Enter a file name (without extension)"
@@ -101,7 +123,7 @@ def create_ui():
         
         generate_btn.click(
             fn=generate_image,
-            inputs=[api_key, prompt, file_name],
+            inputs=[api_key, prompt, input_image, file_name],
             outputs=[output_image, output_text]
         )
         
@@ -109,11 +131,13 @@ def create_ui():
         ## How to use
         1. Enter your Gemini API key (get one from https://ai.google.dev/)
         2. Write a detailed prompt describing the image you want to generate
-        3. (Optional) Provide a file name for your generated image
-        4. Click "Generate Image" and wait for the result
+        3. (Optional) Upload an input image to influence the generation
+        4. (Optional) Provide a file name for your generated image
+        5. Click "Generate Image" and wait for the result
         
         ## Notes
         - The model used is `gemini-2.0-flash-exp-image-generation`
+        - You can use text prompts, input images, or both together
         - All safety filters are set to "BLOCK_NONE" - use responsibly
         - Image generation may take a few seconds to complete
         """)