From e7ce48191024fc960680896fec28be5eaf6d4829 Mon Sep 17 00:00:00 2001 From: xmarre Date: Tue, 31 Mar 2026 02:25:10 +0200 Subject: [PATCH] Cache forecast coefficients and move history to CPU --- comfyui_spectrum/forecast.py | 60 +++++++++++++++++++++++++----------- 1 file changed, 42 insertions(+), 18 deletions(-) diff --git a/comfyui_spectrum/forecast.py b/comfyui_spectrum/forecast.py index a77eac7..0a866ef 100644 --- a/comfyui_spectrum/forecast.py +++ b/comfyui_spectrum/forecast.py @@ -9,7 +9,7 @@ import torch @dataclass(slots=True) class _HistoryEntry: time_coord: float - feature: torch.Tensor + feature_flat: torch.Tensor class ChebyshevSpectrumForecaster: @@ -32,11 +32,20 @@ class ChebyshevSpectrumForecaster: self._feature_shape: Optional[torch.Size] = None self._feature_dtype: Optional[torch.dtype] = None self._device: Optional[torch.device] = None + self._output_device: Optional[torch.device] = None + self._coeff: Optional[torch.Tensor] = None + self._cached_degree: Optional[int] = None + self._cache_dirty = True 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) + if len(self._history) > self.max_history: + self._history = self._history[-self.max_history :] + self._coeff = None + self._cached_degree = None + self._cache_dirty = True @property def feature_shape(self) -> Optional[torch.Size]: @@ -51,15 +60,20 @@ class ChebyshevSpectrumForecaster: if self._feature_shape is None: self._feature_shape = feat.shape self._feature_dtype = feat.dtype - self._device = feat.device + self._device = torch.device("cpu") + self._output_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(float(time_coord), feat)) + feature_flat = feat.reshape(-1).to(device="cpu", dtype=torch.float32, copy=True) + self._history.append(_HistoryEntry(float(time_coord), feature_flat)) if len(self._history) > self.max_history: self._history.pop(0) + self._coeff = None + self._cached_degree = None + self._cache_dirty = True def _build_design(self, coords: torch.Tensor, degree: int) -> torch.Tensor: coords = coords.reshape(-1, 1).to(torch.float32) @@ -84,36 +98,46 @@ class ChebyshevSpectrumForecaster: chol = torch.linalg.cholesky(lhs + jitter * torch.eye(p, device=lhs.device, dtype=lhs.dtype)) return torch.cholesky_solve(rhs, chol) + def _ensure_coeff(self) -> tuple[int, torch.Tensor]: + degree = min(self.degree, len(self._history) - 1) + if not self._cache_dirty and self._coeff is not None and self._cached_degree == degree: + return degree, self._coeff + + coords = torch.tensor([entry.time_coord for entry in self._history], device="cpu", dtype=torch.float32) + features = torch.stack([entry.feature_flat for entry in self._history], dim=0) + design = self._build_design(coords, degree) + self._coeff = self._solve(design, features) + self._cached_degree = degree + self._cache_dirty = False + return degree, self._coeff + def _linear_prediction(self, time_coord: float) -> torch.Tensor: last = self._history[-1] if len(self._history) < 2: - return last.feature.to(torch.float32) + return last.feature_flat prev = self._history[-2] delta_coord = last.time_coord - prev.time_coord if abs(delta_coord) <= 1e-12: - return last.feature.to(torch.float32) + return last.feature_flat k = (float(time_coord) - float(last.time_coord)) / float(delta_coord) - last_f = last.feature.to(torch.float32) - prev_f = prev.feature.to(torch.float32) + last_f = last.feature_flat + prev_f = prev.feature_flat return last_f + k * (last_f - prev_f) def predict(self, time_coord: float, blend_weight: float) -> torch.Tensor: - if self._feature_shape is None or self._feature_dtype is None or self._device is None: + if ( + self._feature_shape is None + or self._feature_dtype is None + or self._device is None + or self._output_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) - coords = torch.tensor([entry.time_coord 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(coords, degree) - coeff = self._solve(design, features) + degree, coeff = self._ensure_coeff() coord_star = torch.tensor([float(time_coord)], device=self._device, dtype=torch.float32) design_star = self._build_design(coord_star, degree) @@ -121,4 +145,4 @@ class ChebyshevSpectrumForecaster: linear = self._linear_prediction(time_coord).reshape(self._feature_shape) out = float(blend_weight) * spectral + (1.0 - float(blend_weight)) * linear - return out.to(dtype=self._feature_dtype) + return out.to(device=self._output_device, dtype=self._feature_dtype)