Allow overriding noise_sampler + reduce code duplication

.gitignore

@@ -0,0 +1,3 @@
+*.bak
+*~
+__pycache__

extra_samplers.py

@@ -4,6 +4,7 @@ import torch
 from torch import nn, FloatTensor
 import torchsde
 from tqdm.auto import trange, tqdm
+import numpy as np
 
 import comfy.sample
 
@@ -45,12 +46,33 @@ def add_schedulers():
         import importlib
         importlib.reload(k_diffusion_sampling)
 
+
 # Noise samplers
+NOISE_SAMPLER_NAMES=("gaussian", "uniform", "brownian", "highres-pyramid", "pyramid", "perlin", "laplacian")
+
+def get_noise_sampler_names(default=None):
+    if not default:
+        return NOISE_SAMPLER_NAMES
+    return (default,) + tuple(n for n in NOISE_SAMPLER_NAMES if n != default)
+
+def mk_noise_sampler(x, fun):
+    return lambda _sigma, _sigma_next: fun(x)
+
+def get_noise_sampler(x, sigmas, noise_sampler_type="brownian", extra_args=None, cpu=False):
+    if noise_sampler_type == "brownian":
+        seed = extra_args.get("seed", None) if extra_args else None
+        sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
+        return BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=seed, cpu=cpu)
+    return mk_noise_sampler(x, NOISE_SAMPLER_HANDLERS.get(noise_sampler_type, uniform_noise_like))
+
 from torch import Generator, Tensor, lerp
 from torch.nn.functional import unfold
 from typing import Callable, Tuple
 from math import pi
 
+def uniform_noise_like(x):
+    return (torch.rand_like(x) - 0.5) * 2 * 1.73
+
 def get_positions(block_shape: Tuple[int, int]) -> Tensor:
     """
     Generate position tensor.
@@ -307,7 +329,8 @@ def green_noise_sampler(x): # This doesn't work properly right now
     print(noise)
     return lambda sigma, sigma_next: noise
 
-def power_noise_sampler(tensor, alpha=2, k=1): # This doesn't work properly right now
+# I'm not sure how this differs from the other implementation but it doesn't seem to be used at present.
+def power_noise_sampler_2(tensor, alpha=2, k=1): # This doesn't work properly right now
     """Generate 1/f noise for a given tensor.
 
     Args:
@@ -330,6 +353,83 @@ def power_noise_sampler(tensor, alpha=2, k=1): # This doesn't work properly righ
     print(variance)
     return lambda sigma, sigma_next: noise / 3
 
