code/learned_gate_standalone.py

"""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))