app.py

# Alternative layout using a separate gr.Image component
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)
    )
    
    # 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():
        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)
    
    # 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(
        fn=chat_handler,
        inputs=[prompt_input, image_input, chat_state],
        outputs=[chatbot, image_input, 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]
    )

if __name__ == "__main__":
    demo.launch()