diff --git a/comfy/k_diffusion/sampling.py b/comfy/k_diffusion/sampling.py index 7af01682..568564dc 100644 --- a/comfy/k_diffusion/sampling.py +++ b/comfy/k_diffusion/sampling.py @@ -808,3 +808,107 @@ def sample_heunpp2(model, x, sigmas, extra_args=None, callback=None, disable=Non d_prime = w1 * d + w2 * d_2 + w3 * d_3 x = x + d_prime * dt return x + +@torch.no_grad() +def dy_sampling_step(x, model, dt, sigma_hat, **extra_args): + + original_shape = x.shape + m, n = original_shape[2] // 2, original_shape[3] // 2 + extra_row = x.shape[2] % 2 == 1 + extra_col = x.shape[3] % 2 == 1 + + if extra_row: + extra_row_content = x[:, :, -1:, :] + x = x[:, :, :-1, :] + if extra_col: + extra_col_content = x[:, :, :, -1:] + x = x[:, :, :, :-1] + + a_list = x.unfold(2, 2, 2).unfold(3, 2, 2).contiguous().view(1, 4, m * n, 2, 2) + c = a_list[:, :, :, 1, 1].view(1, 4, m, n) + + denoised = model(c, sigma_hat * c.new_ones([c.shape[0]]), **extra_args) + d = to_d(c, sigma_hat, denoised) + c = c + d * dt + + d_list = c.view(1, 4, m * n, 1, 1) + a_list[:, :, :, 1, 1] = d_list[:, :, :, 0, 0] + x = a_list.view(1, 4, m, n, 2, 2).permute(0, 1, 2, 4, 3, 5).reshape(1, 4, 2 * m, 2 * n) + + if extra_row or extra_col: + x_expanded = torch.zeros(original_shape, dtype=x.dtype, device=x.device) + x_expanded[:, :, :2 * m, :2 * n] = x + if extra_row: + x_expanded[:, :, -1:, :2 * n + 1] = extra_row_content + if extra_col: + x_expanded[:, :, :2 * m, -1:] = extra_col_content + if extra_row and extra_col: + x_expanded[:, :, -1:, -1:] = extra_col_content[:, :, -1:, :] + x = x_expanded + + return x + +@torch.no_grad() +def sample_euler_dy(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., + s_tmax=float('inf'), s_noise=1.): + """Euler dy. See https://github.com/Koishi-Star/Euler-Smea-Dyn-Sampler.""" + 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): + # print(i) + # i第一步为0 + gamma = max(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) + # print(sigma_hat) + dt = sigmas[i + 1] - sigma_hat + if gamma > 0: + 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 sigmas[i + 1] > 0: + if i // 2 == 1: + x = dy_sampling_step(x, model, dt, sigma_hat, **extra_args) + if callback is not None: + callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat, 'denoised': denoised}) + # Euler method + x = x + d * dt + return x + +@torch.no_grad() +def smea_sampling_step(x, model, dt, sigma_hat, **extra_args): + m, n = x.shape[2], x.shape[3] + x = torch.nn.functional.interpolate(input=x, size=None, scale_factor=(1.25, 1.25), mode='nearest', + align_corners=None,recompute_scale_factor=None) + denoised = model(x, sigma_hat * x.new_ones([x.shape[0]]), **extra_args) + d = to_d(x, sigma_hat, denoised) + x = x + d * dt + x = torch.nn.functional.interpolate(input=x, size=(m,n), scale_factor=None, mode='nearest', + align_corners=None, recompute_scale_factor=None) + return x + +@torch.no_grad() +def sample_euler_smea_dy(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., + s_tmax=float('inf'), s_noise=1.): + """Euler Smea dy. See https://github.com/Koishi-Star/Euler-Smea-Dyn-Sampler.""" + 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): + gamma = max(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) + dt = sigmas[i + 1] - sigma_hat + if gamma > 0: + 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) + # Euler method + x = x + d * dt + if sigmas[i + 1] > 0: + if i + 1 // 2 == 1: + x = dy_sampling_step(x, model, dt, sigma_hat, **extra_args) + if i + 1 // 2 == 0: + x = smea_sampling_step(x, model, dt, sigma_hat, **extra_args) + if callback is not None: + callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat, 'denoised': denoised}) + return x diff --git a/comfy/samplers.py b/comfy/samplers.py index 415a35cc..610bbba1 100644 --- a/comfy/samplers.py +++ b/comfy/samplers.py @@ -504,7 +504,7 @@ class Sampler: sigma = float(sigmas[0]) return math.isclose(max_sigma, sigma, rel_tol=1e-05) or sigma > max_sigma -KSAMPLER_NAMES = ["euler", "euler_ancestral", "heun", "heunpp2","dpm_2", "dpm_2_ancestral", +KSAMPLER_NAMES = ["euler", "euler_ancestral", "euler_dy", "euler_smea_dy", "heun", "heunpp2","dpm_2", "dpm_2_ancestral", "lms", "dpm_fast", "dpm_adaptive", "dpmpp_2s_ancestral", "dpmpp_sde", "dpmpp_sde_gpu", "dpmpp_2m", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "dpmpp_3m_sde", "dpmpp_3m_sde_gpu", "ddpm", "lcm"]