"""Platform-agnostic learned-gate training with HF Hub warm-start. Strips Modal decorators from learned_gate_warm_modal.py. Runs on: - Lightning.ai Studio (interactive or CLI Job) - Any local GPU box - Any cloud VM with PyTorch+CUDA+internet Reads HF_TOKEN from env. Saves to local --output-dir AND pushes to HF Hub (same {arm_tag}/latest.pt convention as the Modal version). On launch, tries to pull latest.pt from HF Hub for this arm and resume. Usage: pip install torch torchvision Pillow numpy 'numpy<2.0' tqdm clean-fid huggingface_hub export HF_TOKEN=hf_... # token with write to {hf_repo} python learned_gate_standalone.py \ --dataset edges2shoes --schedule linear \ --max-iters 30000 --max-wall-secs 10800 \ --base-channels 128 --channel-mults 1,2,4,4 \ --batch-size 128 --output-dir ./out """ from __future__ import annotations import argparse import json import math import os import subprocess import sys import tarfile import time from pathlib import Path PALETTE = [ (228, 26, 28), (55, 126, 184), (77, 175, 74), (152, 78, 163), (255, 127, 0), (255, 217, 47), (166, 86, 40), (247, 129, 191), (102, 194, 165), (179, 179, 179), ] VALID_SCHEDULES = ( "linear", "cosine", "sigmoid", "poly2", "sqrt", "exp", "tanh2", "quartic", "polyk1_5", "polyk1_3", "expl1_5", "expl1_0", "expl2_0", ) VALID_DATASETS = ("colorize_mnist", "edges2shoes", "night2day") PIX2PIX_URLS = { "edges2shoes": "http://efrosgans.eecs.berkeley.edu/pix2pix/datasets/edges2shoes.tar.gz", "night2day": "http://efrosgans.eecs.berkeley.edu/pix2pix/datasets/night2day.tar.gz", } def ensure_dataset_local(ds_name: str, data_root: Path): """Download dataset to data_root if not already present.""" if ds_name == "colorize_mnist": return target = data_root / ds_name if (target / "train").exists(): n = len(list((target / "train").iterdir())) print(f"[data] {ds_name} present at {target} ({n} train imgs)", flush=True) return url = PIX2PIX_URLS.get(ds_name) if url is None: raise RuntimeError(f"no auto-download URL for {ds_name}") data_root.mkdir(parents=True, exist_ok=True) tar_path = data_root / f"{ds_name}.tar.gz" print(f"[data] downloading {url}", flush=True) # Use wget if available, else urllib try: subprocess.run(["wget", "-q", "-O", str(tar_path), url], check=True) except (FileNotFoundError, subprocess.CalledProcessError): import urllib.request urllib.request.urlretrieve(url, tar_path) print(f"[data] downloaded {tar_path.stat().st_size / 1e6:.0f} MB; extracting", flush=True) with tarfile.open(str(tar_path), "r:gz") as tar: tar.extractall(str(data_root)) tar_path.unlink() n = len(list((target / "train").iterdir())) if (target / "train").exists() else 0 print(f"[data] {ds_name} ready ({n} train imgs)", flush=True) def train_one(args): import random import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import Dataset, DataLoader from torchvision import datasets as tvds, transforms from torchvision.utils import save_image from cleanfid import fid from huggingface_hub import HfApi, hf_hub_download # Allow importing model.py and dataset.py from this script's directory sys.path.insert(0, str(Path(__file__).resolve().parent)) from model import ResBlock, SelfAttention, Downsample, Upsample, SinusoidalPosEmb schedule = args.schedule dataset = args.dataset assert schedule in VALID_SCHEDULES, f"bad schedule: {schedule}" assert dataset in VALID_DATASETS, f"bad dataset: {dataset}" channel_mults = tuple(int(x) for x in args.channel_mults.split(",")) attn_res = tuple(int(x) for x in args.attn_res.split(",")) random.seed(args.seed); np.random.seed(args.seed) torch.manual_seed(args.seed); torch.cuda.manual_seed_all(args.seed) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") arm_tag = f"{dataset}__{schedule}" out_root = Path(args.output_dir) / "gate_warm" / arm_tag out_root.mkdir(parents=True, exist_ok=True) ckpt_path = out_root / "latest.pt" print(f"[gate-warm] arm={arm_tag} max_iters={args.max_iters} wall_cap={args.max_wall_secs}s device={device}", flush=True) # ---------------- HF Hub ---------------- hf_token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_TOKEN") or os.environ.get("HUGGINGFACE_HUB_TOKEN") hf_api = HfApi(token=hf_token) if hf_token else HfApi() def hf_ckpt_path_in_repo(): return f"{arm_tag}/latest.pt" def hf_pull_latest() -> bool: try: p = hf_hub_download(repo_id=args.hf_repo, filename=hf_ckpt_path_in_repo(), token=hf_token, repo_type="model") import shutil shutil.copy2(p, ckpt_path) print(f"[hf] resumed: pulled {args.hf_repo}/{hf_ckpt_path_in_repo()}", flush=True) return True except Exception as e: print(f"[hf] no resume ckpt ({type(e).__name__}: {e})", flush=True) return False def hf_push_latest(): try: hf_api.upload_file( path_or_fileobj=str(ckpt_path), path_in_repo=hf_ckpt_path_in_repo(), repo_id=args.hf_repo, repo_type="model", token=hf_token, ) print(f"[hf] pushed -> {args.hf_repo}/{hf_ckpt_path_in_repo()}", flush=True) except Exception as e: print(f"[hf] push failed: {type(e).__name__}: {e}", flush=True) if hf_token: try: hf_api.create_repo(repo_id=args.hf_repo, repo_type="model", private=True, exist_ok=True) except Exception as e: print(f"[hf] create_repo warning: {e}", flush=True) # ---------------- Schedule math ---------------- def schedule_coeffs(t: torch.Tensor, name: str): t = t.clamp(1e-4, 1 - 1e-4) if name == "linear": a = 1.0 - t elif name == "cosine": a = torch.cos(math.pi * t / 2) elif name == "sigmoid": k = 10.0 s = lambda x: 1.0 / (1.0 + torch.exp(-k * (x - 0.5))) s0 = s(torch.zeros_like(t)); s1 = s(torch.ones_like(t)) a = 1.0 - (s(t) - s0) / (s1 - s0) elif name == "poly2": a = (1.0 - t) ** 2 elif name == "sqrt": a = torch.sqrt(1.0 - t) elif name == "exp": lam = 3.0 a = (torch.exp(-lam * t) - math.exp(-lam)) / (1.0 - math.exp(-lam)) elif name == "tanh2": k = 2.0 s = lambda x: 0.5 * (1.0 + torch.tanh(k * (x - 0.5))) s0 = s(torch.zeros_like(t)); s1 = s(torch.ones_like(t)) a = 1.0 - (s(t) - s0) / (s1 - s0) elif name == "quartic": a = 1.0 - t ** 4 elif name == "polyk1_5": a = (1.0 - t) ** 1.5 elif name == "polyk1_3": a = (1.0 - t) ** 1.3 elif name == "expl1_5": lam = 1.5 a = (torch.exp(-lam * t) - math.exp(-lam)) / (1.0 - math.exp(-lam)) elif name == "expl1_0": lam = 1.0 a = (torch.exp(-lam * t) - math.exp(-lam)) / (1.0 - math.exp(-lam)) elif name == "expl2_0": lam = 2.0 a = (torch.exp(-lam * t) - math.exp(-lam)) / (1.0 - math.exp(-lam)) else: raise ValueError(name) return a, 1.0 - a def sigma_ref(t): t = t.clamp(1e-4, 1 - 1e-4) return torch.sqrt(2.0 * t / (1.0 - t)) def t_from_sigma(sigma): s2 = sigma ** 2 return s2 / (2.0 + s2) # ---------------- Model ---------------- sigma_data = args.sigma_data class GatedSimpleUNet(nn.Module): def __init__(self, base_channels, channel_mults, attn_resolutions, t_dim=256, image_size=64, num_heads=4): super().__init__() self.sigma_data = sigma_data self.schedule_name = schedule self.time_mlp = nn.Sequential( SinusoidalPosEmb(t_dim), nn.Linear(t_dim, t_dim * 4), nn.SiLU(), nn.Linear(t_dim * 4, t_dim), ) gate_emb = 64 self.gate_mlp = nn.Sequential( SinusoidalPosEmb(gate_emb), nn.Linear(gate_emb, gate_emb), nn.SiLU(), nn.Linear(gate_emb, 1), ) ch = base_channels self.embed_u = nn.Conv2d(6, ch, 3, padding=1) self.embed_x = nn.Conv2d(6, ch, 3, padding=1) self.enc_blocks = nn.ModuleList() self.enc_attns = nn.ModuleList() self.enc_downs = nn.ModuleList() skip_channels = [ch]; cur_ch = ch; cur_res = image_size for i, mult in enumerate(channel_mults): out_ch = ch * mult self.enc_blocks.append(ResBlock(cur_ch, out_ch, t_dim, 0.0)) cur_ch = out_ch skip_channels.append(cur_ch) if cur_res in attn_resolutions: self.enc_attns.append(SelfAttention(cur_ch, num_heads)) else: self.enc_attns.append(nn.Identity()) if i < len(channel_mults) - 1: self.enc_downs.append(Downsample(cur_ch)) cur_res //= 2 else: self.enc_downs.append(nn.Identity()) self.mid1 = ResBlock(cur_ch, cur_ch, t_dim, 0.0) self.mid_attn = SelfAttention(cur_ch, num_heads) self.mid2 = ResBlock(cur_ch, cur_ch, t_dim, 0.0) self.dec_blocks = nn.ModuleList() self.dec_attns = nn.ModuleList() self.dec_ups = nn.ModuleList() for i in reversed(range(len(channel_mults))): mult = channel_mults[i] out_ch = ch * mult skip_ch = skip_channels.pop() self.dec_blocks.append(ResBlock(cur_ch + skip_ch, out_ch, t_dim, 0.0)) cur_ch = out_ch dec_res = image_size // (2 ** i) if i < len(channel_mults) - 1 else cur_res if dec_res in attn_resolutions: self.dec_attns.append(SelfAttention(cur_ch, num_heads)) else: self.dec_attns.append(nn.Identity()) if i > 0: self.dec_ups.append(Upsample(cur_ch)) else: self.dec_ups.append(nn.Identity()) self.out_norm = nn.GroupNorm(min(32, cur_ch), cur_ch) self.out_conv = nn.Conv2d(cur_ch, 3, 3, padding=1) nn.init.zeros_(self.out_conv.weight); nn.init.zeros_(self.out_conv.bias) def gate(self, t): return torch.sigmoid(self.gate_mlp(t).squeeze(-1)) def trunk(self, h, t_emb): skips = [h] for block, attn, down in zip(self.enc_blocks, self.enc_attns, self.enc_downs): h = block(h, t_emb); h = attn(h); skips.append(h); h = down(h) h = self.mid1(h, t_emb); h = self.mid_attn(h); h = self.mid2(h, t_emb) for block, attn, up in zip(self.dec_blocks, self.dec_attns, self.dec_ups): h = torch.cat([h, skips.pop()], dim=1) h = block(h, t_emb); h = attn(h); h = up(h) h = F.silu(self.out_norm(h)) return self.out_conv(h) def forward(self, u, X_1, sigma): B = u.shape[0]; sd = self.sigma_data sig = sigma.view(-1, 1, 1, 1).float() c_in_u = 1.0 / (sig ** 2 + sd ** 2).sqrt() c_skip = sd ** 2 / (sig ** 2 + sd ** 2) c_out = sig * sd / (sig ** 2 + sd ** 2).sqrt() c_noise = sig.log() / 4.0 t = t_from_sigma(sigma.flatten()) a_t, b_t = schedule_coeffs(t, self.schedule_name) a_b = a_t.view(-1, 1, 1, 1); b_b = b_t.view(-1, 1, 1, 1) X_t = a_b * u + b_b * X_1 var_xt = (a_b ** 2 + b_b ** 2) * sd ** 2 + a_b ** 2 * sig ** 2 c_in_x = 1.0 / var_xt.sqrt() u_in = torch.cat([c_in_u * u, X_1], dim=1) x_in = torch.cat([c_in_x * X_t, torch.zeros_like(X_1)], dim=1) h_u = self.embed_u(u_in); h_x = self.embed_x(x_in) g = self.gate(t).view(-1, 1, 1, 1) h = g * h_u + (1.0 - g) * h_x t_emb = self.time_mlp(c_noise.flatten()) F_out = self.trunk(h, t_emb) D = c_skip * u + c_out * F_out return D # ---------------- Dataset ---------------- class ColorizeMNIST(Dataset): def __init__(self, base_ds, indices): self.base = base_ds; self.indices = list(indices) self.pal = torch.tensor(PALETTE, dtype=torch.float32) / 127.5 - 1.0 def __len__(self): return len(self.indices) def __getitem__(self, i): img, label = self.base[self.indices[i]] img = F.pad(img, (2, 2, 2, 2), value=0.0) gray = img.expand(3, -1, -1) * 2.0 - 1.0 intensity = img * 2.0 - 1.0 color = self.pal[label].view(3, 1, 1) weight = (intensity + 1.0) * 0.5 bg = torch.full_like(gray, -1.0) return gray, bg + (color - bg) * weight data_root = Path(args.data_dir) ensure_dataset_local(dataset, data_root) if dataset == "colorize_mnist": tf = transforms.ToTensor() train_base = tvds.MNIST(str(data_root / "mnist"), train=True, download=True, transform=tf) test_base = tvds.MNIST(str(data_root / "mnist"), train=False, download=True, transform=tf) train_ds = ColorizeMNIST(train_base, indices=range(min(args.n_train, len(train_base)))) eval_ds = ColorizeMNIST(test_base, indices=range(min(args.n_eval, len(test_base)))) else: from dataset import PairedDataset ds_root = str(data_root / dataset) train_ds = PairedDataset(root=ds_root, split="train", image_size=args.image_size, augment=False, format="auto") eval_split = "test" if dataset == "night2day" else "val" eval_ds = PairedDataset(root=ds_root, split=eval_split, image_size=args.image_size, augment=False, format="auto") if args.n_train < len(train_ds): train_ds = torch.utils.data.Subset(train_ds, range(args.n_train)) if args.n_eval < len(eval_ds): eval_ds = torch.utils.data.Subset(eval_ds, range(args.n_eval)) print(f"[data] train={len(train_ds)} eval={len(eval_ds)}", flush=True) g = torch.Generator(); g.manual_seed(args.seed) train_loader = DataLoader(train_ds, batch_size=args.batch_size, shuffle=True, num_workers=2, pin_memory=True, drop_last=True, generator=g) model = GatedSimpleUNet(base_channels=args.base_channels, channel_mults=channel_mults, attn_resolutions=attn_res, image_size=args.image_size).to(device) n_params = sum(p.numel() for p in model.parameters() if p.requires_grad) print(f"[model] params={n_params/1e6:.2f}M base={args.base_channels} mults={channel_mults} attn={attn_res}", flush=True) opt = torch.optim.AdamW(model.parameters(), lr=args.lr, betas=(0.9, 0.999)) # ---------------- Resume from HF ---------------- start_iter = 0 if args.resume_from_hf: if hf_pull_latest() and ckpt_path.exists(): try: state = torch.load(ckpt_path, map_location=device, weights_only=False) model.load_state_dict(state["model"]) opt.load_state_dict(state["opt"]) start_iter = int(state.get("iter", 0)) print(f"[resume] loaded ckpt at iter {start_iter}", flush=True) except Exception as e: print(f"[resume] load failed ({type(e).__name__}: {e}); starting fresh", flush=True) start_iter = 0 def save_ckpt(it: int): state = { "model": model.state_dict(), "opt": opt.state_dict(), "iter": it, "schedule": schedule, "dataset": dataset, "sigma_data": sigma_data, "base_channels": args.base_channels, "channel_mults": list(channel_mults), "attn_res": list(attn_res), "image_size": args.image_size, } try: torch.save(state, ckpt_path) hf_push_latest() except Exception as e: print(f"[ckpt] save failed: {e}", flush=True) # ---------------- Training ---------------- t0 = time.time() iter_count = start_iter print(f"[train] starting at iter={iter_count} target={args.max_iters}", flush=True) P_mean, P_std = args.P_mean, args.P_std while iter_count < args.max_iters: if (time.time() - t0) >= args.max_wall_secs: print(f"[wall-cap] reached {args.max_wall_secs}s at iter={iter_count}; saving and exiting", flush=True) save_ckpt(iter_count) return {"iters_done": iter_count, "fid": None, "wall_secs": time.time() - t0, "reason": "wall_cap"} model.train() for X1, X0 in train_loader: if iter_count >= args.max_iters or (time.time() - t0) >= args.max_wall_secs: break X0 = X0.to(device, non_blocking=True); X1 = X1.to(device, non_blocking=True) B = X0.shape[0] rnd = torch.randn([B, 1, 1, 1], device=device) sigma = (rnd * P_std + P_mean).exp() weight = (sigma ** 2 + sigma_data ** 2) / (sigma * sigma_data) ** 2 z = torch.randn_like(X0) u = X0 + sigma * z opt.zero_grad() D = model(u, X1, sigma.flatten()) loss = (weight * (D - X0) ** 2).mean() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) opt.step() iter_count += 1 if iter_count % args.save_every_iters == 0: elapsed = (time.time() - t0) / 60 print(f"[iter {iter_count}/{args.max_iters}] loss={loss.item():.4f} wall={elapsed:.1f}min", flush=True) save_ckpt(iter_count) print(f"[done-train] iter={iter_count} wall={(time.time()-t0)/60:.1f}min; running final eval", flush=True) save_ckpt(iter_count) model.eval() with torch.no_grad(): t_grid = torch.linspace(1e-3, 1 - 1e-3, 200, device=device) g_vals = model.gate(t_grid).cpu().numpy() a_vals, b_vals = schedule_coeffs(t_grid.cpu(), schedule) sigma_vals = sigma_ref(t_grid.cpu()).numpy() gate_data = { "t": t_grid.cpu().numpy().tolist(), "g": g_vals.tolist(), "a": a_vals.numpy().tolist(), "b": b_vals.numpy().tolist(), "sigma": sigma_vals.tolist(), "schedule": schedule, "dataset": dataset, "iter": iter_count, } (out_root / "gate_curve.json").write_text(json.dumps(gate_data)) print(f"[gate] saved: g(0.01)={g_vals[1]:.3f} g(0.5)={g_vals[100]:.3f} g(0.99)={g_vals[-2]:.3f}", flush=True) @torch.no_grad() def heun_sample(X1, num_steps=35, sigma_min=0.002, sigma_max=80.0, rho=7.0): device_ = X1.device idx = torch.arange(num_steps, dtype=torch.float64, device=device_) t_steps = (sigma_max ** (1 / rho) + idx / (num_steps - 1) * (sigma_min ** (1 / rho) - sigma_max ** (1 / rho))) ** rho t_steps = torch.cat([t_steps, torch.zeros_like(t_steps[:1])]) x_next = torch.randn_like(X1, dtype=torch.float64) * t_steps[0] for i in range(num_steps): t_cur = t_steps[i]; t_next_ = t_steps[i + 1] x_cur = x_next sig = t_cur.to(torch.float32).expand(X1.shape[0]) D = model(x_cur.to(torch.float32), X1, sig).to(torch.float64) d_cur = (x_cur - D) / t_cur x_next = x_cur + (t_next_ - t_cur) * d_cur if i < num_steps - 1: sig2 = t_next_.to(torch.float32).expand(X1.shape[0]) D2 = model(x_next.to(torch.float32), X1, sig2).to(torch.float64) d_prime = (x_next - D2) / t_next_ x_next = x_cur + (t_next_ - t_cur) * 0.5 * (d_cur + d_prime) return x_next.to(torch.float32) out_png = out_root / "samples"; real_png = out_root / "real" out_png.mkdir(exist_ok=True); real_png.mkdir(exist_ok=True) val_loader = DataLoader(eval_ds, batch_size=64, shuffle=False, num_workers=2) count = 0 print(f"[sample] generating {args.fid_samples} samples NFE={args.nfe}", flush=True) for X1_val, X0_val in val_loader: X0_val = X0_val.to(device); X1_val = X1_val.to(device) gen = heun_sample(X1_val, num_steps=args.nfe) gen01 = (gen.clamp(-1, 1) + 1) / 2 real01 = (X0_val.clamp(-1, 1) + 1) / 2 for i in range(gen.shape[0]): if count >= args.fid_samples: break try: save_image(gen01[i], out_png / f"{count:06d}.png") save_image(real01[i], real_png / f"{count:06d}.png") except OSError: pass count += 1 if count >= args.fid_samples: break try: fid_val = float(fid.compute_fid(str(out_png), str(real_png), mode="clean", num_workers=0)) except Exception as e: print(f"[fid] failed: {e}", flush=True); fid_val = float('nan') summary = { "schedule": schedule, "dataset": dataset, "FID": fid_val, "iters_done": iter_count, "nfe": args.nfe, "n_samples": args.fid_samples, "sigma_data": sigma_data, "base_channels": args.base_channels, "channel_mults": list(channel_mults), "attn_res": list(attn_res), "image_size": args.image_size, "n_params_M": n_params / 1e6, "wall_secs": time.time() - t0, } (out_root / "fid_summary.json").write_text(json.dumps(summary, indent=2)) print(f"[done] {arm_tag} FID={fid_val:.3f} iters={iter_count} wall={(time.time()-t0)/60:.1f}min", flush=True) return {"iters_done": iter_count, "fid": fid_val, "wall_secs": time.time() - t0, "reason": "max_iters"} def parse_args(): p = argparse.ArgumentParser() p.add_argument("--schedule", required=True) p.add_argument("--dataset", required=True, choices=VALID_DATASETS) p.add_argument("--image-size", type=int, default=64) p.add_argument("--base-channels", type=int, default=128) p.add_argument("--channel-mults", default="1,2,4,4") p.add_argument("--attn-res", default="16,8") p.add_argument("--batch-size", type=int, default=128) p.add_argument("--sigma-data", type=float, default=0.5) p.add_argument("--P-mean", type=float, default=-1.2) p.add_argument("--P-std", type=float, default=1.2) p.add_argument("--lr", type=float, default=2e-4) p.add_argument("--max-iters", type=int, default=30000) p.add_argument("--max-wall-secs", type=int, default=10800) # 3 hrs default for Lightning free tier p.add_argument("--save-every-iters", type=int, default=1000) p.add_argument("--fid-samples", type=int, default=2000) p.add_argument("--nfe", type=int, default=35) p.add_argument("--n-train", type=int, default=60000) p.add_argument("--n-eval", type=int, default=2000) p.add_argument("--seed", type=int, default=0) p.add_argument("--hf-repo", default="augustander/bci-gate-warm") p.add_argument("--resume-from-hf", action="store_true", default=True) p.add_argument("--no-resume", dest="resume_from_hf", action="store_false") p.add_argument("--output-dir", default="./out") p.add_argument("--data-dir", default="./data") return p.parse_args() if __name__ == "__main__": args = parse_args() result = train_one(args) print(json.dumps(result, indent=2))