Spaces:
Running on Zero
Running on Zero
| """Config, device, model construction and checkpoint helpers shared by the scripts.""" | |
| import argparse | |
| import os | |
| import torch | |
| import yaml | |
| from .thinker import Thinker | |
| from .unet import UNet | |
| def parse_args(extra=None): | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--config", default="configs/paper.yaml") | |
| ap.add_argument("--mode", choices=["hard", "extreme"], default=None, help="overrides config mode") | |
| ap.add_argument("--set", nargs="*", default=[], help="overrides like thinker.steps=100") | |
| for args, kw in (extra or []): | |
| ap.add_argument(*args, **kw) | |
| args = ap.parse_args() | |
| with open(args.config) as f: | |
| cfg = yaml.safe_load(f) | |
| for kv in args.set: | |
| k, v = kv.split("=", 1) | |
| d = cfg | |
| *path, last = k.split(".") | |
| for p in path: | |
| d = d[p] | |
| d[last] = yaml.safe_load(v) | |
| if args.mode: | |
| cfg["mode"] = args.mode | |
| os.makedirs(cfg["run_dir"], exist_ok=True) | |
| return args, cfg | |
| def device(): | |
| if torch.cuda.is_available(): | |
| torch.backends.cuda.matmul.allow_tf32 = torch.backends.cudnn.allow_tf32 = True | |
| torch.backends.cudnn.benchmark = True | |
| return torch.device("cuda") | |
| if torch.backends.mps.is_available(): | |
| return torch.device("mps") | |
| return torch.device("cpu") | |
| def autocast(dev): | |
| """bf16 autocast on CUDA (as in TRM); a no-op elsewhere.""" | |
| return torch.autocast(dev.type, torch.bfloat16, enabled=dev.type == "cuda") | |
| def compile(fn, dev): | |
| return torch.compile(fn) if dev.type == "cuda" else fn | |
| def build_painter(cfg): | |
| p = cfg["painter"] | |
| return UNet(1, 1, tuple(p["channels"]), p["layers_per_block"]) | |
| def build_dm(cfg): | |
| p = cfg["dm"] | |
| return UNet(3, 1, tuple(p["channels"]), p["layers_per_block"]) | |
| def build_thinker(cfg, painter): | |
| t = cfg["thinker"] | |
| return Thinker(painter, dim=t["dim"], heads=t["heads"], layers=t["layers"], expansion=t["expansion"], | |
| grid=t["grid"], enc_hidden=tuple(t["enc_hidden"]), H=t["H"], L=t["L"]) | |
| def load_eval_models(cfg, kind, mode, dev): | |
| """EMA weights in eval mode: {'dm': ...} or {'painter': ..., 'thinker': ...} for kind 'dm' / 'path'.""" | |
| def ema(model, name): | |
| model.load_state_dict(load(ckpt_path(cfg, name), dev)["ema"]) | |
| return model.to(dev).eval() | |
| if kind == "dm": | |
| return {"dm": ema(build_dm(cfg), f"dm_{mode}")} | |
| painter = ema(build_painter(cfg), "painter") | |
| return {"painter": painter, "thinker": ema(build_thinker(cfg, painter), f"thinker_{mode}")} | |
| def n_params(m): | |
| return sum(p.numel() for p in m.parameters()) | |
| def ckpt_path(cfg, name): | |
| return os.path.join(cfg["run_dir"], f"{name}.pt") | |
| def save(path, **state): | |
| tmp = path + ".tmp" | |
| torch.save(state, tmp) | |
| os.replace(tmp, path) | |
| def load(path, dev): | |
| return torch.load(path, map_location=dev, weights_only=False) | |