PrimeTTS: full training pipeline + weights (fine-tune of Inflect-Nano-v1)

README.md

@@ -0,0 +1,117 @@
+---
+license: apache-2.0
+language:
+  - zh
+  - en
+tags:
+  - text-to-speech
+  - tts
+  - onnx
+  - on-device
+  - jetson
+  - telephony
+  - mandarin
+  - taiwanese-mandarin
+base_model: owensong/Inflect-Nano-v1
+library_name: onnxruntime
+pipeline_tag: text-to-speech
+---
+
+# PrimeTTS — tiny bilingual zh-TW + English TTS (8 kHz, CPU)
+
+**PrimeTTS** is an ultra-small **4.63M-parameter** Mandarin (Taiwan) + English text-to-speech model
+that runs **entirely on CPU** via `onnxruntime` and emits **8 kHz** audio — sized for **G.711
+telephony** and **on-device (Jetson-class)** deployment.
+
+It is a **fine-tune of [`owensong/Inflect-Nano-v1`](https://huggingface.co/owensong/Inflect-Nano-v1)**
+(Apache-2.0) — the **same, frozen architecture**, retrained for zh-TW + English. No architecture
+changes, no neural-architecture-search.
+
+- **Acoustic** (`MicroFastSpeech`, ~3.47M): FastSpeech-style depthwise Conv-FFN (**no attention**),
+  external durations + length regulator, frame-pitch, BiGRU, postnet.
+- **Vocoder** (~1.17M): Snake-HiFiGAN. We added an **8 kHz variant** (`snake_8k`: sr 8000, n_fft 512,
+  hop 128, 80 mels) and a **zh-TW + English frontend** (bopomofo + arpabet, one unified phone sequence).
+
+**Live demo:** https://huggingface.co/spaces/Luigi/inflect-nano-zhtw-en-8k-demo
+
+---
+
+## How these weights were obtained — the two things that matter
+
+Inflect-Nano-v1's 4.63M architecture is **not capacity-limited** for this task. We confirmed the
+original English checkpoint scores ~0.05 CER on our eval, yet our first retrains were unintelligible
+(~0.88 CER). The gap was **not** the model size. Two fixes — **architecture frozen** — took held-out
+Mandarin CER from **~0.88 → ~0.06**:
+
+1. **Phone-level forced alignment.** FastSpeech needs a per-phone duration for every training clip.
+   Crude char/letter-CTC alignment (then splitting a char's span across its phones by heuristic) gives
+   *wrong relative phone durations* → the acoustic learns a mis-timed phone→mel map → over-smoothed,
+   garbled output in **every** language. Replace it with **true phone-level forced alignment**
+   (`align_durations_v4.py`: the espeak phoneme-CTC model `facebook/wav2vec2-lv-60-espeak-cv-ft` +
+   `torchaudio.forced_align`, aligning *your frontend's own phone sequence* mapped to IPA). This alone:
+   Mandarin 0.88 → 0.40.
+
+2. **Diverse, well-covered training text.** The model can only pronounce characters/words it has seen.
+   A narrow corpus (≈234 Han chars) left ~39% of held-out characters unseen. Expanding character/word
+   coverage (`select_diverse_text.py`: Tatoeba → OpenCC `s2twp` → greedy coverage) took Mandarin
+   0.40 → **0.06**. Applied symmetrically to English, the same recipe yields a genuinely **bilingual**
+   model (zh-only ≈ 0.13, English ≈ 0.16) in one 4.63M net — no language routing.
+
+Everything else (loss weights, steps, etc.) is essentially Inflect-Nano-v1's defaults.
+
+---
+
+## Reproduce / fine-tune (your own voice or language)
+
+```
+scripts/        full pipeline (frontend, aligner, corpus-gen, train, export, eval)
+inflect_nano/   the trainer (acoustic.py + vocoder.py), forked from Inflect-Nano-v1 (LICENSE included)
+weights/        deployable ONNX (encoder/decoder/vocoder) + meta.json + symbol_table.json
+                + acoustic_zh_v2_60k.pt (PyTorch checkpoint, to resume / fine-tune)
+```
+
+**Step 1 — Teacher corpus.** Synthesize (text → audio) clips in ONE target voice, ASR-gated for
+fidelity. We used **BreezyVoice** (MediaTek, CosyVoice-based; single "mark" voice, does zh + en) with a
+**Breeze-ASR-25** gate (`gen_breezy_corpus.py` — keeps a clip only if its ASR transcript matches the
+intended text, t2s-normalized). Any clean single-speaker TTS/recordings work.
+
+**Step 2 — Diverse text.** `python scripts/select_diverse_text.py --lang zh --n 6000 --out zh.tsv`
+(and `--lang en`). Feed the `.tsv` to step 1. Coverage is the single biggest driver of held-out quality.
+
+**Step 3 — Phone-level alignment.** `python scripts/align_durations_v4.py --manifest corpus/manifest.jsonl
+--out align.jsonl` → per-phone durations from real audio. **This is the key fix — do not skip it.**
+
+**Step 4 — Train acoustic.** `inflect_nano.acoustic --durations-jsonl align.jsonl --vocoder-variant
+snake_8k --sample-rate 8000 ...` (see `scripts/run_bilingual.sh`). Mix languages in one corpus, single
+speaker, for native code-mix. ~60k steps, batch 16.
+
+**Step 5 — Train vocoder.** `inflect_nano.vocoder --train-jsonl voc_rows.jsonl --variant snake_8k ...`
+(see `scripts/run_voc_retrain.sh`). Train on the same diverse audio; higher `--stft-weight` = crisper.
+
+**Step 6 — Export ONNX.** `python scripts/export_8k.py --acoustic-ckpt <pt> --vocoder-ckpt <pt>
+--out-dir onnx/` → `acoustic_encoder.onnx` → numpy length-regulator (`host_regulate`) →
+`acoustic_decoder.onnx` → `vocoder.onnx`. Fully torch-free at inference.
+
+**Step 7 — Evaluate.** `synth_from_text.py` + `assess_big.py` (offline X-ASR CER/WER). Use ≥30 held-out
+sentences — small eval sets are too noisy.
+
+### Inference (torch-free, CPU)
+```python
+import frontend_bopomofo as F, numpy as np, onnxruntime as ort
+o = F.text_to_ids("您好,歡迎使用 PrimeTTS。")     # bopomofo + arpabet -> ids
+# encoder.onnx -> host_regulate (numpy) -> decoder.onnx -> vocoder.onnx -> 8 kHz wav
+```
+See `scripts/synth_from_text.py` for the full ~40-line runtime (also runs as-is on a Jetson Nano CPU).
+
+---
+
+## Credits & licenses
+
+- **Base model / trainer:** [`owensong/Inflect-Nano-v1`](https://huggingface.co/owensong/Inflect-Nano-v1)
+  (Apache-2.0; `inflect_nano/LICENSE.inflect-nano`).
+- **Teacher:** BreezyVoice (MediaTek Research). **Gate ASR:** `Breeze-ASR-25` (MediaTek Research).
+- **Aligner:** `facebook/wav2vec2-lv-60-espeak-cv-ft` + `torchaudio.forced_align`.
+- **Frontend:** `g2pw` (Taiwan bopomofo, polyphone disambiguation) + `g2p_en` (arpabet).
+- **Eval ASR:** sherpa-onnx X-ASR (zh-en Zipformer). **Text:** Tatoeba (CC-BY 2.0 FR).
+
+This repository: **Apache-2.0**.

inflect_nano/LICENSE.inflect-nano

@@ -0,0 +1,201 @@
+Apache License
+Version 2.0, January 2004
+http://www.apache.org/licenses/
+
+TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+1. Definitions.
+
+"License" shall mean the terms and conditions for use, reproduction,
+and distribution as defined by Sections 1 through 9 of this document.
+
+"Licensor" shall mean the copyright owner or entity authorized by
+the copyright owner that is granting the License.
+
+"Legal Entity" shall mean the union of the acting entity and all
+other entities that control, are controlled by, or are under common
+control with that entity. For the purposes of this definition,
+"control" means (i) the power, direct or indirect, to cause the
+direction or management of such entity, whether by contract or
+otherwise, or (ii) ownership of fifty percent (50%) or more of the
+outstanding shares, or (iii) beneficial ownership of such entity.
+
+"You" (or "Your") shall mean an individual or Legal Entity
+exercising permissions granted by this License.
+
+"Source" form shall mean the preferred form for making modifications,
+including but not limited to software source code, documentation
+source, and configuration files.
+
+"Object" form shall mean any form resulting from mechanical
+transformation or translation of a Source form, including but
+not limited to compiled object code, generated documentation,
+and conversions to other media types.
+
+"Work" shall mean the work of authorship, whether in Source or
+Object form, made available under the License, as indicated by a
+copyright notice that is included in or attached to the work
+(an example is provided in the Appendix below).
+
+"Derivative Works" shall mean any work, whether in Source or Object
+form, that is based on (or derived from) the Work and for which the
+editorial revisions, annotations, elaborations, or other modifications
+represent, as a whole, an original work of authorship. For the purposes
+of this License, Derivative Works shall not include works that remain
+separable from, or merely link (or bind by name) to the interfaces of,
+the Work and Derivative Works thereof.
+
+"Contribution" shall mean any work of authorship, including
+the original version of the Work and any modifications or additions
+to that Work or Derivative Works thereof, that is intentionally
+submitted to Licensor for inclusion in the Work by the copyright owner
+or by an individual or Legal Entity authorized to submit on behalf of
+the copyright owner. For the purposes of this definition, "submitted"
+means any form of electronic, verbal, or written communication sent
+to the Licensor or its representatives, including but not limited to
+communication on electronic mailing lists, source code control systems,
+and issue tracking systems that are managed by, or on behalf of, the
+Licensor for the purpose of discussing and improving the Work, but
+excluding communication that is conspicuously marked or otherwise
+designated in writing by the copyright owner as "Not a Contribution."
+
+"Contributor" shall mean Licensor and any individual or Legal Entity
+on behalf of whom a Contribution has been received by Licensor and
+subsequently incorporated within the Work.
+
+2. Grant of Copyright License. Subject to the terms and conditions of
+this License, each Contributor hereby grants to You a perpetual,
+worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+copyright license to reproduce, prepare Derivative Works of,
+publicly display, publicly perform, sublicense, and distribute the
+Work and such Derivative Works in Source or Object form.
+
+3. Grant of Patent License. Subject to the terms and conditions of
+this License, each Contributor hereby grants to You a perpetual,
+worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+(except as stated in this section) patent license to make, have made,
+use, offer to sell, sell, import, and otherwise transfer the Work,
+where such license applies only to those patent claims licensable
+by such Contributor that are necessarily infringed by their
+Contribution(s) alone or by combination of their Contribution(s)
+with the Work to which such Contribution(s) was submitted. If You
+institute patent litigation against any entity (including a
+cross-claim or counterclaim in a lawsuit) alleging that the Work
+or a Contribution incorporated within the Work constitutes direct
+or contributory patent infringement, then any patent licenses
+granted to You under this License for that Work shall terminate
+as of the date such litigation is filed.
+
+4. Redistribution. You may reproduce and distribute copies of the
+Work or Derivative Works thereof in any medium, with or without
+modifications, and in Source or Object form, provided that You
+meet the following conditions:
+
+(a) You must give any other recipients of the Work or
+Derivative Works a copy of this License; and
+
+(b) You must cause any modified files to carry prominent notices
+stating that You changed the files; and
+
+(c) You must retain, in the Source form of any Derivative Works
+that You distribute, all copyright, patent, trademark, and
+attribution notices from the Source form of the Work,
+excluding those notices that do not pertain to any part of
+the Derivative Works; and
+
+(d) If the Work includes a "NOTICE" text file as part of its
+distribution, then any Derivative Works that You distribute must
+include a readable copy of the attribution notices contained
+within such NOTICE file, excluding those notices that do not
+pertain to any part of the Derivative Works, in at least one
+of the following places: within a NOTICE text file distributed
+as part of the Derivative Works; within the Source form or
+documentation, if provided along with the Derivative Works; or,
+within a display generated by the Derivative Works, if and
+wherever such third-party notices normally appear. The contents
+of the NOTICE file are for informational purposes only and
+do not modify the License. You may add Your own attribution
+notices within Derivative Works that You distribute, alongside
+or as an addendum to the NOTICE text from the Work, provided
+that such additional attribution notices cannot be construed
+as modifying the License.
+
+You may add Your own copyright statement to Your modifications and
+may provide additional or different license terms and conditions
+for use, reproduction, or distribution of Your modifications, or
+for any such Derivative Works as a whole, provided Your use,
+reproduction, and distribution of the Work otherwise complies with
+the conditions stated in this License.
+
+5. Submission of Contributions. Unless You explicitly state otherwise,
+any Contribution intentionally submitted for inclusion in the Work
+by You to the Licensor shall be under the terms and conditions of
+this License, without any additional terms or conditions.
+Notwithstanding the above, nothing herein shall supersede or modify
+the terms of any separate license agreement you may have executed
+with Licensor regarding such Contributions.
+
+6. Trademarks. This License does not grant permission to use the trade
+names, trademarks, service marks, or product names of the Licensor,
+except as required for reasonable and customary use in describing the
+origin of the Work and reproducing the content of the NOTICE file.
+
+7. Disclaimer of Warranty. Unless required by applicable law or
+agreed to in writing, Licensor provides the Work (and each
+Contributor provides its Contributions) on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+implied, including, without limitation, any warranties or conditions
+of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+PARTICULAR PURPOSE. You are solely responsible for determining the
+appropriateness of using or redistributing the Work and assume any
+risks associated with Your exercise of permissions under this License.
+
+8. Limitation of Liability. In no event and under no legal theory,
+whether in tort (including negligence), contract, or otherwise,
+unless required by applicable law (such as deliberate and grossly
+negligent acts) or agreed to in writing, shall any Contributor be
+liable to You for damages, including any direct, indirect, special,
+incidental, or consequential damages of any character arising as a
+result of this License or out of the use or inability to use the
+Work (including but not limited to damages for loss of goodwill,
+work stoppage, computer failure or malfunction, or any and all
+other commercial damages or losses), even if such Contributor
+has been advised of the possibility of such damages.
+
+9. Accepting Warranty or Additional Liability. While redistributing
+the Work or Derivative Works thereof, You may choose to offer,
+and charge a fee for, acceptance of support, warranty, indemnity,
+or other liability obligations and/or rights consistent with this
+License. However, in accepting such obligations, You may act only
+on Your own behalf and on Your sole responsibility, not on behalf
+of any other Contributor, and only if You agree to indemnify,
+defend, and hold each Contributor harmless for any liability
+incurred by, or claims asserted against, such Contributor by reason
+of your accepting any such warranty or additional liability.
+
+END OF TERMS AND CONDITIONS
+
+APPENDIX: How to apply the Apache License to your work.
+
+To apply the Apache License to your work, attach the following
+boilerplate notice, with the fields enclosed by brackets "{}"
+replaced with your own identifying information. (Don't include
+the brackets.) The text should be enclosed in the appropriate
+comment syntax for the file format. We also recommend that a
+file or class name and description of purpose be included on the
+same "printed page" as the copyright notice for easier
+identification within third-party archives.
+
+Copyright {yyyy} {name of copyright owner}
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.

inflect_nano/__init__.py

@@ -0,0 +1 @@
+"""Inflect-Nano-v1 runtime package."""

inflect_nano/acoustic.py

@@ -0,0 +1,1295 @@
+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()

inflect_nano/text_cleaning.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+import re
+
+
+_QUOTE_TRANSLATION = str.maketrans(
+    {
+        "\u2018": "'",
+        "\u2019": "'",
+        "\u201c": "",
+        "\u201d": "",
+        "\u2014": ",",
+        "\u2013": ",",
+        ";": ",",
+        ":": ",",
+        "\n": ".",
+    }
+)
+
+
+def clean_tinytts_text(text: str) -> str:
+    """Normalize text into punctuation TinyTTS actually has symbols for."""
+    text = str(text).translate(_QUOTE_TRANSLATION)
+    text = text.replace("...", "…")
+    text = re.sub(r"\s+", " ", text).strip()
+    text = re.sub(r"\s+([,.!?…])", r"\1", text)
+    text = re.sub(r"([,.!?…]){2,}", r"\1", text)
+    return text

inflect_nano/vocoder.py

@@ -0,0 +1,844 @@
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import random
+import time
+from dataclasses import asdict, dataclass
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchaudio
+from torch.nn.utils import remove_weight_norm, spectral_norm, weight_norm
+from torch.utils.data import DataLoader, Dataset
+
+
+@dataclass(frozen=True)
+class HifiGanConfig:
+    variant: str
+    sample_rate: int = 24000
+    n_fft: int = 1024
+    hop_size: int = 256
+    win_size: int = 1024
+    num_mels: int = 80
+    fmin: float = 0.0
+    fmax: float = 12000.0
+    resblock: str = "1"
+    upsample_rates: tuple[int, ...] = (8, 8, 2, 2)
+    upsample_kernel_sizes: tuple[int, ...] = (16, 16, 4, 4)
+    upsample_initial_channel: int = 128
+    resblock_kernel_sizes: tuple[int, ...] = (3, 7, 11)
+    resblock_dilation_sizes: tuple[tuple[int, ...], ...] = ((1, 3, 5), (1, 3, 5), (1, 3, 5))
+    activation: str = "lrelu"
+    conditioning_channels: int = 0
+
+
+def make_config(variant: str) -> HifiGanConfig:
+    if variant == "v2":
+        return HifiGanConfig(variant="v2")
+    if variant == "v2plus":
+        return HifiGanConfig(variant="v2plus", upsample_initial_channel=160)
+    if variant == "v2wide":
+        return HifiGanConfig(variant="v2wide", upsample_initial_channel=176)
+    if variant == "snake_v2mid":
+        return HifiGanConfig(variant="snake_v2mid", upsample_initial_channel=144, activation="snake")
+    if variant == "snake_8k":
+        return HifiGanConfig(variant="snake_8k", sample_rate=8000, n_fft=512, hop_size=128,
+                             win_size=512, num_mels=80, fmin=0.0, fmax=4000.0,
+                             upsample_rates=(8, 4, 2, 2), upsample_kernel_sizes=(16, 8, 4, 4),
+                             upsample_initial_channel=144, activation="snake")
+    if variant == "snake_8k40":
+        # 8kHz variant with 40 mels (vs 80): 0-4kHz over-resolved at 80 -> correlated channels
+        # encourage L1 mean-collapse; 40 mels is the natural count for a 4kHz band (anti over-smoothing).
+        return HifiGanConfig(variant="snake_8k40", sample_rate=8000, n_fft=512, hop_size=128,
+                             win_size=512, num_mels=40, fmin=0.0, fmax=4000.0,
+                             upsample_rates=(8, 4, 2, 2), upsample_kernel_sizes=(16, 8, 4, 4),
+                             upsample_initial_channel=144, activation="snake")
+    if variant == "snake_v2balanced":
+        return HifiGanConfig(variant="snake_v2balanced", upsample_initial_channel=160, activation="snake")
+    if variant == "source_snake_v2balanced":
+        return HifiGanConfig(
+            variant="source_snake_v2balanced",
+            upsample_initial_channel=160,
+            activation="snake",
+            conditioning_channels=5,
+        )
+    if variant == "v3":
+        return HifiGanConfig(
+            variant="v3",
+            resblock="2",
+            upsample_rates=(8, 8, 4),
+            upsample_kernel_sizes=(16, 16, 8),
+            upsample_initial_channel=256,
+            resblock_kernel_sizes=(3, 5, 7),
+            resblock_dilation_sizes=((1, 2), (2, 6), (3, 12)),
+        )
+    raise ValueError(f"Unknown variant: {variant}")
+
+
+def get_padding(kernel_size: int, dilation: int = 1) -> int:
+    return int((kernel_size * dilation - dilation) / 2)
+
+
+class SnakeActivation(nn.Module):
+    def __init__(self, channels: int):
+        super().__init__()
+        self.log_alpha = nn.Parameter(torch.zeros(1, channels, 1))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        alpha = self.log_alpha.exp().clamp(1e-4, 100.0)
+        return x + torch.sin(alpha * x).pow(2) / alpha
+
+
+def make_activation(channels: int, activation: str) -> nn.Module:
+    if activation == "snake":
+        return SnakeActivation(channels)
+    return nn.LeakyReLU(0.1)
+
+
+class ResBlock1(nn.Module):
+    def __init__(self, channels: int, kernel_size: int, dilations: tuple[int, ...], activation: str = "lrelu"):
+        super().__init__()
+        self.convs1 = nn.ModuleList(
+            [
+                weight_norm(
+                    nn.Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=d,
+                        padding=get_padding(kernel_size, d),
+                    )
+                )
+                for d in dilations
+            ]
+        )
+        self.convs2 = nn.ModuleList(
+            [
+                weight_norm(
+                    nn.Conv1d(channels, channels, kernel_size, 1, dilation=1, padding=get_padding(kernel_size, 1))
+                )
+                for _ in dilations
+            ]
+        )
+        self.acts1 = nn.ModuleList([make_activation(channels, activation) for _ in dilations])
+        self.acts2 = nn.ModuleList([make_activation(channels, activation) for _ in dilations])
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        for c1, c2, a1, a2 in zip(self.convs1, self.convs2, self.acts1, self.acts2):
+            y = a1(x)
+            y = c1(y)
+            y = a2(y)
+            y = c2(y)
+            x = x + y
+        return x
+
+    def remove_weight_norm(self) -> None:
+        for layer in list(self.convs1) + list(self.convs2):
+            remove_weight_norm(layer)
+
+
+class ResBlock2(nn.Module):
+    def __init__(self, channels: int, kernel_size: int, dilations: tuple[int, ...], activation: str = "lrelu"):
+        super().__init__()
+        self.convs = nn.ModuleList(
+            [
+                weight_norm(
+                    nn.Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=d,
+                        padding=get_padding(kernel_size, d),
+                    )
+                )
+                for d in dilations
+            ]
+        )
+        self.acts = nn.ModuleList([make_activation(channels, activation) for _ in dilations])
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        for conv, act in zip(self.convs, self.acts):
+            y = act(x)
+            y = conv(y)
+            x = x + y
+        return x
+
+    def remove_weight_norm(self) -> None:
+        for layer in self.convs:
+            remove_weight_norm(layer)
+
+
+class HifiGanGenerator(nn.Module):
+    def __init__(self, cfg: HifiGanConfig):
+        super().__init__()
+        self.cfg = cfg
+        self.num_kernels = len(cfg.resblock_kernel_sizes)
+        self.num_upsamples = len(cfg.upsample_rates)
+        self.conv_pre = weight_norm(
+            nn.Conv1d(cfg.num_mels + cfg.conditioning_channels, cfg.upsample_initial_channel, 7, 1, padding=3)
+        )
+        self.ups = nn.ModuleList()
+        self.up_acts = nn.ModuleList()
+        self.resblocks = nn.ModuleList()
+        resblock_cls = ResBlock1 if cfg.resblock == "1" else ResBlock2
+        for i, (rate, kernel) in enumerate(zip(cfg.upsample_rates, cfg.upsample_kernel_sizes)):
+            in_ch = cfg.upsample_initial_channel // (2**i)
+            out_ch = cfg.upsample_initial_channel // (2 ** (i + 1))
+            self.up_acts.append(make_activation(in_ch, cfg.activation))
+            self.ups.append(
+                weight_norm(
+                    nn.ConvTranspose1d(
+                        in_ch,
+                        out_ch,
+                        kernel,
+                        rate,
+                        padding=(kernel - rate) // 2,
+                    )
+                )
+            )
+            for k, d in zip(cfg.resblock_kernel_sizes, cfg.resblock_dilation_sizes):
+                self.resblocks.append(resblock_cls(out_ch, k, d, cfg.activation))
+        final_ch = cfg.upsample_initial_channel // (2 ** len(cfg.upsample_rates))
+        self.post_act = make_activation(final_ch, cfg.activation)
+        self.conv_post = weight_norm(nn.Conv1d(final_ch, 1, 7, 1, padding=3))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.conv_pre(x)
+        for i, up in enumerate(self.ups):
+            x = self.up_acts[i](x)
+            x = up(x)
+            xs = 0.0
+            for j in range(self.num_kernels):
+                xs = xs + self.resblocks[i * self.num_kernels + j](x)
+            x = xs / self.num_kernels
+        x = self.post_act(x)
+        x = self.conv_post(x)
+        return torch.tanh(x)
+
+    def remove_weight_norm(self) -> None:
+        remove_weight_norm(self.conv_pre)
+        for up in self.ups:
+            remove_weight_norm(up)
+        for block in self.resblocks:
+            block.remove_weight_norm()
+        remove_weight_norm(self.conv_post)
+
+
+def extract_source_features(
+    wav: torch.Tensor,
+    cfg: HifiGanConfig,
+    frames: int,
+    dropout: float = 0.0,
+    noise: float = 0.0,
+) -> torch.Tensor:
+    """Return low-rate F0/voicing features for source-conditioned generators."""
+    pitch = torchaudio.functional.detect_pitch_frequency(
+        wav.detach().cpu(),
+        sample_rate=cfg.sample_rate,
+        frame_time=cfg.hop_size / cfg.sample_rate,
+        win_length=30,
+    ).to(wav.device)
+    if pitch.ndim == 1:
+        pitch = pitch.unsqueeze(0)
+    if pitch.shape[-1] < frames:
+        pitch = F.pad(pitch, (0, frames - pitch.shape[-1]), value=0.0)
+    pitch = pitch[..., :frames]
+    voiced = ((pitch >= 55.0) & (pitch <= 420.0)).float()
+    pitch = pitch.clamp(55.0, 420.0)
+    log_f0 = ((torch.log(pitch) - math.log(140.0)) / 0.45).clamp(-3.0, 3.0) * voiced
+    if noise > 0.0:
+        log_f0 = (log_f0 + torch.randn_like(log_f0) * noise * voiced).clamp(-3.0, 3.0)
+    jump = F.pad((log_f0[..., 1:] - log_f0[..., :-1]).abs(), (1, 0))
+    confidence = torch.exp(-1.5 * jump) * voiced
+    reconstructed_f0 = torch.exp(log_f0 * 0.45 + math.log(140.0))
+    phase = torch.cumsum(2.0 * math.pi * reconstructed_f0 * (cfg.hop_size / cfg.sample_rate), dim=-1)
+    source = torch.stack(
+        [log_f0, voiced, confidence, torch.sin(phase) * confidence, torch.cos(phase) * confidence],
+        dim=1,
+    )
+    if dropout > 0.0:
+        # Drop the complete source sketch for some examples so inference remains
+        # stable when predicted F0 confidence is poor.
+        keep = (torch.rand(source.shape[0], 1, 1, device=source.device) >= dropout).to(source.dtype)
+        source = source * keep
+    return source
+
+
+class DiscriminatorP(nn.Module):
+    def __init__(self, period: int):
+        super().__init__()
+        self.period = period
+        self.convs = nn.ModuleList(
+            [
+                weight_norm(nn.Conv2d(1, 32, (5, 1), (3, 1), padding=(2, 0))),
+                weight_norm(nn.Conv2d(32, 128, (5, 1), (3, 1), padding=(2, 0))),
+                weight_norm(nn.Conv2d(128, 512, (5, 1), (3, 1), padding=(2, 0))),
+                weight_norm(nn.Conv2d(512, 1024, (5, 1), (3, 1), padding=(2, 0))),
+                weight_norm(nn.Conv2d(1024, 1024, (5, 1), 1, padding=(2, 0))),
+            ]
+        )
+        self.conv_post = weight_norm(nn.Conv2d(1024, 1, (3, 1), 1, padding=(1, 0)))
+
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, list[torch.Tensor]]:
+        fmap = []
+        b, c, t = x.shape
+        if t % self.period != 0:
+            x = F.pad(x, (0, self.period - (t % self.period)), mode="reflect")
+            t = x.shape[-1]
+        x = x.view(b, c, t // self.period, self.period)
+        for conv in self.convs:
+            x = F.leaky_relu(conv(x), 0.1)
+            fmap.append(x)
+        x = self.conv_post(x)
+        fmap.append(x)
+        return torch.flatten(x, 1, -1), fmap
+
+
+class MultiPeriodDiscriminator(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.discriminators = nn.ModuleList([DiscriminatorP(p) for p in (2, 3, 5, 7, 11)])
+
+    def forward(self, y: torch.Tensor, y_hat: torch.Tensor):
+        y_d_rs, y_d_gs, fmap_rs, fmap_gs = [], [], [], []
+        for d in self.discriminators:
+            y_d_r, fmap_r = d(y)
+            y_d_g, fmap_g = d(y_hat)
+            y_d_rs.append(y_d_r)
+            y_d_gs.append(y_d_g)
+            fmap_rs.append(fmap_r)
+            fmap_gs.append(fmap_g)
+        return y_d_rs, y_d_gs, fmap_rs, fmap_gs
+
+
+class DiscriminatorS(nn.Module):
+    def __init__(self, use_spectral_norm: bool = False):
+        super().__init__()
+        norm = spectral_norm if use_spectral_norm else weight_norm
+        self.convs = nn.ModuleList(
+            [
+                norm(nn.Conv1d(1, 128, 15, 1, padding=7)),
+                norm(nn.Conv1d(128, 128, 41, 2, groups=4, padding=20)),
+                norm(nn.Conv1d(128, 256, 41, 2, groups=16, padding=20)),
+                norm(nn.Conv1d(256, 512, 41, 4, groups=16, padding=20)),
+                norm(nn.Conv1d(512, 1024, 41, 4, groups=16, padding=20)),
+                norm(nn.Conv1d(1024, 1024, 41, 1, groups=16, padding=20)),
+                norm(nn.Conv1d(1024, 1024, 5, 1, padding=2)),
+            ]
+        )
+        self.conv_post = norm(nn.Conv1d(1024, 1, 3, 1, padding=1))
+
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, list[torch.Tensor]]:
+        fmap = []
+        for conv in self.convs:
+            x = F.leaky_relu(conv(x), 0.1)
+            fmap.append(x)
+        x = self.conv_post(x)
+        fmap.append(x)
+        return torch.flatten(x, 1, -1), fmap
+
+
+class MultiScaleDiscriminator(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.discriminators = nn.ModuleList([DiscriminatorS(True), DiscriminatorS(), DiscriminatorS()])
+        self.meanpools = nn.ModuleList([nn.AvgPool1d(4, 2, padding=2), nn.AvgPool1d(4, 2, padding=2)])
+
+    def forward(self, y: torch.Tensor, y_hat: torch.Tensor):
+        y_d_rs, y_d_gs, fmap_rs, fmap_gs = [], [], [], []
+        for i, d in enumerate(self.discriminators):
+            if i:
+                y = self.meanpools[i - 1](y)
+                y_hat = self.meanpools[i - 1](y_hat)
+            y_d_r, fmap_r = d(y)
+            y_d_g, fmap_g = d(y_hat)
+            y_d_rs.append(y_d_r)
+            y_d_gs.append(y_d_g)
+            fmap_rs.append(fmap_r)
+            fmap_gs.append(fmap_g)
+        return y_d_rs, y_d_gs, fmap_rs, fmap_gs
+
+
+class SpectrogramDiscriminator(nn.Module):
+    def __init__(self):
+        super().__init__()
+        channels = (32, 64, 128, 128)
+        layers: list[nn.Module] = []
+        in_ch = 1
+        for out_ch, stride in zip(channels, ((1, 2), (2, 2), (2, 2), (2, 1))):
+            layers.append(weight_norm(nn.Conv2d(in_ch, out_ch, (5, 5), stride=stride, padding=(2, 2))))
+            in_ch = out_ch
+        self.convs = nn.ModuleList(layers)
+        self.conv_post = weight_norm(nn.Conv2d(in_ch, 1, (3, 3), padding=(1, 1)))
+
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, list[torch.Tensor]]:
+        fmap = []
+        for conv in self.convs:
+            x = F.leaky_relu(conv(x), 0.1)
+            fmap.append(x)
+        x = self.conv_post(x)
+        fmap.append(x)
+        return torch.flatten(x, 1, -1), fmap
+
+
+class MultiResolutionSpectrogramDiscriminator(nn.Module):
+    def __init__(self, fft_sizes: tuple[int, ...] = (256, 512, 1024), hop_sizes: tuple[int, ...] = (64, 128, 256), win_lengths: tuple[int, ...] = (256, 512, 1024)):
+        super().__init__()
+        self.fft_sizes = fft_sizes
+        self.hop_sizes = hop_sizes
+        self.win_lengths = win_lengths
+        self.discriminators = nn.ModuleList([SpectrogramDiscriminator() for _ in fft_sizes])
+
+    def _features(self, wav: torch.Tensor, fft: int, hop: int, win_len: int) -> torch.Tensor:
+        wav = wav.squeeze(1)
+        window = torch.hann_window(win_len, device=wav.device)
+        spec = torch.stft(wav, n_fft=fft, hop_length=hop, win_length=win_len, window=window, return_complex=True)
+        mag = torch.log(spec.abs().clamp_min(1e-5))
+        mean = mag.mean(dim=(1, 2), keepdim=True)
+        std = mag.std(dim=(1, 2), keepdim=True).clamp_min(1e-4)
+        return ((mag - mean) / std).unsqueeze(1)
+
+    def forward(self, y: torch.Tensor, y_hat: torch.Tensor):
+        y_d_rs, y_d_gs, fmap_rs, fmap_gs = [], [], [], []
+        for disc, fft, hop, win_len in zip(self.discriminators, self.fft_sizes, self.hop_sizes, self.win_lengths):
+            y_feat = self._features(y, fft, hop, win_len)
+            y_hat_feat = self._features(y_hat, fft, hop, win_len)
+            y_d_r, fmap_r = disc(y_feat)
+            y_d_g, fmap_g = disc(y_hat_feat)
+            y_d_rs.append(y_d_r)
+            y_d_gs.append(y_d_g)
+            fmap_rs.append(fmap_r)
+            fmap_gs.append(fmap_g)
+        return y_d_rs, y_d_gs, fmap_rs, fmap_gs
+
+
+class MelFrontend(nn.Module):
+    def __init__(self, cfg: HifiGanConfig):
+        super().__init__()
+        self.mel = torchaudio.transforms.MelSpectrogram(
+            sample_rate=cfg.sample_rate,
+            n_fft=cfg.n_fft,
+            win_length=cfg.win_size,
+            hop_length=cfg.hop_size,
+            f_min=cfg.fmin,
+            f_max=cfg.fmax,
+            n_mels=cfg.num_mels,
+            power=1.0,
+            center=True,
+            norm="slaney",
+            mel_scale="slaney",
+        )
+
+    def forward(self, wav: torch.Tensor) -> torch.Tensor:
+        return torch.log(torch.clamp(self.mel(wav), min=1e-5))
+
+
+def load_rows(path: Path, max_rows: int, min_seconds: float, max_seconds: float) -> list[dict]:
+    rows = []
+    with path.open("r", encoding="utf-8-sig") as f:
+        for line in f:
+            if not line.strip():
+                continue
+            row = json.loads(line)
+            audio = Path(str(row.get("target_audio") or ""))
+            text = str(row.get("target_text") or "").strip()
+            dur = float(row.get("target_duration_s") or 0.0)
+            if audio.is_file() and text and min_seconds <= (dur or 4.0) <= max_seconds:
+                rows.append({"audio": str(audio), "text": text, "duration": dur})
+                if max_rows > 0 and len(rows) >= max_rows:
+                    break
+    if not rows:
+        raise RuntimeError(f"No rows loaded from {path}")
+    return rows
+
+
+def load_audio(path: str, sample_rate: int) -> 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)
+    wav = wav.squeeze(0)
+    return wav.clamp(-1, 1)
+
+
+class AudioDataset(Dataset):
+    def __init__(self, rows: list[dict], cfg: HifiGanConfig, segment_size: int, seed: int):
+        self.rows = rows
+        self.cfg = cfg
+        self.segment_size = segment_size
+        self.rng = random.Random(seed)
+
+    def __len__(self) -> int:
+        return len(self.rows)
+
+    def __getitem__(self, idx: int) -> torch.Tensor:
+        wav = load_audio(self.rows[idx]["audio"], self.cfg.sample_rate)
+        if wav.numel() >= self.segment_size:
+            start = self.rng.randint(0, wav.numel() - self.segment_size)
+            return wav[start : start + self.segment_size]
+        return F.pad(wav, (0, self.segment_size - wav.numel()))
+
+
+def feature_loss(fmap_r, fmap_g) -> torch.Tensor:
+    loss = 0.0
+    for dr, dg in zip(fmap_r, fmap_g):
+        for rl, gl in zip(dr, dg):
+            loss = loss + F.l1_loss(rl.detach(), gl)
+    return loss * 2
+
+
+def discriminator_loss(disc_real_outputs, disc_generated_outputs) -> torch.Tensor:
+    loss = 0.0
+    for dr, dg in zip(disc_real_outputs, disc_generated_outputs):
+        loss = loss + torch.mean((1 - dr) ** 2) + torch.mean(dg**2)
+    return loss
+
+
+def generator_loss(disc_outputs) -> torch.Tensor:
+    loss = 0.0
+    for dg in disc_outputs:
+        loss = loss + torch.mean((1 - dg) ** 2)
+    return loss
+
+
+def stft_mag_loss(y_hat: torch.Tensor, y: torch.Tensor, fft_sizes: tuple[int, ...], hop_sizes: tuple[int, ...], win_lengths: tuple[int, ...]) -> torch.Tensor:
+    # Multi-resolution spectral loss catches buzz/shimmer that can hide behind
+    # mel loss, especially for a small generator near convergence.
+    y_hat = y_hat.squeeze(1)
+    y = y.squeeze(1)
+    total = torch.zeros((), device=y.device)
+    for fft, hop, win_len in zip(fft_sizes, hop_sizes, win_lengths):
+        window = torch.hann_window(win_len, device=y.device)
+        pred = torch.stft(y_hat, n_fft=fft, hop_length=hop, win_length=win_len, window=window, return_complex=True)
+        target = torch.stft(y, n_fft=fft, hop_length=hop, win_length=win_len, window=window, return_complex=True)
+        pred_mag = pred.abs().clamp_min(1e-7)
+        target_mag = target.abs().clamp_min(1e-7)
+        sc = torch.linalg.vector_norm(target_mag - pred_mag) / torch.linalg.vector_norm(target_mag).clamp_min(1e-7)
+        log_mag = F.l1_loss(torch.log(pred_mag), torch.log(target_mag))
+        total = total + sc + log_mag
+    return total / max(1, len(fft_sizes))
+
+
+def count_parameters(module: nn.Module) -> int:
+    return sum(p.numel() for p in module.parameters())
+
+
+def jsonable_args(args: argparse.Namespace) -> dict:
+    return {k: str(v) if isinstance(v, Path) else v for k, v in vars(args).items()}
+
+
+def save_checkpoint(
+    path: Path,
+    generator: nn.Module,
+    mpd: nn.Module,
+    msd: nn.Module,
+    optim_g,
+    optim_d,
+    cfg: HifiGanConfig,
+    step: int,
+    args,
+    mrsd: nn.Module | None = None,
+) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    tmp = path.with_suffix(path.suffix + ".tmp")
+    payload = {
+            "generator": generator.state_dict(),
+            "mpd": mpd.state_dict(),
+            "msd": msd.state_dict(),
+            "optim_g": optim_g.state_dict(),
+            "optim_d": optim_d.state_dict(),
+            "config": asdict(cfg),
+            "step": step,
+            "args": jsonable_args(args),
+            "generator_params": count_parameters(generator),
+        }
+    if mrsd is not None:
+        payload["mrsd"] = mrsd.state_dict()
+    torch.save(payload, tmp)
+    tmp.replace(path)
+
+
+def checkpoint_step(path: Path) -> int:
+    stem = path.stem
+    tail = stem.rsplit("-", 1)[-1]
+    return int(tail) if tail.isdigit() else -1
+
+
+def prune_checkpoints(out_dir: Path, variant: str, keep: int) -> None:
+    if keep <= 0:
+        return
+    numbered = [p for p in out_dir.glob(f"hifigan-{variant}-*.pt") if checkpoint_step(p) >= 0]
+    numbered.sort(key=checkpoint_step, reverse=True)
+    for old in numbered[keep:]:
+        old.unlink(missing_ok=True)
+
+
+def latest_checkpoint(out_dir: Path) -> Path | None:
+    numbered = [p for p in out_dir.glob("hifigan-*-*.pt") if checkpoint_step(p) >= 0]
+    if numbered:
+        return max(numbered, key=checkpoint_step)
+    ckpts = sorted(out_dir.glob("hifigan-*-latest.pt"), key=lambda p: p.stat().st_mtime, reverse=True)
+    return ckpts[0] if ckpts else None
+
+
+def partial_load_state(module: nn.Module, state: dict[str, torch.Tensor]) -> tuple[int, int]:
+    current = module.state_dict()
+    patched: dict[str, torch.Tensor] = {}
+    copied = 0
+    skipped = 0
+    for name, target in current.items():
+        source = state.get(name)
+        if source is None:
+            skipped += 1
+            continue
+        if source.shape == target.shape:
+            patched[name] = source
+            copied += 1
+            continue
+        if source.ndim != target.ndim:
+            skipped += 1
+            continue
+        value = target.clone()
+        slices = tuple(slice(0, min(a, b)) for a, b in zip(target.shape, source.shape))
+        value[slices] = source[slices].to(value.device, value.dtype)
+        patched[name] = value
+        copied += 1
+    module.load_state_dict(patched, strict=False)
+    return copied, skipped
+
+
+def train(args: argparse.Namespace) -> None:
+    torch.backends.cudnn.benchmark = True
+    cfg = make_config(args.variant)
+    device = torch.device(args.device)
+    rows = load_rows(args.train_jsonl, args.max_rows, args.min_seconds, args.max_seconds)
+    rng = random.Random(args.seed)
+    rng.shuffle(rows)
+    dataset = AudioDataset(rows, cfg, args.segment_size, args.seed)
+    loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True, drop_last=True, num_workers=args.num_workers)
+    mel_frontend = MelFrontend(cfg).to(device)
+    generator = HifiGanGenerator(cfg).to(device)
+    mpd = MultiPeriodDiscriminator().to(device)
+    msd = MultiScaleDiscriminator().to(device)
+    mrsd = MultiResolutionSpectrogramDiscriminator().to(device) if args.spec_disc_weight > 0.0 else None
+    optim_g = torch.optim.AdamW(generator.parameters(), lr=args.lr, betas=(0.8, 0.99))
+    disc_params = list(mpd.parameters()) + list(msd.parameters())
+    if mrsd is not None:
+        disc_params += list(mrsd.parameters())
+    optim_d = torch.optim.AdamW(disc_params, lr=args.lr, betas=(0.8, 0.99))
+    start_step = 0
+    if args.init_checkpoint and not args.resume:
+        ckpt = torch.load(args.init_checkpoint, map_location=device, weights_only=False)
+        if args.partial_init:
+            copied, skipped = partial_load_state(generator, ckpt["generator"])
+            print(f"Partially initialized generator from {args.init_checkpoint}: copied={copied} skipped={skipped}")
+        else:
+            generator.load_state_dict(ckpt["generator"])
+        if "mpd" in ckpt and "msd" in ckpt:
+            mpd.load_state_dict(ckpt["mpd"])
+            msd.load_state_dict(ckpt["msd"])
+        if mrsd is not None and "mrsd" in ckpt:
+            mrsd.load_state_dict(ckpt["mrsd"])
+        can_load_disc_optim = mrsd is None or "mrsd" in ckpt
+        if not args.partial_init and not args.reset_optim and "optim_g" in ckpt:
+            optim_g.load_state_dict(ckpt["optim_g"])
+        if not args.partial_init and not args.reset_optim and can_load_disc_optim and "optim_d" in ckpt:
+            optim_d.load_state_dict(ckpt["optim_d"])
+        for group in optim_g.param_groups:
+            group["lr"] = args.lr
+        for group in optim_d.param_groups:
+            group["lr"] = args.lr
+        start_step = int(ckpt.get("step") or 0)
+        print(f"Initialized {args.init_checkpoint} at step {start_step}; lr={args.lr:g}")
+    if args.resume:
+        ckpt_path = latest_checkpoint(args.out_dir)
+        if ckpt_path:
+            ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
+            generator.load_state_dict(ckpt["generator"])
+            mpd.load_state_dict(ckpt["mpd"])
+            msd.load_state_dict(ckpt["msd"])
+            if mrsd is not None and "mrsd" in ckpt:
+                mrsd.load_state_dict(ckpt["mrsd"])
+            optim_g.load_state_dict(ckpt["optim_g"])
+            optim_d.load_state_dict(ckpt["optim_d"])
+            for group in optim_g.param_groups:
+                group["lr"] = args.lr
+            for group in optim_d.param_groups:
+                group["lr"] = args.lr
+            start_step = int(ckpt.get("step") or 0)
+            print(f"Resumed {ckpt_path} at step {start_step}; lr={args.lr:g}")
+
+    args.out_dir.mkdir(parents=True, exist_ok=True)
+    prune_checkpoints(args.out_dir, args.variant, args.keep_checkpoints)
+    (args.out_dir / "config.json").write_text(
+        json.dumps(
+            {
+                "config": asdict(cfg),
+                "args": jsonable_args(args),
+                "rows": len(rows),
+                "generator_params": count_parameters(generator),
+                "mpd_params": count_parameters(mpd),
+                "msd_params": count_parameters(msd),
+                "mrsd_params": count_parameters(mrsd) if mrsd is not None else 0,
+            },
+            indent=2,
+        ),
+        encoding="utf-8",
+    )
+    print(f"Variant: {args.variant}")
+    print(f"Rows: {len(rows)}")
+    print(f"Generator params: {count_parameters(generator):,} ({count_parameters(generator)/1_000_000:.3f}M)")
+    print(f"MPD params: {count_parameters(mpd):,} MSD params: {count_parameters(msd):,} (training only)")
+    if mrsd is not None:
+        print(f"MRSD params: {count_parameters(mrsd):,} (training only)")
+    if args.steps == 0:
+        return
+
+    step = start_step
+    started = time.time()
+    try:
+        while step < args.steps:
+            for wav in loader:
+                step += 1
+                y = wav.unsqueeze(1).to(device)
+                with torch.no_grad():
+                    mel = mel_frontend(wav.to(device))
+                    if cfg.conditioning_channels:
+                        source = extract_source_features(
+                            wav.to(device),
+                            cfg,
+                            mel.shape[-1],
+                            dropout=args.source_dropout,
+                            noise=args.source_noise,
+                        )
+                        generator_input = torch.cat([mel, source], dim=1)
+                    else:
+                        generator_input = mel
+                y_hat = generator(generator_input)
+                common = min(y.shape[-1], y_hat.shape[-1])
+                y = y[..., :common]
+                y_hat = y_hat[..., :common]
+                y_mel = mel_frontend(y.squeeze(1))
+                y_hat_mel = mel_frontend(y_hat.squeeze(1))
+
+                optim_d.zero_grad(set_to_none=True)
+                y_df_hat_r, y_df_hat_g, _, _ = mpd(y, y_hat.detach())
+                y_ds_hat_r, y_ds_hat_g, _, _ = msd(y, y_hat.detach())
+                loss_disc = discriminator_loss(y_df_hat_r, y_df_hat_g) + discriminator_loss(y_ds_hat_r, y_ds_hat_g)
+                loss_spec_disc = torch.zeros((), device=device)
+                if mrsd is not None:
+                    y_dm_hat_r, y_dm_hat_g, _, _ = mrsd(y, y_hat.detach())
+                    loss_spec_disc = discriminator_loss(y_dm_hat_r, y_dm_hat_g)
+                    loss_disc = loss_disc + args.spec_disc_weight * loss_spec_disc
+                loss_disc.backward()
+                torch.nn.utils.clip_grad_norm_(disc_params, args.grad_clip)
+                optim_d.step()
+
+                optim_g.zero_grad(set_to_none=True)
+                mel_loss = F.l1_loss(y_mel, y_hat_mel) * args.mel_weight
+                y_df_hat_r, y_df_hat_g, fmap_f_r, fmap_f_g = mpd(y, y_hat)
+                y_ds_hat_r, y_ds_hat_g, fmap_s_r, fmap_s_g = msd(y, y_hat)
+                loss_fm = feature_loss(fmap_f_r, fmap_f_g) + feature_loss(fmap_s_r, fmap_s_g)
+                loss_gen = generator_loss(y_df_hat_g) + generator_loss(y_ds_hat_g)
+                loss_spec_gen = torch.zeros((), device=device)
+                loss_spec_fm = torch.zeros((), device=device)
+                if mrsd is not None:
+                    y_dm_hat_r, y_dm_hat_g, fmap_m_r, fmap_m_g = mrsd(y, y_hat)
+                    loss_spec_gen = generator_loss(y_dm_hat_g)
+                    loss_spec_fm = feature_loss(fmap_m_r, fmap_m_g)
+                wav_l1 = F.l1_loss(y_hat, y) * args.wav_weight
+                stft_loss = torch.zeros((), device=device)
+                if args.stft_weight > 0.0:
+                    stft_loss = stft_mag_loss(y_hat, y, (512, 1024, 2048), (128, 256, 512), (512, 1024, 2048)) * args.stft_weight
+                loss_g = (
+                    mel_loss
+                    + args.fm_weight * loss_fm
+                    + args.adv_weight * loss_gen
+                    + wav_l1
+                    + stft_loss
+                    + args.spec_disc_weight * loss_spec_gen
+                    + args.spec_fm_weight * loss_spec_fm
+                )
+                loss_g.backward()
+                grad_g = torch.nn.utils.clip_grad_norm_(generator.parameters(), args.grad_clip)
+                optim_g.step()
+
+                if step == 1 or step % args.log_interval == 0:
+                    elapsed = max(time.time() - started, 1e-6)
+                    speed = (step - start_step) / elapsed
+                    eta = (args.steps - step) / max(speed, 1e-6)
+                    print(
+                        f"step={step}/{args.steps} g={loss_g.item():.4f} d={loss_disc.item():.4f} "
+                        f"mel={mel_loss.item():.4f} fm={loss_fm.item():.4f} adv={loss_gen.item():.4f} "
+                        f"wav={wav_l1.item():.4f} stft={stft_loss.item():.4f} "
+                        f"sd={loss_spec_disc.item():.4f} sfm={loss_spec_fm.item():.4f} sadv={loss_spec_gen.item():.4f} "
+                        f"grad={float(grad_g):.3f} speed={speed:.3f} step/s eta={eta/60:.1f}m",
+                        flush=True,
+                    )
+                if step % args.save_interval == 0 or step >= args.steps:
+                    prune_checkpoints(args.out_dir, args.variant, max(args.keep_checkpoints - 1, 0))
+                    save_checkpoint(args.out_dir / f"hifigan-{args.variant}-{step}.pt", generator, mpd, msd, optim_g, optim_d, cfg, step, args, mrsd)
+                    save_checkpoint(args.out_dir / f"hifigan-{args.variant}-latest.pt", generator, mpd, msd, optim_g, optim_d, cfg, step, args, mrsd)
+                if step >= args.steps:
+                    break
+    except KeyboardInterrupt:
+        if step > start_step:
+            save_checkpoint(args.out_dir / f"hifigan-{args.variant}-interrupt-{step}.pt", generator, mpd, msd, optim_g, optim_d, cfg, step, args, mrsd)
+            save_checkpoint(args.out_dir / f"hifigan-{args.variant}-latest.pt", generator, mpd, msd, optim_g, optim_d, cfg, step, args, mrsd)
+            print(f"Interrupted. Saved checkpoint at step {step}.", flush=True)
+        raise
+    save_checkpoint(args.out_dir / f"hifigan-{args.variant}-final.pt", generator, mpd, msd, optim_g, optim_d, cfg, step, args, mrsd)
+    print(f"Done. {args.out_dir}")
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser(description="Train exact-ish HiFi-GAN V2/V3 oracle vocoders on corrected Mark audio.")
+    ap.add_argument("--train-jsonl", type=Path, required=True)
+    ap.add_argument("--out-dir", type=Path, required=True)
+    ap.add_argument(
+        "--variant",
+        choices=["v2", "v2plus", "v2wide", "snake_v2mid", "snake_8k", "snake_8k40", "snake_v2balanced", "source_snake_v2balanced", "v3"],
+        required=True,
+    )
+    ap.add_argument("--steps", type=int, default=5000)
+    ap.add_argument("--max-rows", type=int, default=0)
+    ap.add_argument("--min-seconds", type=float, default=1.0)
+    ap.add_argument("--max-seconds", type=float, default=12.0)
+    ap.add_argument("--segment-size", type=int, default=8192)
+    ap.add_argument("--batch-size", type=int, default=8)
+    ap.add_argument("--num-workers", type=int, default=0)
+    ap.add_argument("--lr", type=float, default=2.0e-4)
+    ap.add_argument("--mel-weight", type=float, default=45.0)
+    ap.add_argument("--wav-weight", type=float, default=1.0)
+    ap.add_argument("--fm-weight", type=float, default=1.0)
+    ap.add_argument("--adv-weight", type=float, default=1.0)
+    ap.add_argument("--stft-weight", type=float, default=0.0)
+    ap.add_argument("--spec-disc-weight", type=float, default=0.0, help="Training-only multi-resolution spectrogram adversarial weight.")
+    ap.add_argument("--spec-fm-weight", type=float, default=0.0, help="Training-only spectrogram discriminator feature-matching weight.")
+    ap.add_argument("--source-dropout", type=float, default=0.0, help="Probability of dropping source conditioning per training example.")
+    ap.add_argument("--source-noise", type=float, default=0.0, help="Stddev of normalized log-F0 corruption for source conditioning.")
+    ap.add_argument("--grad-clip", type=float, default=1000.0)
+    ap.add_argument("--log-interval", type=int, default=50)
+    ap.add_argument("--save-interval", type=int, default=1000)
+    ap.add_argument("--keep-checkpoints", type=int, default=12)
+    ap.add_argument("--seed", type=int, default=1234)
+    ap.add_argument("--device", default="cuda")
+    ap.add_argument("--resume", action="store_true")
+    ap.add_argument("--init-checkpoint", type=Path)
+    ap.add_argument("--partial-init", action="store_true", help="Slice-copy compatible generator weights from init-checkpoint into a resized generator.")
+    ap.add_argument("--reset-optim", action="store_true", help="When initializing from a checkpoint, load model/discriminators but start fresh optimizers.")
+    args = ap.parse_args()
+    train(args)
+
+
+if __name__ == "__main__":
+    main()

