Initial import of built repo

.gitignore

@@ -0,0 +1,7 @@
+__pycache__/
+*.pyc
+*.pyo
+*.pyd
+*.so
+*.egg-info/
+.pytest_cache/

LICENSE

@@ -0,0 +1,73 @@
+GNU GENERAL PUBLIC LICENSE
+Version 3, 29 June 2007
+
+Copyright (C) 2007 Free Software Foundation, Inc. <https://fsf.org/>
+Everyone is permitted to copy and distribute verbatim copies
+of this license document, but changing it is not allowed.
+
+Preamble
+
+The GNU General Public License is a free, copyleft license for
+software and other kinds of works.
+
+The licenses for most software and other practical works are designed
+to take away your freedom to share and change the works. By contrast,
+the GNU General Public License is intended to guarantee your freedom to
+share and change all versions of a program--to make sure it remains free
+software for all its users. We, the Free Software Foundation, use the
+GNU General Public License for most of our software; it applies also to
+any other work released this way by its authors. You can apply it to
+your programs, too.
+
+When we speak of free software, we are referring to freedom, not
+price. Our General Public Licenses are designed to make sure that you
+have the freedom to distribute copies of free software (and charge for
+them if you wish), that you receive source code or can get it if you
+want it, that you can change the software or use pieces of it in new
+free programs, and that you know you can do these things.
+
+To protect your rights, we need to prevent others from denying you
+these rights or asking you to surrender the rights. Therefore, you have
+certain responsibilities if you distribute copies of the software, or if
+you modify it: responsibilities to respect the freedom of others.
+
+For example, if you distribute copies of such a program, whether
+gratis or for a fee, you must pass on to the recipients the same
+freedoms that you received. You must make sure that they, too, receive
+or can get the source code. And you must show them these terms so they
+know their rights.
+
+Developers that use the GNU GPL protect your rights with two steps:
+(1) assert copyright on the software, and (2) offer you this License
+giving you legal permission to copy, distribute and/or modify it.
+
+For the developers' and authors' protection, the GPL clearly explains
+that there is no warranty for this free software. For both users' and
+authors' sake, the GPL requires that modified versions be marked as
+changed, so that their problems will not be attributed erroneously to
+authors of previous versions.
+
+Some devices are designed to deny users access to install or run
+modified versions of the software inside them, although the manufacturer
+can do so. This is fundamentally incompatible with the aim of
+protecting users' freedom to change the software. The systematic
+pattern of such abuse occurs in the area of products for individuals to
+use, which is precisely where it is most unacceptable. Therefore, we
+have designed this version of the GPL to prohibit the practice for those
+products. If such problems arise substantially in other domains, we
+stand ready to extend this provision to those domains in future versions
+of the GPL, as needed to protect the freedom of users.
+
+Finally, every program is threatened constantly by software patents.
+States should not allow patents to restrict development and use of
+software on general-purpose computers, but in those that do, we wish to
+avoid the special danger that patents applied to a free program could
+make it effectively proprietary. To prevent this, the GPL assures that
+patents cannot be used to render the program non-free.
+
+The precise terms and conditions for copying, distribution and
+modification follow.
+
+END OF TERMS AND CONDITIONS
+
+For the full license text, see <https://www.gnu.org/licenses/gpl-3.0.txt>.

README.md