+def pyramid_noise_like(size, dtype, layout, generator, device="cpu", discount=0.8):
+    b, c, h, w = size
+    orig_h = h
+    orig_w = w
+    noise = torch.zeros(size=size, dtype=dtype, layout=layout, device=device)
+    r = 1
+    for i in range(5):
+        r *= 2 # Rather than always going 2x,
+        #w, h = max(1, int(w/(r**i))), max(1, int(h/(r**i)))
+        noise += torch.nn.functional.interpolate((torch.normal(mean=0, std=0.5 ** i, size=(b, c, h * r, w * r), dtype=dtype, layout=layout, generator=generator, device=device)), size=(orig_h, orig_w), mode='nearest-exact') * discount**i
+        #if w>=orig_w*16 or h>=orig_h*16: break
+    return noise
+
+def power_noise_sampler(size, dtype, layout, generator, device="cpu", alpha=2, k=1): # This doesn't work properly right now
+    """Generate 1/f noise for a given tensor.
+
+    Args:
+        tensor: The tensor to add noise to.
+        alpha: The parameter that determines the slope of the spectrum.
+        k: A constant.
+
+    Returns:
+        A tensor with the same shape as `tensor` containing 1/f noise.
+    """
+    tensor = torch.randn(size=size, dtype=dtype, layout=layout, generator=generator, device=device)
+    fft = torch.fft.fft2(tensor)
+    freq = torch.arange(1, len(fft) + 1, dtype=torch.float)
+    spectral_density = k / freq**alpha
+    noise = torch.rand(size=size, dtype=dtype, layout=layout, generator=generator, device=device) * spectral_density
+    mean = torch.mean(noise, dim=(-2, -1), keepdim=True).to(tensor.device)
+    std = torch.std(noise, dim=(-2, -1), keepdim=True).to(tensor.device)
+    noise = noise.to(tensor.device).sub_(mean).div_(std)
+    return noise
+
+def prepare_noise(latent_image, seed, noise_type, noise_inds=None): # From `sample.py`
+    """
+    creates random noise given a latent image and a seed.
+    optional arg skip can be used to skip and discard x number of noise generations for a given seed
+    """
+    generator = torch.manual_seed(seed)
+    match noise_type:
+        case "gaussian":
+            noise_func = torch.randn
+        case "uniform":
+            def uniform_rand(*size, **kwargs):
+                return (torch.rand(*size, **kwargs) - 0.5) * 2 * 1.73
+            noise_func = uniform_rand
+        case "pyramid":
+            noise_func = pyramid_noise_like
+        case "power":
+            noise_func = power_noise_sampler
+        case _:
+            noise_func = torch.randn
+    if noise_inds is None:
+        return noise_func(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, generator=generator, device="cpu")
+
+    unique_inds, inverse = np.unique(noise_inds, return_inverse=True)
+    noises = []
+    for i in range(unique_inds[-1]+1):
+        noise = noise_func([1] + list(latent_image.size())[1:], dtype=latent_image.dtype, layout=latent_image.layout, generator=generator, device="cpu")
+        if i in unique_inds:
+            noises.append(noise)
+    noises = [noises[i] for i in inverse]
+    noises = torch.cat(noises, axis=0)
+    return noises
+
+NOISE_SAMPLER_HANDLERS={
+    # Brownian is special-cased.
+    "gaussian": torch.randn_like,
+    "highres-pyramid": highres_pyramid_noise_like,
+    "pyramid": lambda x: pyramid_noise_like(x.size(), x.dtype, x.layout, None, device=x.device),
+    "perlin": rand_perlin_like,
+    "laplacian": rand_laplacian_like,
+    "uniform": uniform_noise_like,
+}
+
+
 # Below this point are extra samplers
 @torch.no_grad()
 def sample_clyb_4m_sde_momentumized(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1.0, s_noise=1., noise_sampler=None, momentum=0.0):
@@ -462,25 +562,8 @@ def sample_ttm_jvp(model, x, sigmas, extra_args=None, callback=None, disable=Non
 
 # Many thanks to Kat + Birch-San for this wonderful sampler implementation! https://github.com/Birch-san/sdxl-play/commits/res/
 from .other_samplers.refined_exp_solver import sample_refined_exp_s
-def sample_res_solver(model, x, sigmas, extra_args=None, callback=None, disable=None, noise_sampler=None, denoise_to_zero=True, simple_phi_calc=False, c2=0.5, ita=torch.Tensor((0.25,)), momentum=0.0):
+def sample_res_solver(model, x, sigmas, extra_args=None, callback=None, disable=None, noise_sampler_type="gaussian", noise_sampler=None, denoise_to_zero=True, simple_phi_calc=False, c2=0.5, ita=torch.Tensor((0.25,)), momentum=0.0):
-    sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
+    return sample_refined_exp_s(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, noise_sampler=noise_sampler or get_noise_sampler(x, sigmas, noise_sampler_type, noise_sampler, extra_args), denoise_to_zero=denoise_to_zero, simple_phi_calc=simple_phi_calc, c2=c2, ita=ita, momentum=momentum)
-    seed = extra_args.get("seed", None)
-    match noise_sampler:
-        case "brownian":
-            noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=seed, cpu=False)
-        case "gaussian":
-            noise_sampler = lambda sigma, sigma_next: torch.randn_like(x)
-        case "uniform":
-            noise_sampler = lambda sigma, sigma_next: (torch.rand_like(x) - 0.5) * 2 * 1.73
-        case "highres-pyramid":
-            noise_sampler = lambda sigma, sigma_next: highres_pyramid_noise_like(x)
-        case "perlin":
-            noise_sampler = lambda sigma, sigma_next: rand_perlin_like(x)
-        case "laplacian":
-            noise_sampler = lambda sigma, sigma_next: rand_laplacian_like(x)
-        case _:
-            noise_sampler = lambda sigma, sigma_next: torch.randn_like(x)
-    return sample_refined_exp_s(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, noise_sampler=noise_sampler, denoise_to_zero=denoise_to_zero, simple_phi_calc=simple_phi_calc, c2=c2, ita=ita, momentum=momentum)
 
 @torch.no_grad()
 def sample_dpmpp_dualsde_momentum(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, r=1/2, momentum=0.0):
