Update spectrum.py

spectrum.py

@@ -78,7 +78,7 @@ class SpectrumConfig:
     def window_cap(self) -> float:
         return max(1.0, self.window_size, float(self.history_points))
 
-def _cheb_basis(u, degree: int) -> torch.Tensor:
+def cheb_basis(u, degree: int) -> torch.Tensor:
     u = torch.as_tensor(u, dtype=torch.float64).reshape(-1)
     cols = [torch.ones_like(u)]
     if degree >= 1:
@@ -87,7 +87,7 @@ def _cheb_basis(u, degree: int) -> torch.Tensor:
         cols.append(2.0 * u * cols[-1] - cols[-2])
     return torch.stack(cols[: degree + 1], dim=-1)          # [n, P]
 
-class _Forecaster:
+class Forecaster:
     """Chebyshev ridge forecaster (prediction = weighted sum of anchors)"""
     def __init__(self, cfg: SpectrumConfig):
         self.cfg = cfg
@@ -116,7 +116,7 @@ class _Forecaster:
         self._chol = None
 
     def _design(self) -> torch.Tensor:
-        return _cheb_basis(self.us, self.cfg.chebyshev_degree)   # [K, P]
+        return cheb_basis(self.us, self.cfg.chebyshev_degree)   # [K, P]
 
     def _factor(self) -> torch.Tensor:
         if self._chol is not None:
@@ -160,7 +160,7 @@ class _Forecaster:
         return w
 
     def weights(self, u: float) -> torch.Tensor:
-        phi = _cheb_basis([u], self.cfg.chebyshev_degree)        # [1, P]
+        phi = cheb_basis([u], self.cfg.chebyshev_degree)        # [1, P]
         solved = torch.cholesky_solve(self._design().T, self._factor())  # [P, K]
         spectral = (phi @ solved).reshape(-1)
         bw = self.cfg.blend_weight
@@ -194,7 +194,7 @@ class _Forecaster:
         den = actual.float().norm().clamp_min(1e-6)
         return float((num / den).item())
 
-class _Branch:
+class Branch:
     """Per-branch scheduler state"""
     def __init__(self, cfg: SpectrumConfig):
         self.cfg = cfg
@@ -202,13 +202,13 @@ class _Branch:
         self.consecutive = 0
         self.actual = 0
         self.forecast = 0
-        self.forecasters: dict[int, _Forecaster] = {}
+        self.forecasters: dict[int, Forecaster] = {}
         self.errors: list[float] = []
 
-    def forecaster(self, sample: int) -> _Forecaster:
+    def forecaster(self, sample: int) -> Forecaster:
         f = self.forecasters.get(sample)
         if f is None:
-            f = self.forecasters[sample] = _Forecaster(self.cfg)
+            f = self.forecasters[sample] = Forecaster(self.cfg)
         return f
 
     def decide(self, step: int, total: int) -> bool:
@@ -234,20 +234,20 @@ class _Branch:
             self.consecutive += 1
 
 
-class _State:
+class State:
     def __init__(self, cfg: SpectrumConfig, scheduler: Any):
         self.cfg = cfg
         self.scheduler = scheduler
         self.original_forward = None
         self.step_index = -1
-        self.branches: dict[Any, _Branch] = {}
+        self.branches: dict[Any, Branch] = {}
         self._sig_counts: dict[Any, int] = {}
         self._log: list[str] = []
 
-    def branch(self, key: Any) -> _Branch:
+    def branch(self, key: Any) -> Branch:
         b = self.branches.get(key)
         if b is None:
-            b = self.branches[key] = _Branch(self.cfg)
+            b = self.branches[key] = Branch(self.cfg)
         return b
 
     def reset_run(self) -> None:
@@ -266,20 +266,13 @@ class _State:
             f"{f}/{a + f} ({speed:.2f}x) mean_forecast_rel_err={mean_err:.4f}"
         )
 
-def _validate_skip_layers(layers, depth: int, name: str = "skip_layers") -> tuple:
+def validate_skip_layers(layers, depth: int, name: str = "skip_layers") -> tuple:
     layers = () if layers is None else tuple(layers)
     if any(type(i) is not int or not 0 <= i < depth for i in layers):
         raise ValueError(f"{name} must contain integer block indices in [0, {depth}).")
     return layers
 
