| """Dynamic Accent - Arm A, written from scratch in PyTorch. |
| |
| Backbone: Parallel Tacotron 2 (Elias et al. 2021, arXiv 2103.14574): non-autoregressive, token durations, |
| differentiable learned upsampling (MLP over start/end offsets -> token->frame weights W), iterative mel loss. |
| Multilingual recipe: Zhang et al. 2019 (arXiv 1907.04448): separate speaker + language embeddings, |
| domain-adversarial speaker classifier (gradient reversal) on text-encoder outputs, residual VAE. |
| |
| Documented deviation (alignment): PT2 trains durations end-to-end with Soft-DTW. Soft-DTW is O(T*N)-sequential and |
| too slow for a 2 h smoke budget, so we use the internal unsupervised alignment-learning objective of |
| Badlani et al. 2022 ("One TTS Alignment To Rule Them All", arXiv 2108.10447): a small text/mel attention trained with |
| a CTC forward-sum loss + beta-binomial prior, hardened with monotonic alignment search (MAS). The hard durations |
| supervise the duration predictor and drive the SAME differentiable PT2 upsampler at train time; predicted durations |
| drive it at inference. No external aligner is used. |
| """ |
| import math |
| import numpy as np |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| from .text import VOCAB, LANGS |
|
|
| N_MELS = 100 |
|
|
|
|
| |
| def seq_mask(lens, maxlen=None): |
| maxlen = maxlen or int(lens.max()) |
| return torch.arange(maxlen, device=lens.device)[None] < lens[:, None] |
|
|
|
|
| def sinusoid(T, d, device): |
| pos = torch.arange(T, device=device).float()[:, None] |
| i = torch.arange(0, d, 2, device=device).float()[None] |
| ang = pos / (10000 ** (i / d)) |
| pe = torch.zeros(T, d, device=device) |
| pe[:, 0::2] = torch.sin(ang) |
| pe[:, 1::2] = torch.cos(ang) |
| return pe |
|
|
|
|
| class GradReverse(torch.autograd.Function): |
| """Gradient reversal. As in Zhang et al. 2019 the reversed gradient sent into the text encoder is clipped |
| (here: its L2 norm is capped at `clip`), which keeps the min-max game from blowing up the encoder.""" |
|
|
| @staticmethod |
| def forward(ctx, x, lam, clip): |
| ctx.lam, ctx.clip = lam, clip |
| return x.view_as(x) |
|
|
| @staticmethod |
| def backward(ctx, g): |
| g = -ctx.lam * g |
| if ctx.clip is not None: |
| n = g.norm() |
| if n > ctx.clip: |
| g = g * (ctx.clip / (n + 1e-12)) |
| return g, None, None |
|
|
|
|
| def grad_reverse(x, lam=1.0, clip=0.5): |
| return GradReverse.apply(x, lam, clip) |
|
|
|
|
| class ConvNorm(nn.Module): |
| def __init__(self, cin, cout, k=5, drop=0.1): |
| super().__init__() |
| self.conv = nn.Conv1d(cin, cout, k, padding=k // 2) |
| self.ln = nn.LayerNorm(cout) |
| self.drop = nn.Dropout(drop) |
|
|
| def forward(self, x, mask): |
| x = x.masked_fill(~mask[..., None], 0) |
| y = self.conv(x.transpose(1, 2)).transpose(1, 2) |
| return self.drop(F.relu(self.ln(y))) |
|
|
|
|
| class FFTBlock(nn.Module): |
| """Pre-LN self-attention + convolutional feed-forward (FastSpeech/PT2-style lightweight block).""" |
|
|
| def __init__(self, d, heads=4, ff=4, k=3, drop=0.1): |
| super().__init__() |
| self.ln1 = nn.LayerNorm(d) |
| self.att = nn.MultiheadAttention(d, heads, dropout=drop, batch_first=True) |
| self.ln2 = nn.LayerNorm(d) |
| self.c1 = nn.Conv1d(d, d * ff, k, padding=k // 2) |
| self.c2 = nn.Conv1d(d * ff, d, 1) |
| self.drop = nn.Dropout(drop) |
|
|
| def forward(self, x, mask): |
| h = self.ln1(x) |
| a, _ = self.att(h, h, h, key_padding_mask=~mask, need_weights=False) |
| x = x + self.drop(a) |
| h = self.ln2(x).masked_fill(~mask[..., None], 0).transpose(1, 2) |
| h = self.c2(F.gelu(self.c1(h))).transpose(1, 2) |
| x = x + self.drop(h) |
| return x.masked_fill(~mask[..., None], 0) |
|
|
|
|
| |
| class TextEncoder(nn.Module): |
| def __init__(self, d, n_lang, layers=6): |
| super().__init__() |
| self.emb = nn.Embedding(VOCAB, d, padding_idx=0) |
| self.lang_in = nn.Embedding(n_lang, d) |
| self.convs = nn.ModuleList([ConvNorm(d, d, 5) for _ in range(3)]) |
| self.blocks = nn.ModuleList([FFTBlock(d) for _ in range(layers)]) |
|
|
| def forward(self, tok, lang, mask): |
| x = self.emb(tok) + self.lang_in(lang) |
| for c in self.convs: |
| x = x + c(x, mask) |
| x = x + sinusoid(x.shape[1], x.shape[2], x.device)[None] |
| for b in self.blocks: |
| x = b(x, mask) |
| return x |
|
|
|
|
| class ResidualVAE(nn.Module): |
| """Global residual (prosody) encoder: mel -> q(z|x); z broadcast to every token (Zhang 2019 / Hsu 2019).""" |
|
|
| def __init__(self, zdim=16, h=256): |
| super().__init__() |
| self.convs = nn.Sequential( |
| nn.Conv1d(N_MELS, h, 3, padding=1), nn.ReLU(), |
| nn.Conv1d(h, h, 3, stride=2, padding=1), nn.ReLU(), |
| nn.Conv1d(h, h, 3, stride=2, padding=1), nn.ReLU()) |
| self.gru = nn.GRU(h, h // 2, batch_first=True, bidirectional=True) |
| self.mu = nn.Linear(h, zdim) |
| self.logvar = nn.Linear(h, zdim) |
|
|
| def forward(self, mel, mel_lens): |
| h = self.convs(mel.transpose(1, 2)).transpose(1, 2) |
| lens = torch.clamp((mel_lens + 3) // 4, min=1) |
| h, _ = self.gru(h) |
| m = seq_mask(lens, h.shape[1]).float()[..., None] |
| h = (h * m).sum(1) / m.sum(1) |
| mu, lv = self.mu(h), self.logvar(h).clamp(-8, 8) |
| z = mu + torch.randn_like(mu) * torch.exp(0.5 * lv) if self.training else mu |
| kl = 0.5 * (mu.pow(2) + lv.exp() - 1 - lv).sum(-1).mean() |
| return z, kl |
|
|
|
|
| class Aligner(nn.Module): |
| """Unsupervised text-mel alignment (Badlani et al. 2022).""" |
|
|
| def __init__(self, d, h=256, temp=0.0005): |
| super().__init__() |
| self.temp = temp |
| self.key = nn.Sequential(nn.Conv1d(d, h * 2, 3, padding=1), nn.ReLU(), nn.Conv1d(h * 2, 80, 1)) |
| self.query = nn.Sequential(nn.Conv1d(N_MELS, h * 2, 3, padding=1), nn.ReLU(), |
| nn.Conv1d(h * 2, h, 1), nn.ReLU(), nn.Conv1d(h, 80, 1)) |
|
|
| def forward(self, txt, mel, tmask, mmask, prior_log): |
| k = self.key(txt.transpose(1, 2)) |
| q = self.query(mel.transpose(1, 2)) |
| q, k = q.float(), k.float() |
| dist = (q.pow(2).sum(1)[:, :, None] + k.pow(2).sum(1)[:, None, :] |
| - 2 * torch.bmm(q.transpose(1, 2), k)).clamp_min(0) |
| logits = -self.temp * dist |
| logits = logits.masked_fill(~tmask[:, None, :], -1e4) |
| logp = F.log_softmax(logits, -1) + prior_log |
| logp = logp.masked_fill(~tmask[:, None, :], -1e4) |
| soft = F.softmax(logp, -1) * mmask[..., None] |
| return logits, logp, soft |
|
|
|
|
| def forward_sum_loss(logits, tlens, mlens, blank_logprob=-1.0): |
| """CTC forward-sum: every token must be visited monotonically. logits [B,T,N] (pre-softmax, masked).""" |
| B, T, N = logits.shape |
| lp = F.pad(logits, (1, 0), value=blank_logprob) |
| lp = F.log_softmax(lp, -1).transpose(0, 1) |
| tgt = torch.arange(1, N + 1, device=logits.device)[None].expand(B, N) |
| return F.ctc_loss(lp.float(), tgt, mlens, tlens, blank=0, reduction="mean", zero_infinity=True) |
|
|
|
|
| def _mas_np(logp, N, T): |
| """Monotonic alignment search on one [T,N] log-prob matrix -> hard [T,N].""" |
| Q = np.full((T, N), -np.inf, dtype=np.float32) |
| Q[0, 0] = logp[0, 0] |
| for t in range(1, T): |
| for n in range(min(N, t + 1)): |
| stay = Q[t - 1, n] |
| move = Q[t - 1, n - 1] if n > 0 else -np.inf |
| Q[t, n] = max(stay, move) + logp[t, n] |
| path = np.zeros((T, N), dtype=np.float32) |
| n = N - 1 |
| for t in range(T - 1, -1, -1): |
| path[t, n] = 1 |
| if t > 0 and n > 0 and (n == t or Q[t - 1, n - 1] > Q[t - 1, n]): |
| n -= 1 |
| return path |
|
|
|
|
| try: |
| import numba |
|
|
| _mas_np = numba.njit(cache=True)(_mas_np) |
| except Exception: |
| pass |
|
|
|
|
| @torch.no_grad() |
| def mas(logp, tlens, mlens): |
| out = torch.zeros_like(logp) |
| lp = logp.float().cpu().numpy() |
| for b in range(lp.shape[0]): |
| T, N = int(mlens[b]), int(tlens[b]) |
| out[b, :T, :N] = torch.from_numpy(_mas_np(np.ascontiguousarray(lp[b, :T, :N]), N, T)) |
| return out |
|
|
|
|
| _prior_cache = {} |
|
|
|
|
| def beta_binomial_prior(N, T, scale=1.0): |
| from scipy.stats import betabinom |
| key = (N, T) |
| if key not in _prior_cache: |
| x = np.arange(N)[None] |
| t = np.arange(1, T + 1)[:, None] |
| pmf = betabinom.pmf(x, N - 1, scale * t, scale * (T + 1 - t)) |
| _prior_cache[key] = torch.from_numpy(np.log(pmf + 1e-8)).float() |
| if len(_prior_cache) > 20000: |
| _prior_cache.clear() |
| return _prior_cache[key] |
|
|
|
|
| def batch_prior(tlens, mlens, N, T, device): |
| p = torch.full((len(tlens), T, N), -1e4) |
| for b in range(len(tlens)): |
| n, t = int(tlens[b]), int(mlens[b]) |
| p[b, :t, :n] = beta_binomial_prior(n, t) |
| return p.to(device) |
|
|
|
|
| class DurationPredictor(nn.Module): |
| def __init__(self, d, h=256): |
| super().__init__() |
| self.convs = nn.ModuleList([ConvNorm(d, h, 3), ConvNorm(h, h, 3), ConvNorm(h, h, 3)]) |
| self.out = nn.Linear(h, 1) |
|
|
| def forward(self, x, mask): |
| for c in self.convs: |
| x = c(x, mask) |
| return self.out(x).squeeze(-1).masked_fill(~mask, 0) |
|
|
|
|
| class LearnedUpsampler(nn.Module): |
| """PT2 differentiable upsampling: W[t,k] = softmax_k MLP(S_tk, E_tk, h_k); O = W H + proj(sum_k W C).""" |
|
|
| def __init__(self, d, hw=16, hc=8): |
| super().__init__() |
| self.se_w = nn.Linear(2, hw) |
| self.h_w = nn.Linear(d, hw) |
| self.w_out = nn.Linear(hw, 1) |
| self.se_c = nn.Linear(2, hc) |
| self.h_c = nn.Linear(d, hc) |
| self.c_out = nn.Linear(hc, d) |
|
|
| def forward(self, H, dur, tmask, T): |
| dur = dur.masked_fill(~tmask, 0) |
| e = torch.cumsum(dur, 1) |
| s = e - dur |
| t = torch.arange(T, device=H.device).float()[None, :, None] + 0.5 |
| S = (t - s[:, None, :]) / 10.0 |
| E = (e[:, None, :] - t) / 10.0 |
| se = torch.stack([S, E], -1) |
| lw = self.w_out(F.silu(self.se_w(se) + self.h_w(H)[:, None])).squeeze(-1) |
| lw = lw.masked_fill(~tmask[:, None, :], -1e4) |
| W = torch.softmax(lw, -1) |
| C = F.silu(self.se_c(se) + self.h_c(H)[:, None]) |
| out = torch.bmm(W, H) + self.c_out(torch.einsum("btn,btnc->btc", W, C)) |
| return out, W |
|
|
|
|
| class MelDecoder(nn.Module): |
| def __init__(self, d, layers=6, n_iter=3): |
| super().__init__() |
| self.blocks = nn.ModuleList([FFTBlock(d) for _ in range(layers)]) |
| self.n_iter = n_iter |
| self.outs = nn.ModuleList([nn.Linear(d, N_MELS) for _ in range(n_iter)]) |
|
|
| def forward(self, x, mask): |
| x = x + sinusoid(x.shape[1], x.shape[2], x.device)[None] |
| mels = [] |
| L = len(self.blocks) |
| for i, b in enumerate(self.blocks): |
| x = b(x, mask) |
| j = i - (L - self.n_iter) |
| if j >= 0: |
| mels.append(self.outs[j](x)) |
| return mels |
|
|
|
|
| |
| class DynamicAccentTTS(nn.Module): |
| def __init__(self, n_spk, n_lang=len(LANGS), d=384, enc_layers=6, dec_layers=6, spk_dim=128, lang_dim=64, |
| zdim=16, adv_hidden=256, mel_mean=None, mel_std=None, align_temp=0.0005): |
| super().__init__() |
| self.n_spk = n_spk |
| self.encoder = TextEncoder(d, n_lang, enc_layers) |
| self.spk_emb = nn.Embedding(n_spk, spk_dim) |
| self.lang_out = nn.Embedding(n_lang, lang_dim) |
| self.vae = ResidualVAE(zdim) |
| self.cond = nn.Linear(d + spk_dim + lang_dim + zdim, d) |
| self.adv = nn.Sequential(nn.LayerNorm(d), nn.Linear(d, adv_hidden), nn.ReLU(), nn.Linear(adv_hidden, n_spk)) |
| self.aligner = Aligner(d, temp=align_temp) |
| self.dur = DurationPredictor(d) |
| self.up = LearnedUpsampler(d) |
| self.dec_spk = nn.Linear(spk_dim, d) |
| self.decoder = MelDecoder(d, dec_layers) |
| self.register_buffer("mel_mean", torch.tensor(mel_mean if mel_mean is not None else -2.5)) |
| self.register_buffer("mel_std", torch.tensor(mel_std if mel_std is not None else 2.2)) |
|
|
| def _cond(self, Ht, spk, lang, z): |
| B, N, _ = Ht.shape |
| s = self.spk_emb(spk)[:, None].expand(B, N, -1) |
| zz = z[:, None].expand(B, N, -1) |
| return self.cond(torch.cat([Ht, s, self.lang_out(lang), zz], -1)), s |
|
|
| def forward(self, tok, lang, tlens, mel, mlens, spk, step=0, adv_lambda=1.0, bin_start=4000, prior=None): |
| tmask, mmask = seq_mask(tlens, tok.shape[1]), seq_mask(mlens, mel.shape[1]) |
| meln = (mel - self.mel_mean) / self.mel_std |
| Ht = self.encoder(tok, lang, tmask) |
| |
| adv_logits = self.adv(grad_reverse(Ht, adv_lambda)) |
| adv_ce = F.cross_entropy(adv_logits[tmask], spk[:, None].expand_as(tok)[tmask]) |
| adv_acc = (adv_logits[tmask].argmax(-1) == spk[:, None].expand_as(tok)[tmask]).float().mean() |
| z, kl = self.vae(meln, mlens) |
| H, _ = self._cond(Ht, spk, lang, z) |
| |
| if prior is None: |
| prior = batch_prior(tlens, mlens, tok.shape[1], mel.shape[1], tok.device) |
| logits, logp, soft = self.aligner(self.encoder.emb(tok) + Ht.detach(), meln, tmask, mmask, prior) |
| l_fs = forward_sum_loss(logits.masked_fill(~mmask[..., None], -1e4) + prior, tlens, mlens) |
| hard = mas(logp, tlens, mlens) |
| l_bin = -(hard * torch.log(soft.clamp_min(1e-8))).sum() / hard.sum() if step >= bin_start else soft.new_zeros(()) |
| dur_t = hard.sum(1) |
| |
| logd = self.dur(H.detach() if step < 1000 else H, tmask) |
| l_dur = F.mse_loss(logd[tmask], torch.log1p(dur_t[tmask])) |
| |
| up, W = self.up(H, dur_t, tmask, mel.shape[1]) |
| x = up + self.dec_spk(self.spk_emb(spk))[:, None] |
| mels = self.decoder(x, mmask) |
| m = mmask[..., None].float() |
| l_mel = sum((F.l1_loss(o, meln, reduction="none") * m).sum() / (m.sum() * N_MELS) for o in mels) / len(mels) |
| return dict(mel=l_mel, dur=l_dur, fs=l_fs, bin=l_bin, kl=kl, adv=adv_ce, adv_acc=adv_acc.detach()), mels[-1] |
|
|
| @torch.no_grad() |
| def infer(self, tok, lang, spk, pace=1.0, z=None): |
| tlens = torch.full((tok.shape[0],), tok.shape[1], device=tok.device) |
| tmask = seq_mask(tlens, tok.shape[1]) |
| Ht = self.encoder(tok, lang, tmask) |
| if z is None: |
| z = torch.zeros(tok.shape[0], self.vae.mu.out_features, device=tok.device) |
| H, _ = self._cond(Ht, spk, lang, z) |
| dur = (torch.expm1(self.dur(H, tmask)).clamp(min=0) * pace).masked_fill(~tmask, 0) |
| T = max(int(dur.sum(1).max().round().item()), 1) |
| up, W = self.up(H, dur, tmask, T) |
| x = up + self.dec_spk(self.spk_emb(spk))[:, None] |
| mel = self.decoder(x, torch.ones(tok.shape[0], T, dtype=torch.bool, device=tok.device))[-1] |
| return mel * self.mel_std + self.mel_mean, dur |
|
|
|
|
| def count_params(m): |
| return sum(p.numel() for p in m.parameters()) |
|
|
