From 448ecdfedd1c87fbfd4833cc9118b1ac01c891c6 Mon Sep 17 00:00:00 2001 From: Clybius Date: Mon, 25 Mar 2024 16:12:05 -0500 Subject: [PATCH] Upgrade Supreme sampler with more options, fix some latent exploding Update README.md Move samplers over into their correct category Add Kornia as a requirement --- README.md | 14 +++- extra_samplers.py | 169 ++++++++++++++++++++++++---------------------- nodes.py | 25 +++---- requirements.txt | 1 + 4 files changed, 117 insertions(+), 92 deletions(-) create mode 100644 requirements.txt diff --git a/README.md b/README.md index 817155b..1d877ae 100644 --- a/README.md +++ b/README.md @@ -6,8 +6,20 @@ * Clyb 4M SDE (A modified DPMPP 3M SDE, with an added SDE, egotisticalized by yours truly) * TTM (Thanks to Kat and Birch-San for the source implementation!) * LCM Custom Noise (Supports different types of noise other than generic gaussian) +* DPMPP 3M SDE with Dynamic ETA (Anneals down towards a minimum eta via a cosine curve) +* Supreme (Many extra functionalities and step methods available) ### Currently included extra K-Sampling nodes: * SamplerCustomNoise (Supports custom noises other than gaussian noise for init noise) -* SamplerCustomNoiseDuo (Same as above, but with an added High-res fix for simplicity. +* SamplerCustomNoiseDuo (Same as above, but with an added High-res fix for simplicity.) * SamplerCustomModelMixtureDuo (Samples with custom noises, and switches between model1 and model2 every step. If you encounter vram errors, try adding/removing `--disable-smart-memory` when launching ComfyUI) + + +#### Supreme Sampler features: +* centralization: Subtracts mean from the denoised latent. This can lead to perceptually sharper results, though may change the perceivable brightness of the image. Conservatively defaults to **(0.02)**. +* normalization: Divides the denoised latent by the standard deviation. Can increase contrast in the image, though may hurt fidelity and coherency at high strengths. Conservatively defaults to **(0.01)**. +* edge_enhancement: Sharpens the latent, and then applies a bilateral blur, leaving the edges sharpened. Conservatively defaults to **(0.05)**. +* perphist: Adds previous denoised variable to the current denoised using perpendicular vector projection. Default of **(0)**, where higher values add the old denoised variable, and negative values subtract the old denoised variable. +* substeps: Amount of times to iterate over each step and average the results. Can be useful for obtaining higher quality at a given step count. Inspired by [ReNoise](https://arxiv.org/pdf/2403.14602v1.pdf). Default of **(2)**. + +Utilizing custom sampling within ComfyUI is encouraged for these samplers! \ No newline at end of file diff --git a/extra_samplers.py b/extra_samplers.py index 2bd5867..35b76a8 100644 --- a/extra_samplers.py +++ b/extra_samplers.py @@ -3,6 +3,7 @@ import math import torch from torch import nn, FloatTensor import torchsde +import kornia from tqdm.auto import trange, tqdm import numpy as np @@ -794,7 +795,7 @@ def sample_dpmpp_3m_sde_dynamic_eta(model, x, sigmas, extra_args=None, callback= return sampler_dpmpp_3m_sde_dynamic_eta(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, eta_max=eta_max, eta_min=eta_min, s_noise=s_noise, noise_sampler=noise_sampler or get_noise_sampler(x, sigmas, noise_sampler_type, noise_sampler, extra_args)) @torch.no_grad() -def sampler_supreme(model, x, sigmas, extra_args=None, callback=None, disable=None, s_noise=1., noise_sampler=None, eta=1.0, step_method="euler", centralization=0.02, normalization=0.01, edge_enhancement=0.5, perphist=-0.15): +def sampler_supreme(model, x, sigmas, extra_args=None, callback=None, disable=None, s_noise=1., noise_sampler=None, eta=1.0, step_method="euler", centralization=0.02, normalization=0.01, edge_enhancement=0.05, perphist=0, substeps=2): """ Supreme Sampler, Euler steps. Based on no paper, purely interesting thoughts. @@ -808,10 +809,11 @@ def sampler_supreme(model, x, sigmas, extra_args=None, callback=None, disable=No s_noise: The noise scale factor. noise_sampler: A custom noise sampler function. eta: Ancestral-ness. - centralization: Subtracts mean from the denoised latent, reduces edge enhancement when between 0-1. + centralization: Subtracts mean from the denoised latent. normalization: Divides the denoised latent by the standard deviation. edge_enhancement: Multiplies the edges by the mean using a laplacian kernel perphist: Adds previous denoised variable to the current denoised using perpendicular vector projection + substeps: Amount of times to iterate over each step and average the results """ extra_args = {} if extra_args is None else extra_args @@ -850,19 +852,11 @@ def sampler_supreme(model, x, sigmas, extra_args=None, callback=None, disable=No eta = eta_min + 0.5 * (eta_max - eta_min) * (1 + math.cos(math.pi * progress)) return eta - def f(x, sigma): - """Function representing the Karras ODE derivative.""" - denoised = model(x, sigma * s_in, **extra_args) - return to_d(x, sigma, denoised) - - old_denoised = None - for i in trange(len(sigmas) - 1, disable=disable): - # DynETA - eta = eta_schedule_cosine_annealing(i, len(sigmas)) - eps_cache = {} - dpm_solver = DPMSolver(model, extra_args) - - denoised = model(x, sigmas[i] * s_in, **extra_args) + def apply_enhancements(x, i, denoised, old_denoised): + if edge_enhancement != 0: + blur = (kornia.filters.joint_bilateral_blur(x, denoised, (3, 3), 0.1, (1.5, 1.5)) - x) # Blurs non-edges + denoised += (kornia.filters.unsharp_mask(denoised, (3, 3), (1.5, 1.5)) - denoised) * (sigmas[i] - sigmas[i + 1]) * edge_enhancement # Sharpens everything + denoised += blur * (sigmas[i] - sigmas[i + 1]) * edge_enhancement # Apply blur to non-edges, thus leaving edges sharpened if centralization != 0: denoised = centralize(denoised, centralization) @@ -872,86 +866,103 @@ def sampler_supreme(model, x, sigmas, extra_args=None, callback=None, disable=No if old_denoised != None and perphist != 0: denoised = perpadd(denoised, old_denoised, x, perphist) - - if old_denoised != None and edge_enhancement != 0: - lap_kern = torch.tensor([[0, -1, 0], [-1, 4, -1], [0, -1, 0]], device=denoised.device, dtype=denoised.dtype).repeat(denoised.shape[1], 1, 1, 1) - denoised = denoised + torch.conv2d(denoised, lap_kern, groups=denoised.shape[1], padding=1) * denoised.mean(dim=(1, 2, 3), keepdim=True) * sigmas[i] * edge_enhancement - if callback is not None: - callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised}) - - eps = (x - denoised) / sigmas[i] - eps_cache = {'eps': eps} + return denoised - match step_method: - case "euler": - sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta) - d = to_d(x, sigmas[i], denoised) - dt = sigma_down - sigmas[i] + renoise_weights = torch.ones(substeps, device=x.device) / substeps - x = x + d * dt - if sigmas[i + 1] > 0: - x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up - case "dpm_1s": # DPM Family - 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}) - x, eps_cache = dpm_solver.dpm_solver_1_step(x, dpm_solver.t(sigmas[i]), dpm_solver.t(sigmas[i + 1]), eps_cache=eps_cache) - case "dpm_2s": - 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}) - x, eps_cache = dpm_solver.dpm_solver_2_step(x, dpm_solver.t(sigmas[i]), dpm_solver.t(sigmas[i + 1]), eps_cache=eps_cache) - case "dpm_3s": - 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}) - x, eps_cache = dpm_solver.dpm_solver_3_step(x, dpm_solver.t(sigmas[i]), dpm_solver.t(sigmas[i + 1]), eps_cache=eps_cache) - case "rk4": # Fourth-order Runge-Kutta method - sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta) - # Calculate the derivative using the model - d = to_d(x, sigmas[i], denoised) - dt = sigma_down - sigmas[i] + orig_model = model + old_denoised = None + for i in trange(len(sigmas) - 1, disable=disable): + def model(x, sigma_s_in, **extra_args): + return apply_enhancements(x, i, orig_model(x, sigma_s_in, **extra_args), old_denoised) + # DynETA + eta = eta_schedule_cosine_annealing(i, len(sigmas)) - # Runge-Kutta steps - k1 = d * dt - k2 = to_d(x + k1 / 2, sigmas[i] + dt / 2, model(x + k1 / 2, (sigmas[i] + dt / 2) * s_in, **extra_args)) * dt - k3 = to_d(x + k2 / 2, sigmas[i] + dt / 2, model(x + k2 / 2, (sigmas[i] + dt / 2) * s_in, **extra_args)) * dt - k4 = to_d(x + k3, sigmas[i] + dt, model(x + k3, (sigmas[i] + dt) * s_in, **extra_args)) * dt + dpm_solver = DPMSolver(model, extra_args) - # Update the sample - x = x + (k1 + 2 * k2 + 2 * k3 + k4) / 6 - if sigmas[i + 1] > 0: - x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up - case "trapezoidal": - if sigmas[i + 1] > 0: - sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta) - dt = sigmas[i + 1] - sigmas[i] + # Renoising iterations + z_avg = torch.zeros_like(x) + sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta) + for k in range(substeps): + z_k = x + eps_cache = {} + denoised = model(z_k, sigmas[i] * s_in, **extra_args) + + if callback is not None: + callback({'x': z_k, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised}) + + eps = (z_k - denoised) / sigmas[i] + eps_cache = {'eps': eps} + + + match step_method if sigmas[i + 1] != 0 else "euler": + case "euler": + d = to_d(z_k, sigmas[i], denoised) + dt = sigma_down - sigmas[i] + + z_k = z_k + d * dt + case "dpm_1s": # DPM Family + 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}) + z_k, eps_cache = dpm_solver.dpm_solver_1_step(z_k, dpm_solver.t(sigmas[i]), dpm_solver.t(sigma_down), eps_cache=eps_cache) + case "dpm_2s": + 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}) + z_k, eps_cache = dpm_solver.dpm_solver_2_step(z_k, dpm_solver.t(sigmas[i]), dpm_solver.t(sigma_down), eps_cache=eps_cache) + case "dpm_3s": + 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}) + z_k, eps_cache = dpm_solver.dpm_solver_3_step(z_k, dpm_solver.t(sigmas[i]), dpm_solver.t(sigma_down), eps_cache=eps_cache) + case "rk4": # Fourth-order Runge-Kutta method # Calculate the derivative using the model - d_i = to_d(x, sigmas[i], denoised) + d = to_d(z_k, sigmas[i], denoised) + dt = sigma_down - sigmas[i] - # Predict the sample at the next sigma using Euler step - x_pred = x + d_i * dt + # Runge-Kutta steps + k1 = d * dt + k2 = to_d(z_k + k1 / 2, sigmas[i] + dt / 2, model(z_k + k1 / 2, (sigmas[i] + dt / 2) * s_in, **extra_args)) * dt + k3 = to_d(z_k + k2 / 2, sigmas[i] + dt / 2, model(z_k + k2 / 2, (sigmas[i] + dt / 2) * s_in, **extra_args)) * dt + k4 = to_d(z_k + k3, sigmas[i] + dt, model(z_k + k3, (sigmas[i] + dt) * s_in, **extra_args)) * dt - # Denoised sample at the next sigma - denoised_i_plus_1 = model(x_pred, sigmas[i + 1] * s_in, **extra_args) + # Update the sample + z_k = z_k + (k1 + 2 * k2 + 2 * k3 + k4) / 6 + case "trapezoidal": + if sigmas[i + 1] > 0: + dt = sigmas[i + 1] - sigmas[i] - # Calculate the derivative at the next sigma - d_i_plus_1 = to_d(x_pred, sigmas[i + 1], denoised_i_plus_1) + # Calculate the derivative using the model + d_i = to_d(z_k, sigmas[i], denoised) - #if callback is not None: - # callback({'x': x, 'i': i, 'sigma': sigmas[i + 1], 'denoised': denoised_i_plus_1}) - dt_2 = sigma_down - sigmas[i] - # Update the sample using the Trapezoidal rule - x = x + dt_2 * (d_i + d_i_plus_1) / 2 - x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up - else: - x = denoised + # Predict the sample at the next sigma using Euler step + x_pred = z_k + d_i * dt + # Denoised sample at the next sigma + denoised_i_plus_1 = model(x_pred, sigmas[i + 1] * s_in, **extra_args) + + # Calculate the derivative at the next sigma + d_i_plus_1 = to_d(x_pred, sigmas[i + 1], denoised_i_plus_1) + + dt_2 = sigma_down - sigmas[i] + # Update the sample using the Trapezoidal rule + z_k = z_k + dt_2 * (d_i + d_i_plus_1) / 2 + else: + z_k = denoised + + z_avg += renoise_weights[k] * z_k + if sigmas[i + 1] > 0: # Random noise for variance on ancestral samplers + z_k = z_k + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up + + x = z_avg + if sigmas[i + 1] > 0: + x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up old_denoised = denoised return x -def sample_supreme(model, x, sigmas, extra_args=None, callback=None, disable=None, s_noise=1., noise_sampler_type="gaussian", noise_sampler=None, eta=1.0, step_method="euler", centralization=0.02, normalization=0.01, edge_enhancement=0.5, perphist=-0.15): - return sampler_supreme(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, s_noise=s_noise, noise_sampler=noise_sampler or get_noise_sampler(x, sigmas, noise_sampler_type, noise_sampler, extra_args), eta=eta, step_method=step_method, centralization=centralization, normalization=normalization, edge_enhancement=edge_enhancement, perphist=perphist) +def sample_supreme(model, x, sigmas, extra_args=None, callback=None, disable=None, s_noise=1., noise_sampler_type="gaussian", noise_sampler=None, eta=1.0, step_method="euler", centralization=0.02, normalization=0.01, edge_enhancement=0.05, perphist=0, substeps=2): + return sampler_supreme(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, s_noise=s_noise, noise_sampler=noise_sampler or get_noise_sampler(x, sigmas, noise_sampler_type, noise_sampler, extra_args), eta=eta, step_method=step_method, centralization=centralization, normalization=normalization, edge_enhancement=edge_enhancement, perphist=perphist, substeps=substeps) # Add your personal samplers below here, just for formatting purposes ;3 diff --git a/nodes.py b/nodes.py index 3258872..aa6f99d 100644 --- a/nodes.py +++ b/nodes.py @@ -20,7 +20,7 @@ class SamplerRES_MOMENTUMIZED: } } RETURN_TYPES = ("SAMPLER",) - CATEGORY = "sampling/custom_sampling" + CATEGORY = "sampling/custom_sampling/samplers" FUNCTION = "get_sampler" @@ -40,7 +40,7 @@ class SamplerDPMPP_DUALSDE_MOMENTUMIZED: } } RETURN_TYPES = ("SAMPLER",) - CATEGORY = "sampling/custom_sampling" + CATEGORY = "sampling/custom_sampling/samplers" FUNCTION = "get_sampler" @@ -58,7 +58,7 @@ class SamplerTTM: } } RETURN_TYPES = ("SAMPLER",) - CATEGORY = "sampling/custom_sampling" + CATEGORY = "sampling/custom_sampling/samplers" FUNCTION = "get_sampler" @@ -75,7 +75,7 @@ class SamplerLCMCustom: } } RETURN_TYPES = ("SAMPLER",) - CATEGORY = "sampling/custom_sampling" + CATEGORY = "sampling/custom_sampling/samplers" FUNCTION = "get_sampler" @@ -94,7 +94,7 @@ class SamplerCLYB_4M_SDE_MOMENTUMIZED: } } RETURN_TYPES = ("SAMPLER",) - CATEGORY = "sampling/custom_sampling" + CATEGORY = "sampling/custom_sampling/samplers" FUNCTION = "get_sampler" @@ -114,7 +114,7 @@ class SamplerEULER_ANCESTRAL_DANCING: } } RETURN_TYPES = ("SAMPLER",) - CATEGORY = "sampling/custom_sampling" + CATEGORY = "sampling/custom_sampling/samplers" FUNCTION = "get_sampler" @@ -133,7 +133,7 @@ class SamplerDPMPP_3M_SDE_DYN_ETA: } } RETURN_TYPES = ("SAMPLER",) - CATEGORY = "sampling/custom_sampling" + CATEGORY = "sampling/custom_sampling/samplers" FUNCTION = "get_sampler" @@ -150,18 +150,19 @@ class SamplerSUPREME: "eta": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step":0.01}), "centralization": ("FLOAT", {"default": 0.02, "min": -1.0, "max": 1.0, "step":0.01}), "normalization": ("FLOAT", {"default": 0.01, "min": -1.0, "max": 1.0, "step":0.01}), - "edge_enhancement": ("FLOAT", {"default": 0.5, "min": -100.0, "max": 100.0, "step":0.01}), - "perphist": ("FLOAT", {"default": -0.15, "min": -5.0, "max": 5.0, "step":0.01}), + "edge_enhancement": ("FLOAT", {"default": 0.05, "min": -100.0, "max": 100.0, "step":0.01}), + "perphist": ("FLOAT", {"default": 0, "min": -5.0, "max": 5.0, "step":0.01}), + "substeps": ("INT", {"default": 2, "min": 1, "max": 100, "step":1}), "s_noise": ("FLOAT", {"default": 1, "min": 0.0, "max": 100.0, "step":0.01}), } } RETURN_TYPES = ("SAMPLER",) - CATEGORY = "sampling/custom_sampling" + CATEGORY = "sampling/custom_sampling/samplers" FUNCTION = "get_sampler" - def get_sampler(self, noise_sampler_type, step_method, eta, centralization, normalization, edge_enhancement, perphist, s_noise): - sampler = comfy.samplers.ksampler("supreme", {"noise_sampler_type": noise_sampler_type, "step_method": step_method, "eta": eta, "centralization": centralization, "normalization": normalization, "edge_enhancement": edge_enhancement, "perphist": perphist, "s_noise": s_noise}) + def get_sampler(self, noise_sampler_type, step_method, eta, centralization, normalization, edge_enhancement, perphist, substeps, s_noise): + sampler = comfy.samplers.ksampler("supreme", {"noise_sampler_type": noise_sampler_type, "step_method": step_method, "eta": eta, "centralization": centralization, "normalization": normalization, "edge_enhancement": edge_enhancement, "perphist": perphist, "substeps": substeps, "s_noise": s_noise}) return (sampler, ) ### Schedulers diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..03ab5b2 --- /dev/null +++ b/requirements.txt @@ -0,0 +1 @@ +kornia