diff --git a/modules/sd_samplers_kdiffusion.py b/modules/sd_samplers_kdiffusion.py index 0aa550b1..5d5310f9 100644 --- a/modules/sd_samplers_kdiffusion.py +++ b/modules/sd_samplers_kdiffusion.py @@ -22,6 +22,7 @@ samplers_k_diffusion = [ ("LMS", "sample_lms", ["k_lms"], {}), ("Heun", "sample_heun", ["k_heun"], {"second_order": True}), ("DPM2", "sample_dpm_2", ["k_dpm_2"], {"scheduler": "karras", "discard_next_to_last_sigma": True, "second_order": True}), + ("Res Multistep", "res_multistep", ["res_multistep"], {}), ("Restart", sd_samplers_extra.restart_sampler, ["restart"], {"scheduler": "karras", "second_order": True}), ("UniPC", sd_samplers_extra.sample_unipc, ["unipc"], {"discard_next_to_last_sigma": True}), ] diff --git a/modules/sd_schedulers.py b/modules/sd_schedulers.py index b090d61b..cfe71286 100644 --- a/modules/sd_schedulers.py +++ b/modules/sd_schedulers.py @@ -95,6 +95,23 @@ def get_align_your_steps_sigmas(n, sigma_min, sigma_max, device): return torch.FloatTensor(sigmas).to(device) +def linear_quadratic(n, sigma_min, sigma_max, device, *, threshold_noise=0.025): + if n == 1: + sigma_schedule = [1.0, 0.0] + else: + linear_steps = n // 2 + linear_sigma_schedule = [i * threshold_noise / linear_steps for i in range(linear_steps)] + threshold_noise_step_diff = linear_steps - threshold_noise * n + quadratic_steps = n - linear_steps + quadratic_coef = threshold_noise_step_diff / (linear_steps * quadratic_steps**2) + linear_coef = threshold_noise / linear_steps - 2 * threshold_noise_step_diff / (quadratic_steps**2) + const = quadratic_coef * (linear_steps**2) + quadratic_sigma_schedule = [quadratic_coef * (i**2) + linear_coef * i + const for i in range(linear_steps, n)] + sigma_schedule = linear_sigma_schedule + quadratic_sigma_schedule + [1.0] + sigma_schedule = [1.0 - x for x in sigma_schedule] + return torch.FloatTensor(sigma_schedule).to(device) * sigma_max + + def kl_optimal(n, sigma_min, sigma_max, device): alpha_min = torch.arctan(torch.tensor(sigma_min, device=device)) alpha_max = torch.arctan(torch.tensor(sigma_max, device=device)) @@ -192,6 +209,7 @@ schedulers = [ Scheduler("simple", "Simple", simple_scheduler, need_inner_model=True), Scheduler("uniform", "Uniform", uniform, need_inner_model=True), Scheduler("sgm_uniform", "SGM Uniform", sgm_uniform, need_inner_model=True, aliases=["SGMUniform"]), + Scheduler("linear_quadratic", "Linear Quadratic", linear_quadratic), Scheduler("kl_optimal", "KL Optimal", kl_optimal), Scheduler("ddim", "DDIM", ddim_scheduler, need_inner_model=True), Scheduler("align_your_steps", "Align Your Steps", get_align_your_steps_sigmas), diff --git a/modules_forge/packages/k_diffusion/sampling.py b/modules_forge/packages/k_diffusion/sampling.py index ea53da4e..bedd18ae 100644 --- a/modules_forge/packages/k_diffusion/sampling.py +++ b/modules_forge/packages/k_diffusion/sampling.py @@ -794,6 +794,46 @@ def sample_dpmpp_3m_sde(model, x, sigmas, extra_args=None, callback=None, disabl return x +@torch.no_grad() +def res_multistep(model, x, sigmas, extra_args=None, callback=None, disable=None, s_noise=1.0, noise_sampler=None): + extra_args = {} if extra_args is None else extra_args + noise_sampler = default_noise_sampler(x) 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 + + 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=0.0) + 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: + d = to_d(x, sigmas[i], denoised) + dt = sigma_down - sigmas[i] + x = x + d * dt + else: + 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) + x = sigma_fn(h) * x + h * (b1 * denoised + b2 * old_denoised) + + if sigmas[i + 1] > 0: + x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up + + old_denoised = denoised + old_sigma_down = sigma_down + return x + + @torch.no_grad() def sample_heunpp2(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0.0, s_tmin=0.0, s_tmax=float("inf"), s_noise=1.0): # From MIT licensed: https://github.com/Carzit/sd-webui-samplers-scheduler/