From 64e8892869605f313f790e682d2de81c60ea98dc Mon Sep 17 00:00:00 2001
From: Chris Tefer <ctefer@gmail.com>
Date: Thu, 5 Dec 2024 19:18:29 -0600
Subject: [PATCH] Added a helper node for both sigmas and schedulers

---
 __init__.py  |  2 ++
 cfe_nodes.py | 41 ++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/__init__.py b/__init__.py
index 37cfca5..2efd0ac 100644
--- a/__init__.py
+++ b/__init__.py
@@ -9,6 +9,7 @@ NODE_CLASS_MAPPINGS = {
     "CFE Flux In Pipe" : CFE_Flux_In_Pipe,
     "CFE Flux Out Pipe" : CFE_Flux_Out_Pipe,
     "CFE FLUX Sampler" : CFE_FLUX_Sampler,
+    "CFE FLUX Sigma Sampler" : CFE_Sigma_Sampler,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
@@ -16,4 +17,5 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "CFE_Flux_In_Pipe":"CFE Flux In Pipe",
     "CFE_Flux_Out_Pipe":"CFE Flux Out Pipe",
     "CFE_FLUX_Sampler":"CFE Flux Sampler",
+    "CFE_Sigma_Sampler":"CFE Flux Sigma Sampler",
 }
diff --git a/cfe_nodes.py b/cfe_nodes.py
index 15eda90..ca9b883 100644
--- a/cfe_nodes.py
+++ b/cfe_nodes.py
@@ -2,6 +2,7 @@ import math
 import torch
 import comfy
 import comfy.model_management
+import comfy.samplers
 import node_helpers
 
 
@@ -156,11 +157,49 @@ class CFE_Flux_Out_Pipe:
     def out_pipe(self, pipe):
         model, clip, vae, cond, sampler, sigmas = pipe
         return (pipe, model, clip, vae, cond, sampler, sigmas)
+    
 
 
+class CFE_Sigma_Sampler:
+    def __init__(self) -> None:
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required" : {
+                "model" : ("MODEL", {"tooltip" : "The model used for denoising the input latent"}),
+                "sampler_select" : (comfy.samplers.SAMPLER_NAMES, {"tooltip" : "The name of the sampler being used"}),
+                "scheduler": (comfy.samplers.SCHEDULER_NAMES, {"tooltip" : "The name of the scheduler being used"}),
+                "steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
+                "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
+            }
+        }
+
+
+    RETURN_TYPES = ("SAMPLER", "SIGMAS")
+    RETURN_NAMES = ("sampler", "sigmas")
+    FUNCTION = "execute"
+    CATEGORY = "CFE"
+
+
+    def execute(self, model, sampler_select, scheduler, steps, denoise):
+        total_steps = steps
+        sigmas = None
+        if denoise < 1.0:
+            if denoise <= 0.0:
+                sigmas = torch.FloatTensor([])
+            total_steps = int(steps/denoise)
+
+        if sigmas is None:
+            sigmas = comfy.samplers.calculate_sigmas(model.get_model_object("model_sampling"), scheduler, total_steps).cpu()
+            sigmas = sigmas[-(steps + 1):]
+
+        return (comfy.samplers.sampler_object(sampler_select), sigmas, )
+
 
 class CFE_FLUX_Sampler:
-    def __init__(self):
+    def __init__(self) -> None:
         pass
 
     @classmethod