Update app.py

app.py

@@ -70,34 +70,61 @@ def generate_image(prompt, output_filename="generated_image"):
 
     return None, "No image generated"
 
-# Create Gradio interface
-with gr.Blocks(title="Image Generation App") as demo:
-    gr.Markdown("# Image Generation with Gemini")
-    gr.Markdown("Enter a prompt to generate an image using Google's Gemini model")
+# Function to handle chat interaction
+def chat_handler(user_input, chat_history):
+    # Add user message to chat history
+    chat_history.append([user_input, None])
     
+    # Generate image based on user input
+    img, status = generate_image(user_input)
+    
+    # Add AI response to chat history
+    if img:
+        chat_history.append([None, img])
+    chat_history.append([None, status])
+    
+    return chat_history, ""
+
+# 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 = gr.Chatbot(
+        label="Chat",
+        height=400,
+        bubble_full_width=False,
+        avatar_images=(None, None)  # No avatars for simplicity
+    )
+    
+    # Input area
     with gr.Row():
-        with gr.Column():
-            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")
-        
-        with gr.Column():
-            image_output = gr.Image(label="Generated Image")
-            text_output = gr.Textbox(label="Status")
+        prompt_input = gr.Textbox(
+            label="",
+            placeholder="Type something",
+            show_label=False,
+            container=False,
+            scale=4
+        )
+        run_btn = gr.Button("Run", scale=1)
+    
+    # State to maintain chat history
+    chat_state = gr.State([])
+
+    # Connect the button to the chat handler
+    run_btn.click(
+        fn=chat_handler,
+        inputs=[prompt_input, chat_state],
+        outputs=[chatbot, prompt_input],
+        _js="() => {return [document.querySelector('input').value, null]}"
+    )
 
-    # Connect the button to the generation function
-    generate_btn.click(
-        fn=generate_image,
-        inputs=[prompt_input, filename_input],
-        outputs=[image_output, text_output]
+    # Also allow Enter key to submit
+    prompt_input.submit(
+        fn=chat_handler,
+        inputs=[prompt_input, chat_state],
+        outputs=[chatbot, prompt_input]
     )
 
 if __name__ == "__main__":