From 2cdbaf5169a631b126542f41432f2484c4f0a608 Mon Sep 17 00:00:00 2001
From: catboxanon <122327233+catboxanon@users.noreply.github.com>
Date: Tue, 14 Jan 2025 19:05:45 -0500
Subject: [PATCH] Add SetFirstSigma node (#6459)

Useful for models utilizing ztSNR. See: https://arxiv.org/abs/2409.15997
---
 comfy_extras/nodes_custom_sampler.py | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

diff --git a/comfy_extras/nodes_custom_sampler.py b/comfy_extras/nodes_custom_sampler.py
index c7ff9a4d8..576fc3b2c 100644
--- a/comfy_extras/nodes_custom_sampler.py
+++ b/comfy_extras/nodes_custom_sampler.py
@@ -231,6 +231,24 @@ class FlipSigmas:
             sigmas[0] = 0.0001
         return (sigmas,)
 
+class SetFirstSigma:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required":
+                    {"sigmas": ("SIGMAS", ),
+                     "sigma": ("FLOAT", {"default": 136.0, "min": 0.0, "max": 20000.0, "step": 0.001, "round": False}),
+                    }
+               }
+    RETURN_TYPES = ("SIGMAS",)
+    CATEGORY = "sampling/custom_sampling/sigmas"
+
+    FUNCTION = "set_first_sigma"
+
+    def set_first_sigma(self, sigmas, sigma):
+        sigmas = sigmas.clone()
+        sigmas[0] = sigma
+        return (sigmas, )
+
 class KSamplerSelect:
     @classmethod
     def INPUT_TYPES(s):
@@ -710,6 +728,7 @@ NODE_CLASS_MAPPINGS = {
     "SplitSigmas": SplitSigmas,
     "SplitSigmasDenoise": SplitSigmasDenoise,
     "FlipSigmas": FlipSigmas,
+    "SetFirstSigma": SetFirstSigma,
 
     "CFGGuider": CFGGuider,
     "DualCFGGuider": DualCFGGuider,