File size: 2,445 Bytes
844ed4d
ed4625b
 
662515a
ed4625b
844ed4d
ed4625b
 
e93189f
844ed4d
 
ed4625b
 
844ed4d
 
 
 
 
ed4625b
19ffd31
662515a
 
 
 
 
 
ed4625b
 
 
 
 
662515a
ed4625b
 
 
 
 
 
844ed4d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ed4625b
 
 
662515a
844ed4d
ed4625b
19ffd31
ed4625b
 
 
662515a
844ed4d
19ffd31
ac3fd77
 
19ffd31
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
# 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()