lewtun HF staff commited on
Commit
9f3851a
1 Parent(s): 9295e2c

Create new file

Browse files
Files changed (1) hide show
  1. app.py +29 -0
app.py ADDED
@@ -0,0 +1,29 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+
3
+ device = "cuda" if torch.cuda.is_available() else "cpu"
4
+ torch_dtype = torch.float16 if device == "cuda" else None
5
+
6
+ from diffusers import StableDiffusionPipeline
7
+
8
+ model_id = "CompVis/stable-diffusion-v1-4"
9
+ pipe = StableDiffusionPipeline.from_pretrained(
10
+ model_id, revision="fp16", torch_dtype=torch_dtype
11
+ ).to(device)
12
+
13
+ def predict(prompt):
14
+ return pipe(prompt).images[0]
15
+
16
+ import gradio as gr
17
+
18
+ gradio_ui = gr.Interface(
19
+ fn=predict,
20
+ title="Stable Diffusion Demo",
21
+ description="Enter a description of an image you'd like to generate!",
22
+ inputs=[
23
+ gr.Textbox(lines=2, label="Paste some text here"),
24
+ ],
25
+ outputs=["image"],
26
+ examples=[["a photograph of an astronaut riding a horse"]],
27
+ )
28
+
29
+ gradio_ui.launch()