97 lines
3.0 KiB
Python
97 lines
3.0 KiB
Python
from __future__ import annotations
|
|
|
|
from .comfyui_spectrum_wan.config import SpectrumWanConfig
|
|
from .comfyui_spectrum_wan.wan import WanSpectrumPatcher
|
|
|
|
|
|
class SpectrumApplyWAN:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"backend": (
|
|
[
|
|
"auto",
|
|
"wan21",
|
|
"wan22_ti2v_5b",
|
|
"wan22_high_noise",
|
|
"wan22_low_noise",
|
|
],
|
|
{"default": "auto"},
|
|
),
|
|
"transition_mode": (
|
|
[
|
|
"separate_fit",
|
|
"bias_shift",
|
|
],
|
|
{"default": "separate_fit"},
|
|
),
|
|
"enabled": ("BOOLEAN", {"default": True}),
|
|
"blend_weight": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"degree": ("INT", {"default": 4, "min": 1, "max": 16, "step": 1}),
|
|
"ridge_lambda": ("FLOAT", {"default": 0.10, "min": 0.0, "max": 10.0, "step": 0.01}),
|
|
"window_size": ("FLOAT", {"default": 2.0, "min": 1.0, "max": 32.0, "step": 0.05}),
|
|
"flex_window": ("FLOAT", {"default": 0.75, "min": 0.0, "max": 16.0, "step": 0.05}),
|
|
"warmup_steps": ("INT", {"default": 5, "min": 0, "max": 64, "step": 1}),
|
|
"history_size": ("INT", {"default": 16, "min": 2, "max": 128, "step": 1}),
|
|
"debug": ("BOOLEAN", {"default": False}),
|
|
"forecaster_cache_mode": (
|
|
[
|
|
"legacy_dense_coeff",
|
|
"low_vram_exact",
|
|
],
|
|
{"default": "legacy_dense_coeff"},
|
|
),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
RETURN_NAMES = ("model",)
|
|
FUNCTION = "apply"
|
|
CATEGORY = "sampling/spectrum"
|
|
|
|
def apply(
|
|
self,
|
|
model,
|
|
backend,
|
|
transition_mode,
|
|
enabled,
|
|
blend_weight,
|
|
degree,
|
|
ridge_lambda,
|
|
window_size,
|
|
flex_window,
|
|
warmup_steps,
|
|
history_size,
|
|
debug,
|
|
forecaster_cache_mode,
|
|
):
|
|
if not enabled:
|
|
return (model,)
|
|
|
|
cfg = SpectrumWanConfig(
|
|
backend=backend,
|
|
transition_mode=transition_mode,
|
|
enabled=enabled,
|
|
blend_weight=blend_weight,
|
|
degree=degree,
|
|
ridge_lambda=ridge_lambda,
|
|
window_size=window_size,
|
|
flex_window=flex_window,
|
|
warmup_steps=warmup_steps,
|
|
history_size=history_size,
|
|
debug=debug,
|
|
forecaster_cache_mode=forecaster_cache_mode,
|
|
).validated()
|
|
return (WanSpectrumPatcher.patch(model, cfg),)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"SpectrumApplyWAN": SpectrumApplyWAN,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"SpectrumApplyWAN": "Spectrum Apply WAN",
|
|
}
|