Update app.py

app.py

@@ -8,37 +8,65 @@ import io
 from PIL import Image
 
 def save_binary_file(file_name, data):
-    with open(file_name, "wb") as f:
-        f.write(data)
+    f = open(file_name, "wb")
+    f.write(data)
+    f.close()
 
 def generate_image(prompt, image=None, output_filename="generated_image"):
-    client = genai.Client(api_key="AIzaSyAQcy3LfrkMy6DqS_8MqftAXu1Bx_ov_E8")
+    # 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)]
+
+    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"],
-        response_mime_type="text/plain"
+        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,
     )
-    
-    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
@@ -47,62 +75,110 @@ def generate_image(prompt, image=None, output_filename="generated_image"):
             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
             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"
 
-def chat_handler(prompt, user_image, chat_history, gallery, output_filename="generated_image"):
+# Function to handle chat interaction
+def chat_handler(prompt, user_image, chat_history, output_filename="generated_image"):
+    # Add the user prompt and image to the chat history
+    user_message_content = []
     if prompt:
-        chat_history.append(("User", prompt))
-    
+        user_message_content.append(prompt)
+    if user_image:
+        user_message_content.append(user_image) # Keep PIL Image in history for display
+    if user_message_content:
+        chat_history.append({"role": "user", "content": user_message_content if len(user_message_content) > 1 else user_message_content[0]})
+
+    # If no input, return early
     if not prompt and not user_image:
-        chat_history.append(("Assistant", "Please provide a prompt or an image."))
-        return chat_history, user_image, None, gallery, ""
-    
+        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)
-    
-    chat_history.append(("Assistant", status))
+
+    assistant_message_content = []
+    assistant_message_content.append(status) # Add text status
+
     if img:
-        gallery.append(img)
-    
-    return chat_history, user_image, img, gallery, ""
+        # Create thumbnail for chatbot
+        thumbnail_size = (100, 100) # Define thumbnail size
+        thumbnail = img.copy()
+        thumbnail.thumbnail(thumbnail_size)
+
+        # Convert thumbnail to base64 for chatbot display
+        buffered = io.BytesIO()
+        thumbnail.save(buffered, format="PNG")
+        thumbnail_base64 = base64.b64encode(buffered.getvalue()).decode()
+        thumbnail_data_uri = f"data:image/png;base64,{thumbnail_base64}"
+        assistant_message_content.append(thumbnail_data_uri) # Add thumbnail to message
+
+    # Add assistant's response to chat history
+    chat_history.append({"role": "assistant", "content": assistant_message_content if len(assistant_message_content) > 1 else assistant_message_content[0]})
 
+    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")
-    
-    chatbot = gr.Chatbot(label="Chat", height=300)
-    gallery_output = gr.Gallery(label="Generated Images", show_label=True, height=200)
-    
+
+    # Chatbot display area for text messages
+    chatbot = gr.Chatbot(
+        label="Chat",
+        height=300,
+        type="messages",
+        avatar_images=(None, None)
+    )
+
+    # Separate image outputs
     with gr.Row():
         uploaded_image_output = gr.Image(label="Uploaded Image")
         generated_image_output = gr.Image(label="Generated Image")
-    
+
+    # Input area
     with gr.Row():
         with gr.Column():
-            image_input = gr.Image(label="Upload Image", type="pil", 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")
+            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([])
-    gallery_state = gr.State([])
-    
+
+    # Connect the button to the chat handler
     generate_btn.click(
         fn=chat_handler,
-        inputs=[prompt_input, image_input, chat_state, gallery_state, filename_input],
-        outputs=[chatbot, uploaded_image_output, generated_image_output, gallery_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, gallery_state, filename_input],
-        outputs=[chatbot, uploaded_image_output, generated_image_output, gallery_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__":
-    demo.launch()
+    demo.launch()