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 audio_dit import AudioDiT 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, split, dev): meta = json.load(open(f"{data}/{split}_meta.json")) mel = np.load(f"{data}/{split}_mel.npy") seq = np.load(f"{data}/{split}_text_seq.npy") pool = np.load(f"{data}/{split}_text_pool.npy") ci = np.array(meta["pair_clip_idx"], dtype=np.int64) return (torch.from_numpy(mel).to(dev), torch.from_numpy(seq).to(dev), torch.from_numpy(pool).to(dev), torch.from_numpy(ci).to(dev), meta) def main(): ap = argparse.ArgumentParser() ap.add_argument("--data", default="/root/data") ap.add_argument("--out", default="/root/runs/audio_v1") ap.add_argument("--steps", type=int, default=90000) ap.add_argument("--batch-size", type=int, default=128) 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("--val-every", type=int, default=2000) ap.add_argument("--val-batch", type=int, default=256) ap.add_argument("--log-every", type=int, default=200) ap.add_argument("--save-every", type=int, default=5000) 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 mel_t, seq_t, pool_t, ci_t, meta = load_split(args.data, "development", dev) vmel_t, vseq_t, vpool_t, vci_t, vmeta = load_split(args.data, "validation", 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() N = ci_t.shape[0] x_res, y_res = meta["x_res"], meta["y_res"] def mel_batch(mel_tensor, ci_tensor, idx): return (mel_tensor[ci_tensor[idx]].float() / 127.5 - 1.0).unsqueeze(1) vg = torch.Generator(device=dev).manual_seed(1234) n_val = min(args.val_batch, vci_t.shape[0]) vidx = torch.arange(n_val, device=dev) vx1 = mel_batch(vmel_t, vci_t, vidx) vs = vseq_t[vidx].float() vp = vpool_t[vidx].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() tb = vt.view(-1, 1, 1, 1) xt = (1 - tb) * vx0 + tb * vx1 with torch.autocast("cuda", dtype=torch.bfloat16): v = ema(xt, vt, vs, vp) return F.mse_loss(v.float(), vx1 - vx0).item() model = AudioDiT(x_res=x_res, y_res=y_res, text_seq_dim=meta["seq_dim"], text_pool_dim=meta["pool_dim"]).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 best_val = float("inf") 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 best_val = ck.get("best_val", float("inf")) B = args.batch_size gen = torch.Generator(device=dev).manual_seed(args.seed) epochs = args.steps * B / N print(f"[train] {N} train pairs ({meta['n_clips']} clips) + {n_val} val, " f"{model.num_params()/1e6:.2f}M params, grid {model.grid_h}x{model.grid_w} " f"({model.grid_h*model.grid_w} tokens)", flush=True) print(f"[train] {args.steps} steps x batch {B} = {epochs:.0f} epochs, 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): idx = torch.randint(0, N, (B,), device=dev, generator=gen) x1 = mel_batch(mel_t, ci_t, idx) s = seq_t[idx].float() p = pool_t[idx].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) 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() tag = "" if vl < best_val: best_val = vl save_file({k: v.contiguous() for k, v in ema.state_dict().items()}, f"{args.out}/model_best.safetensors") json.dump({"step": step + 1, "val_loss": vl}, open(f"{args.out}/best.json", "w")) tag = " (new best, saved model_best.safetensors)" print(f"[s{step+1:06d}] val_loss={vl:.5f} (ema, {n_val} held out){tag}", 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, "best_val": best_val}, 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": {"x_res": x_res, "y_res": y_res, "patch": 16, "dim": 384, "depth": 12, "heads": 6, "text_seq_dim": meta["seq_dim"], "text_pool_dim": meta["pool_dim"]}, "mel": {"sample_rate": meta["sample_rate"], "n_fft": meta["n_fft"], "hop_length": meta["hop_length"], "top_db": meta["top_db"]}, "steps": args.steps, "batch_size": args.batch_size, "train_pairs": N, "train_clips": meta["n_clips"], "best_val": best_val}, open(f"{args.out}/config.json", "w"), indent=2) print("TRAINDONE", flush=True) if __name__ == "__main__": main()