Spaces:
Sleeping
Sleeping
#from: https://research.nvidia.com/labs/toronto-ai/AlignYourSteps/howto.html | |
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 = {"SD1": [14.6146412293, 6.4745760956, 3.8636745985, 2.6946151520, 1.8841921177, 1.3943805092, 0.9642583904, 0.6523686016, 0.3977456272, 0.1515232662, 0.0291671582], | |
"SDXL":[14.6146412293, 6.3184485287, 3.7681790315, 2.1811480769, 1.3405244945, 0.8620721141, 0.5550693289, 0.3798540708, 0.2332364134, 0.1114188177, 0.0291671582], | |
"SVD": [700.00, 54.5, 15.886, 7.977, 4.248, 1.789, 0.981, 0.403, 0.173, 0.034, 0.002]} | |
class AlignYourStepsScheduler: | |
def INPUT_TYPES(s): | |
return {"required": | |
{"model_type": (["SD1", "SDXL", "SVD"], ), | |
"steps": ("INT", {"default": 10, "min": 10, "max": 10000}), | |
"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 = { | |
"AlignYourStepsScheduler": AlignYourStepsScheduler, | |
} | |