#!/usr/bin/env python3 """ rec7_train.py — how rec7 was trained, in one file. This is the joint-reconstruction stage (the one that matters), written to be read. The lineage before it, in brief: 1. code prediction on ~53k self-generated tracks (c0 cross-entropy + acoustic-code cross-entropy through a small causal decoder) -> trunk 2. same, warm-started, with an 11.8k-track community pool added -> trunk 3. a state head fitted on top of the frozen trunk against a teacher that regresses the LM's true hidden states (MSE + 1-cos) -> head 4. THIS STAGE, three times, each warm-started from the last: 4k steps -> 10k steps (+guidance, +window consistency, +EMA) -> 60k steps on a ~22k-audio file stratified real pool -> rec7 Every step of this stage does two things: REAL a real recording -> latents -> trunk+head -> states -> the released depth chain -> the released condition encoder -> ONE denoising step of the released DiT, guided as playback is. Loss: MSE between what the DiT predicts and what the recording's own latents say it should predict. The encoder learns to write states that make the frozen renderer reproduce real audio. ANCHOR a synthetic window (where true codes exist): c0 cross-entropy + head regression to the teacher. Keeps the code head honest while the states are reshaped by REAL. Everything downstream of the encoder is frozen. Gradient reaches the trunk at a tiny learning rate and the head at a small one. Requires the MiniMax-Music3 checkpoint. Data layouts are yours to supply: --synth directory of .pt files with keys z [128,L], codes [T,8] --songs text file, one audio path per line --teacher checkpoint of the hidden-state regressor (stage 3) python rec7_train.py --m3 /path/minimax_music3 --synth pool/ \ --songs songs.txt --init trunk_v5.pt --head head.pt --out rec/ """ from __future__ import annotations import argparse import json import random import time from pathlib import Path import numpy as np import torch import torch.nn.functional as F from rec7_model import (V4Encoder, HHead, FRAMES, LATENT_WINDOW_MAX, LATENT_CHANNELS, SEM_VOCAB, AC_VOCAB, H_DIM, frame_latent_starts, pool_matrix, load_dav, lm_semantic_tables, encode_audio) WIN = 200 # the renderer's window, in frames (690 latents) # ───────────────────────── data ───────────────────────── def real_window(trunk, head, z, t0, dev): """200 frames of states from two stitched 128-frame encoder windows (offsets 0 and 72), plus the recording's own latents for that span. Returns (h [200,4096], x0 [1,128,690], the two window outputs).""" T_tot = int(z.shape[1] / 3.45) - 2 if t0 + WIN > T_tot: return None st = frame_latent_starts(T_tot) acc = torch.zeros(WIN, H_DIM, device=dev) cnt = torch.zeros(WIN, 1, device=dev) wins = [] for off in (0, 72): tt = t0 + off b = st[tt: tt + FRAMES + 1] - st[tt] n = int(b[-1]) if n > LATENT_WINDOW_MAX or st[tt] + n > z.shape[1]: return None lat = torch.zeros(1, LATENT_WINDOW_MAX, LATENT_CHANNELS) lat[0, :n] = z[:, int(st[tt]): int(st[tt]) + n].T.float() pl = torch.from_numpy(pool_matrix(b)).unsqueeze(0) with torch.autocast("cuda", dtype=torch.bfloat16): feats, _ = trunk(lat.to(dev), pl.to(dev)) hw = head(feats[0].float()) wins.append((off, hw)) acc[off: off + FRAMES] += hw cnt[off: off + FRAMES] += 1 lo, hi = int(st[t0]), int(st[t0 + WIN]) x0 = z[:, lo:hi].float().unsqueeze(0).to(dev) return acc / cnt.clamp_min(1), x0, wins def synth_batch(files, n, dev): """n random 128-frame windows from the synthetic pool with their true c0 targets (frame t's window predicts code t+1).""" xs, pls, cts = [], [], [] for f in random.sample(files, min(4 * n, len(files))): if len(xs) >= n: break d = torch.load(f, map_location="cpu", weights_only=False) z, codes = d["z"], d["codes"] T = codes.shape[0] - 1 if T < FRAMES + 4: continue st = frame_latent_starts(T) t0 = random.randrange(0, T - FRAMES) b = st[t0: t0 + FRAMES + 1] - st[t0] n_ = int(b[-1]) if n_ > LATENT_WINDOW_MAX or st[t0] + n_ > z.shape[1]: continue x = torch.zeros(LATENT_WINDOW_MAX, LATENT_CHANNELS) x[:n_] = z[:, int(st[t0]): int(st[t0]) + n_].T.float() xs.append(x) pls.append(torch.from_numpy(pool_matrix(b))) cts.append(codes[t0 + 1: t0 + 1 + FRAMES, 0].long()) if not xs: return None return (torch.stack(xs).to(dev), torch.stack(pls).to(dev), torch.stack(cts).to(dev).clamp(0, SEM_VOCAB - 1)) # ────────────────── states -> renderer condition ────────────────── def condition_from_h(h, depth, W_head, W_embed, cond_enc, tau=1.0): """The 8 streams the condition encoder expects, built from states with gradient. c0: softmax-weighted embedding (soft forward; this is what rec7 trained with). c1..c7: hard pick forward, soft gradient backward, through the released depth decoder's greedy chain.""" dt = torch.bfloat16 hb = h.to(dt) sem = (hb @ W_head.T).float() / tau p0 = F.softmax(sem, -1).to(dt) e0 = p0 @ W_embed seq = [depth.projection(hb).unsqueeze(1), depth.projection(e0).unsqueeze(1)] streams = [hb] for k in range(1, 8): hid = depth(torch.cat(seq, dim=1))[:, -1] streams.append(hid) if k < 7: lg = depth.audio_heads[k - 1](hid).float() / tau pk = F.softmax(lg, -1).to(dt) tbl = depth.audio_embeddings.weight[(k - 1) * AC_VOCAB: k * AC_VOCAB] soft = pk @ tbl hard = tbl[lg.argmax(-1)] emb = hard + (soft - soft.detach()) # straight-through seq.append(depth.projection(emb).unsqueeze(1)) flat = torch.cat(streams, dim=-1).unsqueeze(0) # [1, T, 8*4096] return cond_enc(flat).to(dt) def reconstruction_loss(h, x0, depth, cond_enc, dit, W_head, W_embed, guide=1.7): """One denoising step of the frozen DiT, guided as playback is. Flow matching: x_t = t*x0 + (1-t)*e, target x0 - e, t in [0,1].""" cond = condition_from_h(h, depth, W_head, W_embed, cond_enc) e = torch.randn_like(x0) t = torch.rand(1, device=x0.device) x_t = (t * x0 + (1 - t) * e).to(torch.bfloat16) def pred(c): out = dit(hidden_states=x_t, timestep=t.to(torch.bfloat16), encoder_hidden_states=c, return_dict=False)[0].float() return out.transpose(1, 2) if out.shape != x0.shape else out pc = pred(cond) with torch.no_grad(): pu = pred(torch.zeros_like(cond)) guided = pu + guide * (pc - pu) return F.mse_loss(guided, (x0 - e).float()) def consistency_loss(wins): """The two encoder windows overlap on frames 72..128 of the span; they should agree there (seams between windows are audible).""" (_, a), (_, b) = wins return F.mse_loss(a[72:FRAMES], b[: FRAMES - 72]) # ───────────────────────── training ───────────────────────── def main() -> int: ap = argparse.ArgumentParser() ap.add_argument("--m3", type=Path, required=True) ap.add_argument("--synth", type=Path, required=True) ap.add_argument("--songs", type=Path, required=True) ap.add_argument("--init", type=Path, required=True, help="trunk") ap.add_argument("--head", type=Path, required=True) ap.add_argument("--teacher", type=Path, default=None, help="hidden-state regressor for the anchor; if " "omitted the anchor is c0 cross-entropy only") ap.add_argument("--out", type=Path, required=True) ap.add_argument("--steps", type=int, default=60000) ap.add_argument("--head-lr", type=float, default=3e-5) ap.add_argument("--trunk-lr", type=float, default=5e-6) ap.add_argument("--guide", type=float, default=1.7) ap.add_argument("--cons-w", type=float, default=0.25) ap.add_argument("--anchor-w", type=float, default=1.0) ap.add_argument("--ema", type=float, default=0.999) ap.add_argument("--batch", type=int, default=8) ap.add_argument("--save-every", type=int, default=250) a = ap.parse_args() a.out.mkdir(parents=True, exist_ok=True) dev = "cuda" # ── frozen renderer side ── from diffusers import ModularPipeline pipe = ModularPipeline.from_pretrained(str(a.m3)) pipe.load_components(dtype=torch.bfloat16) pipe.to(dev) depth, cond_enc, dit = (pipe.rvq_depth_decoder, pipe.condition_encoder, pipe.transformer) for m in (depth, cond_enc, dit): m.requires_grad_(False) W_head, W_embed = lm_semantic_tables(a.m3, dev) dav = load_dav(a.m3, dev) # ── the encoder ── trunk = V4Encoder().to(dev) trunk.load_state_dict(torch.load(a.init, map_location="cpu", weights_only=False)["model"]) head = HHead(trunk.pos.shape[-1]).to(dev) head.load_state_dict(torch.load(a.head, map_location="cpu", weights_only=False)["model"]) trunk.train() head.train() teacher = None if a.teacher: from train_hid2 import HiddenEncoder # stage-3 regressor ck = torch.load(a.teacher, map_location="cpu", weights_only=False) cfg = ck.get("cfg", {}) teacher = HiddenEncoder(int(cfg.get("d_model", 512)), int(cfg.get("layers", 8))).to(dev).eval() teacher.load_state_dict(ck["model"], strict=False) teacher.requires_grad_(False) opt = torch.optim.AdamW([ {"params": head.parameters(), "lr": a.head_lr}, {"params": trunk.parameters(), "lr": a.trunk_lr}, ], weight_decay=0.01) params = list(trunk.parameters()) + list(head.parameters()) ema = {k: v.detach().clone().float() for k, v in list(trunk.state_dict().items()) + list(head.state_dict().items())} synth = sorted(a.synth.glob("*.pt")) songs = [Path(l.strip()) for l in a.songs.read_text().splitlines() if l.strip()] zcache = a.out / "zcache" zcache.mkdir(exist_ok=True) def latents(path): """Flow-VAE latents, cached to disk after the first visit.""" c = zcache / (path.stem + ".pt") if c.exists(): return torch.load(c, map_location="cpu", weights_only=False) import torchaudio wav, sr = torchaudio.load(str(path)) if sr != 44100: wav = torchaudio.transforms.Resample(sr, 44100)(wav) if wav.shape[0] == 1: wav = wav.repeat(2, 1) z = encode_audio(dav, wav[:2].float(), dev).to(torch.float16) torch.save(z, c) return z hist, t0 = [], time.time() run = {"real": 0.0, "anchor": 0.0, "n": 0} for step in range(1, a.steps + 1): try: # ── REAL: one window of one recording ── z = latents(random.choice(songs)).float() T_tot = int(z.shape[1] / 3.45) - 2 if T_tot < WIN + 20: continue got = real_window(trunk, head, z, random.randrange(0, T_tot - WIN - 10), dev) if got is None: continue h, x0, wins = got loss = reconstruction_loss(h, x0, depth, cond_enc, dit, W_head, W_embed, a.guide) if a.cons_w > 0: loss = loss + a.cons_w * consistency_loss(wins) opt.zero_grad(set_to_none=True) loss.backward() torch.nn.utils.clip_grad_norm_(params, 1.0) opt.step() # ── ANCHOR: a synthetic batch with true codes ── sb = synth_batch(synth, a.batch, dev) if sb is not None: x, pl, ct = sb with torch.autocast("cuda", dtype=torch.bfloat16): feats, sem = trunk(x, pl) al = F.cross_entropy(sem.float().reshape(-1, SEM_VOCAB), ct.reshape(-1)) if teacher is not None: with torch.no_grad(), torch.autocast( "cuda", dtype=torch.bfloat16): tgt = teacher(x, pl).float() ph = head(feats.float()) al = al + F.mse_loss(ph, tgt) + \ (1 - F.cosine_similarity(ph, tgt, dim=-1).mean()) al = a.anchor_w * al opt.zero_grad(set_to_none=True) al.backward() torch.nn.utils.clip_grad_norm_(params, 1.0) opt.step() run["anchor"] += float(al) # ── EMA ── with torch.no_grad(): for k, v in list(trunk.state_dict().items()) \ + list(head.state_dict().items()): ema[k].mul_(a.ema).add_(v.float(), alpha=1 - a.ema) run["real"] += float(loss) run["n"] += 1 except torch.cuda.OutOfMemoryError: opt.zero_grad(set_to_none=True) torch.cuda.empty_cache() continue if step % 25 == 0: n = max(run["n"], 1) print(f"step {step} | real {run['real'] / n:.4f} | anchor " f"{run['anchor'] / n:.4f} | {(time.time() - t0) / 60:.1f}m", flush=True) hist.append({"step": step, "real": run["real"] / n}) run = {"real": 0.0, "anchor": 0.0, "n": 0} if step % a.save_every == 0: # saved weights are the EMA weights torch.save({"model": {k: ema[k].to(v.dtype) for k, v in trunk.state_dict().items()}, "step": step}, a.out / "trunk_latest.pt") torch.save({"model": {k: ema[k].to(v.dtype) for k, v in head.state_dict().items()}, "step": step}, a.out / "head_latest.pt") (a.out / "history.json").write_text(json.dumps(hist, indent=1)) return 0 if __name__ == "__main__": raise SystemExit(main())