ComfyUI-Spectrum-Proper

Faithful ComfyUI FLUX port of Spectrum from Adaptive Spectral Feature Forecasting for Diffusion Sampling Acceleration.

This repo is intentionally narrow in scope: it implements the FLUX path properly instead of trying to be a half-faithful generic accelerator for every backend.

What this node does

Spectrum Apply Flux patches the native ComfyUI FLUX diffusion model on the MODEL path and applies Spectrum-style forecasting to the final hidden image feature right before final_layer.

That matches the official Spectrum FLUX integration much more closely than forecasting the final denoised output tensor.

In practice the flow is:

1. run a real FLUX forward on selected steps
2. cache the final pre-head hidden feature
3. fit a small Chebyshev ridge regressor online over step index
4. forecast future pre-head features on skipped steps
5. still apply the original FLUX final_layer for the current conditioning

Why this repo exists

The currently circulating ComfyUI Spectrum ports are useful experiments, but they miss important invariants from the paper and the official code.

The main problems I corrected here are:

• Wrong prediction target in one SDXL-style port: it forecasts the whole UNet output instead of the final hidden feature at the model-specific integration point.
• Runtime leakage across model clones in one FLUX port: it closes over a specific runtime object when monkey-patching the shared inner FLUX model.
• Hard-coded 50-step normalization without adapting to the actual detected run length.
• Heuristic pass resets based on timestep direction only, which are brittle in real ComfyUI workflows.
• No clean fallback path for models that share the same patched inner FLUX object but are not actually using Spectrum.

This implementation installs a generic wrapper once on the shared FLUX inner model and looks up the active Spectrum runtime from transformer_options per call. That avoids ghost patching across clones and preserves normal behavior for non-Spectrum models.

Current scope

Supported:

• native ComfyUI FLUX models
• LoRAs on the normal model path
• standard transformer_options patch chains
• standard FLUX control residuals
• ComfyUI runs that sometimes split one logical solver step into multiple internal FLUX calls

For split-step runs, Spectrum now aggregates actual hidden features across the sub-calls and falls back to the real path for any forecast step whose current call shape no longer matches the cached full-batch history. This avoids false run-wide disables while keeping the forecast path conservative.

Not included:

• SDXL
• SD3.5
• video backends

That omission is deliberate. A proper SDXL port in ComfyUI needs either a stable last-block hook in the native UNet path or a maintained fork of the UNet forward. Shipping a brittle pseudo-port would be worse than not shipping one.

Installation

Copy this folder into:

ComfyUI/custom_nodes/ComfyUI-Spectrum-Proper

Restart ComfyUI.

No extra Python dependencies are required beyond what ComfyUI already provides.

Node

Spectrum Apply Flux

Input: MODEL

Output: MODEL

Place it on the FLUX model line:

UNETLoader / CheckpointLoader -> LoRA stack -> Spectrum Apply Flux -> CFGGuider / sampler

Recommended placement:

• after model loading and LoRA application
• before guider/sampler nodes

Parameters

blend_weight

Blend between linear local extrapolation and Chebyshev spectral prediction.

• 1.0 = pure spectral predictor
• 0.0 = pure local linear predictor
• recommended default: 0.5

The official repo notes that a convex blend improves robustness outside the strict paper setting.

degree

Chebyshev degree m.

Recommended default: 4

ridge_lambda

Ridge regularization lambda for the coefficient fit.

Recommended default: 0.1

window_size

Initial interval size before a real forward is required again.

Recommended default: 2.0

flex_window

How much the interval grows after each post-warmup real forward.

This is the ComfyUI-facing equivalent of the adaptive schedule slope used in the official repo.

• 0.75 = paper-style moderate speedup
• 3.0 = more aggressive speedup

warmup_steps

Number of initial real forwards before forecasting is allowed.

Recommended default: 5

tail_actual_steps

Number of final solver steps forced to stay on the real path.

Practical default for quality-sensitive runs: 3

This protects the refinement tail where late-step forecast bias tends to show up first as smoother microdetail, especially in edit-conditioned runs.

max_history

