""" 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): from huggingface_hub import hf_hub_download if checkpoint_path and os.path.exists(checkpoint_path): path = checkpoint_path 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, n_samples=64): """ Run full diagnostic battery. Usage in Colab cell: run(hf_version='v13_imagenet256') run(hf_version='v16_johanna_omega') run(hf_version='v18_johanna_curriculum') run(checkpoint_path='/content/checkpoints/best.pt') """ print("=" * 80) print("UNIVERSAL SVAE DIAGNOSTIC BATTERY") print("=" * 80) model, cfg = load_model(hf_version=hf_version, checkpoint_path=checkpoint_path) # 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 'local' out_path = f'diagnostic_{tag.replace("/","_")}.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 HF_VERSION = 'v16_johanna_omega' # Johanna-small 128 # HF_VERSION = 'v18_johanna_curriculum' # Johanna-tiny 64 # HF_VERSION = None; CHECKPOINT = '/content/checkpoints/best.pt' if __name__ == "__main__": run(hf_version=HF_VERSION)