Update app.py

app.py

@@ -28,16 +28,19 @@ pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, times
 
 # Inference function.
 @spaces.GPU(enable_queue=True)
-def generate_image(prompt, option):
+def generate_image(prompt, option, progress=gr.Progress()):
     global step_loaded
     print(prompt, option)
     ckpt, step = opts[option]
+    progress(0, total=step)
     if step != step_loaded:
         print(f"Switching checkpoint from {step_loaded} to {step}")
         pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing", prediction_type="sample" if step == 1 else "epsilon")
         pipe.unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device=device))
         step_loaded = step
-    return pipe(prompt, num_inference_steps=step, guidance_scale=0).images[0]
+    def inference_callback(p, i, t, kwargs):
+        progress(i+1, total=step)
+    return pipe(prompt, num_inference_steps=step, guidance_scale=0, callback_on_step_end=inference_callback).images[0]
 
 with gr.Blocks(css="style.css") as demo:
     gr.HTML(