Cap for cached real-forward feature points used for the fit.

This is an implementation guard, not a paper hyperparameter. The default is 32, which is still above the expected number of actual cached points for common 50-step runs while avoiding a hidden large-history footprint.

debug

Enables lightweight logging during patch install and a per-run summary of actual vs. forecasted solver steps.

Recommended settings

Safer / closer to the paper's moderate setting

• blend_weight = 0.50
• degree = 4
• ridge_lambda = 0.10
• window_size = 2.0
• flex_window = 0.75
• warmup_steps = 5
• tail_actual_steps = 3

This is intentionally a little more conservative than the paper-style 14-NFE benchmark configuration because the final three solver steps stay actual-only by default. Set tail_actual_steps = 0 when comparing raw NFE against paper-style settings.

More aggressive

• blend_weight = 0.75
• degree = 4
• ridge_lambda = 0.10
• window_size = 2.0
• flex_window = 3.0
• warmup_steps = 5
• tail_actual_steps = 3

Design notes

1. Forecast target is the final hidden FLUX image feature

This repo caches and forecasts the hidden image tokens after the single-stream blocks and before final_layer.

That is the important architectural choice. Forecasting the final model output directly is less faithful to the official FLUX integration and tends to be less stable.

2. Runtime state is per patched model, not per globally monkey-patched inner model

ComfyUI model clones often share the same underlying diffusion model object. If you close over a runtime object when replacing forward_orig, the state can leak between clones.

This repo avoids that by:

• patching the inner FLUX model only once
• storing the active runtime in each cloned model's transformer_options
• looking up the runtime dynamically on every call
• falling back to the original forward_orig when Spectrum is not active

3. Step normalization uses detected schedule length

The paper and official code mostly benchmark 50-step runs. ComfyUI users do not.

This repo normalizes the Chebyshev basis against the detected schedule length from sample_sigmas instead of hard-coding 50 steps.

Known limitations

• This repo currently targets native ComfyUI FLUX only.
• Forecasting is enabled only for sampler paths that are allowlisted as preserving a one-predict_noise-per-solver-step contract: sample_euler, sample_euler_ancestral, sample_euler_flow, sample_lcm, sample_dpmpp_2m_sde, sample_dpmpp_3m_sde, euler_flow, and Flux2JiTSamplerImpl. Other samplers fall back to native ComfyUI FLUX.
• It depends on current ComfyUI FLUX internals staying broadly compatible with the present forward_orig signature.
• It is designed to coexist with standard transformer patch chains, but it is not guaranteed to compose with other custom nodes that also replace FLUX forward_orig directly.
• The scheduler is faithful to the official adaptive-window strategy, but one safety approximation is added: forecasting is held back until enough real points exist to fit the chosen Chebyshev degree.
• The last few refinement steps can also be reserved as actual-only with tail_actual_steps to reduce late-step texture loss.
• No claims are made here about exact paper speedups inside arbitrary ComfyUI workflows. Sampler choice, guidance path, ControlNet usage, resolution, and other wrappers all affect real wall-clock results.

Validation / smoke test

Outside ComfyUI, you can at least validate the scheduler and forecaster math:

cd ComfyUI/custom_nodes/ComfyUI-Spectrum-Proper
python tests/smoke_runtime.py

Expected output:

ok

Repo structure

ComfyUI-Spectrum-Proper/
|-- __init__.py
|-- nodes.py
|-- pyproject.toml
|-- LICENSE
|-- README.md
|-- comfyui_spectrum/
|   |-- __init__.py
|   |-- config.py
|   |-- forecast.py
|   |-- flux.py
|   `-- runtime.py
`-- tests/
    `-- smoke_runtime.py

Credits

• Spectrum official code and the paper Adaptive Spectral Feature Forecasting for Diffusion Sampling Acceleration by Jiaqi Han et al.
• ComfyUI for the native FLUX integration this custom node builds on.
• FLUX by Black Forest Labs as the model family targeted by this port.

License

GPL-3.0-or-later. This is the safest choice because parts of the FLUX forward-path integration are adapted against ComfyUI core internals.