Update app.py

app.py

@@ -6,7 +6,6 @@ from google.genai import types
 import gradio as gr
 import io
 from PIL import Image
-import tempfile
 
 def save_binary_file(file_name, data):
     f = open(file_name, "wb")
@@ -63,7 +62,7 @@ def generate_image(prompt, output_filename="generated_image"):
             filename = f"{output_filename}{file_extension}"
             save_binary_file(filename, inline_data.data)
             
-            # Convert binary data to PIL Image
+            # 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:
@@ -72,7 +71,7 @@ def generate_image(prompt, output_filename="generated_image"):
     return None, "No image generated"
 
 # Function to handle chat interaction
-def chat_handler(user_input, chat_history):
+def chat_handler(user_input, chat_history, current_image):
     # Add user message to chat history
     chat_history.append({"role": "user", "content": user_input})
     
@@ -80,31 +79,29 @@ def chat_handler(user_input, chat_history):
     img, status = generate_image(user_input)
     
     # Add AI response to chat history
-    if img:
-        # Save the PIL Image to a temporary file so Gradio can display it
-        with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmp_file:
-            img.save(tmp_file.name)
-            # Add the image as a file path that Gradio can serve
-            chat_history.append({"role": "assistant", "content": tmp_file.name})
-    
-    # Add the status message
     chat_history.append({"role": "assistant", "content": status})
     
-    return chat_history, ""
+    return chat_history, img, ""
 
 # Create Gradio interface with chatbot layout
 with gr.Blocks(title="Image Editing Chatbot") as demo:
     gr.Markdown("# Image Editing Chatbot")
     gr.Markdown("Type a prompt to generate or edit an image using Google's Gemini model")
     
-    # Chatbot display area
+    # Chatbot display area for text conversation
     chatbot = gr.Chatbot(
         label="Chat",
-        height=400,
+        height=200,
         type="messages",  # Explicitly set to 'messages' format
         avatar_images=(None, None)  # No avatars for simplicity
     )
     
+    # Image display area
+    image_output = gr.Image(
+        label="Generated Image",
+        height=400
+    )
+    
     # Input area
     with gr.Row():
         prompt_input = gr.Textbox(
@@ -118,19 +115,20 @@ with gr.Blocks(title="Image Editing Chatbot") as demo:
     
     # State to maintain chat history
     chat_state = gr.State([])
+    image_state = gr.State(None)
 
     # Connect the button to the chat handler
     run_btn.click(
         fn=chat_handler,
-        inputs=[prompt_input, chat_state],
-        outputs=[chatbot, prompt_input]
+        inputs=[prompt_input, chat_state, image_state],
+        outputs=[chatbot, image_output, prompt_input]
     )
 
     # Also allow Enter key to submit
     prompt_input.submit(
         fn=chat_handler,
-        inputs=[prompt_input, chat_state],
-        outputs=[chatbot, prompt_input]
+        inputs=[prompt_input, chat_state, image_state],
+        outputs=[chatbot, image_output, prompt_input]
     )
 
 if __name__ == "__main__":