Create universal_diagnostic.py

universal_diagnostic.py

@@ -0,0 +1,693 @@
+"""
+Universal SVAE Diagnostic Battery
+===================================
+One script. Any checkpoint. Every dataset.
+
+Usage:
+  python universal_diagnostic.py                          # local best.pt
+  python universal_diagnostic.py --hf v13_imagenet256     # HF version
+  python universal_diagnostic.py --checkpoint /path/to.pt # explicit path
+
+Tests across:
+  - CIFAR-10 (32×32, resized to model native)
+  - MNIST (28×28, resized, grayscale→RGB)
+  - TinyImageNet (64×64)
+  - ImageNet-128 (128×128)
+  - ImageNet-256 (256×256)
+  - 16 noise types (native resolution)
+  - Text bytes (5 sentences)
+  - Piecemeal tiling (4× resolution)
+  - Geometric fingerprint per dataset
+  - Spectrum analysis
+  - Alpha profile
+  - Compression metrics
+"""
+
+import os, sys, json, math, time
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision
+import torchvision.transforms as T
+from collections import defaultdict
+
+DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+# ── Model ────────────────────────────────────────────────────────
+
+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):
+    def __init__(self, D, n_heads=4, max_alpha=0.2, alpha_init=-2.0):
+        super().__init__()
+        self.n_heads = n_heads; self.head_dim = D // n_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_n = self.norm(S)
+        qkv = self.qkv(S_n).reshape(B,N,3,self.n_heads,self.head_dim).permute(2,0,3,1,4)
+        q, k, v = qkv[0], qkv[1], qkv[2]
+        out = (((q @ k.transpose(-2,-1))*self.scale).softmax(-1) @ v).transpose(1,2).reshape(B,N,D)
+        return S * (1.0 + self.alpha.unsqueeze(0).unsqueeze(0) * torch.tanh(self.out_proj(out)))
+
+class PatchSVAE(nn.Module):
+    def __init__(self, V=256, D=16, ps=16, hidden=768, depth=4, n_cross=2):
+        super().__init__()
+        self.matrix_v, self.D, self.patch_size = V, D, ps
+        self.patch_dim = 3*ps*ps; self.mat_dim = V*D
+        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)
+        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)
+        nn.init.orthogonal_(self.enc_out.weight)
+        self.cross_attn = nn.ModuleList([
+            SpectralCrossAttention(D, n_heads=min(4,D)) for _ in range(n_cross)])
+        self.boundary_smooth = BoundarySmooth(channels=3, mid=16)
+
+    def _svd(self, A):
+        orig = 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)
+            eig, V = torch.linalg.eigh(G)
+            eig = eig.flip(-1); V = V.flip(-1)
+            S = torch.sqrt(eig.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), S.to(orig), Vh.to(orig)
+
+    def encode_patches(self, patches):
+        B, N, _ = patches.shape
+        h = F.gelu(self.enc_in(patches.reshape(B*N,-1)))
+        for block in self.enc_blocks: h = h + block(h)
+        M = F.normalize(self.enc_out(h).reshape(B*N, self.matrix_v, self.D), dim=-1)
+        U, S, Vt = self._svd(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_c = S
+        for layer in self.cross_attn: S_c = layer(S_c)
+        return {'U':U, 'S_orig':S, 'S':S_c, 'Vt':Vt, 'M':M}
+
+    def decode_patches(self, U, S, Vt):
+        B, N, V, D = U.shape
+        M_hat = torch.bmm(U.reshape(B*N,V,D)*S.reshape(B*N,D).unsqueeze(1), Vt.reshape(B*N,D,D))
+        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):
+        B, C, H, W = images.shape
+        ps = self.patch_size
+        gh, gw = H//ps, W//ps
+        p = images.reshape(B,C,gh,ps,gw,ps).permute(0,2,4,1,3,5).reshape(B,gh*gw,C*ps*ps)
+        svd = self.encode_patches(p)
+        dec = self.decode_patches(svd['U'], svd['S'], svd['Vt'])
+        dec = dec.reshape(B,gh,gw,3,ps,ps).permute(0,3,1,4,2,5).reshape(B,3,gh*ps,gw*ps)
+        return {'recon': self.boundary_smooth(dec), '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()
+
+
+def load_model(hf_version=None, checkpoint_path=None, hf_file=None):
+    from huggingface_hub import hf_hub_download
+    if checkpoint_path and os.path.exists(checkpoint_path):
+        path = checkpoint_path
+    elif hf_file:
+        path = hf_hub_download(repo_id='AbstractPhil/geolip-SVAE',
+                                filename=hf_file, repo_type='model')
+    elif hf_version:
+        path = hf_hub_download(repo_id='AbstractPhil/geolip-SVAE',
+                                filename=f'{hf_version}/checkpoints/best.pt', repo_type='model')
+    else:
+        path = '/content/checkpoints/best.pt'
+    print(f"  Loading: {path}")
+    ckpt = torch.load(path, map_location='cpu', weights_only=False)
+    cfg = ckpt['config']
+    print(f"  Epoch: {ckpt.get('epoch')}, MSE: {ckpt.get('test_mse','?'):.6f}")
+    print(f"  Config: {cfg}")
+    model = PatchSVAE(V=cfg['V'], D=cfg['D'], ps=cfg['patch_size'],
+                       hidden=cfg['hidden'], depth=cfg['depth'],
+                       n_cross=cfg['n_cross_layers'])
+    model.load_state_dict(ckpt['model_state_dict'], strict=True)
+    model = model.to(DEVICE).eval()
+    print(f"  Params: {sum(p.numel() for p in model.parameters()):,}")
+    return model, cfg
+
+
+# ── Noise Generators ─────────────────────────────────────────────
+
+NOISE_NAMES = {
+    0:'gaussian', 1:'uniform', 2:'uniform_scaled', 3:'poisson',
+    4:'pink', 5:'brown', 6:'salt_pepper', 7:'sparse',
+    8:'block', 9:'gradient', 10:'checkerboard', 11:'mixed',
+    12:'structural', 13:'cauchy', 14:'exponential', 15:'laplace',
+}
+
+def _pink(shape):
+    w = torch.randn(shape); S = torch.fft.rfft2(w)
+    h, ww = shape[-2], shape[-1]
+    fy = torch.fft.fftfreq(h).unsqueeze(-1).expand(-1, ww//2+1)
+    fx = torch.fft.rfftfreq(ww).unsqueeze(0).expand(h, -1)
+    return torch.fft.irfft2(S / torch.sqrt(fx**2 + fy**2).clamp(min=1e-8), s=(h, ww))
+
+def _brown(shape):
+    w = torch.randn(shape); S = torch.fft.rfft2(w)
+    h, ww = shape[-2], shape[-1]
+    fy = torch.fft.fftfreq(h).unsqueeze(-1).expand(-1, ww//2+1)
+    fx = torch.fft.rfftfreq(ww).unsqueeze(0).expand(h, -1)
+    return torch.fft.irfft2(S / (fx**2 + fy**2).clamp(min=1e-8), s=(h, ww))
+
+def generate_noise(noise_type, n, s):
+    rng = np.random.RandomState(42)
+    imgs = []
+    for _ in range(n):
+        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 = rng.uniform(0.5,20.0)
+            img = torch.poisson(torch.full((3,s,s),lam))/lam-1.0
+        elif noise_type == 4: img = _pink((3,s,s)); img = img/(img.std()+1e-8)
+        elif noise_type == 5: img = _brown((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: img = torch.randn(3,s,s)*(torch.rand(3,s,s)>0.9).float()*3
+        elif noise_type == 8:
+            b = rng.randint(2,16); sm = torch.randn(3,s//b+1,s//b+1)
+            img = F.interpolate(sm.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)
+            a = rng.uniform(0,2*math.pi)
+            img = (math.cos(a)*gx+math.sin(a)*gy).unsqueeze(0).expand(3,-1,-1)+torch.randn(3,s,s)*0.5
+        elif noise_type == 10:
+            cs = rng.randint(2,16); cy = torch.arange(s)//cs; cx = torch.arange(s)//cs
+            img = ((cy.unsqueeze(1)+cx.unsqueeze(0))%2).float().unsqueeze(0).expand(3,-1,-1)*2-1+torch.randn(3,s,s)*0.3
+        elif noise_type == 11:
+            alpha = rng.uniform(0.2,0.8)
+            img = alpha*torch.randn(3,s,s)+(1-alpha)*(torch.rand(3,s,s)*2-1)
+        elif noise_type == 12:
+            img = torch.zeros(3,s,s); h2 = s//2
+            img[:,:h2,:h2] = torch.randn(3,h2,h2)
+            img[:,:h2,h2:] = torch.rand(3,h2,h2)*2-1
+            img[:,h2:,:h2] = _pink((3,h2,h2))/2
+            img[:,h2:,h2:] = torch.where(torch.rand(3,h2,h2)>0.5, torch.ones(3,h2,h2), -torch.ones(3,h2,h2))
+        elif noise_type == 13: img = torch.tan(math.pi*(torch.rand(3,s,s)-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: img = torch.randn(3,s,s)
+        imgs.append(img.clamp(-4,4))
+    return torch.stack(imgs)
+
+
+# ── Batched Forward ──────────────────────────────────────────────
+
+def batched_forward(model, images, max_batch=16):
+    """Forward pass in chunks to avoid OOM."""
+    all_recon = []; all_S = []; all_S_orig = []; all_M = []
+    model.eval()
+    with torch.no_grad():
+        for i in range(0, len(images), max_batch):
+            batch = images[i:i+max_batch].to(DEVICE)
+            out = model(batch)
+            all_recon.append(out['recon'].cpu())
+            all_S.append(out['svd']['S'].cpu())
+            all_S_orig.append(out['svd']['S_orig'].cpu())
+            all_M.append(out['svd']['M'].cpu())
+    return {
+        'recon': torch.cat(all_recon),
+        'S': torch.cat(all_S),
+        'S_orig': torch.cat(all_S_orig),
+        'M': torch.cat(all_M),
+    }
+
+
+# ── Dataset Loaders ──────────────────────────────────────────────
+
+def load_dataset_batch(name, s, n=100):
+    """Load n images from a dataset, resized to s×s, normalized.
+    Returns: (tensor [N,3,H,W], mean, std, ds_name)
+    """
+    from datasets import load_dataset as hf_load
+
+    if name == 'cifar10':
+        transform = T.Compose([T.Resize(s), T.ToTensor(),
+                                T.Normalize((0.4914,0.4822,0.4465),(0.2470,0.2435,0.2616))])
+        ds = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
+        imgs = [ds[i][0] for i in range(min(n, len(ds)))]
+        return torch.stack(imgs), (0.4914,0.4822,0.4465), (0.2470,0.2435,0.2616), f'CIFAR-10→{s}'
+
+    elif name == 'mnist':
+        transform = T.Compose([T.Resize(s), T.Grayscale(3), T.ToTensor(),
+                                T.Normalize((0.1307,0.1307,0.1307),(0.3081,0.3081,0.3081))])
+        ds = torchvision.datasets.MNIST(root='./data', train=False, download=True, transform=transform)
+        imgs = [ds[i][0] for i in range(min(n, len(ds)))]
+        return torch.stack(imgs), (0.1307,0.1307,0.1307), (0.3081,0.3081,0.3081), f'MNIST→{s}'
+
+    elif name == 'tiny_imagenet':
+        ds = hf_load('zh-plus/tiny-imagenet', split='valid', streaming=True)
+        transform = T.Compose([T.Resize(s), T.CenterCrop(s), T.ToTensor(),
+                                T.Normalize((0.4802,0.4481,0.3975),(0.2770,0.2691,0.2821))])
+        imgs = []
+        for i, sample in enumerate(ds):
+            imgs.append(transform(sample['image'].convert('RGB')))
+            if i >= n-1: break
+        return torch.stack(imgs), (0.4802,0.4481,0.3975), (0.2770,0.2691,0.2821), f'TinyImageNet→{s}'
+
+    elif name == 'imagenet128':
+        ds = hf_load('benjamin-paine/imagenet-1k-128x128', split='validation', streaming=True)
+        transform = T.Compose([T.Resize(s), T.CenterCrop(s), T.ToTensor(),
+                                T.Normalize((0.485,0.456,0.406),(0.229,0.224,0.225))])
+        imgs = []
+        for i, sample in enumerate(ds):
+            imgs.append(transform(sample['image'].convert('RGB')))
+            if i >= n-1: break
+        return torch.stack(imgs), (0.485,0.456,0.406), (0.229,0.224,0.225), f'ImageNet-128→{s}'
+
+    elif name == 'imagenet256':
+        ds = hf_load('benjamin-paine/imagenet-1k-256x256', split='validation', streaming=True)
+        transform = T.Compose([T.Resize(s), T.CenterCrop(s), T.ToTensor(),
+                                T.Normalize((0.485,0.456,0.406),(0.229,0.224,0.225))])
+        imgs = []
+        for i, sample in enumerate(ds):
+            imgs.append(transform(sample['image'].convert('RGB')))
+            if i >= n-1: break
+        return torch.stack(imgs), (0.485,0.456,0.406), (0.229,0.224,0.225), f'ImageNet-256→{s}'
+
+    raise ValueError(f"Unknown dataset: {name}")
+
+
+# ════════════════════════════════════════════════════════════════
+# TESTS
+# ════════════════════════════════════════════════════════════════
+
+IMAGE_DATASETS = ['cifar10', 'mnist', 'tiny_imagenet', 'imagenet128', 'imagenet256']
+
+
+def test_image_datasets(model, cfg, n=100):
+    """Reconstruction MSE + geometry across all image datasets."""
+    s = cfg['img_size']
+    D = cfg['D']
+    bs = max(4, 64 // max(1, (s // 64) ** 2))
+    print(f"\n{'='*80}")
+    print(f"IMAGE DATASET BATTERY ({s}×{s}, n={n})")
+    print(f"{'='*80}")
+    print(f"  {'dataset':22s} {'MSE':>10s} {'std':>10s} {'min':>10s} {'max':>10s} | "
+          f"{'S0':>6s} {'SD':>6s} {'ratio':>6s} {'erank':>6s}")
+    print("-" * 100)
+
+    results = {}
+    for ds_name in IMAGE_DATASETS:
+        try:
+            imgs, mean, std, label = load_dataset_batch(ds_name, s, n)
+            out = batched_forward(model, imgs, max_batch=bs)
+            mse = F.mse_loss(out['recon'], imgs, reduction='none').mean(dim=(1,2,3))
+            S_mean = out['S'].mean(dim=(0,1))
+            ratio = (S_mean[0] / (S_mean[-1]+1e-8)).item()
+            erank = model.effective_rank(out['S'].reshape(-1, D)).mean().item()
+
+            results[ds_name] = {
+                'label': label,
+                'mse_mean': mse.mean().item(), 'mse_std': mse.std().item(),
+                'mse_min': mse.min().item(), 'mse_max': mse.max().item(),
+                'S0': S_mean[0].item(), 'SD': S_mean[-1].item(),
+                'ratio': ratio, 'erank': erank,
+                'fidelity': (1 - mse.mean()).item() * 100,
+            }
+            print(f"  {label:22s} {mse.mean():10.6f} {mse.std():10.6f} "
+                  f"{mse.min():10.6f} {mse.max():10.6f} | "
+                  f"{S_mean[0]:6.3f} {S_mean[-1]:6.3f} {ratio:6.2f} {erank:6.2f}")
+        except Exception as e:
+            print(f"  {ds_name:22s} FAILED: {e}")
+            results[ds_name] = {'error': str(e)}
+    return results
+
+
+def test_noise_types(model, cfg, n=64):
+    """Per-type noise reconstruction + geometry."""
+    s = cfg['img_size']
+    D = cfg['D']
+    bs = max(4, 64 // max(1, (s // 64) ** 2))
+    print(f"\n{'='*80}")
+    print(f"NOISE TYPE BATTERY ({s}×{s}, n={n})")
+    print(f"{'='*80}")
+    print(f"  {'type':18s} {'MSE':>10s} {'std':>10s} | "
+          f"{'S0':>6s} {'SD':>6s} {'ratio':>6s} {'erank':>6s} | "
+          f"{'byte_acc':>8s} {'±1_acc':>8s}")
+    print("-" * 100)
+
+    results = {}
+    for t in range(16):
+        name = NOISE_NAMES[t]
+        imgs = generate_noise(t, n, s)
+        out = batched_forward(model, imgs, max_batch=bs)
+        mse = F.mse_loss(out['recon'], imgs, reduction='none').mean(dim=(1,2,3))
+        S_mean = out['S'].mean(dim=(0,1))
+        ratio = (S_mean[0] / (S_mean[-1]+1e-8)).item()
+        erank = model.effective_rank(out['S'].reshape(-1, D)).mean().item()
+
+        # Byte accuracy
+        orig_q = ((imgs + 4) / 8 * 255).round().clamp(0,255).long()
+        recon_q = ((out['recon'] + 4) / 8 * 255).round().clamp(0,255).long()
+        byte_acc = (orig_q == recon_q).float().mean().item()
+        byte_1 = ((orig_q - recon_q).abs() <= 1).float().mean().item()
+
+        results[name] = {
+            'mse_mean': mse.mean().item(), 'mse_std': mse.std().item(),
+            'S0': S_mean[0].item(), 'SD': S_mean[-1].item(),
+            'ratio': ratio, 'erank': erank,
+            'byte_exact': byte_acc, 'byte_within1': byte_1,
+        }
+        print(f"  {name:18s} {mse.mean():10.6f} {mse.std():10.6f} | "
+              f"{S_mean[0]:6.3f} {S_mean[-1]:6.3f} {ratio:6.2f} {erank:6.2f} | "
+              f"{byte_acc*100:7.2f}% {byte_1*100:7.2f}%")
+    return results
+
+
+def test_text_bytes(model, cfg):
+    """Text-as-bytes reconstruction."""
+    s = cfg['img_size']
+    print(f"\n{'='*80}")
+    print(f"TEXT BYTE RECONSTRUCTION ({s}×{s})")
+    print(f"{'='*80}")
+
+    texts = [
+        "Hello, world! This is a test of the geometric encoder.",
+        "The quick brown fox jumps over the lazy dog. 0123456789",
+        "import torch; model = PatchSVAE(); output = model(x)",
+        "E = mc² — Albert Einstein, theoretical physicist, 1905",
+        "To be, or not to be, that is the question. — Shakespeare",
+        "∫₀^∞ e^(-x²) dx = √π/2 — Gaussian integral",
+        "01101000 01100101 01101100 01101100 01101111 — binary hello",
+        "SELECT * FROM models WHERE cv BETWEEN 0.20 AND 0.23;",
+    ]
+
+    n_bytes = 3 * s * s
+    results = {}
+    model.eval()
+
+    for text in texts:
+        raw = text.encode('utf-8')
+        actual_len = min(len(raw), n_bytes)
+        padded = (raw + b'\x00' * n_bytes)[:n_bytes]
+
+        arr = np.frombuffer(padded, dtype=np.uint8).copy()
+        tensor = torch.from_numpy(arr).float()
+        tensor = (tensor / 127.5) - 1.0
+        tensor = tensor.reshape(1, 3, s, s).to(DEVICE)
+
+        with torch.no_grad():
+            out = model(tensor)
+            recon = out['recon']
+            mse = F.mse_loss(recon, tensor).item()
+
+        orig_b = ((tensor.squeeze(0).cpu().flatten() + 1.0) * 127.5).round().clamp(0,255).byte()
+        recon_b = ((recon.squeeze(0).cpu().flatten() + 1.0) * 127.5).round().clamp(0,255).byte()
+        exact_acc = (orig_b[:actual_len] == recon_b[:actual_len]).float().mean().item()
+        recovered = recon_b[:actual_len].numpy().tobytes().decode('utf-8', errors='replace')
+
+        results[text[:40]] = {'mse': mse, 'byte_acc': exact_acc}
+        print(f"\n  In:  '{text[:60]}'")
+        print(f"  Out: '{recovered[:60]}'")
+        print(f"  MSE: {mse:.6f}  Byte: {exact_acc*100:.1f}%")
+    return results
+
+
+def test_piecemeal(model, cfg):
+    """Piecemeal tiling at 4× resolution."""
+    s = cfg['img_size']
+    src = max(256, s * 4)
+    bs = max(2, 32 // max(1, (s // 64) ** 2))
+    print(f"\n{'='*80}")
+    print(f"PIECEMEAL {src}→{s} TILED RECONSTRUCTION")
+    print(f"{'='*80}")
+    model.eval()
+    results = {}
+
+    test_types = [0, 1, 4, 6, 13]  # Gaussian, Uniform, Pink, Salt-pepper, Cauchy
+    with torch.no_grad():
+        for t in test_types:
+            img_src = generate_noise(t, 1, src).squeeze(0)
+            tiles = []
+            gh, gw = src // s, src // s
+            for gy in range(gh):
+                for gx in range(gw):
+                    tiles.append(img_src[:, gy*s:(gy+1)*s, gx*s:(gx+1)*s])
+
+            # Batch tiles through model
+            all_recon = []
+            tile_t = torch.stack(tiles)
+            for i in range(0, len(tile_t), bs):
+                batch = tile_t[i:i+bs].to(DEVICE)
+                out = model(batch)
+                all_recon.append(out['recon'].cpu())
+            recon_tiles = torch.cat(all_recon)
+
+            recon_full = torch.zeros(3, src, src)
+            idx = 0
+            for gy in range(gh):
+                for gx in range(gw):
+                    recon_full[:, gy*s:(gy+1)*s, gx*s:(gx+1)*s] = recon_tiles[idx]
+                    idx += 1
+
+            mse = F.mse_loss(recon_full, img_src).item()
+            results[NOISE_NAMES[t]] = mse
+            print(f"  {NOISE_NAMES[t]:18s}: {gh*gw} tiles, MSE={mse:.6f}")
+    return results
+
+
+def test_signal_survival(model, cfg, n=32):
+    """Signal energy survival and SNR per dataset."""
+    s = cfg['img_size']
+    bs = max(4, 64 // max(1, (s // 64) ** 2))
+    print(f"\n{'='*80}")
+    print(f"SIGNAL ENERGY SURVIVAL")
+    print(f"{'='*80}")
+    print(f"  {'source':22s} {'survival':>10s} {'SNR_dB':>10s} {'orig_E':>10s} {'recon_E':>10s}")
+    print("-" * 70)
+
+    results = {}
+
+    # Image datasets
+    for ds_name in IMAGE_DATASETS:
+        try:
+            imgs, _, _, label = load_dataset_batch(ds_name, s, n)
+            out = batched_forward(model, imgs, max_batch=bs)
+            orig_E = (imgs**2).mean().item()
+            recon_E = (out['recon']**2).mean().item()
+            err_E = ((imgs - out['recon'])**2).mean().item()
+            survival = recon_E / (orig_E + 1e-8) * 100
+            snr = 10 * math.log10(orig_E / (err_E + 1e-8))
+            results[ds_name] = {'survival': survival, 'snr': snr}
+            print(f"  {label:22s} {survival:9.1f}% {snr:9.1f}dB {orig_E:10.4f} {recon_E:10.4f}")
+        except:
+            pass
+
+    # Key noise types
+    for t in [0, 4, 6, 13]:
+        imgs = generate_noise(t, n, s)
+        out = batched_forward(model, imgs, max_batch=bs)
+        orig_E = (imgs**2).mean().item()
+        recon_E = (out['recon']**2).mean().item()
+        err_E = ((imgs - out['recon'])**2).mean().item()
+        survival = recon_E / (orig_E + 1e-8) * 100
+        snr = 10 * math.log10(orig_E / (err_E + 1e-8))
+        results[NOISE_NAMES[t]] = {'survival': survival, 'snr': snr}
+        print(f"  noise/{NOISE_NAMES[t]:17s} {survival:9.1f}% {snr:9.1f}dB {orig_E:10.4f} {recon_E:10.4f}")
+
+    return results
+
+
+def test_alpha_profile(model):
+    """Cross-attention alpha analysis."""
+    print(f"\n{'='*80}")
+    print("ALPHA PROFILE")
+    print(f"{'='*80}")
+    results = {}
+    for li, layer in enumerate(model.cross_attn):
+        alpha = layer.alpha.detach().cpu()
+        results[f'layer_{li}'] = {
+            'mean': alpha.mean().item(), 'max': alpha.max().item(),
+            'min': alpha.min().item(), 'std': alpha.std().item(),
+            'values': alpha.tolist(),
+        }
+        print(f"\n  Layer {li}: mean={alpha.mean():.5f}  max={alpha.max():.5f}  "
+              f"min={alpha.min():.5f}  std={alpha.std():.6f}")
+        bar_scale = 50 / (alpha.max().item() + 1e-8)
+        for d in range(len(alpha)):
+            bar = "█" * int(alpha[d].item() * bar_scale)
+            print(f"    α[{d:2d}]: {alpha[d]:.5f}  {bar}")
+    return results
+
+
+def test_compression(model, cfg):
+    """Compression metrics."""
+    s = cfg['img_size']
+    D = cfg['D']; ps = cfg['patch_size']
+    n_patches = (s // ps) ** 2
+    input_vals = 3 * s * s
+    latent_vals = D * n_patches
+    ratio = input_vals / latent_vals
+    print(f"\n{'='*80}")
+    print("COMPRESSION METRICS")
+    print(f"{'='*80}")
+    print(f"  Input:   {s}×{s}×3 = {input_vals:,} values")
+    print(f"  Latent:  {D}×{n_patches} = {latent_vals:,} omega tokens")
+    print(f"  Ratio:   {ratio:.1f}:1")
+    for bits in [8, 16, 32]:
+        ib = input_vals * (bits//8)
+        lb = latent_vals * (bits//8)
+        print(f"  {bits}-bit: input={ib/1024:.1f}KB  latent={lb/1024:.1f}KB  ratio={ib/lb:.1f}:1")
+    return {'input_values': input_vals, 'latent_values': latent_vals, 'ratio': ratio}
+
+
+def test_reconstruction_grid(model, cfg):
+    """Save visual grid: 5 image datasets + 4 key noise types."""
+    s = cfg['img_size']
+    print(f"\n{'='*80}")
+    print("RECONSTRUCTION GRID")
+    print(f"{'='*80}")
+    import matplotlib
+    matplotlib.use('Agg')
+    import matplotlib.pyplot as plt
+
+    rows = []
+    labels = []
+
+    # Image datasets
+    for ds_name in IMAGE_DATASETS:
+        try:
+            imgs, mean, std, label = load_dataset_batch(ds_name, s, 2)
+            mean_t = torch.tensor(mean).reshape(1,3,1,1)
+            std_t = torch.tensor(std).reshape(1,3,1,1)
+            out = batched_forward(model, imgs[:1], max_batch=1)
+            orig_vis = (imgs[:1] * std_t + mean_t).clamp(0,1)
+            recon_vis = (out['recon'][:1] * std_t + mean_t).clamp(0,1)
+            rows.append((orig_vis[0], recon_vis[0]))
+            labels.append(label)
+        except:
+            pass
+
+    # Key noise types
+    for t in [0, 6, 13, 4]:
+        imgs = generate_noise(t, 1, s)
+        out = batched_forward(model, imgs, max_batch=1)
+        o = imgs[0].clamp(-3,3); r = out['recon'][0].clamp(-3,3)
+        o = (o - o.min())/(o.max()-o.min()+1e-8)
+        r = (r - r.min())/(r.max()-r.min()+1e-8)
+        rows.append((o, r))
+        labels.append(f'noise/{NOISE_NAMES[t]}')
+
+    n_rows = len(rows)
+    fig, axes = plt.subplots(n_rows, 3, figsize=(9, n_rows*3))
+    if n_rows == 1: axes = axes.reshape(1, -1)
+    for i, (orig, recon) in enumerate(rows):
+        diff = (orig - recon).abs().clamp(0,1)
+        axes[i,0].imshow(orig.permute(1,2,0).numpy())
+        axes[i,1].imshow(recon.permute(1,2,0).numpy())
+        axes[i,2].imshow((diff * 5).clamp(0,1).permute(1,2,0).numpy())
+        axes[i,0].set_ylabel(labels[i], fontsize=8)
+        for j in range(3): axes[i,j].axis('off')
+    axes[0,0].set_title('Original', fontsize=9)
+    axes[0,1].set_title('Recon', fontsize=9)
+    axes[0,2].set_title('|Err|×5', fontsize=9)
+    plt.tight_layout()
+    fname = 'universal_diagnostic_grid.png'
+    plt.savefig(fname, dpi=150, bbox_inches='tight')
+    print(f"  Saved: {fname}")
+    plt.close()
+
+
+# ════════════════════════════════════════════════════════════════
+# MAIN
+# ════════════════════════════════════════════════════════════════
+
+def run(hf_version=None, checkpoint_path=None, hf_file=None, n_samples=64):
+    """
+    Run full diagnostic battery.
+
+    Usage in Colab cell:
+        run(hf_version='v13_imagenet256')                                          # best.pt
+        run(hf_file='v18_johanna_curriculum/checkpoints/epoch_0300.pt')            # specific file
+        run(checkpoint_path='/content/checkpoints/best.pt')                         # local
+    """
+    print("=" * 80)
+    print("UNIVERSAL SVAE DIAGNOSTIC BATTERY")
+    print("=" * 80)
+
+    model, cfg = load_model(hf_version=hf_version, checkpoint_path=checkpoint_path, hf_file=hf_file)
+
+    # Infer img_size if not in config
+    if 'img_size' not in cfg:
+        ds = cfg.get('dataset', '')
+        if '256' in ds: cfg['img_size'] = 256
+        elif '128' in ds: cfg['img_size'] = 128
+        elif 'tiny' in ds: cfg['img_size'] = 64
+        elif 'cifar' in ds: cfg['img_size'] = 32
+        else: cfg['img_size'] = 64
+        print(f"  Inferred img_size={cfg['img_size']} from dataset='{ds}'")
+
+    s = cfg['img_size']
+    n = min(n_samples, max(16, 100 // max(1, (s // 64) ** 2)))
+    print(f"  Resolution: {s}×{s}, samples_per_test: {n}")
+
+    results = {'config': cfg}
+
+    results['image_datasets'] = test_image_datasets(model, cfg, n=n)
+    results['noise_types'] = test_noise_types(model, cfg, n=n)
+    results['text'] = test_text_bytes(model, cfg)
+    results['piecemeal'] = test_piecemeal(model, cfg)
+    results['signal_survival'] = test_signal_survival(model, cfg, n=n)
+    results['alpha'] = test_alpha_profile(model)
+    results['compression'] = test_compression(model, cfg)
+    test_reconstruction_grid(model, cfg)
+
+    tag = hf_version or (hf_file or 'local').replace('/','_').replace('.pt','')
+    out_path = f'diagnostic_{tag}.json'
+    with open(out_path, 'w') as f:
+        json.dump(results, f, indent=2, default=str)
+    print(f"\n  Results: {out_path}")
+    print(f"\n{'='*80}")
+    print("DIAGNOSTIC COMPLETE")
+    print(f"{'='*80}")
+    return results
+
+
+# ── CONFIG: Change this per run ──────────────────────────────────
+# Uncomment the model you want to diagnose:
+
+# HF_VERSION = 'v13_imagenet256'        # Fresnel-base 256 (best.pt)
+# HF_VERSION = 'v16_johanna_omega'      # Johanna-small 128 (best.pt)
+# HF_VERSION = 'v18_johanna_curriculum' # Johanna-tiny 64 (best.pt = ep10 pre-crash!)
+
+# For specific checkpoints use hf_file:
+HF_FILE = 'v18_johanna_curriculum/checkpoints/epoch_0300.pt'  # Johanna-tiny FULL curriculum
+
+if __name__ == "__main__":
+    if 'HF_FILE' in dir() and HF_FILE:
+        run(hf_file=HF_FILE)
+    else:
+        run(hf_version=HF_VERSION)