298977887/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-03-15 14:09:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add a total_steps value to sampler callback.

Browse Source
This commit is contained in:
comfyanonymous 2023-05-03 12:58:10 -04:00
parent 93c64afaa9
commit 908dc1d5a8
4 changed files with 11 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
comfy/extra_samplers/uni_pc.py
View File

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

8
comfy/samplers.py
View File

@ -623,7 +623,8 @@ class KSampler:
ddim_callback = None ddim_callback = None
if callback is not 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 = DDIMSampler(self.model, device=self.device)
sampler.make_schedule_timesteps(ddim_timesteps=timesteps, verbose=False) sampler.make_schedule_timesteps(ddim_timesteps=timesteps, verbose=False)
@ -654,13 +655,14 @@ class KSampler:
noise = noise * sigmas[0] noise = noise * sigmas[0]
k_callback = None k_callback = None
total_steps = len(sigmas) - 1
if callback is not None: 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: if latent_image is not None:
noise += latent_image noise += latent_image
if self.sampler == "dpm_fast": 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": 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) 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: else:

4
comfy/utils.py
View File

@ -100,7 +100,9 @@ class ProgressBar:
self.current = 0 self.current = 0
self.hook = PROGRESS_BAR_HOOK 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: if value > self.total:
value = self.total value = self.total
self.current = value self.current = value

4
nodes.py
View File

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