From 575b027bbefaeab446a02923088decb75c895e9b Mon Sep 17 00:00:00 2001
From: silveroxides <ishimarukaito@gmail.com>
Date: Wed, 9 Jul 2025 17:06:58 +0200
Subject: [PATCH] add __init__.py with the node

---
 __init__.py | 62 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 62 insertions(+)
 create mode 100644 __init__.py

diff --git a/__init__.py b/__init__.py
new file mode 100644
index 0000000..7db885a
--- /dev/null
+++ b/__init__.py
@@ -0,0 +1,62 @@
+import torch
+import numpy
+from comfy.samplers import SchedulerHandler, SCHEDULER_HANDLERS, SCHEDULER_NAMES
+
+def power_shift_scheduler(model_sampling, steps, power=2.0, midpoint_shift=1.0):
+    total_timesteps = (len(model_sampling.sigmas) - 1)
+    x = numpy.linspace(0, 1, steps, endpoint=False)
+    x = x**midpoint_shift
+
+    ts_normalized = (1 - x**power)**power
+    ts = numpy.rint(ts_normalized * total_timesteps)
+
+    sigs = []
+    last_t = -1
+    for t in ts:
+        t_int = min(int(t), total_timesteps)
+        if t_int != last_t:
+            sigs.append(float(model_sampling.sigmas[t_int]))
+        last_t = t_int
+        
+    sigs.append(0.0)
+    return torch.FloatTensor(sigs)
+
+class PowerShiftSchedulerNode:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required":
+                    {"model": ("MODEL",),
+                     "steps": ("INT", {"default": 20, "min": 3, "max": 1000}),
+                     "power": ("FLOAT", {"default": 2.0, "min": 0.0, "max": 5.0, "step": 0.001}),
+                     "midpoint_shift": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.001}),
+                     "denoise": ("FLOAT", {"default": 1.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, power, midpoint_shift, denoise):
+        total_steps = steps
+        if denoise < 1.0:
+            total_steps = int(steps/denoise)
+
+        sigmas = power_shift_scheduler(model.get_model_object("model_sampling"), total_steps, power, midpoint_shift).cpu()
+        sigmas = sigmas[-(steps + 1):]
+
+        return (sigmas, )
+
+scheduler_name = "power_shift"
+if scheduler_name not in SCHEDULER_HANDLERS:
+    scheduler_handler = SchedulerHandler(handler=power_shift_scheduler, use_ms=True)
+    SCHEDULER_HANDLERS[scheduler_name] = scheduler_handler
+    if scheduler_name not in SCHEDULER_NAMES:
+        SCHEDULER_NAMES.append(scheduler_name)
+
+NODE_CLASS_MAPPINGS = {
+    "PowerShiftScheduler": PowerShiftSchedulerNode,
+}
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "PowerShiftScheduler": "Power Shift Scheduler",
+}
\ No newline at end of file