ComfyUI Spectrum WAN Proper

Native ComfyUI custom node repo that ports Spectrum from Adaptive Spectral Feature Forecasting for Diffusion Sampling Acceleration to the WAN video backend with backend-specific handlers for:

• Wan 2.1
• Wan 2.2 TI2V 5B
• Wan 2.2 14B High-Noise expert
• Wan 2.2 14B Low-Noise expert

This is intentionally one WAN-focused repository instead of a generic "Spectrum for everything" package.

What this repo does

Spectrum is a training-free diffusion acceleration method that caches and forecasts denoiser features over time instead of running the expensive network at every sampling step.

For the native ComfyUI WAN backend, this repo:

1. patches the internal WanModel.forward_orig(...) path,
2. caches the final hidden transformer feature after the last WAN block,
3. fits an online Chebyshev + ridge regression forecaster,
4. predicts the final hidden feature on skipped steps,
5. runs the normal WAN head + unpatchify on the predicted feature.

That is the WAN analogue of the official Spectrum design: forecast the final block output, not arbitrary module activations or the outer model wrapper output.

Why WAN gets backend-specific handlers

Wan 2.1 and Wan 2.2 are close enough to live in one repo, but not close enough to be treated as one undifferentiated backend.

Wan 2.1 / Wan 2.2 5B

These are handled as single-model WAN paths.

Wan 2.2 14B

Wan 2.2 14B uses separate high-noise and low-noise expert models in the native ComfyUI workflows.

This repo therefore treats them as separate Spectrum targets:

• one patched high-noise expert model
• one patched low-noise expert model

This is deliberate. It avoids pretending the two experts share one stationary feature trajectory when ComfyUI actually loads and samples them as distinct diffusion models.

Installation

Install via Comfyui-Manager

or

Clone or copy this repository into your ComfyUI custom_nodes directory:

git clone https://github.com/xmarre/ComfyUI-Spectrum-WAN-Proper ComfyUI/custom_nodes/ComfyUI-Spectrum-WAN-Proper

Then restart ComfyUI.

No additional dependencies are required beyond normal ComfyUI requirements.

Node

Spectrum Apply WAN

Inputs

• model — ComfyUI MODEL
• backend — backend handler selection
  • auto
  • wan21
  • wan22_ti2v_5b
  • wan22_high_noise
  • wan22_low_noise
• transition_mode — expert transition behavior
  • separate_fit — default behavior; each model/expert fits and forecasts independently
  • bias_shift — experimental Wan 2.2 high-noise → low-noise handoff for expert workflows
• enabled — enable/disable Spectrum patching
• blend_weight — w in the official implementation; blend between local linear extrapolation and Chebyshev prediction
• degree — Chebyshev polynomial degree M
• ridge_lambda — ridge regularization λ
• window_size — initial practical scheduling window
• flex_window — amount added to the window after each actual forward
• warmup_steps — number of initial actual forwards
• history_size — number of actual WAN features to retain for fitting
• tail_actual_steps — number of final known-schedule solver steps forced to stay on the actual WAN path
• forecaster_cache_mode — Chebyshev cache implementation mode
  • low_vram_exact — default; avoids the dense coefficient cache and applies equivalent history weights chunk-by-chunk at prediction time
  • legacy_dense_coeff — opt-in fast path for systems with abundant VRAM; preserves the original dense coefficient cache path
• debug — emit runtime step decisions to stderr

Output

• patched MODEL

Recommended defaults

These defaults keep the reference-style cadence conservative while reducing WAN feature-cache VRAM pressure:

blend_weight = 1.0
degree = 4
ridge_lambda = 0.10
window_size = 2.0
flex_window = 0.75
warmup_steps = 5
history_size = 16
tail_actual_steps = 1
transition_mode = separate_fit
forecaster_cache_mode = low_vram_exact

Why history_size = 16 instead of 100?

The official reference implementation uses a large history cap, but the recommended adaptive WAN settings in the paper consume only 14 or 10 actual network evaluations in the important regimes. A history cap of 16 therefore preserves all actual points in those standard settings while materially reducing memory pressure for WAN video features.

