from __future__ import annotations import argparse import json import math import random import sys import time from dataclasses import asdict, dataclass from pathlib import Path import torch from torch import nn import torch.nn.functional as F import torchaudio SCRIPT_ROOT = Path(__file__).resolve().parent PROJECT_ROOT = SCRIPT_ROOT.parents[0] FRONTEND_ROOT = PROJECT_ROOT / "third_party" / "tiny_tts_frontend" sys.path = [str(FRONTEND_ROOT), str(SCRIPT_ROOT)] + [p for p in sys.path if p] from inflect_nano.vocoder import HifiGanConfig, HifiGanGenerator, MelFrontend, make_config @dataclass class MicroFastSpeechConfig: vocab_size: int = 256 tone_size: int = 16 lang_size: int = 4 n_mels: int = 80 hidden: int = 168 encoder_layers: int = 5 decoder_layers: int = 6 decoder_ff_mult: int = 3 kernel_size: int = 7 speaker_count: int = 2 speaker_dim: int = 64 dropout: float = 0.08 sample_rate: int = 24000 max_frames: int = 1400 postnet_scale: float = 0.10 use_frame_pitch: bool = True abs_frame_bins: int = 512 use_contextual_predictors: bool = False use_group_duration_planner: bool = False def count_parameters(model: nn.Module) -> int: return sum(p.numel() for p in model.parameters()) def load_rows(path: Path, max_rows: int = 0) -> list[dict]: rows = [] with path.open("r", encoding="utf-8") as f: for line in f: if line.strip(): row = json.loads(line) if Path(str(row.get("target_audio") or "")).is_file(): rows.append(row) if max_rows and len(rows) >= max_rows: break if not rows: raise RuntimeError(f"No usable rows in {path}") return rows def load_audio(path: str, sample_rate: int, max_seconds: float) -> torch.Tensor: import soundfile as _sf # avoid torchaudio.load (needs torchcodec/ffmpeg on torch>=2.1) _a, sr = _sf.read(path, dtype="float32", always_2d=True) wav = torch.from_numpy(_a.T) # [ch, T] if wav.shape[0] > 1: wav = wav.mean(dim=0, keepdim=True) if sr != sample_rate: wav = torchaudio.functional.resample(wav, sr, sample_rate) return wav[:, : int(sample_rate * max_seconds)].squeeze(0).clamp(-1.0, 1.0) def fit_durations(durations: list[int], target_frames: int) -> list[int]: if sum(durations) == target_frames: return list(durations) total = max(1, sum(durations)) raw = [max(0.0, d * target_frames / total) for d in durations] out = [int(math.floor(x)) for x in raw] order = sorted(((raw[i] - out[i], i) for i in range(len(out))), reverse=True) for _, idx in order[: max(0, target_frames - sum(out))]: out[idx] += 1 while sum(out) > target_frames: idx = max(range(len(out)), key=lambda i: out[i]) out[idx] -= 1 return out def pad_1d(items: list[torch.Tensor], value: float = 0.0) -> torch.Tensor: max_len = max(x.numel() for x in items) out = torch.full((len(items), max_len), value, dtype=items[0].dtype) for i, item in enumerate(items): out[i, : item.numel()] = item return out def pad_2d(items: list[torch.Tensor], value: float = 0.0) -> torch.Tensor: max_len = max(x.shape[0] for x in items) dim = items[0].shape[1] out = torch.full((len(items), max_len, dim), value, dtype=items[0].dtype) for i, item in enumerate(items): out[i, : item.shape[0]] = item return out def pad_mels(items: list[torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]: max_len = max(x.shape[-1] for x in items) n_mels = items[0].shape[0] out = torch.zeros(len(items), n_mels, max_len, dtype=items[0].dtype) mask = torch.zeros(len(items), max_len, dtype=torch.bool) for i, mel in enumerate(items): frames = mel.shape[-1] out[i, :, :frames] = mel mask[i, :frames] = True return out, mask def pad_wavs(items: list[torch.Tensor], frames: list[int], hop_size: int) -> torch.Tensor: max_len = max(max(1, int(frame_count)) * hop_size for frame_count in frames) out = torch.zeros(len(items), max_len, dtype=items[0].dtype) for i, (wav, frame_count) in enumerate(zip(items, frames)): length = max(1, int(frame_count)) * hop_size cropped = wav[:length] out[i, : cropped.numel()] = cropped return out def aggregate_token_features(mel: torch.Tensor, durations: list[int]) -> tuple[torch.Tensor, torch.Tensor]: # mel: [80, frames], log-mel from the exact V2+ frontend. frames = mel.shape[-1] amp = torch.exp(mel).clamp_min(1e-5) energy_frame = mel.mean(dim=0) bins = torch.linspace(0.0, 1.0, mel.shape[0], device=mel.device).view(-1, 1) bright_frame = (amp * bins).sum(dim=0) / amp.sum(dim=0).clamp_min(1e-5) energies = [] brights = [] pos = 0 for dur in durations: end = min(frames, pos + max(0, int(dur))) if end > pos: energies.append(energy_frame[pos:end].mean()) brights.append(bright_frame[pos:end].mean()) else: energies.append(torch.zeros((), device=mel.device, dtype=mel.dtype)) brights.append(torch.zeros((), device=mel.device, dtype=mel.dtype)) pos = end return torch.stack(energies), torch.stack(brights) def aggregate_token_pitch(pitch_frame: torch.Tensor, durations: list[int]) -> torch.Tensor: # pitch_frame: [2, frames] with normalized log-f0 and voiced flag. frames = pitch_frame.shape[-1] out = [] pos = 0 for dur in durations: end = min(frames, pos + max(0, int(dur))) if end > pos: span = pitch_frame[:, pos:end] voiced = span[1].mean() voiced_mask = span[1] > 0.5 if bool(voiced_mask.any()): log_f0 = span[0, voiced_mask].mean() else: log_f0 = torch.zeros((), dtype=pitch_frame.dtype) out.append(torch.stack([log_f0, voiced])) else: out.append(torch.zeros(2, dtype=pitch_frame.dtype)) pos = end return torch.stack(out, dim=0) def extract_pitch_features(wav: torch.Tensor, sample_rate: int, frames: int) -> torch.Tensor: # Returns [2, frames]: normalized log-f0 and voiced flag. The detector can # produce octave spikes, so clip to speech range and median-smooth lightly. pitch = torchaudio.functional.detect_pitch_frequency( wav.unsqueeze(0).cpu(), sample_rate, frame_time=256 / sample_rate, ).squeeze(0) if pitch.numel() < frames: pitch = F.pad(pitch, (0, frames - pitch.numel()), value=0.0) pitch = pitch[:frames] voiced = ((pitch >= 55.0) & (pitch <= 420.0)).float() pitch = pitch.clamp(55.0, 420.0) # Median filter over 5 frames to reduce spurious jumps. if pitch.numel() >= 5: padded = F.pad(pitch.view(1, 1, -1), (2, 2), mode="replicate") windows = padded.unfold(-1, 5, 1).squeeze(0).squeeze(0) pitch = windows.median(dim=-1).values log_f0 = (torch.log(pitch) - math.log(140.0)) / 0.45 log_f0 = log_f0.clamp(-3.0, 3.0) * voiced return torch.stack([log_f0, voiced], dim=0) class ConvFFNBlock(nn.Module): def __init__(self, hidden: int, kernel_size: int, dropout: float, ff_mult: int = 4) -> None: super().__init__() pad = kernel_size // 2 self.norm1 = nn.LayerNorm(hidden) self.depth = nn.Conv1d(hidden, hidden * 2, kernel_size, padding=pad, groups=hidden) self.point = nn.Conv1d(hidden, hidden, 1) self.drop = nn.Dropout(dropout) self.norm2 = nn.LayerNorm(hidden) self.ff = nn.Sequential( nn.Linear(hidden, hidden * ff_mult), nn.SiLU(), nn.Dropout(dropout), nn.Linear(hidden * ff_mult, hidden), ) def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Tensor: y = self.norm1(x).transpose(1, 2) a, b = self.depth(y).chunk(2, dim=1) y = self.point(a * torch.sigmoid(b)).transpose(1, 2) x = x + self.drop(y) x = x + self.drop(self.ff(self.norm2(x))) if mask is not None: x = x * mask.unsqueeze(-1) return x class MicroFastSpeech(nn.Module): def __init__(self, cfg: MicroFastSpeechConfig) -> None: super().__init__() self.cfg = cfg # Phone id 0 is a real inserted blank/silence token from TinyTTS, not # padding. Padding is tracked by duration masks instead. self.phone = nn.Embedding(cfg.vocab_size, cfg.hidden) self.tone = nn.Embedding(cfg.tone_size, cfg.hidden) self.lang = nn.Embedding(cfg.lang_size, cfg.hidden) self.speaker = nn.Embedding(cfg.speaker_count, cfg.speaker_dim) self.speaker_proj = nn.Linear(cfg.speaker_dim, cfg.hidden) self.encoder = nn.ModuleList([ConvFFNBlock(cfg.hidden, cfg.kernel_size, cfg.dropout) for _ in range(cfg.encoder_layers)]) self.duration_head = nn.Sequential(nn.LayerNorm(cfg.hidden), nn.Linear(cfg.hidden, cfg.hidden), nn.SiLU(), nn.Linear(cfg.hidden, 1)) self.energy_head = nn.Sequential(nn.LayerNorm(cfg.hidden), nn.Linear(cfg.hidden, cfg.hidden // 2), nn.SiLU(), nn.Linear(cfg.hidden // 2, 1)) self.bright_head = nn.Sequential(nn.LayerNorm(cfg.hidden), nn.Linear(cfg.hidden, cfg.hidden // 2), nn.SiLU(), nn.Linear(cfg.hidden // 2, 1)) self.pitch_head = nn.Sequential(nn.LayerNorm(cfg.hidden), nn.Linear(cfg.hidden, cfg.hidden), nn.SiLU(), nn.Linear(cfg.hidden, 2)) self.group_duration_delta = nn.Linear(cfg.hidden, 1) if cfg.use_group_duration_planner else None if self.group_duration_delta is not None: nn.init.zeros_(self.group_duration_delta.weight) nn.init.zeros_(self.group_duration_delta.bias) self.predictor_context = ( ConvFFNBlock(cfg.hidden, 5, cfg.dropout, 2) if cfg.use_contextual_predictors else nn.Identity() ) self.duration_delta = nn.Linear(cfg.hidden, 1) if cfg.use_contextual_predictors else None self.energy_delta = nn.Linear(cfg.hidden, 1) if cfg.use_contextual_predictors else None self.bright_delta = nn.Linear(cfg.hidden, 1) if cfg.use_contextual_predictors else None self.pitch_delta = nn.Linear(cfg.hidden, 2) if cfg.use_contextual_predictors else None if cfg.use_contextual_predictors: for layer in (self.duration_delta, self.energy_delta, self.bright_delta, self.pitch_delta): nn.init.zeros_(layer.weight) nn.init.zeros_(layer.bias) self.energy_proj = nn.Linear(1, cfg.hidden) self.bright_proj = nn.Linear(1, cfg.hidden) self.pitch_proj = nn.Sequential(nn.Linear(2, cfg.hidden), nn.SiLU(), nn.Linear(cfg.hidden, cfg.hidden)) self.abs_frame = nn.Embedding(cfg.abs_frame_bins, cfg.hidden) self.frame_proj = nn.Sequential(nn.Linear(8, cfg.hidden), nn.SiLU(), nn.Linear(cfg.hidden, cfg.hidden)) self.local_ctx = nn.Sequential( nn.Linear(cfg.hidden * 3, cfg.hidden * 2), nn.SiLU(), nn.Linear(cfg.hidden * 2, cfg.hidden), ) self.decoder = nn.ModuleList([ConvFFNBlock(cfg.hidden, cfg.kernel_size, cfg.dropout, cfg.decoder_ff_mult) for _ in range(cfg.decoder_layers)]) self.frame_gru = nn.GRU(cfg.hidden, cfg.hidden // 2, num_layers=1, batch_first=True, bidirectional=True) self.mel_head = nn.Sequential(nn.LayerNorm(cfg.hidden), nn.Linear(cfg.hidden, cfg.hidden), nn.SiLU(), nn.Linear(cfg.hidden, cfg.n_mels)) self.postnet = nn.Sequential( nn.Conv1d(cfg.n_mels, cfg.hidden, 5, padding=2), nn.Tanh(), nn.Conv1d(cfg.hidden, cfg.hidden, 5, padding=2), nn.Tanh(), nn.Conv1d(cfg.hidden, cfg.n_mels, 5, padding=2), ) def encode(self, phone: torch.Tensor, tone: torch.Tensor, lang: torch.Tensor, speaker: torch.Tensor, token_mask: torch.Tensor) -> torch.Tensor: x = self.phone(phone) + self.tone(tone.clamp_max(self.cfg.tone_size - 1)) + self.lang(lang.clamp_max(self.cfg.lang_size - 1)) x = x + self.speaker_proj(self.speaker(speaker)).unsqueeze(1) x = x * token_mask.unsqueeze(-1) for block in self.encoder: x = block(x, token_mask) return x def regulate(self, encoded: torch.Tensor, durations: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: batch_frames = [] batch_meta = [] lengths = [] device = encoded.device for b in range(encoded.shape[0]): reps = [] meta = [] durs = durations[b].long().clamp_min(0) token_count = max(1, int((durs > 0).sum().item())) for i, dur_t in enumerate(durs.tolist()): dur = int(dur_t) if dur <= 0: continue reps.append(encoded[b, i].view(1, -1).expand(dur, -1)) rel = torch.linspace(0.0, 1.0, dur, device=device) token_pos = torch.full((dur,), i / max(1, token_count - 1), device=device) log_dur = torch.full((dur,), math.log1p(dur) / 6.0, device=device) inv_rel = 1.0 - rel center = 1.0 - torch.abs(rel * 2.0 - 1.0) meta.append( torch.stack( [ rel, inv_rel, center, torch.sin(rel * math.pi), torch.cos(rel * math.pi), token_pos, log_dur, torch.full_like(rel, dur / 40.0), ], dim=-1, ) ) if reps: frames = torch.cat(reps, dim=0) frame_meta = torch.cat(meta, dim=0) else: frames = encoded[b, :1] frame_meta = torch.zeros(1, 8, device=device) batch_frames.append(frames[: self.cfg.max_frames]) batch_meta.append(frame_meta[: self.cfg.max_frames]) lengths.append(min(frames.shape[0], self.cfg.max_frames)) max_len = max(lengths) out = torch.zeros(encoded.shape[0], max_len, encoded.shape[-1], device=device) meta_out = torch.zeros(encoded.shape[0], max_len, 8, device=device) mask = torch.zeros(encoded.shape[0], max_len, dtype=torch.bool, device=device) for b, frames in enumerate(batch_frames): n = min(frames.shape[0], max_len) out[b, :n] = frames[:n] meta_out[b, :n] = batch_meta[b][:n] mask[b, :n] = True return out, meta_out, mask def add_local_context(self, encoded: torch.Tensor, durations: torch.Tensor) -> torch.Tensor: device = encoded.device batch_frames = [] for b in range(encoded.shape[0]): reps = [] durs = durations[b].long().clamp_min(0) for i, dur_t in enumerate(durs.tolist()): dur = int(dur_t) if dur <= 0: continue prev_i = max(0, i - 1) next_i = min(encoded.shape[1] - 1, i + 1) ctx = torch.cat([encoded[b, prev_i], encoded[b, i], encoded[b, next_i]], dim=-1) reps.append(ctx.view(1, -1).expand(dur, -1)) if reps: frames = torch.cat(reps, dim=0) else: frames = torch.zeros(1, encoded.shape[-1] * 3, device=device) batch_frames.append(frames[: self.cfg.max_frames]) max_len = max(x.shape[0] for x in batch_frames) ctx_out = torch.zeros(encoded.shape[0], max_len, encoded.shape[-1] * 3, device=device) for b, frames in enumerate(batch_frames): ctx_out[b, : frames.shape[0]] = frames return self.local_ctx(ctx_out) def expand_token_feature(self, feature: torch.Tensor, durations: torch.Tensor) -> torch.Tensor: device = feature.device batch_frames = [] for b in range(feature.shape[0]): reps = [] durs = durations[b].long().clamp_min(0) for i, dur_t in enumerate(durs.tolist()): dur = int(dur_t) if dur <= 0: continue reps.append(feature[b, i].view(1, -1).expand(dur, -1)) if reps: frames = torch.cat(reps, dim=0) else: frames = torch.zeros(1, feature.shape[-1], device=device) batch_frames.append(frames[: self.cfg.max_frames]) max_len = max(x.shape[0] for x in batch_frames) out = torch.zeros(feature.shape[0], max_len, feature.shape[-1], device=device) for b, frames in enumerate(batch_frames): out[b, : frames.shape[0]] = frames return out def forward( self, phone: torch.Tensor, tone: torch.Tensor, lang: torch.Tensor, speaker: torch.Tensor, durations: torch.Tensor, energy_target: torch.Tensor | None = None, bright_target: torch.Tensor | None = None, pitch_frame: torch.Tensor | None = None, predicted_prosody_mix: float = 0.0, detach_mixed_predictions: bool = True, ) -> dict[str, torch.Tensor]: token_mask = durations.gt(0) encoded = self.encode(phone, tone, lang, speaker, token_mask) log_dur, energy_pred, bright_pred, pitch_pred = self.predict_prosody(encoded, token_mask) mixed_energy_pred = energy_pred.detach() if detach_mixed_predictions else energy_pred mixed_bright_pred = bright_pred.detach() if detach_mixed_predictions else bright_pred if energy_target is not None: energy = torch.lerp(energy_target, mixed_energy_pred, predicted_prosody_mix) else: energy = energy_pred if bright_target is not None: bright = torch.lerp(bright_target, mixed_bright_pred, predicted_prosody_mix) else: bright = bright_pred conditioned = encoded + self.energy_proj(energy.unsqueeze(-1)) + self.bright_proj(bright.unsqueeze(-1)) frames, frame_meta, frame_mask = self.regulate(conditioned, durations) x = frames + self.frame_proj(frame_meta) + self.add_local_context(conditioned, durations) pos = torch.arange(x.shape[1], device=x.device) pos = torch.div(pos * self.cfg.abs_frame_bins, max(1, self.cfg.max_frames), rounding_mode="floor").clamp_max( self.cfg.abs_frame_bins - 1 ) x = x + self.abs_frame(pos).unsqueeze(0) if self.cfg.use_frame_pitch: if pitch_frame is not None: pitch_frame = pitch_frame[:, :, : x.shape[1]].transpose(1, 2) if pitch_frame.shape[1] < x.shape[1]: pitch_frame = F.pad(pitch_frame, (0, 0, 0, x.shape[1] - pitch_frame.shape[1])) if predicted_prosody_mix > 0.0: mixed_pitch_pred = pitch_pred.detach() if detach_mixed_predictions else pitch_pred predicted_pitch_frame = self.expand_token_feature(mixed_pitch_pred, durations)[:, : x.shape[1]] pitch_frame = torch.lerp(pitch_frame, predicted_pitch_frame, predicted_prosody_mix) else: pitch_frame = self.expand_token_feature(pitch_pred, durations)[:, : x.shape[1]] x = x + self.pitch_proj(pitch_frame) for block in self.decoder: x = block(x, frame_mask) x = x + self.frame_gru(x)[0] mel = self.mel_head(x).transpose(1, 2) mel = mel + self.cfg.postnet_scale * self.postnet(mel) group_log_dur, group_mask = self.group_log_durations(phone, log_dur, encoded) return { "mel": mel, "frame_mask": frame_mask, "log_dur": log_dur, "group_log_dur": group_log_dur, "group_mask": group_mask, "energy": energy_pred, "bright": bright_pred, "pitch": pitch_pred, "token_mask": token_mask, } def predict_prosody( self, encoded: torch.Tensor, token_mask: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: log_dur = self.duration_head(encoded).squeeze(-1) energy = self.energy_head(encoded).squeeze(-1) bright = self.bright_head(encoded).squeeze(-1) pitch = self.pitch_head(encoded) if self.cfg.use_contextual_predictors: context = self.predictor_context(encoded, token_mask) log_dur = log_dur + self.duration_delta(context).squeeze(-1) energy = energy + self.energy_delta(context).squeeze(-1) bright = bright + self.bright_delta(context).squeeze(-1) pitch = pitch + self.pitch_delta(context) return log_dur, energy, bright, pitch def group_log_durations( self, phone: torch.Tensor, log_dur: torch.Tensor, encoded: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: """Predict stable blank-plus-phone region durations at visible phones.""" base_dur = torch.expm1(log_dur).clamp_min(0.05) grouped = torch.zeros_like(log_dur) group_mask = torch.zeros_like(phone, dtype=torch.bool) delta = self.group_duration_delta(encoded).squeeze(-1) if self.group_duration_delta is not None else None for batch_index in range(phone.shape[0]): pending: list[torch.Tensor] = [] last_visible: int | None = None for token_index in range(phone.shape[1]): pending.append(base_dur[batch_index, token_index]) if int(phone[batch_index, token_index].item()) != 0: value = torch.stack(pending).sum() if delta is not None: value = value * torch.exp(delta[batch_index, token_index].clamp(-1.5, 1.5)) grouped[batch_index, token_index] = torch.log1p(value) group_mask[batch_index, token_index] = True pending = [] last_visible = token_index if pending and last_visible is not None: value = torch.expm1(grouped[batch_index, last_visible]) + torch.stack(pending).sum() grouped[batch_index, last_visible] = torch.log1p(value) return grouped, group_mask def apply_group_duration_plan( self, phone: torch.Tensor, log_dur: torch.Tensor, encoded: torch.Tensor, length_scale: float, max_duration: int ) -> torch.Tensor: base = torch.expm1(log_dur).clamp_min(0.05) group_log, _ = self.group_log_durations(phone, log_dur, encoded) planned = torch.zeros_like(base, dtype=torch.long) for batch_index in range(phone.shape[0]): pending: list[int] = [] last_visible: int | None = None for token_index in range(phone.shape[1]): pending.append(token_index) if int(phone[batch_index, token_index].item()) != 0: target = max(len(pending), int(round(float(torch.expm1(group_log[batch_index, token_index]) * length_scale)))) weights = base[batch_index, pending] remaining = target - len(pending) raw = weights / weights.sum().clamp_min(1e-6) * remaining allocated = torch.ones_like(raw, dtype=torch.long) + torch.floor(raw).long() remainder = target - int(allocated.sum().item()) if remainder > 0: order = torch.argsort(raw - torch.floor(raw), descending=True) allocated[order[:remainder]] += 1 planned[batch_index, pending] = allocated pending = [] last_visible = token_index if pending and last_visible is not None: planned[batch_index, last_visible] += max(1, int(round(float(base[batch_index, pending].sum() * length_scale)))) return planned.clamp(0, max_duration) @torch.no_grad() def infer( self, phone: torch.Tensor, tone: torch.Tensor, lang: torch.Tensor, speaker: torch.Tensor, length_scale: float = 1.0, min_duration: int = 1, max_duration: int = 80, pitch_scale: float = 1.0, energy_scale: float = 1.0, smooth_predictors: bool = False, ) -> torch.Tensor: # In single-sample inference there is no padded tail; id 0 remains the # explicit blank/pause token and must keep duration. token_mask = torch.ones_like(phone, dtype=torch.bool) encoded = self.encode(phone, tone, lang, speaker, token_mask) log_dur, energy, bright, pitch = self.predict_prosody(encoded, token_mask) if self.group_duration_delta is not None: durations = self.apply_group_duration_plan(phone, log_dur, encoded, length_scale, max_duration) durations = durations.masked_fill(~token_mask, 0) else: pred_dur = torch.expm1(log_dur).clamp(0, max_duration) * length_scale durations = torch.round(pred_dur).long().clamp_min(min_duration).masked_fill(~token_mask, 0) energy = energy * energy_scale pitch = torch.stack([pitch[..., 0] * pitch_scale, pitch[..., 1].clamp(0.0, 1.0)], dim=-1) if smooth_predictors and phone.shape[1] >= 3: energy = F.avg_pool1d(energy.unsqueeze(1), 3, stride=1, padding=1).squeeze(1) bright = F.avg_pool1d(bright.unsqueeze(1), 3, stride=1, padding=1).squeeze(1) pitch_t = pitch.transpose(1, 2) pitch = F.avg_pool1d(pitch_t, 3, stride=1, padding=1).transpose(1, 2) conditioned = encoded + self.energy_proj(energy.unsqueeze(-1)) + self.bright_proj(bright.unsqueeze(-1)) frames, frame_meta, frame_mask = self.regulate(conditioned, durations) x = frames + self.frame_proj(frame_meta) + self.add_local_context(conditioned, durations) pos = torch.arange(x.shape[1], device=x.device) pos = torch.div(pos * self.cfg.abs_frame_bins, max(1, self.cfg.max_frames), rounding_mode="floor").clamp_max( self.cfg.abs_frame_bins - 1 ) x = x + self.abs_frame(pos).unsqueeze(0) if self.cfg.use_frame_pitch: pitch_frame = self.expand_token_feature(pitch, durations)[:, : x.shape[1]] x = x + self.pitch_proj(pitch_frame) for block in self.decoder: x = block(x, frame_mask) x = x + self.frame_gru(x)[0] mel = self.mel_head(x).transpose(1, 2) mel = mel + self.cfg.postnet_scale * self.postnet(mel) return mel def collate(batch: list[dict], cfg: MicroFastSpeechConfig, mel_frontend: MelFrontend, device: torch.device, max_seconds: float, hop_size: int): phones = [torch.LongTensor(x["phone_ids"]) for x in batch] tones = [torch.LongTensor(x["tone_ids"]) for x in batch] langs = [torch.LongTensor(x["lang_ids"]) for x in batch] durations_raw = [list(map(int, x["hifigan_durations"])) for x in batch] speakers = torch.LongTensor([int(x["speaker_id"]) for x in batch]) phone = pad_1d(phones, 0).long() tone = pad_1d(tones, 0).long() lang = pad_1d(langs, 0).long() mels = [] durations = [] energies = [] brights = [] pitches = [] token_pitches = [] wavs = [] frame_counts = [] with torch.no_grad(): for row, dur in zip(batch, durations_raw): wav_1d = load_audio(str(row["target_audio"]), cfg.sample_rate, max_seconds) wav = wav_1d.unsqueeze(0).to(device) mel = mel_frontend(wav).squeeze(0).detach().cpu() dur = fit_durations(dur[: len(row["phone_ids"])], min(mel.shape[-1], cfg.max_frames)) mel = mel[:, : sum(dur)] energy, bright = aggregate_token_features(mel, dur) pitch = extract_pitch_features(wav_1d, cfg.sample_rate, mel.shape[-1]) token_pitch = aggregate_token_pitch(pitch, dur) mels.append(mel) durations.append(torch.LongTensor(dur)) energies.append(energy) brights.append(bright) pitches.append(pitch) token_pitches.append(token_pitch) wavs.append(wav_1d) frame_counts.append(mel.shape[-1]) duration = pad_1d(durations, 0).long() energy = pad_1d(energies, 0.0).float() bright = pad_1d(brights, 0.0).float() token_pitch = pad_2d(token_pitches, 0.0).float() target_mel, frame_mask = pad_mels(mels) pitch_frame, _ = pad_mels(pitches) target_wav = pad_wavs(wavs, frame_counts, hop_size) return ( phone.to(device), tone.to(device), lang.to(device), speakers.to(device), duration.to(device), energy.to(device), bright.to(device), token_pitch.to(device), target_mel.to(device), frame_mask.to(device), pitch_frame.to(device), target_wav.to(device), ) def prepare_row_features( row: dict, cfg: MicroFastSpeechConfig, mel_frontend: MelFrontend, device: torch.device, max_seconds: float, ) -> dict: dur = list(map(int, row["hifigan_durations"])) wav_1d = load_audio(str(row["target_audio"]), cfg.sample_rate, max_seconds) with torch.no_grad(): wav = wav_1d.unsqueeze(0).to(device) mel = mel_frontend(wav).squeeze(0).detach().cpu() dur = fit_durations(dur[: len(row["phone_ids"])], min(mel.shape[-1], cfg.max_frames)) mel = mel[:, : sum(dur)] energy, bright = aggregate_token_features(mel, dur) pitch = extract_pitch_features(wav_1d, cfg.sample_rate, mel.shape[-1]) token_pitch = aggregate_token_pitch(pitch, dur) return { "phone": torch.LongTensor(row["phone_ids"]), "tone": torch.LongTensor(row["tone_ids"]), "lang": torch.LongTensor(row["lang_ids"]), "speaker": int(row["speaker_id"]), "duration": torch.LongTensor(dur), "energy": energy.float(), "bright": bright.float(), "token_pitch": token_pitch.float(), "target_mel": mel.float(), "pitch_frame": pitch.float(), "target_wav": wav_1d.float(), "frame_count": int(mel.shape[-1]), } def collate_prepared(batch: list[dict], device: torch.device, hop_size: int): phone = pad_1d([x["phone"] for x in batch], 0).long() tone = pad_1d([x["tone"] for x in batch], 0).long() lang = pad_1d([x["lang"] for x in batch], 0).long() speakers = torch.LongTensor([int(x["speaker"]) for x in batch]) duration = pad_1d([x["duration"] for x in batch], 0).long() energy = pad_1d([x["energy"] for x in batch], 0.0).float() bright = pad_1d([x["bright"] for x in batch], 0.0).float() token_pitch = pad_2d([x["token_pitch"] for x in batch], 0.0).float() target_mel, frame_mask = pad_mels([x["target_mel"] for x in batch]) pitch_frame, _ = pad_mels([x["pitch_frame"] for x in batch]) target_wav = pad_wavs([x["target_wav"] for x in batch], [int(x["frame_count"]) for x in batch], hop_size) return ( phone.to(device), tone.to(device), lang.to(device), speakers.to(device), duration.to(device), energy.to(device), bright.to(device), token_pitch.to(device), target_mel.to(device), frame_mask.to(device), pitch_frame.to(device), target_wav.to(device), ) def masked_l1(pred: torch.Tensor, target: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: common = min(pred.shape[-1], target.shape[-1], mask.shape[-1]) pred = pred[..., :common] target = target[..., :common] mask = mask[:, :common].unsqueeze(1) return (torch.abs(pred - target) * mask).sum() / (mask.sum() * pred.shape[1]).clamp_min(1.0) def masked_mse(pred: torch.Tensor, target: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: common = min(pred.shape[-1], target.shape[-1], mask.shape[-1]) pred = pred[..., :common] target = target[..., :common] mask = mask[:, :common].unsqueeze(1) return (((pred - target) ** 2) * mask).sum() / (mask.sum() * pred.shape[1]).clamp_min(1.0) def weighted_frame_l1(pred: torch.Tensor, target: torch.Tensor, wmap: torch.Tensor, valid_count: torch.Tensor) -> torch.Tensor: """L1 with a PER-FRAME weight map wmap [B,T] (0 in pad). Normalized by valid_count*n_mels so that wmap == frame_mask reproduces masked_l1 exactly. Used for per-language mel weighting (rank 4).""" common = min(pred.shape[-1], target.shape[-1], wmap.shape[-1]) p = pred[..., :common]; t = target[..., :common]; w = wmap[:, :common].unsqueeze(1) return (torch.abs(p - t) * w).sum() / (valid_count * pred.shape[1]).clamp_min(1.0) def weighted_frame_mse(pred: torch.Tensor, target: torch.Tensor, wmap: torch.Tensor, valid_count: torch.Tensor) -> torch.Tensor: common = min(pred.shape[-1], target.shape[-1], wmap.shape[-1]) p = pred[..., :common]; t = target[..., :common]; w = wmap[:, :common].unsqueeze(1) return (((p - t) ** 2) * w).sum() / (valid_count * pred.shape[1]).clamp_min(1.0) def masked_delta_loss(pred: torch.Tensor, target: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: common = min(pred.shape[-1], target.shape[-1], mask.shape[-1]) if common < 2: return torch.zeros((), device=pred.device) dp = pred[..., 1:common] - pred[..., : common - 1] dt = target[..., 1:common] - target[..., : common - 1] dm = (mask[:, 1:common] & mask[:, : common - 1]).unsqueeze(1) return (torch.abs(dp - dt) * dm).sum() / (dm.sum() * pred.shape[1]).clamp_min(1.0) def masked_accel_loss(pred: torch.Tensor, target: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: common = min(pred.shape[-1], target.shape[-1], mask.shape[-1]) if common < 3: return torch.zeros((), device=pred.device) dp = pred[..., 2:common] - 2.0 * pred[..., 1 : common - 1] + pred[..., : common - 2] dt = target[..., 2:common] - 2.0 * target[..., 1 : common - 1] + target[..., : common - 2] dm = (mask[:, 2:common] & mask[:, 1 : common - 1] & mask[:, : common - 2]).unsqueeze(1) return (torch.abs(dp - dt) * dm).sum() / (dm.sum() * pred.shape[1]).clamp_min(1.0) def token_mse(pred: torch.Tensor, target: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: common = min(pred.shape[1], target.shape[1], mask.shape[1]) pred = pred[:, :common] target = target[:, :common] mask = mask[:, :common] return (((pred - target) ** 2) * mask).sum() / mask.sum().clamp_min(1.0) def token_mse_nd(pred: torch.Tensor, target: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: common = min(pred.shape[1], target.shape[1], mask.shape[1]) pred = pred[:, :common] target = target[:, :common] mask = mask[:, :common].unsqueeze(-1) return (((pred - target) ** 2) * mask).sum() / (mask.sum() * pred.shape[-1]).clamp_min(1.0) def group_duration_targets(phone: torch.Tensor, durations: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: grouped = torch.zeros_like(durations, dtype=torch.float32) mask = torch.zeros_like(phone, dtype=torch.bool) for batch_index in range(phone.shape[0]): pending: list[torch.Tensor] = [] last_visible: int | None = None for token_index in range(phone.shape[1]): pending.append(durations[batch_index, token_index].float()) if int(phone[batch_index, token_index].item()) != 0: grouped[batch_index, token_index] = torch.log1p(torch.stack(pending).sum()) mask[batch_index, token_index] = True pending = [] last_visible = token_index if pending and last_visible is not None: value = torch.expm1(grouped[batch_index, last_visible]) + torch.stack(pending).sum() grouped[batch_index, last_visible] = torch.log1p(value) return grouped, mask def masked_wav_l1(pred: torch.Tensor, target: torch.Tensor, frame_mask: torch.Tensor, hop_size: int) -> torch.Tensor: if pred.dim() == 3: pred = pred.squeeze(1) common = min(pred.shape[-1], target.shape[-1], frame_mask.shape[-1] * hop_size) pred = pred[:, :common] target = target[:, :common] sample_mask = frame_mask.repeat_interleave(hop_size, dim=1)[:, :common].to(pred.dtype) return (torch.abs(pred - target) * sample_mask).sum() / sample_mask.sum().clamp_min(1.0) def load_frozen_vocoder(path: Path, device: torch.device) -> tuple[HifiGanGenerator, HifiGanConfig]: ckpt = torch.load(path, map_location=device, weights_only=False) cfg_payload = ckpt.get("config") or {"variant": "v2plus"} cfg = HifiGanConfig(**cfg_payload) if isinstance(cfg_payload, dict) else cfg_payload vocoder = HifiGanGenerator(cfg).to(device) vocoder.load_state_dict(ckpt["generator"]) vocoder.eval() for param in vocoder.parameters(): param.requires_grad_(False) return vocoder, cfg def load_model_state_flexible(model: nn.Module, state: dict[str, torch.Tensor]) -> tuple[int, int]: current = model.state_dict() compatible = {key: value for key, value in state.items() if key in current and current[key].shape == value.shape} model.load_state_dict(compatible, strict=False) return len(compatible), len(state) - len(compatible) def set_trainable_by_mode(model: MicroFastSpeech, mode: str) -> None: if mode == "all": for param in model.parameters(): param.requires_grad_(True) return for param in model.parameters(): param.requires_grad_(False) prefixes: tuple[str, ...] if mode == "duration": prefixes = ("phone.", "tone.", "lang.", "speaker.", "speaker_proj.", "encoder.", "duration_head.") elif mode == "predictors": prefixes = ( "phone.", "tone.", "lang.", "speaker.", "speaker_proj.", "encoder.", "duration_head.", "energy_head.", "bright_head.", "pitch_head.", ) elif mode == "heads": prefixes = ( "duration_head.", "energy_head.", "bright_head.", "pitch_head.", "predictor_context.", "duration_delta.", "energy_delta.", "bright_delta.", "pitch_delta.", ) elif mode == "contextual": prefixes = ( "predictor_context.", "duration_delta.", "energy_delta.", "bright_delta.", "pitch_delta.", ) elif mode == "group_duration": prefixes = ("group_duration_delta.",) elif mode == "decoder_adapt": prefixes = ( "energy_proj.", "bright_proj.", "pitch_proj.", "abs_frame.", "frame_proj.", "local_ctx.", "decoder.", "frame_gru.", "mel_head.", "postnet.", ) else: raise ValueError(f"Unknown trainable mode: {mode}") for name, param in model.named_parameters(): if name.startswith(prefixes): param.requires_grad_(True) def latest_checkpoint(out_dir: Path) -> Path | None: found = [] for path in out_dir.glob("inflect-micro-fastspeech-*.pt"): tail = path.stem.rsplit("-", 1)[-1] if tail.isdigit(): found.append((int(tail), path)) return max(found)[1] if found else None def save_checkpoint(path: Path, model: nn.Module, optim, cfg: MicroFastSpeechConfig, step: int, args, speakers: dict[str, int]) -> None: path.parent.mkdir(parents=True, exist_ok=True) tmp = path.with_suffix(path.suffix + ".tmp") torch.save( { "model": model.state_dict(), "optim": optim.state_dict(), "config": asdict(cfg), "step": step, "speakers": speakers, "args": {k: str(v) if isinstance(v, Path) else v for k, v in vars(args).items()}, "params": count_parameters(model), }, tmp, ) tmp.replace(path) class MelDiscriminator(nn.Module): """Training-only multi-conv mel discriminator (LSGAN). Discarded at inference, so the exported ONNX graph is byte-for-byte unchanged. Adversarial + feature-matching loss push the predicted mel onto the real-mel manifold, countering the L1/L2 regression-to-the-mean over-smoothing floor (Ren et al. ACL 2022; GANSpeech IS2021).""" def __init__(self, n_mels: int, ch: int = 64) -> None: super().__init__() from torch.nn.utils import weight_norm as wn self.convs = nn.ModuleList([ wn(nn.Conv1d(n_mels, ch, 5, 1, 2)), wn(nn.Conv1d(ch, ch, 5, 2, 2)), wn(nn.Conv1d(ch, ch * 2, 5, 2, 2)), wn(nn.Conv1d(ch * 2, ch * 2, 5, 2, 2)), ]) self.post = wn(nn.Conv1d(ch * 2, 1, 3, 1, 1)) def forward(self, mel: torch.Tensor): # mel [B, n_mels, T] fmaps = [] x = mel for c in self.convs: x = F.leaky_relu(c(x), 0.1) fmaps.append(x) return self.post(x), fmaps class MelDiscriminator2D(nn.Module): """RANK 5: training-only 2D time-frequency mel discriminator (spectral-norm). Treats the mel as an image [B,1,n_mels,T] so it judges JOINT time-frequency texture (formant structure), not just per-frame spectra like the Conv1d disc that failed in M11. Discarded at inference -> ONNX unchanged.""" def __init__(self, n_mels: int = 80) -> None: super().__init__() from torch.nn.utils import spectral_norm as sn chs = [1, 32, 64, 128, 256, 256] self.convs = nn.ModuleList([ sn(nn.Conv2d(chs[i], chs[i + 1], 5, 2, 2)) for i in range(5) ]) self.post = sn(nn.Conv2d(256, 1, 3, 1, 1)) def forward(self, mel: torch.Tensor): # mel [B, n_mels, T] -> image [B,1,n_mels,T] x = mel.unsqueeze(1) fmaps = [] for c in self.convs: x = F.leaky_relu(c(x), 0.2) fmaps.append(x) return self.post(x), fmaps def train(args: argparse.Namespace) -> None: device = torch.device(args.device) rows = load_rows(args.durations_jsonl, args.max_rows) speakers = {voice: idx for idx, voice in enumerate(sorted({str(r.get("voice_id") or "mark") for r in rows}))} max_phone_id = max(max(map(int, r["phone_ids"])) for r in rows) max_tone_id = max(max(map(int, r["tone_ids"])) for r in rows) max_lang_id = max(max(map(int, r["lang_ids"])) for r in rows) cfg = MicroFastSpeechConfig( vocab_size=max(256, max_phone_id + 1), tone_size=max(16, max_tone_id + 1), lang_size=max(4, max_lang_id + 1), speaker_count=max(2, len(speakers)), hidden=args.hidden, encoder_layers=args.encoder_layers, decoder_layers=args.decoder_layers, decoder_ff_mult=args.decoder_ff_mult, max_frames=args.max_frames, postnet_scale=args.postnet_scale, abs_frame_bins=args.abs_frame_bins, use_contextual_predictors=args.contextual_predictors, use_group_duration_planner=args.group_duration_planner, sample_rate=args.sample_rate, n_mels=make_config(args.vocoder_variant).num_mels, # auto-match acoustic mel count to vocoder variant ) for row in rows: row["speaker_id"] = speakers[str(row.get("voice_id") or "mark")] random.Random(args.seed).shuffle(rows) model = MicroFastSpeech(cfg).to(device) start_step = 0 if args.init_checkpoint and not args.resume: ckpt = torch.load(args.init_checkpoint, map_location=device, weights_only=False) copied, skipped = load_model_state_flexible(model, ckpt["model"]) print(f"Initialized model from {args.init_checkpoint} ({copied} tensors copied, {skipped} skipped)") set_trainable_by_mode(model, args.trainable) trainable_params = [param for param in model.parameters() if param.requires_grad] optim = torch.optim.AdamW(trainable_params, lr=args.lr, betas=(0.9, 0.98), weight_decay=args.weight_decay) mel_disc = None disc_optim = None if getattr(args, "mel_gan_weight", 0.0) > 0.0 or getattr(args, "mel_fm_weight", 0.0) > 0.0: if getattr(args, "gan_2d", False): mel_disc = MelDiscriminator2D(cfg.n_mels).to(device) else: mel_disc = MelDiscriminator(cfg.n_mels).to(device) disc_optim = torch.optim.AdamW(mel_disc.parameters(), lr=args.disc_lr, betas=(0.5, 0.9)) print(f"Mel-GAN ON ({'2D' if getattr(args,'gan_2d',False) else '1D'}): adv={args.mel_gan_weight} " f"fm={'auto' if getattr(args,'gan_fm_auto',False) else args.mel_fm_weight} warmup={args.gan_warmup_steps} " f"r1={getattr(args,'gan_r1_gamma',0.0)} crop={getattr(args,'gan_crop',0)} " f"disc_lr={args.disc_lr} disc_params={sum(p.numel() for p in mel_disc.parameters()):,}", flush=True) if args.resume: ckpt_path = latest_checkpoint(args.out_dir) if ckpt_path: ckpt = torch.load(ckpt_path, map_location=device, weights_only=False) model.load_state_dict(ckpt["model"]) optim.load_state_dict(ckpt["optim"]) start_step = int(ckpt.get("step") or 0) print(f"Resumed {ckpt_path} at step {start_step}") hifi_cfg = make_config(args.vocoder_variant) mel_frontend = MelFrontend(hifi_cfg).to(device) assert hifi_cfg.sample_rate == cfg.sample_rate, f"vocoder sr {hifi_cfg.sample_rate} != cfg sr {cfg.sample_rate}" prepared_rows = None if args.preload_features: print("Preloading audio/mel/pitch features...", flush=True) prepared_rows = [prepare_row_features(row, cfg, mel_frontend, device, args.max_seconds) for row in rows] total_frames = sum(int(row["frame_count"]) for row in prepared_rows) print(f"Preloaded {len(prepared_rows)} rows ({total_frames:,} frames)", flush=True) consistency_vocoder = None if args.vocoder_checkpoint: consistency_vocoder, consistency_cfg = load_frozen_vocoder(args.vocoder_checkpoint, device) if consistency_cfg.hop_size != hifi_cfg.hop_size: raise RuntimeError(f"Vocoder hop mismatch: {consistency_cfg.hop_size} != {hifi_cfg.hop_size}") print(f"Loaded frozen vocoder consistency checkpoint: {args.vocoder_checkpoint}") if (args.vocoder_wav_weight > 0.0 or args.vocoder_mel_weight > 0.0) and consistency_vocoder is None: raise RuntimeError("--vocoder-checkpoint is required when vocoder consistency losses are enabled") args.out_dir.mkdir(parents=True, exist_ok=True) (args.out_dir / "config.json").write_text( json.dumps({"config": asdict(cfg), "speakers": speakers, "rows": len(rows), "params": count_parameters(model)}, indent=2), encoding="utf-8", ) print(f"Rows: {len(rows)}") print(f"Speakers: {speakers}") print(f"Acoustic params: {count_parameters(model):,} ({count_parameters(model)/1_000_000:.3f}M)") print(f"Trainable params: {sum(p.numel() for p in model.parameters() if p.requires_grad):,} mode={args.trainable}") print(f"Total with V2+ vocoder: {(count_parameters(model)+1_426_842):,} ({(count_parameters(model)+1_426_842)/1_000_000:.3f}M)") rng = random.Random(args.seed + start_step) # language-balanced sampling pool: replicate English-row indices by --en-upsample en_up = max(1, int(round(args.en_upsample))) sample_pool = list(range(len(rows))) if en_up > 1: en_idx = [i for i, r in enumerate(rows) if str(r.get("id") or "").startswith("en")] sample_pool = sample_pool + en_idx * (en_up - 1) print(f"lang-balance: {len(en_idx)} en rows upsampled x{en_up} -> pool {len(sample_pool)} " f"(en exposure ~{100*len(en_idx)*en_up/len(sample_pool):.0f}%)") step = start_step started = time.time() while step < args.steps: source_rows = prepared_rows if prepared_rows is not None else rows batch = [source_rows[sample_pool[rng.randrange(len(sample_pool))]] for _ in range(args.batch_size)] if prepared_rows is not None: phone, tone, lang, speaker, durations, energy_t, bright_t, pitch_token_t, target_mel, frame_mask, pitch_frame, target_wav = collate_prepared( batch, device, hifi_cfg.hop_size ) else: phone, tone, lang, speaker, durations, energy_t, bright_t, pitch_token_t, target_mel, frame_mask, pitch_frame, target_wav = collate( batch, cfg, mel_frontend, device, args.max_seconds, hifi_cfg.hop_size ) out = model(phone, tone, lang, speaker, durations, energy_t, bright_t, pitch_frame) token_mask = out["token_mask"] log_dur_t = torch.log1p(durations.float()) group_log_dur_t, group_mask = group_duration_targets(phone, durations) mel_l1 = masked_l1(out["mel"], target_mel, frame_mask) mel_mse = masked_mse(out["mel"], target_mel, frame_mask) # RANK 4: per-language mel-loss weighting. Down-weight English frames so they stop # crowding out Chinese in the 4.6M budget. At 1:1 reproduces M7 exactly (gated off). _lw_zh = getattr(args, "lang_loss_zh_weight", 1.0) _lw_en = getattr(args, "lang_loss_en_weight", 1.0) if _lw_zh != 1.0 or _lw_en != 1.0: T = out["mel"].shape[-1] lang_frame = model.expand_token_feature(lang.unsqueeze(-1).float(), durations).squeeze(-1) # [B,Tf] common = min(T, lang_frame.shape[-1], frame_mask.shape[-1]) fm = frame_mask[:, :common].to(out["mel"].dtype) lf = lang_frame[:, :common] wmap = fm * torch.where(lf < 0.5, float(_lw_zh), float(_lw_en)) # lang 0 = zh vc = frame_mask[:, :common].sum() mel_l1 = weighted_frame_l1(out["mel"], target_mel, wmap, vc) mel_mse = weighted_frame_mse(out["mel"], target_mel, wmap, vc) delta = masked_delta_loss(out["mel"], target_mel, frame_mask) accel = masked_accel_loss(out["mel"], target_mel, frame_mask) dur_loss = token_mse(out["log_dur"], log_dur_t, token_mask) group_dur_loss = token_mse(out["group_log_dur"], group_log_dur_t, group_mask) energy_loss = token_mse(out["energy"], energy_t, token_mask) bright_loss = token_mse(out["bright"], bright_t, token_mask) pitch_loss = token_mse_nd(out["pitch"], pitch_token_t, token_mask) predicted_prosody_mel_loss = torch.zeros((), device=device) predicted_prosody_delta_loss = torch.zeros((), device=device) if args.predicted_prosody_mel_weight > 0.0 or args.predicted_prosody_delta_weight > 0.0: # Train the predictor heads against the acoustic result they produce at # inference, while retaining reference durations so this path remains # differentiable and isolates prosody exposure bias. predicted_conditioning = model(phone, tone, lang, speaker, durations) if args.predicted_prosody_mel_weight > 0.0: predicted_prosody_mel_loss = masked_l1(predicted_conditioning["mel"], target_mel, frame_mask) if args.predicted_prosody_delta_weight > 0.0: predicted_prosody_delta_loss = masked_delta_loss(predicted_conditioning["mel"], target_mel, frame_mask) robust_prosody_mel_loss = torch.zeros((), device=device) robust_prosody_delta_loss = torch.zeros((), device=device) if args.robust_prosody_mel_weight > 0.0 or args.robust_prosody_delta_weight > 0.0: robust_conditioning = model( phone, tone, lang, speaker, durations, energy_t, bright_t, pitch_frame, predicted_prosody_mix=args.robust_prosody_mix, detach_mixed_predictions=True, ) if args.robust_prosody_mel_weight > 0.0: robust_prosody_mel_loss = masked_l1(robust_conditioning["mel"], target_mel, frame_mask) if args.robust_prosody_delta_weight > 0.0: robust_prosody_delta_loss = masked_delta_loss(robust_conditioning["mel"], target_mel, frame_mask) voc_wav_loss = torch.zeros((), device=device) voc_mel_loss = torch.zeros((), device=device) voc_mrstft_loss = torch.zeros((), device=device) _mrstft_w = getattr(args, "vocoder_mrstft_weight", 0.0) if consistency_vocoder is not None and (args.vocoder_wav_weight > 0.0 or args.vocoder_mel_weight > 0.0 or _mrstft_w > 0.0): pred_wav = consistency_vocoder(out["mel"].clamp(-12.0, 2.0)) if args.vocoder_wav_weight > 0.0: voc_wav_loss = masked_wav_l1(pred_wav, target_wav, frame_mask, hifi_cfg.hop_size) if args.vocoder_mel_weight > 0.0: pred_recon_mel = mel_frontend(pred_wav.squeeze(1)) voc_mel_loss = masked_l1(pred_recon_mel, target_mel, frame_mask) if _mrstft_w > 0.0: # RANK 3: multi-resolution STFT magnitude loss through the frozen vocoder. # 8kHz-appropriate FFT sizes (<=512; >512 over-resolves a 4kHz-Nyquist signal). from .vocoder import stft_mag_loss _hop = hifi_cfg.hop_size _common = min(pred_wav.shape[-1], target_wav.shape[-1], frame_mask.shape[-1] * _hop) _sm = frame_mask.repeat_interleave(_hop, dim=1)[:, :_common].to(pred_wav.dtype) _pw = pred_wav.squeeze(1)[:, :_common] * _sm _tw = (target_wav.squeeze(1) if target_wav.dim() == 3 else target_wav)[:, :_common] * _sm voc_mrstft_loss = stft_mag_loss(_pw, _tw, (128, 256, 512), (32, 64, 128), (128, 256, 512)) # ramp the MR-STFT weight in from --mrstft-warmup-steps over 4000 steps (limit early vocoder-quirk exploitation) mrstft_eff = _mrstft_w * min(1.0, max(0.0, (step - args.mrstft_warmup_steps) / 4000.0)) if _mrstft_w > 0.0 else 0.0 loss = ( mel_l1 + args.mse_weight * mel_mse + args.delta_weight * delta + args.accel_weight * accel + args.duration_weight * dur_loss + args.group_duration_weight * group_dur_loss + args.energy_weight * energy_loss + args.bright_weight * bright_loss + args.pitch_weight * pitch_loss + args.predicted_prosody_mel_weight * predicted_prosody_mel_loss + args.predicted_prosody_delta_weight * predicted_prosody_delta_loss + args.robust_prosody_mel_weight * robust_prosody_mel_loss + args.robust_prosody_delta_weight * robust_prosody_delta_loss + args.vocoder_wav_weight * voc_wav_loss + args.vocoder_mel_weight * voc_mel_loss + mrstft_eff * voc_mrstft_loss ) gan_g = torch.zeros((), device=device) gan_fm = torch.zeros((), device=device) gan_d = torch.zeros((), device=device) if mel_disc is not None and step >= args.gan_warmup_steps: T = out["mel"].shape[-1] m = frame_mask[:, :T].unsqueeze(1).to(out["mel"].dtype) real = target_mel[..., :T] * m fake = out["mel"] * m # RANK 5: optional random time-crop (2D disc judges local TF texture; stabilizes + speeds). _crop = getattr(args, "gan_crop", 0) if _crop > 0 and T > _crop: _s = int(torch.randint(0, T - _crop + 1, (1,)).item()) real = real[..., _s:_s + _crop]; fake = fake[..., _s:_s + _crop] # Discriminator step (LSGAN) with optional lazy R1 on real mels every 16 steps. _r1 = getattr(args, "gan_r1_gamma", 0.0) do_r1 = _r1 > 0.0 and (step % 16 == 0) if do_r1: real = real.detach().requires_grad_(True) d_real, _ = mel_disc(real) d_fake_d, _ = mel_disc(fake.detach()) gan_d = 0.5 * ((d_real - 1.0) ** 2).mean() + 0.5 * (d_fake_d ** 2).mean() if do_r1: gp = torch.autograd.grad(d_real.sum(), real, create_graph=True)[0] gan_d = gan_d + (_r1 / 2.0) * gp.pow(2).flatten(1).sum(1).mean() disc_optim.zero_grad(set_to_none=True) gan_d.backward() torch.nn.utils.clip_grad_norm_(mel_disc.parameters(), args.grad_clip) disc_optim.step() # Generator adversarial + feature-matching (real features detached). d_fake_g, feats_fake = mel_disc(fake) _, feats_real = mel_disc(real.detach()) gan_g = ((d_fake_g - 1.0) ** 2).mean() gan_fm = sum(F.l1_loss(ff, fr.detach()) for ff, fr in zip(feats_fake, feats_real)) / len(feats_fake) # auto-FM scaling: lambda_FM = (recon / FM).detach().clamp[0,50] (GANSpeech-style; FM does the work) if getattr(args, "gan_fm_auto", False): fm_w = (mel_l1.detach() / (gan_fm.detach() + 1e-8)).clamp(0.0, 50.0) else: fm_w = args.mel_fm_weight loss = loss + args.mel_gan_weight * gan_g + fm_w * gan_fm optim.zero_grad(set_to_none=True) loss.backward() grad = torch.nn.utils.clip_grad_norm_(trainable_params, args.grad_clip) optim.step() step += 1 if step == 1 or step % args.log_interval == 0: elapsed = max(1e-6, time.time() - started) speed = (step - start_step) / elapsed eta = (args.steps - step) / max(1e-6, speed) print( f"step={step}/{args.steps} loss={loss.item():.4f} mel={mel_l1.item():.4f} " f"mse={mel_mse.item():.4f} delta={delta.item():.4f} accel={accel.item():.4f} dur={dur_loss.item():.4f} " f"gdur={group_dur_loss.item():.4f} " f"energy={energy_loss.item():.4f} bright={bright_loss.item():.4f} " f"pitch={pitch_loss.item():.4f} pmel={predicted_prosody_mel_loss.item():.4f} " f"pdelta={predicted_prosody_delta_loss.item():.4f} rmel={robust_prosody_mel_loss.item():.4f} " f"rdelta={robust_prosody_delta_loss.item():.4f} vwav={voc_wav_loss.item():.4f} " f"vmel={voc_mel_loss.item():.4f} mrstft={voc_mrstft_loss.item():.4f} ganG={gan_g.item():.4f} ganFM={gan_fm.item():.4f} ganD={gan_d.item():.4f} grad={float(grad):.2f} " f"speed={speed:.3f} step/s eta={eta/60:.1f}m", flush=True, ) if step % args.save_interval == 0 or step >= args.steps: save_checkpoint(args.out_dir / f"inflect-micro-fastspeech-{step}.pt", model, optim, cfg, step, args, speakers) save_checkpoint(args.out_dir / "inflect-micro-fastspeech-latest.pt", model, optim, cfg, step, args, speakers) print(f"Done. {args.out_dir}") def main() -> None: ap = argparse.ArgumentParser(description="Train Inflect Micro duration-conditioned acoustic model.") ap.add_argument("--durations-jsonl", type=Path, required=True) ap.add_argument("--out-dir", type=Path, required=True) ap.add_argument("--vocoder-variant", type=str, default="v2plus") ap.add_argument("--sample-rate", type=int, default=24000) ap.add_argument("--en-upsample", type=float, default=1.0, help="Oversample English rows (id starts 'en') by this factor in the " "training sampler, to balance a zh-dominant bilingual corpus.") ap.add_argument("--max-rows", type=int, default=0) ap.add_argument("--steps", type=int, default=20000) ap.add_argument("--batch-size", type=int, default=6) ap.add_argument("--lr", type=float, default=2.0e-4) ap.add_argument("--weight-decay", type=float, default=1.0e-4) ap.add_argument("--hidden", type=int, default=168) ap.add_argument("--encoder-layers", type=int, default=5) ap.add_argument("--decoder-layers", type=int, default=6) ap.add_argument("--decoder-ff-mult", type=int, default=3) ap.add_argument("--max-seconds", type=float, default=12.0) ap.add_argument("--max-frames", type=int, default=1400) ap.add_argument("--mse-weight", type=float, default=0.25) ap.add_argument("--delta-weight", type=float, default=0.18) # Training-only mel GAN (anti over-smoothing). Discriminator discarded at inference -> ONNX unchanged. ap.add_argument("--mel-gan-weight", type=float, default=0.0, help="generator adversarial loss weight (0=off)") ap.add_argument("--mel-fm-weight", type=float, default=0.0, help="feature-matching loss weight") ap.add_argument("--disc-lr", type=float, default=2.0e-4, help="mel discriminator learning rate") ap.add_argument("--gan-warmup-steps", type=int, default=2000, help="steps of pure recon before GAN kicks in") # RANK 5: corrected GAN — 2D TF discriminator + auto-FM + R1 + crop ap.add_argument("--gan-2d", action="store_true", help="use 2D time-frequency mel discriminator (spectral-norm)") ap.add_argument("--gan-fm-auto", action="store_true", help="auto-scale feature-matching weight = (recon/FM).clamp[0,50]") ap.add_argument("--gan-r1-gamma", type=float, default=0.0, help="lazy R1 gradient-penalty gamma (every 16 steps)") ap.add_argument("--gan-crop", type=int, default=0, help="random time-crop width for the disc (0=off)") # RANK 3: multi-resolution STFT loss through the frozen consistency vocoder (anti over-smoothing, loss-only) ap.add_argument("--vocoder-mrstft-weight", type=float, default=0.0, help="MR-STFT-through-vocoder loss weight (0=off)") ap.add_argument("--mrstft-warmup-steps", type=int, default=4000, help="step at which MR-STFT ramp begins") # RANK 4: per-language mel-loss weighting (anti capacity-interference). 1:1 = M7 (gated off). ap.add_argument("--lang-loss-zh-weight", type=float, default=1.0, help="mel-loss weight on zh frames") ap.add_argument("--lang-loss-en-weight", type=float, default=1.0, help="mel-loss weight on en frames") ap.add_argument("--accel-weight", type=float, default=0.0) ap.add_argument("--duration-weight", type=float, default=0.08) ap.add_argument("--group-duration-weight", type=float, default=0.0) ap.add_argument("--energy-weight", type=float, default=0.04) ap.add_argument("--bright-weight", type=float, default=0.04) ap.add_argument("--pitch-weight", type=float, default=0.04) ap.add_argument("--predicted-prosody-mel-weight", type=float, default=0.0) ap.add_argument("--predicted-prosody-delta-weight", type=float, default=0.0) ap.add_argument("--robust-prosody-mix", type=float, default=0.0) ap.add_argument("--robust-prosody-mel-weight", type=float, default=0.0) ap.add_argument("--robust-prosody-delta-weight", type=float, default=0.0) ap.add_argument("--grad-clip", type=float, default=5.0) ap.add_argument("--postnet-scale", type=float, default=0.10) ap.add_argument("--abs-frame-bins", type=int, default=512) ap.add_argument("--init-checkpoint", type=Path) ap.add_argument("--vocoder-checkpoint", type=Path) ap.add_argument("--vocoder-wav-weight", type=float, default=0.0) ap.add_argument("--vocoder-mel-weight", type=float, default=0.0) ap.add_argument("--save-interval", type=int, default=2000) ap.add_argument("--log-interval", type=int, default=50) ap.add_argument("--seed", type=int, default=42) ap.add_argument("--resume", action="store_true") ap.add_argument("--preload-features", action="store_true", help="Cache decoded audio, mels, pitch, and token features in RAM before training.") ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") ap.add_argument( "--trainable", choices=["all", "duration", "predictors", "heads", "contextual", "group_duration", "decoder_adapt"], default="all", ) ap.add_argument("--contextual-predictors", action="store_true") ap.add_argument("--group-duration-planner", action="store_true") args = ap.parse_args() train(args) if __name__ == "__main__": main()