fix: progress updater

.gitignore

@@ -1,4 +1,5 @@
 .env
 __pycache__/
 *.mp4
-*.jpg
+*.jpg
+test.py

app.py

@@ -4,13 +4,16 @@ import torchvision
 from diffusers import I2VGenXLPipeline, DiffusionPipeline
 from torchvision.transforms.functional import to_tensor
 from PIL import Image
+from utils import create_progress_updater
 
 if gr.NO_RELOAD:
-    n_steps = 50
+    n_sdxl_steps = 50
+    n_i2v_steps = 50
     high_noise_frac = 0.8
     negative_prompt = "Distorted, discontinuous, Ugly, blurry, low resolution, motionless, static, disfigured, disconnected limbs, Ugly faces, incomplete arms"
     generator = torch.manual_seed(8888)
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    total_steps = n_sdxl_steps + n_i2v_steps
     print("Device:", device)
 
     base = DiffusionPipeline.from_pretrained(
@@ -41,17 +44,27 @@ def generate(prompt: str, progress=gr.Progress()):
     progress((0, 100), desc="Starting..")
     image = base(
         prompt=prompt,
-        num_inference_steps=n_steps,
+        num_inference_steps=n_sdxl_steps,
         denoising_end=high_noise_frac,
         output_type="latent",
-        callback_on_step_end=lambda p, s, t, d: progress((s, 100), desc="Generating first frame..."),
+        callback_on_step_end=create_progress_updater(
+            start=0, 
+            total=total_steps, 
+            desc="Generating first frame...", 
+            progress=progress,
+        ),
     ).images[0]
     image = refiner(
         prompt=prompt,
-        num_inference_steps=n_steps,
+        num_inference_steps=n_sdxl_steps,
         denoising_start=high_noise_frac,
         image=image,
-         callback_on_step_end=lambda p, s, t, d: progress((s+40, 100), desc="Refining first frame..."),
+        callback_on_step_end=create_progress_updater(
+            start=n_sdxl_steps * high_noise_frac,
+            total=total_steps,
+            desc="Refining first frame...",
+            progress=progress,
+        ),
     ).images[0]
     image = to_tensor(image)
     frames: list[Image.Image] = pipeline(
@@ -62,7 +75,12 @@ def generate(prompt: str, progress=gr.Progress()):
         guidance_scale=9.0,
         generator=generator,
         decode_chunk_size=10,
-        callback_on_step_end=lambda p, s, t, d: progress((s+50, 100), desc="Generating video..."),
+        callback_on_step_end=create_progress_updater(
+            start=n_sdxl_steps,
+            total=total_steps,
+            desc="Generating video...",
+            progress=progress,
+        ),
     ).frames[0]
     frames = [to_tensor(frame.convert("RGB")).mul(255).byte().permute(1, 2, 0) for frame in frames]
     frames = torch.stack(frames)

utils.py

@@ -0,0 +1,7 @@
+from gradio import Progress
+
+def create_progress_updater(start: int, total: int, desc: str, progress: Progress):
+    def updater(pipe, step, timestep, callback_kwargs):
+        progress((step + start, total), desc=desc)
+        return callback_kwargs
+    return updater