Update app.py

app.py

@@ -1,4 +1,109 @@
-# Alternative layout using a separate gr.Image component
+import base64
+import os
+import mimetypes
+from google import genai
+from google.genai import types
+import gradio as gr
+import io
+from PIL import Image
+
+def save_binary_file(file_name, data):
+    f = open(file_name, "wb")
+    f.write(data)
+    f.close()
+
+def generate_image(prompt, image=None, output_filename="generated_image"):
+    # Initialize client with the API key
+    client = genai.Client(
+        api_key="AIzaSyAQcy3LfrkMy6DqS_8MqftAXu1Bx_ov_E8",
+    )
+
+    model = "gemini-2.0-flash-exp-image-generation"
+    parts = [types.Part.from_text(text=prompt)]
+    
+    # If an image is provided, add it to the content
+    if image:
+        # Convert PIL Image to bytes
+        img_byte_arr = io.BytesIO()
+        image.save(img_byte_arr, format="PNG")
+        img_bytes = img_byte_arr.getvalue()
+        # Add the image as a Part with inline_data
+        parts.append({
+            "inline_data": {
+                "mime_type": "image/png",
+                "data": img_bytes
+            }
+        })
+
+    contents = [
+        types.Content(
+            role="user",
+            parts=parts,
+        ),
+    ]
+    generate_content_config = types.GenerateContentConfig(
+        temperature=1,
+        top_p=0.95,
+        top_k=40,
+        max_output_tokens=8192,
+        response_modalities=[
+            "image",
+            "text",
+        ],
+        safety_settings=[
+            types.SafetySetting(
+                category="HARM_CATEGORY_CIVIC_INTEGRITY",
+                threshold="OFF",
+            ),
+        ],
+        response_mime_type="text/plain",
+    )
+
+    # Generate the content
+    response = client.models.generate_content_stream(
+        model=model,
+        contents=contents,
+        config=generate_content_config,
+    )
+
+    # Process the response
+    for chunk in response:
+        if not chunk.candidates or not chunk.candidates[0].content or not chunk.candidates[0].content.parts:
+            continue
+        if chunk.candidates[0].content.parts[0].inline_data:
+            inline_data = chunk.candidates[0].content.parts[0].inline_data
+            file_extension = mimetypes.guess_extension(inline_data.mime_type)
+            filename = f"{output_filename}{file_extension}"
+            save_binary_file(filename, inline_data.data)
+            
+            # Convert binary data to PIL Image for Gradio display
+            img = Image.open(io.BytesIO(inline_data.data))
+            return img, f"Image saved as {filename}"
+        else:
+            return None, chunk.text
+
+    return None, "No image generated"
+
+# Function to handle chat interaction
+def chat_handler(prompt, user_image, chat_history, output_filename="generated_image"):
+    # Add the prompt to the chat history
+    if prompt:
+        chat_history.append({"role": "user", "content": prompt})
+    
+    # If no input, return early
+    if not prompt and not user_image:
+        chat_history.append({"role": "assistant", "content": "Please provide a prompt or an image."})
+        return chat_history, user_image, None, ""
+
+    # Generate image based on user input
+    img, status = generate_image(prompt or "Generate an image", user_image, output_filename)
+    
+    # Add the status message to the chat history
+    chat_history.append({"role": "assistant", "content": status})
+    
+    return chat_history, user_image, img, ""
+
+# Create Gradio interface
 with gr.Blocks(title="Image Editing Chatbot") as demo:
     gr.Markdown("# Image Editing Chatbot")
     gr.Markdown("Upload an image and/or type a prompt to generate or edit an image using Google's Gemini model")
@@ -18,54 +123,40 @@ with gr.Blocks(title="Image Editing Chatbot") as demo:
     
     # Input area
     with gr.Row():
-        image_input = gr.Image(
-            label="Upload Image",
-            type="pil",
-            scale=1,
-            height=100
-        )
-        prompt_input = gr.Textbox(
-            label="",
-            placeholder="Type something",
-            show_label=False,
-            container=False,
-            scale=3
-        )
-        run_btn = gr.Button("Run", scale=1)
+        with gr.Column():
+            image_input = gr.Image(
+                label="Upload Image",
+                type="pil",
+                scale=1,
+                height=100
+            )
+            prompt_input = gr.Textbox(
+                label="Prompt",
+                placeholder="Enter your image description here...",
+                lines=3
+            )
+            filename_input = gr.Textbox(
+                label="Output Filename",
+                value="generated_image",
+                placeholder="Enter desired filename (without extension)"
+            )
+            generate_btn = gr.Button("Generate Image")
     
     # State to maintain chat history
     chat_state = gr.State([])
 
-    def chat_handler(user_input, user_image, chat_history):
-        # Add user message to chat history
-        if user_input:
-            chat_history.append({"role": "user", "content": user_input})
-        
-        # If no input, return early
-        if not user_input and not user_image:
-            chat_history.append({"role": "assistant", "content": "Please provide a prompt or an image."})
-            return chat_history, None, user_image, None, ""
-
-        # Generate image based on user input
-        img, status = generate_image(user_input or "Generate an image", user_image)
-        
-        # Add AI response to chat history
-        chat_history.append({"role": "assistant", "content": status})
-        
-        return chat_history, None, user_image, img, ""
-
     # Connect the button to the chat handler
-    run_btn.click(
+    generate_btn.click(
         fn=chat_handler,
-        inputs=[prompt_input, image_input, chat_state],
-        outputs=[chatbot, image_input, uploaded_image_output, generated_image_output, prompt_input]
+        inputs=[prompt_input, image_input, chat_state, filename_input],
+        outputs=[chatbot, uploaded_image_output, generated_image_output, prompt_input]
     )
 
     # Also allow Enter key to submit
     prompt_input.submit(
         fn=chat_handler,
-        inputs=[prompt_input, image_input, chat_state],
-        outputs=[chatbot, image_input, uploaded_image_output, generated_image_output, prompt_input]
+        inputs=[prompt_input, image_input, chat_state, filename_input],
+        outputs=[chatbot, uploaded_image_output, generated_image_output, prompt_input]
     )
 
 if __name__ == "__main__":