OptimalSteps scheduler improvements

comfy_extras/nodes_optimalsteps.py

@@ -1,9 +1,9 @@
 # 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.
@@ -18,39 +18,65 @@ def loglinear_interp(t_steps, num_steps):
     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],
-}
+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], "SDXL": [12.1769, 9.9182, 7.0887, 4.5944, 2.2473, 0.9020, 0.2872, 0.0738, 0.0197, 0.0020, 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",)
+        return {
+            "required": {
+                "model_type": (["FLUX", "Wan", "SDXL"], ),
+                "steps": ("INT", {
+                    "default": 20,
+                    "min": 1,
+                    "max": 1000
+                }),
+                "denoise": ("FLOAT", {
+                    "default": 1.0,
+                    "min": 0.0,
+                    "max": 1.0,
+                    "step": 0.01
+                }),
+            },
+            "optional": {
+                "custom_sigmas": ("STRING", {
+                    "default": "",
+                    "placeholder": "Comma-separated sigma values"
+                }),
+            }
+        }
+
+    RETURN_TYPES = ("SIGMAS", )
     CATEGORY = "sampling/custom_sampling/schedulers"
 
     FUNCTION = "get_sigmas"
 
-    def get_sigmas(self, model_type, steps, denoise):
+    def get_sigmas(self, model_type, steps, denoise, custom_sigmas=""):
         total_steps = steps
         if denoise < 1.0:
             if denoise <= 0.0:
-                return (torch.FloatTensor([]),)
+                return (torch.FloatTensor([]), )
             total_steps = round(steps * denoise)
 
-        sigmas = NOISE_LEVELS[model_type][:]
+        if custom_sigmas:
+            # Parse the custom_sigmas string into a list of floats
+            try:
+                sigmas = [float(s.strip()) for s in custom_sigmas.split(",") if s.strip()]
+            except ValueError:
+                raise ValueError("Invalid custom_sigmas format. Ensure it is a comma-separated list of numbers.")
+        else:
+            # Use the predefined NOISE_LEVELS
+            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,
 }