nodes/Clybius-ComfyUI-Extra-Samplers
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Upgrade Supreme sampler with more options, fix some latent exploding

Browse Source 
Update README.md
Move samplers over into their correct category
Add Kornia as a requirement
This commit is contained in:
Clybius
2024-03-25 16:12:05 -05:00
parent 6f66795962
commit 448ecdfedd
4 changed files with 117 additions and 92 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+13 -1
View File
@@ -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
+88 -77
View File
@@ -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)
@@ -873,85 +867,102 @@ 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
return denoised
if callback is not None:
callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
renoise_weights = torch.ones(substeps, device=x.device) / substeps
eps = (x - denoised) / sigmas[i]
eps_cache = {'eps': eps}
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))
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]
dpm_solver = DPMSolver(model, extra_args)
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]
# 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 = {}
# 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
denoised = model(z_k, sigmas[i] * s_in, **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]
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
nodes.py
+13 -12
View File
@@ -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
requirements.txt
+1
View File
@@ -0,0 +1 @@
kornia