import gradio as gr
import torch
from diffusers import StableDiffusionPipeline
from torch import autocast
from transformers import pipeline

device = "cuda" if torch.cuda.is_available() else "cpu"
model_id = "stable-diffusion-v1-4"
translator = pipeline("translation", model="Helsinki-NLP/opus-mt-zh-en")
pipe = StableDiffusionPipeline.from_pretrained(
    model_id,
    use_auth_token=True,
).to(device)

def stablediffusion(prompt,negative_prompt,H,W,steps,guidance_scale):
    prompt = translator(prompt)[0]["translation_text"]
    image = pipe(prompt, num_inference_steps=steps, guidance_scale=guidance_scale,height=H,width=W,negative_prompt=negative_prompt).images[0]

    return image

interface = gr.Interface(fn=stablediffusion,
                         inputs=[gr.inputs.Textbox(lines=5,placeholder="中文描述>>>"),
                                 gr.inputs.Textbox(lines=5,placeholder="消极中文描述>>>"),
                                 gr.inputs.Slider(0,2048,label="高度",default=512),
                                 gr.inputs.Slider(0,2048,label="宽度",default=512),
                                 gr.inputs.Slider(0,100,label="生成的步数",default=10,step=1),
                                 gr.inputs.Slider(0.,10.,label="强度",default=7.5),],
                         outputs=gr.inputs.Image(label="生成的图片"),
                         description="中文的stable-diffusion模型",
                         examples=[['美丽的天空,青色的草地上有着马儿在奔跑',"",512,512,20,7.5],
                                   ['美丽的少女戴着帽子',"不规则的脸,不正确的手",512,512,10,8.]],)

interface.launch()