diff --git a/comfy/extra_samplers/uni_pc.py b/comfy/extra_samplers/uni_pc.py index 78bab593..2ff10caf 100644 --- a/comfy/extra_samplers/uni_pc.py +++ b/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: diff --git a/comfy/samplers.py b/comfy/samplers.py index b30fc3d9..dcf93cca 100644 --- a/comfy/samplers.py +++ b/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: diff --git a/comfy/utils.py b/comfy/utils.py index 7f3c3978..f1ff9779 100644 --- a/comfy/utils.py +++ b/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 diff --git a/nodes.py b/nodes.py index 90c943fe..c2bc3685 100644 --- a/nodes.py +++ b/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,