scripts/align_durations_v4.py

@@ -0,0 +1,181 @@
+#!/usr/bin/env python3
+"""v4 aligner: PHONE-LEVEL forced alignment via espeak phoneme-CTC.
+
+Root-cause fix. v1/v3 did CHAR-level CTC then heuristically SPLIT each char span
+across its phones -> wrong RELATIVE phone durations -> over-smoothed acoustic.
+v4 aligns OUR frontend's phone sequence DIRECTLY:
+  each phone -> one IPA token in facebook/wav2vec2-lv-60-espeak-cv-ft's vocab
+  -> torchaudio.forced_align(emissions, our_phone_ipa_targets) -> per-phone frames.
+No splitting. Durations are contiguous (cover all frames); fit_durations() rescales
+to true mel length at train time, so only RELATIVE proportions matter -- and those
+are now correct from the phone recognizer.
+
+Non-alignable phones (SP / punctuation / OOV) carry a small relative value.
+"""
+from __future__ import annotations
+import argparse, json, sys
+import numpy as np, torch, soundfile as sf, librosa
+sys.path.insert(0, "/home/luigi/jetson-tts/mossnano/zhtw8k")
+
+ESPEAK = "facebook/wav2vec2-lv-60-espeak-cv-ft"
+
+# ARPABET (en, stress stripped) -> single IPA token present in the espeak vocab.
+ARPA2IPA = {
+    'AA':'ɑ','AE':'æ','AH':'ʌ','AO':'ɔ','AW':'aʊ','AY':'aɪ','B':'b','CH':'tʃ',
+    'D':'d','DH':'ð','EH':'ɛ','ER':'ɚ','EY':'eɪ','F':'f','G':'ɡ','HH':'h',
+    'IH':'ɪ','IY':'i','JH':'dʒ','K':'k','L':'l','M':'m','N':'n','NG':'ŋ',
+    'OW':'oʊ','OY':'ɔɪ','P':'p','R':'ɹ','S':'s','SH':'ʃ','T':'t','TH':'θ',
+    'UH':'ʊ','UW':'u','V':'v','W':'w','Y':'j','Z':'z','ZH':'ʒ',
+}
+# bopomofo (zh) -> toneless IPA token in the espeak Mandarin inventory (best-effort).
+BOPO2IPA = {
+    'ㄅ':'p','ㄆ':'pʰ','ㄇ':'m','ㄈ':'f','ㄉ':'t','ㄊ':'tʰ','ㄋ':'n','ㄌ':'l',
+    'ㄍ':'k','ㄎ':'kʰ','ㄏ':'x','ㄐ':'tɕ','ㄑ':'tɕh','ㄒ':'ɕ','ㄓ':'tʃ','ㄔ':'tʃʰ',
+    'ㄕ':'ʂ','ㄖ':'ʐ','ㄗ':'ts','ㄘ':'tsh','ㄙ':'s',
+    'ㄧ':'i','ㄨ':'u','ㄩ':'y','ㄚ':'a','ㄛ':'o','ㄜ':'ɤ','ㄝ':'e','ㄞ':'ai',
+    'ㄟ':'ei','ㄠ':'au','ㄡ':'ou','ㄢ':'a','ㄣ':'ə','ㄤ':'ɑ','ㄥ':'ə','ㄦ':'ɚ',
+    'ㄭ':'ɨ',
+}
+
+
+def build_frontend_with_charidx():
+    import frontend_bopomofo as F
+    import re
+    F._lazy()
+    def run(text):
+        text = F._zh_num(text)
+        bopo = F._g2pw(text)[0]
+        chars = list(text)
+        phones, tones, langs = [], [], []
+        i = 0
+        while i < len(chars):
+            b = bopo[i] if i < len(bopo) else None
+            ch = chars[i]
+            if b is not None:
+                units, tone = F._split_syllable(b)
+                for u in units:
+                    phones.append(u); tones.append(min(tone,5)); langs.append(0)
+                i += 1
+            elif re.match(r'[A-Za-z]', ch):
+                j = i
+                while j < len(chars) and re.match(r"[A-Za-z']", chars[j]): j += 1
+                for p in F._g2pen(''.join(chars[i:j])):
+                    p = p.strip()
+                    if not p: continue
+                    if p[-1].isdigit(): st=int(p[-1]); p=p[:-1]
+                    else: st=0
+                    if p in F.SYM2ID:
+                        phones.append(p); tones.append(st); langs.append(1)
+                phones.append('SP'); tones.append(0); langs.append(1)
+                i = j
+            else:
+                if ch in F.PUNCT:
+                    phones.append(ch); tones.append(0); langs.append(0)
+                elif ch.strip()=='' and phones and phones[-1]!='SP':
+                    phones.append('SP'); tones.append(0); langs.append(0)
+                i += 1
+        return text, phones, tones, langs
+    return run, F
+
+
+class Aligner:
+    def __init__(self, device):
+        from transformers import Wav2Vec2ForCTC
+        from huggingface_hub import hf_hub_download
+        self.vocab = json.load(open(hf_hub_download(ESPEAK, "vocab.json")))
+        self.blank = self.vocab["<pad>"]
+        self.model = Wav2Vec2ForCTC.from_pretrained(ESPEAK).to(device).eval()
+        self.device = device
+
+    def emissions(self, audio16k):
+        iv = torch.from_numpy(audio16k).float().unsqueeze(0).to(self.device)
+        with torch.inference_mode():
+            logits = self.model(iv).logits[0]              # [T,V]
+        return torch.log_softmax(logits, dim=-1).cpu()
+
+    def durations(self, audio16k, phones, langs):
+        """Return integer relative duration per phone (contiguous CTC spans)."""
+        from torchaudio.functional import forced_align
+        emis = self.emissions(audio16k)
+        T = emis.shape[0]
+        # target = alignable phones mapped to a vocab token id
+        tgt_ids, tgt_pi = [], []
+        for pi, (p, lg) in enumerate(zip(phones, langs)):
+            tok = ARPA2IPA.get(p) if lg == 1 else BOPO2IPA.get(p)
+            if tok is not None and tok in self.vocab:
+                tgt_ids.append(self.vocab[tok]); tgt_pi.append(pi)
+        n = len(phones); dur = [0.0]*n
+        if len(tgt_ids) < 1 or len(tgt_ids) > T:
+            # fall back: uniform
+            for pi in range(n): dur[pi] = 1.0
+            return [max(1,int(round(x))) for x in dur], 0
+        tokens = torch.tensor([tgt_ids], dtype=torch.int32)
+        try:
+            aligned, _ = forced_align(emis.unsqueeze(0), tokens, blank=self.blank)
+        except Exception:
+            for pi in range(n): dur[pi] = 1.0
+            return [max(1,int(round(x))) for x in dur], 0
+        path = aligned[0].tolist()
+        # first frame of each emitted target token (in order)
+        starts = []; prev = self.blank
+        for fi, tk in enumerate(path):
+            if tk != self.blank and tk != prev:
+                starts.append(fi)
+            prev = tk
+        L = min(len(starts), len(tgt_pi))
+        # contiguous span per alignable phone: [start_k, start_{k+1})
+        for k in range(L):
+            s = starts[k]
+            e = starts[k+1] if k+1 < len(starts) else T
+            dur[tgt_pi[k]] = max(1.0, float(e - s))
+        # leading silence -> give to first alignable phone
+        if L >= 1 and starts[0] > 0:
+            dur[tgt_pi[0]] += starts[0]
+        # non-alignable phones (SP / punct / OOV): small relative value
+        anchored = [d for d in dur if d > 0]
+        base = (sum(anchored)/len(anchored)) if anchored else 4.0
+        for pi in range(n):
+            if dur[pi] == 0:
+                p = phones[pi]
+                dur[pi] = base*0.5 if p == 'SP' else base*0.25
+        return [max(1,int(round(x))) for x in dur], len(tgt_pi)
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--manifest", required=True)
+    ap.add_argument("--out", required=True)
+    ap.add_argument("--limit", type=int, default=0)
+    ap.add_argument("--device", default="cuda")
+    args = ap.parse_args()
+    dev = args.device if torch.cuda.is_available() else "cpu"
+    run, F = build_frontend_with_charidx()
+    al = Aligner(dev)
+    rows = [json.loads(l) for l in open(args.manifest)]
+    if args.limit: rows = rows[:args.limit]
+    out = open(args.out, "w", encoding="utf-8"); n_ok = 0
+    for r in rows:
+        wav = r.get("wav") or r.get("target_audio")
+        try:
+            a, sr = sf.read(wav, dtype="float32")
+            if a.ndim > 1: a = a.mean(1)
+            a16 = librosa.resample(a, orig_sr=sr, target_sr=16000) if sr != 16000 else a
+            _, phones, tones, langs = run(r["text"])
+            durs, n_anchor = al.durations(a16, phones, langs)
+            phone_ids = [F.SYM2ID.get(p, F.SYM2ID['UNK']) for p in phones]
+            out.write(json.dumps({
+                "id": r["id"], "text": r["text"], "target_audio": wav,
+                "phone_ids": phone_ids, "tone_ids": tones, "lang_ids": langs,
+                "hifigan_durations": durs, "speaker_id": 0,
+            }, ensure_ascii=False)+"\n")
+            n_ok += 1
+            if n_ok <= 5:
+                print(f"  {r['id']} n_ph={len(phones)} anchored={n_anchor} "
+                      f"dur_sum={sum(durs)} dur[:18]={durs[:18]}", flush=True)
+        except Exception as e:
+            print(f"  [skip {r.get('id')}] {type(e).__name__}: {e}", flush=True)
+    out.close(); print(f"DONE aligned {n_ok} rows -> {args.out}")
+
+
+if __name__ == "__main__":
+    main()

