AP123 commited on
Commit
0024d7b
1 Parent(s): 31e6e36

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +40 -0
app.py ADDED
@@ -0,0 +1,40 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ from diffusers import StableDiffusionXLPipeline, EulerDiscreteScheduler
4
+ from huggingface_hub import hf_hub_download
5
+ from PIL import Image
6
+ import io
7
+ import spaces
8
+
9
+ @spaces.GPU
10
+ def generate_image(prompt):
11
+ base = "stabilityai/stable-diffusion-xl-base-1.0"
12
+ repo = "ByteDance/SDXL-Lightning"
13
+ ckpt = "sdxl_lightning_4step_unet.pth"
14
+
15
+ # Load model
16
+ pipe = StableDiffusionXLPipeline.from_pretrained(base, torch_dtype=torch.float16, variant="fp16").to("cuda")
17
+ pipe.unet.load_state_dict(torch.load(hf_hub_download(repo, ckpt), map_location="cuda"))
18
+ pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")
19
+
20
+ # Generate image
21
+ image = pipe(prompt, num_inference_steps=4, guidance_scale=0).images[0]
22
+
23
+ return image
24
+
25
+ #gradio
26
+ description = """
27
+ This demo utilizes the SDXL-Lightning model by ByteDance, which is a fast text-to-image generative model capable of producing high-quality images in 4 steps.
28
+ As a community effort, this demo was put together by AngryPenguin. Link to model: https://huggingface.co/ByteDance/SDXL-Lightning
29
+ """
30
+
31
+ demo = gr.Interface(
32
+ fn=generate_image,
33
+ inputs="text",
34
+ outputs="image",
35
+ title="Text-to-Image with SDXL Lightning ⚡",
36
+ description=description
37
+ )
38
+
39
+ demo.queue(max_size=20)
40
+ demo.launch()