From ebecee7f6886e7d76080c900ccdb6c812700e8a0 Mon Sep 17 00:00:00 2001 From: yoinked Date: Sat, 29 Jun 2024 13:54:07 -0700 Subject: [PATCH] "fix" the sampler --- comfy/k_diffusion/sampling.py | 68 ++++++++++++++++++++++++++++++----- 1 file changed, 59 insertions(+), 9 deletions(-) diff --git a/comfy/k_diffusion/sampling.py b/comfy/k_diffusion/sampling.py index a112ec352..bf7dd5fa3 100644 --- a/comfy/k_diffusion/sampling.py +++ b/comfy/k_diffusion/sampling.py @@ -7,6 +7,7 @@ import torchsde from tqdm.auto import trange, tqdm from . import utils +from . import deis import comfy.model_patcher def append_zero(x): @@ -946,9 +947,58 @@ def sample_ipndm_v(model, x, sigmas, extra_args=None, callback=None, disable=Non 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'): + 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 = [] + for i in trange(len(sigmas) - 1, disable=disable): + t_cur = sigmas[i] + t_next = sigmas[i + 1] + + x_cur = x_next + + denoised = 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()) + + return x_next @torch.no_grad() -def sample_euler_pp(model, x, sigmas, extra_args=None, callback=None, disable=None): +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] @@ -963,16 +1013,16 @@ def sample_euler_pp(model, x, sigmas, extra_args=None, callback=None, disable=No for i in trange(len(sigmas) - 1, disable=disable): sigma_hat = sigmas[i] denoised = model(x, sigma_hat * s_in, **extra_args) - d = to_d(x - denoised + temp[0], sigma_hat, denoised) + 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}) dt = sigmas[i + 1] - sigma_hat # Euler method - x = x + d * dt + x = denoised + d * sigmas[i + 1] return x @torch.no_grad() -def sample_euler_ancestral_pp(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None): +def sample_euler_ancestral_cfg_pp(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None): """Ancestral sampling with Euler method steps.""" extra_args = {} if extra_args is None else extra_args noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler @@ -991,15 +1041,15 @@ def sample_euler_ancestral_pp(model, x, sigmas, extra_args=None, callback=None, sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta) if callback is not None: callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised}) - d = to_d(x - denoised + temp[0], sigmas[i], denoised) + d = to_d(x, sigmas[i], temp[0]) # Euler method dt = sigma_down - sigmas[i] - x = x + d * dt + 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_2s_ancestral_pp(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None): +def sample_dpmpp_2s_ancestral_pp_cfg(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None): """Ancestral sampling with DPM-Solver++(2S) second-order steps.""" extra_args = {} if extra_args is None else extra_args noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler @@ -1023,9 +1073,9 @@ def sample_dpmpp_2s_ancestral_pp(model, x, sigmas, extra_args=None, callback=Non callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised}) if sigma_down == 0: # Euler method - d = to_d(x - denoised + temp[0], sigmas[i], denoised) + d = to_d(x, sigmas[i], temp[0]) dt = sigma_down - sigmas[i] - x = x + d * dt + x = denoised + d * sigma_down else: # DPM-Solver++(2S) t, t_next = t_fn(sigmas[i]), t_fn(sigma_down)