waterplayfire commited on
Commit
8771154
·
1 Parent(s): 7fe319b
Files changed (1) hide show
  1. app.py +21 -0
app.py ADDED
@@ -0,0 +1,21 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import pipeline
3
+
4
+
5
+
6
+ model_id = "CompVis/stable-diffusion-v1-4"
7
+ device = "cuda"
8
+
9
+
10
+ pipeline = pipeline(task="TEXT2IMAGE", model=model_id)
11
+ pipeeline = pipeeline(device)
12
+ def predict(prompt):
13
+ predictions = pipeline(prompt)
14
+ return {p["label"]: p["score"] for p in predictions}
15
+
16
+ gr.Interface(
17
+ predict,
18
+ inputs=gr.inputs.Label(label='prompt'),
19
+ outputs=gr.outputs.Image(),
20
+ title="Image",
21
+ ).launch()