@@ -582,73 +665,19 @@ def sample_dpmpp_dualsde_momentum(model, x, sigmas, extra_args=None, callback=No
             h_1, h_2, h_3 = h, h_1, h_2
     return x
 
-def sample_dpmpp_dualsdemomentum(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, r=1/2, momentum=0.0):
+def sample_dpmpp_dualsdemomentum(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler_type="gaussian", noise_sampler=None, r=1/2, momentum=0.0):
-    sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
+    return sample_dpmpp_dualsde_momentum(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, eta=eta, s_noise=s_noise, noise_sampler=noise_sampler or get_noise_sampler(x, sigmas, noise_sampler_type, noise_sampler, extra_args), r=r, momentum=momentum)
-    seed = extra_args.get("seed", None)
-    match noise_sampler:
-        case "brownian":
-            noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=seed, cpu=False)
-        case "gaussian":
-            noise_sampler = lambda sigma, sigma_next: torch.randn_like(x)
-        case "uniform":
-            noise_sampler = lambda sigma, sigma_next: (torch.rand_like(x) - 0.5) * 2 * 1.73
-        case "perlin":
-            noise_sampler = lambda sigma, sigma_next: rand_perlin_like(x)
-        case "laplacian":
-            noise_sampler = lambda sigma, sigma_next: rand_laplacian_like(x)
-        case _:
-            noise_sampler = lambda sigma, sigma_next: torch.randn_like(x)
-    return sample_dpmpp_dualsde_momentum(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, eta=eta, s_noise=s_noise, noise_sampler=noise_sampler, r=r, momentum=momentum)
 
 from .other_samplers.sample_ttm import sample_ttm_jvp
-def sample_ttmcustom(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None):
+def sample_ttmcustom(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler_type="gaussian",noise_sampler=None):
-    sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
+    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 or get_noise_sampler(x, sigmas, noise_sampler_type, noise_sampler, extra_args))
-    seed = extra_args.get("seed", None)
-    match noise_sampler:
-        case "gaussian":
-            noise_sampler = lambda sigma, sigma_next: torch.randn_like(x)
-        case "uniform":
-            noise_sampler = lambda sigma, sigma_next: (torch.rand_like(x) - 0.5) * 2 * 1.73
-        case "brownian":
-            noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=seed, cpu=False)
-        case _:
-            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)
 
 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_type="gaussian", noise_sampler=None):
-    sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
+    return sample_lcm(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, noise_sampler=noise_sampler or get_noise_sampler(x, sigmas, noise_sampler_type, noise_sampler, extra_args))
-    seed = extra_args.get("seed", None)
-    match noise_sampler:
-        case "gaussian":
-            noise_sampler = lambda sigma, sigma_next: torch.randn_like(x)
-        case "uniform":
-            noise_sampler = lambda sigma, sigma_next: (torch.rand_like(x) - 0.5) * 2 * 1.73
-        case "brownian":
-            noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=seed, cpu=False)
-        case _:
-            noise_sampler = lambda sigma, sigma_next: torch.randn_like(x)
-    return sample_lcm(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, noise_sampler=noise_sampler)
 
