Create 3_johanna_model_trainer.py

3_johanna_model_trainer.py

@@ -0,0 +1,1212 @@
+"""
+Johanna F-class Miniature Trainer
+==================================
+Minimum-viable-battery research target. Sweeps small PatchSVAE
+configurations to find the floor where self-assembly breaks.
+
+Naming: johanna-F-S{img_size}-V{V}-D{D}-h{hidden}-d{depth}-p{patch}
+
+Uploads to: AbstractPhil/geolip-svae-batteries
+Structure:  {config_name}/checkpoints/epoch_NNNN.pt
+            {config_name}/tensorboard/...
+            {config_name}/config.json
+            {config_name}/final_report.json
+
+Trains from scratch (no pretrained init — that would mask failure modes).
+Gaussian-only foundation (Tier 0) for fast MSE floor discovery.
+Full 16-type dataset available via --all_types flag for battery-behavior
+verification on viable candidates.
+
+Diagnostic battery captured per report step:
+  Recon:          train_recon, test_mse (per-noise-type when all_types)
+  Geometry:       row_cv, ratio, erank, S₀, S_min, S_delta (binding)
+  Alpha:          alpha mean/std (per cross-attn layer)
+  CV health:      in 0.13-0.30 band?, proximity to target
+  Stability:      grad_norm, recon_w, prox
+  Training:       lr, epoch_time, batch_time
+
+All scalars logged to TensorBoard. Full run JSON at finish.
+"""
+
+import os
+import sys
+import json
+import math
+import time
+import argparse
+from dataclasses import dataclass, asdict, field
+from typing import Optional, List, Dict
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+try:
+    from tqdm.auto import tqdm
+    _HAS_TQDM = True
+except ImportError:
+    _HAS_TQDM = False
+
+
+# ── HuggingFace auth from Colab secrets (optional) ──────────────
+
+try:
+    from google.colab import userdata
+    os.environ["HF_TOKEN"] = userdata.get('HF_TOKEN')
+    from huggingface_hub import login
+    login(token=os.environ["HF_TOKEN"])
+except Exception:
+    pass
+
+
+# ── SVD Backend (fp64 internals, matches SVAE lineage) ──────────
+
+try:
+    from geolip_core.linalg.eigh import FLEigh, _FL_MAX_N
+    _HAS_FL = True
+except ImportError:
+    _HAS_FL = False
+
+
+def _gram_eigh_svd(A):
+    """Gram + torch.linalg.eigh in fp64 with diagonal regularization."""
+    orig_dtype = A.dtype
+    with torch.amp.autocast('cuda', enabled=False):
+        A_d = A.double()
+        G = torch.bmm(A_d.transpose(1, 2), A_d)
+        G.diagonal(dim1=-2, dim2=-1).add_(1e-12)  # conditioning regularizer
+        eigenvalues, V = torch.linalg.eigh(G)
+        eigenvalues = eigenvalues.flip(-1)
+        V = V.flip(-1)
+        S = torch.sqrt(eigenvalues.clamp(min=1e-24))
+        U = torch.bmm(A_d, V) / S.unsqueeze(1).clamp(min=1e-16)
+        Vh = V.transpose(-2, -1).contiguous()
+    return U.to(orig_dtype), S.to(orig_dtype), Vh.to(orig_dtype)
+
+
+def _svd_fp64(A):
+    """Auto-dispatch SVD with fp64 internals."""
+    B, M, N = A.shape
+    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()
+            G = torch.bmm(A_d.transpose(1, 2), A_d)
+            G.diagonal(dim1=-2, dim2=-1).add_(1e-12)
+            eigenvalues, V = FLEigh()(G.float())
+            eigenvalues = eigenvalues.double().flip(-1)
+            V = V.double().flip(-1)
+            S = torch.sqrt(eigenvalues.clamp(min=1e-24))
+            U = torch.bmm(A_d, V) / S.unsqueeze(1).clamp(min=1e-16)
+            Vh = V.transpose(-2, -1).contiguous()
+        return U.to(orig_dtype), S.to(orig_dtype), Vh.to(orig_dtype)
+    else:
+        return _gram_eigh_svd(A)
+
+
+# ── Cayley-Menger CV (fp64 determinant) ──────────────────────────
+
+def cayley_menger_vol2(points):
+    B, N, D = points.shape
+    pts = points.double()
+    gram = torch.bmm(pts, pts.transpose(1, 2))
+    norms = torch.diagonal(gram, dim1=1, dim2=2)
+    d2 = F.relu(norms.unsqueeze(2) + norms.unsqueeze(1) - 2 * gram)
+    cm = torch.zeros(B, N + 1, N + 1, device=points.device, dtype=torch.float64)
+    cm[:, 0, 1:] = 1.0
+    cm[:, 1:, 0] = 1.0
+    cm[:, 1:, 1:] = d2
+    k = N - 1
+    sign = (-1.0) ** (k + 1)
+    fact = math.factorial(k)
+    return sign * torch.linalg.det(cm) / ((2 ** k) * (fact ** 2))
+
+
+def cv_of(emb, n_samples=200):
+    """CV of pentachoron volumes for a single (V, D) embedding."""
+    if emb.dim() != 2 or emb.shape[0] < 5:
+        return 0.0
+    N, D = emb.shape
+    pool = min(N, 512)
+    indices = torch.stack([torch.randperm(pool, device=emb.device)[:5]
+                           for _ in range(n_samples)])
+    vol2 = cayley_menger_vol2(emb[:pool][indices])
+    valid = vol2 > 1e-20
+    if valid.sum() < 10:
+        return 0.0
+    vols = vol2[valid].sqrt()
+    return (vols.std() / (vols.mean() + 1e-8)).item()
+
+
+# ── Omega Noise Dataset (16 types) ───────────────────────────────
+
+class OmegaNoiseDataset(torch.utils.data.Dataset):
+    """16 noise types at arbitrary resolution, with optional type-restriction.
+
+    Args:
+        size: dataset length
+        img_size: spatial resolution
+        seed_rotate_every: re-seed every N calls (prevents epoch repetition)
+        allowed_types: iterable of int type indices, or None for all 16.
+                       For Gaussian-only foundation use allowed_types=[0].
+    """
+    N_TYPES = 16
+
+    def __init__(self, size=1_000_000, img_size=128,
+                 seed_rotate_every=1000, allowed_types=None):
+        self.size = size
+        self.img_size = img_size
+        self.seed_rotate_every = seed_rotate_every
+        self._rng = np.random.RandomState(42)
+        self._call_count = 0
+        self.allowed_types = (list(allowed_types) if allowed_types
+                              else list(range(self.N_TYPES)))
+
+    def __len__(self):
+        return self.size
+
+    def _rotate_seed(self):
+        self._call_count += 1
+        if self._call_count % self.seed_rotate_every == 0:
+            new_seed = int.from_bytes(os.urandom(4), 'big')
+            self._rng = np.random.RandomState(new_seed)
+            torch.manual_seed(new_seed)
+
+    def _pink_noise(self, shape):
+        white = torch.randn(shape)
+        S = torch.fft.rfft2(white)
+        h, w = shape[-2], shape[-1]
+        fy = torch.fft.fftfreq(h).unsqueeze(-1).expand(-1, w // 2 + 1)
+        fx = torch.fft.rfftfreq(w).unsqueeze(0).expand(h, -1)
+        f = torch.sqrt(fx**2 + fy**2).clamp(min=1e-8)
+        return torch.fft.irfft2(S / f, s=(h, w))
+
+    def _brown_noise(self, shape):
+        white = torch.randn(shape)
+        S = torch.fft.rfft2(white)
+        h, w = shape[-2], shape[-1]
+        fy = torch.fft.fftfreq(h).unsqueeze(-1).expand(-1, w // 2 + 1)
+        fx = torch.fft.rfftfreq(w).unsqueeze(0).expand(h, -1)
+        f = (fx**2 + fy**2).clamp(min=1e-8)
+        return torch.fft.irfft2(S / f, s=(h, w))
+
+    def __getitem__(self, idx):
+        self._rotate_seed()
+        s = self.img_size
+        noise_type = self.allowed_types[idx % len(self.allowed_types)]
+
+        if noise_type == 0:
+            img = torch.randn(3, s, s)
+        elif noise_type == 1:
+            img = torch.rand(3, s, s) * 2 - 1
+        elif noise_type == 2:
+            img = (torch.rand(3, s, s) - 0.5) * 4
+        elif noise_type == 3:
+            lam = self._rng.uniform(0.5, 20.0)
+            img = torch.poisson(torch.full((3, s, s), lam)) / lam - 1.0
+        elif noise_type == 4:
+            img = self._pink_noise((3, s, s))
+            img = img / (img.std() + 1e-8)
+        elif noise_type == 5:
+            img = self._brown_noise((3, s, s))
+            img = img / (img.std() + 1e-8)
+        elif noise_type == 6:
+            img = torch.where(torch.rand(3, s, s) > 0.5,
+                              torch.ones(3, s, s) * 2, torch.ones(3, s, s) * -2)
+            img = img + torch.randn(3, s, s) * 0.1
+        elif noise_type == 7:
+            mask = torch.rand(3, s, s) > 0.9
+            img = torch.randn(3, s, s) * mask.float() * 3
+        elif noise_type == 8:
+            block = self._rng.randint(2, 16)
+            small = torch.randn(3, s // block + 1, s // block + 1)
+            img = F.interpolate(small.unsqueeze(0), size=s, mode='nearest').squeeze(0)
+        elif noise_type == 9:
+            gy = torch.linspace(-2, 2, s).unsqueeze(1).expand(s, s)
+            gx = torch.linspace(-2, 2, s).unsqueeze(0).expand(s, s)
+            angle = self._rng.uniform(0, 2 * math.pi)
+            grad = math.cos(angle) * gx + math.sin(angle) * gy
+            img = grad.unsqueeze(0).expand(3, -1, -1) + torch.randn(3, s, s) * 0.5
+        elif noise_type == 10:
+            check_size = self._rng.randint(2, 16)
+            coords_y = torch.arange(s) // check_size
+            coords_x = torch.arange(s) // check_size
+            checker = ((coords_y.unsqueeze(1) + coords_x.unsqueeze(0)) % 2).float() * 2 - 1
+            img = checker.unsqueeze(0).expand(3, -1, -1) + torch.randn(3, s, s) * 0.3
+        elif noise_type == 11:
+            a = torch.randn(3, s, s)
+            b = torch.rand(3, s, s) * 2 - 1
+            alpha = self._rng.uniform(0.2, 0.8)
+            img = alpha * a + (1 - alpha) * b
+        elif noise_type == 12:
+            img = torch.zeros(3, s, s)
+            h2, w2 = s // 2, s // 2
+            img[:, :h2, :w2] = torch.randn(3, h2, w2)
+            img[:, :h2, w2:] = torch.rand(3, h2, w2) * 2 - 1
+            img[:, h2:, :w2] = self._pink_noise((3, h2, w2)) / 2
+            sp = torch.where(torch.rand(3, h2, w2) > 0.5,
+                             torch.ones(3, h2, w2), -torch.ones(3, h2, w2))
+            img[:, h2:, w2:] = sp
+        elif noise_type == 13:
+            u = torch.rand(3, s, s)
+            img = torch.tan(math.pi * (u - 0.5)).clamp(-3, 3)
+        elif noise_type == 14:
+            img = torch.empty(3, s, s).exponential_(1.0) - 1.0
+        elif noise_type == 15:
+            u = torch.rand(3, s, s) - 0.5
+            img = -torch.sign(u) * torch.log1p(-2 * u.abs())
+        else:
+            raise ValueError(f"Unknown noise type {noise_type}")
+
+        return img.clamp(-4, 4).float(), noise_type
+
+
+# ── Patch Utils + Small Components ───────────────────────────────
+
+def extract_patches(images, patch_size):
+    B, C, H, W = images.shape
+    gh, gw = H // patch_size, W // patch_size
+    p = images.reshape(B, C, gh, patch_size, gw, patch_size)
+    p = p.permute(0, 2, 4, 1, 3, 5)
+    return p.reshape(B, gh * gw, C * patch_size * patch_size), gh, gw
+
+
+def stitch_patches(patches, gh, gw, patch_size):
+    B = patches.shape[0]
+    p = patches.reshape(B, gh, gw, 3, patch_size, patch_size)
+    p = p.permute(0, 3, 1, 4, 2, 5)
+    return p.reshape(B, 3, gh * patch_size, gw * patch_size)
+
+
+class BoundarySmooth(nn.Module):
+    def __init__(self, channels=3, mid=16):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Conv2d(channels, mid, 3, padding=1), nn.GELU(),
+            nn.Conv2d(mid, channels, 3, padding=1))
+        nn.init.zeros_(self.net[-1].weight)
+        nn.init.zeros_(self.net[-1].bias)
+
+    def forward(self, x):
+        return x + self.net(x)
+
+
+class SpectralCrossAttention(nn.Module):
+    """Multiplicative cross-attention on S vectors with bounded alpha."""
+    def __init__(self, D, n_heads=4, max_alpha=0.2, alpha_init=-2.0):
+        super().__init__()
+        heads = min(n_heads, D) if D % n_heads != 0 else n_heads
+        # ensure divisibility
+        while D % heads != 0 and heads > 1:
+            heads -= 1
+        self.n_heads = heads
+        self.head_dim = D // heads
+        self.max_alpha = max_alpha
+        self.qkv = nn.Linear(D, 3 * D)
+        self.out_proj = nn.Linear(D, D)
+        self.norm = nn.LayerNorm(D)
+        self.scale = self.head_dim ** -0.5
+        self.alpha_logits = nn.Parameter(torch.full((D,), alpha_init))
+
+    @property
+    def alpha(self):
+        return self.max_alpha * torch.sigmoid(self.alpha_logits)
+
+    def forward(self, S):
+        B, N, D = S.shape
+        S_normed = self.norm(S)
+        qkv = self.qkv(S_normed).reshape(B, N, 3, self.n_heads, self.head_dim)
+        qkv = qkv.permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        attn = attn.softmax(dim=-1)
+        out = (attn @ v).transpose(1, 2).reshape(B, N, D)
+        gate = torch.tanh(self.out_proj(out))
+        return S * (1.0 + self.alpha.unsqueeze(0).unsqueeze(0) * gate)
+
+
+# ── PatchSVAE-F (miniature) ──────────────────────────────────────
+
+class PatchSVAE_F(nn.Module):
+    """F-class miniature PatchSVAE.
+
+    Architecture matches Fresnel/Johanna reference verbatim, sized down.
+    All defense-stack mechanisms preserved:
+      - Sphere-normalized rows (F.normalize after enc_out)
+      - fp64 SVD with 1e-24 / 1e-16 floors and 1e-12 diag reg
+      - Orthogonal init on enc_out
+      - Multiplicative cross-attention with bounded alpha (max 0.2)
+      - Zero-initialized boundary smoothing
+      - No BatchNorm, no Dropout
+    """
+    def __init__(self, matrix_v=64, D=8, patch_size=16, hidden=128,
+                 depth=1, n_cross_layers=1, n_heads=4,
+                 max_alpha=0.2, alpha_init=-2.0):
+        super().__init__()
+        self.matrix_v = matrix_v
+        self.D = D
+        self.patch_size = patch_size
+        self.patch_dim = 3 * patch_size * patch_size
+        self.mat_dim = matrix_v * D
+        self.hidden = hidden
+        self.depth = depth
+        self.n_cross_layers = n_cross_layers
+
+        self.enc_in = nn.Linear(self.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)
+        nn.init.orthogonal_(self.enc_out.weight)
+
+        self.dec_in = nn.Linear(self.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.cross_attn = nn.ModuleList([
+            SpectralCrossAttention(D, n_heads=n_heads,
+                                   max_alpha=max_alpha, alpha_init=alpha_init)
+            for _ in range(n_cross_layers)])
+        self.boundary_smooth = BoundarySmooth(channels=3, mid=16)
+
+    def encode_patches(self, patches):
+        B, N, _ = patches.shape
+        flat = patches.reshape(B * N, -1)
+        h = F.gelu(self.enc_in(flat))
+        for block in self.enc_blocks:
+            h = h + block(h)
+        M = self.enc_out(h).reshape(B * N, self.matrix_v, self.D)
+        M = F.normalize(M, dim=-1)
+        U, S, Vt = _svd_fp64(M)
+        U = U.reshape(B, N, self.matrix_v, self.D)
+        S = S.reshape(B, N, self.D)
+        Vt = Vt.reshape(B, N, self.D, self.D)
+        M = M.reshape(B, N, self.matrix_v, self.D)
+        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}
+
+    def decode_patches(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)
+        M_hat = torch.bmm(U_flat * S_flat.unsqueeze(1), Vt_flat)
+        h = F.gelu(self.dec_in(M_hat.reshape(B * N, -1)))
+        for block in self.dec_blocks:
+            h = h + block(h)
+        return self.dec_out(h).reshape(B, N, -1)
+
+    def forward(self, images):
+        patches, gh, gw = extract_patches(images, self.patch_size)
+        svd = self.encode_patches(patches)
+        decoded = self.decode_patches(svd['U'], svd['S'], svd['Vt'])
+        recon = stitch_patches(decoded, gh, gw, self.patch_size)
+        recon = self.boundary_smooth(recon)
+        return {'recon': recon, 'svd': svd, 'gh': gh, 'gw': gw}
+
+    @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()
+
+
+# ── Config dataclass ─────────────────────────────────────────────
+
+@dataclass
+class RunConfig:
+    """F-class run configuration.
+
+    Naming convention: johanna-F-S{img_size}-V{V}-D{D}-h{hidden}-d{depth}-p{patch}
+    """
+    # Architecture (the sweep axes)
+    matrix_v: int = 64
+    D: int = 8
+    patch_size: int = 16
+    hidden: int = 128
+    depth: int = 1
+    n_cross_layers: int = 1
+    n_heads: int = 4
+    max_alpha: float = 0.2
+    alpha_init: float = -2.0
+
+    # Training
+    img_size: int = 128
+    batch_size: int = 128
+    lr: float = 1e-3
+    epochs: int = 20
+    weight_decay: float = 0.0   # Phil: always pure Adam
+
+    # Loss / soft-hand
+    # Soft-hand guides against CV-EMA — geometric coherence signal.
+    # CV (Cayley-Menger pentachoron volume CV) measures whether the
+    # sphere-normalized rows are arranged with geometric consistency.
+    # We don't force CV toward a specific value — we track its own EMA
+    # and reward CV being near its own trajectory (auto-attenuating).
+    #
+    # This is the correct reading from Phil's SVAE lineage:
+    #   - CV tells us the geometry is COHERENT (relational guidepost)
+    #   - Recon MSE tells us the arrangement is VALID (reversible)
+    #   - Both are needed; recon alone doesn't guarantee geometric structure
+    use_cv_ema: bool = True
+    cv_ema_alpha: float = 0.01         # EMA smoothing (slow is better)
+    cv_alignment_epochs: int = 1       # Pure-MSE epochs before soft-hand activates
+    cv_measure_every: int = 25         # Measure CV every N batches
+    cv_sigma_scale: float = 0.3        # Proximity width = sigma_scale × cv_ema
+    boost: float = 0.5                 # Max recon-weight boost when CV near EMA
+    cross_attn_clip: float = 0.5
+
+    # Data
+    allowed_types: List[int] = field(default_factory=lambda: [0])  # Gaussian only
+    train_size: int = 100_000
+    val_size: int = 2000
+    num_workers: int = 4
+
+    # Reporting
+    report_every: int = 500          # TB log cadence (batches)
+    major_report_every: int = 10     # Console major-report cadence (epochs)
+    save_every: int = 5              # Checkpoint save cadence (epochs)
+    seed: int = 42
+
+    # Upload
+    hf_repo: str = "AbstractPhil/geolip-svae-batteries"
+    upload: bool = True
+
+    def name(self) -> str:
+        return (f"johanna-F-S{self.img_size}-V{self.matrix_v}"
+                f"-D{self.D}-h{self.hidden}-d{self.depth}"
+                f"-p{self.patch_size}")
+
+
+# ── Training runner ──────────────────────────────────────────────
+
+def run(cfg: RunConfig, out_root: str = "/content/johanna_F_runs"):
+    torch.manual_seed(cfg.seed)
+    np.random.seed(cfg.seed)
+    torch.set_float32_matmul_precision('high')
+
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    run_name = cfg.name()
+    run_dir = os.path.join(out_root, run_name)
+    ckpt_dir = os.path.join(run_dir, "checkpoints")
+    tb_dir = os.path.join(run_dir, "tensorboard")
+    os.makedirs(ckpt_dir, exist_ok=True)
+    os.makedirs(tb_dir, exist_ok=True)
+
+    # Save config snapshot
+    with open(os.path.join(run_dir, "config.json"), 'w') as f:
+        json.dump(asdict(cfg), f, indent=2)
+
+    # Model
+    model = PatchSVAE_F(
+        matrix_v=cfg.matrix_v, D=cfg.D, patch_size=cfg.patch_size,
+        hidden=cfg.hidden, depth=cfg.depth,
+        n_cross_layers=cfg.n_cross_layers, n_heads=cfg.n_heads,
+        max_alpha=cfg.max_alpha, alpha_init=cfg.alpha_init,
+    ).to(device)
+
+    n_params = sum(p.numel() for p in model.parameters())
+    n_patches = (cfg.img_size // cfg.patch_size) ** 2
+
+    # TensorBoard
+    from torch.utils.tensorboard import SummaryWriter
+    writer = SummaryWriter(tb_dir)
+
+    # HF
+    hf_enabled = False
+    api = None
+    if cfg.upload:
+        try:
+            from huggingface_hub import HfApi
+            api = HfApi()
+            api.whoami()
+            hf_enabled = True
+        except Exception as e:
+            print(f"  HF upload disabled: {e}")
+
+    def upload_file(local_path, remote_name):
+        if not hf_enabled:
+            return
+        try:
+            api.upload_file(
+                path_or_fileobj=local_path,
+                path_in_repo=f"{run_name}/{remote_name}",
+                repo_id=cfg.hf_repo, repo_type="model")
+        except Exception as e:
+            print(f"  HF upload failed ({remote_name}): {e}")
+
+    def upload_folder(local_dir, remote_prefix):
+        if not hf_enabled:
+            return
+        try:
+            api.upload_folder(
+                folder_path=local_dir,
+                path_in_repo=f"{run_name}/{remote_prefix}",
+                repo_id=cfg.hf_repo, repo_type="model")
+        except Exception as e:
+            print(f"  HF folder upload failed ({remote_prefix}): {e}")
+
+    # Datasets
+    train_ds = OmegaNoiseDataset(size=cfg.train_size, img_size=cfg.img_size,
+                                 allowed_types=cfg.allowed_types)
+    val_ds = OmegaNoiseDataset(size=cfg.val_size, img_size=cfg.img_size,
+                               allowed_types=cfg.allowed_types)
+    train_loader = torch.utils.data.DataLoader(
+        train_ds, batch_size=cfg.batch_size, shuffle=True,
+        num_workers=cfg.num_workers, pin_memory=True, drop_last=True)
+    test_loader = torch.utils.data.DataLoader(
+        val_ds, batch_size=cfg.batch_size, shuffle=False,
+        num_workers=cfg.num_workers, pin_memory=True)
+
+    # Optimizer (pure Adam per Phil preference)
+    opt = torch.optim.Adam(model.parameters(), lr=cfg.lr,
+                           weight_decay=cfg.weight_decay)
+    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=cfg.epochs)
+
+    # Header
+    print("=" * 100)
+    print(f"RUN: {run_name}")
+    print(f"  V={cfg.matrix_v}, D={cfg.D}, hidden={cfg.hidden}, depth={cfg.depth}, "
+          f"patch={cfg.patch_size}, n_cross={cfg.n_cross_layers}")
+    print(f"  patches/image={n_patches}, params={n_params:,}, batch={cfg.batch_size}")
+    print(f"  types={cfg.allowed_types}, epochs={cfg.epochs}, lr={cfg.lr}")
+    print(f"  HF repo={cfg.hf_repo}, upload={hf_enabled}")
+    print("=" * 100)
+    print(f" {'ep':>3} {'batch':>7} | {'loss':>7} {'recon':>7} {'test':>7} {'r_obs':>6} | "
+          f"{'S0':>6} {'SD':>6} {'ratio':>5} {'erank':>5} | "
+          f"{'cv':>6} {'cv_ema':>6} | "
+          f"{'prox':>5} {'rw':>4} {'S_del':>6} {'a_m':>5} | "
+          f"{'grad':>7} {'ph':>4}")
+    print("-" * 130)
+
+    best_mse = float('inf')
+    global_batch = 0
+    last_cv = 0.0
+    cv_ema = None            # the guidance signal — CV's own trajectory
+    recon_ema_obs = None     # observable only, used for logging
+    last_prox = 1.0
+    in_alignment = cfg.use_cv_ema
+    history = []
+
+    # Save initial config upload (creates repo structure early)
+    if hf_enabled:
+        upload_file(os.path.join(run_dir, "config.json"), "config.json")
+
+    for epoch in range(1, cfg.epochs + 1):
+        model.train()
+        tot_loss = tot_recon = n_seen = 0
+        epoch_max_grad = 0.0
+        t0 = time.time()
+
+        # Alignment phase: pure MSE, CV-EMA accumulates
+        if cfg.use_cv_ema:
+            in_alignment = epoch <= cfg.cv_alignment_epochs
+
+        # Progress bar for this epoch's batches
+        if _HAS_TQDM:
+            batch_iter = tqdm(train_loader, desc=f"ep {epoch:3d}/{cfg.epochs}",
+                              leave=False, dynamic_ncols=True,
+                              bar_format='{l_bar}{bar:20}{r_bar}')
+        else:
+            batch_iter = train_loader
+
+        for batch_idx, (images, _) in enumerate(batch_iter):
+            images = images.to(device, non_blocking=True)
+            opt.zero_grad()
+            out = model(images)
+            recon_loss = F.mse_loss(out['recon'], images)
+            recon_val = recon_loss.item()
+
+            with torch.no_grad():
+                # Track recon_ema as observable (for logging only)
+                if recon_ema_obs is None:
+                    recon_ema_obs = recon_val
+                else:
+                    recon_ema_obs = 0.99 * recon_ema_obs + 0.01 * recon_val
+
+                # Measure CV and update CV-EMA (the GUIDANCE signal)
+                if batch_idx % cfg.cv_measure_every == 0:
+                    current_cv = cv_of(out['svd']['M'][0, 0])
+                    if current_cv > 0:
+                        last_cv = current_cv
+                        if cv_ema is None:
+                            cv_ema = current_cv
+                        else:
+                            cv_ema = ((1.0 - cfg.cv_ema_alpha) * cv_ema
+                                      + cfg.cv_ema_alpha * current_cv)
+
+                    # Soft-hand proximity: is CURRENT CV near the EMA trajectory?
+                    # This measures geometric coherence — whether the arrangement
+                    # is settling (CV near its trend) or thrashing (CV deviating).
+                    if cv_ema is not None and cv_ema > 1e-6:
+                        sigma_adapt = max(cfg.cv_sigma_scale * cv_ema, 1e-6)
+                        delta = last_cv - cv_ema
+                        last_prox = math.exp(-(delta ** 2) / (2 * sigma_adapt ** 2))
+
+            # Soft-hand: boost recon weight when geometry is coherent.
+            # No penalty term — we never fight the model toward a specific CV.
+            # We only reward the recon path when geometry is relationally stable.
+            if in_alignment:
+                loss = recon_loss
+                recon_w = 1.0
+            else:
+                recon_w = 1.0 + cfg.boost * last_prox
+                loss = recon_w * recon_loss
+
+            loss.backward()
+
+            # Clip cross-attn only (rest stays free per SVAE recipe)
+            torch.nn.utils.clip_grad_norm_(
+                model.cross_attn.parameters(), max_norm=cfg.cross_attn_clip)
+
+            # Measure total grad norm for stability tracking
+            total_grad = sum(
+                p.grad.pow(2).sum().item()
+                for p in model.parameters() if p.grad is not None
+            ) ** 0.5
+            epoch_max_grad = max(epoch_max_grad, total_grad)
+
+            opt.step()
+
+            tot_loss += loss.item() * images.size(0)
+            tot_recon += recon_val * images.size(0)
+            n_seen += images.size(0)
+            global_batch += 1
+
+            # Update tqdm postfix with live stats
+            if _HAS_TQDM:
+                cvema_s = f"{cv_ema:.3f}" if cv_ema is not None else "---"
+                batch_iter.set_postfix_str(
+                    f"r={recon_val:.4f} cv={last_cv:.3f} ema={cvema_s} "
+                    f"px={last_prox:.2f} rw={recon_w:.2f}", refresh=False)
+
+            # Periodic TB + history snapshot (silent — no console print at batch level)
+            if global_batch % cfg.report_every == 0:
+                model.eval()
+                with torch.no_grad():
+                    test_imgs, _ = next(iter(test_loader))
+                    test_imgs = test_imgs.to(device)
+                    t_out = model(test_imgs)
+                    test_mse = F.mse_loss(t_out['recon'], test_imgs).item()
+
+                    # Geometry
+                    S_batch = t_out['svd']['S']            # (B, N, D)
+                    S_orig = t_out['svd']['S_orig']
+                    S_mean = S_batch.mean(dim=(0, 1))      # (D,)
+                    S0 = S_mean[0].item()
+                    SD = S_mean[-1].item()
+                    ratio = S0 / (SD + 1e-8)
+                    erank = PatchSVAE_F.effective_rank(
+                        S_batch.reshape(-1, cfg.D)).mean().item()
+                    s_delta = (S_batch - S_orig).abs().mean().item()
+
+                    # Alpha across cross-attn layers
+                    if cfg.n_cross_layers > 0:
+                        alphas = [layer.alpha.detach() for layer in model.cross_attn]
+                        alpha_mean = torch.stack([a.mean() for a in alphas]).mean().item()
+                        alpha_std = torch.stack([a.std() for a in alphas]).mean().item()
+                    else:
+                        alpha_mean = 0.0
+                        alpha_std = 0.0
+
+                    # CV band check
+                    cv_in_band = 0.13 <= last_cv <= 0.30
+
+                # TB logs (always)
+                writer.add_scalar('train/loss', tot_loss / n_seen, global_batch)
+                writer.add_scalar('train/recon', tot_recon / n_seen, global_batch)
+                writer.add_scalar('test/mse', test_mse, global_batch)
+                if recon_ema_obs is not None:
+                    writer.add_scalar('stab/recon_ema_obs', recon_ema_obs, global_batch)
+                writer.add_scalar('geo/S0', S0, global_batch)
+                writer.add_scalar('geo/SD', SD, global_batch)
+                writer.add_scalar('geo/ratio', ratio, global_batch)
+                writer.add_scalar('geo/erank', erank, global_batch)
+                writer.add_scalar('geo/row_cv', last_cv, global_batch)
+                if cv_ema is not None:
+                    writer.add_scalar('geo/cv_ema', cv_ema, global_batch)
+                writer.add_scalar('geo/cv_in_band', float(cv_in_band), global_batch)
+                writer.add_scalar('geo/S_delta', s_delta, global_batch)
+                writer.add_scalar('cross_attn/alpha_mean', alpha_mean, global_batch)
+                writer.add_scalar('cross_attn/alpha_std', alpha_std, global_batch)
+                writer.add_scalar('stab/prox', last_prox, global_batch)
+                writer.add_scalar('stab/recon_w', recon_w, global_batch)
+                writer.add_scalar('stab/epoch_max_grad', epoch_max_grad, global_batch)
+                writer.add_scalar('stab/lr', opt.param_groups[0]['lr'], global_batch)
+                writer.add_scalar('stab/in_alignment', float(in_alignment), global_batch)
+
+                history.append({
+                    'epoch': epoch, 'global_batch': global_batch,
+                    'train_recon': tot_recon / n_seen,
+                    'test_mse': test_mse,
+                    'recon_ema_obs': recon_ema_obs,
+                    'S0': S0, 'SD': SD, 'ratio': ratio, 'erank': erank,
+                    'row_cv': last_cv, 'cv_ema': cv_ema,
+                    'cv_in_band': cv_in_band,
+                    'S_delta': s_delta,
+                    'alpha_mean': alpha_mean, 'alpha_std': alpha_std,
+                    'grad_max': epoch_max_grad,
+                    'in_alignment': in_alignment,
+                    'prox': last_prox, 'recon_w': recon_w,
+                })
+
+                if test_mse < best_mse:
+                    best_mse = test_mse
+                    # Save best silently (upload only at save_every)
+                    torch.save({
+                        'epoch': epoch, 'test_mse': test_mse,
+                        'global_batch': global_batch,
+                        'model_state_dict': model.state_dict(),
+                        'config': asdict(cfg),
+                    }, os.path.join(ckpt_dir, 'best.pt'))
+
+                model.train()
+
+        if _HAS_TQDM:
+            batch_iter.close()
+
+        sched.step()
+        epoch_time = time.time() - t0
+
+        # Full-eval at epoch boundary
+        model.eval()
+        tot_test = test_n = 0
+        with torch.no_grad():
+            for t_imgs, _ in test_loader:
+                t_imgs = t_imgs.to(device)
+                t_out = model(t_imgs)
+                tot_test += F.mse_loss(t_out['recon'], t_imgs).item() * t_imgs.size(0)
+                test_n += t_imgs.size(0)
+        epoch_test_mse = tot_test / test_n
+        writer.add_scalar('epoch/test_mse', epoch_test_mse, epoch)
+        writer.add_scalar('epoch/time_s', epoch_time, epoch)
+        writer.add_scalar('epoch/max_grad', epoch_max_grad, epoch)
+
+        if epoch_test_mse < best_mse:
+            best_mse = epoch_test_mse
+            torch.save({
+                'epoch': epoch, 'test_mse': epoch_test_mse,
+                'global_batch': global_batch,
+                'model_state_dict': model.state_dict(),
+                'config': asdict(cfg),
+            }, os.path.join(ckpt_dir, 'best.pt'))
+
+        # MAJOR REPORT: every 10 epochs + first + last
+        is_major_report = (
+            epoch == 1 or epoch == cfg.epochs or
+            epoch % cfg.major_report_every == 0
+        )
+        if is_major_report:
+            cvema_s = f"{cv_ema:.4f}" if cv_ema is not None else "---"
+            rema_s = f"{recon_ema_obs:.4f}" if recon_ema_obs is not None else "---"
+            print(f"  ep {epoch:3d}/{cfg.epochs}: "
+                  f"test_mse={epoch_test_mse:.6f} best={best_mse:.6f} "
+                  f"| cv_ema={cvema_s} recon_ema={rema_s} "
+                  f"| max_grad={epoch_max_grad:.2f} "
+                  f"| {epoch_time:.1f}s"
+                  f"{'  [ALGN]' if in_alignment else '  [HAND]'}")
+        else:
+            # Minimal per-epoch print (one line)
+            print(f"  ep {epoch:3d}: mse={epoch_test_mse:.4f} "
+                  f"grad={epoch_max_grad:.1f} {epoch_time:.1f}s")
+
+        # Periodic save + upload
+        if epoch % cfg.save_every == 0 or epoch == cfg.epochs:
+            ep_path = os.path.join(ckpt_dir, f"epoch_{epoch:04d}.pt")
+            torch.save({
+                'epoch': epoch, 'test_mse': epoch_test_mse,
+                'global_batch': global_batch,
+                'model_state_dict': model.state_dict(),
+                'optimizer_state_dict': opt.state_dict(),
+                'scheduler_state_dict': sched.state_dict(),
+                'config': asdict(cfg),
+            }, ep_path)
+            writer.flush()
+
+            if hf_enabled:
+                upload_file(ep_path, f"checkpoints/epoch_{epoch:04d}.pt")
+                best_path = os.path.join(ckpt_dir, 'best.pt')
+                if os.path.exists(best_path):
+                    upload_file(best_path, "checkpoints/best.pt")
+                upload_folder(tb_dir, "tensorboard")
+
+    writer.close()
+
+    # Final report
+    final = {
+        'run_name': run_name,
+        'config': asdict(cfg),
+        'n_params': n_params,
+        'n_patches': n_patches,
+        'best_test_mse': best_mse,
+        'final_epoch_mse': epoch_test_mse,
+        'final_cv_ema': cv_ema,
+        'final_recon_ema_obs': recon_ema_obs,
+        'final_S0': history[-1]['S0'] if history else None,
+        'final_erank': history[-1]['erank'] if history else None,
+        'final_row_cv': history[-1]['row_cv'] if history else None,
+        'final_cv_in_band': history[-1]['cv_in_band'] if history else None,
+        'final_S_delta': history[-1]['S_delta'] if history else None,
+        'final_alpha_mean': history[-1]['alpha_mean'] if history else None,
+        'history': history,
+    }
+    report_path = os.path.join(run_dir, "final_report.json")
+    with open(report_path, 'w') as f:
+        json.dump(final, f, indent=2)
+    if hf_enabled:
+        upload_file(report_path, "final_report.json")
+
+    print(f"\n  RUN COMPLETE: {run_name}")
+    print(f"  Best test MSE: {best_mse:.6f}")
+    cvema = f"{cv_ema:.4f}" if cv_ema is not None else "n/a"
+    rema = f"{recon_ema_obs:.4f}" if recon_ema_obs is not None else "n/a"
+    print(f"  Final: S0={final['final_S0']:.3f}, erank={final['final_erank']:.3f}, "
+          f"cv={final['final_row_cv']:.4f} (cv_ema={cvema}), recon_ema_obs={rema}")
+    print(f"  Checkpoints: {ckpt_dir}")
+    print(f"  TensorBoard: {tb_dir}")
+    print(f"  Report: {report_path}")
+    return final
+
+
+# ── Smoke test entry point ───────────────────────────────────────
+
+def smoke():
+    """Minimum viable smoke test — tiny config, 3 epochs, Gaussian only.
+
+    Soft-hand guides against CV-EMA (geometric coherence signal).
+      Epoch 1: alignment (pure MSE, cv_ema accumulates)
+      Epochs 2-3: hand active (proximity against cv_ema)
+    """
+    cfg = RunConfig(
+        matrix_v=32, D=4, patch_size=16, hidden=64, depth=1,
+        n_cross_layers=1, n_heads=2,
+        img_size=64,
+        batch_size=64,
+        lr=1e-3,
+        epochs=3,
+        allowed_types=[0],
+        train_size=5000,
+        val_size=500,
+        num_workers=2,
+        report_every=20,
+        save_every=1,
+        use_cv_ema=True,
+        cv_alignment_epochs=1,
+        cv_ema_alpha=0.05,      # faster EMA for short smoke
+        cv_sigma_scale=0.3,
+        upload=True,
+    )
+    return run(cfg)
+
+
+# ── F-class sweep ───────────────────────────────────────────────
+
+def _johanna_base_cfg(**overrides):
+    """Base config for F-class (miniature) battery experiments.
+
+    Defaults are genuinely small — 16 patches per image, modest V/D,
+    thin substrate. Override any axis you're sweeping.
+
+    Reasonable overridable axes:
+        matrix_v, D, patch_size, hidden, depth, n_cross_layers,
+        img_size, batch_size, lr, epochs, allowed_types
+    """
+    base = dict(
+        # F-class defaults (small, stackable)
+        matrix_v=64, D=8, patch_size=16,
+        n_cross_layers=1, n_heads=4,
+        max_alpha=0.2, alpha_init=-2.0,
+
+        # Training conventions from Johanna lineage
+        img_size=64,
+        batch_size=128,
+        lr=1e-4,
+        epochs=30,
+        weight_decay=0.0,
+
+        allowed_types=list(range(16)),
+        train_size=1_280_000,
+        val_size=10_000,
+        num_workers=4,
+
+        report_every=500,
+        save_every=5,
+
+        # Soft-hand on CV-EMA (geometric coherence guidepost)
+        use_cv_ema=True,
+        cv_ema_alpha=0.01,
+        cv_alignment_epochs=2,
+        cv_sigma_scale=0.3,
+        cv_measure_every=50,
+        boost=0.5,
+        cross_attn_clip=0.5,
+
+        upload=True,
+    )
+    base.update(overrides)
+    return RunConfig(**base)
+
+
+SWEEP_F_CLASS = [
+    # ═══════════════════════════════════════════════════════════════════
+    #  F-class sweep: the experimental nursery.
+    #
+    #  F-class models are NOT expected to succeed. They are research
+    #  specimens — most will collapse, some will barely function,
+    #  a few might surprise us. The sweep's job is to catalog failure
+    #  modes at small scale, not to find a winning config.
+    #
+    #  A-class = Johanna/Fresnel (17M, teachable workhorses)
+    #  S-class = Freckles (2.5M, superior recon, too disorderly to teach)
+    #  F-class = this sweep (0.03M-0.6M, expected to fail)
+    #
+    #  Axes varied deliberately:
+    #    - TINY overall (V,D,hidden,depth all small)
+    #    - SMALL patchworks (few patches per image — less info to work with)
+    #    - LARGE patchworks with SMALL internals (many patches + weak substrate)
+    #    - UNUSUAL shapes (V >> D, D >> V equivalents, unusual depth ratios)
+    #
+    #  Naming: johanna-F-S{S}-V{V}-D{D}-h{hidden}-d{depth}-p{patch}
+    #  All runs log to HF `AbstractPhil/geolip-svae-batteries/<run_name>/`
+    # ═══════════════════════════════════════════════════════════════════
+
+    # ── TIER 1: D=16 spine (Johanna's universal-attractor dim, small scale) ──
+    # Question: does D=16 carry through when everything around it shrinks?
+    _johanna_base_cfg(img_size=64,  matrix_v=64,  D=16, hidden=64,  depth=1, patch_size=16),   # 16 patches, ~250K
+    _johanna_base_cfg(img_size=64,  matrix_v=64,  D=16, hidden=64,  depth=1, patch_size=8),    # 64 patches, ~176K
+    _johanna_base_cfg(img_size=64,  matrix_v=128, D=16, hidden=128, depth=1, patch_size=8),    # larger V+h, ~645K
+
+    # ── TIER 2: D-sweep at matched substrate (V=64, h=64, d=1, p=16) ──
+    # Does battery behavior survive D<16? (answers an unasked question of omega paper)
+    _johanna_base_cfg(img_size=64,  matrix_v=64,  D=8,  hidden=64,  depth=1, patch_size=16),   # D=8
+    _johanna_base_cfg(img_size=64,  matrix_v=64,  D=4,  hidden=64,  depth=1, patch_size=16),   # D=4
+    _johanna_base_cfg(img_size=64,  matrix_v=64,  D=2,  hidden=64,  depth=1, patch_size=16),   # D=2 — likely collapse
+
+    # ── TIER 3: Substrate axis at (V=64, D=8) ──
+    # Which axis carries the self-assembly work — width or depth?
+    _johanna_base_cfg(img_size=64,  matrix_v=64,  D=8,  hidden=128, depth=1, patch_size=16),   # wider substrate
+    _johanna_base_cfg(img_size=64,  matrix_v=64,  D=8,  hidden=64,  depth=2, patch_size=16),   # deeper substrate
+    _johanna_base_cfg(img_size=64,  matrix_v=64,  D=8,  hidden=32,  depth=1, patch_size=16),   # starved width
+
+    # ── TIER 4: Patchwork sweeps — big patchworks with small internals ──
+    # Many weak cells vs few strong cells. Tests the "cells are SVAE-shaped
+    # functions, not batteries" framing.
+    _johanna_base_cfg(img_size=64,  matrix_v=32,  D=4,  hidden=32,  depth=1, patch_size=4),    # 256 patches, tiny internals
+    _johanna_base_cfg(img_size=64,  matrix_v=32,  D=4,  hidden=64,  depth=1, patch_size=4),    # 256 patches, slightly more substrate
+    _johanna_base_cfg(img_size=32,  matrix_v=32,  D=4,  hidden=32,  depth=1, patch_size=2, batch_size=256),  # 256 patches @ S=32
+
+    # ── TIER 5: Small patchworks (few patches, stronger cells) ──
+    # Inverse: few big cells. Easier for a wrapper to channel, but each
+    # cell carries more load.
+    _johanna_base_cfg(img_size=32,  matrix_v=64,  D=8,  hidden=128, depth=1, patch_size=16, batch_size=256),  # 4 patches
+    _johanna_base_cfg(img_size=32,  matrix_v=64,  D=8,  hidden=64,  depth=1, patch_size=16, batch_size=256),  # 4 patches thinner
+
+    # ── TIER 6: Unusual shapes (chaos measurements) ──
+    # V >> D and D >> V ratios well outside Johanna's 16:1 ratio.
+    _johanna_base_cfg(img_size=64,  matrix_v=256, D=2,  hidden=64,  depth=1, patch_size=16),   # V=128×D ratio
+    _johanna_base_cfg(img_size=64,  matrix_v=16,  D=16, hidden=64,  depth=1, patch_size=16),   # V=D (square matrix)
+    _johanna_base_cfg(img_size=64,  matrix_v=8,   D=16, hidden=64,  depth=1, patch_size=16),   # V<D (wide matrix)
+
+    # ── TIER 7: Extreme smallness (almost-certain collapse) ──
+    # Lower bound of what can be built. Mostly exists as a "what does total
+    # collapse look like" reference.
+    _johanna_base_cfg(img_size=32,  matrix_v=16,  D=4,  hidden=16,  depth=1, patch_size=8, batch_size=256),   # ~20K params
+    _johanna_base_cfg(img_size=16,  matrix_v=8,   D=2,  hidden=8,   depth=1, patch_size=4, batch_size=256),   # absurdly small
+]
+
+
+# ── Run existence check (for auto-resume across sessions) ────────
+
+def _hf_run_exists(run_name: str, hf_repo: str) -> bool:
+    """Check if this run has already been completed on HF.
+    A run is considered complete if `<run_name>/final_report.json` exists.
+    """
+    try:
+        from huggingface_hub import HfApi
+        api = HfApi()
+        # list_repo_files returns paths in the repo. Check for final_report.json.
+        files = api.list_repo_files(repo_id=hf_repo, repo_type="model")
+        marker = f"{run_name}/final_report.json"
+        return marker in files
+    except Exception as e:
+        # If we can't check (auth, network, etc), assume not complete and run.
+        print(f"[HF-check] could not verify completion of {run_name}: {e}")
+        return False
+
+
+def sweep(configs=None, out_root="/content/johanna_F_runs",
+          skip_on_error=True, skip_completed=True):
+    """Run a list of RunConfigs sequentially. One per config, full training each.
+
+    Args:
+        configs: List of RunConfig. Defaults to SWEEP_F_CLASS.
+        out_root: Local output directory.
+        skip_on_error: If True, log the error and continue to next config.
+        skip_completed: If True, skip configs whose final_report.json is
+            already on HF (for auto-resume across sessions).
+    """
+    if configs is None:
+        configs = SWEEP_F_CLASS
+
+    print("\n" + "#" * 100)
+    print(f"# F-CLASS SWEEP: {len(configs)} configurations")
+    print(f"#   Most will collapse. This is expected. Collapse is a data point.")
+    print(f"#   Results go to HF: {configs[0].hf_repo if configs else '(n/a)'}")
+    print(f"#   Use a separate readout cell to pool JSON metrics afterward.")
+    print("#" * 100)
+
+    # Pre-flight: check which runs already exist on HF
+    completed = set()
+    if skip_completed and configs and configs[0].upload:
+        print("\n[preflight] checking HF for already-completed runs...")
+        for cfg in configs:
+            if _hf_run_exists(cfg.name(), cfg.hf_repo):
+                completed.add(cfg.name())
+                print(f"  ✓ already complete: {cfg.name()}")
+        print(f"[preflight] {len(completed)}/{len(configs)} already done — will skip\n")
+
+    for i, cfg in enumerate(configs):
+        status = "SKIP" if cfg.name() in completed else "RUN "
+        print(f"#  {status} {i+1:2d}. {cfg.name()}")
+    print("#" * 100 + "\n")
+
+    results = []
+    for i, cfg in enumerate(configs):
+        if cfg.name() in completed:
+            results.append({
+                'run_name': cfg.name(),
+                'skipped': True,
+                'reason': 'already completed on HF',
+                'config': asdict(cfg),
+            })
+            continue
+
+        print(f"\n\n{'▓' * 100}")
+        print(f"▓ [{i+1}/{len(configs)}] {cfg.name()}")
+        print(f"{'▓' * 100}")
+        try:
+            final = run(cfg, out_root=out_root)
+            results.append(final)
+        except Exception as e:
+            print(f"\n[!] RUN FAILED: {cfg.name()}")
+            print(f"[!] {type(e).__name__}: {e}")
+            import traceback
+            traceback.print_exc()
+            results.append({'run_name': cfg.name(), 'error': str(e),
+                            'error_type': type(e).__name__,
+                            'config': asdict(cfg)})
+            if not skip_on_error:
+                raise
+
+    # Write combined summary
+    summary_path = os.path.join(out_root, "sweep_summary.json")
+    os.makedirs(out_root, exist_ok=True)
+    summary = {
+        'n_runs': len(results),
+        'n_succeeded': sum(1 for r in results
+                           if 'error' not in r and not r.get('skipped')),
+        'n_skipped': sum(1 for r in results if r.get('skipped')),
+        'n_errored': sum(1 for r in results if 'error' in r),
+        'runs': [
+            {
+                'run_name': r.get('run_name'),
+                'skipped': r.get('skipped', False),
+                'error': r.get('error'),
+                'best_test_mse': r.get('best_test_mse'),
+                'final_cv_ema': r.get('final_cv_ema'),
+                'final_recon_ema_obs': r.get('final_recon_ema_obs'),
+                'final_S0': r.get('final_S0'),
+                'final_erank': r.get('final_erank'),
+                'final_row_cv': r.get('final_row_cv'),
+                'final_alpha_mean': r.get('final_alpha_mean'),
+                'n_params': r.get('n_params'),
+                'n_patches': r.get('n_patches'),
+            } for r in results
+        ],
+    }
+    with open(summary_path, 'w') as f:
+        json.dump(summary, f, indent=2)
+    print(f"\n\n{'#' * 100}")
+    print(f"# SWEEP DONE: {summary['n_succeeded']} succeeded, "
+          f"{summary['n_skipped']} skipped, {summary['n_errored']} errored")
+    print(f"# Summary: {summary_path}")
+    print(f"{'#' * 100}\n")
+
+    # Upload summary to HF
+    if configs and configs[0].upload:
+        try:
+            from huggingface_hub import HfApi
+            api = HfApi()
+            api.whoami()
+            api.upload_file(
+                path_or_fileobj=summary_path,
+                path_in_repo="sweep_summary.json",
+                repo_id=configs[0].hf_repo, repo_type="model")
+            print(f"[HF] Uploaded sweep summary to {configs[0].hf_repo}/sweep_summary.json")
+        except Exception as e:
+            print(f"[HF] Summary upload failed: {e}")
+
+    return results
+
+
+def _in_notebook():
+    """Detect Jupyter / Colab environment."""
+    try:
+        from IPython import get_ipython
+        shell = get_ipython().__class__.__name__
+        return shell in ('ZMQInteractiveShell', 'Shell', 'Google.Colab')
+    except Exception:
+        return False
+
+
+def parse_args(argv=None):
+    p = argparse.ArgumentParser()
+    p.add_argument('--smoke', action='store_true', help='Run smoke test')
+    p.add_argument('--V', type=int, default=64)
+    p.add_argument('--D', type=int, default=8)
+    p.add_argument('--hidden', type=int, default=128)
+    p.add_argument('--depth', type=int, default=1)
+    p.add_argument('--patch', type=int, default=16)
+    p.add_argument('--img_size', type=int, default=128)
+    p.add_argument('--batch', type=int, default=128)
+    p.add_argument('--epochs', type=int, default=20)
+    p.add_argument('--lr', type=float, default=1e-3)
+    p.add_argument('--n_cross', type=int, default=1)
+    p.add_argument('--all_types', action='store_true')
+    p.add_argument('--no_upload', action='store_true')
+    return p.parse_args(argv)
+
+
+def _cli_main(argv=None):
+    args = parse_args(argv)
+    if args.smoke:
+        return smoke()
+    cfg = RunConfig(
+        matrix_v=args.V, D=args.D, patch_size=args.patch,
+        hidden=args.hidden, depth=args.depth,
+        n_cross_layers=args.n_cross,
+        img_size=args.img_size,
+        batch_size=args.batch,
+        epochs=args.epochs, lr=args.lr,
+        allowed_types=list(range(16)) if args.all_types else [0],
+        upload=not args.no_upload,
+    )
+    return run(cfg)
+
+
+if __name__ == "__main__" and not _in_notebook():
+    _cli_main()
+elif __name__ == "__main__":
+    # In Colab / Jupyter: do nothing automatically.
+    print("johanna_F_trainer loaded in notebook mode.")
+    print("  → smoke test:           smoke()")
+    print("  → single config:        run(RunConfig(matrix_v=256, D=16, ...))")
+    print("  → full F-class sweep:   sweep()")
+    print("  → custom sweep list:    sweep([cfg1, cfg2, ...])")