From be8e07720f6a3cdc7d177fd483ab8d9e2bdb2dd3 Mon Sep 17 00:00:00 2001
From: BlakeOne <30273164+BlakeOne@users.noreply.github.com>
Date: Sun, 31 Mar 2024 17:58:08 -0700
Subject: [PATCH] Create node.py

---
 node.py | 44 ++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 44 insertions(+)
 create mode 100644 node.py

diff --git a/node.py b/node.py
new file mode 100644
index 0000000..79d75af
--- /dev/null
+++ b/node.py
@@ -0,0 +1,44 @@
+import torch
+import comfy.model_management
+import comfy.samplers
+
+class SchedulerMixer:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "model": ("MODEL",),
+                "steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
+                "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
+                "normal": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}),
+                "karras": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}),
+                "exponential": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}),
+                "sgm_uniform": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}),
+                "simple": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}),
+                "ddim_uniform": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}),
+            }
+        }
+    
+    RETURN_TYPES = ("SIGMAS",)
+    CATEGORY = "sampling/custom_sampling/schedulers"
+    FUNCTION = "get_sigmas"
+
+    def get_sigmas(self, model, steps, denoise, normal, karras, exponential, sgm_uniform, simple, ddim_uniform):
+        total_steps = steps
+        if denoise < 1.0:
+            total_steps = int(steps / denoise)
+
+        comfy.model_management.load_models_gpu([model])
+
+        scheduler_weights = [normal, karras, exponential, sgm_uniform, simple, ddim_uniform]
+        scheduler_names = ["normal", "karras", "exponential", "sgm_uniform", "simple", "ddim_uniform"]
+
+        mixed_sigmas = torch.zeros((steps + 1,), device="cpu", dtype=torch.float)
+        for weight, name in zip(scheduler_weights, scheduler_names):                        
+            if weight > 0.0:                
+                sigmas = comfy.samplers.calculate_sigmas_scheduler(model.model, name, total_steps).cpu()
+                sigmas = sigmas[-(steps + 1):]                
+                mixed_sigmas += sigmas * weight
+
+        return (mixed_sigmas,)
+