import gradio as gr
from LdmZhPipeline import LDMZhTextToImagePipeline
import torch
import numpy as np
from PIL import Image

device = "cuda" if torch.cuda.is_available() else "cpu"
model_id = "alibaba-pai/pai-diffusion-food-large-zh"

pipe_text2img = LDMZhTextToImagePipeline.from_pretrained(model_id)
pipe_text2img = pipe_text2img.to(device)

def infer_text2img(prompt, guide, steps):
    output = pipe_text2img(prompt, guidance_scale=guide, num_inference_steps=steps, use_sr=True)
    images = output.images[0]
    return images

with gr.Blocks() as demo:
    examples = [
                ["小炒黄牛肉"], 
                ["韩式炸鸡"]
                ]
    with gr.Row():
        with gr.Column(scale=0.5, ):
            image_out = gr.Image(label = '输出(output)')
        with gr.Column(scale=1, ):
            prompt = gr.Textbox(label = '提示词(prompt)')
            submit_btn = gr.Button("生成图像(Generate)")
            with gr.Row(scale=0.5 ):
                guide = gr.Slider(2, 15, value = 7, label = '文本引导强度(guidance scale)')
                steps = gr.Slider(10, 50, value = 20, step = 1, label = '迭代次数(inference steps)')
                ex = gr.Examples(examples, fn=infer_text2img, inputs=[prompt, guide, steps], outputs=image_out)
        submit_btn.click(fn = infer_text2img, inputs = [prompt, guide, steps], outputs = image_out)

demo.queue(concurrency_count=1, max_size=8).launch()