XGBooster commited on
Commit
f6b52a7
·
1 Parent(s): dc05727

Upload app.py

Browse files
Files changed (1) hide show
  1. app.py +46 -0
app.py ADDED
@@ -0,0 +1,46 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ from diffusers import StableDiffusionPipeline
4
+ import os
5
+
6
+ def get_token() -> str:
7
+ return os.environ.get("HUGGING_FACE_TOKEN")
8
+
9
+ def save_images(images: list) -> list:
10
+
11
+ output_files_names = []
12
+ for id, image in enumerate(images):
13
+ filename = f"output{id}.png"
14
+ image.save(filename)
15
+ output_files_names.append(filename)
16
+
17
+ return output_files_names
18
+
19
+
20
+
21
+ def create_img(prompt :str, number_output_requested: int) -> list:
22
+ AUTH_TOKEN = get_token()
23
+ generator = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4",
24
+ revision="fp16",
25
+ torch_dtype=torch.float16,
26
+ use_auth_token=AUTH_TOKEN)
27
+ generator.to("cuda")
28
+ prompt = [prompt] * number_output_requested
29
+ with torch.autocast("cuda"):
30
+ images = generator(prompt).images
31
+ output_paths = save_images(images)
32
+ return output_paths
33
+
34
+ diffusers_app = gr.Interface(
35
+ fn=create_img,
36
+ inputs =
37
+ [
38
+ gr.Textbox(label="Write your prompt below", placeholder = "A squirrel bench pressing 200 kg"),
39
+ gr.Slider(value=1, minimum=1, maximum=8, step=1, label="Number of pictures to generate")
40
+ ],
41
+ outputs = gr.Gallery(label="Generated Images").style(grid=[2]),
42
+ title="Text to Image with Stable Diffusion",
43
+ description="This is a basic app to generate pictures with Stable Diffusion."
44
+ )
45
+
46
+ diffusers_app.launch(debug=True)