Upload lrf_v3.py with huggingface_hub

lrf_v3.py

@@ -0,0 +1,450 @@
+#!/usr/bin/env python3
+"""
+LatentRecurrentFlow v3 — Complete Self-Contained Training & Generation
+======================================================================
+
+This is one file. Run it and it:
+1. Downloads TAESD (pre-trained tiny VAE, 2.4M params, frozen)
+2. Pre-computes CIFAR-10 latents (~4 min on CPU)
+3. Trains a 1.5M-param recursive flow-matching denoiser (30 epochs, ~60 min CPU)
+4. Generates class-conditional 32×32 images and saves them
+5. Saves everything to a HuggingFace repo
+
+No custom VAE training. No grey images. Works end-to-end.
+"""
+
+import math, os, sys, time, json
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader, TensorDataset
+from einops import rearrange
+from typing import Optional, Dict, Any
+import numpy as np
+
+DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
+print(f"Device: {DEVICE}")
+
+# =============================================================================
+# 1. MODEL ARCHITECTURE
+# =============================================================================
+
+class RMSNorm(nn.Module):
+    def __init__(self, d, eps=1e-6):
+        super().__init__()
+        self.eps = eps
+        self.w = nn.Parameter(torch.ones(d))
+    def forward(self, x):
+        return x * (x.float().pow(2).mean(-1, keepdim=True) + self.eps).rsqrt().type_as(x) * self.w
+
+class SwiGLU(nn.Module):
+    def __init__(self, d, hd=None):
+        super().__init__()
+        hd = hd or ((d * 8 // 3 + 7) // 8 * 8)
+        self.w1 = nn.Linear(d, hd, bias=False)
+        self.w2 = nn.Linear(hd, d, bias=False)
+        self.w3 = nn.Linear(d, hd, bias=False)
+    def forward(self, x):
+        return self.w2(F.silu(self.w1(x)) * self.w3(x))
+
+class SpatialMixer(nn.Module):
+    """Multi-head self-attention + depthwise conv for 2D locality."""
+    def __init__(self, d, nh=4, hd=32):
+        super().__init__()
+        self.nh, self.hd = nh, hd
+        inner = nh * hd
+        self.qkv = nn.Linear(d, 3 * inner, bias=False)
+        self.out = nn.Linear(inner, d, bias=False)
+        self.gate = nn.Sequential(nn.Linear(d, inner, bias=False), nn.SiLU())
+        self.dw = nn.Conv2d(inner, inner, 3, padding=1, groups=inner, bias=False)
+        self.norm = RMSNorm(inner)
+
+    def forward(self, x, h, w):
+        B, N, D = x.shape
+        q, k, v = self.qkv(x).chunk(3, dim=-1)
+        q = rearrange(q, 'b n (h d) -> b h n d', h=self.nh)
+        k = rearrange(k, 'b n (h d) -> b h n d', h=self.nh)
+        v = rearrange(v, 'b n (h d) -> b h n d', h=self.nh)
+        a = (q @ k.transpose(-2, -1)) * (self.hd ** -0.5)
+        o = rearrange(F.softmax(a, -1) @ v, 'b h n d -> b n (h d)')
+        o = self.norm(o)
+        # 2D locality
+        inner = self.nh * self.hd
+        lc = self.dw(rearrange(x[:, :, :inner], 'b (h w) d -> b d h w', h=h, w=w))
+        lc = rearrange(lc, 'b d h w -> b (h w) d')
+        return self.out(self.gate(x) * o + 0.1 * lc)
+
+class DenoiseBlock(nn.Module):
+    """AdaLN-modulated SpatialMixer + cross-attn + SwiGLU."""
+    def __init__(self, d, cd, nh=4, hd=32, fm=2.67):
+        super().__init__()
+        self.n1 = RMSNorm(d); self.n2 = RMSNorm(d)
+        self.mix = SpatialMixer(d, nh, hd)
+        self.ffn = SwiGLU(d, int(d * fm))
+        self.mod = nn.Sequential(nn.SiLU(), nn.Linear(cd, 6 * d))
+        # Cross-attn to condition tokens
+        self.cn = RMSNorm(d)
+        self.cq = nn.Linear(d, d, bias=False)
+        self.ckv = nn.Linear(cd, 2 * d, bias=False)
+        self.co = nn.Linear(d, d, bias=False)
+        self.cs = nn.Parameter(torch.zeros(1))
+
+    def forward(self, x, c, ctx=None, h=4, w=4):
+        s1, h1, g1, s2, h2, g2 = self.mod(c).chunk(6, -1)
+        xn = self.n1(x) * (1 + s1.unsqueeze(1)) + h1.unsqueeze(1)
+        x = x + g1.unsqueeze(1) * self.mix(xn, h, w)
+        if ctx is not None:
+            xc = self.cn(x); q = self.cq(xc)
+            k, v = self.ckv(ctx).chunk(2, -1)
+            a = F.softmax(q @ k.transpose(-2,-1) * (q.shape[-1]**-0.5), -1)
+            x = x + self.cs.tanh() * self.co(a @ v)
+        xn = self.n2(x) * (1 + s2.unsqueeze(1)) + h2.unsqueeze(1)
+        x = x + g2.unsqueeze(1) * self.ffn(xn)
+        return x
+
+class RecursiveCore(nn.Module):
+    """N shared blocks × T recursions with abstract state + IFT training."""
+    def __init__(self, lc=4, d=128, cd=128, nb=4, nh=4, hd=32,
+                 ti=2, to=1, fm=2.67, ift=False):
+        super().__init__()
+        self.nb, self.ti, self.to_, self.ift = nb, ti, to, ift
+        self.inp = nn.Linear(lc, d)
+        self.tmb = nn.Sequential(nn.Linear(256, cd), nn.SiLU(), nn.Linear(cd, cd))
+        self.blk = nn.ModuleList([DenoiseBlock(d, cd, nh, hd, fm) for _ in range(nb)])
+        self.ag = nn.Parameter(torch.tensor(0.0))
+        self.ap = nn.Sequential(nn.Linear(d, d, bias=False), nn.SiLU(), nn.Linear(d, d, bias=False))
+        self.se = nn.Embedding(to * ti + 1, cd)
+        self.on = RMSNorm(d)
+        self.op = nn.Linear(d, lc)
+        nn.init.zeros_(self.op.weight); nn.init.zeros_(self.op.bias)
+
+    def _sinemb(self, t, d=256):
+        h = d // 2
+        f = torch.exp(torch.arange(h, device=t.device).float() * -(math.log(10000) / h))
+        a = t.unsqueeze(-1) * f.unsqueeze(0)
+        return torch.cat([a.sin(), a.cos()], -1)
+
+    def _blocks(self, z, c, ctx, h, w):
+        for b in self.blk: z = b(z, c, ctx, h, w)
+        return z
+
+    def _refine(self, z, cb, ctx, h, w):
+        za = z.mean(1, keepdim=True).expand_as(z)
+        s = 0
+        for j in range(self.to_):
+            za = za + self.ag.tanh() * self.ap(z.mean(1, keepdim=True).expand_as(z))
+            for i in range(self.ti):
+                se = self.se(torch.tensor([s], device=z.device)).expand(z.shape[0], -1)
+                zn = self._blocks(z + za, cb + se, ctx, h, w)
+                z = z + 0.5 * (zn - z)
+                s += 1
+        return z
+
+    def forward(self, zt, t, tg=None, ic=None):
+        B, C, H, W = zt.shape
+        z = self.inp(rearrange(zt, 'b c h w -> b (h w) c'))
+        if ic is not None:
+            z = z + self.inp(rearrange(ic, 'b c h w -> b (h w) c'))
+        c = self.tmb(self._sinemb(t))
+        if tg is not None: c = c + tg
+        if self.training and self.ift and self.to_ > 1:
+            with torch.no_grad():
+                for _ in range(self.to_ - 1): z = self._refine(z, c, None, H, W)
+            z = self._refine(z, c, None, H, W)
+        else:
+            z = self._refine(z, c, None, H, W)
+        return rearrange(self.op(self.on(z)), 'b (h w) c -> b c h w', h=H, w=W)
+
+class LRF(nn.Module):
+    """Complete model: RecursiveCore + class conditioner."""
+    def __init__(self, cfg=None):
+        super().__init__()
+        cfg = cfg or self.default()
+        self.cfg = cfg
+        nc = cfg.get('nc', 10)
+        self.core = RecursiveCore(
+            lc=cfg['lc'], d=cfg['d'], cd=cfg['cd'], nb=cfg['nb'],
+            nh=cfg['nh'], hd=cfg['hd'], ti=cfg['ti'], to=cfg['to'],
+            fm=cfg.get('fm', 2.67), ift=cfg.get('ift', False))
+        self.cemb = nn.Embedding(nc + 1, cfg['cd'])
+        self.null = nc
+
+    @staticmethod
+    def default():
+        return dict(lc=4, d=128, cd=128, nb=4, nh=4, hd=32, ti=2, to=1, fm=2.0, ift=False, nc=10)
+
+    def predict_v(self, zt, t, cls=None, cfg_drop=0.0):
+        B = zt.shape[0]
+        if cls is not None:
+            if self.training and cfg_drop > 0:
+                m = torch.rand(B, device=zt.device) < cfg_drop
+                cls = cls.clone(); cls[m] = self.null
+            c = self.cemb(cls)
+        else:
+            c = self.cemb(torch.full((B,), self.null, device=zt.device, dtype=torch.long))
+        return self.core(zt, t, tg=c)
+
+    def count(self):
+        return sum(p.numel() for p in self.parameters())
+
+# =============================================================================
+# 2. FLOW SCHEDULER
+# =============================================================================
+
+class FlowScheduler:
+    def add_noise(self, z0, eps, t):
+        t = t.view(-1, 1, 1, 1)
+        return (1 - t) * z0 + t * eps
+
+    def velocity(self, z0, eps):
+        return eps - z0
+
+    def sample_t(self, B, dev):
+        return torch.rand(B, device=dev).clamp(1e-4, 1 - 1e-4)
+
+    @torch.no_grad()
+    def sample(self, model, shape, cls=None, steps=50, cfg=3.0, dev='cpu'):
+        z = torch.randn(shape, device=dev)
+        ts = torch.linspace(1, 0, steps + 1, device=dev)
+        for i in range(steps):
+            tb = torch.full((shape[0],), ts[i].item(), device=dev)
+            dt = ts[i] - ts[i + 1]
+            if cfg > 1.0 and cls is not None:
+                vc = model.predict_v(z, tb, cls)
+                vu = model.predict_v(z, tb, None)
+                v = vu + cfg * (vc - vu)
+            else:
+                v = model.predict_v(z, tb, cls)
+            z = z - dt * v
+        return z
+
+# =============================================================================
+# 3. DATA + TAESD
+# =============================================================================
+
+def get_taesd(dev='cpu'):
+    from diffusers import AutoencoderTiny
+    vae = AutoencoderTiny.from_pretrained('madebyollin/taesd', torch_dtype=torch.float32)
+    vae.eval().to(dev)
+    for p in vae.parameters(): p.requires_grad_(False)
+    return vae
+
+def get_cifar(root='/app/data'):
+    import torchvision, torchvision.transforms as T
+    tf = T.Compose([T.ToTensor(), T.Normalize([.5]*3, [.5]*3)])
+    tr = torchvision.datasets.CIFAR10(root, True, tf, download=True)
+    te = torchvision.datasets.CIFAR10(root, False, tf, download=True)
+    return tr, te
+
+def precompute(vae, ds, bs=256, dev='cpu', cache=None):
+    if cache and os.path.exists(cache):
+        print(f"  Loading cached latents from {cache}", flush=True)
+        d = torch.load(cache, weights_only=True)
+        return d['lat'], d['lab']
+    dl = DataLoader(ds, bs, shuffle=False, num_workers=0)
+    lats, labs = [], []
+    t0 = time.time()
+    with torch.no_grad():
+        for i, (img, lab) in enumerate(dl):
+            lats.append(vae.encode(img.to(dev)).latents.cpu())
+            labs.append(lab)
+            if (i+1) % 50 == 0 or i == 0:
+                print(f"    batch {i+1}/{len(dl)} ({time.time()-t0:.0f}s)", flush=True)
+    lats, labs = torch.cat(lats), torch.cat(labs)
+    if cache:
+        os.makedirs(os.path.dirname(cache) or '.', exist_ok=True)
+        torch.save({'lat': lats, 'lab': labs}, cache)
+    print(f"  Done: {lats.shape}, mean={lats.mean():.3f}, std={lats.std():.3f}", flush=True)
+    return lats, labs
+
+# =============================================================================
+# 4. TRAINING
+# =============================================================================
+
+def train(epochs=30, bs=64, lr=3e-4, dev=DEVICE, out='/app/lrf_out'):
+    os.makedirs(out, exist_ok=True)
+    print("=" * 60, flush=True)
+    print("LatentRecurrentFlow v3 — Training on CIFAR-10", flush=True)
+    print("=" * 60, flush=True)
+
+    # VAE
+    print("\n[1/5] Loading TAESD...", flush=True)
+    vae = get_taesd(dev)
+    print(f"  TAESD: {sum(p.numel() for p in vae.parameters()):,} params (frozen)", flush=True)
+
+    # Data
+    print("\n[2/5] Loading CIFAR-10 + precomputing latents...", flush=True)
+    tr, te = get_cifar()
+    tr_lat, tr_lab = precompute(vae, tr, 256, dev, f'{out}/cache_train.pt')
+    te_lat, te_lab = precompute(vae, te, 256, dev, f'{out}/cache_test.pt')
+
+    # Verify VAE works
+    print("\n[2b] VAE sanity check...", flush=True)
+    with torch.no_grad():
+        imgs = torch.stack([tr[i][0] for i in range(8)]).to(dev)
+        rec = vae.decode(vae.encode(imgs).latents).sample
+        mse = F.mse_loss(rec, imgs).item()
+        print(f"  Recon MSE = {mse:.4f} (should be <0.1)", flush=True)
+        save_grid(torch.cat([imgs[:4].cpu(), rec[:4].cpu()]), f'{out}/vae_check.png', 4)
+
+    # Model
+    print("\n[3/5] Creating model...", flush=True)
+    cfg = LRF.default()
+    model = LRF(cfg).to(dev)
+    print(f"  Params: {model.count():,}", flush=True)
+    print(f"  Depth: {cfg['to']}×{cfg['ti']}×{cfg['nb']} = {cfg['to']*cfg['ti']*cfg['nb']} eff. layers", flush=True)
+
+    # Train
+    print(f"\n[4/5] Training {epochs} epochs, bs={bs}, lr={lr}...", flush=True)
+    sched = FlowScheduler()
+    opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01, betas=(0.9, 0.95))
+    lrs = torch.optim.lr_scheduler.CosineAnnealingLR(opt, epochs * (len(tr_lat)//bs), lr*0.01)
+    ema = {n: p.clone().detach() for n, p in model.named_parameters()}
+    losses = []
+
+    dl = DataLoader(TensorDataset(tr_lat, tr_lab), bs, shuffle=True, drop_last=True)
+    t0 = time.time()
+    for ep in range(epochs):
+        model.train()
+        el = 0; nb = 0
+        for lat, lab in dl:
+            lat, lab = lat.to(dev), lab.to(dev)
+            B = lat.shape[0]
+            t = sched.sample_t(B, dev)
+            eps = torch.randn_like(lat)
+            zt = sched.add_noise(lat, eps, t)
+            vp = model.predict_v(zt, t, lab, cfg_drop=0.1)
+            vt = sched.velocity(lat, eps)
+            lps = (vp - vt).pow(2).mean([1,2,3])
+            w = 1.0 / (t * (1-t) + 0.01); w = w / w.mean()
+            loss = (lps * w).mean()
+            opt.zero_grad(); loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+            opt.step(); lrs.step()
+            with torch.no_grad():
+                for n, p in model.named_parameters():
+                    ema[n].mul_(0.999).add_(p, alpha=0.001)
+            el += loss.item(); nb += 1
+        al = el / nb; losses.append(al)
+        elapsed = time.time() - t0
+        if (ep+1) % 5 == 0 or ep == 0:
+            print(f"  Ep {ep+1:3d}/{epochs}: loss={al:.4f}, lr={opt.param_groups[0]['lr']:.1e}, "
+                  f"time={elapsed:.0f}s", flush=True)
+        if (ep+1) % 10 == 0 or ep == epochs - 1:
+            # Sample with EMA
+            bak = {n: p.clone() for n, p in model.named_parameters()}
+            with torch.no_grad():
+                for n, p in model.named_parameters(): p.copy_(ema[n])
+            model.eval()
+            samps = gen(model, vae, sched, dev, 16, 20, 2.0)
+            save_grid(samps, f'{out}/ep{ep+1:03d}.png', 4)
+            with torch.no_grad():
+                for n, p in model.named_parameters(): p.copy_(bak[n])
+
+    # Final: swap to EMA permanently
+    with torch.no_grad():
+        for n, p in model.named_parameters(): p.copy_(ema[n])
+    model.eval()
+
+    # Save checkpoint
+    torch.save({'state': model.state_dict(), 'cfg': cfg, 'losses': losses}, f'{out}/model.pt')
+
+    # Final generation
+    print(f"\n[5/5] Generating final samples...", flush=True)
+    classes = ['airplane','auto','bird','cat','deer','dog','frog','horse','ship','truck']
+    all_s = []
+    for ci in range(10):
+        s = gen(model, vae, sched, dev, 4, 50, 3.0, ci)
+        all_s.append(s)
+        std = s.std().item()
+        print(f"  {classes[ci]:10s}: std={std:.3f} {'✓' if std > 0.1 else '✗ FLAT'}", flush=True)
+    save_grid(torch.cat(all_s), f'{out}/final.png', 4)
+
+    # Loss plot
+    try:
+        import matplotlib; matplotlib.use('Agg'); import matplotlib.pyplot as plt
+        plt.figure(figsize=(8,3)); plt.plot(losses); plt.xlabel('Epoch'); plt.ylabel('Loss')
+        plt.title('Training Loss'); plt.grid(True, alpha=.3)
+        plt.savefig(f'{out}/loss.png', dpi=100, bbox_inches='tight'); plt.close()
+    except: pass
+
+    print(f"\n{'='*60}", flush=True)
+    print(f"DONE. Best loss: {min(losses):.4f}. Files in {out}/", flush=True)
+    print(f"{'='*60}", flush=True)
+    return model, vae, losses
+
+def gen(model, vae, sched, dev, n=8, steps=20, cfg=2.0, cls_id=None):
+    cls = torch.full((n,), cls_id if cls_id is not None else 0, dtype=torch.long, device=dev)
+    if cls_id is None: cls = torch.randint(0, 10, (n,), device=dev)
+    z = sched.sample(model, (n, 4, 4, 4), cls, steps, cfg, dev)
+    with torch.no_grad(): imgs = vae.decode(z.to(dev)).sample.clamp(-1, 1)
+    return imgs.cpu()
+
+def save_grid(imgs, path, nr=8):
+    from PIL import Image
+    imgs = ((imgs + 1) / 2).clamp(0, 1)
+    import torchvision
+    g = torchvision.utils.make_grid(imgs, nrow=nr, padding=2)
+    arr = (g.permute(1,2,0).numpy() * 255).astype(np.uint8)
+    Image.fromarray(arr).save(path)
+    print(f"  Saved: {path}", flush=True)
+
+# =============================================================================
+# 5. NOTEBOOK CONVERSION (Jupyter-compatible cells as functions)
+# =============================================================================
+
+def notebook_cell_1_setup():
+    """Cell 1: Install & import."""
+    print("Installing dependencies...")
+    os.system("pip install -q torch torchvision einops diffusers safetensors huggingface_hub matplotlib pillow")
+    print("Done.")
+
+def notebook_cell_2_architecture():
+    """Cell 2: Show architecture details."""
+    cfg = LRF.default()
+    model = LRF(cfg)
+    print(f"LatentRecurrentFlow Architecture")
+    print(f"================================")
+    print(f"Latent channels: {cfg['lc']} (TAESD)")
+    print(f"Model dim:       {cfg['d']}")
+    print(f"Shared blocks:   {cfg['nb']}")
+    print(f"Recursions:      {cfg['to']}×{cfg['ti']} = {cfg['to']*cfg['ti']}")
+    print(f"Effective depth: {cfg['to']*cfg['ti']*cfg['nb']} layers")
+    print(f"Total params:    {model.count():,}")
+    print(f"FP32 size:       {model.count()*4/1e6:.1f} MB")
+    print(f"INT8 size:       {model.count()/1e6:.1f} MB")
+    return model
+
+def notebook_cell_3_train():
+    """Cell 3: Full training loop."""
+    return train(epochs=30, bs=64, lr=3e-4, out='/app/lrf_out')
+
+def notebook_cell_4_generate(model, vae):
+    """Cell 4: Generate and display images."""
+    sched = FlowScheduler()
+    classes = ['airplane','auto','bird','cat','deer','dog','frog','horse','ship','truck']
+    for ci, cn in enumerate(classes):
+        imgs = gen(model, vae, sched, DEVICE, 4, 50, 3.0, ci)
+        save_grid(imgs, f'/app/lrf_out/class_{cn}.png', 4)
+    print("All class images generated!")
+
+def notebook_cell_5_push(repo_id='krystv/LatentRecurrentFlow'):
+    """Cell 5: Push to HuggingFace Hub."""
+    from huggingface_hub import HfApi
+    api = HfApi()
+    out = '/app/lrf_out'
+    for f in os.listdir(out):
+        if f.endswith(('.pt', '.png', '.json')):
+            fp = os.path.join(out, f)
+            if os.path.getsize(fp) < 50_000_000:  # Skip huge files
+                api.upload_file(path_or_fileobj=fp, path_in_repo=f'v3/{f}',
+                                repo_id=repo_id, repo_type='model')
+                print(f"  Uploaded v3/{f}")
+    print(f"Done! See https://huggingface.co/{repo_id}")
+
+# =============================================================================
+# MAIN
+# =============================================================================
+
+if __name__ == '__main__':
+    model, vae, losses = train(epochs=30, bs=64, lr=3e-4, out='/app/lrf_out')