"""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

"""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))

"""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))

"""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))

"""Calculates the noise level (sigma_down) to step down to and the amount

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

if seed is not None:

if x.device == torch.device("cpu"):

seed += 1

generator = torch.Generator(device=x.device)

generator.manual_seed(seed)

else:

generator = None

"""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)

"""A noise sampler backed by a torchsde.BrownianTree.

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))

"""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()

"""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()

"""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)

"""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]])

for i in trange(len(sigmas) - 1, 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

denoised = 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

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)

"""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]])

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})

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

"""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_noise = s_noise * getattr(model.inner_model.model_patcher.get_model_object('model_sampling'), "noise_scale", 1.0)

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})

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

"""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]])

for i in trange(len(sigmas) - 1, 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

denoised = 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 = 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

"""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]])

for i in trange(len(sigmas) - 1, 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

denoised = 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 = 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

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)

"""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]])

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], 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 = 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

"""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_noise = s_noise * getattr(model.inner_model.model_patcher.get_model_object('model_sampling'), "noise_scale", 1.0)

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)

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 = 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

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

"""CFG++ Euler with variable-step AB2 history and optional sigma-zero extrapolation."""

extra_args = {} if extra_args is None else extra_args

model_sampling = model.inner_model.model_patcher.get_model_object("model_sampling")

lambda_fn = partial(sigma_to_half_log_snr, model_sampling=model_sampling)

s_in = x.new_ones([x.shape[0]])

sigmas_cpu = sigmas.detach().cpu().numpy()

derivatives = []

denoised_history = []

old_uncond_d = None

uncond_denoised = None

def post_cfg_function(args):

nonlocal uncond_denoised

uncond_denoised = args["uncond_denoised"] if args["uncond"] is not None else args["cond_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)

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})

alpha_s = sigmas[i] * lambda_fn(sigmas[i]).exp()

alpha_t = sigmas[i + 1] * lambda_fn(sigmas[i + 1]).exp() if sigmas[i + 1] != 0 else sigmas[i + 1].new_ones([])

current_uncond_d = to_d(x, sigmas[i], alpha_s * uncond_denoised)

uncond_d = current_uncond_d

dt = sigmas[i + 1] - sigmas[i]

if i > 0 and uncond_history_weight:

step_ratio = dt / (sigmas[i] - sigmas[i - 1])

uncond_d = uncond_d + uncond_history_weight * step_ratio * (current_uncond_d - old_uncond_d)

euler_step = alpha_t * denoised + sigmas[i + 1] * uncond_d - x

d = euler_step / dt

derivatives.append(d)

if len(derivatives) > 2:

derivatives.pop(0)

if len(derivatives) == 1:

step = euler_step

else:

coeffs = [linear_multistep_coeff(2, sigmas_cpu, i, j) for j in range(2)]

history_step = sum(coeff * derivative for coeff, derivative in zip(coeffs, reversed(derivatives)))

step = torch.lerp(euler_step, history_step, history_weight)

x = x + step

if sigmas[i + 1] == 0 and zero_weight is not None and denoised_history:

if zero_order == 2 and len(denoised_history) > 1:

sigma_0, sigma_1, sigma_2 = sigmas[i - 2], sigmas[i - 1], sigmas[i]

weight_0 = sigma_1 * sigma_2 / ((sigma_0 - sigma_1) * (sigma_0 - sigma_2))

weight_1 = sigma_0 * sigma_2 / ((sigma_1 - sigma_0) * (sigma_1 - sigma_2))

weight_2 = sigma_0 * sigma_1 / ((sigma_2 - sigma_0) * (sigma_2 - sigma_1))

zero_prediction = weight_0 * denoised_history[-2] + weight_1 * denoised_history[-1] + weight_2 * denoised

else:

denoised_slope = (denoised - denoised_history[-1]) / (sigmas[i] - sigmas[i - 1])

zero_prediction = denoised - sigmas[i] * denoised_slope

x = torch.lerp(x, zero_prediction, zero_weight)

denoised_history.append(denoised)

if len(denoised_history) > 2:

denoised_history.pop(0)

old_uncond_d = current_uncond_d

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 = []

for i in trange(len(sigmas) - 1, disable=disable):

denoised = 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)))

"""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

"""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, *args, **kwargs):

if key in eps_cache:

return eps_cache[key], eps_cache

sigma = self.sigma(t) * x.new_ones([x.shape[0]])

eps = (x - self.model(x, sigma, *args, **self.extra_args, **kwargs)) / 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):

eps_cache = {} if eps_cache is None else eps_cache

h = t_next - t

eps, eps_cache = self.eps(eps_cache, 'eps', x, t)

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):

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)

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):

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)

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.

eps, eps_cache = self.eps(eps_cache, 'eps', x, t)

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)

elif orders[i] == 2:

x, eps_cache = self.dpm_solver_2_step(x, t, t_next_, eps_cache=eps_cache)

else:

x, eps_cache = self.dpm_solver_3_step(x, t, t_next_, eps_cache=eps_cache)

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.

eps, eps_cache = self.eps(eps_cache, 'eps', x, s)

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)

x_high, eps_cache = self.dpm_solver_2_step(x, s, t_, eps_cache=eps_cache)

else:

x_low, eps_cache = self.dpm_solver_2_step(x, s, t_, r1=1 / 3, eps_cache=eps_cache)

x_high, eps_cache = self.dpm_solver_3_step(x, s, t_, eps_cache=eps_cache)

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})

"""DPM-Solver-Fast (fixed step size). See https://arxiv.org/abs/2206.00927."""

if sigma_min <= 0 or sigma_max <= 0:

raise ValueError('sigma_min and sigma_max must not be 0')

with tqdm(total=n, disable=disable) as pbar:

dpm_solver = DPMSolver(model, extra_args, eps_callback=pbar.update)

if callback is not None:

dpm_solver.info_callback = lambda info: callback({'sigma': dpm_solver.sigma(info['t']), 'sigma_hat': dpm_solver.sigma(info['t_up']), **info})