From 689eb76f4df19908a52adf4b23a8863d39b74bcf Mon Sep 17 00:00:00 2001 From: Haoming Date: Tue, 9 Sep 2025 10:12:18 +0800 Subject: [PATCH] cfgpp --- backend/sampling/sampling_function.py | 5 +- modules/sd_samplers.py | 2 + modules/sd_samplers_cfg_denoiser.py | 3 +- modules/sd_samplers_cfgpp.py | 305 ++++++++++++++++++++++++++ modules_forge/forge_alter_samplers.py | 51 +++++ 5 files changed, 363 insertions(+), 3 deletions(-) create mode 100644 modules/sd_samplers_cfgpp.py create mode 100644 modules_forge/forge_alter_samplers.py diff --git a/backend/sampling/sampling_function.py b/backend/sampling/sampling_function.py index fcca658b..9c8511dd 100644 --- a/backend/sampling/sampling_function.py +++ b/backend/sampling/sampling_function.py @@ -328,7 +328,7 @@ def sampling_function_inner(model, x, timestep, uncond, cond, cond_scale, model_ return cfg_result -def sampling_function(self, denoiser_params, cond_scale, cond_composition): +def sampling_function(self, denoiser_params, cond_scale, cond_composition, extra_model_options=None): unet_patcher = self.inner_model.inner_model.forge_objects.unet model = unet_patcher.model control = unet_patcher.controlnet_linked_list @@ -337,7 +337,8 @@ def sampling_function(self, denoiser_params, cond_scale, cond_composition): timestep = denoiser_params.sigma uncond = compile_conditions(denoiser_params.text_uncond) cond = compile_weighted_conditions(denoiser_params.text_cond, cond_composition) - model_options = unet_patcher.model_options + model_options = (unet_patcher.model_options or {}).copy() + model_options.update(extra_model_options or {}) seed = self.p.seeds[0] if extra_concat_condition is not None: diff --git a/modules/sd_samplers.py b/modules/sd_samplers.py index 20c2da5e..26e57d99 100644 --- a/modules/sd_samplers.py +++ b/modules/sd_samplers.py @@ -4,6 +4,7 @@ import functools import logging from modules import sd_samplers_common, sd_samplers_kdiffusion, sd_samplers_lcm, sd_samplers_timesteps, sd_schedulers, shared +from modules_forge import forge_alter_samplers # imports for functions that previously were here and are used by other modules from modules.sd_samplers_common import sample_to_image, samples_to_image_grid # noqa: F401 @@ -12,6 +13,7 @@ all_samplers = [ *sd_samplers_kdiffusion.samplers_data_k_diffusion, *sd_samplers_timesteps.samplers_data_timesteps, *sd_samplers_lcm.samplers_data_lcm, + *forge_alter_samplers.samplers_data_alter, ] all_samplers_map = {x.name: x for x in all_samplers} diff --git a/modules/sd_samplers_cfg_denoiser.py b/modules/sd_samplers_cfg_denoiser.py index 10081f1e..81596f03 100644 --- a/modules/sd_samplers_cfg_denoiser.py +++ b/modules/sd_samplers_cfg_denoiser.py @@ -194,7 +194,8 @@ class CFGDenoiser(torch.nn.Module): if shared.opts.s_min_uncond_all: self.p.extra_generation_params["NGMS all steps"] = shared.opts.s_min_uncond_all - denoised, cond_pred, uncond_pred = sampling_function(self, denoiser_params=denoiser_params, cond_scale=cond_scale, cond_composition=cond_composition) + extra_model_options = kwargs.get("model_options", {}) + denoised, cond_pred, uncond_pred = sampling_function(self, denoiser_params=denoiser_params, cond_scale=cond_scale, cond_composition=cond_composition, extra_model_options=extra_model_options) if self.need_last_noise_uncond: self.last_noise_uncond = (x - uncond_pred) / sigma[:, None, None, None] diff --git a/modules/sd_samplers_cfgpp.py b/modules/sd_samplers_cfgpp.py new file mode 100644 index 00000000..bdbf8dec --- /dev/null +++ b/modules/sd_samplers_cfgpp.py @@ -0,0 +1,305 @@ +import torch +from k_diffusion.sampling import BrownianTreeNoiseSampler, default_noise_sampler, get_ancestral_step, to_d +from tqdm.auto import trange + +from backend.patcher.base import set_model_options_post_cfg_function + + +def _sigma_fn(t): + return t.neg().exp() + + +def _t_fn(sigma): + return sigma.log().neg() + + +@torch.no_grad() +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 + + 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"] = 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, sigma_hat, temp[0]) + if callback is not None: + callback( + { + "x": x, + "i": i, + "sigma": sigmas[i], + "sigma_hat": sigma_hat, + "denoised": denoised, + } + ) + x = denoised + d * sigmas[i + 1] + return x + + +@torch.no_grad() +def sample_euler_ancestral_cfg_pp(model, x, sigmas, extra_args=None, callback=None, disable=None, noise_sampler=None): + eta = 1.0 + s_noise = 1.0 + 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"] = 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, sigmas[i], temp[0]) + x = denoised + d * sigma_down + if sigmas[i + 1] > 0: + x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up + return x + + +@torch.no_grad() +def sample_dpmpp_sde_cfg_pp(model, x, sigmas, extra_args=None, callback=None, disable=None, noise_sampler=None): + eta = 1.0 + s_noise = 1.0 + r = 0.5 + + if len(sigmas) <= 1: + return x + + sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max() + noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=extra_args.get("seed", None), cpu=True) if noise_sampler is None else noise_sampler + 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"] = 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) + if callback is not None: + callback( + { + "x": x, + "i": i, + "sigma": sigmas[i], + "sigma_hat": sigmas[i], + "denoised": denoised, + } + ) + + if sigmas[i + 1] == 0: + d = to_d(x, sigmas[i], temp[0]) + x = denoised + d * sigmas[i + 1] + else: + t, t_next = _t_fn(sigmas[i]), _t_fn(sigmas[i + 1]) + h = t_next - t + s = t + h * r + fac = 1 / (2 * r) + + sd, su = get_ancestral_step(_sigma_fn(t), _sigma_fn(s), eta) + s_ = _t_fn(sd) + x_2 = (_sigma_fn(s_) / _sigma_fn(t)) * x - (t - s_).expm1() * denoised + x_2 = x_2 + noise_sampler(_sigma_fn(t), _sigma_fn(s)) * s_noise * su + denoised_2 = model(x_2, _sigma_fn(s) * s_in, **extra_args) + + sd, su = get_ancestral_step(_sigma_fn(t), _sigma_fn(t_next), eta) + denoised_d = (1 - fac) * temp[0] + fac * temp[0] + x = denoised_2 + to_d(x, sigmas[i], denoised_d) * sd + x = x + noise_sampler(_sigma_fn(t), _sigma_fn(t_next)) * s_noise * su + return x + + +@torch.no_grad() +def sample_dpmpp_2m_cfg_pp(model, x, sigmas, extra_args=None, callback=None, disable=None): + extra_args = {} if extra_args is None else extra_args + s_in = x.new_ones([x.shape[0]]) + + old_uncond_denoised = None + uncond_denoised = None + + def post_cfg_function(args): + nonlocal uncond_denoised + uncond_denoised = args["uncond_denoised"] + return args["denoised"] + + model_options = extra_args.get("model_options", {}).copy() + extra_args["model_options"] = set_model_options_post_cfg_function(model_options, post_cfg_function, disable_cfg1_optimization=True) + + for i in trange(len(sigmas) - 1, disable=disable): + denoised = model(x, sigmas[i] * s_in, **extra_args) + if callback is not None: + callback( + { + "x": x, + "i": i, + "sigma": sigmas[i], + "sigma_hat": sigmas[i], + "denoised": denoised, + } + ) + t, t_next = _t_fn(sigmas[i]), _t_fn(sigmas[i + 1]) + h = t_next - t + if old_uncond_denoised is None or sigmas[i + 1] == 0: + denoised_mix = -torch.exp(-h) * uncond_denoised + else: + h_last = t - _t_fn(sigmas[i - 1]) + r = h_last / h + denoised_mix = -torch.exp(-h) * uncond_denoised - torch.expm1(-h) * (1 / (2 * r)) * (denoised - old_uncond_denoised) + x = denoised + denoised_mix + torch.exp(-h) * x + old_uncond_denoised = uncond_denoised + return x + + +@torch.no_grad() +def sample_dpmpp_2m_sde_cfg_pp(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1.0, s_noise=1.0, noise_sampler=None, solver_type="midpoint"): + if len(sigmas) <= 1: + return x + + if solver_type not in {"heun", "midpoint"}: + raise ValueError('solver_type must be "heun" or "midpoint"') + + seed = extra_args.get("seed", None) + sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max() + noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=seed, cpu=True) if noise_sampler is None else noise_sampler + extra_args = {} if extra_args is None else extra_args + s_in = x.new_ones([x.shape[0]]) + + old_denoised = None + h_last = None + h = None + + 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"] = set_model_options_post_cfg_function(model_options, post_cfg_function, disable_cfg1_optimization=True) + + for i in trange(len(sigmas) - 1, disable=disable): + denoised = model(x, sigmas[i] * s_in, **extra_args) + if callback is not None: + callback({"x": x, "i": i, "sigma": sigmas[i], "sigma_hat": sigmas[i], "denoised": denoised}) + if sigmas[i + 1] == 0: + x = denoised + else: + t, s = -sigmas[i].log(), -sigmas[i + 1].log() + h = s - t + eta_h = eta * h + + x = sigmas[i + 1] / sigmas[i] * (-eta_h).exp() * (x + (denoised - temp[0])) + (-h - eta_h).expm1().neg() * denoised + + if old_denoised is not None: + r = h_last / h + if solver_type == "heun": + x = x + ((-h - eta_h).expm1().neg() / (-h - eta_h) + 1) * (1 / r) * (denoised - old_denoised) + elif solver_type == "midpoint": + x = x + 0.5 * (-h - eta_h).expm1().neg() * (1 / r) * (denoised - old_denoised) + + if eta: + x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * sigmas[i + 1] * (-2 * eta_h).expm1().neg().sqrt() * s_noise + + old_denoised = denoised + h_last = h + return x + + +@torch.no_grad() +def sample_dpmpp_3m_sde_cfg_pp(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1.0, s_noise=1.0, noise_sampler=None): + if len(sigmas) <= 1: + return x + + seed = extra_args.get("seed", None) + sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max() + noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=seed, cpu=True) if noise_sampler is None else noise_sampler + extra_args = {} if extra_args is None else extra_args + s_in = x.new_ones([x.shape[0]]) + + denoised_1, denoised_2 = None, None + h, h_1, h_2 = None, None, None + + 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"] = set_model_options_post_cfg_function(model_options, post_cfg_function, disable_cfg1_optimization=True) + + for i in trange(len(sigmas) - 1, disable=disable): + denoised = model(x, sigmas[i] * s_in, **extra_args) + if callback is not None: + callback( + { + "x": x, + "i": i, + "sigma": sigmas[i], + "sigma_hat": sigmas[i], + "denoised": denoised, + } + ) + if sigmas[i + 1] == 0: + x = denoised + else: + t, s = -sigmas[i].log(), -sigmas[i + 1].log() + h = s - t + h_eta = h * (eta + 1) + + x = torch.exp(-h_eta) * (x + (denoised - temp[0])) + (-h_eta).expm1().neg() * denoised + + if h_2 is not None: + r0 = h_1 / h + r1 = h_2 / h + d1_0 = (denoised - denoised_1) / r0 + d1_1 = (denoised_1 - denoised_2) / r1 + d1 = d1_0 + (d1_0 - d1_1) * r0 / (r0 + r1) + d2 = (d1_0 - d1_1) / (r0 + r1) + phi_2 = h_eta.neg().expm1() / h_eta + 1 + phi_3 = phi_2 / h_eta - 0.5 + x = x + phi_2 * d1 - phi_3 * d2 + elif h_1 is not None: + r = h_1 / h + d = (denoised - denoised_1) / r + phi_2 = h_eta.neg().expm1() / h_eta + 1 + x = x + phi_2 * d + + if eta: + x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * sigmas[i + 1] * (-2 * h * eta).expm1().neg().sqrt() * s_noise + + denoised_1, denoised_2 = denoised, denoised_1 + h_1, h_2 = h, h_1 + return x diff --git a/modules_forge/forge_alter_samplers.py b/modules_forge/forge_alter_samplers.py new file mode 100644 index 00000000..103cd89f --- /dev/null +++ b/modules_forge/forge_alter_samplers.py @@ -0,0 +1,51 @@ +import logging +from typing import Callable + +from modules import sd_samplers_cfgpp, sd_samplers_common, sd_samplers_kdiffusion + + +class AlterSampler(sd_samplers_kdiffusion.KDiffusionSampler): + def __init__(self, sd_model, sampler_name): + sampler_function: Callable = getattr(sd_samplers_cfgpp, f"sample_{sampler_name}", None) + if sampler_function is None: + raise ValueError(f"Unknown sampler: {sampler_name}") + + super().__init__(sampler_function, sd_model, None) + + def sample(self, p, *args, **kwargs): + if p.cfg_scale > 2.0: + logging.warning("CFG between 1.0 ~ 2.0 is recommended when using CFG++ samplers") + return super().sample(p, *args, **kwargs) + + def sample_img2img(self, p, *args, **kwargs): + if p.cfg_scale > 2.0: + logging.warning("CFG between 1.0 ~ 2.0 is recommended when using CFG++ samplers") + return super().sample_img2img(p, *args, **kwargs) + + +def build_constructor(sampler_key: str) -> Callable: + def constructor(model): + return AlterSampler(model, sampler_key) + + return constructor + + +def create_cfg_pp_sampler(sampler_name: str, sampler_key: str) -> "sd_samplers_common.SamplerData": + config = {} + base_name = sampler_name.removesuffix(" CFG++") + for name, _, _, params in sd_samplers_kdiffusion.samplers_k_diffusion: + if name == base_name: + config = params.copy() + break + + return sd_samplers_common.SamplerData(sampler_name, build_constructor(sampler_key=sampler_key), [sampler_key], config) + + +samplers_data_alter = [ + create_cfg_pp_sampler("DPM++ 2M CFG++", "dpmpp_2m_cfg_pp"), + create_cfg_pp_sampler("DPM++ SDE CFG++", "dpmpp_sde_cfg_pp"), + create_cfg_pp_sampler("DPM++ 2M SDE CFG++", "dpmpp_2m_sde_cfg_pp"), + create_cfg_pp_sampler("DPM++ 3M SDE CFG++", "dpmpp_3m_sde_cfg_pp"), + create_cfg_pp_sampler("Euler a CFG++", "euler_ancestral_cfg_pp"), + create_cfg_pp_sampler("Euler CFG++", "euler_cfg_pp"), +]