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

@@ -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!

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):
+    def apply_enhancements(x, i, denoised, old_denoised):
-        """Function representing the Karras ODE derivative."""
+        if edge_enhancement != 0:
-        denoised = model(x, sigma * s_in, **extra_args)
+            blur = (kornia.filters.joint_bilateral_blur(x, denoised, (3, 3), 0.1, (1.5, 1.5)) - x) # Blurs non-edges
-        return to_d(x, sigma, denoised)
+            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
-    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)
 
         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:
+        return denoised
-            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
-        
-        eps = (x - denoised) / sigmas[i]
-        eps_cache = {'eps': eps}
 
-        match step_method:
+    renoise_weights = torch.ones(substeps, device=x.device) / substeps
-            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]
 
-                x = x + d * dt
+    orig_model = model
-                if sigmas[i + 1] > 0:
+    old_denoised = None
-                    x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up
+    for i in trange(len(sigmas) - 1, disable=disable):
-            case "dpm_1s": # DPM Family
+        def model(x, sigma_s_in, **extra_args):
-                if callback is not None:
+            return apply_enhancements(x, i, orig_model(x, sigma_s_in, **extra_args), old_denoised)
-                    dpm_solver.info_callback = lambda info: callback({'sigma': dpm_solver.sigma(info['t']), 'sigma_hat': dpm_solver.sigma(info['t_up']), **info})
+        # DynETA
-                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)
+        eta = eta_schedule_cosine_annealing(i, len(sigmas))
-            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]
 
-                # Runge-Kutta steps
+        dpm_solver = DPMSolver(model, extra_args)
-                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
 
-                # Update the sample
+        # Renoising iterations
-                x = x + (k1 + 2 * k2 + 2 * k3 + k4) / 6
+        z_avg = torch.zeros_like(x)
-                if sigmas[i + 1] > 0:
+        sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta)
-                    x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up
+        for k in range(substeps):
-            case "trapezoidal":
+            z_k = x
-                if sigmas[i + 1] > 0:
+            eps_cache = {}
-                    sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta)
-                    dt = sigmas[i + 1] - sigmas[i]
 
+            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
+                    # Runge-Kutta steps
-                    x_pred = x + d_i * dt
+                    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
+                    # Update the sample
-                    denoised_i_plus_1 = model(x_pred, sigmas[i + 1] * s_in, **extra_args)
+                    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
+                        # Calculate the derivative using the model
-                    d_i_plus_1 = to_d(x_pred, sigmas[i + 1], denoised_i_plus_1)
+                        d_i = to_d(z_k, sigmas[i], denoised)
 
-                    #if callback is not None:
+                        # Predict the sample at the next sigma using Euler step
-                    #    callback({'x': x, 'i': i, 'sigma': sigmas[i + 1], 'denoised': denoised_i_plus_1})
+                        x_pred = z_k + d_i * dt
-                    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
 
+                        # 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):
+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)
+    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
 

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}),
+                     "edge_enhancement": ("FLOAT", {"default": 0.05, "min": -100.0, "max": 100.0, "step":0.01}),
-                     "perphist": ("FLOAT", {"default": -0.15, "min": -5.0, "max": 5.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):
+    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, "s_noise": 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

requirements.txt

@@ -0,0 +1 @@
+kornia