-def sample_clyb_4m_sde(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler="brownian", momentum=0.0):
+def sample_clyb_4m_sde(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler_type="brownian", noise_sampler=None, momentum=0.0):
-    sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
+    return sample_clyb_4m_sde_momentumized(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, eta=eta, s_noise=s_noise, noise_sampler=noise_sampler or get_noise_sampler(x, sigmas, noise_sampler_type, noise_sampler, extra_args), momentum=momentum)
-    seed = extra_args.get("seed", None)
-    match noise_sampler:
-        case "brownian":
-            noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=seed, cpu=False)
-        case "gaussian":
-            noise_sampler = lambda sigma, sigma_next: torch.randn_like(x)
-        case "uniform":
-            noise_sampler = lambda sigma, sigma_next: (torch.rand_like(x) - 0.5) * 2 * 1.73
-        case "highres-pyramid":
-            noise_sampler = lambda sigma, sigma_next: highres_pyramid_noise_like(x)
-        case "perlin":
-            noise_sampler = lambda sigma, sigma_next: rand_perlin_like(x)
-        case "laplacian":
-            noise_sampler = lambda sigma, sigma_next: rand_laplacian_like(x)
-        case _:
-            noise_sampler = lambda sigma, sigma_next: (torch.rand_like(x) - 0.5) * 2 * 1.73
-    return sample_clyb_4m_sde_momentumized(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, eta=eta, s_noise=s_noise, noise_sampler=noise_sampler, momentum=momentum)
 
 
 # This code works, but I'm currently experimenting with different methods
@@ -698,25 +727,8 @@ def sampler_euler_ancestral_dancing(model, x, sigmas, extra_args=None, callback=
 
     return x
 
-def sample_euler_ancestral_dancing(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler="gaussian", leap=2, eta_dance=1.0):
+def sample_euler_ancestral_dancing(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler_type="gaussian", noise_sampler=None, leap=2, eta_dance=1.0):
-    sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
+    return sampler_euler_ancestral_dancing(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, eta=eta, s_noise=s_noise, noise_sampler=noise_sampler or get_noise_sampler(x, sigmas, noise_sampler_type, noise_sampler, extra_args), leap=leap, eta_dance=eta_dance)
-    seed = extra_args.get("seed", None)
-    match noise_sampler:
-        case "brownian":
-            noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=seed, cpu=False)
-        case "gaussian":
-            noise_sampler = lambda sigma, sigma_next: torch.randn_like(x)
-        case "uniform":
-            noise_sampler = lambda sigma, sigma_next: (torch.rand_like(x) - 0.5) * 2 * 1.73
-        case "highres-pyramid":
-            noise_sampler = lambda sigma, sigma_next: highres_pyramid_noise_like(x)
-        case "perlin":
-            noise_sampler = lambda sigma, sigma_next: rand_perlin_like(x)
-        case "laplacian":
-            noise_sampler = lambda sigma, sigma_next: rand_laplacian_like(x)
-        case _:
-            noise_sampler = lambda sigma, sigma_next: (torch.rand_like(x) - 0.5) * 2 * 1.73
-    return sampler_euler_ancestral_dancing(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, eta=eta, s_noise=s_noise, noise_sampler=noise_sampler, leap=leap, eta_dance=eta_dance)
 
 
 @torch.no_grad()
@@ -778,25 +790,8 @@ def sampler_dpmpp_3m_sde_dynamic_eta(model, x, sigmas, extra_args=None, callback
         h_1, h_2 = h, h_1
     return x
 
-def sample_dpmpp_3m_sde_dynamic_eta(model, x, sigmas, extra_args=None, callback=None, disable=None, eta_max=1.0, eta_min=0.0, s_noise=1., noise_sampler="brownian"):
+def sample_dpmpp_3m_sde_dynamic_eta(model, x, sigmas, extra_args=None, callback=None, disable=None, eta_max=1.0, eta_min=0.0, s_noise=1., noise_sampler_type="brownian", noise_sampler=None):
-    sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
+    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))
-    seed = extra_args.get("seed", None)
-    match noise_sampler:
-        case "brownian":
-            noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=seed, cpu=False)
-        case "gaussian":
-            noise_sampler = lambda sigma, sigma_next: torch.randn_like(x)
-        case "uniform":
-            noise_sampler = lambda sigma, sigma_next: (torch.rand_like(x) - 0.5) * 2 * 1.73
-        case "highres-pyramid":
-            noise_sampler = lambda sigma, sigma_next: highres_pyramid_noise_like(x)
-        case "perlin":
-            noise_sampler = lambda sigma, sigma_next: rand_perlin_like(x)
-        case "laplacian":
-            noise_sampler = lambda sigma, sigma_next: rand_laplacian_like(x)
-        case _:
-            noise_sampler = lambda sigma, sigma_next: (torch.rand_like(x) - 0.5) * 2 * 1.73
-    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)
 
 # Add your personal samplers below here, just for formatting purposes ;3
 

