Update app.py

app.py

@@ -8,65 +8,37 @@ import io
 from PIL import Image
 
 def save_binary_file(file_name, data):
-    f = open(file_name, "wb")
-    f.write(data)
-    f.close()
+    with open(file_name, "wb") as f:
+        f.write(data)
 
 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",
-    )
-
+    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",
-        ],
-        safety_settings=[
-            types.SafetySetting(
-                category="HARM_CATEGORY_CIVIC_INTEGRITY",
-                threshold="OFF",
-            ),
-        ],
-        response_mime_type="text/plain",
+        response_modalities=["image", "text"],
+        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
+    
+    response = client.models.generate_content_stream(model=model, contents=contents, config=generate_content_config)
+    
     for chunk in response:
         if not chunk.candidates or not chunk.candidates[0].content or not chunk.candidates[0].content.parts:
             continue
@@ -76,88 +48,61 @@ def generate_image(prompt, image=None, output_filename="generated_image"):
             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
+def chat_handler(prompt, user_image, chat_history, gallery, output_filename="generated_image"):
     if prompt:
-        chat_history.append({"role": "user", "content": prompt})
+        chat_history.append(("User", 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
+        chat_history.append(("Assistant", "Please provide a prompt or an image."))
+        return chat_history, user_image, None, gallery, ""
+    
     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})
+    chat_history.append(("Assistant", status))
+    if img:
+        gallery.append(img)
     
-    return chat_history, user_image, img, ""
+    return chat_history, user_image, img, gallery, ""
 
-# 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 display area for text messages
-    chatbot = gr.Chatbot(
-        label="Chat",
-        height=300,
-        type="messages",
-        avatar_images=(None, None)
-    )
+    chatbot = gr.Chatbot(label="Chat", height=300)
+    gallery_output = gr.Gallery(label="Generated Images", show_label=True, height=200)
     
-    # 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",
-                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)"
-            )
+            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")
             generate_btn = gr.Button("Generate Image")
     
-    # State to maintain chat history
     chat_state = gr.State([])
-
-    # Connect the button to the chat handler
+    gallery_state = gr.State([])
+    
     generate_btn.click(
         fn=chat_handler,
-        inputs=[prompt_input, image_input, chat_state, filename_input],
-        outputs=[chatbot, uploaded_image_output, generated_image_output, prompt_input]
+        inputs=[prompt_input, image_input, chat_state, gallery_state, filename_input],
+        outputs=[chatbot, uploaded_image_output, generated_image_output, gallery_output, prompt_input]
     )
-
-    # Also allow Enter key to submit
+    
     prompt_input.submit(
         fn=chat_handler,
-        inputs=[prompt_input, image_input, chat_state, filename_input],
-        outputs=[chatbot, uploaded_image_output, generated_image_output, prompt_input]
+        inputs=[prompt_input, image_input, chat_state, gallery_state, filename_input],
+        outputs=[chatbot, uploaded_image_output, generated_image_output, gallery_output, prompt_input]
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()