| |
| """ |
| 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}") |
|
|
| |
| |
| |
|
|
| 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) |
| |
| 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)) |
| |
| 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()) |
|
|
| |
| |
| |
|
|
| 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 |
|
|
| |
| |
| |
|
|
| 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 |
|
|
| |
| |
| |
|
|
| 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) |
|
|
| |
| 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) |
|
|
| |
| 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') |
|
|
| |
| 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) |
|
|
| |
| 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) |
|
|
| |
| 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: |
| |
| 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]) |
|
|
| |
| with torch.no_grad(): |
| for n, p in model.named_parameters(): p.copy_(ema[n]) |
| model.eval() |
|
|
| |
| torch.save({'state': model.state_dict(), 'cfg': cfg, 'losses': losses}, f'{out}/model.pt') |
|
|
| |
| 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) |
|
|
| |
| 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) |
|
|
| |
| |
| |
|
|
| 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: |
| 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}") |
|
|
| |
| |
| |
|
|
| if __name__ == '__main__': |
| model, vae, losses = train(epochs=30, bs=64, lr=3e-4, out='/app/lrf_out') |
|
|