scripts/assess_big.py

@@ -0,0 +1,38 @@
+#!/usr/bin/env python3
+"""De-noised eval: score a synth dir against eval_big.jsonl (36 held-out sentences).
+Run in moss-nano-venv. Usage: python assess_big.py --synth-dir m7_eval_big
+Pairs with: synth_from_text.py --onnx-dir <m>_onnx --out-dir <m>_eval_big --texts eval_big.jsonl
+Reports aggregate zh CER (zh + mix categories) and en WER, plus per-category, for low-noise comparison."""
+import argparse, json, sys
+ZT = "/home/luigi/jetson-tts/mossnano/zhtw8k"
+sys.path.insert(0, ZT)
+import xasr_offline as X
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--synth-dir", required=True)
+    ap.add_argument("--tag", default="")
+    args = ap.parse_args()
+    lang = {r["id"]: r["lang"] for r in (json.loads(l) for l in open(f"{ZT}/eval_big.jsonl"))}
+    text = {r["id"]: r["text"] for r in (json.loads(l) for l in open(f"{ZT}/eval_big.jsonl"))}
+    rows = [json.loads(l) for l in open(f"{args.synth_dir}/synth.jsonl") if l.strip()]
+    cat = {"zh": [], "mix": [], "en": []}
+    for r in rows:
+        hyp = X.asr(r["wav"])
+        sc = X.score(text[r["id"]], hyp)
+        v = sc if not isinstance(sc, dict) else sc.get("cer", sc.get("wer"))
+        cat[lang[r["id"]]].append(v)
+
+    def avg(xs):
+        return sum(xs) / len(xs) if xs else float("nan")
+    zh_cer = avg(cat["zh"] + cat["mix"])  # CER over zh + code-mix
+    en_wer = avg(cat["en"])
+    print(f"[{args.tag}] N={len(rows)}  zh-only={avg(cat['zh']):.3f}  mix={avg(cat['mix']):.3f}  "
+          f"en={avg(cat['en']):.3f}")
+    print(f"[{args.tag}] AGGREGATE  zh_CER(zh+mix)={zh_cer:.3f}  en_WER={en_wer:.3f}  "
+          f"(n_zh={len(cat['zh'])} n_mix={len(cat['mix'])} n_en={len(cat['en'])})")
+
+
+if __name__ == "__main__":
+    main()

