From ea613ed2e328f290b94f2bc0c6ef0e3512341968 Mon Sep 17 00:00:00 2001
From: catboxanon <122327233+catboxanon@users.noreply.github.com>
Date: Thu, 1 May 2025 18:14:27 -0400
Subject: [PATCH] Rename, add interpolation functions

Co-authored-by: liesen <liesen.dev@gmail.com>
---
 comfy_extras/nodes_custom_sampler.py | 45 +++++++++++++++++-----------
 1 file changed, 27 insertions(+), 18 deletions(-)

diff --git a/comfy_extras/nodes_custom_sampler.py b/comfy_extras/nodes_custom_sampler.py
index 6c65ba4b6..47f8b8eac 100644
--- a/comfy_extras/nodes_custom_sampler.py
+++ b/comfy_extras/nodes_custom_sampler.py
@@ -1,3 +1,4 @@
+import math
 import comfy.samplers
 import comfy.sample
 from comfy.k_diffusion import sampling as k_diffusion_sampling
@@ -249,45 +250,53 @@ class SetFirstSigma:
         sigmas[0] = sigma
         return (sigmas, )
 
-class ExpandSigmas:
+class ExtendIntermediateSigmas:
     @classmethod
     def INPUT_TYPES(s):
         return {"required":
                     {"sigmas": ("SIGMAS", ),
                      "steps": ("INT", {"default": 2, "min": 1, "max": 100}),
-                     "sigma_min": ("FLOAT", {"default": 12.0, "min":  0.0, "max": 20000.0, "step": 0.01, "round": False}),
-                     "sigma_max": ("FLOAT", {"default": -1.0, "min": -1.0, "max": 20000.0, "step": 0.01, "round": False}),
+                     "start_at_sigma": ("FLOAT", {"default": -1.0, "min": -1.0, "max": 20000.0, "step": 0.01, "round": False}),
+                     "end_at_sigma": ("FLOAT", {"default": 12.0, "min":  0.0, "max": 20000.0, "step": 0.01, "round": False}),
+                     "spacing": (['linear', 'cosine', 'sine'],),
                     }
                }
     RETURN_TYPES = ("SIGMAS",)
     CATEGORY = "sampling/custom_sampling/sigmas"
 
-    FUNCTION = "expand"
+    FUNCTION = "extend"
 
-    def expand(self, sigmas, steps, sigma_max, sigma_min):
-        if sigma_max < 0:
-            sigma_max = float("inf")
+    def extend(self, sigmas, steps, start_at_sigma, end_at_sigma, spacing):
+        if start_at_sigma < 0:
+            start_at_sigma = float("inf")
 
-        expanded_sigmas = []
+        interpolator = {
+            'linear': lambda x: x,
+            'cosine': lambda x: torch.sin(x*math.pi/2),
+            'sine':   lambda x: 1 - torch.cos(x*math.pi/2)
+        }[spacing]
+
+        # linear space for our interpolation function
+        x = torch.linspace(0, 1, steps + 1, device=sigmas.device)[1:-1]
+
+        extended_sigmas = []
         for i in range(len(sigmas) - 1):
             sigma_current = sigmas[i]
             sigma_next = sigmas[i+1]
 
-            expanded_sigmas.append(sigma_current)
+            extended_sigmas.append(sigma_current)
 
-            if sigma_min <= sigma_current <= sigma_max:
-                #XXX: Might be nice to interpolate the sigmas with different methods.
-                #     Would require writing this not as a loop.
-                interpolated_steps = torch.linspace(sigma_current, sigma_next, steps + 1, device=sigmas.device)[1:-1]
-                expanded_sigmas.extend(interpolated_steps.tolist())
+            if end_at_sigma <= sigma_current <= start_at_sigma:
+                interpolated_steps = interpolator(x) * (sigma_next - sigma_current) + sigma_current
+                extended_sigmas.extend(interpolated_steps.tolist())
 
         # Add the last sigma value
         if len(sigmas) > 0:
-            expanded_sigmas.append(sigmas[-1])
+            extended_sigmas.append(sigmas[-1])
 
-        expanded_sigmas = torch.FloatTensor(expanded_sigmas)
+        extended_sigmas = torch.FloatTensor(extended_sigmas)
 
-        return (expanded_sigmas,)
+        return (extended_sigmas,)
 
 class KSamplerSelect:
     @classmethod
@@ -775,7 +784,7 @@ NODE_CLASS_MAPPINGS = {
     "SplitSigmasDenoise": SplitSigmasDenoise,
     "FlipSigmas": FlipSigmas,
     "SetFirstSigma": SetFirstSigma,
-    "ExpandSigmas": ExpandSigmas,
+    "ExtendIntermediateSigmas": ExtendIntermediateSigmas,
 
     "CFGGuider": CFGGuider,
     "DualCFGGuider": DualCFGGuider,