Add a total_steps value to sampler callback.

comfy/extra_samplers/uni_pc.py

@@ -767,7 +767,7 @@ class UniPC:
                                 model_x = self.model_fn(x, vec_t)
                             model_prev_list[-1] = model_x
                 if callback is not None:
-                    callback(step_index, model_prev_list[-1], x)
+                    callback(step_index, model_prev_list[-1], x, steps)
         else:
             raise NotImplementedError()
         if denoise_to_zero:

comfy/samplers.py

@@ -623,7 +623,8 @@ class KSampler:
 
                 ddim_callback = None
                 if callback is not None:
-                    ddim_callback = lambda pred_x0, i: callback(i, pred_x0, None)
+                    total_steps = len(timesteps) - 1
+                    ddim_callback = lambda pred_x0, i: callback(i, pred_x0, None, total_steps)
 
                 sampler = DDIMSampler(self.model, device=self.device)
                 sampler.make_schedule_timesteps(ddim_timesteps=timesteps, verbose=False)
@@ -654,13 +655,14 @@ class KSampler:
                 noise = noise * sigmas[0]
 
                 k_callback = None
+                total_steps = len(sigmas) - 1
                 if callback is not None:
-                    k_callback = lambda x: callback(x["i"], x["denoised"], x["x"])
+                    k_callback = lambda x: callback(x["i"], x["denoised"], x["x"], total_steps)
 
                 if latent_image is not None:
                     noise += latent_image
                 if self.sampler == "dpm_fast":
-                    samples = k_diffusion_sampling.sample_dpm_fast(self.model_k, noise, sigma_min, sigmas[0], self.steps, extra_args=extra_args, callback=k_callback, disable=disable_pbar)
+                    samples = k_diffusion_sampling.sample_dpm_fast(self.model_k, noise, sigma_min, sigmas[0], total_steps, extra_args=extra_args, callback=k_callback, disable=disable_pbar)
                 elif self.sampler == "dpm_adaptive":
                     samples = k_diffusion_sampling.sample_dpm_adaptive(self.model_k, noise, sigma_min, sigmas[0], extra_args=extra_args, callback=k_callback, disable=disable_pbar)
                 else:

comfy/utils.py

@@ -100,7 +100,9 @@ class ProgressBar:
         self.current = 0
         self.hook = PROGRESS_BAR_HOOK
 
-    def update_absolute(self, value):
+    def update_absolute(self, value, total=None):
+        if total is not None:
+            self.total = total
         if value > self.total:
             value = self.total
         self.current = value

nodes.py

@@ -816,8 +816,8 @@ def common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive,
         noise_mask = latent["noise_mask"]
 
     pbar = comfy.utils.ProgressBar(steps)
-    def callback(step, x0, x):
-        pbar.update_absolute(step + 1)
+    def callback(step, x0, x, total_steps):
+        pbar.update_absolute(step + 1, total_steps)
 
     samples = comfy.sample.sample(model, noise, steps, cfg, sampler_name, scheduler, positive, negative, latent_image,
                                   denoise=denoise, disable_noise=disable_noise, start_step=start_step, last_step=last_step,