Update modeling_patchsvae.py

modeling_patchsvae.py

@@ -1,19 +1,20 @@
 """PatchSVAE model for HuggingFace AutoModel.
 
 Usage:
-    from transformers import AutoModel
+    from transformers import AutoConfig, AutoModel
+
+    config = AutoConfig.from_pretrained("AbstractPhil/svae-fresnel-128", trust_remote_code=True)
     model = AutoModel.from_pretrained("AbstractPhil/svae-fresnel-128", trust_remote_code=True)
 
     # Full reconstruction
-    output = model(images)  # output.recon, output.svd
-
-    # Encode to omega tokens (latent)
-    latent = model.encode(images)  # (B, D, gh, gw) = (B, 16, 8, 8)
+    output = model(images)
+    recon = output["recon"]          # (B, 3, 128, 128)
+    latent = output["latent"]        # (B, 16, 8, 8) omega tokens
 
-    # Decode from omega tokens
-    recon = model.decode(latent)   # (B, 3, 128, 128)
+    # Encode to omega tokens
+    omega = model.encode(images)     # (B, 16, 8, 8)
 
-    # Encode to full SVD components
+    # Full SVD decomposition
     svd = model.encode_full(images)  # dict with U, S, Vt, M per patch
 """
 
@@ -21,22 +22,22 @@ import math
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from dataclasses import dataclass
-from typing import Optional, Dict
+from typing import Optional, Dict, Union
 from transformers import PreTrainedModel
 from .configuration_patchsvae import PatchSVAEConfig
 
 
-# ── SVD Backend ──────────────────────────────────────────────────
+# ── SVD Backend (self-contained, no external deps required) ──────
 
 try:
     from geolip_core.linalg.eigh import FLEigh, _FL_MAX_N
-    HAS_FL = True
+    _HAS_FL = True
 except ImportError:
-    HAS_FL = False
+    _HAS_FL = False
 
 
-def _gram_eigh_svd_fp64(A):
+def _gram_eigh_svd(A):
+    """Thin SVD via Gram + eigh in fp64."""
     orig_dtype = A.dtype
     with torch.amp.autocast('cuda', enabled=False):
         A_d = A.double()
@@ -51,8 +52,9 @@ def _gram_eigh_svd_fp64(A):
 
 
 def _svd_fp64(A):
+    """Auto-dispatch: FL eigh for N<=12, Gram eigh otherwise."""
     B, M, N = A.shape
-    if HAS_FL and N <= _FL_MAX_N and A.is_cuda:
+    if _HAS_FL and N <= _FL_MAX_N and A.is_cuda:
         orig_dtype = A.dtype
         with torch.amp.autocast('cuda', enabled=False):
             A_d = A.double()
@@ -65,30 +67,29 @@ def _svd_fp64(A):
             Vh = V.transpose(-2, -1).contiguous()
         return U.to(orig_dtype), S.to(orig_dtype), Vh.to(orig_dtype)
     else:
-        return _gram_eigh_svd_fp64(A)
+        return _gram_eigh_svd(A)
 
 
 # ── Patch Utilities ──────────────────────────────────────────────
 
-def extract_patches(images, patch_size):
+def _extract_patches(images, patch_size):
     B, C, H, W = images.shape
     gh, gw = H // patch_size, W // patch_size
-    patches = images.reshape(B, C, gh, patch_size, gw, patch_size)
-    patches = patches.permute(0, 2, 4, 1, 3, 5)
-    patches = patches.reshape(B, gh * gw, C * patch_size * patch_size)
-    return patches, gh, gw
+    x = images.reshape(B, C, gh, patch_size, gw, patch_size)
+    x = x.permute(0, 2, 4, 1, 3, 5)
+    return x.reshape(B, gh * gw, C * patch_size * patch_size), gh, gw
 
 
-def stitch_patches(patches, gh, gw, patch_size):
+def _stitch_patches(patches, gh, gw, patch_size):
     B = patches.shape[0]
-    patches = patches.reshape(B, gh, gw, 3, patch_size, patch_size)
-    patches = patches.permute(0, 3, 1, 4, 2, 5)
-    return patches.reshape(B, 3, gh * patch_size, gw * patch_size)
+    x = patches.reshape(B, gh, gw, 3, patch_size, patch_size)
+    x = x.permute(0, 3, 1, 4, 2, 5)
+    return x.reshape(B, 3, gh * patch_size, gw * patch_size)
 
 
 # ── Components ───────────────────────────────────────────────────
 