nodes.py

@@ -1,83 +1,17 @@
 from .other_samplers.refined_exp_solver import sample_refined_exp_s
+from .extra_samplers import get_noise_sampler_names, prepare_noise
 import comfy.samplers
 import comfy.sample
 from comfy.k_diffusion import sampling as k_diffusion_sampling
 import latent_preview
 import torch
-import numpy as np
 from tqdm.auto import trange
 
-def pyramid_noise_like(size, dtype, layout, generator, device="cpu", discount=0.8):
-    b, c, h, w = size
-    orig_h = h
-    orig_w = w
-    noise = torch.zeros(size=size, dtype=dtype, layout=layout, device=device)
-    r = 1
-    for i in range(5):
-        r *= 2 # Rather than always going 2x,
-        #w, h = max(1, int(w/(r**i))), max(1, int(h/(r**i)))
-        noise += torch.nn.functional.interpolate((torch.normal(mean=0, std=0.5 ** i, size=(b, c, h * r, w * r), dtype=dtype, layout=layout, generator=generator, device=device)), size=(orig_h, orig_w), mode='nearest-exact') * discount**i
-        #if w>=orig_w*16 or h>=orig_h*16: break
-    return noise
-
-def power_noise_sampler(size, dtype, layout, generator, device="cpu", alpha=2, k=1): # This doesn't work properly right now
-    """Generate 1/f noise for a given tensor.
-
-    Args:
-        tensor: The tensor to add noise to.
-        alpha: The parameter that determines the slope of the spectrum.
-        k: A constant.
-
-    Returns:
-        A tensor with the same shape as `tensor` containing 1/f noise.
-    """
-    tensor = torch.randn(size=size, dtype=dtype, layout=layout, generator=generator, device=device)
-    fft = torch.fft.fft2(tensor)
-    freq = torch.arange(1, len(fft) + 1, dtype=torch.float)
-    spectral_density = k / freq**alpha
-    noise = torch.rand(size=size, dtype=dtype, layout=layout, generator=generator, device=device) * spectral_density
-    mean = torch.mean(noise, dim=(-2, -1), keepdim=True).to(tensor.device)
-    std = torch.std(noise, dim=(-2, -1), keepdim=True).to(tensor.device)
-    noise = noise.to(tensor.device).sub_(mean).div_(std)
-    return noise
-
-def prepare_noise(latent_image, seed, noise_type, noise_inds=None): # From `sample.py`
-    """
-    creates random noise given a latent image and a seed.
-    optional arg skip can be used to skip and discard x number of noise generations for a given seed
-    """
-    generator = torch.manual_seed(seed)
-    match noise_type:
-        case "gaussian":
-            noise_func = torch.randn
-        case "uniform":
-            def uniform_rand(*size, **kwargs):
-                return (torch.rand(*size, **kwargs) - 0.5) * 2 * 1.73
-            noise_func = uniform_rand
-        case "pyramid":
-            noise_func = pyramid_noise_like
-        case "power":
-            noise_func = power_noise_sampler
-        case _:
-            noise_func = torch.randn
-    if noise_inds is None:
-        return noise_func(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, generator=generator, device="cpu")
-
-    unique_inds, inverse = np.unique(noise_inds, return_inverse=True)
-    noises = []
-    for i in range(unique_inds[-1]+1):
-        noise = noise_func([1] + list(latent_image.size())[1:], dtype=latent_image.dtype, layout=latent_image.layout, generator=generator, device="cpu")
-        if i in unique_inds:
-            noises.append(noise)
-    noises = [noises[i] for i in inverse]
-    noises = torch.cat(noises, axis=0)
-    return noises
-
 class SamplerRES_MOMENTUMIZED:
     @classmethod
     def INPUT_TYPES(s):
         return {"required":
-                    {"noise_sampler_type": (["gaussian", "uniform", "brownian", "highres-pyramid", "perlin", "laplacian"], ),
+                    {"noise_sampler_type": (get_noise_sampler_names(), ),
                      "momentum": ("FLOAT", {"default": 0.5, "min": -1.0, "max": 1.0, "step":0.01}),
                      "denoise_to_zero": ("BOOLEAN", {"default": True}),
                      "simple_phi_calc": ("BOOLEAN", {"default": False}),
@@ -91,14 +25,14 @@ class SamplerRES_MOMENTUMIZED:
     FUNCTION = "get_sampler"
 
     def get_sampler(self, noise_sampler_type, momentum, denoise_to_zero, simple_phi_calc, ita, c2):
-        sampler = comfy.samplers.ksampler("res_momentumized", {"noise_sampler": noise_sampler_type, "denoise_to_zero": denoise_to_zero, "simple_phi_calc": simple_phi_calc, "c2": c2, "ita": torch.Tensor((ita,)), "momentum": momentum})
+        sampler = comfy.samplers.ksampler("res_momentumized", {"noise_sampler_type": noise_sampler_type, "denoise_to_zero": denoise_to_zero, "simple_phi_calc": simple_phi_calc, "c2": c2, "ita": torch.Tensor((ita,)), "momentum": momentum})
         return (sampler, )
 
 class SamplerDPMPP_DUALSDE_MOMENTUMIZED:
     @classmethod
     def INPUT_TYPES(s):
         return {"required":
-                    {"noise_sampler_type": (["gaussian", "uniform", "brownian", "perlin", "laplacian"], ),
+                    {"noise_sampler_type": (get_noise_sampler_names(), ),
                      "momentum": ("FLOAT", {"default": 0.5, "min": -1.0, "max": 1.0, "step":0.01}),
                      "eta": ("FLOAT", {"default": 1, "min": 0.0, "max": 100.0, "step":0.01}),
                      "s_noise": ("FLOAT", {"default": 1, "min": 0.0, "max": 100.0, "step":0.01}),
@@ -111,14 +45,14 @@ class SamplerDPMPP_DUALSDE_MOMENTUMIZED:
     FUNCTION = "get_sampler"
 
     def get_sampler(self, noise_sampler_type, momentum, eta, s_noise, r,):
-        sampler = comfy.samplers.ksampler("dpmpp_dualsde_momentumized", {"noise_sampler": noise_sampler_type, "eta": eta, "s_noise": s_noise, "r": r, "momentum": momentum})
+        sampler = comfy.samplers.ksampler("dpmpp_dualsde_momentumized", {"noise_sampler_type": noise_sampler_type, "eta": eta, "s_noise": s_noise, "r": r, "momentum": momentum})
         return (sampler, )
 
 class SamplerTTM:
     @classmethod
     def INPUT_TYPES(s):
         return {"required":
-                    {"noise_sampler_type": (["gaussian", "uniform", "brownian"], ),
+                    {"noise_sampler_type": (get_noise_sampler_names(), ),
                      "eta": ("FLOAT", {"default": 1, "min": 0.0, "max": 100.0, "step":0.01}),
                      "s_noise": ("FLOAT", {"default": 1, "min": 0.0, "max": 100.0, "step":0.01}),
                       }
@@ -129,7 +63,7 @@ class SamplerTTM:
     FUNCTION = "get_sampler"
 
     def get_sampler(self, noise_sampler_type, eta, s_noise):
-        sampler = comfy.samplers.ksampler("ttm", {"noise_sampler": noise_sampler_type, "eta": eta, "s_noise": s_noise})
+        sampler = comfy.samplers.ksampler("ttm", {"noise_sampler_type": noise_sampler_type, "eta": eta, "s_noise": s_noise})
         return (sampler, )
 
 
@@ -137,7 +71,7 @@ class SamplerLCMCustom:
     @classmethod
     def INPUT_TYPES(s):
         return {"required":
-                    {"noise_sampler_type": (["gaussian", "uniform", "brownian"], ),
+                    {"noise_sampler_type": (get_noise_sampler_names(), ),
                       }
                }
     RETURN_TYPES = ("SAMPLER",)
@@ -146,14 +80,14 @@ class SamplerLCMCustom:
     FUNCTION = "get_sampler"
 
     def get_sampler(self, noise_sampler_type):
-        sampler = comfy.samplers.ksampler("lcm_custom_noise", {"noise_sampler": noise_sampler_type})
+        sampler = comfy.samplers.ksampler("lcm_custom_noise", {"noise_sampler_type": noise_sampler_type})
         return (sampler, )
 
 class SamplerCLYB_4M_SDE_MOMENTUMIZED:
     @classmethod
     def INPUT_TYPES(s):
         return {"required":
-                    {"noise_sampler_type": (["gaussian", "uniform", "brownian", "highres-pyramid", "perlin", "laplacian"], ),
+                    {"noise_sampler_type": (get_noise_sampler_names(default="brownian"), ),
                      "momentum": ("FLOAT", {"default": 0.5, "min": -1.0, "max": 1.0, "step":0.01}),
                      "eta": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step":0.01}),
                      "s_noise": ("FLOAT", {"default": 1, "min": 0.0, "max": 100.0, "step":0.01}),
@@ -165,14 +99,14 @@ class SamplerCLYB_4M_SDE_MOMENTUMIZED:
     FUNCTION = "get_sampler"
 
     def get_sampler(self, noise_sampler_type, eta, s_noise, momentum):
-        sampler = comfy.samplers.ksampler("clyb_4m_sde_momentumized", {"noise_sampler": noise_sampler_type, "eta": eta, "s_noise": s_noise, "momentum": momentum})
+        sampler = comfy.samplers.ksampler("clyb_4m_sde_momentumized", {"noise_sampler_type": noise_sampler_type, "eta": eta, "s_noise": s_noise, "momentum": momentum})
         return (sampler, )
 
 class SamplerEULER_ANCESTRAL_DANCING:
     @classmethod
     def INPUT_TYPES(s):
         return {"required":
-                    {"noise_sampler_type": (["gaussian", "uniform", "brownian", "highres-pyramid", "perlin", "laplacian"], ),
+                    {"noise_sampler_type": (get_noise_sampler_names(), ),
                      "eta": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step":0.01}),
                      "eta_dance": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step":0.01}),
                      "s_noise": ("FLOAT", {"default": 1, "min": 0.0, "max": 100.0, "step":0.01}),
@@ -185,14 +119,14 @@ class SamplerEULER_ANCESTRAL_DANCING:
     FUNCTION = "get_sampler"
 
     def get_sampler(self, noise_sampler_type, eta, s_noise, leap, eta_dance):
-        sampler = comfy.samplers.ksampler("euler_ancestral_dancing", {"noise_sampler": noise_sampler_type, "eta": eta, "s_noise": s_noise, "leap": leap, "eta_dance": eta_dance})
+        sampler = comfy.samplers.ksampler("euler_ancestral_dancing", {"noise_sampler_type": noise_sampler_type, "eta": eta, "s_noise": s_noise, "leap": leap, "eta_dance": eta_dance})
         return (sampler, )
 
 class SamplerDPMPP_3M_SDE_DYN_ETA:
     @classmethod
     def INPUT_TYPES(s):
         return {"required":
-                    {"noise_sampler_type": (["gaussian", "uniform", "brownian", "highres-pyramid", "perlin", "laplacian"], ),
+                    {"noise_sampler_type": (get_noise_sampler_names(default="brownian"), ),
                      "eta_max": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step":0.01}),
                      "eta_min": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 100.0, "step":0.01}),
                      "s_noise": ("FLOAT", {"default": 1, "min": 0.0, "max": 100.0, "step":0.01}),
@@ -204,7 +138,7 @@ class SamplerDPMPP_3M_SDE_DYN_ETA:
     FUNCTION = "get_sampler"
 
     def get_sampler(self, noise_sampler_type, eta_max, eta_min, s_noise):
-        sampler = comfy.samplers.ksampler("dpmpp_3m_sde_dynamic_eta", {"noise_sampler": noise_sampler_type, "eta_max": eta_max, "eta_min": eta_min, "s_noise": s_noise})
+        sampler = comfy.samplers.ksampler("dpmpp_3m_sde_dynamic_eta", {"noise_sampler_type": noise_sampler_type, "eta_max": eta_max, "eta_min": eta_min, "s_noise": s_noise})
         return (sampler, )
 
 ### Schedulers