1213 lines
56 KiB
Python
1213 lines
56 KiB
Python
import math
|
||
from functools import partial
|
||
|
||
import logging
|
||
from scipy import integrate
|
||
import torch
|
||
from torch import nn
|
||
import torchsde
|
||
from tqdm.auto import trange, tqdm
|
||
|
||
from comfy.k_diffusion import utils
|
||
from comfy.k_diffusion import deis
|
||
|
||
import comfy.model_patcher
|
||
import comfy.model_sampling
|
||
|
||
from comfy.samplers import KSamplerX0Inpaint, Sampler
|
||
|
||
|
||
# =====================================================================
|
||
# EasyCache state container
|
||
# ---------------------------------------------------------------------
|
||
# 参考 EasyCache4HunyuanVideo/easycache_sample_video.py：
|
||
# 1) cache = denoised - x （上一次实算时输出与输入的"变换向量"）
|
||
# 跳过实算时直接 denoised = x + cache 复用。
|
||
# 2) 用一阶线性灵敏度 k = |Δoutput| / |Δinput| 估计跳过该 step 的预测误差
|
||
# pred_change = k * |Δx| / |prev_output|
|
||
# 累加为 accumulated_error；< threshold 时继续跳过，否则强制实算并清零。
|
||
# 3) 前 ret_steps 步必算（启动保护），最后一步必算（收尾保护）。
|
||
# 由于一些 sampler（Heun/DPM-2 等）每 step 内会多次调用 model(...)，
|
||
# 这里通过 begin_step/model_call/end_step 三个接口把 EasyCache 状态绑定到外层
|
||
# step 索引 i：本 step 的所有 model 调用共享同一个 should_calc 决策。
|
||
# =====================================================================
|
||
class EasyCacheState:
|
||
def __init__(self, num_steps, threshold=0.025, ret_steps=5):
|
||
self.num_steps = int(num_steps)
|
||
self.threshold = float(threshold)
|
||
self.ret_steps = int(ret_steps)
|
||
|
||
self.cnt = 0
|
||
self.should_calc = True
|
||
self.accumulated_error = 0.0
|
||
|
||
self.cache = None # 上一次实算的 (denoised - x)
|
||
self.k = None # 灵敏度 |Δoutput|/|Δinput|
|
||
|
||
self.previous_raw_input = None
|
||
self.prev_prev_raw_input = None
|
||
self.previous_output = None
|
||
|
||
self._step_calc_done = False
|
||
|
||
def begin_step(self, i, total_steps, raw_input):
|
||
self.cnt = i
|
||
self._step_calc_done = False
|
||
# 启动 / 收尾保护
|
||
if self.cnt < self.ret_steps or self.cnt >= total_steps - 1:
|
||
self.should_calc = True
|
||
self.accumulated_error = 0.0
|
||
return
|
||
if (self.previous_raw_input is None or self.previous_output is None
|
||
or self.cache is None or self.k is None):
|
||
self.should_calc = True
|
||
return
|
||
raw_input_change = (raw_input - self.previous_raw_input).abs().mean()
|
||
output_norm = self.previous_output.abs().mean().clamp(min=1e-8)
|
||
pred_change = float(self.k) * float(raw_input_change / output_norm)
|
||
self.accumulated_error += pred_change
|
||
if self.accumulated_error < self.threshold:
|
||
self.should_calc = False
|
||
else:
|
||
self.should_calc = True
|
||
self.accumulated_error = 0.0
|
||
|
||
def model_call(self, model, x, sigma, **extra_args):
|
||
if (not self.should_calc) and (self.cache is not None):
|
||
return x + self.cache
|
||
denoised = model(x, sigma, **extra_args)
|
||
# 仅在每个 step 第一次实算时更新缓存与灵敏度
|
||
if not self._step_calc_done:
|
||
new_cache = denoised - x
|
||
if self.previous_output is not None and self.prev_prev_raw_input is not None:
|
||
output_change = (denoised - self.previous_output).abs().mean()
|
||
input_change = (self.previous_raw_input - self.prev_prev_raw_input).abs().mean().clamp(min=1e-8)
|
||
self.k = float(output_change / input_change)
|
||
self.cache = new_cache
|
||
self.previous_output = denoised.detach().clone()
|
||
self._step_calc_done = True
|
||
return denoised
|
||
|
||
def end_step(self, raw_input):
|
||
self.prev_prev_raw_input = self.previous_raw_input
|
||
self.previous_raw_input = raw_input.detach().clone()
|
||
|
||
|
||
class KSAMPLER(Sampler):
|
||
def __init__(self, sampler_function, threshold, ret_steps, extra_options={}, inpaint_options={}):
|
||
# parameters for EasyCache
|
||
self.threshold = threshold
|
||
self.ret_steps = ret_steps
|
||
self.sampler_function = sampler_function
|
||
self.extra_options = extra_options
|
||
self.inpaint_options = inpaint_options
|
||
|
||
def sample(self, model_wrap, sigmas, extra_args, callback, noise, latent_image=None, denoise_mask=None, disable_pbar=False):
|
||
extra_args["denoise_mask"] = denoise_mask
|
||
model_k = KSamplerX0Inpaint(model_wrap, sigmas)
|
||
model_k.latent_image = latent_image
|
||
if self.inpaint_options.get("random", False): #TODO: Should this be the default?
|
||
generator = torch.manual_seed(extra_args.get("seed", 41) + 1)
|
||
model_k.noise = torch.randn(noise.shape, generator=generator, device="cpu").to(noise.dtype).to(noise.device)
|
||
else:
|
||
model_k.noise = noise
|
||
|
||
noise = model_wrap.inner_model.model_sampling.noise_scaling(sigmas[0], noise, latent_image, self.max_denoise(model_wrap, sigmas))
|
||
|
||
k_callback = None
|
||
total_steps = len(sigmas) - 1
|
||
if callback is not None:
|
||
k_callback = lambda x: callback(x["i"], x["denoised"], x["x"], total_steps)
|
||
|
||
samples = self.sampler_function(model_k, noise, sigmas, threshold=self.threshold, ret_steps=self.ret_steps, extra_args=extra_args, callback=k_callback, disable=disable_pbar, **self.extra_options)
|
||
samples = model_wrap.inner_model.model_sampling.inverse_noise_scaling(sigmas[-1], samples)
|
||
return samples
|
||
|
||
def append_zero(x):
|
||
return torch.cat([x, x.new_zeros([1])])
|
||
|
||
def get_sigmas_karras(n, sigma_min, sigma_max, rho=7., device='cpu'):
|
||
"""Constructs the noise schedule of Karras et al. (2022)."""
|
||
ramp = torch.linspace(0, 1, n, device=device)
|
||
min_inv_rho = sigma_min ** (1 / rho)
|
||
max_inv_rho = sigma_max ** (1 / rho)
|
||
sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho
|
||
return append_zero(sigmas).to(device)
|
||
|
||
def get_sigmas_exponential(n, sigma_min, sigma_max, device='cpu'):
|
||
"""Constructs an exponential noise schedule."""
|
||
sigmas = torch.linspace(math.log(sigma_max), math.log(sigma_min), n, device=device).exp()
|
||
return append_zero(sigmas)
|
||
|
||
def get_sigmas_polyexponential(n, sigma_min, sigma_max, rho=1., device='cpu'):
|
||
"""Constructs an polynomial in log sigma noise schedule."""
|
||
ramp = torch.linspace(1, 0, n, device=device) ** rho
|
||
sigmas = torch.exp(ramp * (math.log(sigma_max) - math.log(sigma_min)) + math.log(sigma_min))
|
||
return append_zero(sigmas)
|
||
|
||
def get_sigmas_vp(n, beta_d=19.9, beta_min=0.1, eps_s=1e-3, device='cpu'):
|
||
"""Constructs a continuous VP noise schedule."""
|
||
t = torch.linspace(1, eps_s, n, device=device)
|
||
sigmas = torch.sqrt(torch.special.expm1(beta_d * t ** 2 / 2 + beta_min * t))
|
||
return append_zero(sigmas)
|
||
|
||
def get_sigmas_laplace(n, sigma_min, sigma_max, mu=0., beta=0.5, device='cpu'):
|
||
"""Constructs the noise schedule proposed by Tiankai et al. (2024). """
|
||
epsilon = 1e-5 # avoid log(0)
|
||
x = torch.linspace(0, 1, n, device=device)
|
||
clamp = lambda x: torch.clamp(x, min=sigma_min, max=sigma_max)
|
||
lmb = mu - beta * torch.sign(0.5-x) * torch.log(1 - 2 * torch.abs(0.5-x) + epsilon)
|
||
sigmas = clamp(torch.exp(lmb))
|
||
return sigmas
|
||
|
||
def to_d(x, sigma, denoised):
|
||
"""Converts a denoiser output to a Karras ODE derivative."""
|
||
return (x - denoised) / utils.append_dims(sigma, x.ndim)
|
||
|
||
def get_ancestral_step(sigma_from, sigma_to, eta=1.):
|
||
"""Calculates the noise level (sigma_down) to step down to and the amount
|
||
of noise to add (sigma_up) when doing an ancestral sampling step."""
|
||
if not eta:
|
||
return sigma_to, 0.
|
||
sigma_up = min(sigma_to, eta * (sigma_to ** 2 * (sigma_from ** 2 - sigma_to ** 2) / sigma_from ** 2) ** 0.5)
|
||
sigma_down = (sigma_to ** 2 - sigma_up ** 2) ** 0.5
|
||
return sigma_down, sigma_up
|
||
|
||
def default_noise_sampler(x, seed=None):
|
||
if seed is not None:
|
||
generator = torch.Generator(device=x.device)
|
||
generator.manual_seed(seed)
|
||
else:
|
||
generator = None
|
||
|
||
return lambda sigma, sigma_next: torch.randn(x.size(), dtype=x.dtype, layout=x.layout, device=x.device, generator=generator)
|
||
|
||
class BatchedBrownianTree:
|
||
"""A wrapper around torchsde.BrownianTree that enables batches of entropy."""
|
||
|
||
def __init__(self, x, t0, t1, seed=None, **kwargs):
|
||
self.cpu_tree = kwargs.pop("cpu", True)
|
||
t0, t1, self.sign = self.sort(t0, t1)
|
||
w0 = kwargs.pop('w0', None)
|
||
if w0 is None:
|
||
w0 = torch.zeros_like(x)
|
||
self.batched = False
|
||
if seed is None:
|
||
seed = (torch.randint(0, 2 ** 63 - 1, ()).item(),)
|
||
elif isinstance(seed, (tuple, list)):
|
||
if len(seed) != x.shape[0]:
|
||
raise ValueError("Passing a list or tuple of seeds to BatchedBrownianTree requires a length matching the batch size.")
|
||
self.batched = True
|
||
w0 = w0[0]
|
||
else:
|
||
seed = (seed,)
|
||
if self.cpu_tree:
|
||
t0, w0, t1 = t0.detach().cpu(), w0.detach().cpu(), t1.detach().cpu()
|
||
self.trees = tuple(torchsde.BrownianTree(t0, w0, t1, entropy=s, **kwargs) for s in seed)
|
||
|
||
@staticmethod
|
||
def sort(a, b):
|
||
return (a, b, 1) if a < b else (b, a, -1)
|
||
|
||
def __call__(self, t0, t1):
|
||
t0, t1, sign = self.sort(t0, t1)
|
||
device, dtype = t0.device, t0.dtype
|
||
if self.cpu_tree:
|
||
t0, t1 = t0.detach().cpu().float(), t1.detach().cpu().float()
|
||
w = torch.stack([tree(t0, t1) for tree in self.trees]).to(device=device, dtype=dtype) * (self.sign * sign)
|
||
return w if self.batched else w[0]
|
||
|
||
class BrownianTreeNoiseSampler:
|
||
"""A noise sampler backed by a torchsde.BrownianTree.
|
||
|
||
Args:
|
||
x (Tensor): The tensor whose shape, device and dtype to use to generate
|
||
random samples.
|
||
sigma_min (float): The low end of the valid interval.
|
||
sigma_max (float): The high end of the valid interval.
|
||
seed (int or List[int]): The random seed. If a list of seeds is
|
||
supplied instead of a single integer, then the noise sampler will
|
||
use one BrownianTree per batch item, each with its own seed.
|
||
transform (callable): A function that maps sigma to the sampler's
|
||
internal timestep.
|
||
"""
|
||
|
||
def __init__(self, x, sigma_min, sigma_max, seed=None, transform=lambda x: x, cpu=False):
|
||
self.transform = transform
|
||
t0, t1 = self.transform(torch.as_tensor(sigma_min)), self.transform(torch.as_tensor(sigma_max))
|
||
self.tree = BatchedBrownianTree(x, t0, t1, seed, cpu=cpu)
|
||
|
||
def __call__(self, sigma, sigma_next):
|
||
t0, t1 = self.transform(torch.as_tensor(sigma)), self.transform(torch.as_tensor(sigma_next))
|
||
return self.tree(t0, t1) / (t1 - t0).abs().sqrt()
|
||
|
||
def sigma_to_half_log_snr(sigma, model_sampling):
|
||
"""Convert sigma to half-logSNR log(alpha_t / sigma_t)."""
|
||
if isinstance(model_sampling, comfy.model_sampling.CONST):
|
||
# log((1 - t) / t) = log((1 - sigma) / sigma)
|
||
return sigma.logit().neg()
|
||
return sigma.log().neg()
|
||
|
||
def half_log_snr_to_sigma(half_log_snr, model_sampling):
|
||
"""Convert half-logSNR log(alpha_t / sigma_t) to sigma."""
|
||
if isinstance(model_sampling, comfy.model_sampling.CONST):
|
||
# 1 / (1 + exp(half_log_snr))
|
||
return half_log_snr.neg().sigmoid()
|
||
return half_log_snr.neg().exp()
|
||
|
||
def offset_first_sigma_for_snr(sigmas, model_sampling, percent_offset=1e-4):
|
||
"""Adjust the first sigma to avoid invalid logSNR."""
|
||
if len(sigmas) <= 1:
|
||
return sigmas
|
||
if isinstance(model_sampling, comfy.model_sampling.CONST):
|
||
if sigmas[0] >= 1:
|
||
sigmas = sigmas.clone()
|
||
sigmas[0] = model_sampling.percent_to_sigma(percent_offset)
|
||
return sigmas
|
||
|
||
def ei_h_phi_1(h: torch.Tensor) -> torch.Tensor:
|
||
"""Compute the result of h*phi_1(h) in exponential integrator methods."""
|
||
return torch.expm1(h)
|
||
|
||
def ei_h_phi_2(h: torch.Tensor) -> torch.Tensor:
|
||
"""Compute the result of h*phi_2(h) in exponential integrator methods."""
|
||
return (torch.expm1(h) - h) / h
|
||
|
||
def linear_multistep_coeff(order, t, i, j):
|
||
if order - 1 > i:
|
||
raise ValueError(f'Order {order} too high for step {i}')
|
||
def fn(tau):
|
||
prod = 1.
|
||
for k in range(order):
|
||
if j == k:
|
||
continue
|
||
prod *= (tau - t[i - k]) / (t[i - j] - t[i - k])
|
||
return prod
|
||
return integrate.quad(fn, t[i], t[i + 1], epsrel=1e-4)[0]
|
||
|
||
class PIDStepSizeController:
|
||
"""A PID controller for ODE adaptive step size control."""
|
||
def __init__(self, h, pcoeff, icoeff, dcoeff, order=1, accept_safety=0.81, eps=1e-8):
|
||
self.h = h
|
||
self.b1 = (pcoeff + icoeff + dcoeff) / order
|
||
self.b2 = -(pcoeff + 2 * dcoeff) / order
|
||
self.b3 = dcoeff / order
|
||
self.accept_safety = accept_safety
|
||
self.eps = eps
|
||
self.errs = []
|
||
|
||
def limiter(self, x):
|
||
return 1 + math.atan(x - 1)
|
||
|
||
def propose_step(self, error):
|
||
inv_error = 1 / (float(error) + self.eps)
|
||
if not self.errs:
|
||
self.errs = [inv_error, inv_error, inv_error]
|
||
self.errs[0] = inv_error
|
||
factor = self.errs[0] ** self.b1 * self.errs[1] ** self.b2 * self.errs[2] ** self.b3
|
||
factor = self.limiter(factor)
|
||
accept = factor >= self.accept_safety
|
||
if accept:
|
||
self.errs[2] = self.errs[1]
|
||
self.errs[1] = self.errs[0]
|
||
self.h *= factor
|
||
return accept
|
||
|
||
class DPMSolver(nn.Module):
|
||
"""DPM-Solver. See https://arxiv.org/abs/2206.00927."""
|
||
|
||
def __init__(self, model,extra_args=None, eps_callback=None, info_callback=None):
|
||
super().__init__()
|
||
self.model = model
|
||
self.extra_args = {} if extra_args is None else extra_args
|
||
self.eps_callback = eps_callback
|
||
self.info_callback = info_callback
|
||
|
||
def t(self, sigma):
|
||
return -sigma.log()
|
||
|
||
def sigma(self, t):
|
||
return t.neg().exp()
|
||
|
||
def eps(self, eps_cache, key, x, t, final_step = False, *args, **kwargs):
|
||
if key in eps_cache:
|
||
return eps_cache[key], eps_cache
|
||
sigma = self.sigma(t) * x.new_ones([x.shape[0]])
|
||
model_output = self.model(x, sigma, *args, **self.extra_args, **kwargs)
|
||
eps = (x - model_output) / self.sigma(t)
|
||
if self.eps_callback is not None:
|
||
self.eps_callback()
|
||
return eps, {key: eps, **eps_cache}
|
||
|
||
def dpm_solver_1_step(self, x, t, t_next, eps_cache=None,final_step=False):
|
||
eps_cache = {} if eps_cache is None else eps_cache
|
||
h = t_next - t
|
||
eps, eps_cache = self.eps(eps_cache, 'eps', x, t,final_step=final_step)
|
||
x_1 = x - self.sigma(t_next) * h.expm1() * eps
|
||
return x_1, eps_cache
|
||
|
||
def dpm_solver_2_step(self, x, t, t_next, r1=1 / 2, eps_cache=None,final_step=False):
|
||
eps_cache = {} if eps_cache is None else eps_cache
|
||
h = t_next - t
|
||
eps, eps_cache = self.eps(eps_cache, 'eps', x, t)
|
||
s1 = t + r1 * h
|
||
u1 = x - self.sigma(s1) * (r1 * h).expm1() * eps
|
||
eps_r1, eps_cache = self.eps(eps_cache, 'eps_r1', u1, s1,final_step=final_step)
|
||
x_2 = x - self.sigma(t_next) * h.expm1() * eps - self.sigma(t_next) / (2 * r1) * h.expm1() * (eps_r1 - eps)
|
||
return x_2, eps_cache
|
||
|
||
def dpm_solver_3_step(self, x, t, t_next, r1=1 / 3, r2=2 / 3, eps_cache=None,final_step=False):
|
||
eps_cache = {} if eps_cache is None else eps_cache
|
||
h = t_next - t
|
||
eps, eps_cache = self.eps(eps_cache, 'eps', x, t)
|
||
s1 = t + r1 * h
|
||
s2 = t + r2 * h
|
||
u1 = x - self.sigma(s1) * (r1 * h).expm1() * eps
|
||
eps_r1, eps_cache = self.eps(eps_cache, 'eps_r1', u1, s1)
|
||
u2 = x - self.sigma(s2) * (r2 * h).expm1() * eps - self.sigma(s2) * (r2 / r1) * ((r2 * h).expm1() / (r2 * h) - 1) * (eps_r1 - eps)
|
||
eps_r2, eps_cache = self.eps(eps_cache, 'eps_r2', u2, s2,final_step=final_step)
|
||
x_3 = x - self.sigma(t_next) * h.expm1() * eps - self.sigma(t_next) / r2 * (h.expm1() / h - 1) * (eps_r2 - eps)
|
||
return x_3, eps_cache
|
||
|
||
def dpm_solver_fast(self, x, t_start, t_end, nfe, eta=0., s_noise=1., noise_sampler=None):
|
||
noise_sampler = default_noise_sampler(x, seed=self.extra_args.get("seed", None)) if noise_sampler is None else noise_sampler
|
||
if not t_end > t_start and eta:
|
||
raise ValueError('eta must be 0 for reverse sampling')
|
||
|
||
m = math.floor(nfe / 3) + 1
|
||
ts = torch.linspace(t_start, t_end, m + 1, device=x.device)
|
||
|
||
if nfe % 3 == 0:
|
||
orders = [3] * (m - 2) + [2, 1]
|
||
else:
|
||
orders = [3] * (m - 1) + [nfe % 3]
|
||
|
||
for i in range(len(orders)):
|
||
eps_cache = {}
|
||
t, t_next = ts[i], ts[i + 1]
|
||
if eta:
|
||
sd, su = get_ancestral_step(self.sigma(t), self.sigma(t_next), eta)
|
||
t_next_ = torch.minimum(t_end, self.t(sd))
|
||
su = (self.sigma(t_next) ** 2 - self.sigma(t_next_) ** 2) ** 0.5
|
||
else:
|
||
t_next_, su = t_next, 0.
|
||
|
||
final_step = True if t_next_ > t_end - 1e-5 else False
|
||
|
||
eps, eps_cache = self.eps(eps_cache, 'eps', x, t, final_step=final_step)
|
||
denoised = x - self.sigma(t) * eps
|
||
if self.info_callback is not None:
|
||
self.info_callback({'x': x, 'i': i, 't': ts[i], 't_up': t, 'denoised': denoised})
|
||
|
||
if orders[i] == 1:
|
||
x, eps_cache = self.dpm_solver_1_step(x, t, t_next_, eps_cache=eps_cache, final_step=final_step)
|
||
elif orders[i] == 2:
|
||
x, eps_cache = self.dpm_solver_2_step(x, t, t_next_, eps_cache=eps_cache, final_step=final_step)
|
||
else:
|
||
x, eps_cache = self.dpm_solver_3_step(x, t, t_next_, eps_cache=eps_cache, final_step=final_step)
|
||
|
||
x = x + su * s_noise * noise_sampler(self.sigma(t), self.sigma(t_next))
|
||
|
||
return x
|
||
|
||
def dpm_solver_adaptive(self, x, t_start, t_end, order=3, rtol=0.05, atol=0.0078, h_init=0.05, pcoeff=0., icoeff=1., dcoeff=0., accept_safety=0.81, eta=0., s_noise=1., noise_sampler=None):
|
||
noise_sampler = default_noise_sampler(x, seed=self.extra_args.get("seed", None)) if noise_sampler is None else noise_sampler
|
||
if order not in {2, 3}:
|
||
raise ValueError('order should be 2 or 3')
|
||
forward = t_end > t_start
|
||
if not forward and eta:
|
||
raise ValueError('eta must be 0 for reverse sampling')
|
||
h_init = abs(h_init) * (1 if forward else -1)
|
||
atol = torch.tensor(atol)
|
||
rtol = torch.tensor(rtol)
|
||
s = t_start
|
||
x_prev = x
|
||
accept = True
|
||
pid = PIDStepSizeController(h_init, pcoeff, icoeff, dcoeff, 1.5 if eta else order, accept_safety)
|
||
info = {'steps': 0, 'nfe': 0, 'n_accept': 0, 'n_reject': 0}
|
||
|
||
while s < t_end - 1e-5 if forward else s > t_end + 1e-5:
|
||
eps_cache = {}
|
||
t = torch.minimum(t_end, s + pid.h) if forward else torch.maximum(t_end, s + pid.h)
|
||
if eta:
|
||
sd, su = get_ancestral_step(self.sigma(s), self.sigma(t), eta)
|
||
t_ = torch.minimum(t_end, self.t(sd))
|
||
su = (self.sigma(t) ** 2 - self.sigma(t_) ** 2) ** 0.5
|
||
else:
|
||
t_, su = t, 0.
|
||
|
||
if forward:
|
||
final_step = True if t_ > t_end - 1e-5 else False
|
||
else:
|
||
final_step = True if t_ < t_end + 1e-5 else False
|
||
|
||
eps, eps_cache = self.eps(eps_cache, 'eps', x, s, final_step=final_step)
|
||
denoised = x - self.sigma(s) * eps
|
||
|
||
if order == 2:
|
||
x_low, eps_cache = self.dpm_solver_1_step(x, s, t_, eps_cache=eps_cache, final_step=final_step)
|
||
x_high, eps_cache = self.dpm_solver_2_step(x, s, t_, eps_cache=eps_cache, final_step=final_step)
|
||
else:
|
||
x_low, eps_cache = self.dpm_solver_2_step(x, s, t_, r1=1 / 3, eps_cache=eps_cache, final_step=final_step)
|
||
x_high, eps_cache = self.dpm_solver_3_step(x, s, t_, eps_cache=eps_cache, final_step=final_step)
|
||
delta = torch.maximum(atol, rtol * torch.maximum(x_low.abs(), x_prev.abs()))
|
||
error = torch.linalg.norm((x_low - x_high) / delta) / x.numel() ** 0.5
|
||
accept = pid.propose_step(error)
|
||
if accept:
|
||
x_prev = x_low
|
||
x = x_high + su * s_noise * noise_sampler(self.sigma(s), self.sigma(t))
|
||
s = t
|
||
info['n_accept'] += 1
|
||
else:
|
||
info['n_reject'] += 1
|
||
info['nfe'] += order
|
||
info['steps'] += 1
|
||
|
||
if self.info_callback is not None:
|
||
self.info_callback({'x': x, 'i': info['steps'] - 1, 't': s, 't_up': s, 'denoised': denoised, 'error': error, 'h': pid.h, **info})
|
||
|
||
return x, info
|
||
|
||
@torch.no_grad()
|
||
def sample_euler(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1., threshold=0.025, ret_steps=5):
|
||
"""Implements Algorithm 2 (Euler steps) from Karras et al. (2022)."""
|
||
extra_args = {} if extra_args is None else extra_args
|
||
s_in = x.new_ones([x.shape[0]])
|
||
#------------add for EasyCache------------
|
||
total_steps = len(sigmas) - 1
|
||
ec = EasyCacheState(total_steps, threshold=threshold, ret_steps=ret_steps)
|
||
#-----------------------------------------
|
||
for i in trange(total_steps, disable=disable):
|
||
if s_churn > 0:
|
||
gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.
|
||
sigma_hat = sigmas[i] * (gamma + 1)
|
||
else:
|
||
gamma = 0
|
||
sigma_hat = sigmas[i]
|
||
if gamma > 0:
|
||
eps = torch.randn_like(x) * s_noise
|
||
x = x + eps * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5
|
||
#------------add for EasyCache------------
|
||
x_step_in = x
|
||
ec.begin_step(i, total_steps, x_step_in)
|
||
denoised = ec.model_call(model, x, sigma_hat * s_in, **extra_args)
|
||
#-----------------------------------------
|
||
d = to_d(x, 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
|
||
ec.end_step(x_step_in)
|
||
return x
|
||
|
||
@torch.no_grad()
|
||
def sample_heun(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1., threshold=0.025, ret_steps=5):
|
||
"""Implements Algorithm 2 (Heun steps) from Karras et al. (2022)."""
|
||
extra_args = {} if extra_args is None else extra_args
|
||
s_in = x.new_ones([x.shape[0]])
|
||
#------------add for EasyCache------------
|
||
total_steps = len(sigmas) - 1
|
||
ec = EasyCacheState(total_steps, threshold=threshold, ret_steps=ret_steps)
|
||
#-----------------------------------------
|
||
for i in trange(total_steps, disable=disable):
|
||
if s_churn > 0:
|
||
gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.
|
||
sigma_hat = sigmas[i] * (gamma + 1)
|
||
else:
|
||
gamma = 0
|
||
sigma_hat = sigmas[i]
|
||
|
||
sigma_hat = sigmas[i] * (gamma + 1)
|
||
if gamma > 0:
|
||
eps = torch.randn_like(x) * s_noise
|
||
x = x + eps * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5
|
||
#------------add for EasyCache------------
|
||
x_step_in = x
|
||
ec.begin_step(i, total_steps, x_step_in)
|
||
denoised = ec.model_call(model, x, sigma_hat * s_in, **extra_args)
|
||
#-----------------------------------------
|
||
d = to_d(x, 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
|
||
if sigmas[i + 1] == 0:
|
||
# Euler method
|
||
x = x + d * dt
|
||
else:
|
||
# Heun's method
|
||
x_2 = x + d * dt
|
||
denoised_2 = ec.model_call(model, x_2, sigmas[i + 1] * s_in, **extra_args)
|
||
d_2 = to_d(x_2, sigmas[i + 1], denoised_2)
|
||
d_prime = (d + d_2) / 2
|
||
x = x + d_prime * dt
|
||
ec.end_step(x_step_in)
|
||
return x
|
||
|
||
@torch.no_grad()
|
||
def sample_dpm_2(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1., threshold=0.025, ret_steps=5):
|
||
"""A sampler inspired by DPM-Solver-2 and Algorithm 2 from Karras et al. (2022)."""
|
||
extra_args = {} if extra_args is None else extra_args
|
||
s_in = x.new_ones([x.shape[0]])
|
||
#------------add for EasyCache------------
|
||
total_steps = len(sigmas) - 1
|
||
ec = EasyCacheState(total_steps, threshold=threshold, ret_steps=ret_steps)
|
||
#-----------------------------------------
|
||
for i in trange(total_steps, disable=disable):
|
||
if s_churn > 0:
|
||
gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.
|
||
sigma_hat = sigmas[i] * (gamma + 1)
|
||
else:
|
||
gamma = 0
|
||
sigma_hat = sigmas[i]
|
||
|
||
if gamma > 0:
|
||
eps = torch.randn_like(x) * s_noise
|
||
x = x + eps * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5
|
||
#------------add for EasyCache------------
|
||
x_step_in = x
|
||
ec.begin_step(i, total_steps, x_step_in)
|
||
denoised = ec.model_call(model, x, sigma_hat * s_in, **extra_args)
|
||
#-----------------------------------------
|
||
d = to_d(x, sigma_hat, denoised)
|
||
if callback is not None:
|
||
callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat, 'denoised': denoised})
|
||
if sigmas[i + 1] == 0:
|
||
# Euler method
|
||
dt = sigmas[i + 1] - sigma_hat
|
||
x = x + d * dt
|
||
else:
|
||
# DPM-Solver-2
|
||
sigma_mid = sigma_hat.log().lerp(sigmas[i + 1].log(), 0.5).exp()
|
||
dt_1 = sigma_mid - sigma_hat
|
||
dt_2 = sigmas[i + 1] - sigma_hat
|
||
x_2 = x + d * dt_1
|
||
denoised_2 = ec.model_call(model, x_2, sigma_mid * s_in, **extra_args)
|
||
d_2 = to_d(x_2, sigma_mid, denoised_2)
|
||
x = x + d_2 * dt_2
|
||
ec.end_step(x_step_in)
|
||
return x
|
||
|
||
@torch.no_grad()
|
||
def sample_lms(model, x, sigmas, extra_args=None, callback=None, disable=None, order=4, threshold=0.025, ret_steps=5):
|
||
extra_args = {} if extra_args is None else extra_args
|
||
s_in = x.new_ones([x.shape[0]])
|
||
sigmas_cpu = sigmas.detach().cpu().numpy()
|
||
ds = []
|
||
#------------add for EasyCache------------
|
||
total_steps = len(sigmas) - 1
|
||
ec = EasyCacheState(total_steps, threshold=threshold, ret_steps=ret_steps)
|
||
#-----------------------------------------
|
||
for i in trange(total_steps, disable=disable):
|
||
#------------add for EasyCache------------
|
||
x_step_in = x
|
||
ec.begin_step(i, total_steps, x_step_in)
|
||
denoised = ec.model_call(model, x, sigmas[i] * s_in, **extra_args)
|
||
#-----------------------------------------
|
||
d = to_d(x, sigmas[i], denoised)
|
||
ds.append(d)
|
||
if len(ds) > order:
|
||
ds.pop(0)
|
||
if callback is not None:
|
||
callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
|
||
if sigmas[i + 1] == 0:
|
||
# Denoising step
|
||
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)))
|
||
ec.end_step(x_step_in)
|
||
return x
|
||
|
||
@torch.no_grad()
|
||
def sample_dpm_fast(model, x, sigma_min, sigma_max, n, extra_args=None, callback=None, disable=None, eta=0., s_noise=1., noise_sampler=None, threshold=0.025, ret_steps=5):
|
||
"""DPM-Solver-Fast (fixed step size). EasyCache 注入到底层 model 调用。"""
|
||
if sigma_min <= 0 or sigma_max <= 0:
|
||
raise ValueError('sigma_min and sigma_max must not be 0')
|
||
#------------add for EasyCache------------
|
||
ec = EasyCacheState(n, threshold=threshold, ret_steps=ret_steps)
|
||
class _ECWrappedModel:
|
||
def __init__(self, base_model):
|
||
self._base = base_model
|
||
def __call__(self, x, sigma, **kw):
|
||
return ec.model_call(self._base, x, sigma, **kw)
|
||
cur_step = {"i": 0}
|
||
#-----------------------------------------
|
||
with tqdm(total=n, disable=disable) as pbar:
|
||
wrapped = _ECWrappedModel(model)
|
||
dpm_solver = DPMSolver(wrapped, extra_args=extra_args, eps_callback=pbar.update)
|
||
def _info(info):
|
||
i = info.get("i", cur_step["i"])
|
||
ec.begin_step(i, n, x)
|
||
ec.end_step(x)
|
||
cur_step["i"] = i + 1
|
||
if callback is not None:
|
||
callback({'sigma': dpm_solver.sigma(info['t']), 'sigma_hat': dpm_solver.sigma(info['t_up']), **info})
|
||
dpm_solver.info_callback = _info
|
||
return dpm_solver.dpm_solver_fast(x, dpm_solver.t(torch.tensor(sigma_max)), dpm_solver.t(torch.tensor(sigma_min)), n, eta, s_noise, noise_sampler)
|
||
|
||
@torch.no_grad()
|
||
def sample_dpm_adaptive(model, x, sigma_min, sigma_max, n, extra_args=None, callback=None, disable=None, order=3, rtol=0.05, atol=0.0078, h_init=0.05, pcoeff=0., icoeff=1., dcoeff=0., accept_safety=0.81, eta=0., s_noise=1., noise_sampler=None, return_info=False, threshold=0.025, ret_steps=5):
|
||
"""DPM-Solver-12 / 23 (adaptive step size). EasyCache 注入到底层 model 调用。"""
|
||
if sigma_min <= 0 or sigma_max <= 0:
|
||
raise ValueError('sigma_min and sigma_max must not be 0')
|
||
#------------add for EasyCache------------
|
||
ec = EasyCacheState(max(n, 1), threshold=threshold, ret_steps=ret_steps)
|
||
class _ECWrappedModel:
|
||
def __init__(self, base_model):
|
||
self._base = base_model
|
||
def __call__(self, x, sigma, **kw):
|
||
return ec.model_call(self._base, x, sigma, **kw)
|
||
cur_step = {"i": 0}
|
||
#-----------------------------------------
|
||
with tqdm(disable=disable) as pbar:
|
||
wrapped = _ECWrappedModel(model)
|
||
dpm_solver = DPMSolver(wrapped, extra_args=extra_args, eps_callback=pbar.update)
|
||
def _info(info):
|
||
i = info.get("i", cur_step["i"])
|
||
ec.begin_step(i, max(n, 1), x)
|
||
ec.end_step(x)
|
||
cur_step["i"] = i + 1
|
||
if callback is not None:
|
||
callback({'sigma': dpm_solver.sigma(info['t']), 'sigma_hat': dpm_solver.sigma(info['t_up']), **info})
|
||
dpm_solver.info_callback = _info
|
||
x_out, info = dpm_solver.dpm_solver_adaptive(x, dpm_solver.t(torch.tensor(sigma_max)), dpm_solver.t(torch.tensor(sigma_min)), order, rtol, atol, h_init, pcoeff, icoeff, dcoeff, accept_safety, eta, s_noise, noise_sampler)
|
||
if return_info:
|
||
return x_out, info
|
||
return x_out
|
||
|
||
@torch.no_grad()
|
||
def sample_dpmpp_2m(model, x, sigmas, extra_args=None, callback=None, disable=None, threshold=0.025, ret_steps=5):
|
||
"""DPM-Solver++(2M)."""
|
||
extra_args = {} if extra_args is None else extra_args
|
||
s_in = x.new_ones([x.shape[0]])
|
||
sigma_fn = lambda t: t.neg().exp()
|
||
t_fn = lambda sigma: sigma.log().neg()
|
||
old_denoised = None
|
||
#------------add for EasyCache------------
|
||
total_steps = len(sigmas) - 1
|
||
ec = EasyCacheState(total_steps, threshold=threshold, ret_steps=ret_steps)
|
||
#-----------------------------------------
|
||
for i in trange(total_steps, disable=disable):
|
||
#------------add for EasyCache------------
|
||
x_step_in = x
|
||
ec.begin_step(i, total_steps, x_step_in)
|
||
denoised = ec.model_call(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_denoised is None or sigmas[i + 1] == 0:
|
||
x = (sigma_fn(t_next) / sigma_fn(t)) * x - (-h).expm1() * denoised
|
||
else:
|
||
h_last = t - t_fn(sigmas[i - 1])
|
||
r = h_last / h
|
||
denoised_d = (1 + 1 / (2 * r)) * denoised - (1 / (2 * r)) * old_denoised
|
||
x = (sigma_fn(t_next) / sigma_fn(t)) * x - (-h).expm1() * denoised_d
|
||
old_denoised = denoised
|
||
ec.end_step(x_step_in)
|
||
return x
|
||
|
||
@torch.no_grad()
|
||
def sample_heunpp2(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1., threshold=0.025, ret_steps=5):
|
||
# From MIT licensed: https://github.com/Carzit/sd-webui-samplers-scheduler/
|
||
extra_args = {} if extra_args is None else extra_args
|
||
s_in = x.new_ones([x.shape[0]])
|
||
s_end = sigmas[-1]
|
||
#------------add for EasyCache------------
|
||
total_steps = len(sigmas) - 1
|
||
ec = EasyCacheState(total_steps, threshold=threshold, ret_steps=ret_steps)
|
||
#-----------------------------------------
|
||
for i in trange(total_steps, disable=disable):
|
||
gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.
|
||
eps = torch.randn_like(x) * s_noise
|
||
sigma_hat = sigmas[i] * (gamma + 1)
|
||
if gamma > 0:
|
||
x = x + eps * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5
|
||
#------------add for EasyCache------------
|
||
x_step_in = x
|
||
ec.begin_step(i, total_steps, x_step_in)
|
||
denoised = ec.model_call(model, x, sigma_hat * s_in, **extra_args)
|
||
#-----------------------------------------
|
||
d = to_d(x, 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
|
||
if sigmas[i + 1] == s_end:
|
||
# Euler method
|
||
x = x + d * dt
|
||
elif sigmas[i + 2] == s_end:
|
||
|
||
# Heun's method
|
||
x_2 = x + d * dt
|
||
denoised_2 = ec.model_call(model, x_2, sigmas[i + 1] * s_in, **extra_args)
|
||
d_2 = to_d(x_2, sigmas[i + 1], denoised_2)
|
||
|
||
w = 2 * sigmas[0]
|
||
w2 = sigmas[i+1]/w
|
||
w1 = 1 - w2
|
||
|
||
d_prime = d * w1 + d_2 * w2
|
||
|
||
|
||
x = x + d_prime * dt
|
||
|
||
else:
|
||
# Heun++
|
||
x_2 = x + d * dt
|
||
denoised_2 = ec.model_call(model, x_2, sigmas[i + 1] * s_in, **extra_args)
|
||
d_2 = to_d(x_2, sigmas[i + 1], denoised_2)
|
||
dt_2 = sigmas[i + 2] - sigmas[i + 1]
|
||
|
||
x_3 = x_2 + d_2 * dt_2
|
||
denoised_3 = ec.model_call(model, x_3, sigmas[i + 2] * s_in, **extra_args)
|
||
d_3 = to_d(x_3, sigmas[i + 2], denoised_3)
|
||
|
||
w = 3 * sigmas[0]
|
||
w2 = sigmas[i + 1] / w
|
||
w3 = sigmas[i + 2] / w
|
||
w1 = 1 - w2 - w3
|
||
|
||
d_prime = w1 * d + w2 * d_2 + w3 * d_3
|
||
x = x + d_prime * dt
|
||
ec.end_step(x_step_in)
|
||
return x
|
||
|
||
#From https://github.com/zju-pi/diff-sampler/blob/main/diff-solvers-main/solvers.py
|
||
#under Apache 2 license
|
||
def sample_ipndm(model, x, sigmas, extra_args=None, callback=None, disable=None, max_order=4, threshold=0.025, ret_steps=5):
|
||
extra_args = {} if extra_args is None else extra_args
|
||
s_in = x.new_ones([x.shape[0]])
|
||
x_next = x
|
||
buffer_model = []
|
||
#------------add for EasyCache------------
|
||
total_steps = len(sigmas) - 1
|
||
ec = EasyCacheState(total_steps, threshold=threshold, ret_steps=ret_steps)
|
||
#-----------------------------------------
|
||
for i in trange(total_steps, disable=disable):
|
||
t_cur = sigmas[i]
|
||
t_next = sigmas[i + 1]
|
||
|
||
x_cur = x_next
|
||
|
||
#------------add for EasyCache------------
|
||
x_step_in = x_cur
|
||
ec.begin_step(i, total_steps, x_step_in)
|
||
denoised = ec.model_call(model, x_cur, t_cur * s_in, **extra_args)
|
||
#-----------------------------------------
|
||
if callback is not None:
|
||
callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
|
||
|
||
d_cur = (x_cur - denoised) / t_cur
|
||
|
||
order = min(max_order, i+1)
|
||
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
|
||
elif order == 2: # Use one history point.
|
||
x_next = x_cur + (t_next - t_cur) * (3 * d_cur - buffer_model[-1]) / 2
|
||
elif order == 3: # Use two history points.
|
||
x_next = x_cur + (t_next - t_cur) * (23 * d_cur - 16 * buffer_model[-1] + 5 * buffer_model[-2]) / 12
|
||
elif order == 4: # Use three history points.
|
||
x_next = x_cur + (t_next - t_cur) * (55 * d_cur - 59 * buffer_model[-1] + 37 * buffer_model[-2] - 9 * buffer_model[-3]) / 24
|
||
|
||
if len(buffer_model) == max_order - 1:
|
||
for k in range(max_order - 2):
|
||
buffer_model[k] = buffer_model[k+1]
|
||
buffer_model[-1] = d_cur
|
||
else:
|
||
buffer_model.append(d_cur)
|
||
ec.end_step(x_step_in)
|
||
return x_next
|
||
|
||
#From https://github.com/zju-pi/diff-sampler/blob/main/diff-solvers-main/solvers.py
|
||
#under Apache 2 license
|
||
def sample_ipndm_v(model, x, sigmas, extra_args=None, callback=None, disable=None, max_order=4, threshold=0.025, ret_steps=5):
|
||
extra_args = {} if extra_args is None else extra_args
|
||
s_in = x.new_ones([x.shape[0]])
|
||
x_next = x
|
||
t_steps = sigmas
|
||
buffer_model = []
|
||
#------------add for EasyCache------------
|
||
total_steps = len(sigmas) - 1
|
||
ec = EasyCacheState(total_steps, threshold=threshold, ret_steps=ret_steps)
|
||
#-----------------------------------------
|
||
for i in trange(total_steps, disable=disable):
|
||
t_cur = sigmas[i]
|
||
t_next = sigmas[i + 1]
|
||
x_cur = x_next
|
||
#------------add for EasyCache------------
|
||
x_step_in = x_cur
|
||
ec.begin_step(i, total_steps, x_step_in)
|
||
denoised = ec.model_call(model, x_cur, t_cur * s_in, **extra_args)
|
||
#-----------------------------------------
|
||
if callback is not None:
|
||
callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
|
||
d_cur = (x_cur - denoised) / t_cur
|
||
order = min(max_order, i+1)
|
||
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
|
||
elif order == 2: # Use one history point.
|
||
h_n = (t_next - t_cur)
|
||
h_n_1 = (t_cur - t_steps[i-1])
|
||
coeff1 = (2 + (h_n / h_n_1)) / 2
|
||
coeff2 = -(h_n / h_n_1) / 2
|
||
x_next = x_cur + (t_next - t_cur) * (coeff1 * d_cur + coeff2 * buffer_model[-1])
|
||
elif order == 3: # Use two history points.
|
||
h_n = (t_next - t_cur)
|
||
h_n_1 = (t_cur - t_steps[i-1])
|
||
h_n_2 = (t_steps[i-1] - t_steps[i-2])
|
||
temp = (1 - h_n / (3 * (h_n + h_n_1)) * (h_n * (h_n + h_n_1)) / (h_n_1 * (h_n_1 + h_n_2))) / 2
|
||
coeff1 = (2 + (h_n / h_n_1)) / 2 + temp
|
||
coeff2 = -(h_n / h_n_1) / 2 - (1 + h_n_1 / h_n_2) * temp
|
||
coeff3 = temp * h_n_1 / h_n_2
|
||
x_next = x_cur + (t_next - t_cur) * (coeff1 * d_cur + coeff2 * buffer_model[-1] + coeff3 * buffer_model[-2])
|
||
elif order == 4: # Use three history points.
|
||
h_n = (t_next - t_cur)
|
||
h_n_1 = (t_cur - t_steps[i-1])
|
||
h_n_2 = (t_steps[i-1] - t_steps[i-2])
|
||
h_n_3 = (t_steps[i-2] - t_steps[i-3])
|
||
temp1 = (1 - h_n / (3 * (h_n + h_n_1)) * (h_n * (h_n + h_n_1)) / (h_n_1 * (h_n_1 + h_n_2))) / 2
|
||
temp2 = ((1 - h_n / (3 * (h_n + h_n_1))) / 2 + (1 - h_n / (2 * (h_n + h_n_1))) * h_n / (6 * (h_n + h_n_1 + h_n_2))) \
|
||
* (h_n * (h_n + h_n_1) * (h_n + h_n_1 + h_n_2)) / (h_n_1 * (h_n_1 + h_n_2) * (h_n_1 + h_n_2 + h_n_3))
|
||
coeff1 = (2 + (h_n / h_n_1)) / 2 + temp1 + temp2
|
||
coeff2 = -(h_n / h_n_1) / 2 - (1 + h_n_1 / h_n_2) * temp1 - (1 + (h_n_1 / h_n_2) + (h_n_1 * (h_n_1 + h_n_2) / (h_n_2 * (h_n_2 + h_n_3)))) * temp2
|
||
coeff3 = temp1 * h_n_1 / h_n_2 + ((h_n_1 / h_n_2) + (h_n_1 * (h_n_1 + h_n_2) / (h_n_2 * (h_n_2 + h_n_3))) * (1 + h_n_2 / h_n_3)) * temp2
|
||
coeff4 = -temp2 * (h_n_1 * (h_n_1 + h_n_2) / (h_n_2 * (h_n_2 + h_n_3))) * h_n_1 / h_n_2
|
||
x_next = x_cur + (t_next - t_cur) * (coeff1 * d_cur + coeff2 * buffer_model[-1] + coeff3 * buffer_model[-2] + coeff4 * buffer_model[-3])
|
||
if len(buffer_model) == max_order - 1:
|
||
for k in range(max_order - 2):
|
||
buffer_model[k] = buffer_model[k+1]
|
||
buffer_model[-1] = d_cur.detach()
|
||
else:
|
||
buffer_model.append(d_cur.detach())
|
||
ec.end_step(x_step_in)
|
||
|
||
return x_next
|
||
|
||
#From https://github.com/zju-pi/diff-sampler/blob/main/diff-solvers-main/solvers.py
|
||
#under Apache 2 license
|
||
@torch.no_grad()
|
||
def sample_deis(model, x, sigmas, extra_args=None, callback=None, disable=None, max_order=3, deis_mode='tab', threshold=0.025, ret_steps=5):
|
||
extra_args = {} if extra_args is None else extra_args
|
||
s_in = x.new_ones([x.shape[0]])
|
||
x_next = x
|
||
t_steps = sigmas
|
||
coeff_list = deis.get_deis_coeff_list(t_steps, max_order, deis_mode=deis_mode)
|
||
buffer_model = []
|
||
#------------add for EasyCache------------
|
||
total_steps = len(sigmas) - 1
|
||
ec = EasyCacheState(total_steps, threshold=threshold, ret_steps=ret_steps)
|
||
#-----------------------------------------
|
||
for i in trange(total_steps, disable=disable):
|
||
t_cur = sigmas[i]
|
||
t_next = sigmas[i + 1]
|
||
x_cur = x_next
|
||
#------------add for EasyCache------------
|
||
x_step_in = x_cur
|
||
ec.begin_step(i, total_steps, x_step_in)
|
||
denoised = ec.model_call(model, x_cur, t_cur * s_in, **extra_args)
|
||
#-----------------------------------------
|
||
if callback is not None:
|
||
callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
|
||
d_cur = (x_cur - denoised) / t_cur
|
||
order = min(max_order, i+1)
|
||
if t_next <= 0:
|
||
order = 1
|
||
if order == 1: # First Euler step.
|
||
x_next = x_cur + (t_next - t_cur) * d_cur
|
||
elif order == 2: # Use one history point.
|
||
coeff_cur, coeff_prev1 = coeff_list[i]
|
||
x_next = x_cur + coeff_cur * d_cur + coeff_prev1 * buffer_model[-1]
|
||
elif order == 3: # Use two history points.
|
||
coeff_cur, coeff_prev1, coeff_prev2 = coeff_list[i]
|
||
x_next = x_cur + coeff_cur * d_cur + coeff_prev1 * buffer_model[-1] + coeff_prev2 * buffer_model[-2]
|
||
elif order == 4: # Use three history points.
|
||
coeff_cur, coeff_prev1, coeff_prev2, coeff_prev3 = coeff_list[i]
|
||
x_next = x_cur + coeff_cur * d_cur + coeff_prev1 * buffer_model[-1] + coeff_prev2 * buffer_model[-2] + coeff_prev3 * buffer_model[-3]
|
||
if len(buffer_model) == max_order - 1:
|
||
for k in range(max_order - 2):
|
||
buffer_model[k] = buffer_model[k+1]
|
||
buffer_model[-1] = d_cur.detach()
|
||
else:
|
||
buffer_model.append(d_cur.detach())
|
||
ec.end_step(x_step_in)
|
||
return x_next
|
||
|
||
@torch.no_grad()
|
||
def res_multistep(model, x, sigmas, extra_args=None, callback=None, disable=None, s_noise=1., noise_sampler=None, eta=1., cfg_pp=False, threshold=0.025, ret_steps=5):
|
||
extra_args = {} if extra_args is None else extra_args
|
||
seed = extra_args.get("seed", None)
|
||
noise_sampler = default_noise_sampler(x, seed=seed) if noise_sampler is None else noise_sampler
|
||
s_in = x.new_ones([x.shape[0]])
|
||
sigma_fn = lambda t: t.neg().exp()
|
||
t_fn = lambda sigma: sigma.log().neg()
|
||
phi1_fn = lambda t: torch.expm1(t) / t
|
||
phi2_fn = lambda t: (phi1_fn(t) - 1.0) / t
|
||
old_sigma_down = None
|
||
old_denoised = None
|
||
uncond_denoised = None
|
||
def post_cfg_function(args):
|
||
nonlocal uncond_denoised
|
||
uncond_denoised = args["uncond_denoised"]
|
||
return args["denoised"]
|
||
if cfg_pp:
|
||
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)
|
||
#------------add for EasyCache------------
|
||
total_steps = len(sigmas) - 1
|
||
ec = EasyCacheState(total_steps, threshold=threshold, ret_steps=ret_steps)
|
||
#-----------------------------------------
|
||
for i in trange(total_steps, disable=disable):
|
||
#------------add for EasyCache------------
|
||
x_step_in = x
|
||
ec.begin_step(i, total_steps, x_step_in)
|
||
denoised = ec.model_call(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})
|
||
if sigma_down == 0 or old_denoised is None:
|
||
# Euler method
|
||
if cfg_pp:
|
||
d = to_d(x, sigmas[i], uncond_denoised)
|
||
x = denoised + d * sigma_down
|
||
else:
|
||
d = to_d(x, sigmas[i], denoised)
|
||
dt = sigma_down - sigmas[i]
|
||
x = x + d * dt
|
||
else:
|
||
# Second order multistep method in https://arxiv.org/pdf/2308.02157
|
||
t, t_old, t_next, t_prev = t_fn(sigmas[i]), t_fn(old_sigma_down), t_fn(sigma_down), t_fn(sigmas[i - 1])
|
||
h = t_next - t
|
||
c2 = (t_prev - t_old) / h
|
||
phi1_val, phi2_val = phi1_fn(-h), phi2_fn(-h)
|
||
b1 = torch.nan_to_num(phi1_val - phi2_val / c2, nan=0.0)
|
||
b2 = torch.nan_to_num(phi2_val / c2, nan=0.0)
|
||
if cfg_pp:
|
||
x = x + (denoised - uncond_denoised)
|
||
x = sigma_fn(h) * x + h * (b1 * uncond_denoised + b2 * old_denoised)
|
||
else:
|
||
x = sigma_fn(h) * x + h * (b1 * denoised + b2 * old_denoised)
|
||
# Noise addition
|
||
if sigmas[i + 1] > 0:
|
||
x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up
|
||
if cfg_pp:
|
||
old_denoised = uncond_denoised
|
||
else:
|
||
old_denoised = denoised
|
||
old_sigma_down = sigma_down
|
||
ec.end_step(x_step_in)
|
||
return x
|
||
|
||
@torch.no_grad()
|
||
def sample_res_multistep(model, x, sigmas, extra_args=None, callback=None, disable=None, s_noise=1., noise_sampler=None, threshold=0.025, ret_steps=5):
|
||
return res_multistep(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, s_noise=s_noise, noise_sampler=noise_sampler, eta=0., cfg_pp=False, threshold=threshold, ret_steps=ret_steps)
|
||
|
||
@torch.no_grad()
|
||
def sample_gradient_estimation(model, x, sigmas, extra_args=None, callback=None, disable=None, ge_gamma=2., cfg_pp=False, threshold=0.025, ret_steps=5):
|
||
"""Gradient-estimation sampler. Paper: https://openreview.net/pdf?id=o2ND9v0CeK"""
|
||
extra_args = {} if extra_args is None else extra_args
|
||
s_in = x.new_ones([x.shape[0]])
|
||
old_d = None
|
||
uncond_denoised = None
|
||
def post_cfg_function(args):
|
||
nonlocal uncond_denoised
|
||
uncond_denoised = args["uncond_denoised"]
|
||
return args["denoised"]
|
||
if cfg_pp:
|
||
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)
|
||
#------------add for EasyCache------------
|
||
total_steps = len(sigmas) - 1
|
||
ec = EasyCacheState(total_steps, threshold=threshold, ret_steps=ret_steps)
|
||
#-----------------------------------------
|
||
for i in trange(total_steps, disable=disable):
|
||
#------------add for EasyCache------------
|
||
x_step_in = x
|
||
ec.begin_step(i, total_steps, x_step_in)
|
||
denoised = ec.model_call(model, x, sigmas[i] * s_in, **extra_args)
|
||
#-----------------------------------------
|
||
if cfg_pp:
|
||
d = to_d(x, sigmas[i], uncond_denoised)
|
||
else:
|
||
d = to_d(x, sigmas[i], denoised)
|
||
if callback is not None:
|
||
callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
|
||
dt = sigmas[i + 1] - sigmas[i]
|
||
if sigmas[i + 1] == 0:
|
||
# Denoising step
|
||
x = denoised
|
||
else:
|
||
# Euler method
|
||
if cfg_pp:
|
||
x = denoised + d * sigmas[i + 1]
|
||
else:
|
||
x = x + d * dt
|
||
if i >= 1:
|
||
# Gradient estimation
|
||
d_bar = (ge_gamma - 1) * (d - old_d)
|
||
x = x + d_bar * dt
|
||
old_d = d
|
||
ec.end_step(x_step_in)
|
||
return x
|
||
|
||
#ancestral、sde方法，可多样性
|
||
@torch.no_grad()
|
||
def sample_euler_ancestral(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, threshold=0.025, ret_steps=5):
|
||
if isinstance(model.inner_model.inner_model.model_sampling, comfy.model_sampling.CONST):
|
||
return sample_euler_ancestral_RF(model, x, sigmas, extra_args, callback, disable, eta, s_noise, noise_sampler, threshold=threshold, ret_steps=ret_steps)
|
||
"""Ancestral sampling with Euler method steps."""
|
||
extra_args = {} if extra_args is None else extra_args
|
||
seed = extra_args.get("seed", None)
|
||
noise_sampler = default_noise_sampler(x, seed=seed) if noise_sampler is None else noise_sampler
|
||
s_in = x.new_ones([x.shape[0]])
|
||
#------------add for EasyCache------------
|
||
total_steps = len(sigmas) - 1
|
||
ec = EasyCacheState(total_steps, threshold=threshold, ret_steps=ret_steps)
|
||
#-----------------------------------------
|
||
for i in trange(total_steps, disable=disable):
|
||
#------------add for EasyCache------------
|
||
x_step_in = x
|
||
ec.begin_step(i, total_steps, x_step_in)
|
||
denoised = ec.model_call(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})
|
||
|
||
if sigma_down == 0:
|
||
x = denoised
|
||
else:
|
||
d = to_d(x, sigmas[i], denoised)
|
||
# Euler method
|
||
dt = sigma_down - sigmas[i]
|
||
x = x + d * dt + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up
|
||
ec.end_step(x_step_in)
|
||
return x
|
||
|
||
@torch.no_grad()
|
||
def sample_euler_ancestral_RF(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1.0, s_noise=1., noise_sampler=None, threshold=0.025, ret_steps=5):
|
||
"""Ancestral sampling with Euler method steps."""
|
||
extra_args = {} if extra_args is None else extra_args
|
||
seed = extra_args.get("seed", None)
|
||
noise_sampler = default_noise_sampler(x, seed=seed) if noise_sampler is None else noise_sampler
|
||
s_in = x.new_ones([x.shape[0]])
|
||
#------------add for EasyCache------------
|
||
total_steps = len(sigmas) - 1
|
||
ec = EasyCacheState(total_steps, threshold=threshold, ret_steps=ret_steps)
|
||
#-----------------------------------------
|
||
for i in trange(total_steps, disable=disable):
|
||
#------------add for EasyCache------------
|
||
x_step_in = x
|
||
ec.begin_step(i, total_steps, x_step_in)
|
||
denoised = ec.model_call(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})
|
||
|
||
if sigmas[i + 1] == 0:
|
||
x = denoised
|
||
else:
|
||
downstep_ratio = 1 + (sigmas[i + 1] / sigmas[i] - 1) * eta
|
||
sigma_down = sigmas[i + 1] * downstep_ratio
|
||
alpha_ip1 = 1 - sigmas[i + 1]
|
||
alpha_down = 1 - sigma_down
|
||
renoise_coeff = (sigmas[i + 1]**2 - sigma_down**2 * alpha_ip1**2 / alpha_down**2)**0.5
|
||
# Euler method
|
||
sigma_down_i_ratio = sigma_down / sigmas[i]
|
||
x = sigma_down_i_ratio * x + (1 - sigma_down_i_ratio) * denoised
|
||
if eta > 0:
|
||
x = (alpha_ip1 / alpha_down) * x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * renoise_coeff
|
||
ec.end_step(x_step_in)
|
||
return x
|
||
|
||
@torch.no_grad()
|
||
def sample_dpm_2_ancestral(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, threshold=0.025, ret_steps=5):
|
||
if isinstance(model.inner_model.inner_model.model_sampling, comfy.model_sampling.CONST):
|
||
return sample_dpm_2_ancestral_RF(model, x, sigmas, extra_args, callback, disable, eta, s_noise, noise_sampler, threshold=threshold, ret_steps=ret_steps)
|
||
"""Ancestral sampling with DPM-Solver second-order steps."""
|
||
extra_args = {} if extra_args is None else extra_args
|
||
seed = extra_args.get("seed", None)
|
||
noise_sampler = default_noise_sampler(x, seed=seed) if noise_sampler is None else noise_sampler
|
||
s_in = x.new_ones([x.shape[0]])
|
||
#------------add for EasyCache------------
|
||
total_steps = len(sigmas) - 1
|
||
ec = EasyCacheState(total_steps, threshold=threshold, ret_steps=ret_steps)
|
||
#-----------------------------------------
|
||
for i in trange(total_steps, disable=disable):
|
||
#------------add for EasyCache------------
|
||
x_step_in = x
|
||
ec.begin_step(i, total_steps, x_step_in)
|
||
denoised = ec.model_call(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], denoised)
|
||
if sigma_down == 0:
|
||
# Euler method
|
||
dt = sigma_down - sigmas[i]
|
||
x = x + d * dt
|
||
else:
|
||
# DPM-Solver-2
|
||
sigma_mid = sigmas[i].log().lerp(sigma_down.log(), 0.5).exp()
|
||
dt_1 = sigma_mid - sigmas[i]
|
||
dt_2 = sigma_down - sigmas[i]
|
||
x_2 = x + d * dt_1
|
||
denoised_2 = ec.model_call(model, x_2, sigma_mid * s_in, **extra_args)
|
||
d_2 = to_d(x_2, sigma_mid, denoised_2)
|
||
x = x + d_2 * dt_2
|
||
x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up
|
||
ec.end_step(x_step_in)
|
||
return x
|
||
|
||
@torch.no_grad()
|
||
def sample_dpm_2_ancestral_RF(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, threshold=0.025, ret_steps=5):
|
||
"""Ancestral sampling with DPM-Solver second-order steps."""
|
||
extra_args = {} if extra_args is None else extra_args
|
||
seed = extra_args.get("seed", None)
|
||
noise_sampler = default_noise_sampler(x, seed=seed) if noise_sampler is None else noise_sampler
|
||
s_in = x.new_ones([x.shape[0]])
|
||
#------------add for EasyCache------------
|
||
total_steps = len(sigmas) - 1
|
||
ec = EasyCacheState(total_steps, threshold=threshold, ret_steps=ret_steps)
|
||
#-----------------------------------------
|
||
for i in trange(total_steps, disable=disable):
|
||
#------------add for EasyCache------------
|
||
x_step_in = x
|
||
ec.begin_step(i, total_steps, x_step_in)
|
||
denoised = ec.model_call(model, x, sigmas[i] * s_in, **extra_args)
|
||
#-----------------------------------------
|
||
downstep_ratio = 1 + (sigmas[i+1]/sigmas[i] - 1) * eta
|
||
sigma_down = sigmas[i+1] * downstep_ratio
|
||
alpha_ip1 = 1 - sigmas[i+1]
|
||
alpha_down = 1 - sigma_down
|
||
renoise_coeff = (sigmas[i+1]**2 - sigma_down**2*alpha_ip1**2/alpha_down**2)**0.5
|
||
|
||
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], denoised)
|
||
if sigma_down == 0:
|
||
# Euler method
|
||
dt = sigma_down - sigmas[i]
|
||
x = x + d * dt
|
||
else:
|
||
# DPM-Solver-2
|
||
sigma_mid = sigmas[i].log().lerp(sigma_down.log(), 0.5).exp()
|
||
dt_1 = sigma_mid - sigmas[i]
|
||
dt_2 = sigma_down - sigmas[i]
|
||
x_2 = x + d * dt_1
|
||
denoised_2 = ec.model_call(model, x_2, sigma_mid * s_in, **extra_args)
|
||
d_2 = to_d(x_2, sigma_mid, denoised_2)
|
||
x = x + d_2 * dt_2
|
||
x = (alpha_ip1/alpha_down) * x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * renoise_coeff
|
||
ec.end_step(x_step_in)
|
||
return x
|