File size: 8,859 Bytes
3139867 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 | """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=48)
ap.add_argument("--dec-pw", type=int, default=144)
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=20)
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()
|