From 2bf051fda87cfa94e5c99bbd88fc7f1434e9e1a2 Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Thu, 28 Sep 2023 00:30:45 -0400
Subject: [PATCH] Add a basic node to generate sigmas from scheduler.

---
 comfy_extras/nodes_custom_sampler.py | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/comfy_extras/nodes_custom_sampler.py b/comfy_extras/nodes_custom_sampler.py
index 1c587db..aafde8f 100644
--- a/comfy_extras/nodes_custom_sampler.py
+++ b/comfy_extras/nodes_custom_sampler.py
@@ -4,6 +4,26 @@ from comfy.k_diffusion import sampling as k_diffusion_sampling
 import latent_preview
 import torch
 
+
+class BasicScheduler:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required":
+                    {"model": ("MODEL",),
+                     "scheduler": (comfy.samplers.SCHEDULER_NAMES, ),
+                     "steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
+                      }
+               }
+    RETURN_TYPES = ("SIGMAS",)
+    CATEGORY = "_for_testing/custom_sampling"
+
+    FUNCTION = "get_sigmas"
+
+    def get_sigmas(self, model, scheduler, steps):
+        sigmas = comfy.samplers.calculate_sigmas_scheduler(model.model, scheduler, steps).cpu()
+        return (sigmas, )
+
+
 class KarrasScheduler:
     @classmethod
     def INPUT_TYPES(s):
@@ -95,4 +115,5 @@ NODE_CLASS_MAPPINGS = {
     "SamplerCustom": SamplerCustom,
     "KarrasScheduler": KarrasScheduler,
     "KSamplerSelect": KSamplerSelect,
+    "BasicScheduler": BasicScheduler,
 }