File size: 756 Bytes
0579aab
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
import torch

device = "cuda" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.float16 if device == "cuda" else None

from diffusers import StableDiffusionPipeline

model_id = "CompVis/stable-diffusion-v1-4"
pipe = StableDiffusionPipeline.from_pretrained(
    model_id, revision="fp16", torch_dtype=torch_dtype
).to(device)

def predict(prompt):
    return pipe(prompt).images[0]
    
import gradio as gr

gradio_ui = gr.Interface(
    fn=predict,
    title="Stable Diffusion Demo",
    description="Enter a description of an image you'd like to generate!",
    inputs=[
        gr.Textbox(lines=2, label="Paste some text here"),
    ],
    outputs=["image"],
    examples=[["a photograph of an astronaut riding a horse"]],
)

gradio_ui.launch()