mirror of
https://git.datalinker.icu/comfyanonymous/ComfyUI
synced 2026-09-06 18:17:07 +08:00
83 lines
2.6 KiB
Python
83 lines
2.6 KiB
Python
# 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], "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", "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, custom_sigmas=""):
|
|
total_steps = steps
|
|
if denoise < 1.0:
|
|
if denoise <= 0.0:
|
|
return (torch.FloatTensor([]), )
|
|
total_steps = round(steps * denoise)
|
|
|
|
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,
|
|
}
|