commit 6807027e0101acd25cf5a2811f652e69dbd78803 Author: xmarre Date: Tue Mar 17 05:44:36 2026 +0100 Initial import of built repo diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..94bb7b2 --- /dev/null +++ b/.gitignore @@ -0,0 +1,7 @@ +__pycache__/ +*.pyc +*.pyo +*.pyd +*.so +*.egg-info/ +.pytest_cache/ diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..5de65ad --- /dev/null +++ b/LICENSE @@ -0,0 +1,73 @@ +GNU GENERAL PUBLIC LICENSE +Version 3, 29 June 2007 + +Copyright (C) 2007 Free Software Foundation, Inc. +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 . diff --git a/README.md b/README.md new file mode 100644 index 0000000..44bb80a --- /dev/null +++ b/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. diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..39a8c6b --- /dev/null +++ b/__init__.py @@ -0,0 +1,3 @@ +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/comfyui_spectrum/__init__.py b/comfyui_spectrum/__init__.py new file mode 100644 index 0000000..5d1597a --- /dev/null +++ b/comfyui_spectrum/__init__.py @@ -0,0 +1,9 @@ +from .config import SpectrumConfig +from .flux import FluxSpectrumPatcher +from .runtime import SpectrumRuntime + +__all__ = [ + "FluxSpectrumPatcher", + "SpectrumConfig", + "SpectrumRuntime", +] diff --git a/comfyui_spectrum/config.py b/comfyui_spectrum/config.py new file mode 100644 index 0000000..48d5610 --- /dev/null +++ b/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) diff --git a/comfyui_spectrum/flux.py b/comfyui_spectrum/flux.py new file mode 100644 index 0000000..e475fb7 --- /dev/null +++ b/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 diff --git a/comfyui_spectrum/forecast.py b/comfyui_spectrum/forecast.py new file mode 100644 index 0000000..c40d8b6 --- /dev/null +++ b/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) diff --git a/comfyui_spectrum/runtime.py b/comfyui_spectrum/runtime.py new file mode 100644 index 0000000..1ab0a93 --- /dev/null +++ b/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 diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..0fcd7d9 --- /dev/null +++ b/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", +} diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..f822469 --- /dev/null +++ b/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 = "" diff --git a/tests/smoke_runtime.py b/tests/smoke_runtime.py new file mode 100644 index 0000000..eb8288e --- /dev/null +++ b/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()