Upload train_gpu.py with huggingface_hub

train_gpu.py

@@ -0,0 +1,174 @@
+#!/usr/bin/env python3
+"""
+LRF Extended Training — more epochs on CPU with cached latents.
+Uses the same proven architecture from v3, just trains longer.
+Pushes results to HF Hub.
+"""
+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
+import numpy as np
+
+DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+# Reuse the exact architecture from lrf_v3.py
+sys.path.insert(0, '/app')
+from lrf_v3 import LRF, FlowScheduler, get_taesd, get_cifar, precompute, save_grid, gen
+
+def main():
+    OUT = '/app/lrf_extended'
+    REPO = 'krystv/LatentRecurrentFlow'
+    os.makedirs(OUT, exist_ok=True)
+
+    EPOCHS = 100  # 3x more than v3
+    BS = 128      # Bigger batch
+    LR = 5e-4     # Slightly higher LR for faster convergence
+
+    print("=" * 60, flush=True)
+    print(f"LRF Extended Training — {EPOCHS} epochs, bs={BS}", flush=True)
+    print(f"Device: {DEVICE}", flush=True)
+    print("=" * 60, flush=True)
+
+    # VAE + Data (use cached latents from v3 if available)
+    print("\n[1] Loading TAESD + CIFAR-10...", flush=True)
+    vae = get_taesd(DEVICE)
+    tr, te = get_cifar()
+
+    # Check for cached latents from previous run
+    cache_dir = '/app/lrf_out'
+    if os.path.exists(f'{cache_dir}/cache_train.pt'):
+        print("  Using cached latents from v3 run!", flush=True)
+        tr_lat, tr_lab = precompute(vae, tr, 256, DEVICE, f'{cache_dir}/cache_train.pt')
+    else:
+        tr_lat, tr_lab = precompute(vae, tr, 256, DEVICE, f'{OUT}/cache_train.pt')
+
+    # Model — use the proven fast config
+    print("\n[2] Creating model...", flush=True)
+    cfg = LRF.default()
+    model = LRF(cfg).to(DEVICE)
+    print(f"  {model.count():,} params", flush=True)
+
+    # Try to warm-start from v3 checkpoint
+    v3_ckpt = '/app/lrf_out/model.pt'
+    if os.path.exists(v3_ckpt):
+        print(f"  Warm-starting from {v3_ckpt}", flush=True)
+        ckpt = torch.load(v3_ckpt, map_location=DEVICE, weights_only=False)
+        model.load_state_dict(ckpt['state'])
+        prev_losses = ckpt.get('losses', [])
+        print(f"  Previous best loss: {min(prev_losses):.4f}", flush=True)
+    else:
+        prev_losses = []
+
+    # Train
+    print(f"\n[3] Training {EPOCHS} epochs...", flush=True)
+    sched = FlowScheduler()
+    opt = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=0.01, betas=(0.9, 0.95))
+    total_steps = EPOCHS * (len(tr_lat) // BS)
+    lr_sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, total_steps, LR * 0.01)
+    ema = {n: p.clone().detach() for n, p in model.named_parameters()}
+    losses = list(prev_losses)  # Continue loss history
+
+    dl = DataLoader(TensorDataset(tr_lat, tr_lab), BS, shuffle=True, drop_last=True, num_workers=0)
+    best_loss = min(losses) if losses else 999
+    t0 = time.time()
+
+    for ep in range(EPOCHS):
+        model.train()
+        el, nb = 0, 0
+        for lat, lab in dl:
+            lat, lab = lat.to(DEVICE), lab.to(DEVICE)
+            B = lat.shape[0]
+            t = sched.sample_t(B, DEVICE)
+            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(); lr_sched.step()
+            with torch.no_grad():
+                for n, p in model.named_parameters():
+                    ema[n].mul_(0.9995).add_(p, alpha=0.0005)
+            el += loss.item(); nb += 1
+
+        al = el / nb
+        losses.append(al)
+        if al < best_loss: best_loss = al
+        elapsed = time.time() - t0
+
+        if (ep+1) % 10 == 0 or ep == 0 or ep == EPOCHS-1:
+            print(f"  Ep {ep+1:3d}/{EPOCHS}: loss={al:.4f} best={best_loss:.4f} "
+                  f"lr={opt.param_groups[0]['lr']:.1e} {elapsed:.0f}s", flush=True)
+
+        # Sample every 25 epochs
+        if (ep+1) % 25 == 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, DEVICE, 16, 20, 2.5)
+            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 EMA
+    with torch.no_grad():
+        for n, p in model.named_parameters(): p.copy_(ema[n])
+    model.eval()
+
+    # Final generation
+    print(f"\n[4] Final generation...", 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, DEVICE, 8, 50, 3.0, ci)
+        all_s.append(s)
+        print(f"  {classes[ci]:10s}: std={s.std():.3f}", flush=True)
+    save_grid(torch.cat(all_s), f'{OUT}/final.png', 8)
+
+    # Save
+    torch.save({'state': model.state_dict(), 'cfg': cfg, 'losses': losses}, f'{OUT}/model.pt')
+
+    # Loss plot
+    try:
+        import matplotlib; matplotlib.use('Agg'); import matplotlib.pyplot as plt
+        plt.figure(figsize=(10,4))
+        plt.plot(losses, 'b-', alpha=0.7)
+        if prev_losses:
+            plt.axvline(x=len(prev_losses), color='r', linestyle='--', alpha=0.5, label='Extended training start')
+            plt.legend()
+        plt.xlabel('Epoch'); plt.ylabel('Loss')
+        plt.title(f'LRF Training (best={best_loss:.4f})')
+        plt.grid(True, alpha=0.3)
+        plt.savefig(f'{OUT}/loss.png', dpi=150, bbox_inches='tight'); plt.close()
+    except: pass
+
+    # Push to Hub
+    print(f"\n[5] Pushing to Hub...", flush=True)
+    from huggingface_hub import HfApi
+    api = HfApi()
+    for f in sorted(os.listdir(OUT)):
+        fp = os.path.join(OUT, f)
+        if f.endswith(('.pt', '.png')) and os.path.getsize(fp) < 100_000_000 and 'cache' not in f:
+            api.upload_file(path_or_fileobj=fp, path_in_repo=f'gpu_trained/{f}',
+                            repo_id=REPO, repo_type='model')
+            print(f"    Uploaded gpu_trained/{f}", flush=True)
+
+    # Upload train script
+    api.upload_file(path_or_fileobj='/app/train_extended.py', path_in_repo='train_gpu.py',
+                    repo_id=REPO, repo_type='model')
+    print(f"    Uploaded train_gpu.py", flush=True)
+
+    print(f"\n{'='*60}", flush=True)
+    print(f"DONE! Best loss: {best_loss:.4f}", flush=True)
+    print(f"See: https://huggingface.co/{REPO}", flush=True)
+    print(f"{'='*60}", flush=True)
+
+if __name__ == '__main__':
+    main()