mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2025-07-27 08:16:44 +00:00
85 lines
2.2 KiB
Python
85 lines
2.2 KiB
Python
# from: https://research.nvidia.com/labs/toronto-ai/AlignYourSteps/howto.html
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
from comfy_api.latest import io
|
|
|
|
|
|
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)
|
|
|
|
return np.exp(new_ys)[::-1].copy()
|
|
|
|
|
|
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(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id="AlignYourStepsScheduler_V3",
|
|
category="sampling/custom_sampling/schedulers",
|
|
inputs=[
|
|
io.Combo.Input("model_type", options=["SD1", "SDXL", "SVD"]),
|
|
io.Int.Input("steps", default=10, min=1, max=10000),
|
|
io.Float.Input("denoise", default=1.0, min=0.0, max=1.0, step=0.01),
|
|
],
|
|
outputs=[io.Sigmas.Output()],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, model_type, steps, denoise) -> io.NodeOutput:
|
|
total_steps = steps
|
|
if denoise < 1.0:
|
|
if denoise <= 0.0:
|
|
return io.NodeOutput(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 io.NodeOutput(torch.FloatTensor(sigmas))
|
|
|
|
|
|
NODES_LIST: list[type[io.ComfyNode]] = [
|
|
AlignYourStepsScheduler,
|
|
]
|