import gradio as gr
from PIL import Image

from diffusers import (
    StableDiffusionPipeline,
    StableDiffusionImg2ImgPipeline,
    StableDiffusionInpaintPipeline,
)

device="cuda"
model_id = "IDEA-CCNL/Taiyi-Stable-Diffusion-1B-Chinese-v0.1"

pipe_text2img = StableDiffusionPipeline.from_pretrained(model_id).to(device)
pipe_img2img = StableDiffusionImg2ImgPipeline(**pipe_text2img.components)
pipe_inpaint = StableDiffusionInpaintPipeline(**pipe_text2img.components)

# io1_text2img = gr.Interface.load("models/IDEA-CCNL/Taiyi-Stable-Diffusion-1B-Chinese-v0.1")

def resize(w_val,l_val,img):
  img = Image.open(img)
  img = img.resize((w_val,l_val), Image.Resampling.LANCZOS)
  return img

def infer(prompt, guide, steps, width, height): 
    output = pipe_text2img(prompt, guidance_scale=guide, num_inference_steps=steps, width=width, height=height)
    # output = io1_text2img(prompt, guidance_scale=guide, num_inference_steps=steps, width=width, height=height)
    image = output.images[0]
    return image

with gr.Blocks() as demo:
    examples = [
                ["飞流直下三千尺, 疑是银河落九天, 瀑布, 插画"], 
                ["东临碣石, 以观沧海, 波涛汹涌, 插画"], 
                ["孤帆远影碧空尽，惟见长江天际流,油画"],
                ["女孩背影, 日落, 唯美插画"],
                ]
    with gr.Row():
        with gr.Column(scale=2, ):
            output = gr.Image(label = '输出(output)')
        
        with gr.Column(scale=1, ):
            guide = gr.Slider(2, 15, value = 7, label = '文本引导强度(guidance scale)')
            steps = gr.Slider(10, 30, value = 20, step = 1, label = '迭代次数(inference steps)')
            width = gr.Slider(256, 768, value = 512, step = 64, label = '宽度(width)')
            height = gr.Slider(256, 768, value = 512, step = 64, label = '高度(height)')
            
            prompt = gr.Textbox(label = '提示词(prompt)')
            submit_btn = gr.Button("生成图片(Generate)").style(margin=False, rounded=(False, True, True, False), full_width=False,)

            ex = gr.Examples(examples, fn=infer, inputs=[prompt, guide, steps, width, height], outputs=output)

        submit_btn.click(fn = infer, inputs = [prompt, guide, steps, width, height], outputs = output)
        
demo.queue(concurrency_count=10).launch()