测试加入进度支持

app.py

@@ -58,7 +58,7 @@ def infer(input_image, prompt, seed=42, randomize_seed=False, guidance_scale=2.5
 
     progress(0,desc="Starting")
 
-    def callback_fn(step: int, timestep: int, latents: torch.FloatTensor):
+    def callback_fn(pipe, step, timestep, callback_kwargs):
         print(f"[Step {step}] Timestep: {timestep}")
         progress_value = (step+1.0)/num_inference_steps
         progress(progress_value, desc=f"Image generating, {step + 1}/{num_inference_steps} steps")
@@ -75,7 +75,7 @@ def infer(input_image, prompt, seed=42, randomize_seed=False, guidance_scale=2.5
             width = input_image.size[0],
             height = input_image.size[1],
             num_inference_steps=steps,
-            callback=callback_fn,
+            callback_on_step_end=callback_fn,
             generator=torch.Generator().manual_seed(seed),
         ).images[0]
     else:
@@ -83,7 +83,7 @@ def infer(input_image, prompt, seed=42, randomize_seed=False, guidance_scale=2.5
             prompt=prompt,
             guidance_scale=guidance_scale,
             num_inference_steps=steps,
-            callback=callback_fn,
+            callback_on_step_end=callback_fn,
             generator=torch.Generator().manual_seed(seed),
         ).images[0]