huai's picture
Update app.py
2e13497
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()