import gradio as gr
import subprocess
from chat import chat_vivogpt
from draw import draw_vivogpt

def result_for_chat(ask):
    anser = chat_vivogpt(ask)
    return anser

def result_for_draw(height, width, prompt, mode):
    images_url = draw_vivogpt(height, width, prompt, mode)
    return images_url

def handle_button1_click(input_text):
    answer = result_for_chat(input_text)
    return answer

def handle_button2_click(height, width, prompt, mode):
    if mode=="v3":
        m=0
    if mode=="v4":
        m=1
    answer = result_for_draw(height, width, prompt, m)
    return answer



with gr.Blocks() as demo:
    with gr.Tab(label="chat"):
        with gr.Column():
            txt1 = gr.Textbox(label="output", lines=10)
            txt2 = gr.Textbox(label="input", scale=8)
            button1 = gr.Button(value="send", scale=1)
            button1.click(fn=handle_button1_click, inputs=txt2, outputs=txt1)

    with gr.Tab(label="txt-image"):
        with gr.Row():
            with gr.Column():
                image = gr.Image(label="Displayed Image")
            with gr.Column():
                with gr.Row():
                    prompt_draw = gr.Textbox(label="请输入您对图片的描述", scale=8)
                    modes = gr.Dropdown(
                        choices=['v3', 'v4'],
                        label="模型选择",
                        value='v3'  # 设置默认选择
                    )
                    height = gr.Slider(minimum=400, maximum=1200, step=8, label="长", value=512)
                    width = gr.Slider(minimum=400, maximum=1200, step=8, label="宽", value=512)
                    button2 = gr.Button(value="send")
                button2.click(fn=handle_button2_click, inputs=[height, width, prompt_draw, modes], outputs=image)

demo.launch()