PeterL1n commited on
Commit
e935aa8
1 Parent(s): f2aa102

Fix bug and improve speed

Browse files
Files changed (1) hide show
  1. app.py +10 -10
app.py CHANGED
@@ -18,6 +18,7 @@ checkpoints = {
18
  "4-Step" : ["sdxl_lightning_4step_unet.safetensors", 4],
19
  "8-Step" : ["sdxl_lightning_8step_unet.safetensors", 8],
20
  }
 
21
 
22
 
23
  # Ensure model and scheduler are initialized in GPU-enabled function
@@ -49,18 +50,17 @@ if SAFETY_CHECKER:
49
  # Function
50
  @spaces.GPU(enable_queue=True)
51
  def generate_image(prompt, ckpt):
 
 
52
 
53
  checkpoint = checkpoints[ckpt][0]
54
- num_inference_steps = checkpoints[ckpt][1]
55
-
56
- if num_inference_steps==1:
57
- # Ensure sampler uses "trailing" timesteps and "sample" prediction type for 1-step inference.
58
- pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing", prediction_type="sample")
59
- else:
60
- # Ensure sampler uses "trailing" timesteps.
61
- pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")
62
 
63
- pipe.unet.load_state_dict(load_file(hf_hub_download(repo, checkpoint), device="cuda"))
64
  results = pipe(prompt, num_inference_steps=num_inference_steps, guidance_scale=0)
65
 
66
  if SAFETY_CHECKER:
@@ -84,7 +84,7 @@ with gr.Blocks(css="style.css") as demo:
84
  gr.Markdown(description)
85
  with gr.Group():
86
  with gr.Row():
87
- prompt = gr.Textbox(label='Enter you image prompt:', scale=8)
88
  ckpt = gr.Dropdown(label='Select inference steps',choices=['1-Step', '2-Step', '4-Step', '8-Step'], value='4-Step', interactive=True)
89
  submit = gr.Button(scale=1, variant='primary')
90
  img = gr.Image(label='SDXL-Lightning Generated Image')
 
18
  "4-Step" : ["sdxl_lightning_4step_unet.safetensors", 4],
19
  "8-Step" : ["sdxl_lightning_8step_unet.safetensors", 8],
20
  }
21
+ loaded = None
22
 
23
 
24
  # Ensure model and scheduler are initialized in GPU-enabled function
 
50
  # Function
51
  @spaces.GPU(enable_queue=True)
52
  def generate_image(prompt, ckpt):
53
+ global loaded
54
+ print(prompt, ckpt)
55
 
56
  checkpoint = checkpoints[ckpt][0]
57
+ num_inference_steps = checkpoints[ckpt][1]
58
+
59
+ if loaded != num_inference_steps:
60
+ pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing", prediction_type="sample" if num_inference_steps==1 else "epsilon")
61
+ pipe.unet.load_state_dict(load_file(hf_hub_download(repo, checkpoint), device="cuda"))
62
+ loaded = num_inference_steps
 
 
63
 
 
64
  results = pipe(prompt, num_inference_steps=num_inference_steps, guidance_scale=0)
65
 
66
  if SAFETY_CHECKER:
 
84
  gr.Markdown(description)
85
  with gr.Group():
86
  with gr.Row():
87
+ prompt = gr.Textbox(label='Enter your prompt (English)', scale=8)
88
  ckpt = gr.Dropdown(label='Select inference steps',choices=['1-Step', '2-Step', '4-Step', '8-Step'], value='4-Step', interactive=True)
89
  submit = gr.Button(scale=1, variant='primary')
90
  img = gr.Image(label='SDXL-Lightning Generated Image')