From ed42dfe45c3bf39b4482b0501836c9a34b02aebc Mon Sep 17 00:00:00 2001 From: yoinked Date: Thu, 19 Sep 2024 17:50:17 -0700 Subject: [PATCH] should be here instead lol --- comfy/samplers.py | 4 +--- comfy_extras/nodes_custom_sampler.py | 22 ++++++++++++++++++++++ 2 files changed, 23 insertions(+), 3 deletions(-) diff --git a/comfy/samplers.py b/comfy/samplers.py index ef9dab3de..a20f65d4a 100644 --- a/comfy/samplers.py +++ b/comfy/samplers.py @@ -729,7 +729,7 @@ def sample(model, noise, positive, negative, cfg, device, sampler, sigmas, model return cfg_guider.sample(noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed) -SCHEDULER_NAMES = ["normal", "karras", "exponential", "sgm_uniform", "simple", "ddim_uniform", "beta", "laplace"] +SCHEDULER_NAMES = ["normal", "karras", "exponential", "sgm_uniform", "simple", "ddim_uniform", "beta"] SAMPLER_NAMES = KSAMPLER_NAMES + ["ddim", "uni_pc", "uni_pc_bh2"] def calculate_sigmas(model_sampling, scheduler_name, steps): @@ -747,8 +747,6 @@ def calculate_sigmas(model_sampling, scheduler_name, steps): sigmas = normal_scheduler(model_sampling, steps, sgm=True) elif scheduler_name == "beta": sigmas = beta_scheduler(model_sampling, steps) - elif scheduler_name == "laplace": - sigmas = k_diffusion_sampling.get_sigmas_laplace(n=steps, sigma_min=float(model_sampling.sigma_min), sigma_max=float(model_sampling.sigma_max)) else: logging.error("error invalid scheduler {}".format(scheduler_name)) return sigmas diff --git a/comfy_extras/nodes_custom_sampler.py b/comfy_extras/nodes_custom_sampler.py index 219975e27..773361760 100644 --- a/comfy_extras/nodes_custom_sampler.py +++ b/comfy_extras/nodes_custom_sampler.py @@ -90,6 +90,27 @@ class PolyexponentialScheduler: sigmas = k_diffusion_sampling.get_sigmas_polyexponential(n=steps, sigma_min=sigma_min, sigma_max=sigma_max, rho=rho) return (sigmas, ) +class LaplaceScheduler: + @classmethod + def INPUT_TYPES(s): + return {"required": + {"steps": ("INT", {"default": 20, "min": 1, "max": 10000}), + "sigma_max": ("FLOAT", {"default": 14.614642, "min": 0.0, "max": 5000.0, "step":0.01, "round": False}), + "sigma_min": ("FLOAT", {"default": 0.0291675, "min": 0.0, "max": 5000.0, "step":0.01, "round": False}), + "mu": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step":0.01, "round": False}), + "beta": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step":0.01, "round": False}), + } + } + RETURN_TYPES = ("SIGMAS",) + CATEGORY = "sampling/custom_sampling/schedulers" + + FUNCTION = "get_sigmas" + + def get_sigmas(self, steps, sigma_max, sigma_min, mu, beta): + sigmas = k_diffusion_sampling.get_sigmas_laplace(n=steps, sigma_min=sigma_min, sigma_max=sigma_max, mu=mu, beta=beta) + return (sigmas, ) + + class SDTurboScheduler: @classmethod def INPUT_TYPES(s): @@ -673,6 +694,7 @@ NODE_CLASS_MAPPINGS = { "KarrasScheduler": KarrasScheduler, "ExponentialScheduler": ExponentialScheduler, "PolyexponentialScheduler": PolyexponentialScheduler, + "LaplaceScheduler": LaplaceScheduler, "VPScheduler": VPScheduler, "BetaSamplingScheduler": BetaSamplingScheduler, "SDTurboScheduler": SDTurboScheduler,