Skip to content

Commit

Permalink
Implement my alternative take on CFG++ as the euler_pp sampler.
Browse files Browse the repository at this point in the history
Add euler_ancestral_pp which is the ancestral version of euler with the
same modification.
  • Loading branch information
comfyanonymous committed Jun 25, 2024
1 parent 90aebb6 commit 69d710e
Show file tree
Hide file tree
Showing 3 changed files with 56 additions and 26 deletions.
55 changes: 54 additions & 1 deletion comfy/k_diffusion/sampling.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
from tqdm.auto import trange, tqdm

from . import utils

import comfy.model_patcher

def append_zero(x):
return torch.cat([x, x.new_zeros([1])])
Expand Down Expand Up @@ -945,3 +945,56 @@ def sample_ipndm_v(model, x, sigmas, extra_args=None, callback=None, disable=Non
buffer_model.append(d_cur.detach())

return x_next


@torch.no_grad()
def sample_euler_pp(model, x, sigmas, extra_args=None, callback=None, disable=None):
extra_args = {} if extra_args is None else extra_args

temp = [0]
def post_cfg_function(args):
temp[0] = args["uncond_denoised"]
return args["denoised"]

model_options = extra_args.get("model_options", {}).copy()
extra_args["model_options"] = comfy.model_patcher.set_model_options_post_cfg_function(model_options, post_cfg_function, disable_cfg1_optimization=True)

s_in = x.new_ones([x.shape[0]])
for i in trange(len(sigmas) - 1, disable=disable):
sigma_hat = sigmas[i]
denoised = model(x, sigma_hat * s_in, **extra_args)
d = to_d(x - denoised + temp[0], sigma_hat, denoised)
if callback is not None:
callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat, 'denoised': denoised})
dt = sigmas[i + 1] - sigma_hat
# Euler method
x = x + d * dt
return x

@torch.no_grad()
def sample_euler_ancestral_pp(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None):
"""Ancestral sampling with Euler method steps."""
extra_args = {} if extra_args is None else extra_args
noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler

temp = [0]
def post_cfg_function(args):
temp[0] = args["uncond_denoised"]
return args["denoised"]

model_options = extra_args.get("model_options", {}).copy()
extra_args["model_options"] = comfy.model_patcher.set_model_options_post_cfg_function(model_options, post_cfg_function, disable_cfg1_optimization=True)

s_in = x.new_ones([x.shape[0]])
for i in trange(len(sigmas) - 1, disable=disable):
denoised = model(x, sigmas[i] * s_in, **extra_args)
sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta)
if callback is not None:
callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
d = to_d(x - denoised + temp[0], sigmas[i], denoised)
# Euler method
dt = sigma_down - sigmas[i]
x = x + d * dt
if sigmas[i + 1] > 0:
x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up
return x
2 changes: 1 addition & 1 deletion comfy/samplers.py
Original file line number Diff line number Diff line change
Expand Up @@ -537,7 +537,7 @@ def max_denoise(self, model_wrap, sigmas):
sigma = float(sigmas[0])
return math.isclose(max_sigma, sigma, rel_tol=1e-05) or sigma > max_sigma

KSAMPLER_NAMES = ["euler", "euler_ancestral", "heun", "heunpp2","dpm_2", "dpm_2_ancestral",
KSAMPLER_NAMES = ["euler", "euler_pp", "euler_ancestral", "euler_ancestral_pp", "heun", "heunpp2","dpm_2", "dpm_2_ancestral",
"lms", "dpm_fast", "dpm_adaptive", "dpmpp_2s_ancestral", "dpmpp_sde", "dpmpp_sde_gpu",
"dpmpp_2m", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "dpmpp_3m_sde", "dpmpp_3m_sde_gpu", "ddpm", "lcm",
"ipndm", "ipndm_v"]
Expand Down
25 changes: 1 addition & 24 deletions comfy_extras/nodes_advanced_samplers.py
Original file line number Diff line number Diff line change
Expand Up @@ -82,29 +82,6 @@ def post_cfg_function(args):
x = denoised + sigmas[i + 1] * d
return x

@torch.no_grad()
def sample_euler_cfgpp_alt(model, x, sigmas, extra_args=None, callback=None, disable=None):
extra_args = {} if extra_args is None else extra_args

temp = [0]
def post_cfg_function(args):
temp[0] = args["uncond_denoised"]
return args["denoised"]

model_options = extra_args.get("model_options", {}).copy()
extra_args["model_options"] = comfy.model_patcher.set_model_options_post_cfg_function(model_options, post_cfg_function, disable_cfg1_optimization=True)

s_in = x.new_ones([x.shape[0]])
for i in trange(len(sigmas) - 1, disable=disable):
sigma_hat = sigmas[i]
denoised = model(x, sigma_hat * s_in, **extra_args)
d = to_d(x - denoised + temp[0], sigma_hat, denoised)
if callback is not None:
callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat, 'denoised': denoised})
dt = sigmas[i + 1] - sigma_hat
# Euler method
x = x + d * dt
return x

class SamplerEulerCFGpp:
@classmethod
Expand All @@ -122,7 +99,7 @@ def get_sampler(self, version):
if version == "regular":
sampler = comfy.samplers.KSAMPLER(sample_euler_cfgpp)
else:
sampler = comfy.samplers.KSAMPLER(sample_euler_cfgpp_alt)
sampler = comfy.samplers.ksampler("euler_pp")
return (sampler, )

NODE_CLASS_MAPPINGS = {
Expand Down

0 comments on commit 69d710e

Please sign in to comment.