from __future__ import annotations import argparse import copy import json import math import os import time import numpy as np import torch import torch.nn.functional as F from safetensors.torch import save_file from voxel_dit import VoxelDiT def cosine_lr(step, total, base, floor, warmup): if step < warmup: return base * (step + 1) / max(1, warmup) t = (step - warmup) / max(1, total - warmup) return floor + 0.5 * (base - floor) * (1 + math.cos(math.pi * t)) def load_split(data, filtered, n_val, seed=0): meta = json.load(open(f"{data}/meta.json")) keep = np.array(meta["keep"], bool) idx = np.nonzero(keep)[0] if filtered else np.arange(len(keep)) perm = np.random.RandomState(seed).permutation(len(idx)) idx = idx[perm] val, tr = idx[:n_val], idx[n_val:] packed = np.load(f"{data}/voxels_packed.npy") seq = np.load(f"{data}/text_seq.npy") pool = np.load(f"{data}/text_pool.npy") pack = lambda s: (np.unpackbits(packed[s], axis=1), seq[s], pool[s]) return pack(tr), pack(val), len(keep) def main(): ap = argparse.ArgumentParser() ap.add_argument("--data", default="/root/data") ap.add_argument("--out", default="/root/runs/voxel_v1") ap.add_argument("--steps", type=int, default=90000) ap.add_argument("--batch-size", type=int, default=256) ap.add_argument("--lr", type=float, default=2e-4) ap.add_argument("--min-lr", type=float, default=1e-6) ap.add_argument("--warmup", type=int, default=500) ap.add_argument("--cfg-dropout", type=float, default=0.1) ap.add_argument("--ema", type=float, default=0.9999) ap.add_argument("--filtered", type=int, default=1) ap.add_argument("--val-size", type=int, default=1024) ap.add_argument("--val-every", type=int, default=2000) ap.add_argument("--log-every", type=int, default=200) ap.add_argument("--save-every", type=int, default=10000) ap.add_argument("--resume", default="") ap.add_argument("--seed", type=int, default=0) args = ap.parse_args() dev = "cuda" os.makedirs(args.out, exist_ok=True) torch.manual_seed(args.seed) torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True (vox, seq, pool), (vvox, vseq, vpool), total_meshes = load_split( args.data, args.filtered, args.val_size) M = vox.shape[0] vox_t = torch.from_numpy(vox).to(dev) seq_t = torch.from_numpy(seq).to(dev) pool_t = torch.from_numpy(pool).to(dev) null_seq = torch.from_numpy(np.load(f"{args.data}/null_seq.npy")).to(dev).float() null_pool = torch.from_numpy(np.load(f"{args.data}/null_pool.npy")).to(dev).float() vg = torch.Generator(device=dev).manual_seed(1234) vx1 = torch.from_numpy(vvox).to(dev).view(-1, 1, 32, 32, 32).float() * 2 - 1 vs = torch.from_numpy(vseq).to(dev).float() vp = torch.from_numpy(vpool).to(dev).float() vx0 = torch.randn(vx1.shape, device=dev, generator=vg) vt = torch.sigmoid(torch.randn(vx1.shape[0], device=dev, generator=vg)) @torch.no_grad() def val_loss(): ema.eval() tot, n = 0.0, 0 for j in range(0, vx1.shape[0], args.batch_size): sl = slice(j, j + args.batch_size) m = vx1[sl].shape[0] tb = vt[sl].view(-1, 1, 1, 1, 1) xt = (1 - tb) * vx0[sl] + tb * vx1[sl] with torch.autocast("cuda", dtype=torch.bfloat16): v = ema(xt, vt[sl], vs[sl], vp[sl]) tot += F.mse_loss(v.float(), vx1[sl] - vx0[sl]).item() * m n += m return tot / n model = VoxelDiT().to(dev) ema = copy.deepcopy(model).eval() for p in ema.parameters(): p.requires_grad_(False) opt = torch.optim.AdamW(model.parameters(), lr=args.lr, betas=(0.9, 0.99), weight_decay=0.0) start = 0 if args.resume and os.path.exists(args.resume): ck = torch.load(args.resume, map_location=dev) model.load_state_dict(ck["model"]) ema.load_state_dict(ck["ema"]) opt.load_state_dict(ck["opt"]) start = ck["step"] + 1 B = args.batch_size gen = torch.Generator(device=dev).manual_seed(args.seed) epochs = args.steps * B / M print(f"[train] {M} train + {vx1.shape[0]} val of {total_meshes} captioned " f"({'filtered' if args.filtered else 'unfiltered'}), " f"{model.num_params()/1e6:.2f}M params", flush=True) print(f"[train] {args.steps} steps x batch {B} = {epochs:.0f} epochs, " f"resume from {start}", flush=True) logf = open(f"{args.out}/log.jsonl", "a") running, t0 = 0.0, time.time() for step in range(start, args.steps): i = torch.randint(0, M, (B,), device=dev, generator=gen) x1 = vox_t[i].view(B, 1, 32, 32, 32).float() * 2 - 1 s = seq_t[i].float() p = pool_t[i].float() drop = torch.rand(B, device=dev, generator=gen) < args.cfg_dropout s = torch.where(drop[:, None, None], null_seq, s) p = torch.where(drop[:, None], null_pool, p) x0 = torch.randn(x1.shape, device=dev, generator=gen) u = torch.randn(B, device=dev, generator=gen) t = torch.sigmoid(u) tb = t.view(-1, 1, 1, 1, 1) xt = (1 - tb) * x0 + tb * x1 target = x1 - x0 lr = cosine_lr(step, args.steps, args.lr, args.min_lr, args.warmup) for g in opt.param_groups: g["lr"] = lr with torch.autocast("cuda", dtype=torch.bfloat16): v = model(xt, t, s, p) loss = F.mse_loss(v.float(), target) opt.zero_grad(set_to_none=True) loss.backward() gn = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) opt.step() with torch.no_grad(): d = 1 - args.ema for pe, pm in zip(ema.parameters(), model.parameters()): pe.mul_(args.ema).add_(pm.detach(), alpha=d) running += loss.item() if (step + 1) % args.log_every == 0: avg = running / args.log_every el = time.time() - t0 sps = args.log_every / el eta = (args.steps - step - 1) / sps / 3600 print(f"[s{step+1:06d}] loss={avg:.5f} lr={lr:.2e} gnorm={gn:.2f} " f"{sps:.2f} steps/s eta={eta:.2f}h", flush=True) logf.write(json.dumps({"step": step + 1, "loss": avg, "lr": lr, "gnorm": float(gn), "steps_per_s": sps}) + "\n") logf.flush() running, t0 = 0.0, time.time() if (step + 1) % args.val_every == 0 or step + 1 == args.steps: vl = val_loss() print(f"[s{step+1:06d}] val_loss={vl:.5f} (ema, {args.val_size} held out)", flush=True) logf.write(json.dumps({"step": step + 1, "val_loss": vl}) + "\n") logf.flush() t0 = time.time() if (step + 1) % args.save_every == 0 or step + 1 == args.steps: torch.save({"model": model.state_dict(), "ema": ema.state_dict(), "opt": opt.state_dict(), "step": step}, f"{args.out}/ckpt.pt") save_file({k: v.contiguous() for k, v in ema.state_dict().items()}, f"{args.out}/model.safetensors") print(f"[s{step+1:06d}] saved", flush=True) json.dump({"dit": {"vox_size": 32, "patch": 4, "dim": 384, "depth": 12, "heads": 6, "text_dim": 512}, "steps": args.steps, "batch_size": args.batch_size, "meshes": M, "filtered": bool(args.filtered)}, open(f"{args.out}/config.json", "w"), indent=2) print("TRAINDONE", flush=True) if __name__ == "__main__": main()