From 9ee6ca99d88d77678e7306dab2f1f0f092d8ed43 Mon Sep 17 00:00:00 2001
From: JNP <50867151+bebebe666@users.noreply.github.com>
Date: Sun, 13 Apr 2025 08:33:36 +0800
Subject: [PATCH] add_optimalsteps (#7584)

Co-authored-by: bebebe666 <jianningpei@tencent.com>
---
 comfy_extras/nodes_optimalsteps.py | 56 ++++++++++++++++++++++++++++++
 nodes.py                           |  1 +
 2 files changed, 57 insertions(+)
 create mode 100644 comfy_extras/nodes_optimalsteps.py

diff --git a/comfy_extras/nodes_optimalsteps.py b/comfy_extras/nodes_optimalsteps.py
new file mode 100644
index 00000000..f6928199
--- /dev/null
+++ b/comfy_extras/nodes_optimalsteps.py
@@ -0,0 +1,56 @@
+# from https://github.com/bebebe666/OptimalSteps
+
+
+import numpy as np
+import torch
+
+def loglinear_interp(t_steps, num_steps):
+    """
+    Performs log-linear interpolation of a given array of decreasing numbers.
+    """
+    xs = np.linspace(0, 1, len(t_steps))
+    ys = np.log(t_steps[::-1])
+
+    new_xs = np.linspace(0, 1, num_steps)
+    new_ys = np.interp(new_xs, xs, ys)
+
+    interped_ys = np.exp(new_ys)[::-1].copy()
+    return interped_ys
+
+
+NOISE_LEVELS = {"FLUX": [0.9968, 0.9886, 0.9819, 0.975, 0.966, 0.9471, 0.9158, 0.8287, 0.5512, 0.2808, 0.001],
+"Wan":[1.0, 0.997, 0.995, 0.993, 0.991, 0.989, 0.987, 0.985, 0.98, 0.975, 0.973, 0.968, 0.96, 0.946, 0.927, 0.902, 0.864, 0.776, 0.539, 0.208, 0.001],
+}
+
+class OptimalStepsScheduler:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required":
+                    {"model_type": (["FLUX", "Wan"], ),
+                     "steps": ("INT", {"default": 20, "min": 3, "max": 1000}),
+                     "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_type, steps, denoise):
+        total_steps = steps
+        if denoise < 1.0:
+            if denoise <= 0.0:
+                return (torch.FloatTensor([]),)
+            total_steps = round(steps * denoise)
+
+        sigmas = NOISE_LEVELS[model_type][:]
+        if (steps + 1) != len(sigmas):
+            sigmas = loglinear_interp(sigmas, steps + 1)
+
+        sigmas = sigmas[-(total_steps + 1):]
+        sigmas[-1] = 0
+        return (torch.FloatTensor(sigmas), )
+
+NODE_CLASS_MAPPINGS = {
+    "OptimalStepsScheduler": OptimalStepsScheduler,
+}
diff --git a/nodes.py b/nodes.py
index e2893e83..e66b5c71 100644
--- a/nodes.py
+++ b/nodes.py
@@ -2280,6 +2280,7 @@ def init_builtin_extra_nodes():
         "nodes_hunyuan3d.py",
         "nodes_primitive.py",
         "nodes_cfg.py",
+        "nodes_optimalsteps.py"
     ]
 
     import_failed = []