scripts/export_8k.py

@@ -0,0 +1,136 @@
+#!/usr/bin/env python3
+"""Export a trained zh-TW/en 8k Inflect-Nano (acoustic + snake_8k vocoder) to ONNX.
+Config-driven from the train() checkpoints. FastSpeech split:
+  encoder.onnx -> numpy host_regulate -> decoder.onnx -> vocoder.onnx.
+Validates full-pipeline parity vs torch. Run in moss-train-venv."""
+from __future__ import annotations
+import argparse, sys, math, json
+from pathlib import Path
+import numpy as np, torch
+
+REPO = "/tmp/inflect-nano"
+sys.path.insert(0, REPO)
+from inflect_nano.acoustic import MicroFastSpeech, MicroFastSpeechConfig
+from inflect_nano.vocoder import HifiGanGenerator, make_config
+
+
+class EncoderHead(torch.nn.Module):
+    def __init__(self, m): super().__init__(); self.m = m
+    def forward(self, phone, tone, lang, speaker):
+        m = self.m
+        tok = torch.ones_like(phone, dtype=torch.bool)
+        enc = m.encode(phone, tone, lang, speaker, tok)
+        log_dur, energy, bright, pitch = m.predict_prosody(enc, tok)
+        dur = (torch.exp(log_dur) - 1.0).clamp(0, 80).round().clamp_min(1).long()
+        cond = enc + m.energy_proj(energy.unsqueeze(-1)) + m.bright_proj(bright.unsqueeze(-1))
+        pitch = torch.stack([pitch[..., 0], pitch[..., 1].clamp(0, 1)], dim=-1)
+        return cond, dur, pitch
+
+
+class DecoderHead(torch.nn.Module):
+    def __init__(self, m): super().__init__(); self.m = m
+    def forward(self, frames, frame_meta, local_ctx_raw, abs_pos, pitch_frame, frame_mask):
+        m = self.m
+        x = frames + m.frame_proj(frame_meta) + m.local_ctx(local_ctx_raw)
+        x = x + m.abs_frame(abs_pos)
+        if m.cfg.use_frame_pitch:
+            x = x + m.pitch_proj(pitch_frame)
+        for blk in m.decoder:
+            x = blk(x, frame_mask)
+        x = x + m.frame_gru(x)[0]
+        mel = m.mel_head(x).transpose(1, 2)
+        return mel + m.cfg.postnet_scale * m.postnet(mel)
+
+
+def host_regulate(cond, dur, pitch, abs_bins, max_frames):
+    c = cond[0]; d = dur[0].astype(np.int64); d[d < 0] = 0
+    T, H = c.shape
+    frames = np.repeat(c, d, axis=0); F = frames.shape[0]
+    tok = np.repeat(np.arange(T), d); starts = np.cumsum(d) - d
+    within = np.arange(F) - starts[tok]; dpf = d[tok].astype(np.float32)
+    rel = (within / np.maximum(dpf - 1, 1)).astype(np.float32)
+    tc = max(1, int((d > 0).sum())); token_pos = (tok / max(1, tc - 1)).astype(np.float32)
+    ld = (np.log1p(dpf) / 6.0).astype(np.float32); center = 1.0 - np.abs(rel * 2 - 1)
+    fm = np.stack([rel, 1 - rel, center, np.sin(rel*np.pi), np.cos(rel*np.pi), token_pos, ld, dpf/40.0], -1).astype(np.float32)
+    prev = np.concatenate([c[:1], c[:-1]], 0); nxt = np.concatenate([c[1:], c[-1:]], 0)
+    lc = np.repeat(np.concatenate([prev, c, nxt], -1), d, axis=0).astype(np.float32)
+    pos = np.arange(F); abs_pos = np.minimum(pos*abs_bins//max(1, max_frames), abs_bins-1).astype(np.int64)
+    pf = np.repeat(pitch[0], d, axis=0).astype(np.float32)
+    return {"frames": frames[None].astype(np.float32), "frame_meta": fm[None], "local_ctx_raw": lc[None],
+            "abs_pos": abs_pos[None], "pitch_frame": pf[None], "frame_mask": np.ones((1, F), bool)}
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--acoustic-ckpt", required=True)
+    ap.add_argument("--vocoder-ckpt", required=True)
+    ap.add_argument("--out-dir", required=True)
+    ap.add_argument("--symbol-table", default="/home/luigi/jetson-tts/mossnano/zhtw8k/symbol_table.json")
+    args = ap.parse_args()
+    import onnxruntime as ort
+    OUT = Path(args.out_dir); OUT.mkdir(parents=True, exist_ok=True)
+    dev = torch.device("cpu")
+
+    ac = torch.load(args.acoustic_ckpt, map_location=dev, weights_only=False)
+    cfg = MicroFastSpeechConfig(**ac["config"])
+    m = MicroFastSpeech(cfg); m.load_state_dict(ac["model"]); m.eval()
+    enc, dec = EncoderHead(m).eval(), DecoderHead(m).eval()
+    print(f"acoustic: sr={cfg.sample_rate} vocab={cfg.vocab_size} tone={cfg.tone_size} lang={cfg.lang_size} "
+          f"abs_bins={cfg.abs_frame_bins} max_frames={cfg.max_frames}")
+
+    vc = torch.load(args.vocoder_ckpt, map_location=dev, weights_only=False)
+    vcfg = make_config(vc["config"]["variant"])
+    vm = HifiGanGenerator(vcfg); vm.load_state_dict(vc["generator"]); vm.remove_weight_norm(); vm.eval()
+    assert vcfg.sample_rate == cfg.sample_rate
+
+    # sample input: a short valid id sequence (plumbing test)
+    T = 40
+    g = torch.Generator().manual_seed(0)
+    phone = torch.randint(1, min(80, cfg.vocab_size), (1, T), generator=g)
+    tone = torch.randint(0, cfg.tone_size, (1, T), generator=g)
+    lang = torch.randint(0, cfg.lang_size, (1, T), generator=g)
+    spk = torch.zeros(1, dtype=torch.long)
+
+    with torch.no_grad():
+        cond, dur, pitch = enc(phone, tone, lang, spk)
+        reg = host_regulate(cond.numpy(), dur.numpy(), pitch.numpy(), cfg.abs_frame_bins, cfg.max_frames)
+        bt = tuple(torch.from_numpy(reg[k]).clone() for k in ["frames","frame_meta","local_ctx_raw","abs_pos","pitch_frame","frame_mask"])
+        mel_split = dec(*bt)
+        mel_ref = m.infer(phone, tone, lang, spk)
+    print(f"mel parity max_abs_diff={float((mel_ref-mel_split).abs().max()):.2e}")
+
+    torch.onnx.export(enc, (phone, tone, lang, spk), str(OUT/"acoustic_encoder.onnx"),
+        input_names=["phone","tone","lang","speaker"], output_names=["conditioned","durations","pitch"],
+        dynamic_axes={"phone":{1:"T"},"tone":{1:"T"},"lang":{1:"T"},"conditioned":{1:"T"},"durations":{1:"T"},"pitch":{1:"T"}},
+        opset_version=17, dynamo=False)
+    bn = ["frames","frame_meta","local_ctx_raw","abs_pos","pitch_frame","frame_mask"]
+    torch.onnx.export(dec, bt, str(OUT/"acoustic_decoder.onnx"),
+        input_names=bn, output_names=["mel"],
+        dynamic_axes={**{n:{1:"F"} for n in bn}, "mel":{2:"F"}}, opset_version=17, dynamo=False)
+    dummy = torch.randn(1, vcfg.num_mels, 60)  # match vocoder mel count (40 for snake_8k40, 80 otherwise)
+    torch.onnx.export(vm, dummy, str(OUT/"vocoder.onnx"), input_names=["mel"], output_names=["wav"],
+        dynamic_axes={"mel":{2:"frames"},"wav":{2:"samples"}}, opset_version=17, dynamo=False)
+
+    # full-pipeline ONNX parity
+    sA = ort.InferenceSession(str(OUT/"acoustic_encoder.onnx"), providers=["CPUExecutionProvider"])
+    sB = ort.InferenceSession(str(OUT/"acoustic_decoder.onnx"), providers=["CPUExecutionProvider"])
+    sV = ort.InferenceSession(str(OUT/"vocoder.onnx"), providers=["CPUExecutionProvider"])
+    oc, od, op = sA.run(None, {"phone":phone.numpy(),"tone":tone.numpy(),"lang":lang.numpy(),"speaker":spk.numpy()})
+    reg2 = host_regulate(oc, od, op, cfg.abs_frame_bins, cfg.max_frames)
+    feeds = {n:(reg2[n].astype(np.float32) if reg2[n].dtype!=bool else reg2[n]) for n in bn}
+    feeds["abs_pos"] = reg2["abs_pos"].astype(np.int64)
+    mel_onnx = sB.run(None, feeds)[0]
+    wav_onnx = sV.run(None, {"mel": mel_onnx.astype(np.float32)})[0]
+    with torch.inference_mode(): wav_ref = vm(mel_ref).numpy()
+    n = min(wav_ref.shape[-1], wav_onnx.shape[-1])
+    print(f"FULL-PIPELINE wav parity max_abs_diff={float(np.abs(wav_ref[...,:n]-wav_onnx[...,:n]).max()):.2e}")
+    # save metadata for the Nano runtime
+    json.dump({"sample_rate":cfg.sample_rate,"abs_frame_bins":cfg.abs_frame_bins,"max_frames":cfg.max_frames,
+               "hop_size":vcfg.hop_size,"n_mels":cfg.n_mels,"use_frame_pitch":cfg.use_frame_pitch},
+              open(OUT/"meta.json","w"), indent=1)
+    print("sizes(KB):", {f.name: f.stat().st_size//1024 for f in OUT.glob("*.onnx")})
+    print("EXPORT_OK", OUT)
+
+
+if __name__ == "__main__":
+    main()

scripts/frontend_bopomofo.py

@@ -0,0 +1,102 @@
+"""zh-TW/en unified frontend for the Inflect-Nano retrain.
+zh chars -> bopomofo (g2pw, Taiwan readings) -> zhuyin symbol units + tone (1-5);
+en words -> arpabet (g2p_en) + stress; one sequence, per-phone language id (ZH/EN).
+"""
+from __future__ import annotations
+import re
+from g2pw import G2PWConverter
+from g2p_en import G2p
+
+# 37 standard zhuyin symbols (U+3105..U+3129)
+ZHUYIN = [chr(c) for c in range(0x3105, 0x312A)]
+ARPABET = ['AA','AE','AH','AO','AW','AY','B','CH','D','DH','EH','ER','EY','F','G','HH',
+           'IH','IY','JH','K','L','M','N','NG','OW','OY','P','R','S','SH','T','TH',
+           'UH','UW','V','W','Y','Z','ZH']
+PUNCT = [',', '.', '?', '!', '…', '-', "'"]
+SPECIAL = ['_blank', '_pad', 'UNK', 'SP']            # SP = inter-word/space pause
+SYMBOLS = SPECIAL + ZHUYIN + ARPABET + PUNCT
+SYM2ID = {s: i for i, s in enumerate(SYMBOLS)}
+LANG = {'ZH': 0, 'EN': 1}                            # per-phone language id
+
+_g2pw = None
+_g2pen = None
+_zh_num = None
+
+_ZH_DIGIT = {"0":"零","1":"一","2":"二","3":"三","4":"四","5":"五","6":"六","7":"七","8":"八","9":"九"}
+
+def _make_zh_normalizer():
+    import cn2an
+    def norm(text):
+        # long digit runs (>=5, e.g. phone/order numbers) -> digit-by-digit zh; else cardinal
+        def repl(m):
+            d = m.group(0)
+            if len(d) >= 5:
+                return "".join(_ZH_DIGIT[c] for c in d)
+            try:
+                return cn2an.an2cn(d, "low")
+            except Exception:
+                return "".join(_ZH_DIGIT[c] for c in d)
+        text = re.sub(r"\d+", repl, text)
+        text = text.replace("，", ",").replace("。", ".").replace("？", "?").replace("！", "!")
+        return text
+    return norm
+
+def _lazy():
+    global _g2pw, _g2pen, _zh_num
+    if _g2pw is None:
+        _g2pw = G2PWConverter()
+        _g2pen = G2p()
+        _zh_num = _make_zh_normalizer()
+
+def _split_syllable(syl: str):
+    """'ㄓㄨㄢ3' -> (['ㄓ','ㄨ','ㄢ'], tone 3)."""
+    tone = 0
+    if syl and syl[-1].isdigit():
+        tone = int(syl[-1]); syl = syl[:-1]
+    units = [c for c in syl if c in SYM2ID]
+    return units, tone
+
+def text_to_phones(text: str):
+    _lazy()
+    text = _zh_num(text)                              # numbers->zh words, normalize punct
+    bopo = _g2pw(text)[0]                             # per-char bopomofo or None
+    chars = list(text)
+    phones, tones, langs = [], [], []
+    i = 0
+    while i < len(chars):
+        b = bopo[i] if i < len(bopo) else None
+        ch = chars[i]
+        if b is not None:                             # zh char
+            units, tone = _split_syllable(b)
+            for u in units:
+                phones.append(u); tones.append(min(tone, 5)); langs.append(LANG['ZH'])
+            i += 1
+        elif re.match(r'[A-Za-z]', ch):               # English run -> g2p_en
+            j = i
+            while j < len(chars) and re.match(r"[A-Za-z']", chars[j]):
+                j += 1
+            word = ''.join(chars[i:j])
+            for p in _g2pen(word):
+                p = p.strip()
+                if not p:
+                    continue
+                stress = 0
+                if p[-1].isdigit():
+                    stress = int(p[-1]); p = p[:-1]
+                if p in SYM2ID:
+                    phones.append(p); tones.append(stress); langs.append(LANG['EN'])
+            phones.append('SP'); tones.append(0); langs.append(LANG['EN'])
+            i = j
+        else:                                         # punctuation / space / other
+            if ch in PUNCT:
+                phones.append(ch); tones.append(0); langs.append(LANG['ZH'])
+            elif ch.strip() == '':
+                if phones and phones[-1] != 'SP':
+                    phones.append('SP'); tones.append(0); langs.append(LANG['ZH'])
+            i += 1
+    return phones, tones, langs
+
+def text_to_ids(text: str):
+    phones, tones, langs = text_to_phones(text)
+    ids = [SYM2ID.get(p, SYM2ID['UNK']) for p in phones]
+    return {"phones": phones, "phone_ids": ids, "tone_ids": tones, "lang_ids": langs}

scripts/gen_breezy_corpus.py

@@ -0,0 +1,124 @@
+#!/usr/bin/env python3
+"""Generate the BreezyVoice teacher corpus for distillation into Inflect-Nano.
+Clean short reference (no prompt-leak) + INLINE ASR GATE: each clip is transcribed
+(faster-whisper) and kept only if Han-CER vs intended text is below threshold; else
+retried, then skipped. Writes 22.05kHz wav + manifest {id,text,wav,dur,cer}. Resumable.
+Run in .venv-breezy with PYTHONPATH=BreezyVoice repo.
+"""
+from __future__ import annotations
+import argparse, json, os, sys, time
+import soundfile as sf, numpy as np
+
+BV = "/home/luigi/jetson-tts/third_party/BreezyVoice"
+ZT = "/home/luigi/jetson-tts/mossnano/zhtw8k"
+sys.path.insert(0, BV)
+from single_inference import CustomCosyVoice, get_bopomofo_rare
+from g2pw import G2PWConverter
+from cosyvoice.utils.file_utils import load_wav
+from faster_whisper import WhisperModel
+
+REF_AUDIO = f"{ZT}/ref/ref_clean.wav"
+REF_TEXT = open(f"{ZT}/ref/ref_clean.txt").read().strip()
+
+
+import re as _re
+import opencc as _opencc
+_T2S = _opencc.OpenCC('t2s')  # normalize trad<->simp so the gate scores REAL pronunciation, not script
+def han(s): return "".join(c for c in s if "一" <= c <= "鿿")
+def is_zh(text): return bool(_re.search(r"[一-鿿]", text))
+def _lev(r, h):
+    if not r: return 0.0
+    d = list(range(len(h)+1))
+    for i in range(1, len(r)+1):
+        prev = d[0]; d[0] = i
+        for j in range(1, len(h)+1):
+            cur = d[j]; d[j] = min(d[j]+1, d[j-1]+1, prev+(r[i-1] != h[j-1])); prev = cur
+    return d[len(h)]/len(r)
+def _enwords(s): return _re.findall(r"[a-z']+", s.lower())
+def score(ref, hyp):
+    """Han-CER for zh/mix; word-error-rate for pure-English. Lower = better."""
+    if is_zh(ref):
+        return _lev(han(_T2S.convert(ref)), han(_T2S.convert(hyp)))
+    return _lev(_enwords(ref), _enwords(hyp))   # word-level Levenshtein ratio
+
+
+def read_tsv(path, limit, skip):
+    rows = []
+    for line in open(path, encoding="utf-8"):
+        line = line.rstrip("\n")
+        if not line.strip(): continue
+        parts = line.split("\t")
+        rows.append((parts[0] if len(parts) > 1 else f"utt{len(rows):06d}", parts[-1].strip()))
+    rows = rows[skip:]
+    return rows[:limit] if limit else rows
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--model", default="/home/luigi/jetson-tts/models/BreezyVoice")
+    ap.add_argument("--corpus", default="/home/luigi/jetson-tts/data/text/train.tsv")
+    ap.add_argument("--out-dir", required=True)
+    ap.add_argument("--limit", type=int, default=0)
+    ap.add_argument("--skip", type=int, default=0)
+    ap.add_argument("--min-sec", type=float, default=0.8)
+    ap.add_argument("--max-sec", type=float, default=20.0)
+    ap.add_argument("--cer-thresh", type=float, default=0.30)
+    ap.add_argument("--retries", type=int, default=3)
+    args = ap.parse_args()
+
+    os.makedirs(args.out_dir, exist_ok=True)
+    man_path = os.path.join(args.out_dir, "manifest.jsonl")
+    done = set()
+    if os.path.exists(man_path):
+        for l in open(man_path):
+            try: done.add(json.loads(l)["id"])
+            except Exception: pass
+    print(f"resuming: {len(done)} done | ref: {REF_TEXT}")
+
+    cv = CustomCosyVoice(args.model); conv = G2PWConverter()
+    asr = WhisperModel("SoybeanMilk/faster-whisper-Breeze-ASR-25", device="cuda", compute_type="float16")  # zh-TW gate (traditional output)
+    ref_bopo = get_bopomofo_rare(cv.frontend.text_normalize_new(REF_TEXT, split=False), conv)
+    ref_wav = load_wav(REF_AUDIO, 16000)
+
+    rows = read_tsv(args.corpus, args.limit, args.skip)
+    todo = [(u, t) for u, t in rows if u not in done]
+    print(f"to synth: {len(todo)} / {len(rows)}")
+    mf = open(man_path, "a", encoding="utf-8")
+    t0 = time.time(); n_ok = 0; n_skip = 0; tot = 0.0
+    for i, (utt, text) in enumerate(todo):
+        bopo = get_bopomofo_rare(cv.frontend.text_normalize_new(text, split=False), conv)
+        best = None; best_cer = 9.9
+        for attempt in range(args.retries):
+            try:
+                out = cv.inference_zero_shot_no_normalize(bopo, ref_bopo, ref_wav)
+                w = out["tts_speech"].squeeze().cpu().numpy().astype(np.float32)
+            except Exception as e:
+                print(f"  [err {utt}] {e}"); continue
+            dur = len(w) / 22050
+            if dur < args.min_sec or dur > args.max_sec: continue
+            tmp = os.path.join(args.out_dir, f".{utt}.tmp.wav"); sf.write(tmp, w, 22050)
+            segs, _ = asr.transcribe(tmp, language=("zh" if is_zh(text) else "en"), beam_size=5)
+            c = score(text, "".join(s.text for s in segs))
+            if c < best_cer: best_cer = c; best = (w, dur, tmp)
+            if c <= args.cer_thresh: break
+        if best is None or best_cer > args.cer_thresh:
+            n_skip += 1
+            if best and os.path.exists(best[2]): os.remove(best[2])
+            if n_skip <= 20 or n_skip % 50 == 0:
+                print(f"  [SKIP {utt}] best_cer={best_cer:.2f} | {text[:30]}")
+            continue
+        w, dur, tmp = best
+        wp = os.path.join(args.out_dir, f"{utt}.wav"); os.replace(tmp, wp)
+        mf.write(json.dumps({"id": utt, "text": text, "wav": wp, "dur": round(dur, 3),
+                             "cer": round(best_cer, 3)}, ensure_ascii=False) + "\n"); mf.flush()
+        n_ok += 1; tot += dur
+        if (i + 1) % 25 == 0:
+            el = time.time() - t0
+            print(f"  {i+1}/{len(todo)} ok={n_ok} skip={n_skip} audio={tot/60:.1f}min "
+                  f"{el/(i+1):.2f}s/clip eta={(len(todo)-i-1)*el/(i+1)/60:.0f}min")
+    mf.close()
+    print(f"DONE ok={n_ok} skip={n_skip} audio={tot/60:.1f}min in {(time.time()-t0)/60:.0f}min")
+
+
+if __name__ == "__main__":
+    main()

scripts/probe_forced_synth.py

@@ -0,0 +1,50 @@
+#!/usr/bin/env python3
+"""RANK-1 acoustic probe: synth N training clips with FORCED (ground-truth) durations through an
+ONNX dir, isolating the phones->mel mapping from the duration predictor + g2p frontend.
+Reads phone/tone/lang ids + GT durations directly from m2_align.jsonl (no frontend).
+Run in moss-train-venv. Then ASR the wavs (probe_forced_asr via xasr_offline) -> CER.
+Pairs the 0.80(forced)-vs-0.145(GT-mel) acoustic gap to a single number per config."""
+import argparse, json, sys
+from pathlib import Path
+import numpy as np, soundfile as sf, onnxruntime as ort
+ZT = "/home/luigi/jetson-tts/mossnano/zhtw8k"
+sys.path.insert(0, ZT)
+from synth_from_text import host_regulate
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--onnx-dir", required=True)
+    ap.add_argument("--out-dir", required=True)
+    ap.add_argument("--n", type=int, default=30)
+    ap.add_argument("--align-jsonl", default=f"{ZT}/m2_align.jsonl")
+    args = ap.parse_args()
+    meta = json.load(open(f"{args.onnx_dir}/meta.json"))
+    so = ort.SessionOptions(); so.intra_op_num_threads = 4
+    sA = ort.InferenceSession(f"{args.onnx_dir}/acoustic_encoder.onnx", so, providers=["CPUExecutionProvider"])
+    sB = ort.InferenceSession(f"{args.onnx_dir}/acoustic_decoder.onnx", so, providers=["CPUExecutionProvider"])
+    sV = ort.InferenceSession(f"{args.onnx_dir}/vocoder.onnx", so, providers=["CPUExecutionProvider"])
+    bn = ["frames", "frame_meta", "local_ctx_raw", "abs_pos", "pitch_frame", "frame_mask"]
+    Path(args.out_dir).mkdir(parents=True, exist_ok=True)
+    rows = [json.loads(l) for l in open(args.align_jsonl) if l.strip()][:args.n]
+    out = open(f"{args.out_dir}/synth.jsonl", "w")
+    for i, r in enumerate(rows):
+        phone = np.array([r["phone_ids"]], np.int64); tone = np.array([r["tone_ids"]], np.int64)
+        lang = np.array([r["lang_ids"]], np.int64); spk = np.zeros(1, np.int64)
+        cond, _dur_pred, pitch = sA.run(None, {"phone": phone, "tone": tone, "lang": lang, "speaker": spk})
+        # substitute GT (forced) durations, rescaled so total ~ predicted length (stable regulator)
+        df = np.array([r["hifigan_durations"]], np.float32)
+        df = df * (_dur_pred.sum() / max(1.0, df.sum()))
+        reg = host_regulate(cond, df, pitch, meta["abs_frame_bins"], meta["max_frames"])
+        feeds = {n: (reg[n].astype(np.float32) if reg[n].dtype != bool else reg[n]) for n in bn}
+        feeds["abs_pos"] = reg["abs_pos"].astype(np.int64)
+        mel = sB.run(None, feeds)[0]
+        wav = sV.run(None, {"mel": mel.astype(np.float32)})[0].reshape(-1)
+        wp = f"{args.out_dir}/p{i:03d}.wav"; sf.write(wp, wav, meta["sample_rate"])
+        out.write(json.dumps({"id": f"p{i:03d}", "text": r["text"], "wav": wp}, ensure_ascii=False) + "\n")
+    out.close()
+    print(f"PROBE SYNTH DONE {len(rows)} clips -> {args.out_dir}/synth.jsonl")
+
+
+if __name__ == "__main__":
+    main()

scripts/run_bilingual.sh

@@ -0,0 +1,55 @@
+#!/usr/bin/env bash
+# BILINGUAL NATIVE experiment: apply the proven recipe (v4 phone-align + diverse data) symmetrically.
+# Merge zh-expanded corpus (15423, mark voice, incl code-mix) + diverse English-expanded (mark voice)
+# -> one balanced single-voice bilingual corpus -> train ONE 4.63M model good at zh AND en (+code-mix).
+# Run when English synth done + a GPU free. GPU via $1 (default 0).
+set -e
+ZT=/home/luigi/jetson-tts/mossnano/zhtw8k; LOG=$ZT/bilingual.log
+TRAIN=/home/luigi/moss-train-venv/bin/python
+GPU=${1:-0}
+VOC=$ZT/m2_vocoder_8k/hifigan-snake_8k-latest.pt
+exec >>"$LOG" 2>&1
+echo "===== BILINGUAL START $(date) GPU=$GPU ====="
+
+# 1) canonical-fix English-expand manifest wav paths from id
+$TRAIN - <<PY
+import json, os
+p="$ZT/teacher_corpus_en_expand/manifest.jsonl"; D="$ZT/teacher_corpus_en_expand"
+rows=[json.loads(l) for l in open(p) if l.strip()]
+ok=0
+for r in rows: r["wav"]=f"{D}/{r['id']}.wav"; ok+=os.path.exists(r["wav"])
+open(p,"w").write("\n".join(json.dumps(r,ensure_ascii=False) for r in rows)+"\n")
+print(f"[en-expand manifest] {len(rows)} rows, {ok} files exist")
+PY
+
+# 2) v4 phone-align the English-expand clips (espeak handles English natively)
+echo "[align en-expand] $(date)"
+CUDA_VISIBLE_DEVICES=$GPU $TRAIN $ZT/align_durations_v4.py \
+  --manifest $ZT/teacher_corpus_en_expand/manifest.jsonl \
+  --out $ZT/en_expand_v4_align.jsonl --device cuda
+
+# 3) merge zh-expanded + en-expanded -> bilingual corpus (all mark voice, speaker 0)
+$TRAIN - <<PY
+import json
+def fix(p):
+    rows=[json.loads(l) for l in open(p) if l.strip()]
+    for r in rows:
+        if not r["target_audio"].startswith("/"): r["target_audio"]="/home/luigi/jetson-tts/"+r["target_audio"]
+    return rows
+a=fix("$ZT/m_v4_v2_align.jsonl"); b=fix("$ZT/en_expand_v4_align.jsonl")
+out="$ZT/m_bilingual_align.jsonl"
+open(out,"w").write("\n".join(json.dumps(r,ensure_ascii=False) for r in a+b)+"\n")
+import re
+en=sum(1 for r in a+b if all(ord(c)<128 for c in r["text"]))
+print(f"[bilingual] {len(a)} + {len(b)} = {len(a)+len(b)} rows -> {out} | ~{en} pure-ASCII(en) rows")
+PY
+
+# 4) train the bilingual model (en-upsample 1: English now abundant via expansion)
+echo "[bilingual acoustic] $(date)"
+cd /tmp/inflect-nano
+CUDA_VISIBLE_DEVICES=$GPU $TRAIN -m inflect_nano.acoustic --durations-jsonl $ZT/m_bilingual_align.jsonl \
+  --out-dir $ZT/bili_acoustic_8k --vocoder-variant snake_8k --sample-rate 8000 \
+  --steps 60000 --batch-size 16 --lr 2e-4 --max-frames 1000 --en-upsample 1 \
+  --vocoder-checkpoint $VOC --vocoder-mel-weight 1.0 \
+  --save-interval 5000 --log-interval 200 --device cuda
+echo "===== BILINGUAL TRAIN DONE $(date) ====="

scripts/run_v4zh_v2.sh

@@ -0,0 +1,54 @@
+#!/usr/bin/env bash
+# v4zh-v2: retrain on DIVERSITY-EXPANDED corpus (existing 9042 + ~6500 broad-char Tatoeba/zh-TW).
+# Tests the data-diversity hypothesis: held-out zh should drop from 0.40 toward the 0.21 seen-text floor.
+# Keeps the proven v4 phone-level alignment. Run when teacher synth + GPU are free. GPU via $1 (default 1).
+set -e
+ZT=/home/luigi/jetson-tts/mossnano/zhtw8k; LOG=$ZT/v4zh_v2.log
+TRAIN=/home/luigi/moss-train-venv/bin/python
+GPU=${1:-1}
+VOC=$ZT/m2_vocoder_8k/hifigan-snake_8k-latest.pt
+exec >>"$LOG" 2>&1
+echo "===== V4ZH_V2 START $(date) GPU=$GPU ====="
+
+# 1) canonical-fix expansion manifest wav paths from id (idempotent, no doubling)
+$TRAIN - <<PY
+import json, os
+p="$ZT/teacher_corpus_expand/manifest.jsonl"; D="$ZT/teacher_corpus_expand"
+rows=[json.loads(l) for l in open(p) if l.strip()]
+ok=0
+for r in rows:
+    r["wav"]=f"{D}/{r['id']}.wav"; ok+=os.path.exists(r["wav"])
+open(p,"w").write("\n".join(json.dumps(r,ensure_ascii=False) for r in rows)+"\n")
+print(f"[expand manifest] {len(rows)} rows, {ok} files exist")
+PY
+
+# 2) v4-align the expansion clips (phone-level)
+echo "[align expand] $(date)"
+CUDA_VISIBLE_DEVICES=$GPU $TRAIN $ZT/align_durations_v4.py \
+  --manifest $ZT/teacher_corpus_expand/manifest.jsonl \
+  --out $ZT/m_v4_expand_align.jsonl --device cuda
+
+# 3) abspath-fix + merge with existing v4 alignment
+$TRAIN - <<PY
+import json
+R="/home/luigi/jetson-tts/"
+def fix(p):
+    rows=[json.loads(l) for l in open(p) if l.strip()]
+    for r in rows:
+        if not r["target_audio"].startswith("/"): r["target_audio"]=R+r["target_audio"]
+    return rows
+a=fix("$ZT/m_v4_align.jsonl"); b=fix("$ZT/m_v4_expand_align.jsonl")
+out="$ZT/m_v4_v2_align.jsonl"
+open(out,"w").write("\n".join(json.dumps(r,ensure_ascii=False) for r in a+b)+"\n")
+print(f"[merged] {len(a)} + {len(b)} = {len(a)+len(b)} -> {out}")
+PY
+
+# 4) retrain on merged corpus (same recipe, frozen arch)
+echo "[v4zh-v2 acoustic] $(date)"
+cd /tmp/inflect-nano
+CUDA_VISIBLE_DEVICES=$GPU $TRAIN -m inflect_nano.acoustic --durations-jsonl $ZT/m_v4_v2_align.jsonl \
+  --out-dir $ZT/v4zh_v2_acoustic_8k --vocoder-variant snake_8k --sample-rate 8000 \
+  --steps 60000 --batch-size 16 --lr 2e-4 --max-frames 1000 --en-upsample 2 \
+  --vocoder-checkpoint $VOC --vocoder-mel-weight 1.0 \
+  --save-interval 5000 --log-interval 200 --device cuda
+echo "===== V4ZH_V2 TRAIN DONE $(date) ====="

scripts/run_voc_retrain.sh

@@ -0,0 +1,14 @@
+#!/usr/bin/env bash
+# Retrain the 8k Snake-HiFiGAN vocoder on the FULL diverse corpus (~21k, zh+en mark voice),
+# longer + higher MR-STFT weight -> crisper, less-buzzy 8kHz audio (raises the GT-mel->audio ceiling).
+set -e
+ZT=/home/luigi/jetson-tts/mossnano/zhtw8k; LOG=$ZT/voc_retrain.log
+TRAIN=/home/luigi/moss-train-venv/bin/python
+exec >>"$LOG" 2>&1
+echo "===== VOC RETRAIN START $(date) ====="
+cd /tmp/inflect-nano
+CUDA_VISIBLE_DEVICES=1 $TRAIN -m inflect_nano.vocoder --train-jsonl $ZT/bili_voc_rows.jsonl \
+  --out-dir $ZT/bili_vocoder_8k --variant snake_8k --steps 40000 --batch-size 16 \
+  --segment-size 8192 --min-seconds 0.8 --max-seconds 20 --num-workers 4 --stft-weight 2.5 \
+  --save-interval 5000 --log-interval 200 --device cuda
+echo "===== VOC RETRAIN DONE $(date) ====="

scripts/select_diverse_text.py

@@ -0,0 +1,107 @@
+#!/usr/bin/env python3
+"""Select a DIVERSE training-text set to maximize phoneme/character coverage.
+
+Why this matters: at 4.63M params the model is NOT capacity-limited — but it can only
+pronounce characters/words it has SEEN. A narrow corpus (e.g. a few hundred Han chars)
+leaves most held-out characters unseen -> garbled output. This script builds a broad,
+coverage-maximizing text set from Tatoeba so the teacher (and then the student) cover
+the common vocabulary.
+
+zh-TW: Tatoeba `cmn` -> OpenCC s2twp (Taiwan traditional) -> greedy Han-CHAR coverage.
+en   : Tatoeba `eng` -> greedy WORD coverage (English phones are few; word/prosody variety matters).
+
+Usage:
+  python select_diverse_text.py --lang zh --n 6000 --out expand_zh.tsv
+  python select_diverse_text.py --lang en --n 6000 --out expand_en.tsv
+Then feed the .tsv (id<TAB>text) to gen_breezy_corpus.py to synthesize the teacher audio.
+
+Deps: requests/urllib (download), opencc (zh only). Tatoeba dumps are CC-BY 2.0 FR.
+"""
+import argparse, bz2, os, re, random, urllib.request
+
+TATOEBA = "https://downloads.tatoeba.org/exports/per_language/{lang}/{lang}_sentences.tsv.bz2"
+HAN = lambda s: set(c for c in s if "一" <= c <= "鿿")
+
+
+def download(lang):
+    f = f"{lang}_sentences.tsv"
+    if not os.path.exists(f):
+        url = TATOEBA.format(lang=lang)
+        print("downloading", url)
+        urllib.request.urlretrieve(url, f + ".bz2")
+        with bz2.open(f + ".bz2", "rt", encoding="utf-8") as i, open(f, "w", encoding="utf-8") as o:
+            for line in i:
+                o.write(line)
+    return f
+
+
+def select_zh(path, n, seed=42):
+    import opencc
+    cc = opencc.OpenCC("s2twp")  # simplified -> Taiwan traditional (with phrase conversion)
+    allowed = set("，。！？、：；…")
+    seen_t, cands = set(), []
+    for l in open(path, encoding="utf-8"):
+        p = l.rstrip("\n").split("\t")
+        if len(p) < 3:
+            continue
+        t = cc.convert(p[2].strip()).replace(",", "，").replace("!", "！").replace("?", "？")
+        h = HAN(t)
+        if not (6 <= len(h) <= 26):
+            continue
+        if any(("一" <= c <= "鿿") is False and c not in allowed for c in t):
+            continue
+        if t not in seen_t:
+            seen_t.add(t); cands.append(t)
+    return greedy_cover(cands, HAN, n, seed)
+
+
+def select_en(path, n, seed=42):
+    words = lambda s: set(re.findall(r"[a-z']+", s.lower()))
+    seen_t, cands = set(), []
+    for l in open(path, encoding="utf-8"):
+        p = l.rstrip("\n").split("\t")
+        if len(p) < 3:
+            continue
+        t = p[2].strip()
+        if not re.fullmatch(r"[A-Za-z0-9 ,.\-'?!]+", t):
+            continue
+        w = re.findall(r"[A-Za-z']+", t)
+        if not (4 <= len(w) <= 14) or any(len(x) > 15 for x in w):
+            continue
+        if t not in seen_t:
+            seen_t.add(t); cands.append(t)
+    return greedy_cover(cands, words, n, seed)
+
+
+def greedy_cover(cands, unit, n, seed):
+    """Greedy max-coverage of `unit(text)` items, then random top-up to n for frequency."""
+    random.seed(seed); random.shuffle(cands)
+    covered, selected, rest = set(), [], []
+    cands.sort(key=lambda t: len(unit(t) - covered), reverse=True)
+    for t in cands:
+        if len(unit(t) - covered) >= 1 and len(selected) < n:
+            selected.append(t); covered |= unit(t)
+        else:
+            rest.append(t)
+    random.shuffle(rest)
+    selected += rest[: max(0, n - len(selected))]
+    print(f"selected {len(selected)} sentences | unique units covered: {len(covered)}")
+    return selected
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--lang", choices=["zh", "en"], required=True)
+    ap.add_argument("--n", type=int, default=6000)
+    ap.add_argument("--out", required=True)
+    args = ap.parse_args()
+    path = download("cmn" if args.lang == "zh" else "eng")
+    sents = select_zh(path, args.n) if args.lang == "zh" else select_en(path, args.n)
+    with open(args.out, "w", encoding="utf-8") as o:
+        for i, t in enumerate(sents):
+            o.write(f"{args.lang}e{i:05d}\t{t}\n")
+    print("wrote", args.out, len(sents))
+
+
+if __name__ == "__main__":
+    main()

scripts/symbol_table.json

@@ -0,0 +1,186 @@
+{
+ "symbols": [
+  "_blank",
+  "_pad",
+  "UNK",
+  "SP",
+  "ㄅ",
+  "ㄆ",
+  "ㄇ",
+  "ㄈ",
+  "ㄉ",
+  "ㄊ",
+  "ㄋ",
+  "ㄌ",
+  "ㄍ",
+  "ㄎ",
+  "ㄏ",
+  "ㄐ",
+  "ㄑ",
+  "ㄒ",
+  "ㄓ",
+  "ㄔ",
+  "ㄕ",
+  "ㄖ",
+  "ㄗ",
+  "ㄘ",
+  "ㄙ",
+  "ㄚ",
+  "ㄛ",
+  "ㄜ",
+  "ㄝ",
+  "ㄞ",
+  "ㄟ",
+  "ㄠ",
+  "ㄡ",
+  "ㄢ",
+  "ㄣ",
+  "ㄤ",
+  "ㄥ",
+  "ㄦ",
+  "ㄧ",
+  "ㄨ",
+  "ㄩ",
+  "AA",
+  "AE",
+  "AH",
+  "AO",
+  "AW",
+  "AY",
+  "B",
+  "CH",
+  "D",
+  "DH",
+  "EH",
+  "ER",
+  "EY",
+  "F",
+  "G",
+  "HH",
+  "IH",
+  "IY",
+  "JH",
+  "K",
+  "L",
+  "M",
+  "N",
+  "NG",
+  "OW",
+  "OY",
+  "P",
+  "R",
+  "S",
+  "SH",
+  "T",
+  "TH",
+  "UH",
+  "UW",
+  "V",
+  "W",
+  "Y",
+  "Z",
+  "ZH",
+  ",",
+  ".",
+  "?",
+  "!",
+  "…",
+  "-",
+  "'"
+ ],
+ "sym2id": {
+  "_blank": 0,
+  "_pad": 1,
+  "UNK": 2,
+  "SP": 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,
+  "AA": 41,
+  "AE": 42,
+  "AH": 43,
+  "AO": 44,
+  "AW": 45,
+  "AY": 46,
+  "B": 47,
+  "CH": 48,
+  "D": 49,
+  "DH": 50,
+  "EH": 51,
+  "ER": 52,
+  "EY": 53,
+  "F": 54,
+  "G": 55,
+  "HH": 56,
+  "IH": 57,
+  "IY": 58,
+  "JH": 59,
+  "K": 60,
+  "L": 61,
+  "M": 62,
+  "N": 63,
+  "NG": 64,
+  "OW": 65,
+  "OY": 66,
+  "P": 67,
+  "R": 68,
+  "S": 69,
+  "SH": 70,
+  "T": 71,
+  "TH": 72,
+  "UH": 73,
+  "UW": 74,
+  "V": 75,
+  "W": 76,
+  "Y": 77,
+  "Z": 78,
+  "ZH": 79,
+  ",": 80,
+  ".": 81,
+  "?": 82,
+  "!": 83,
+  "…": 84,
+  "-": 85,
+  "'": 86
+ },
+ "lang": {
+  "ZH": 0,
+  "EN": 1
+ },
+ "n_symbols": 87,
+ "max_tone": 5
+}

scripts/synth_from_text.py

@@ -0,0 +1,68 @@
+#!/usr/bin/env python3
+"""End-to-end synthesis from text via the exported 8k ONNX pipeline:
+text -> bopomofo+arpabet frontend -> ids -> encoder.onnx -> numpy host_regulate
+-> decoder.onnx -> vocoder.onnx -> 8kHz wav. Run in moss-train-venv (g2pw+ort).
+Used for M1 eval (synthesize zh-TW/en/code-mix test sentences). X-ASR scoring is a
+separate step in moss-nano-venv (xasr_offline.py) on the produced wavs."""
+from __future__ import annotations
+import argparse, json, sys
+from pathlib import Path
+import numpy as np, soundfile as sf, onnxruntime as ort
+
+ZT = "/home/luigi/jetson-tts/mossnano/zhtw8k"
+sys.path.insert(0, ZT)
+import frontend_bopomofo as F  # g2pw bopomofo + g2p_en arpabet -> ids
+
+
+def host_regulate(cond, dur, pitch, abs_bins, max_frames):
+    c = cond[0]; d = dur[0].astype(np.int64); d[d < 0] = 0
+    T, H = c.shape
+    frames = np.repeat(c, d, axis=0); Fn = frames.shape[0]
+    tok = np.repeat(np.arange(T), d); starts = np.cumsum(d) - d
+    within = np.arange(Fn) - starts[tok]; dpf = d[tok].astype(np.float32)
+    rel = (within / np.maximum(dpf - 1, 1)).astype(np.float32)
+    tc = max(1, int((d > 0).sum())); token_pos = (tok / max(1, tc - 1)).astype(np.float32)
+    ld = (np.log1p(dpf) / 6.0).astype(np.float32); center = 1.0 - np.abs(rel * 2 - 1)
+    fm = np.stack([rel, 1 - rel, center, np.sin(rel*np.pi), np.cos(rel*np.pi), token_pos, ld, dpf/40.0], -1).astype(np.float32)
+    prev = np.concatenate([c[:1], c[:-1]], 0); nxt = np.concatenate([c[1:], c[-1:]], 0)
+    lc = np.repeat(np.concatenate([prev, c, nxt], -1), d, axis=0).astype(np.float32)
+    pos = np.arange(Fn); ap = np.minimum(pos*abs_bins//max(1, max_frames), abs_bins-1).astype(np.int64)
+    pf = np.repeat(pitch[0], d, axis=0).astype(np.float32)
+    return {"frames": frames[None].astype(np.float32), "frame_meta": fm[None], "local_ctx_raw": lc[None],
+            "abs_pos": ap[None], "pitch_frame": pf[None], "frame_mask": np.ones((1, Fn), bool)}
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--onnx-dir", required=True)
+    ap.add_argument("--out-dir", required=True)
+    ap.add_argument("--texts", required=True, help="jsonl with {id,text}")
+    args = ap.parse_args()
+    meta = json.load(open(f"{args.onnx_dir}/meta.json"))
+    so = ort.SessionOptions(); so.intra_op_num_threads = 4
+    sA = ort.InferenceSession(f"{args.onnx_dir}/acoustic_encoder.onnx", so, providers=["CPUExecutionProvider"])
+    sB = ort.InferenceSession(f"{args.onnx_dir}/acoustic_decoder.onnx", so, providers=["CPUExecutionProvider"])
+    sV = ort.InferenceSession(f"{args.onnx_dir}/vocoder.onnx", so, providers=["CPUExecutionProvider"])
+    Path(args.out_dir).mkdir(parents=True, exist_ok=True)
+    sr = meta["sample_rate"]; bn = ["frames","frame_meta","local_ctx_raw","abs_pos","pitch_frame","frame_mask"]
+    rows = [json.loads(l) for l in open(args.texts) if l.strip()]
+    out_manifest = open(f"{args.out_dir}/synth.jsonl", "w")
+    for r in rows:
+        o = F.text_to_ids(r["text"])
+        phone = np.array([o["phone_ids"]], np.int64); tone = np.array([o["tone_ids"]], np.int64); lang = np.array([o["lang_ids"]], np.int64)
+        spk = np.zeros(1, np.int64)
+        cond, dur, pitch = sA.run(None, {"phone": phone, "tone": tone, "lang": lang, "speaker": spk})
+        reg = host_regulate(cond, dur, pitch, meta["abs_frame_bins"], meta["max_frames"])
+        feeds = {n: (reg[n].astype(np.float32) if reg[n].dtype != bool else reg[n]) for n in bn}
+        feeds["abs_pos"] = reg["abs_pos"].astype(np.int64)
+        mel = sB.run(None, feeds)[0]
+        wav = sV.run(None, {"mel": mel.astype(np.float32)})[0].reshape(-1)
+        wp = f"{args.out_dir}/{r['id']}.wav"; sf.write(wp, wav, sr)
+        out_manifest.write(json.dumps({"id": r["id"], "text": r["text"], "wav": wp, "dur": round(len(wav)/sr, 2)}, ensure_ascii=False) + "\n")
+        print(f"  {r['id']}: {len(wav)/sr:.1f}s -> {wp}")
+    out_manifest.close()
+    print(f"DONE synth -> {args.out_dir}/synth.jsonl")
+
+
+if __name__ == "__main__":
+    main()

scripts/xasr_offline.py

@@ -0,0 +1,56 @@
+"""Offline X-ASR (zh-en zipformer transducer, 2026-06-03) scorer — accurate code-mix
+assessment. Reuses the model's own raw-ort decode (int64-patched) + OpenCC T2S norm.
+Run in moss-nano-venv (onnxruntime 1.27 + kaldi_native_fbank + opencc + soundfile/librosa)."""
+import os, re, sys
+MDIR="/home/luigi/jetson-tts/models/xasr-offline/sherpa-onnx-x-asr-zipformer-transducer-zh-en-2026-06-03"
+_cwd=os.getcwd(); os.chdir(MDIR); sys.path.insert(0, MDIR)
+import importlib.util
+_spec=importlib.util.spec_from_file_location("xasr_test", f"{MDIR}/test_onnx.py")
+T=importlib.util.module_from_spec(_spec); _spec.loader.exec_module(T)
+os.chdir(_cwd)
+import opencc, cn2an
+_ZD={"0":"零","1":"一","2":"二","3":"三","4":"四","5":"五","6":"六","7":"七","8":"八","9":"九"}
+def _numnorm(t):
+    def r(m):
+        d=m.group(0)
+        if len(d)>=5: return "".join(_ZD[c] for c in d)
+        try: return cn2an.an2cn(d,"low")
+        except: return "".join(_ZD[c] for c in d)
+    return re.sub(r"\d+",r,t)
+_t2s=opencc.OpenCC('t2s'); _model=None; _id2tok=None
+def _ensure():
+    global _model,_id2tok
+    if _model is None:
+        cwd=os.getcwd(); os.chdir(MDIR)
+        _id2tok=T.load_tokens("./tokens.txt"); _model=T.load_model(use_int8=False)
+        os.chdir(cwd)
+    return _model,_id2tok
+def asr(wav):
+    import soundfile as sf, librosa, numpy as np
+    m,id2tok=_ensure()
+    samples,sr=sf.read(wav,dtype="float32")
+    if samples.ndim>1: samples=samples.mean(1)
+    if sr!=16000: samples=librosa.resample(samples,orig_sr=sr,target_sr=16000,res_type="soxr_hq")
+    if len(samples)<64000: samples=np.concatenate([samples,np.zeros(64000-len(samples),np.float32)])  # >=4s pad: short clips break encoder convs
+    feats=T.compute_feat(samples=samples.astype(np.float32), sample_rate=16000)
+    blank=0; hyp=[blank]*m.context_size
+    dout=m.run_decoder(hyp); enc=m.run_encoder(feats[None])
+    for k in range(enc.shape[1]):
+        jo=m.run_joiner(enc[0,k:k+1], dout); tid=int(jo.argmax())
+        if tid!=blank:
+            hyp.append(tid); dout=m.run_decoder(hyp[-m.context_size:])
+    toks=[id2tok[i] for i in hyp[m.context_size:]]
+    return "".join(toks).replace("▁"," ").strip()
+def han(s): return "".join(c for c in s if "一"<=c<="鿿")
+def _lev(r,h):
+    if not r: return 0.0
+    d=list(range(len(h)+1))
+    for i in range(1,len(r)+1):
+        p=d[0]; d[0]=i
+        for j in range(1,len(h)+1):
+            c=d[j]; d[j]=min(d[j]+1,d[j-1]+1,p+(r[i-1]!=h[j-1])); p=c
+    return d[len(h)]/len(r)
+def score(ref,hyp):
+    is_zh=bool(re.search(r"[一-鿿]",ref))
+    if is_zh: return _lev(han(_t2s.convert(_numnorm(ref))),han(_t2s.convert(hyp)))
+    return _lev(re.findall(r"[a-z']+",ref.lower()),re.findall(r"[a-z']+",hyp.lower()))

weights/acoustic_decoder.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6ea97d8caf52af44a6372aa6cb445de6b70690561d02d2cf8514c019fbad91bc
+size 8147039

weights/acoustic_encoder.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:605048511cd3c8ab429768a8b11f7e2d45705379c9e6ff91200e88ab8664c18f
+size 5782030

weights/acoustic_zh_v2_60k.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:559970e6a52d5ea9b7aaecf592835f8c589fe8b706223125a62a5e4c618ed670
+size 41829011

weights/meta.json

@@ -0,0 +1,8 @@
+{
+ "sample_rate": 8000,
+ "abs_frame_bins": 512,
+ "max_frames": 1000,
+ "hop_size": 128,
+ "n_mels": 80,
+ "use_frame_pitch": true
+}

weights/symbol_table.json

@@ -0,0 +1,186 @@
+{
+ "symbols": [
+  "_blank",
+  "_pad",
+  "UNK",
+  "SP",
+  "ㄅ",
+  "ㄆ",
+  "ㄇ",
+  "ㄈ",
+  "ㄉ",
+  "ㄊ",
+  "ㄋ",
+  "ㄌ",
+  "ㄍ",
+  "ㄎ",
+  "ㄏ",
+  "ㄐ",
+  "ㄑ",
+  "ㄒ",
+  "ㄓ",
+  "ㄔ",
+  "ㄕ",
+  "ㄖ",
+  "ㄗ",
+  "ㄘ",
+  "ㄙ",
+  "ㄚ",
+  "ㄛ",
+  "ㄜ",
+  "ㄝ",
+  "ㄞ",
+  "ㄟ",
+  "ㄠ",
+  "ㄡ",
+  "ㄢ",
+  "ㄣ",
+  "ㄤ",
+  "ㄥ",
+  "ㄦ",
+  "ㄧ",
+  "ㄨ",
+  "ㄩ",
+  "AA",
+  "AE",
+  "AH",
+  "AO",
+  "AW",
+  "AY",
+  "B",
+  "CH",
+  "D",
+  "DH",
+  "EH",
+  "ER",
+  "EY",
+  "F",
+  "G",
+  "HH",
+  "IH",
+  "IY",
+  "JH",
+  "K",
+  "L",
+  "M",
+  "N",
+  "NG",
+  "OW",
+  "OY",
+  "P",
+  "R",
+  "S",
+  "SH",
+  "T",
+  "TH",
+  "UH",
+  "UW",
+  "V",
+  "W",
+  "Y",
+  "Z",
+  "ZH",
+  ",",
+  ".",
+  "?",
+  "!",
+  "…",
+  "-",
+  "'"
+ ],
+ "sym2id": {
+  "_blank": 0,
+  "_pad": 1,
+  "UNK": 2,
+  "SP": 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,
+  "AA": 41,
+  "AE": 42,
+  "AH": 43,
+  "AO": 44,
+  "AW": 45,
+  "AY": 46,
+  "B": 47,
+  "CH": 48,
+  "D": 49,
+  "DH": 50,
+  "EH": 51,
+  "ER": 52,
+  "EY": 53,
+  "F": 54,
+  "G": 55,
+  "HH": 56,
+  "IH": 57,
+  "IY": 58,
+  "JH": 59,
+  "K": 60,
+  "L": 61,
+  "M": 62,
+  "N": 63,
+  "NG": 64,
+  "OW": 65,
+  "OY": 66,
+  "P": 67,
+  "R": 68,
+  "S": 69,
+  "SH": 70,
+  "T": 71,
+  "TH": 72,
+  "UH": 73,
+  "UW": 74,
+  "V": 75,
+  "W": 76,
+  "Y": 77,
+  "Z": 78,
+  "ZH": 79,
+  ",": 80,
+  ".": 81,
+  "?": 82,
+  "!": 83,
+  "…": 84,
+  "-": 85,
+  "'": 86
+ },
+ "lang": {
+  "ZH": 0,
+  "EN": 1
+ },
+ "n_symbols": 87,
+ "max_tone": 5
+}

weights/vocoder.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fa20f208f9d006f07ad00965809aea72f461e5280f8592412f6d371095a2636f
+size 4705722