-class BoundarySmooth(nn.Module):
+class _BoundarySmooth(nn.Module):
     def __init__(self, channels=3, mid=16):
         super().__init__()
         self.net = nn.Sequential(
@@ -103,7 +104,7 @@ class BoundarySmooth(nn.Module):
         return x + self.net(x)
 
 
-class SpectralCrossAttention(nn.Module):
+class _SpectralCrossAttention(nn.Module):
     def __init__(self, D, n_heads=4, max_alpha=0.2, alpha_init=-2.0):
         super().__init__()
         self.n_heads = n_heads
@@ -130,80 +131,68 @@ class SpectralCrossAttention(nn.Module):
         attn = attn.softmax(dim=-1)
         out = (attn @ v).transpose(1, 2).reshape(B, N, D)
         gate = torch.tanh(self.out_proj(out))
-        alpha = self.alpha
-        return S * (1.0 + alpha.unsqueeze(0).unsqueeze(0) * gate)
-
-
-# ── Output ───────────────────────────────────────────────────────
-
-@dataclass
-class PatchSVAEOutput:
-    """Output from PatchSVAE forward pass."""
-    recon: torch.Tensor           # (B, 3, H, W) reconstructed image
-    latent: torch.Tensor          # (B, D, gh, gw) omega tokens
-    svd: Optional[Dict] = None    # full SVD components if requested
+        return S * (1.0 + self.alpha.unsqueeze(0).unsqueeze(0) * gate)
 
 
 # ── Model ────────────────────────────────────────────────────────
 
 class PatchSVAEModel(PreTrainedModel):
-    """Patch-based SVD Autoencoder — The Geometric Engine.
+    """Patch-based SVD Autoencoder — The Fresnel Geometric Compression Lens.
 
     Decomposes images into patches, encodes each to a sphere-normalized
     matrix, performs SVD, coordinates spectra via cross-attention,
-    and reconstructs.
+    and reconstructs with 99.993% fidelity.
 
     The spectral vectors S form omega tokens: modality-agnostic,
     geometrically structured, universal representations.
     """
     config_class = PatchSVAEConfig
-    supports_gradient_checkpointing = False
+    _tied_weights_keys = []
 
     def __init__(self, config: PatchSVAEConfig):
         super().__init__(config)
-        self.config = config
 
         V = config.matrix_v
         D = config.D
         hidden = config.hidden
         depth = config.depth
         ps = config.patch_size
-
-        self.patch_dim = 3 * ps * ps
-        self.mat_dim = V * D
+        patch_dim = 3 * ps * ps
+        mat_dim = V * D
 
         # Encoder
-        self.enc_in = nn.Linear(self.patch_dim, hidden)
+        self.enc_in = nn.Linear(patch_dim, hidden)
         self.enc_blocks = nn.ModuleList([
             nn.Sequential(nn.LayerNorm(hidden), nn.Linear(hidden, hidden),
                           nn.GELU(), nn.Linear(hidden, hidden))
             for _ in range(depth)
         ])
-        self.enc_out = nn.Linear(hidden, self.mat_dim)
+        self.enc_out = nn.Linear(hidden, mat_dim)
         nn.init.orthogonal_(self.enc_out.weight)
 
         # Decoder
-        self.dec_in = nn.Linear(self.mat_dim, hidden)
+        self.dec_in = nn.Linear(mat_dim, hidden)
         self.dec_blocks = nn.ModuleList([
             nn.Sequential(nn.LayerNorm(hidden), nn.Linear(hidden, hidden),
                           nn.GELU(), nn.Linear(hidden, hidden))
             for _ in range(depth)
         ])
-        self.dec_out = nn.Linear(hidden, self.patch_dim)
+        self.dec_out = nn.Linear(hidden, patch_dim)
 
         # Cross-attention
         self.cross_attn = nn.ModuleList([
-            SpectralCrossAttention(D, n_heads=min(4, D),
-                                   max_alpha=config.max_alpha,
-                                   alpha_init=config.alpha_init)
+            _SpectralCrossAttention(D, n_heads=min(4, D),
+                                    max_alpha=config.max_alpha,
+                                    alpha_init=config.alpha_init)
             for _ in range(config.n_cross_layers)
         ])
 
         # Boundary smoothing
-        self.boundary_smooth = BoundarySmooth(channels=3, mid=16)
+        self.boundary_smooth = _BoundarySmooth(channels=3, mid=16)
+
+        self.post_init()
 
-    def _encode_patches(self, patches):
-        """Encode flat patches to SVD components."""
+    def _encode_patches_to_svd(self, patches):
         B, N, _ = patches.shape
         V, D = self.config.matrix_v, self.config.D
 
@@ -212,7 +201,7 @@ class PatchSVAEModel(PreTrainedModel):
         for block in self.enc_blocks:
             h = h + block(h)
         M = self.enc_out(h).reshape(B * N, V, D)
-        M = F.normalize(M, dim=-1)  # rows to S^(D-1)
+        M = F.normalize(M, dim=-1)
 
         U, S, Vt = _svd_fp64(M)
 
@@ -221,17 +210,14 @@ class PatchSVAEModel(PreTrainedModel):
         Vt = Vt.reshape(B, N, D, D)
         M = M.reshape(B, N, V, D)
 
-        # Cross-patch spectral coordination
         S_coord = S
         for layer in self.cross_attn:
             S_coord = layer(S_coord)
 
-        return {'U': U, 'S_orig': S, 'S': S_coord, 'Vt': Vt, 'M': M}
+        return {"U": U, "S_orig": S, "S": S_coord, "Vt": Vt, "M": M}
 
-    def _decode_patches(self, U, S, Vt):
-        """Decode from SVD components to flat patches."""
+    def _decode_from_svd(self, U, S, Vt):
         B, N, V, D = U.shape
-
         U_flat = U.reshape(B * N, V, D)
         S_flat = S.reshape(B * N, D)
         Vt_flat = Vt.reshape(B * N, D, D)
@@ -240,98 +226,96 @@ class PatchSVAEModel(PreTrainedModel):
         h = F.gelu(self.dec_in(M_hat.reshape(B * N, -1)))
         for block in self.dec_blocks:
             h = h + block(h)
-        patches = self.dec_out(h)
-        return patches.reshape(B, N, -1)
+        return self.dec_out(h).reshape(B, N, -1)
 
-    def encode(self, images: torch.Tensor) -> torch.Tensor:
+    def encode(self, pixel_values: torch.Tensor) -> torch.Tensor:
         """Encode images to omega tokens (spatial latent).
 
         Args:
-            images: (B, 3, H, W) normalized images
+            pixel_values: (B, 3, H, W) normalized images
 
         Returns:
-            latent: (B, D, gh, gw) spectral latent — the omega tokens
-                    For 128×128 with 16×16 patches: (B, 16, 8, 8)
+            (B, D, gh, gw) spectral latent — omega tokens
+            For 128×128: (B, 16, 8, 8) = 1024 values, 48:1 compression
         """
         ps = self.config.patch_size
-        patches, gh, gw = extract_patches(images, ps)
-        svd = self._encode_patches(patches)
-        S = svd['S']  # (B, N, D)
+        patches, gh, gw = _extract_patches(pixel_values, ps)
+        svd = self._encode_patches_to_svd(patches)
+        S = svd["S"]  # (B, N, D)
         return S.permute(0, 2, 1).reshape(S.shape[0], self.config.D, gh, gw)
 
-    def encode_full(self, images: torch.Tensor) -> Dict:
-        """Encode images to full SVD decomposition per patch.
+    def encode_full(self, pixel_values: torch.Tensor) -> Dict:
+        """Encode to full SVD decomposition per patch.
 
-        Args:
-            images: (B, 3, H, W) normalized images
-
-        Returns:
-            dict with U, S_orig, S (coordinated), Vt, M per patch,
-            plus gh, gw grid dimensions.
+        Returns dict with U, S_orig, S, Vt, M, gh, gw.
         """
         ps = self.config.patch_size
-        patches, gh, gw = extract_patches(images, ps)
-        svd = self._encode_patches(patches)
-        svd['gh'] = gh
-        svd['gw'] = gw
+        patches, gh, gw = _extract_patches(pixel_values, ps)
+        svd = self._encode_patches_to_svd(patches)
+        svd["gh"] = gh
+        svd["gw"] = gw
         return svd
 
-    def decode(self, latent: torch.Tensor, U: torch.Tensor = None,
-               Vt: torch.Tensor = None) -> torch.Tensor:
-        """Decode from omega tokens back to images.
-
-        If U and Vt are provided, uses them for full reconstruction.
-        If only latent is provided, uses identity directions (lossy).
+    def decode(self, latent: torch.Tensor,
+               U: Optional[torch.Tensor] = None,
+               Vt: Optional[torch.Tensor] = None) -> torch.Tensor:
+        """Decode from omega tokens to images.
 
         Args:
             latent: (B, D, gh, gw) spectral latent
-            U: (B, N, V, D) optional left singular vectors
-            Vt: (B, N, D, D) optional right singular vectors
+            U: optional (B, N, V, D) for lossless reconstruction
+            Vt: optional (B, N, D, D) for lossless reconstruction
 
         Returns:
-            recon: (B, 3, H, W) reconstructed image
+            (B, 3, H, W) reconstructed image
         """
         B, D, gh, gw = latent.shape
         N = gh * gw
-        S = latent.reshape(B, D, N).permute(0, 2, 1)  # (B, N, D)
+        S = latent.reshape(B, D, N).permute(0, 2, 1)
 
         if U is None or Vt is None:
-            # Lossy decode: construct approximate U, Vt from identity
             V = self.config.matrix_v
-            U = torch.eye(V, D, device=latent.device).unsqueeze(0).unsqueeze(0)
-            U = U.expand(B, N, -1, -1)
-            Vt = torch.eye(D, device=latent.device).unsqueeze(0).unsqueeze(0)
-            Vt = Vt.expand(B, N, -1, -1)
-
-        decoded = self._decode_patches(U, S, Vt)
-        recon = stitch_patches(decoded, gh, gw, self.config.patch_size)
-        recon = self.boundary_smooth(recon)
-        return recon
-
-    def forward(self, pixel_values: torch.Tensor, **kwargs) -> PatchSVAEOutput:
+            U = torch.eye(V, D, device=latent.device, dtype=latent.dtype)
+            U = U.unsqueeze(0).unsqueeze(0).expand(B, N, -1, -1)
+            Vt = torch.eye(D, device=latent.device, dtype=latent.dtype)
+            Vt = Vt.unsqueeze(0).unsqueeze(0).expand(B, N, -1, -1)
+
+        decoded = self._decode_from_svd(U, S, Vt)
+        recon = _stitch_patches(decoded, gh, gw, self.config.patch_size)
+        return self.boundary_smooth(recon)
+
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        **kwargs,
+    ) -> Dict[str, torch.Tensor]:
         """Full encode → SVD → coordinate → decode pipeline.
 
         Args:
             pixel_values: (B, 3, H, W) normalized images
 
         Returns:
-            PatchSVAEOutput with recon, latent, and svd components
+            dict with "recon", "latent", "svd" keys
         """
         ps = self.config.patch_size
-        patches, gh, gw = extract_patches(pixel_values, ps)
-        svd = self._encode_patches(patches)
-        decoded = self._decode_patches(svd['U'], svd['S'], svd['Vt'])
-        recon = stitch_patches(decoded, gh, gw, ps)
+        patches, gh, gw = _extract_patches(pixel_values, ps)
+        svd = self._encode_patches_to_svd(patches)
+        decoded = self._decode_from_svd(svd["U"], svd["S"], svd["Vt"])
+        recon = _stitch_patches(decoded, gh, gw, ps)
         recon = self.boundary_smooth(recon)
 
-        # Omega tokens as spatial latent
-        S = svd['S']
+        S = svd["S"]
         latent = S.permute(0, 2, 1).reshape(S.shape[0], self.config.D, gh, gw)
 
-        return PatchSVAEOutput(recon=recon, latent=latent, svd=svd)
+        return {"recon": recon, "latent": latent, "svd": svd}
 
     @staticmethod
     def effective_rank(S):
         p = S / (S.sum(-1, keepdim=True) + 1e-8)
         p = p.clamp(min=1e-8)
-        return (-(p * p.log()).sum(-1)).exp()
+        return (-(p * p.log()).sum(-1)).exp()
+
+
+# Register for AutoClass — this is what makes AutoModel.from_pretrained work
+PatchSVAEConfig.register_for_auto_class()
+PatchSVAEModel.register_for_auto_class("AutoModel")