import gradio as gr
from diffusers import StableDiffusionPipeline
# from PIL import Image


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

# pipe_img2img = StableDiffusionImg2ImgPipeline.from_pretrained(model_id, tokenizer=tokenizer, text_encoder=text_encoder, vae=vae, unet=unet).to(device)
pipe_text2img = StableDiffusionPipeline.from_pretrained(model_id).to(device)

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)
    image = output.images[0]
    return image

gr.Interface(fn=infer, inputs=
    [
        # gr.Image(source="upload", type="filepath", label="原始图像"), 
        gr.Textbox(label = '提示词(prompt)'),
        gr.Slider(2, 15, value = 7, label = '文本引导强度(guidance scale)'),
        gr.Slider(10, 30, value = 20, step = 1, label = '迭代次数(inference steps)'),
        gr.Slider(256, 768, value = 512, step = 64, label = '宽度(width)'),
        gr.Slider(256, 768, value = 512, step = 64, label = '高度(height)'),
    ],outputs='image').queue(concurrency_count=10).launch()