mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2025-03-15 14:09:36 +00:00
Add a total_steps value to sampler callback.
This commit is contained in:
parent
93c64afaa9
commit
908dc1d5a8
@ -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:
|
||||||
|
@ -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:
|
||||||
|
@ -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
4
nodes.py
@ -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,
|
||||||
|
Loading…
Reference in New Issue
Block a user