import gradio as gr
from ppdiffusers import StableDiffusionPipeline

def generate_image(model, prompt, width, height, num_inference_steps, guidance_scale):

    if model == "shanshui_gen_style":
        pipe = StableDiffusionPipeline.from_pretrained("megemini/shanshui_gen_style", from_hf_hub=True)
    elif model == "shanshui_style":
        pipe = StableDiffusionPipeline.from_pretrained("megemini/shanshui_style", from_hf_hub=True)
    else:
        raise

    image = pipe(
        prompt, 
        num_inference_steps=100, 
        guidance_scale=7.5, 
        height=height, 
        width=width,).images[0]
    return image

demo = gr.Blocks()

with demo:
    gr.Markdown(
        r"### 【Hackathon】基于PaddleNLP PPDiffusers 训练 AIGC 趣味模型"
    )
    gr.Markdown(
        r"""
        [【Hackathon】基于PaddleNLP PPDiffusers 训练 AIGC 趣味模型](https://github.com/PaddlePaddle/community/blob/master/hackthon_4th/%E3%80%90PaddlePaddle%20Hackathon%204%E3%80%91%20%E6%A8%A1%E5%9E%8B%E5%A5%97%E4%BB%B6%E5%BC%80%E6%BA%90%E8%B4%A1%E7%8C%AE%E4%BB%BB%E5%8A%A1%E5%90%88%E9%9B%86.md#no105%E5%9F%BA%E4%BA%8Epaddlenlp-ppdiffusers-%E8%AE%AD%E7%BB%83-aigc-%E8%B6%A3%E5%91%B3%E6%A8%A1%E5%9E%8B-)
        
        
        """
    )
    with gr.Row():
        with gr.Column():
            with gr.Row():
                model = gr.Dropdown(["shanshui_gen_style", "shanshui_style"], label="Model", info="The model to generate image.")
            with gr.Row():
                prompt = gr.Textbox(label='prompt')
            with gr.Row():
                width = gr.Slider(128, 768, value=512, step=8, label="Width", info="The width of image.")
                height = gr.Slider(128, 768, value=512, step=8, label="Height", info="The height of image.")
            with gr.Row():
                num_inference_steps = gr.Textbox(label='num inference steps')
                guidance_scale = gr.Textbox(label='guidance scale')
            with gr.Row():
                btn = gr.Button(value="Run")
        with gr.Column():
            with gr.Row():
                output = gr.Image()

    gr.Examples(
        [
            [
                "shanshui_gen_style",
                "A fantasy landscape in <shanshui-gen-style>",
                512,
                288,
                100,
                7.5,
            ],
            [
                "shanshui_style",
                "A fantasy landscape in <shanshui-style>",
                512,
                288,
                100,
                7.5,
            ],
        ],
        [model, prompt, width, height, num_inference_steps, guidance_scale]
    )

    btn.click(
        generate_image, 
        [model, prompt, width, height, num_inference_steps, guidance_scale], 
        output)



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