mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-05-30 19:07:25 +08:00
174208df6b
Some checks are pending
Python Linting / Run Ruff (push) Waiting to run
Python Linting / Run Pylint (push) Waiting to run
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.10, [self-hosted Linux], stable) (push) Waiting to run
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.11, [self-hosted Linux], stable) (push) Waiting to run
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.12, [self-hosted Linux], stable) (push) Waiting to run
Full Comfy CI Workflow Runs / test-unix-nightly (12.1, , linux, 3.11, [self-hosted Linux], nightly) (push) Waiting to run
Execution Tests / test (macos-latest) (push) Waiting to run
Execution Tests / test (ubuntu-latest) (push) Waiting to run
Execution Tests / test (windows-latest) (push) Waiting to run
Test server launches without errors / test (push) Waiting to run
Unit Tests / test (macos-latest) (push) Waiting to run
Unit Tests / test (ubuntu-latest) (push) Waiting to run
Unit Tests / test (windows-2022) (push) Waiting to run
* Move dataset/text nodes to text category * Rename category utils into utilities * Rename category api node into partner * Move categories conditioning, latent, sampling, model_patches, training, etc. under model category * Dispatch partner nodes in to 3d, audio, image, text, video categories * Move PreviewAny node to utilities category
72 lines
2.4 KiB
Python
72 lines
2.4 KiB
Python
# from https://github.com/bebebe666/OptimalSteps
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
from typing_extensions import override
|
|
from comfy_api.latest import ComfyExtension, 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)
|
|
|
|
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],
|
|
"Chroma": [0.992, 0.99, 0.988, 0.985, 0.982, 0.978, 0.973, 0.968, 0.961, 0.953, 0.943, 0.931, 0.917, 0.9, 0.881, 0.858, 0.832, 0.802, 0.769, 0.731, 0.69, 0.646, 0.599, 0.55, 0.501, 0.451, 0.402, 0.355, 0.311, 0.27, 0.232, 0.199, 0.169, 0.143, 0.12, 0.101, 0.084, 0.07, 0.058, 0.048, 0.001],
|
|
}
|
|
|
|
class OptimalStepsScheduler(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="OptimalStepsScheduler",
|
|
category="model/sampling/schedulers",
|
|
inputs=[
|
|
io.Combo.Input("model_type", options=["FLUX", "Wan", "Chroma"]),
|
|
io.Int.Input("steps", default=20, min=3, max=1000),
|
|
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))
|
|
|
|
|
|
class OptimalStepsExtension(ComfyExtension):
|
|
@override
|
|
async def get_node_list(self) -> list[type[io.ComfyNode]]:
|
|
return [
|
|
OptimalStepsScheduler,
|
|
]
|
|
|
|
|
|
async def comfy_entrypoint() -> OptimalStepsExtension:
|
|
return OptimalStepsExtension()
|