matt-le-kat commited on
Commit
cb85ac7
1 Parent(s): fe6a365

Create new file

Browse files
Files changed (1) hide show
  1. app.py +25 -0
app.py ADDED
@@ -0,0 +1,25 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+
3
+ import torch
4
+ from torch import autocast
5
+ from diffusers import StableDiffusionPipeline
6
+ from datasets import load_dataset
7
+ from PIL import Image
8
+ import re
9
+
10
+ model_id = "CompVis/stable-diffusion-v1-4"
11
+ device = "cuda"
12
+
13
+ pipe = StableDiffusionPipeline.from_pretrained(model_id, use_auth_token=True)
14
+ pipe = pipe.to(device)
15
+
16
+ def infer(prompt):
17
+ with autocast("cuda"):
18
+ images_list = pipe(
19
+ [prompt])['sample']
20
+ return images_list
21
+ prompt = 'a perfect photo of a sunset on a Greek beach'
22
+
23
+ intf = gr.Interface(fn = infer, inputs = prompt, outputs = gr.outputs.Image())
24
+
25
+ intf.launch()