@@ -0,0 +1,221 @@
+# 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
+
+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:
+
+```text
+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:
+
+```text
+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 `λ` 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`
+
+### `max_history`
+Cap for cached real-forward feature points used for the fit.
+
+This is an implementation guard, not a paper hyperparameter. With standard FLUX schedules it is usually far above the number of actual cached points anyway.
+
+### `debug`
+Enables lightweight logging during patch install.
+
+## 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`
+
+### More aggressive
+
+- `blend_weight = 0.75`
+- `degree = 4`
+- `ridge_lambda = 0.10`
+- `window_size = 2.0`
+- `flex_window = 3.0`
+- `warmup_steps = 5`
+
+## 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**.
+- 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.
+- 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:
+
+```bash
+cd ComfyUI/custom_nodes/ComfyUI-Spectrum-Proper
+python tests/smoke_runtime.py
+```
+
+Expected output:
+
+```text
+ok
+```
+
+## Repo structure
+
+```text
+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 paper and official code by Jiaqi Han et al.
+- ComfyUI FLUX integration details adapted against current native ComfyUI FLUX internals
+
+## 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.

__init__.py

@@ -0,0 +1,3 @@
+from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
+
+__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]

comfyui_spectrum/__init__.py

@@ -0,0 +1,9 @@
+from .config import SpectrumConfig
+from .flux import FluxSpectrumPatcher
+from .runtime import SpectrumRuntime
+
+__all__ = [
+    "FluxSpectrumPatcher",
+    "SpectrumConfig",
+    "SpectrumRuntime",
+]

comfyui_spectrum/config.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, asdict
+
+
+@dataclass(slots=True)
+class SpectrumConfig:
+    enabled: bool = True
+    backend: str = "flux"
+    blend_weight: float = 0.50
+    degree: int = 4
+    ridge_lambda: float = 0.10
+    window_size: float = 2.0
+    flex_window: float = 0.75
+    warmup_steps: int = 5
+    max_history: int = 128
+    debug: bool = False
+
+    def validate(self) -> "SpectrumConfig":
+        if self.backend != "flux":
+            raise ValueError(f"Unsupported backend: {self.backend!r}")
+        if not (0.0 <= float(self.blend_weight) <= 1.0):
+            raise ValueError("blend_weight must be in [0, 1].")
+        if int(self.degree) < 1:
+            raise ValueError("degree must be >= 1.")
+        if float(self.ridge_lambda) < 0.0:
+            raise ValueError("ridge_lambda must be >= 0.")
+        if float(self.window_size) < 1.0:
+            raise ValueError("window_size must be >= 1.")
+        if float(self.flex_window) < 0.0:
+            raise ValueError("flex_window must be >= 0.")
+        if int(self.warmup_steps) < 0:
+            raise ValueError("warmup_steps must be >= 0.")
+        if int(self.max_history) < max(8, int(self.degree) + 1):
+            raise ValueError("max_history must be at least max(8, degree + 1).")
+        return self
+
+    def to_dict(self) -> dict:
+        return asdict(self)

comfyui_spectrum/flux.py

@@ -0,0 +1,351 @@
+from __future__ import annotations
+
+import logging
+from typing import Any, Dict, Optional, Sequence, Tuple
+
+import torch
+
+from .config import SpectrumConfig
+from .runtime import SpectrumRuntime
+
+LOG = logging.getLogger(__name__)
+
+
+def _clone_model(model: Any) -> Any:
+    return model.clone() if hasattr(model, "clone") else model
+
+
+def _ensure_model_options(model: Any) -> Dict[str, Any]:
+    if not hasattr(model, "model_options") or model.model_options is None:
+        model.model_options = {}
+    return model.model_options
+
+
+def _ensure_transformer_options(model: Any) -> Dict[str, Any]:
+    options = _ensure_model_options(model)
+    if "transformer_options" not in options or options["transformer_options"] is None:
+        options["transformer_options"] = {}
+    return options["transformer_options"]
+
+
+def locate_flux_inner_model(model: Any) -> Tuple[Optional[Any], Optional[str]]:
+    outer = getattr(model, "model", None)
+    if outer is not None and hasattr(outer, "diffusion_model"):
+        return outer.diffusion_model, "model.diffusion_model"
+    if hasattr(model, "diffusion_model"):
+        return model.diffusion_model, "diffusion_model"
+    return None, None
+
+
+def _invert_slices(slices: Sequence[Tuple[int, int]], length: int):
+    sorted_slices = sorted(slices)
+    result = []
+    current = 0
+    for start, end in sorted_slices:
+        if current < start:
+            result.append((current, start))
+        current = max(current, end)
+    if current < length:
+        result.append((current, length))
+    return result
+
+
+def _install_generic_flux_wrapper(inner: Any) -> None:
+    if getattr(inner, "_spectrum_forward_orig_installed", False):
+        return
+
+    original_forward_orig = inner.forward_orig
+    setattr(inner, "_spectrum_original_forward_orig", original_forward_orig)
+
+    def spectrum_forward_orig(
+        self,
+        img,
+        img_ids,
+        txt,
+        txt_ids,
+        timesteps,
+        y,
+        guidance=None,
+        control=None,
+        timestep_zero_index=None,
+        transformer_options=None,
+        attn_mask=None,
+    ):
+        options = transformer_options or {}
+        runtime = options.get("spectrum_runtime") if isinstance(options, dict) else None
+        if runtime is None or not getattr(runtime.cfg, "enabled", False):
+            return self._spectrum_original_forward_orig(
+                img,
+                img_ids,
+                txt,
+                txt_ids,
+                timesteps,
+                y,
+                guidance=guidance,
+                control=control,
+                timestep_zero_index=timestep_zero_index,
+                transformer_options=options,
+                attn_mask=attn_mask,
+            )
+
+        return _run_flux_forward_with_spectrum(
+            self,
+            runtime,
+            img=img,
+            img_ids=img_ids,
+            txt=txt,
+            txt_ids=txt_ids,
+            timesteps=timesteps,
+            y=y,
+            guidance=guidance,
+            control=control,
+            timestep_zero_index=timestep_zero_index,
+            transformer_options=options,
+            attn_mask=attn_mask,
+        )
+
+    inner.forward_orig = spectrum_forward_orig.__get__(inner, type(inner))
+    setattr(inner, "_spectrum_forward_orig_installed", True)
+
+
+def _run_flux_forward_with_spectrum(
+    inner: Any,
+    runtime: SpectrumRuntime,
+    *,
+    img: torch.Tensor,
+    img_ids: torch.Tensor,
+    txt: torch.Tensor,
+    txt_ids: torch.Tensor,
+    timesteps: torch.Tensor,
+    y: torch.Tensor,
+    guidance: Optional[torch.Tensor],
+    control: Optional[dict],
+    timestep_zero_index=None,
+    transformer_options: Dict[str, Any],
+    attn_mask: Optional[torch.Tensor],
+) -> torch.Tensor:
+    from comfy.ldm.flux.layers import timestep_embedding
+
+    transformer_options = (transformer_options or {}).copy()
+    patches = transformer_options.get("patches", {})
+    patches_replace = transformer_options.get("patches_replace", {})
+
+    decision = runtime.begin_step(transformer_options, timesteps)
+    step_idx = decision["step_idx"]
+    total_steps = decision["total_steps"]
+    actual_forward = decision["actual_forward"]
+
+    if img.ndim != 3 or txt.ndim != 3:
+        raise ValueError("Input img and txt tensors must have 3 dimensions.")
+
+    img = inner.img_in(img)
+    vec = inner.time_in(timestep_embedding(timesteps, 256).to(img.dtype))
+
+    if inner.params.guidance_embed and guidance is not None:
+        vec = vec + inner.guidance_in(timestep_embedding(guidance, 256).to(img.dtype))
+
+    if inner.vector_in is not None:
+        if y is None:
+            y = torch.zeros((img.shape[0], inner.params.vec_in_dim), device=img.device, dtype=img.dtype)
+        vec = vec + inner.vector_in(y[:, : inner.params.vec_in_dim])
+
+    if inner.txt_norm is not None:
+        txt = inner.txt_norm(txt)
+    txt = inner.txt_in(txt)
+
+    if "post_input" in patches:
+        for patch in patches["post_input"]:
+            out = patch(
+                {
+                    "img": img,
+                    "txt": txt,
+                    "img_ids": img_ids,
+                    "txt_ids": txt_ids,
+                    "transformer_options": transformer_options,
+                }
+            )
+            img = out["img"]
+            txt = out["txt"]
+            img_ids = out["img_ids"]
+            txt_ids = out["txt_ids"]
+
+    if img_ids is not None:
+        ids = torch.cat((txt_ids, img_ids), dim=1)
+        pe = inner.pe_embedder(ids)
+    else:
+        pe = None
+
+    vec_orig = vec
+    txt_vec = vec
+    extra_kwargs = {}
+    modulation_dims = None
+
+    if timestep_zero_index is not None:
+        modulation_dims = []
+        batch = vec.shape[0] // 2
+        vec_orig = vec_orig.reshape(2, batch, vec.shape[1]).movedim(0, 1)
+        for start, end in _invert_slices(timestep_zero_index, img.shape[1]):
+            modulation_dims.append((start, end, 0))
+        for start, end in timestep_zero_index:
+            modulation_dims.append((start, end, 1))
+        extra_kwargs["modulation_dims_img"] = modulation_dims
+        txt_vec = vec[:batch]
+
+    if not actual_forward and runtime.forecaster.ready(runtime.min_fit_points):
+        pred_feature = runtime.forecaster.predict(
+            step_index=step_idx,
+            total_steps=total_steps,
+            blend_weight=runtime.cfg.blend_weight,
+        )
+        final_kwargs = {}
+        if modulation_dims is not None:
+            final_kwargs["modulation_dims"] = modulation_dims
+        return inner.final_layer(pred_feature.to(img.dtype), vec_orig, **final_kwargs)
+
+    if inner.params.global_modulation:
+        vec = (inner.double_stream_modulation_img(vec_orig), inner.double_stream_modulation_txt(txt_vec))
+
+    blocks_replace = patches_replace.get("dit", {})
+    transformer_options["total_blocks"] = len(inner.double_blocks)
+    transformer_options["block_type"] = "double"
+
+    for block_index, block in enumerate(inner.double_blocks):
+        transformer_options["block_index"] = block_index
+        if ("double_block", block_index) in blocks_replace:
+            def block_wrap(args):
+                out = {}
+                out["img"], out["txt"] = block(
+                    img=args["img"],
+                    txt=args["txt"],
+                    vec=args["vec"],
+                    pe=args["pe"],
+                    attn_mask=args.get("attn_mask"),
+                    transformer_options=args.get("transformer_options"),
+                    **extra_kwargs,
+                )
+                return out
+
+            out = blocks_replace[("double_block", block_index)](
+                {
+                    "img": img,
+                    "txt": txt,
+                    "vec": vec,
+                    "pe": pe,
+                    "attn_mask": attn_mask,
+                    "transformer_options": transformer_options,
+                },
+                {"original_block": block_wrap},
+            )
+            txt = out["txt"]
+            img = out["img"]
+        else:
+            img, txt = block(
+                img=img,
+                txt=txt,
+                vec=vec,
+                pe=pe,
+                attn_mask=attn_mask,
+                transformer_options=transformer_options,
+                **extra_kwargs,
+            )
+
+        if control is not None:
+            control_i = control.get("input")
+            if control_i is not None and block_index < len(control_i):
+                add = control_i[block_index]
+                if add is not None:
+                    img[:, : add.shape[1]] += add
+
+    if img.dtype == torch.float16:
+        img = torch.nan_to_num(img, nan=0.0, posinf=65504, neginf=-65504)
+
+    img = torch.cat((txt, img), dim=1)
+
+    if inner.params.global_modulation:
+        vec, _ = inner.single_stream_modulation(vec_orig)
+
+    extra_kwargs = {}
+    if modulation_dims is not None:
+        modulation_dims_combined = [
+            (0 if start == 0 else start + txt.shape[1], end + txt.shape[1], kind)
+            for start, end, kind in modulation_dims
+        ]
+        extra_kwargs["modulation_dims"] = modulation_dims_combined
+
+    transformer_options["total_blocks"] = len(inner.single_blocks)
+    transformer_options["block_type"] = "single"
+    transformer_options["img_slice"] = [txt.shape[1], img.shape[1]]
+
+    for block_index, block in enumerate(inner.single_blocks):
+        transformer_options["block_index"] = block_index
+        if ("single_block", block_index) in blocks_replace:
+            def block_wrap(args):
+                out = {}
+                out["img"] = block(
+                    args["img"],
+                    vec=args["vec"],
+                    pe=args["pe"],
+                    attn_mask=args.get("attn_mask"),
+                    transformer_options=args.get("transformer_options"),
+                    **extra_kwargs,
+                )
+                return out
+
+            out = blocks_replace[("single_block", block_index)](
+                {
+                    "img": img,
+                    "vec": vec,
+                    "pe": pe,
+                    "attn_mask": attn_mask,
+                    "transformer_options": transformer_options,
+                },
+                {"original_block": block_wrap},
+            )
+            img = out["img"]
+        else:
+            img = block(
+                img,
+                vec=vec,
+                pe=pe,
+                attn_mask=attn_mask,
+                transformer_options=transformer_options,
+                **extra_kwargs,
+            )
+
+        if control is not None:
+            control_o = control.get("output")
+            if control_o is not None and block_index < len(control_o):
+                add = control_o[block_index]
+                if add is not None:
+                    img[:, txt.shape[1] : txt.shape[1] + add.shape[1], ...] += add
+
+    prehead_feature = img[:, txt.shape[1] :, ...]
+    runtime.forecaster.update(step_idx, prehead_feature)
+
+    final_kwargs = {}
+    if modulation_dims is not None:
+        final_kwargs["modulation_dims"] = modulation_dims
+    return inner.final_layer(prehead_feature, vec_orig, **final_kwargs)
+
+
+class FluxSpectrumPatcher:
+    @staticmethod
+    def patch(model: Any, cfg: SpectrumConfig) -> Any:
+        cfg = cfg.validate()
+        patched = _clone_model(model)
+        transformer_options = _ensure_transformer_options(patched)
+        runtime = SpectrumRuntime(cfg)
+        transformer_options["spectrum_runtime"] = runtime
+        transformer_options["spectrum_enabled"] = cfg.enabled
+        transformer_options["spectrum_backend"] = "flux"
+        transformer_options["spectrum_cfg"] = cfg.to_dict()
+
+        inner, path = locate_flux_inner_model(patched)
+        if inner is None:
+            raise RuntimeError(
+                "Could not locate ComfyUI Flux diffusion model. This node currently supports native FLUX models only."
+            )
+        _install_generic_flux_wrapper(inner)
+        if cfg.debug:
+            LOG.warning("Spectrum installed on %s", path)
+        return patched

comfyui_spectrum/forecast.py

@@ -0,0 +1,127 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import List, Optional
+
+import torch
+
+
+@dataclass(slots=True)
+class _HistoryEntry:
+    step_index: int
+    feature: torch.Tensor
+
+
+class ChebyshevSpectrumForecaster:
+    """Online Chebyshev forecaster with optional linear blending.
+
+    The forecaster operates on the final hidden feature of the denoiser, not on
+    the denoised output itself. This matches the official Spectrum integration
+    strategy for FLUX and is the main reason this port stays on the model path
+    instead of wrapping the whole sampler output.
+    """
+
+    def __init__(self, degree: int = 4, ridge_lambda: float = 0.1, max_history: int = 128):
+        self.degree = int(degree)
+        self.ridge_lambda = float(ridge_lambda)
+        self.max_history = int(max_history)
+        self.reset()
+
+    def reset(self) -> None:
+        self._history: List[_HistoryEntry] = []
+        self._feature_shape: Optional[torch.Size] = None
+        self._feature_dtype: Optional[torch.dtype] = None
+        self._device: Optional[torch.device] = None
+
+    def configure(self, degree: int, ridge_lambda: float, max_history: int) -> None:
+        self.degree = int(degree)
+        self.ridge_lambda = float(ridge_lambda)
+        self.max_history = int(max_history)
+
+    def ready(self, min_points: Optional[int] = None) -> bool:
+        needed = max(2, int(min_points) if min_points is not None else self.degree + 1)
+        return len(self._history) >= needed
+
+    def update(self, step_index: int, feature: torch.Tensor) -> None:
+        feat = feature.detach()
+        if self._feature_shape is None:
+            self._feature_shape = feat.shape
+            self._feature_dtype = feat.dtype
+            self._device = feat.device
+        elif feat.shape != self._feature_shape:
+            raise ValueError(
+                f"Spectrum feature shape changed from {tuple(self._feature_shape)} to {tuple(feat.shape)}."
+            )
+
+        self._history.append(_HistoryEntry(int(step_index), feat))
+        if len(self._history) > self.max_history:
+            self._history.pop(0)
+
+    def _tau(self, step_indices: torch.Tensor, total_steps: int) -> torch.Tensor:
+        span = max(int(total_steps) - 1, 1)
+        return (step_indices.to(torch.float32) / float(span)) * 2.0 - 1.0
+
+    def _build_design(self, taus: torch.Tensor, degree: int) -> torch.Tensor:
+        taus = taus.reshape(-1, 1)
+        cols = [torch.ones((taus.shape[0], 1), device=taus.device, dtype=torch.float32)]
+        if degree >= 1:
+            cols.append(taus.to(torch.float32))
+            for _ in range(2, degree + 1):
+                cols.append(2.0 * taus * cols[-1] - cols[-2])
+        return torch.cat(cols[: degree + 1], dim=1)
+
+    def _solve(self, design: torch.Tensor, features: torch.Tensor) -> torch.Tensor:
+        p = design.shape[1]
+        lhs = design.transpose(0, 1) @ design
+        if self.ridge_lambda > 0.0:
+            lhs = lhs + self.ridge_lambda * torch.eye(p, device=design.device, dtype=design.dtype)
+        rhs = design.transpose(0, 1) @ features
+        try:
+            chol = torch.linalg.cholesky(lhs)
+        except RuntimeError:
+            diag_mean = lhs.diag().mean() if lhs.numel() else torch.tensor(1.0, device=lhs.device)
+            jitter = max(float(diag_mean.item()) * 1e-6, 1e-8)
+            chol = torch.linalg.cholesky(lhs + jitter * torch.eye(p, device=lhs.device, dtype=lhs.dtype))
+        return torch.cholesky_solve(rhs, chol)
+
+    def _linear_prediction(self, step_index: int) -> torch.Tensor:
+        last = self._history[-1]
+        if len(self._history) < 2:
+            return last.feature.to(torch.float32)
+
+        prev = self._history[-2]
+        delta_steps = last.step_index - prev.step_index
+        if delta_steps <= 0:
+            return last.feature.to(torch.float32)
+
+        k = (float(step_index) - float(last.step_index)) / float(delta_steps)
+        last_f = last.feature.to(torch.float32)
+        prev_f = prev.feature.to(torch.float32)
+        return last_f + k * (last_f - prev_f)
+
+    def predict(self, step_index: int, total_steps: int, blend_weight: float) -> torch.Tensor:
+        if self._feature_shape is None or self._feature_dtype is None or self._device is None:
+            raise RuntimeError("Spectrum forecaster has no cached feature history.")
+        if not self.ready():
+            raise RuntimeError("Spectrum forecaster is not ready yet.")
+
+        degree = min(self.degree, len(self._history) - 1)
+        steps = torch.tensor([entry.step_index for entry in self._history], device=self._device, dtype=torch.float32)
+        features = torch.stack(
+            [entry.feature.reshape(-1).to(torch.float32) for entry in self._history],
+            dim=0,
+        )
+
+        design = self._build_design(self._tau(steps, total_steps), degree)
+        coeff = self._solve(design, features)
+
+        tau_star = self._tau(
+            torch.tensor([int(step_index)], device=self._device, dtype=torch.float32),
+            total_steps,
+        )
+        design_star = self._build_design(tau_star, degree)
+        spectral = (design_star @ coeff).reshape(self._feature_shape)
+
+        linear = self._linear_prediction(step_index).reshape(self._feature_shape)
+        out = float(blend_weight) * spectral + (1.0 - float(blend_weight)) * linear
+        return out.to(dtype=self._feature_dtype)

comfyui_spectrum/runtime.py

@@ -0,0 +1,164 @@
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional
+
+import torch
+
+from .config import SpectrumConfig
+from .forecast import ChebyshevSpectrumForecaster
+
+
+@dataclass(slots=True)
+class RuntimeStats:
+    actual_forward_count: int = 0
+    forecasted_count: int = 0
+    total_steps: int = 0
+    current_window: float = 0.0
+    last_sigma: Optional[float] = None
+    run_id: int = 0
+
+
+class SpectrumRuntime:
+    def __init__(self, cfg: SpectrumConfig):
+        self.cfg = cfg.validate()
+        self.forecaster = ChebyshevSpectrumForecaster(
+            degree=self.cfg.degree,
+            ridge_lambda=self.cfg.ridge_lambda,
+            max_history=self.cfg.max_history,
+        )
+        self._last_schedule_signature: Optional[tuple] = None
+        self.run_id = 0
+        self.stats = RuntimeStats(current_window=float(self.cfg.window_size))
+        self.reset_cycle(reset_schedule=False)
+
+    @property
+    def min_fit_points(self) -> int:
+        return max(2, self.cfg.degree + 1)
+
+    def reset_cycle(self, reset_schedule: bool = False) -> None:
+        self.step_idx = 0
+        self.curr_ws = float(self.cfg.window_size)
+        self.num_consecutive_cached_steps = 0
+        self.decisions_by_sigma: Dict[float, Dict[str, Any]] = {}
+        self.seen_sigmas: List[float] = []
+        self.cycle_finished = False
+        self.forecaster.reset()
+        if reset_schedule:
+            self._last_schedule_signature = None
+        self.stats.current_window = float(self.cfg.window_size)
+
+    def reset_all(self) -> None:
+        self.run_id += 1
+        self.stats = RuntimeStats(current_window=float(self.cfg.window_size), run_id=self.run_id)
+        self.reset_cycle(reset_schedule=True)
+
+    def _schedule_signature(self, transformer_options: Dict[str, Any]) -> Optional[tuple]:
+        sample_sigmas = transformer_options.get("sample_sigmas")
+        if sample_sigmas is None:
+            return None
+        try:
+            values = sample_sigmas.detach().float().cpu().flatten().tolist()
+            return tuple(round(float(v), 8) for v in values)
+        except Exception:
+            return None
+
+    def _ensure_run_sync(self, transformer_options: Dict[str, Any]) -> None:
+        signature = self._schedule_signature(transformer_options)
+        if signature is None:
+            return
+        if self._last_schedule_signature is None:
+            self._last_schedule_signature = signature
+            self.stats.total_steps = max(len(signature) - 1, 1)
+            return
+        if signature != self._last_schedule_signature:
+            self.run_id += 1
+            self._last_schedule_signature = signature
+            self.stats.actual_forward_count = 0
+            self.stats.forecasted_count = 0
+            self.stats.total_steps = max(len(signature) - 1, 1)
+            self.stats.run_id = self.run_id
+            self.reset_cycle(reset_schedule=False)
+
+    def num_steps(self) -> int:
+        if self.stats.total_steps > 0:
+            return self.stats.total_steps
+        return 50
+
+    def _sigma_key(self, transformer_options: Dict[str, Any], timesteps: torch.Tensor) -> float:
+        sigmas = transformer_options.get("sigmas")
+        if sigmas is not None:
+            try:
+                return round(float(sigmas.detach().flatten()[0].item()), 8)
+            except Exception:
+                pass
+        try:
+            return round(float(timesteps.detach().flatten()[0].item()), 8)
+        except Exception:
+            return float(self.step_idx)
+
+    def _finish_cycle_if_needed(self) -> None:
+        if len(self.seen_sigmas) >= self.num_steps() and not self.cycle_finished:
+            self.cycle_finished = True
+
+    def _restart_cycle(self) -> None:
+        self.run_id += 1
+        self.stats.actual_forward_count = 0
+        self.stats.forecasted_count = 0
+        self.stats.run_id = self.run_id
+        self.reset_cycle(reset_schedule=False)
+
+    def _should_restart_on_sigma(self, sigma: float) -> bool:
+        if not self.seen_sigmas:
+            return False
+        if sigma != self.seen_sigmas[0]:
+            return False
+        return len(self.seen_sigmas) > 1
+
+    def begin_step(self, transformer_options: Dict[str, Any], timesteps: torch.Tensor) -> Dict[str, Any]:
+        transformer_options = transformer_options or {}
+        self._ensure_run_sync(transformer_options)
+
+        sigma = self._sigma_key(transformer_options, timesteps)
+        self.stats.last_sigma = sigma
+        self._finish_cycle_if_needed()
+
+        if self.cycle_finished or self._should_restart_on_sigma(sigma):
+            self._restart_cycle()
+
+        if sigma in self.decisions_by_sigma:
+            return self.decisions_by_sigma[sigma]
+
+        step_idx = len(self.seen_sigmas)
+        self.seen_sigmas.append(sigma)
+
+        actual_forward = True
+        if step_idx >= self.cfg.warmup_steps and self.forecaster.ready(self.min_fit_points):
+            ws_floor = max(1, int(math.floor(self.curr_ws)))
+            actual_forward = ((self.num_consecutive_cached_steps + 1) % ws_floor) == 0
+
+        if not self.forecaster.ready(self.min_fit_points):
+            actual_forward = True
+
+        if actual_forward:
+            self.num_consecutive_cached_steps = 0
+            if step_idx >= self.cfg.warmup_steps:
+                self.curr_ws = round(self.curr_ws + float(self.cfg.flex_window), 6)
+            self.stats.actual_forward_count += 1
+        else:
+            self.num_consecutive_cached_steps += 1
+            self.stats.forecasted_count += 1
+
+        self.step_idx = step_idx
+        self.stats.current_window = self.curr_ws
+
+        decision = {
+            "sigma": sigma,
+            "step_idx": step_idx,
+            "total_steps": self.num_steps(),
+            "actual_forward": actual_forward,
+            "run_id": self.run_id,
+        }
+        self.decisions_by_sigma[sigma] = decision
+        return decision

nodes.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+
+from .comfyui_spectrum import FluxSpectrumPatcher, SpectrumConfig
+
+
+class SpectrumApplyFlux:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "model": ("MODEL",),
+                "enabled": ("BOOLEAN", {"default": True}),
+                "blend_weight": ("FLOAT", {"default": 0.50, "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": 16.0, "step": 0.05}),
+                "flex_window": ("FLOAT", {"default": 0.75, "min": 0.0, "max": 8.0, "step": 0.05}),
+                "warmup_steps": ("INT", {"default": 5, "min": 0, "max": 32, "step": 1}),
+                "max_history": ("INT", {"default": 128, "min": 8, "max": 512, "step": 1}),
+                "debug": ("BOOLEAN", {"default": False}),
+            }
+        }
+
+    RETURN_TYPES = ("MODEL",)
+    RETURN_NAMES = ("model",)
+    FUNCTION = "apply"
+    CATEGORY = "sampling/spectrum"
+
+    def apply(
+        self,
+        model,
+        enabled,
+        blend_weight,
+        degree,
+        ridge_lambda,
+        window_size,
+        flex_window,
+        warmup_steps,
+        max_history,
+        debug,
+    ):
+        if not enabled:
+            return (model,)
+
+        cfg = SpectrumConfig(
+            enabled=bool(enabled),
+            backend="flux",
+            blend_weight=float(blend_weight),
+            degree=int(degree),
+            ridge_lambda=float(ridge_lambda),
+            window_size=float(window_size),
+            flex_window=float(flex_window),
+            warmup_steps=int(warmup_steps),
+            max_history=int(max_history),
+            debug=bool(debug),
+        ).validate()
+        return (FluxSpectrumPatcher.patch(model, cfg),)
+
+
+NODE_CLASS_MAPPINGS = {
+    "SpectrumApplyFlux": SpectrumApplyFlux,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "SpectrumApplyFlux": "Spectrum Apply Flux",
+}

pyproject.toml

@@ -0,0 +1,18 @@
+[project]
+name = "comfyui-spectrum-proper"
+description = "Faithful FLUX-focused ComfyUI port of Spectrum spectral feature forecasting."
+version = "0.1.0"
+requires-python = ">=3.10"
+dependencies = []
+classifiers = [
+  "Operating System :: OS Independent",
+  "Programming Language :: Python :: 3",
+  "Programming Language :: Python :: 3.10",
+  "Programming Language :: Python :: 3.11",
+  "License :: OSI Approved :: GNU General Public License v3 or later (GPLv3+)",
+]
+
+[tool.comfy]
+PublisherId = "local"
+DisplayName = "ComfyUI-Spectrum-Proper"
+Icon = ""

tests/smoke_runtime.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+import torch
+
+ROOT = Path(__file__).resolve().parents[1]
+if str(ROOT) not in sys.path:
+    sys.path.insert(0, str(ROOT))
+
+from comfyui_spectrum.config import SpectrumConfig
+from comfyui_spectrum.runtime import SpectrumRuntime
+
+
+def main() -> None:
+    cfg = SpectrumConfig(
+        blend_weight=0.5,
+        degree=4,
+        ridge_lambda=0.1,
+        window_size=2.0,
+        flex_window=0.75,
+        warmup_steps=5,
+        max_history=128,
+    ).validate()
+    runtime = SpectrumRuntime(cfg)
+    sample_sigmas = torch.linspace(1.0, 0.0, 51)
+    transformer_options = {"sample_sigmas": sample_sigmas}
+
+    # Warm up with five actual features.
+    for i in range(5):
+        decision = runtime.begin_step(transformer_options, torch.tensor([sample_sigmas[i]]))
+        assert decision["actual_forward"] is True
+        runtime.forecaster.update(i, torch.randn(1, 8, 4))
+
+    # Scheduler should now be able to forecast.
+    decision = runtime.begin_step(transformer_options, torch.tensor([sample_sigmas[5]]))
+    assert "actual_forward" in decision
+    assert runtime.forecaster.predict(
+        step_index=5,
+        total_steps=decision["total_steps"],
+        blend_weight=cfg.blend_weight,
+    ).shape == (1, 8, 4)
+
+    print("ok")
+
+
+if __name__ == "__main__":
+    main()