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

Use the ComfyUI version of k_diffusion.sampling

Browse Source
This commit is contained in:
blepping
2024-02-14 03:38:26 -07:00
parent e58d971bec
commit 50799c9474
2 changed files with 8 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
extra_samplers.py
+6 -7
View File
@@ -8,8 +8,7 @@ from tqdm.auto import trange, tqdm
import comfy.sample import comfy.sample
import k_diffusion.sampling from comfy.k_diffusion.sampling import BrownianTreeNoiseSampler, PIDStepSizeController, get_ancestral_step, to_d, default_noise_sampler
from k_diffusion.sampling import BrownianTreeNoiseSampler, PIDStepSizeController, get_ancestral_step, to_d, default_noise_sampler
import random import random
# The following function adds the samplers during initialization, in __init__.py # The following function adds the samplers during initialization, in __init__.py
@@ -251,7 +250,7 @@ def highres_pyramid_noise_like(x, discount=0.7):
u = torch.nn.Upsample(size=(orig_h, orig_w), mode='bilinear') u = torch.nn.Upsample(size=(orig_h, orig_w), mode='bilinear')
noise = (torch.rand_like(x) - 0.5) * 2 * 1.73 # Start with scaled uniform noise noise = (torch.rand_like(x) - 0.5) * 2 * 1.73 # Start with scaled uniform noise
for i in range(4): for i in range(4):
r = random.random()*2+2 # Rather than always going 2x, r = random.random()*2+2 # Rather than always going 2x,
h, w = min(orig_h*15, int(h*(r**i))), min(orig_w*15, int(w*(r**i))) h, w = min(orig_h*15, int(h*(r**i))), min(orig_w*15, int(w*(r**i)))
noise += u(torch.randn(b, c, h, w).to(x)) * discount**i noise += u(torch.randn(b, c, h, w).to(x)) * discount**i
if h>=orig_h*15 or w>=orig_w*15: break # Lowest resolution is 1x1 if h>=orig_h*15 or w>=orig_w*15: break # Lowest resolution is 1x1
@@ -322,7 +321,7 @@ def sample_clyb_4m_sde_momentumized(model, x, sigmas, extra_args=None, callback=
The expression for d1 is derived from the extrapolation formula given in the paper “Diffusion Monte Carlo with stochastic Hamiltonians” by M. Foulkes, L. Mitas, R. Needs, and G. Rajagopal. The formula is given as follows: The expression for d1 is derived from the extrapolation formula given in the paper “Diffusion Monte Carlo with stochastic Hamiltonians” by M. Foulkes, L. Mitas, R. Needs, and G. Rajagopal. The formula is given as follows:
d1 = d1_0 + (d1_0 - d1_1) * r2 / (r2 + r1) + ((d1_0 - d1_1) * r2 / (r2 + r1) - (d1_1 - d1_2) * r1 / (r0 + r1)) * r2 / ((r2 + r1) * (r0 + r1)) d1 = d1_0 + (d1_0 - d1_1) * r2 / (r2 + r1) + ((d1_0 - d1_1) * r2 / (r2 + r1) - (d1_1 - d1_2) * r1 / (r0 + r1)) * r2 / ((r2 + r1) * (r0 + r1))
(if this is an incorrect citing, we blame Google's Bard and OpenAI's ChatGPT for this and NOT me :^) ) (if this is an incorrect citing, we blame Google's Bard and OpenAI's ChatGPT for this and NOT me :^) )
where d1_0, d1_1, and d1_2 are defined as follows: where d1_0, d1_1, and d1_2 are defined as follows:
d1_0 = (denoised - denoised_1) / r2 d1_0 = (denoised - denoised_1) / r2
d1_1 = (denoised_1 - denoised_2) / r1 d1_1 = (denoised_1 - denoised_2) / r1
@@ -403,7 +402,7 @@ def sample_clyb_4m_sde_momentumized(model, x, sigmas, extra_args=None, callback=
if eta: if eta:
x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * sigmas[i + 1] * (-2 * h * eta).expm1().neg().sqrt() * s_noise x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * sigmas[i + 1] * (-2 * h * eta).expm1().neg().sqrt() * s_noise
denoised_1, denoised_2, denoised_3 = denoised, denoised_1, denoised_2 denoised_1, denoised_2, denoised_3 = denoised, denoised_1, denoised_2
h_1, h_2, h_3 = h, h_1, h_2 h_1, h_2, h_3 = h, h_1, h_2
return x return x
@@ -596,7 +595,7 @@ def sample_ttmcustom(model, x, sigmas, extra_args=None, callback=None, disable=N
noise_sampler = lambda sigma, sigma_next: torch.randn_like(x) noise_sampler = lambda sigma, sigma_next: torch.randn_like(x)
return sample_ttm_jvp(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, eta=eta, s_noise=s_noise, noise_sampler=noise_sampler) return sample_ttm_jvp(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, eta=eta, s_noise=s_noise, noise_sampler=noise_sampler)
from k_diffusion.sampling import sample_lcm from comfy.k_diffusion.sampling import sample_lcm
def sample_lcmcustom(model, x, sigmas, extra_args=None, callback=None, disable=None, noise_sampler=None): def sample_lcmcustom(model, x, sigmas, extra_args=None, callback=None, disable=None, noise_sampler=None):
sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max() sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
seed = extra_args.get("seed", None) seed = extra_args.get("seed", None)
@@ -638,4 +637,4 @@ extra_samplers = {
"clyb_4m_sde_momentumized": sample_clyb_4m_sde, "clyb_4m_sde_momentumized": sample_clyb_4m_sde,
"ttm": sample_ttmcustom, "ttm": sample_ttmcustom,
"lcm_custom_noise": sample_lcmcustom, "lcm_custom_noise": sample_lcmcustom,
} }
other_samplers/sample_ttm.py
+2 -2
View File
@@ -1,4 +1,4 @@
from k_diffusion.sampling import default_noise_sampler, to_d from comfy.k_diffusion.sampling import default_noise_sampler, to_d
from tqdm import trange from tqdm import trange
import torch import torch
from torch import enable_grad from torch import enable_grad
@@ -115,4 +115,4 @@ def sample_lcm_ttm_jvp(model, x, sigmas, extra_args=None, callback=None, disable
if sigmas[i + 1] > 0: if sigmas[i + 1] > 0:
x = x + sigmas[i + 1] * noise_sampler(sigmas[i], sigmas[i + 1]) x = x + sigmas[i + 1] * noise_sampler(sigmas[i], sigmas[i + 1])
return x return x