import gradio as gr

def generate_multi_model(input_img):
    output_img = input_img
    return output_img

def train_one_shot(title, context, img):
    return f"Title:{title}\nContext:{context}\n...{img}"

def generate_one_shot(src_img, ref_img):
    output_img = src_img
    return output_img

def train_zero_shot(title, context, img):
    return f"Title:{title}\nContext:{context}\n...{img}"

def generate_zero_shot(src_img, ref_prompt):
    output_img = src_img
    return output_img

with gr.Blocks() as demo:
    # 顶部文字
    gr.Markdown("# MMFS")

    # 多个tab
    with gr.Tabs():

        with gr.TabItem("multi-model"):
            multi_input_img = gr.Image(shape=(100, 100),  label="请上传人像图片")
            multi_model_button = gr.Button("随机风格化")
            
            multi_output_img = gr.Image(shape=(100, 100),  label="人像风格化图")


        with gr.TabItem("one-shot"):
            with gr.Row():
                one_shot_src_img = gr.Image(shape=(100, 100),  label="请上传内容图片")
                one_shot_ref_img = gr.Image(shape=(100, 100),  label="请上传风格图片")
            with gr.Row():
                one_shot_train_button = gr.Button("训练模型")
                one_shot_test_button = gr.Button("风格化")

            one_shot_output_img = gr.Image(shape=(100, 100),  label="one-shot风格化图")

        
        with gr.TabItem("zero-shot"):
            zero_shot_src_img = gr.Image(shape=(100, 100),  label="请上传内容图片")
            zero_shot_ref_prompt = gr.Textbox(label="prompt", lines=1, placeholder="请输入参考风格描述（英文）")
            with gr.Row():
                zero_shot_train_button = gr.Button("训练模型")
                zero_shot_test_button = gr.Button("风格化")

            zero_shot_output_img = gr.Image(shape=(100, 100),  label="zero-shot风格化图")


    multi_model_button.click(fn=generate_multi_model, inputs=multi_input_img, outputs=multi_output_img)
    one_shot_test_button.click(fn=generate_one_shot, inputs=[one_shot_src_img, one_shot_ref_img], outputs=one_shot_output_img)
    zero_shot_test_button.click(fn=generate_zero_shot, inputs=[zero_shot_src_img, zero_shot_ref_prompt], outputs=zero_shot_output_img)

demo.launch()