AP123 commited on
Commit
d017ab0
1 Parent(s): 21f1280

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +6 -7
app.py CHANGED
@@ -9,14 +9,14 @@ base = "stabilityai/stable-diffusion-xl-base-1.0"
9
  repo = "ByteDance/SDXL-Lightning"
10
  ckpt = "sdxl_lightning_4step_unet.pth"
11
 
12
- # Model Initialization
13
- pipe = StableDiffusionXLPipeline.from_pretrained(base, torch_dtype=torch.float16, variant="fp16").to("cuda")
14
- pipe.unet.load_state_dict(torch.load(hf_hub_download(repo, ckpt), map_location="cuda"))
15
- pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")
16
-
17
  # Function
18
  @spaces.GPU
19
  def generate_image(prompt):
 
 
 
 
 
20
  image = pipe(prompt, num_inference_steps=4, guidance_scale=0).images[0]
21
  return image
22
 
@@ -34,5 +34,4 @@ demo = gr.Interface(
34
  description=description
35
  )
36
 
37
- demo.queue(max_size=20)
38
- demo.launch()
 
9
  repo = "ByteDance/SDXL-Lightning"
10
  ckpt = "sdxl_lightning_4step_unet.pth"
11
 
 
 
 
 
 
12
  # Function
13
  @spaces.GPU
14
  def generate_image(prompt):
15
+ # Ensure model and scheduler are initialized in GPU-enabled function
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
  image = pipe(prompt, num_inference_steps=4, guidance_scale=0).images[0]
21
  return image
22
 
 
34
  description=description
35
  )
36
 
37
+ demo.launch()