That is an explicit practical approximation in this repo.

Why tail_actual_steps = 1?

tail_actual_steps reserves the final known-schedule solver steps for real WAN forwards instead of forecasted features. This mirrors the FLUX Spectrum node's tail guard and protects the refinement tail, where forecast bias is most likely to appear as softened microdetail or late-step texture drift.

The guard is only applied when ComfyUI exposes a known schedule length through sample_sigmas. When a sampler path does not expose that length, WAN disables forecast decisions and keeps actual forwards only, because hidden feature history cannot be safely bounded across unrelated runs without a real schedule boundary.

Forecaster cache modes

Spectrum Apply WAN exposes two Chebyshev cache modes:

• low_vram_exact avoids materializing the dense (degree + 1, flat_feature_size) coefficient cache on the feature device and is the default.
• legacy_dense_coeff preserves the original implementation behavior and remains available as an explicit fast option when VRAM is abundant.

low_vram_exact keeps the same ridge/Chebyshev forecast formulation and is intended for workloads where the dense coefficient cache causes unnecessary VRAM pressure. It should preserve expected output quality, but it is not guaranteed to be bit-identical to legacy_dense_coeff because the two paths quantize intermediate values at different points.

Usage

Wan 2.1

Place the node directly after the WAN diffusion model loader:

Load Diffusion Model (Wan 2.1)
  -> Spectrum Apply WAN (backend = wan21, transition_mode = separate_fit)
  -> sampler

Wan 2.2 TI2V 5B

Load Diffusion Model (Wan 2.2 TI2V 5B)
  -> Spectrum Apply WAN (backend = wan22_ti2v_5b, transition_mode = separate_fit)
  -> sampler

Wan 2.2 14B T2V / I2V

Apply the node separately to each expert:

Load Diffusion Model (Wan 2.2 high-noise 14B)
  -> Spectrum Apply WAN (backend = wan22_high_noise, transition_mode = separate_fit or bias_shift)

Load Diffusion Model (Wan 2.2 low-noise 14B)
  -> Spectrum Apply WAN (backend = wan22_low_noise, transition_mode = separate_fit or bias_shift)

Use those patched models in place of the original high-noise / low-noise expert models inside the existing ComfyUI Wan 2.2 workflow.

Transition modes

Spectrum Apply WAN supports two expert-transition modes:

• separate_fit keeps the default per-expert reset / re-fit behavior.
• bias_shift is an experimental Wan 2.2 high-noise to low-noise handoff that forces the first low-noise step actual, computes a 1-step bias correction, refreshes that bias on later actual low-noise refresh steps, and uses the transferred high-noise predictor on forecast-eligible low-noise steps.

For Wan 2.2 expert workflows, set transition_mode = bias_shift on both expert nodes if you want the experimental handoff. The high-noise expert publishes the transfer state and the low-noise expert consumes it. If the handoff is missing or incompatible, the low-noise expert falls back to the normal per-expert fit path.

bias_shift is only intended for the Wan 2.2 expert pair (wan22_high_noise and wan22_low_noise). In practice, use separate_fit for Wan 2.1 and Wan 2.2 TI2V 5B.

Spectrum Apply WAN uses current-call runtime injection through the patched outer model, so later bypassed/original-model calls with fresh transformer_options do not inherit Spectrum state from an earlier patched run.

spectrum_wan_run_token must stay consistent across both phases, and spectrum_wan_global_step_override is only needed if the low phase is not contiguous from the published high-phase boundary.

Implementation notes

Native hook point

The patch targets the internal WAN path after the final transformer block loop and before the WAN head / unpatchify.

This is the correct Spectrum target for WAN because the paper’s practical recommendation is to forecast the final block output rather than maintaining per-block caches.

Patch scope

The patch keeps ComfyUI’s existing WAN features intact:

• patches_replace block replacement support
• I2V image context handling
• reference latent path via ref_conv
• normal WAN head and unpatchify behavior

Runtime scoping

