From 72009e76961d2d955e43586917bb12de763465ed Mon Sep 17 00:00:00 2001
From: yoinked <pepe.dannyboy@gmail.com>
Date: Thu, 19 Sep 2024 17:41:26 -0700
Subject: [PATCH] add laplace scheduler [2407.03297]

---
 comfy/k_diffusion/sampling.py | 11 +++++++++++
 comfy/samplers.py             |  4 +++-
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/comfy/k_diffusion/sampling.py b/comfy/k_diffusion/sampling.py
index ae56aee8b..e9e4edcc6 100644
--- a/comfy/k_diffusion/sampling.py
+++ b/comfy/k_diffusion/sampling.py
@@ -44,6 +44,17 @@ def get_sigmas_vp(n, beta_d=19.9, beta_min=0.1, eps_s=1e-3, device='cpu'):
     return append_zero(sigmas)
 
 
+def get_sigmas_laplace(n, sigma_min, sigma_max, mu=0., beta=0.5, device='cpu'):
+    """Constructs the noise schedule proposed by Tiankai et al. (2024). """
+    epsilon = 1e-5 # avoid log(0)
+    x = torch.linspace(0, 1, n, device=device)
+    clamp = lambda x: torch.clamp(x, min=sigma_min, max=sigma_max)
+    lmb = mu - beta * torch.sign(0.5-x) * torch.log(1 - 2 * torch.abs(0.5-x) + epsilon)
+    sigmas = clamp(torch.exp(lmb))
+    return sigmas
+
+
+
 def to_d(x, sigma, denoised):
     """Converts a denoiser output to a Karras ODE derivative."""
     return (x - denoised) / utils.append_dims(sigma, x.ndim)
diff --git a/comfy/samplers.py b/comfy/samplers.py
index a20f65d4a..ef9dab3de 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"]
+SCHEDULER_NAMES = ["normal", "karras", "exponential", "sgm_uniform", "simple", "ddim_uniform", "beta", "laplace"]
 SAMPLER_NAMES = KSAMPLER_NAMES + ["ddim", "uni_pc", "uni_pc_bh2"]
 
 def calculate_sigmas(model_sampling, scheduler_name, steps):
@@ -747,6 +747,8 @@ 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