Create freckles_256_trainer.py

freckles_256_trainer.py

@@ -0,0 +1,333 @@
+"""
+Freckles High-Resolution Noise Training
+=========================================
+Same 2.5M param model (V=48, D=4, ps=4), scaled to larger images.
+Initialized from v40 Freckles 64×64 weights — patch-level weights transfer directly.
+
+256×256:  4096 patches (64×64 grid)  — batch=16
+512×512: 16384 patches (128×128 grid) — batch=4
+
+Cross-attention over N patches is O(N²). Batch sizes adjusted accordingly.
+"""
+
+import os
+import math
+import time
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from tqdm import tqdm
+
+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
+
+
+# ═══════════════════════════════════════════════════════════════
+# NOISE GENERATORS (16 types)
+# ═══════════════════════════════════════════════════════════════
+
+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 _gen_noise(t, s, rng):
+    if t == 0: return torch.randn(3, s, s)
+    elif t == 1: return torch.rand(3, s, s) * 2 - 1
+    elif t == 2: return (torch.rand(3, s, s) - 0.5) * 4
+    elif t == 3:
+        lam = rng.uniform(0.5, 20.0)
+        return torch.poisson(torch.full((3, s, s), lam)) / lam - 1.0
+    elif t == 4:
+        img = _pink((3, s, s)); return img / (img.std() + 1e-8)
+    elif t == 5:
+        img = _brown((3, s, s)); return img / (img.std() + 1e-8)
+    elif t == 6:
+        return torch.where(torch.rand(3, s, s) > 0.5,
+                          torch.ones(3, s, s) * 2, -torch.ones(3, s, s) * 2) + torch.randn(3, s, s) * 0.1
+    elif t == 7:
+        return torch.randn(3, s, s) * (torch.rand(3, s, s) > 0.9).float() * 3
+    elif t == 8:
+        b = rng.randint(2, max(3, s // 4))
+        sm = torch.randn(3, s // b + 1, s // b + 1)
+        return F.interpolate(sm.unsqueeze(0), size=s, mode='nearest').squeeze(0)
+    elif t == 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)
+        return (math.cos(a) * gx + math.sin(a) * gy).unsqueeze(0).expand(3, -1, -1) + torch.randn(3, s, s) * 0.5
+    elif t == 10:
+        cs = rng.randint(2, max(3, s // 4))
+        cy = torch.arange(s) // cs; cx = torch.arange(s) // cs
+        return ((cy.unsqueeze(1) + cx.unsqueeze(0)) % 2).float().unsqueeze(0).expand(3, -1, -1) * 2 - 1 + torch.randn(3, s, s) * 0.3
+    elif t == 11:
+        alpha = rng.uniform(0.2, 0.8)
+        return alpha * torch.randn(3, s, s) + (1 - alpha) * (torch.rand(3, s, s) * 2 - 1)
+    elif t == 12:
+        img = torch.zeros(3, s, s); h2 = s // 2; w2 = s // 2
+        img[:, :h2, :w2] = torch.randn(3, h2, w2)
+        img[:, :h2, w2:s] = torch.rand(3, h2, s - w2) * 2 - 1
+        img[:, h2:s, :w2] = _pink((3, s - h2, w2)) / 2
+        img[:, h2:s, w2:s] = torch.where(torch.rand(3, s - h2, s - w2) > 0.5,
+                                           torch.ones(3, s - h2, s - w2), -torch.ones(3, s - h2, s - w2))
+        return img
+    elif t == 13:
+        return torch.tan(math.pi * (torch.rand(3, s, s) - 0.5)).clamp(-3, 3)
+    elif t == 14:
+        return torch.empty(3, s, s).exponential_(1.0) - 1.0
+    elif t == 15:
+        u = torch.rand(3, s, s) - 0.5
+        return -torch.sign(u) * torch.log1p(-2 * u.abs())
+    return torch.randn(3, s, s)
+
+
+class OmegaNoiseDataset(torch.utils.data.Dataset):
+    def __init__(self, size=500000, img_size=256):
+        self.size = size
+        self.img_size = img_size
+        self._rng = np.random.RandomState(42)
+        self._call_count = 0
+    def __len__(self):
+        return self.size
+    def __getitem__(self, idx):
+        self._call_count += 1
+        if self._call_count % 1000 == 0:
+            self._rng = np.random.RandomState(int.from_bytes(os.urandom(4), 'big'))
+            torch.manual_seed(int.from_bytes(os.urandom(4), 'big'))
+        noise_type = idx % 16
+        img = _gen_noise(noise_type, self.img_size, self._rng).clamp(-4, 4)
+        return img.float(), noise_type
+
+
+NOISE_NAMES = {
+    0: 'gaussian', 1: 'uniform', 2: 'uniform_sc', 3: 'poisson',
+    4: 'pink', 5: 'brown', 6: 'salt_pepper', 7: 'sparse',
+    8: 'block', 9: 'gradient', 10: 'checker', 11: 'mixed',
+    12: 'structural', 13: 'cauchy', 14: 'exponential', 15: 'laplace',
+}
+
+
+# ═══════════════════════════════════════════════════════════════
+# PER-TYPE EVAL
+# ═══════════════════════════════════════════════════════════════
+
+def eval_per_type(model, img_size, device, n_per=16):
+    rng = np.random.RandomState(99)
+    model.eval()
+    results = {}
+    with torch.no_grad():
+        for t in range(16):
+            imgs = torch.stack([_gen_noise(t, img_size, rng).clamp(-4, 4)
+                                for _ in range(n_per)]).to(device)
+            out = model(imgs)
+            results[t] = F.mse_loss(out['recon'], imgs).item()
+    return results
+
+
+# ═══════════════════════════════════════════════════════════════
+# TRAINING
+# ═══════════════════════════════════════════════════════════════
+
+PRESETS = {
+    '256': dict(
+        img_size=256,
+        batch_size=64,
+        ds_size=1280000,
+        val_size=12800,
+        epochs=1,
+        lr=1e-4,
+        hf_version='v41_freckles_256',
+        save_every=1,
+    ),
+    '512': dict(
+        img_size=512,
+        batch_size=12,
+        ds_size=1280000,
+        val_size=12800,
+        epochs=1,
+        lr=1e-4,
+        hf_version='v42_freckles_512',
+        save_every=1,
+    ),
+}
+
+
+def train(preset='256', device='cuda'):
+    from geolip_svae import load_model
+
+    cfg = PRESETS[preset]
+    img_size = cfg['img_size']
+    batch_size = cfg['batch_size']
+    epochs = cfg['epochs']
+    lr = cfg['lr']
+    hf_version = cfg['hf_version']
+    save_every = cfg['save_every']
+    ps = 4
+
+    device = torch.device(device if torch.cuda.is_available() else 'cpu')
+    n_patches = (img_size // ps) ** 2
+
+    print("\n" + "=" * 70)
+    print(f"FRECKLES {img_size}×{img_size} — High-Resolution Noise Training")
+    print("=" * 70)
+
+    # ── Load from v40 Freckles ──
+    print("  Loading Freckles v40 (64×64 trained)...")
+    model, base_cfg = load_model(hf_version='v40_freckles_noise', device=device)
+    model.train()
+    for p in model.parameters():
+        p.requires_grad = True
+
+    n_params = sum(p.numel() for p in model.parameters())
+    print(f"  Params: {n_params:,} (from v40, trainable)")
+    print(f"  Resolution: {img_size}×{img_size}")
+    print(f"  Patches: {n_patches} ({img_size//ps}×{img_size//ps} grid)")
+    print(f"  SVD: ({base_cfg['V']},{base_cfg['D']}), ps={ps}")
+    print(f"  Batch: {batch_size}, lr={lr}, epochs={epochs}")
+    print(f"  Cross-attn sequence length: {n_patches}")
+    print(f"  Estimated attn memory: ~{n_patches**2 * 2 * 4 / 1e6:.0f}MB per sample")
+    print("=" * 70)
+
+    opt = torch.optim.Adam(model.parameters(), lr=lr)
+    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=epochs)
+
+    train_ds = OmegaNoiseDataset(size=cfg['ds_size'], img_size=img_size)
+    val_ds = OmegaNoiseDataset(size=cfg['val_size'], img_size=img_size)
+    train_loader = torch.utils.data.DataLoader(
+        train_ds, batch_size=batch_size, shuffle=True,
+        num_workers=8, pin_memory=True, drop_last=True)
+    val_loader = torch.utils.data.DataLoader(
+        val_ds, batch_size=batch_size, shuffle=False,
+        num_workers=2, pin_memory=True)
+
+    save_dir = f'/content/freckles_{img_size}_checkpoints'
+    hf_repo = 'AbstractPhil/geolip-SVAE'
+    os.makedirs(save_dir, exist_ok=True)
+
+    hf_enabled = False
+    api = None
+    try:
+        from huggingface_hub import HfApi
+        api = HfApi(); api.whoami(); hf_enabled = True
+        print(f"  HuggingFace: {hf_repo}/{hf_version}")
+    except:
+        pass
+
+    best_mse = float('inf')
+    D = base_cfg['D']
+
+    for epoch in range(1, epochs + 1):
+        model.train()
+        total_loss, total_recon, n = 0, 0, 0
+        t0 = time.time()
+
+        pbar = tqdm(train_loader, desc=f"Ep {epoch}/{epochs}",
+                    bar_format='{l_bar}{bar:20}{r_bar}')
+        for batch_idx, (images, _) in enumerate(pbar):
+            images = images.to(device)
+            opt.zero_grad()
+            out = model(images)
+            recon_loss = F.mse_loss(out['recon'], images)
+
+            # Pure recon loss — no CV penalty (Freckles doesn't need it)
+            loss = recon_loss
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
+            opt.step()
+
+            total_loss += loss.item() * len(images)
+            total_recon += recon_loss.item() * len(images)
+            n += len(images)
+            pbar.set_postfix_str(f"mse={recon_loss.item():.6f}")
+
+        sched.step()
+        epoch_time = time.time() - t0
+
+        # ── Eval ──
+        model.eval()
+        val_mse, val_n = 0, 0
+        with torch.no_grad():
+            for imgs, _ in val_loader:
+                imgs = imgs.to(device)
+                out = model(imgs)
+                val_mse += F.mse_loss(out['recon'], imgs).item() * len(imgs)
+                val_n += len(imgs)
+        val_mse /= val_n
+
+        # Geometry
+        with torch.no_grad():
+            sample = next(iter(val_loader))[0][:min(8, batch_size)].to(device)
+            out = model(sample)
+            S_mean = out['svd']['S'].mean(dim=(0, 1))
+            erank = model.effective_rank(out['svd']['S'].reshape(-1, D)).mean().item()
+
+        # Per-type (smaller sample for speed at high res)
+        type_mse = eval_per_type(model, img_size, device, n_per=8)
+        type_str = " ".join(f"{NOISE_NAMES[t][:4]}={v:.4f}" for t, v in sorted(type_mse.items()))
+
+        print(f"  ep{epoch:3d} | recon={total_recon/n:.6f} val={val_mse:.6f} | "
+              f"S0={S_mean[0]:.3f} SD={S_mean[-1]:.3f} er={erank:.2f} | {epoch_time:.0f}s")
+        print(f"         {type_str}")
+
+        # ── Checkpoint ──
+        ckpt = {
+            'epoch': epoch, 'val_mse': val_mse,
+            'model_state_dict': model.state_dict(),
+            'config': {
+                'V': base_cfg['V'], 'D': D, 'patch_size': ps,
+                'hidden': base_cfg['hidden'], 'depth': base_cfg['depth'],
+                'n_cross_layers': base_cfg['n_cross_layers'],
+                'n_heads': 2, 'smooth_mid': 8,
+                'img_size': img_size,
+            },
+        }
+
+        if val_mse < best_mse:
+            best_mse = val_mse
+            torch.save(ckpt, os.path.join(save_dir, 'best.pt'))
+
+        if epoch % save_every == 0:
+            path = os.path.join(save_dir, f'epoch_{epoch:04d}.pt')
+            torch.save(ckpt, path)
+            if hf_enabled:
+                try:
+                    api.upload_file(path_or_fileobj=path,
+                                    path_in_repo=f"{hf_version}/checkpoints/{os.path.basename(path)}",
+                                    repo_id=hf_repo, repo_type="model")
+                    api.upload_file(path_or_fileobj=os.path.join(save_dir, 'best.pt'),
+                                    path_in_repo=f"{hf_version}/checkpoints/best.pt",
+                                    repo_id=hf_repo, repo_type="model")
+                    print(f"  ☁️  Uploaded ep{epoch}")
+                except Exception as e:
+                    print(f"  ⚠️  Upload: {e}")
+
+    print(f"\n  FRECKLES {img_size}×{img_size} COMPLETE")
+    print(f"  Best val MSE: {best_mse:.6f}")
+    return model
+
+
+if __name__ == "__main__":
+    import sys
+    torch.set_float32_matmul_precision('high')
+
+    # CLI: python freckles_hires.py 256
+    # Colab: just set PRESET below
+    PRESET = '256'  # ← change to '512' for the other run
+    train(preset=PRESET)