mirrors/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-07-30 17:56:26 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add the denoising step to several samplers (#8780)

Browse Source
This commit is contained in:
chaObserv 2025-07-04 07:20:53 +08:00 committed by GitHub
parent f74fc4d927
commit f41f323c52
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 24 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

36
comfy/k_diffusion/sampling.py
View File

@ -412,9 +412,13 @@ def sample_lms(model, x, sigmas, extra_args=None, callback=None, disable=None, o
ds.pop(0) ds.pop(0)
if callback is not None: if callback is not None:
callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised}) callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
cur_order = min(i + 1, order) if sigmas[i + 1] == 0:
coeffs = [linear_multistep_coeff(cur_order, sigmas_cpu, i, j) for j in range(cur_order)] # Denoising step
x = x + sum(coeff * d for coeff, d in zip(coeffs, reversed(ds))) x = denoised
else:
cur_order = min(i + 1, order)
coeffs = [linear_multistep_coeff(cur_order, sigmas_cpu, i, j) for j in range(cur_order)]
x = x + sum(coeff * d for coeff, d in zip(coeffs, reversed(ds)))
return x return x
@ -1067,7 +1071,9 @@ def sample_ipndm(model, x, sigmas, extra_args=None, callback=None, disable=None,
d_cur = (x_cur - denoised) / t_cur d_cur = (x_cur - denoised) / t_cur
order = min(max_order, i+1) order = min(max_order, i+1)
if order == 1: # First Euler step. if t_next == 0: # Denoising step
x_next = denoised
elif order == 1: # First Euler step.
x_next = x_cur + (t_next - t_cur) * d_cur x_next = x_cur + (t_next - t_cur) * d_cur
elif order == 2: # Use one history point. elif order == 2: # Use one history point.
x_next = x_cur + (t_next - t_cur) * (3 * d_cur - buffer_model[-1]) / 2 x_next = x_cur + (t_next - t_cur) * (3 * d_cur - buffer_model[-1]) / 2
@ -1085,6 +1091,7 @@ def sample_ipndm(model, x, sigmas, extra_args=None, callback=None, disable=None,
return x_next return x_next
#From https://github.com/zju-pi/diff-sampler/blob/main/diff-solvers-main/solvers.py #From https://github.com/zju-pi/diff-sampler/blob/main/diff-solvers-main/solvers.py
#under Apache 2 license #under Apache 2 license
def sample_ipndm_v(model, x, sigmas, extra_args=None, callback=None, disable=None, max_order=4): def sample_ipndm_v(model, x, sigmas, extra_args=None, callback=None, disable=None, max_order=4):
@ -1108,7 +1115,9 @@ def sample_ipndm_v(model, x, sigmas, extra_args=None, callback=None, disable=Non
d_cur = (x_cur - denoised) / t_cur d_cur = (x_cur - denoised) / t_cur
order = min(max_order, i+1) order = min(max_order, i+1)
if order == 1: # First Euler step. if t_next == 0: # Denoising step
x_next = denoised
elif order == 1: # First Euler step.
x_next = x_cur + (t_next - t_cur) * d_cur x_next = x_cur + (t_next - t_cur) * d_cur
elif order == 2: # Use one history point. elif order == 2: # Use one history point.
h_n = (t_next - t_cur) h_n = (t_next - t_cur)
@ -1148,6 +1157,7 @@ def sample_ipndm_v(model, x, sigmas, extra_args=None, callback=None, disable=Non
return x_next return x_next
#From https://github.com/zju-pi/diff-sampler/blob/main/diff-solvers-main/solvers.py #From https://github.com/zju-pi/diff-sampler/blob/main/diff-solvers-main/solvers.py
#under Apache 2 license #under Apache 2 license
@torch.no_grad() @torch.no_grad()
@ -1198,6 +1208,7 @@ def sample_deis(model, x, sigmas, extra_args=None, callback=None, disable=None,
return x_next return x_next
@torch.no_grad() @torch.no_grad()
def sample_euler_cfg_pp(model, x, sigmas, extra_args=None, callback=None, disable=None): def sample_euler_cfg_pp(model, x, sigmas, extra_args=None, callback=None, disable=None):
extra_args = {} if extra_args is None else extra_args extra_args = {} if extra_args is None else extra_args
@ -1404,6 +1415,7 @@ def sample_res_multistep_ancestral(model, x, sigmas, extra_args=None, callback=N
def sample_res_multistep_ancestral_cfg_pp(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None): def sample_res_multistep_ancestral_cfg_pp(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None):
return res_multistep(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, s_noise=s_noise, noise_sampler=noise_sampler, eta=eta, cfg_pp=True) return res_multistep(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, s_noise=s_noise, noise_sampler=noise_sampler, eta=eta, cfg_pp=True)
@torch.no_grad() @torch.no_grad()
def sample_gradient_estimation(model, x, sigmas, extra_args=None, callback=None, disable=None, ge_gamma=2., cfg_pp=False): def sample_gradient_estimation(model, x, sigmas, extra_args=None, callback=None, disable=None, ge_gamma=2., cfg_pp=False):
"""Gradient-estimation sampler. Paper: https://openreview.net/pdf?id=o2ND9v0CeK""" """Gradient-estimation sampler. Paper: https://openreview.net/pdf?id=o2ND9v0CeK"""
@ -1430,19 +1442,19 @@ def sample_gradient_estimation(model, x, sigmas, extra_args=None, callback=None,
if callback is not None: if callback is not None:
callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised}) callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
dt = sigmas[i + 1] - sigmas[i] dt = sigmas[i + 1] - sigmas[i]
if i == 0: if sigmas[i + 1] == 0:
# Denoising step
x = denoised
else:
# Euler method # Euler method
if cfg_pp: if cfg_pp:
x = denoised + d * sigmas[i + 1] x = denoised + d * sigmas[i + 1]
else: else:
x = x + d * dt x = x + d * dt
else:
# Gradient estimation if i >= 1:
if cfg_pp: # Gradient estimation
d_bar = (ge_gamma - 1) * (d - old_d) d_bar = (ge_gamma - 1) * (d - old_d)
x = denoised + d * sigmas[i + 1] + d_bar * dt
else:
d_bar = ge_gamma * d + (1 - ge_gamma) * old_d
x = x + d_bar * dt x = x + d_bar * dt
old_d = d old_d = d
return x return x