Runtime activation is resolved from the current call's transformer_options, with the patched outer apply_model(...) path injecting the active Spectrum runtime for that execution.

That keeps Spectrum tied to the currently patched model call instead of reusing sticky inner-model state across later fresh or bypassed calls.

Schedule tracking

The runtime tracks runs from transformer_options["sample_sigmas"] and transformer_options["sigmas"] when available.

It also separates stream state by cond_or_uncond signature when ComfyUI provides one, so it does not accidentally reuse one runtime state across mismatched conditional/unconditional call patterns.

Fit memory handling

WAN hidden features are very large. To keep the implementation practical, the ridge solve is implemented in feature-dimension chunks, which avoids the worst transient float32 allocations from a naive (K, F) full-matrix solve.

The default low_vram_exact mode avoids caching the final dense coefficient tensor on the feature device.

The optional legacy_dense_coeff mode preserves the original dense-cache fast prediction path for systems with abundant VRAM. low_vram_exact keeps the same ridge/Chebyshev formulation, stores only the small solver state, and applies the equivalent history weights chunk-by-chunk at prediction time. That materially reduces avoidable VRAM overhead while keeping the same forecasting mechanism.

Assumptions, caveats, and limitations

Targeted backend scope

This repo targets the native ComfyUI WAN backend path, i.e. WanModel-style diffusion backbones.

It is not a wrapper for LTX, Hunyuan, Cosmos, or third-party WAN wrappers.

Wan 2.2 expert handling

For Wan 2.2 14B, this repo does not attempt to create one merged forecaster across both experts.

Instead, each expert model gets its own Spectrum runtime and its own cached feature history. That matches the native ComfyUI workflow structure and is the safer design.

Internal-block custom patches on skipped steps

Any custom node that depends on executing internal WAN blocks on every step will naturally not see those block-level effects on skipped steps. Spectrum avoids those full block executions by design.

Memory pressure

Even with final-block-only caching, WAN video features are large. history_size, tail_actual_steps, frame count, resolution, hidden width, and forecaster_cache_mode all affect VRAM pressure and speedup.

low_vram_exact reduces the avoidable VRAM overhead from the dense Chebyshev coefficient cache, but it does not remove the core working-set cost of GPU-resident feature history and chunked prediction buffers.

Practical approximation versus the paper

This repo is faithful to the core method, but makes two explicit practical adaptations:

1. history_size = 16 by default instead of keeping a much larger history cap.
2. tail_actual_steps = 1 by default to keep the final known-schedule refinement step on the actual path.
3. The adaptive schedule follows the practical official implementation behavior used by public Spectrum code (window_size + flex_window) rather than exposing the paper’s triangular-step formula directly as the user-facing control surface.

Future ComfyUI changes

This repo reproduces the current internal WAN forward_orig(...) structure closely. If ComfyUI changes the internal WAN model implementation later, this node may need a maintenance update.

Repo structure

ComfyUI-Spectrum-WAN-Proper/
├── __init__.py
├── nodes.py
├── pyproject.toml
├── LICENSE
├── README.md
├── comfyui_spectrum_wan/
│   ├── __init__.py
│   ├── config.py
│   ├── forecast.py
│   ├── handlers.py
│   ├── runtime.py
│   └── wan.py
└── tests/
    ├── smoke_runtime.py
    ├── test_forecast.py
    ├── test_runtime_debug_logging.py
    ├── test_smoke_bias_shift_transition.py
    └── test_wan_runtime_resolution.py

Tests

Lightweight non-ComfyUI tests are included, including forecast-path parity coverage for the legacy dense coefficient mode and the opt-in exact low-VRAM mode.

Smoke test:

PYTHONPATH=. python tests/smoke_runtime.py

Full local test suite:

PYTHONPATH=. pytest -q

Expected smoke-test output:

ok

References and credits

• Adaptive Spectral Feature Forecasting for Diffusion Sampling Acceleration
• Spectrum project page
• Official Spectrum code

This repository is an unofficial ComfyUI adaptation of the original Spectrum method for native WAN backends. Credit for the underlying method, paper, and reference implementation goes to the original Spectrum authors.