| """Train the decoder student on the deployment distribution. |
| |
| The pack holds, per row: the acoustic student's mels (the exact inference |
| input), the teacher's phoneme ids/durations, and — as audio targets — full |
| vocoder renders of those exact mels. The objective is therefore "reproduce the |
| vocoder's output on the mels you will actually receive", and there is no |
| acoustic->decoder distribution gap to fix later. |
| |
| This is the full supervision stack used for the released weights: |
| |
| waveform L1 (wav-weight 1.0) |
| multi-resolution spectral loss |
| sig-temporal (cosine-gram) (0.4) temporal texture |
| feature-hint (channel moments) (0.05) distillation from the vocoder |
| quiet-ceiling (0.010) stops quiet-frame hiss |
| high-band-excess (0.1) penalizes HF overshoot |
| MRD adversarial + feature match (0.075 / 0.75 from step 1500) |
| |
| Requires `vocos` (pip install vocos) for the feature hooks; the pack itself |
| needs no vocoder at train time. |
| |
| Stages used for the release: |
| recovery --mix-prob 0.0 (the deployment pack IS the target distribution) |
| ... then continue on the same pack for as long as the curve rises |
| (the released decoder ran to 115k steps; see train/README.md) |
| """ |
| 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 as nn |
| import torch.nn.functional as F |
|
|
| from microtts.models import DecoderStudent |
|
|
|
|
| |
| |
| |
| class TeacherFeats: |
| """Capture the vocoder's internal activations at chosen stages.""" |
|
|
| def __init__(self, vocos): |
| self.feats = {} |
| self.handles = [] |
| b = vocos.backbone |
| self.handles.append(b.embed.register_forward_hook(self._h("embed"))) |
| self.handles.append(b.convnext[3].register_forward_hook(self._h("cn3"))) |
| self.handles.append(b.final_layer_norm.register_forward_hook(self._h("final"))) |
|
|
| def _h(self, key): |
| def fn(mod, inp, out): |
| self.feats[key] = out |
| return fn |
|
|
| def clear(self): |
| self.feats = {} |
|
|
|
|
| def standardized_time_series(v): |
| mean = v.mean(dim=1, keepdim=True) |
| std = v.std(dim=1, keepdim=True).clamp_min(1e-4) |
| return (v - mean) / std |
|
|
|
|
| def temporal_cosine_gram(value): |
| """Per-timestep channel-vector self-similarity [B,T,T].""" |
| if value.ndim != 3: |
| raise RuntimeError(f"expected [B, C, T], got {value.shape}") |
| series = value.transpose(1, 2) |
| series = series - series.mean(dim=-1, keepdim=True) |
| series = F.normalize(series, dim=-1, eps=1e-6) |
| return torch.matmul(series, series.transpose(1, 2)) |
|
|
|
|
| def signature_temporal_loss(student_feats, teacher_feats, keys=("embed", "cn3")): |
| losses = [] |
| for key in keys: |
| s = student_feats.get(key) |
| t = teacher_feats.get(key) |
| if s is None or t is None: |
| continue |
| if s.ndim == 3: |
| s = s[0] |
| if t.ndim == 3: |
| t = t[0] |
| T = min(s.shape[-1], t.shape[-1]) |
| if T < 4: |
| continue |
| s = s[:, :T].unsqueeze(0).float() |
| t = t[:, :T].unsqueeze(0).float() |
| losses.append(F.l1_loss(temporal_cosine_gram(s), temporal_cosine_gram(t))) |
| if not losses: |
| return torch.zeros((), device=next(iter(student_feats.values())).device) |
| return sum(losses) / len(losses) |
|
|
|
|
| def channel_moment_hint_loss(student_feats, teacher_feats): |
| """Standardized channel-mean + log-RMS envelopes.""" |
| losses = [] |
| for key in ("embed", "cn3"): |
| s = student_feats.get(key) |
| t = teacher_feats.get(key) |
| if s is None or t is None: |
| continue |
| if s.ndim == 3: |
| s = s[0] |
| if t.ndim == 3: |
| t = t[0] |
| if s.shape[-1] != t.shape[-1]: |
| T = min(s.shape[-1], t.shape[-1]) |
| s, t = s[:, :T], t[:, :T] |
| s_mean = standardized_time_series(s.mean(dim=0, keepdim=True)) |
| t_mean = standardized_time_series(t.mean(dim=0, keepdim=True)) |
| s_rms = standardized_time_series(torch.log1p(torch.sqrt(s.square().mean(dim=0, keepdim=True) + 1e-6))) |
| t_rms = standardized_time_series(torch.log1p(torch.sqrt(t.square().mean(dim=0, keepdim=True) + 1e-6))) |
| losses.append(F.l1_loss(s_mean, t_mean)) |
| losses.append(F.l1_loss(s_rms, t_rms)) |
| if not losses: |
| return torch.zeros((), device=next(iter(student_feats.values())).device) |
| return sum(losses) / len(losses) |
|
|
|
|
| |
| |
| |
| class SpectralResolutionDiscriminator(nn.Module): |
| def __init__(self, n_fft, hop_length, win_length, channels=(32, 32, 32, 32)): |
| super().__init__() |
| self.n_fft, self.hop_length, self.win_length = n_fft, hop_length, win_length |
| ks, st, pd = (3, 9), (1, 2), (1, 4) |
| layers = [] |
| in_ch = 1 |
| for out_ch in channels: |
| layers.append(nn.Conv2d(in_ch, out_ch, kernel_size=ks, stride=st, padding=pd)) |
| in_ch = out_ch |
| self.layers = nn.ModuleList(layers) |
| self.post = nn.Conv2d(in_ch, 1, kernel_size=ks, padding=pd) |
|
|
| def spec(self, x): |
| if x.shape[-1] < self.n_fft: |
| return None |
| w = torch.hann_window(self.n_fft, device=x.device, dtype=x.dtype) |
| s = torch.stft(x[0], self.n_fft, self.hop_length, win_length=self.win_length, |
| window=w, return_complex=True) |
| return torch.log1p(s.abs()).unsqueeze(0).unsqueeze(0) |
|
|
| def forward(self, x): |
| h = self.spec(x) |
| if h is None: |
| return None, None |
| feats = [] |
| for layer in self.layers: |
| h = F.leaky_relu(layer(h), 0.1) |
| feats.append(h) |
| score = self.post(h) |
| feats.append(score) |
| return score, feats |
|
|
|
|
| class MRD(nn.Module): |
| def __init__(self, configs=((512, 128, 512), (1024, 256, 1024), (2048, 512, 2048))): |
| super().__init__() |
| self.discs = nn.ModuleList([SpectralResolutionDiscriminator(*c) for c in configs]) |
|
|
| def forward(self, x): |
| scores, feats = [], [] |
| for d in self.discs: |
| s, f = d(x) |
| if s is None: |
| continue |
| scores.append(s) |
| feats.append(f) |
| return scores, feats |
|
|
|
|
| def generator_lsgan_loss(fake_scores): |
| if not fake_scores: |
| return torch.zeros(()) |
| total = sum((s - 1.0).square().mean() for s in fake_scores) |
| return total / len(fake_scores) |
|
|
|
|
| def discriminator_lsgan_loss(real_scores, fake_scores): |
| if not real_scores: |
| return torch.zeros(()) |
| total = sum((r - 1.0).square().mean() + f.square().mean() |
| for r, f in zip(real_scores, fake_scores)) |
| return total / len(real_scores) |
|
|
|
|
| def feature_matching_loss(real_feats, fake_feats): |
| losses = [] |
| for rf, ff in zip(real_feats, fake_feats): |
| for a, b in zip(rf, ff): |
| losses.append(F.l1_loss(b, a.detach())) |
| if not losses: |
| return torch.zeros(()) |
| return sum(losses) / len(losses) |
|
|
|
|
| |
| |
| |
| 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 high_band_excess_loss(prediction, target, sample_rate=24000, high_band_hz=6000.0, margin_db=1.0): |
| """Penalize HF energy ratio above the target's ratio times a margin.""" |
| n_fft, hop = 1024, 256 |
| window = torch.hann_window(n_fft, device=prediction.device, dtype=prediction.dtype) |
| ps = torch.stft(prediction[0], n_fft, hop, win_length=n_fft, window=window, return_complex=True) |
| ts = torch.stft(target[0], n_fft, hop, win_length=n_fft, window=window, return_complex=True) |
| nyq = sample_rate / 2.0 |
| freqs = torch.linspace(0.0, nyq, n_fft // 2 + 1, device=prediction.device, dtype=prediction.dtype) |
| mask = (freqs >= high_band_hz).to(dtype=prediction.dtype).unsqueeze(-1) |
| p_pow = ps.abs().square() |
| t_pow = ts.abs().square() |
| pr = torch.sum(p_pow * mask) / torch.sum(p_pow).clamp_min(1e-10) |
| tr = torch.sum(t_pow * mask) / torch.sum(t_pow).clamp_min(1e-10) |
| margin_ratio = 10.0 ** (margin_db / 10.0) |
| excess = torch.log(pr.clamp_min(1e-10)) - torch.log((tr * margin_ratio).clamp_min(1e-10)) |
| return F.relu(excess).square() |
|
|
|
|
| def framed_rms(x, frame_size=1024, frame_hop=256): |
| if x.ndim == 2: |
| x = x[0] |
| if x.ndim == 1: |
| x = x.unsqueeze(0) |
| if x.shape[-1] < frame_size: |
| return x.square().mean(dim=-1).sqrt().unsqueeze(-1) |
| frames = x.unfold(-1, frame_size, frame_hop) |
| return frames.square().mean(dim=-1).sqrt().mean(dim=0, keepdim=True) |
|
|
|
|
| def quiet_ceiling_loss(prediction, target, quantile=0.25, frame_size=1024, frame_hop=256, margin_db=0.5): |
| """Penalize predictions louder than the target on the target's quiet frames.""" |
| tr = framed_rms(target, frame_size, frame_hop) |
| pr = framed_rms(prediction, frame_size, frame_hop) |
| T = min(tr.shape[-1], pr.shape[-1]) |
| tr, pr = tr[..., :T], pr[..., :T] |
| thr = torch.quantile(tr.flatten(), quantile) |
| mask = (torch.log1p(tr) <= thr).to(prediction.dtype) |
| db = 20.0 / np.log(10.0) |
| excess_db = db * (torch.log(pr + 1e-8) - torch.log(tr + 1e-8)) - margin_db |
| return torch.sum(F.relu(excess_db) * mask) / mask.sum().clamp_min(1.0) / 20.0 |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--pack", default="packs/deploy") |
| ap.add_argument("--steps", type=int, default=20000) |
| ap.add_argument("--batch", type=int, default=4) |
| ap.add_argument("--lr", type=float, default=2e-4) |
| ap.add_argument("--device", default="cuda") |
| ap.add_argument("--tag", default="dec") |
| ap.add_argument("--mix-prob", type=float, default=0.0, |
| help="fraction of items that use live acoustic mels instead " |
| "of the pack's (the deployment pack needs 0.0)") |
| ap.add_argument("--init", required=True, |
| help="teacher-initialized decoder checkpoint (init_decoder.py)") |
| ap.add_argument("--resume-step", type=int, default=0) |
| ap.add_argument("--vocos-target", action="store_true", |
| help="render targets live through the vocoder instead of " |
| "reading pack audio (the released pack stores audio)") |
| ap.add_argument("--noise-zeros", action="store_true", |
| help="zeros instead of randn for the decoder noise input " |
| "(matches inference; the released recipe)") |
| ap.add_argument("--dec-dim", type=int, default=67) |
| ap.add_argument("--dec-pw", type=int, default=197) |
| ap.add_argument("--dec-blocks", type=int, default=4) |
| ap.add_argument("--ac-hidden", type=int, default=64) |
| ap.add_argument("--ac-ckpt", default="acoustic.pt") |
| ap.add_argument("--dur-hidden", type=int, default=12) |
| ap.add_argument("--dur-ckpt", default="duration.pt") |
| |
| ap.add_argument("--feature-hint-weight", type=float, default=0.05) |
| ap.add_argument("--quiet-ceiling-weight", type=float, default=0.010) |
| ap.add_argument("--high-band-excess-weight", type=float, default=0.1) |
| ap.add_argument("--adv-weight", type=float, default=0.075) |
| ap.add_argument("--adv-feature-weight", type=float, default=0.75) |
| ap.add_argument("--adv-lr", type=float, default=2e-4) |
| ap.add_argument("--adv-start-step", type=int, default=1500) |
| ap.add_argument("--const-lr", action="store_true", |
| help="constant lr (the released recipe)") |
| ap.add_argument("--betas", type=str, default="0.9,0.999") |
| ap.add_argument("--wav-weight", type=float, default=1.0) |
| ap.add_argument("--sig-temporal-weight", type=float, default=0.4) |
| ap.add_argument("--loud-weight", type=float, default=0.0) |
| 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}", flush=True) |
|
|
| from microtts.models import AcousticStudent, DurationStudent |
|
|
| 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).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 {args.init})", flush=True) |
|
|
| from vocos import Vocos |
| teacher = Vocos.from_pretrained("charactr/vocos-mel-24khz").to(device).eval() |
| for p in teacher.parameters(): |
| p.requires_grad_(False) |
| tf = TeacherFeats(teacher) |
| print("vocoder loaded with hooks", flush=True) |
|
|
| mrd = MRD().to(device) |
| betas = tuple(float(x) for x in args.betas.split(",")) |
| opt = torch.optim.AdamW(model.parameters(), lr=args.lr, betas=betas) |
| opt_d = torch.optim.AdamW(mrd.parameters(), lr=args.adv_lr, betas=betas) |
| sched = None if args.const_lr else torch.optim.lr_scheduler.ExponentialLR(opt, gamma=0.9998) |
| if sched is not None and args.resume_step > 0: |
| for _ in range(args.resume_step): |
| sched.step() |
|
|
| RUN = os.path.join("runs", args.tag) |
| os.makedirs(RUN, exist_ok=True) |
| t0 = time.time() |
| for step in range(args.resume_step + 1, args.steps + 1): |
| idx = np.random.choice(len(files), min(args.batch, len(files)), replace=False) |
| cnt = 0 |
| content_total = None |
| opt.zero_grad() |
| gan_pair = None |
| for ii, i in enumerate(idx): |
| d = np.load(files[i]) |
| mel_t = torch.from_numpy(d["mel"]).to(device) |
| audio = None if args.vocos_target else torch.from_numpy(d["audio"]).to(device) |
| T = mel_t.shape[1] |
| seg = 64 |
| s = 0 |
| if T > seg: |
| s = np.random.randint(0, T - seg) |
| 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) |
| 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 = mel_s[:, s_s:s_s + seg_s] |
| if args.vocos_target: |
| with torch.no_grad(): |
| aud_full = teacher.decode(mel_s.unsqueeze(0))[0].float() |
| a0, a1 = s_s * 256, (s_s + mel.shape[1]) * 256 |
| if a1 > aud_full.shape[0]: |
| a1 = aud_full.shape[0] |
| if a1 - a0 < 1024: |
| continue |
| audio = aud_full[a0:a1] |
| else: |
| 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] |
| else: |
| mel = mel_t[:, s:s + seg] if T > seg else mel_t |
| if args.vocos_target: |
| with torch.no_grad(): |
| aud_full = teacher.decode(mel_t.unsqueeze(0))[0].float() |
| a0, a1 = s * 256, (s + mel.shape[1]) * 256 |
| if a1 > aud_full.shape[0]: |
| a1 = aud_full.shape[0] |
| if a1 - a0 < 1024: |
| continue |
| audio = aud_full[a0:a1] |
| else: |
| audio = audio[s * 256:(s + mel.shape[1]) * 256] |
| noise = torch.zeros(4, mel.shape[1], device=device) if args.noise_zeros \ |
| else torch.randn(4, mel.shape[1], device=device) |
| wav = model.synthesize(mel, noise) |
| n_w = min(wav.shape[-1], audio.shape[-1]) |
| wav_c, aud_c = wav[:n_w].unsqueeze(0), audio[:n_w].unsqueeze(0) |
| if gan_pair is None: |
| gan_pair = (wav_c, aud_c) |
|
|
| l_wav = F.l1_loss(wav[:n_w], audio[:n_w]) |
| l_stft = stft_loss(wav_c, aud_c, 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) |
|
|
| |
| l_hint = torch.zeros((), device=device) |
| l_sigt = torch.zeros((), device=device) |
| need_feats = (args.feature_hint_weight > 0 or args.sig_temporal_weight > 0) and ii == 0 |
| if need_feats: |
| with torch.no_grad(): |
| tf.clear() |
| teacher.decode(mel.unsqueeze(0)) |
| t_feats = {k: v.detach() for k, v in tf.feats.items()} |
| s_feats = {} |
| x = model.embed(mel) + model.noise_adapter(noise) |
| s_feats["embed"] = x |
| x2 = model.norm(x.transpose(0, 1)).transpose(0, 1) |
| cn_idx = min(3, len(model.blocks) - 1) |
| for bi, b in enumerate(model.blocks): |
| x2 = b(x2) |
| if bi == cn_idx: |
| s_feats["cn3"] = x2 |
| t_embed = t_feats.get("embed") |
| if t_embed is not None: |
| t_embed = t_embed[0] |
| Tt = min(t_embed.shape[-1], s_feats["embed"].shape[-1]) |
| t_feats["embed"] = t_embed[:, :Tt].unsqueeze(0) |
| s_feats["embed"] = s_feats["embed"][:, :Tt].unsqueeze(0) |
| t_cn3 = t_feats.get("cn3") |
| if t_cn3 is not None: |
| t_cn3 = t_cn3[0] |
| Tt = min(t_cn3.shape[-1], s_feats["cn3"].shape[-1]) |
| t_feats["cn3"] = t_cn3[:, :Tt].unsqueeze(0) |
| s_feats["cn3"] = s_feats["cn3"][:, :Tt].unsqueeze(0) |
| l_hint = channel_moment_hint_loss(s_feats, t_feats) |
| if args.sig_temporal_weight > 0: |
| l_sigt = signature_temporal_loss(s_feats, t_feats) |
|
|
| l_hf = high_band_excess_loss(wav_c, aud_c, 24000, 6000.0, 1.0) \ |
| if args.high_band_excess_weight > 0 else torch.zeros((), device=device) |
| l_quiet = quiet_ceiling_loss(wav_c[0], aud_c[0], 0.25, 1024, 256, 0.5) \ |
| if args.quiet_ceiling_weight > 0 else torch.zeros((), device=device) |
|
|
| item = (args.wav_weight * l_wav + 1.0 * l_stft + args.loud_weight * l_loud |
| + args.feature_hint_weight * l_hint |
| + args.sig_temporal_weight * l_sigt |
| + args.high_band_excess_weight * l_hf |
| + args.quiet_ceiling_weight * l_quiet) / max(args.batch, 1) |
| content_total = item if content_total is None else content_total + item |
| cnt += 1 |
|
|
| if cnt == 0 or content_total is None: |
| continue |
|
|
| content_total.backward(retain_graph=True) |
| l_adv = torch.zeros((), device=device) |
| l_adv_feat = torch.zeros((), device=device) |
| if step >= args.adv_start_step and gan_pair is not None: |
| wav_g, aud_g = gan_pair |
| nw = min(wav_g.shape[-1], aud_g.shape[-1]) |
| wav_g, aud_g = wav_g[..., :nw], aud_g[..., :nw] |
| opt_d.zero_grad() |
| real_scores, _ = mrd(aud_g) |
| fake_scores, _ = mrd(wav_g.detach()) |
| d_loss = discriminator_lsgan_loss(real_scores, fake_scores) |
| d_loss.backward() |
| torch.nn.utils.clip_grad_norm_(mrd.parameters(), 5.0) |
| opt_d.step() |
| for p in mrd.parameters(): |
| p.requires_grad_(False) |
| fake_scores_g, fake_feats_g = mrd(wav_g) |
| _, real_feats_g = mrd(aud_g) |
| l_adv = generator_lsgan_loss(fake_scores_g) |
| l_adv_feat = feature_matching_loss(real_feats_g, fake_feats_g) |
| for p in mrd.parameters(): |
| p.requires_grad_(True) |
| gan_loss = (args.adv_weight * l_adv + args.adv_feature_weight * l_adv_feat) / max(args.batch, 1) |
| gan_loss.backward() |
|
|
| torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0) |
| opt.step() |
| if sched is not None: |
| 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"hint {l_hint.item():.4f} sigt {l_sigt.item():.4f} hf {l_hf.item():.4f} " |
| f"quiet {l_quiet.item():.5f} adv {l_adv.item():.3f} fm {l_adv_feat.item():.3f} " |
| f"| {step/dt:.1f} sps", flush=True) |
| if step % 500 == 0: |
| torch.save(model.state_dict(), os.path.join(RUN, "last.pt")) |
| with open(os.path.join(RUN, "last_step.txt"), "w") as f: |
| f.write(str(step)) |
|
|
| torch.save(model.state_dict(), os.path.join(RUN, "final.pt")) |
| print("DECODER DONE") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
