"""Z-mix decoder training: recovery on teacher mels, then a 50/50 student mix. This is the load-bearing stage for the acoustic -> decoder handoff. A decoder trained only on teacher mels is brittle to the acoustic student's (smoother) output at synthesis time; mixing student mels into the training data with --mix-prob 0.5 fixes it. The schedule that works: 1. recovery --mix-prob 0.0 (teacher mels only) 2. z-mix --mix-prob 0.5 (half student mels) Use the same --init checkpoint for both stages and pass the acoustic and duration checkpoints with --ac-ckpt / --dur-ckpt (they are frozen and only serve to synthesize training mels). """ import argparse import glob import os import time import numpy as np import torch import sys from pathlib import Path sys.path.insert(0, str(Path(__file__).resolve().parents[1])) import torch.nn.functional as F from microtts.models import DurationStudent, AcousticStudent, DecoderStudent def stft_loss(y, yh, device): loss = 0.0 n = min(y.shape[-1], yh.shape[-1]) for n_fft, hop in [(512, 128), (1024, 256), (2048, 512)]: if n < n_fft: continue # skip FFTs larger than the input Y = torch.stft(y[..., :n], n_fft, hop, window=torch.hann_window(n_fft, device=device), return_complex=True, center=False) Yh = torch.stft(yh[..., :n], n_fft, hop, window=torch.hann_window(n_fft, device=device), return_complex=True, center=False) T = min(Y.shape[-1], Yh.shape[-1]) mag, mag_r = Y[..., :T].abs(), Yh[..., :T].abs() sc = (mag - mag_r).norm(dim=-2) / (mag.norm(dim=-2) + mag_r.norm(dim=-2) + 1e-6) lm = (mag.clamp_min(1e-6).log10() - mag_r.clamp_min(1e-6).log10()).abs().mean() loss = loss + sc.mean() + lm return loss / 3.0 def main(): ap = argparse.ArgumentParser() ap.add_argument("--pack", default="packs/train") ap.add_argument("--steps", type=int, default=10000) ap.add_argument("--batch", type=int, default=4) ap.add_argument("--lr", type=float, default=1e-4) ap.add_argument("--device", default="cuda") ap.add_argument("--tag", default="dec_zmix") ap.add_argument("--mix-prob", type=float, default=0.5) ap.add_argument("--phone-aligned", action="store_true", help="crop student-mel windows phone-aligned to the audio " "(fixes train/inference timeline desync for the mix stage)") ap.add_argument("--init", default="decoder_init.pt", help="teacher-initialized decoder checkpoint") ap.add_argument("--dec-dim", type=int, default=40) ap.add_argument("--dec-pw", type=int, default=120) ap.add_argument("--dec-blocks", type=int, default=3) ap.add_argument("--ac-hidden", type=int, default=28) ap.add_argument("--ac-ckpt", default="acoustic.pt") ap.add_argument("--dur-hidden", type=int, default=12) ap.add_argument("--dur-depth", type=int, default=3) ap.add_argument("--dur-ckpt", default="duration.pt") args = ap.parse_args() device = torch.device(args.device if torch.cuda.is_available() else "cpu") files = sorted(glob.glob(os.path.join(args.pack, "*.npz"))) print(f"pack: {len(files)} rows, device {device}") # load acoustic student (frozen) for student-mel generation ac = AcousticStudent(hidden=args.ac_hidden).to(device) ac.load_state_dict(torch.load(args.ac_ckpt, map_location=device)) ac.eval() dur = DurationStudent(hidden=args.dur_hidden, depth=args.dur_depth).to(device) dur.load_state_dict(torch.load(args.dur_ckpt, map_location=device)) dur.eval() model = DecoderStudent(dim=args.dec_dim, pw_hidden=args.dec_pw, blocks=args.dec_blocks).to(device) model.load_state_dict(torch.load(args.init, map_location=device)) n = sum(p.numel() for p in model.parameters()) print(f"decoder params: {n:,} (init from {args.init})") opt = torch.optim.AdamW(model.parameters(), lr=args.lr, betas=(0.8, 0.99)) sched = torch.optim.lr_scheduler.ExponentialLR(opt, gamma=0.9998) RUN = os.path.join("runs", args.tag) os.makedirs(RUN, exist_ok=True) t0 = time.time() for step in range(1, args.steps + 1): idx = np.random.choice(len(files), min(args.batch, len(files)), replace=False) total = 0.0 cnt = 0 opt.zero_grad() for i in idx: d = np.load(files[i]) mel_t = torch.from_numpy(d["mel"]).to(device) # teacher mel [100, T] audio = torch.from_numpy(d["audio"]).to(device) T = mel_t.shape[1] seg = 64 if T > seg: s = np.random.randint(0, T - seg) mel_t = mel_t[:, s:s + seg] audio = audio[s * 256:(s + seg) * 256] # 50/50: teacher mel or acoustic-student mel if np.random.rand() < args.mix_prob: ids = torch.from_numpy(d["phone_ids"]).long().to(device) with torch.no_grad(): dd = dur(ids) mel_s = ac(ids, dd) # student mel [100, T_full] if args.phone_aligned: # Phone-aligned crop: choose a phone range, take student mel # frames from the STUDENT timeline and audio from the TEACHER # timeline for the SAME phones (student/teacher frame counts # differ; same-offset cropping desynchronises them). w = np.rint(d["durations"].astype(np.float64)).astype(np.int64) # teacher per-phone dd_np = np.rint(dd.detach().cpu().numpy().reshape(-1)).astype(np.int64) cum_s = np.concatenate([[0], np.cumsum(np.maximum(dd_np, 1))]) cum_t = np.concatenate([[0], np.cumsum(np.maximum(w, 1))]) P = len(dd_np) if P < 2: continue # find a phone window [j1, j2) with >= 16 student frames j1 = np.random.randint(0, max(1, P - 1)) j2 = j1 + 1 while j2 < P and cum_s[j2] - cum_s[j1] < seg: j2 += 1 if cum_s[j2] - cum_s[j1] < 16: continue x_s, y_s = int(cum_s[j1]), int(cum_s[j2]) x_t, y_t = int(cum_t[j1]), int(cum_t[j2]) mel_s = mel_s[:, x_s:y_s] a0, a1 = x_t * 256, y_t * 256 if a1 > audio.shape[0]: a1 = audio.shape[0] if a1 - a0 < 1024: continue audio = audio[a0:a1] mel = mel_s else: seg_s = min(seg, mel_s.shape[1]) if seg_s < 16: # need >= 16 frames for the 2048-FFT loss continue s_s = np.random.randint(0, max(1, mel_s.shape[1] - seg_s + 1)) mel_s = mel_s[:, s_s:s_s + seg_s] a0, a1 = s_s * 256, (s_s + seg_s) * 256 if a1 > audio.shape[0]: a1 = audio.shape[0] if a1 - a0 < 1024: continue audio = audio[a0:a1] mel = mel_s else: mel = mel_t noise = torch.randn(4, mel.shape[1], device=device) wav = model.synthesize(mel, noise) n_w = min(wav.shape[-1], audio.shape[-1]) l_wav = F.l1_loss(wav[:n_w], audio[:n_w]) l_stft = stft_loss(wav[:n_w].unsqueeze(0), audio[:n_w].unsqueeze(0), device) rms_p = wav[:n_w].square().mean().clamp_min(1e-8).sqrt() rms_t = audio[:n_w].square().mean().clamp_min(1e-8).sqrt() l_loud = F.l1_loss(rms_p, rms_t) total = total + 5.0 * l_wav + 1.0 * l_stft + 3.0 * l_loud cnt += 1 if cnt == 0: continue loss = total / cnt loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) opt.step() sched.step() if step % 100 == 0 or step == 1: dt = time.time() - t0 print(f"step {step}/{args.steps} | wav {l_wav.item():.4f} stft {l_stft.item():.3f} " f"loud {l_loud.item():.4f} | {step/dt:.1f} sps", flush=True) if step % 1000 == 0: torch.save(model.state_dict(), os.path.join(RUN, "last.pt")) torch.save(model.state_dict(), os.path.join(RUN, "final.pt")) print("Z-MIX DECODER DONE") if __name__ == "__main__": main()