| """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
|
| 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}")
|
|
|
|
|
| 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)
|
| 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]
|
|
|
| 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)
|
| if args.phone_aligned:
|
|
|
|
|
|
|
|
|
| w = np.rint(d["durations"].astype(np.float64)).astype(np.int64)
|
| 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
|
|
|
| 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:
|
| 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()
|
|
|
