diff --git a/comfy_extras/nodes_optimalsteps.py b/comfy_extras/nodes_optimalsteps.py index f6928199b06..9ce6e8af32c 100644 --- a/comfy_extras/nodes_optimalsteps.py +++ b/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,32 +18,57 @@ 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) @@ -51,6 +76,7 @@ def get_sigmas(self, model_type, steps, denoise): sigmas[-1] = 0 return (torch.FloatTensor(sigmas), ) + NODE_CLASS_MAPPINGS = { "OptimalStepsScheduler": OptimalStepsScheduler, }