diff --git a/comfy/k_diffusion/sampling.py b/comfy/k_diffusion/sampling.py index 11db46d94..3fa275870 100644 --- a/comfy/k_diffusion/sampling.py +++ b/comfy/k_diffusion/sampling.py @@ -1316,14 +1316,19 @@ def sample_euler_cfg_pp(model, x, sigmas, extra_args=None, callback=None, disabl @torch.no_grad() def sample_dpmpp_2s_ancestral_cfg_pp(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.""" + model_sampling = model.inner_model.model_patcher.get_model_object('model_sampling') + if isinstance(model_sampling, comfy.model_sampling.CONST): + return sample_dpmpp_2s_ancestral_RF_cfg_pp(model, x, sigmas, extra_args, callback, disable, eta, s_noise, noise_sampler) + 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_noise = s_noise * getattr(model_sampling, "noise_scale", 1.0) - temp = [0] + uncond_denoised = None def post_cfg_function(args): - temp[0] = args["uncond_denoised"] + nonlocal uncond_denoised + uncond_denoised = args["uncond_denoised"] return args["denoised"] model_options = extra_args.get("model_options", {}).copy() @@ -1335,26 +1340,94 @@ def sample_dpmpp_2s_ancestral_cfg_pp(model, x, sigmas, extra_args=None, callback for i in trange(len(sigmas) - 1, disable=disable): denoised = model(x, sigmas[i] * s_in, **extra_args) + guidance_vector = denoised - uncond_denoised 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: # Euler method - d = to_d(x, sigmas[i], temp[0]) - x = denoised + d * sigma_down + d = to_d(x, sigmas[i], uncond_denoised) + dt = sigma_down - sigmas[i] + x = x + d * dt else: # DPM-Solver++(2S) t, t_next = t_fn(sigmas[i]), t_fn(sigma_down) - # r = torch.sinh(1 + (2 - eta) * (t_next - t) / (t - t_fn(sigma_up))) works only on non-cfgpp, weird r = 1 / 2 h = t_next - t s = t + r * h - x_2 = (sigma_fn(s) / sigma_fn(t)) * (x + (denoised - temp[0])) - (-h * r).expm1() * denoised - denoised_2 = model(x_2, sigma_fn(s) * s_in, **extra_args) - x = (sigma_fn(t_next) / sigma_fn(t)) * (x + (denoised - temp[0])) - (-h).expm1() * denoised_2 + x_2 = (sigma_fn(s) / sigma_fn(t)) * x - (-h * r).expm1() * uncond_denoised + _ = model(x_2, sigma_fn(s) * s_in, **extra_args) + uncond_denoised_2 = uncond_denoised + x = (sigma_fn(t_next) / sigma_fn(t)) * x - (-h).expm1() * uncond_denoised_2 # Noise addition if sigmas[i + 1] > 0: x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up + # Not RF: alpha = 1 + x += guidance_vector + return x + +def sample_dpmpp_2s_ancestral_RF_cfg_pp(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 + 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) + + uncond_denoised = None + def post_cfg_function(args): + nonlocal uncond_denoised + uncond_denoised = args["uncond_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, disable_cfg1_optimization=True) + + s_in = x.new_ones([x.shape[0]]) + sigma_fn = lambda lbda: (lbda.exp() + 1) ** -1 + lambda_fn = lambda sigma: ((1-sigma)/sigma).log() + + # logged_x = x.unsqueeze(0) + + for i in trange(len(sigmas) - 1, disable=disable): + denoised = model(x, sigmas[i] * s_in, **extra_args) + guidance_vector = denoised - uncond_denoised + 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 + # 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: + # Euler method + d = to_d(x, sigmas[i], uncond_denoised) + dt = sigma_down - sigmas[i] + x = x + d * dt + else: + # DPM-Solver++(2S) + if sigmas[i] == 1.0: + sigma_s = 0.9999 + else: + t_i, t_down = lambda_fn(sigmas[i]), lambda_fn(sigma_down) + r = 1 / 2 + h = t_down - t_i + s = t_i + r * h + sigma_s = sigma_fn(s) + # sigma_s = sigmas[i+1] + sigma_s_i_ratio = sigma_s / sigmas[i] + u = sigma_s_i_ratio * x + (1 - sigma_s_i_ratio) * uncond_denoised + _ = model(u, sigma_s * s_in, **extra_args) + uncond_D_i = uncond_denoised + sigma_down_i_ratio = sigma_down / sigmas[i] + x = sigma_down_i_ratio * x + (1 - sigma_down_i_ratio) * uncond_D_i + # print("sigma_i", sigmas[i], "sigma_ip1", sigmas[i+1],"sigma_down", sigma_down, "sigma_down_i_ratio", sigma_down_i_ratio, "sigma_s_i_ratio", sigma_s_i_ratio, "renoise_coeff", renoise_coeff) + # Noise addition + if sigmas[i + 1] > 0 and eta > 0: + x = (alpha_ip1/alpha_down) * x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * renoise_coeff + ### RF: alpha = 1 - t = 1 - sigma + x += (1 - sigmas[i + 1]) * guidance_vector + # logged_x = torch.cat((logged_x, x.unsqueeze(0)), dim=0) return x @torch.no_grad()