-
-def _coord(step: int, total: int) -> float:
-    if total <= 1:
-        return 0.0
-    return max(-1.0, min(1.0, 2.0 * step / (total - 1) - 1.0))
-
-
-def _resolve_step(state: _State, timestep: torch.Tensor) -> tuple[int, int]:
+def resolve_step(state: State, timestep: torch.Tensor) -> tuple[int, int]:
     sched = state.scheduler
     timesteps = getattr(sched, "timesteps", None) if sched is not None else None
     if timesteps is None or len(timesteps) == 0:
@@ -289,7 +282,7 @@ def _resolve_step(state: _State, timestep: torch.Tensor) -> tuple[int, int]:
     idx = int((ts.float() - t).abs().argmin().item())
     return idx, len(ts)
 
-def _spectrum_forward(
+def spectrum_forward(
     self,
     hidden_states: torch.Tensor,
     timestep,
@@ -302,7 +295,7 @@ def _spectrum_forward(
     License: Apache-2.0
     https://huggingface.co/nebulette/moss-spectrum-diffusers
     """
-    state: Optional[_State] = getattr(self, "_spectrum_state", None)
+    state: Optional[State] = getattr(self, "_spectrum_state", None)
     if state is None or state.scheduler is None:
         return state.original_forward(
             hidden_states, timestep, encoder_hidden_states,
@@ -311,7 +304,7 @@ def _spectrum_forward(
         )
 
     cfg = state.cfg
-    skip = _validate_skip_layers(skip_layers, len(self.layers))
+    skip = validate_skip_layers(skip_layers, len(self.layers))
 
     if hidden_states.ndim != 4 or hidden_states.shape[1] != self.config.latent_channels:
         raise ValueError("hidden_states must be [B, latent_channels, H, W].")
@@ -326,7 +319,7 @@ def _spectrum_forward(
         raise ValueError("encoder_attention_mask must be a binary [B, S] mask.")
 
     # --- step bookkeeping ------------------------------------------------- #
-    step, total = _resolve_step(state, timestep)
+    step, total = resolve_step(state, timestep)
     if step < 0:  # no schedule available -> stay exact
         return state.original_forward(
             hidden_states, timestep, encoder_hidden_states,
@@ -345,7 +338,7 @@ def _spectrum_forward(
     branch = state.branch((sig, occurrence))
 
     actual = branch.decide(step, total)
-    u = _coord(step, total)
+    u = 0.0 if total <= 1 else max(-1.0, min(1.0, 2.0 * step / (total - 1) - 1.0))
 
     # shared timestep conditioning (identical to the original)
     times = torch.as_tensor(timestep, device=hidden_states.device, dtype=torch.float32).reshape(-1).expand(batch)
@@ -374,7 +367,7 @@ def _spectrum_forward(
             # skip embeddings, text blocks and all layers
             try:
                 feat = forecaster.predict(u)
-            except Exception:  
+            except Exception:
                 # numerical trouble -> fall back to a real step
                 actual = True
                 branch.note(True, step)
@@ -428,11 +421,11 @@ def _spectrum_forward(
         return (sample,)
     return Transformer2DModelOutput(sample=sample)
 
-    
+
 def apply_spectrum(pipe, config: Optional[SpectrumConfig] = None):
     """
     Patch ``pipe`` in place so the MOSS DiT runs under Spectrum scheduling.
-    
+
     License: Apache-2.0
     https://huggingface.co/nebulette/moss-spectrum-diffusers
     """
@@ -446,11 +439,11 @@ def apply_spectrum(pipe, config: Optional[SpectrumConfig] = None):
     if getattr(transformer, "_spectrum_state", None) is not None:
         remove_spectrum(pipe)
 
-    state = _State(cfg, getattr(pipe, "scheduler", None))
+    state = State(cfg, getattr(pipe, "scheduler", None))
     state.original_forward = transformer.forward
 
     transformer._spectrum_state = state
-    transformer.forward = types.MethodType(_spectrum_forward, transformer)
+    transformer.forward = types.MethodType(spectrum_forward, transformer)
     return pipe