Fix Audio object handling in HF datasets (has .array attr, not a plain dict)

- inflect_nano/train_smooth.py: enhanced training with multi-res STFT loss,
  adversarial mel discriminator, vocoder consistency loss, deeper residual
  postnet, and cosine LR schedule
- preprocess_dataset.py: convert HF datasets, local dirs, or LJSpeech CSVs
  to the durations.jsonl format needed by training
- inference.py: add --smooth-prosody, --mel-smooth-sigma, --lowpass-hz flags
  for zero-cost inference-time quality improvements
- test_inference.py: smoke tests for model loading and synthesis
- colab_smooth_finetune.ipynb: Colab notebook for T4 GPU fine-tuning
- requirements.txt: add numba, scipy, datasets

.gitattributes

@@ -1,4 +1,4 @@
 *.pt filter=lfs diff=lfs merge=lfs -text
 *.wav filter=lfs diff=lfs merge=lfs -text
 assets/inflect-nano-banner.png filter=lfs diff=lfs merge=lfs -text
-tiny_tts/text/cmudict_cache.pickle filter=lfs diff=lfs merge=lfs -text
+third_party/tiny_tts_frontend/tiny_tts/text/cmudict_cache.pickle filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,4 @@
+__pycache__/
+*.pyc
+*.wav
+.venv/

README.md

@@ -21,6 +21,8 @@ library_name: pytorch
 
 # Inflect-Nano-v1
 
+Edit 06/17/2026 -- I'm really happy to see that this model is doing decently! If more people find it useful, I might consider training a v2 with a much larger budget. But for now, if you would like to see a v2, just like/favourite this model to get more people see it! Thank you for everyone for checking out this model. 
+
 **Inflect-Nano-v1 is a tiny English text-to-speech model with 4.63M total inference parameters, including its vocoder.**
 
 It is not trying to beat large TTS models. It is a small, local, complete text-to-waveform stack built to test how far ultra-lightweight speech synthesis can go.
@@ -99,6 +101,20 @@ weights/inflect_nano_v1_acoustic.pt
 weights/inflect_nano_v1_vocoder.pt
 ```
 
+## Repo Layout
+
+```text
+weights/                         model weights
+examples/                        audio examples
+assets/                          README banner
+inflect_nano/                    runtime model code
+third_party/tiny_tts_frontend/   vendored text frontend used for English G2P/token IDs
+inference.py                     simple CLI inference
+app.py                           local Gradio demo
+```
+
+The model itself is in `weights/`. The vendored frontend is included only so the released model can reproduce the same text normalization and tokenization path.
+
 ## What Makes It Different
 
 Many small TTS projects depend on a separate larger vocoder. Inflect-Nano-v1 includes the vocoder in the published inference stack, so the full text-to-waveform path stays under 5M parameters.
@@ -159,4 +175,4 @@ Use it as a tiny-model research/demo release, not as a production TTS engine.
 
 Apache-2.0.
 
-This repository includes a small third-party English text frontend for tokenization/G2P compatibility. Its license is included as `TINY_TTS_LICENSE`.
+This repository includes a small third-party English text frontend for tokenization/G2P compatibility. Its license is included at `third_party/tiny_tts_frontend/LICENSE`.

colab_run.py

@@ -0,0 +1,170 @@
+"""
+Colab runner for Inflect-Nano smoothness fine-tuning.
+
+Single script that handles the full pipeline with graceful checkpointing —
+safe to re-run after interruptions.  Each step is skipped if its output
+already exists.
+
+Usage in Colab:
+    !wget https://your-server/colab_run.py
+    !python colab_run.py
+"""
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+import time
+from pathlib import Path
+
+# ---------------------------------------------------------------------------
+# Config — change these
+# ---------------------------------------------------------------------------
+REPO_URL = "https://git.quietidiot.com/michael-treadgold/speech-nano.git"
+REPO_DIR = Path("/content/speech-nano")
+DATASET = "MikhailT/cmu-arctic"
+SPEAKER = "rms"                # "rms", "bdl", "jmk", "awb", "ksp" (CMU ARCTIC male speakers)
+TRAIN_STEPS = 5000
+BATCH_SIZE = 8
+
+# ---------------------------------------------------------------------------
+# Step helpers
+# ---------------------------------------------------------------------------
+def step(name: str) -> None:
+    print(f"\n{'='*60}\n  {name}\n{'='*60}")
+
+def run(cmd: str, **kwargs) -> bool:
+    print(f"  $ {cmd}")
+    result = subprocess.run(cmd, shell=True, **kwargs)
+    return result.returncode == 0
+
+def already(path: Path | str) -> bool:
+    p = Path(path)
+    ok = p.exists()
+    if ok:
+        print(f"  ⏭  Already exists: {p}")
+    return ok
+
+def pip_installed(pkg: str) -> bool:
+    result = subprocess.run(
+        [sys.executable, "-c", f"import {pkg}"],
+        capture_output=True, text=True, timeout=30,
+    )
+    return result.returncode == 0
+
+
+def main() -> None:
+    start = time.time()
+
+    # ---- Step 1: Clone or pull repo ----
+    step("1/6  Clone / pull repo")
+    if REPO_DIR.exists():
+        print(f"  Repo exists, pulling latest...")
+        subprocess.run(["git", "-C", str(REPO_DIR), "pull"], check=False)
+        subprocess.run(["git", "-C", str(REPO_DIR), "lfs", "pull"], check=False)
+    else:
+        run(f"git clone {REPO_URL} {REPO_DIR}")
+    sys.path.insert(0, str(REPO_DIR))
+    sys.path.insert(0, str(REPO_DIR / "third_party" / "tiny_tts_frontend"))
+    os.chdir(REPO_DIR)
+    print(f"  Working dir: {REPO_DIR}")
+
+    # ---- Step 2: Install dependencies ----
+    step("2/6  Install dependencies")
+    pkgs = "torch torchaudio soundfile numpy g2p_en transformers gradio numba scipy datasets"
+    if not pip_installed("g2p_en"):
+        run(f"{sys.executable} -m pip install -q {pkgs}")
+    else:
+        print("  ⏭  Already installed (g2p_en found)")
+
+    import nltk
+    nltk.download("averaged_perceptron_tagger_eng", quiet=True)
+    nltk.download("cmudict", quiet=True)
+    print("  NLTK data OK")
+
+    # ---- Step 3: Download model weights ----
+    step("3/6  Download model weights")
+    WEIGHTS_DIR = REPO_DIR / "weights"
+    AC_WEIGHTS = WEIGHTS_DIR / "inflect_nano_v1_acoustic.pt"
+    VO_WEIGHTS = WEIGHTS_DIR / "inflect_nano_v1_vocoder.pt"
+
+    if AC_WEIGHTS.exists() and AC_WEIGHTS.stat().st_size > 1000:
+        print(f"  ⏭  Acoustic weights OK ({AC_WEIGHTS.stat().st_size / 1e6:.1f} MB)")
+    else:
+        run(f"{sys.executable} download_weights.py")
+
+    # ---- Step 4: Preprocess dataset ----
+    step("4/6  Preprocess dataset")
+    JSONL = Path(f"/content/durations_{SPEAKER}.jsonl")
+
+    if JSONL.exists() and JSONL.stat().st_size > 100:
+        print(f"  ⏭  JSONL exists ({JSONL.stat().st_size / 1e6:.2f} MB)")
+    else:
+        ok = run(
+            f"{sys.executable} preprocess_dataset.py hf "
+            f"--dataset {DATASET} --split {SPEAKER} "
+            f"--out {JSONL} --voice-id mark"
+        )
+        if not ok:
+            print("  ❌  Preprocessing failed. Check error above.")
+            sys.exit(1)
+
+    # ---- Step 5: Train ----
+    step("5/6  Train smoothness model")
+    OUT_DIR = Path("/content/checkpoints/smooth-v1")
+    LATEST_CKPT = OUT_DIR / "inflect-smooth-latest.pt"
+
+    import torch
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    print(f"  Device: {device}  ({torch.cuda.get_device_name(0) if device == 'cuda' else 'CPU'})")
+
+    if LATEST_CKPT.exists():
+        print(f"  ⏭  Checkpoint exists — adding {TRAIN_STEPS} more steps")
+        resume_flag = "--resume"
+    else:
+        print(f"  Starting fresh — {TRAIN_STEPS} steps")
+        resume_flag = ""
+
+    ok = run(
+        f"{sys.executable} -m inflect_nano.train_smooth "
+        f"--durations-jsonl {JSONL} "
+        f"--out-dir {OUT_DIR} "
+        f"--init-checkpoint {AC_WEIGHTS} "
+        f"--vocoder-checkpoint {VO_WEIGHTS} "
+        f"--steps {TRAIN_STEPS} "
+        f"--batch-size {BATCH_SIZE} "
+        f"--device {device} "
+        f"{resume_flag}"
+    )
+    if not ok:
+        print("  ❌  Training failed. Check error above.")
+        sys.exit(1)
+
+    # ---- Step 6: Generate sample ----
+    step("6/6  Generate sample")
+    from inference import load_acoustic, load_vocoder, synthesize
+    import soundfile as sf
+    import numpy as np
+
+    ac_path = LATEST_CKPT if LATEST_CKPT.exists() else AC_WEIGHTS
+    acoustic, speakers, ap = load_acoustic(ac_path, torch.device(device))
+    vocoder, vp = load_vocoder(VO_WEIGHTS, torch.device(device))
+    print(f"  Acoustic: {ap:,} params  Vocoder: {vp:,} params  Total: {ap+vp:,}")
+
+    text = "Every man is destined to die, but his work echoes through the ages."
+    audio = synthesize(
+        text, acoustic, vocoder, speakers, torch.device(device),
+        smooth_prosody=True, mel_smooth_sigma=0.8,
+    )
+    out = Path("/content/sample_smooth.wav")
+    sf.write(str(out), audio, 24000, subtype="PCM_16")
+
+    elapsed = (time.time() - start) / 60
+    print(f"\n{'='*60}")
+    print(f"  Done in {elapsed:.0f}m.  Sample: {out}  ({audio.size/24000:.1f}s)")
+    print(f"  Checkpoints: {OUT_DIR}")
+    print(f"{'='*60}")
+
+
+if __name__ == "__main__":
+    main()

colab_smooth_finetune.ipynb

@@ -0,0 +1,47 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Inflect-Nano Smoothness Fine-Tuning\n",
+    "Single-click pipeline. Each cell is safe to re-run \u2014 it skips steps already completed.\n",
+    "\n",
+    "**Runtime \u2192 Change runtime type \u2192 T4 GPU**"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {
+    "id": "runner"
+   },
+   "source": [
+    "# Download the runner script and execute it\n",
+    "import os, sys, subprocess\n",
+    "from pathlib import Path\n",
+    "\n",
+    "REPO = Path(\"/content/speech-nano\")\n",
+    "if not REPO.exists():\n",
+    "    subprocess.run([\"git\", \"clone\", \"https://git.quietidiot.com/michael-treadgold/speech-nano.git\", str(REPO)], check=False)\n",
+    "else:\n",
+    "    subprocess.run([\"git\", \"-C\", str(REPO), \"pull\"], check=False)\n",
+    "\n",
+    "sys.path.insert(0, str(REPO))\n",
+    "os.chdir(REPO)\n",
+    "\n",
+    "# Run the pipeline\n",
+    "exec(open(\"colab_run.py\").read())"
+   ],
+   "execution_count": null,
+   "outputs": []
+  }
+ ],
+ "metadata": {
+  "accelerator": "GPU",
+  "colab": {
+   "provenance": []
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

download_weights.py

@@ -0,0 +1,72 @@
+"""Download Inflect-Nano-v1 model weights from HuggingFace.
+
+Run this once after cloning the repo:
+    python download_weights.py
+"""
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+# Use requests or urllib — try both
+try:
+    import requests
+    _has_requests = True
+except ImportError:
+    import urllib.request
+    import urllib.error
+    _has_requests = False
+
+WEIGHTS_DIR = Path(__file__).resolve().parent / "weights"
+HF_BASE = "https://huggingface.co/owensong/Inflect-Nano-v1/resolve/main/weights"
+
+FILES = [
+    "inflect_nano_v1_acoustic.pt",
+    "inflect_nano_v1_vocoder.pt",
+]
+
+
+def download_file(url: str, dest: Path) -> None:
+    dest.parent.mkdir(parents=True, exist_ok=True)
+    print(f"Downloading {dest.name}...", end=" ", flush=True)
+
+    if _has_requests:
+        resp = requests.get(url, stream=True)
+        resp.raise_for_status()
+        total = int(resp.headers.get("content-length", 0))
+        with dest.open("wb") as f:
+            downloaded = 0
+            for chunk in resp.iter_content(chunk_size=8192):
+                f.write(chunk)
+                downloaded += len(chunk)
+                if total:
+                    pct = downloaded * 100 // total
+                    print(f"\rDownloading {dest.name}... {pct}%", end="", flush=True)
+        print(" done.")
+    else:
+        urllib.request.urlretrieve(url, str(dest))
+        print("done.")
+
+
+def main() -> None:
+    missing = [f for f in FILES if not (WEIGHTS_DIR / f).is_file()]
+    if not missing:
+        print("All weights already present.")
+        return
+
+    print(f"Missing {len(missing)} weight file(s). Downloading from HuggingFace...\n")
+    for filename in missing:
+        url = f"{HF_BASE}/{filename}"
+        dest = WEIGHTS_DIR / filename
+        try:
+            download_file(url, dest)
+        except Exception as e:
+            print(f"\n  FAILED: {e}")
+            print(f"  Manual download: {url}")
+            sys.exit(1)
+
+    print(f"\nAll weights downloaded to {WEIGHTS_DIR}")
+
+
+if __name__ == "__main__":
+    main()

inference.py

@@ -10,16 +10,18 @@ import soundfile as sf
 import torch
 
 REPO_ROOT = Path(__file__).resolve().parent
+VENDORED_FRONTEND = REPO_ROOT / "third_party" / "tiny_tts_frontend"
 sys.path.insert(0, str(REPO_ROOT))
+sys.path.insert(0, str(VENDORED_FRONTEND))
 
 from tiny_tts.nn import commons
 from tiny_tts.text import phonemes_to_ids
 from tiny_tts.text.english import grapheme_to_phoneme, normalize_text
 from tiny_tts.utils import ADD_BLANK
 
-from tinytts_text_cleaning import clean_tinytts_text
+from inflect_nano.text_cleaning import clean_tinytts_text
-from train_hifigan_oracle_v1 import HifiGanGenerator, make_config
+from inflect_nano.vocoder import HifiGanGenerator, make_config
-from train_inflect_micro_fastspeech_v3_pitch import MicroFastSpeech, MicroFastSpeechConfig
+from inflect_nano.acoustic import MicroFastSpeech, MicroFastSpeechConfig
 
 
 DEFAULT_ACOUSTIC = REPO_ROOT / "weights" / "inflect_nano_v1_acoustic.pt"
@@ -76,6 +78,31 @@ def normalize_audio(audio: np.ndarray, target_rms_db: float = -20.0, peak_db: fl
     return np.clip(audio, -1.0, 1.0)
 
 
+def smooth_mel(mel: torch.Tensor, sigma: float = 1.0) -> torch.Tensor:
+    """Gaussian temporal smoothing on mel frames to reduce frame-to-frame jitter."""
+    if sigma <= 0:
+        return mel
+    kernel_size = int(2 * math.ceil(2 * sigma) + 1)
+    if kernel_size < 3:
+        return mel
+    kernel = torch.exp(-0.5 * (torch.arange(kernel_size, device=mel.device, dtype=mel.dtype) - kernel_size // 2) ** 2 / sigma**2)
+    kernel = kernel / kernel.sum()
+    # [B, n_mels, T] -> pad last dim (time), then conv1d over time
+    pad = kernel_size // 2
+    mel_padded = torch.nn.functional.pad(mel, (pad, pad), mode="replicate")
+    kernel_expanded = kernel.view(1, 1, -1).expand(mel.shape[1], 1, -1)
+    return torch.nn.functional.conv1d(mel_padded, kernel_expanded, groups=mel.shape[1])
+
+
+def apply_lowpass(wav: np.ndarray, cutoff_hz: float, sample_rate: int = 24000) -> np.ndarray:
+    """Simple low-pass filter to reduce vocoder buzz above cutoff."""
+    if cutoff_hz <= 0 or cutoff_hz >= sample_rate / 2:
+        return wav
+    from scipy import signal
+    sos = signal.butter(4, cutoff_hz, btype="low", fs=sample_rate, output="sos")
+    return signal.sosfiltfilt(sos, wav).astype(np.float32)
+
+
 @torch.inference_mode()
 def synthesize(
     text: str,
@@ -86,6 +113,9 @@ def synthesize(
     length_scale: float = 1.0,
     pitch_scale: float = 1.0,
     energy_scale: float = 1.0,
+    smooth_prosody: bool = False,
+    mel_smooth_sigma: float = 0.0,
+    lowpass_hz: float = 0.0,
 ) -> np.ndarray:
     phone, tone, lang = text_to_tokens(text)
     phone = phone.unsqueeze(0).to(device)
@@ -100,9 +130,15 @@ def synthesize(
         length_scale=float(length_scale),
         pitch_scale=float(pitch_scale),
         energy_scale=float(energy_scale),
+        smooth_predictors=smooth_prosody,
     )
+    if mel_smooth_sigma > 0:
+        mel = smooth_mel(mel, mel_smooth_sigma)
     wav = vocoder(mel).squeeze().detach().cpu().numpy()
-    return normalize_audio(wav)
+    wav = normalize_audio(wav)
+    if lowpass_hz > 0:
+        wav = apply_lowpass(wav, lowpass_hz)
+    return wav
 
 
 def main() -> None:
@@ -115,6 +151,12 @@ def main() -> None:
     ap.add_argument("--length-scale", type=float, default=1.0)
     ap.add_argument("--pitch-scale", type=float, default=1.0)
     ap.add_argument("--energy-scale", type=float, default=1.0)
+    ap.add_argument("--smooth-prosody", action="store_true",
+                    help="Apply 3-frame averaging to pitch/energy/brightness (reduces jitter)")
+    ap.add_argument("--mel-smooth-sigma", type=float, default=0.0,
+                    help="Gaussian temporal smooth on mel frames (0.5-1.5 reduces buzzing)")
+    ap.add_argument("--lowpass-hz", type=float, default=0.0,
+                    help="Low-pass cutoff in Hz (e.g. 8000-11000 reduces vocoder buzz)")
     args = ap.parse_args()
 
     device = torch.device(args.device)
@@ -129,6 +171,9 @@ def main() -> None:
         length_scale=args.length_scale,
         pitch_scale=args.pitch_scale,
         energy_scale=args.energy_scale,
+        smooth_prosody=args.smooth_prosody,
+        mel_smooth_sigma=args.mel_smooth_sigma,
+        lowpass_hz=args.lowpass_hz,
     )
     args.out.parent.mkdir(parents=True, exist_ok=True)
     sf.write(str(args.out), audio, 24000, subtype="PCM_16")

inflect_nano/__init__.py

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

inflect_nano/acoustic.py

@@ -16,10 +16,10 @@ import torchaudio
 
 SCRIPT_ROOT = Path(__file__).resolve().parent
 PROJECT_ROOT = SCRIPT_ROOT.parents[0]
-TINY_ROOT = PROJECT_ROOT / "third_party" / "tiny-tts"
+FRONTEND_ROOT = PROJECT_ROOT / "third_party" / "tiny_tts_frontend"
-sys.path = [str(TINY_ROOT), str(SCRIPT_ROOT)] + [p for p in sys.path if p]
+sys.path = [str(FRONTEND_ROOT), str(SCRIPT_ROOT)] + [p for p in sys.path if p]
 
-from train_hifigan_oracle_v1 import HifiGanConfig, HifiGanGenerator, MelFrontend
+from inflect_nano.vocoder import HifiGanConfig, HifiGanGenerator, MelFrontend
 
 
 @dataclass

inflect_nano/train_smooth.py

@@ -0,0 +1,594 @@
+"""
+Enhanced smoothness training for Inflect-Nano acoustic model.
+
+Key additions over the base training:
+  1. Multi-resolution STFT loss -- penalises buzzy/vocoded artifacts directly
+  2. Deeper residual postnet with per-layer skip connections
+  3. Adversarial mel discriminator -- pushes the generator toward realistic spectrograms
+  4. Vocoder consistency loss enabled with sensible defaults
+  5. Tuned hyperparameters for smoother prosody and spectral continuity
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import random
+import time
+from dataclasses import dataclass
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from inflect_nano.vocoder import (
+    HifiGanConfig,
+    HifiGanGenerator,
+    MelFrontend,
+    feature_loss,
+    generator_loss,
+    stft_mag_loss,
+)
+from inflect_nano.acoustic import (
+    ConvFFNBlock,
+    MicroFastSpeech,
+    MicroFastSpeechConfig,
+    collate,
+    collate_prepared,
+    count_parameters,
+    fit_durations,
+    group_duration_targets,
+    load_audio,
+    load_frozen_vocoder,
+    load_model_state_flexible,
+    load_rows,
+    masked_accel_loss,
+    masked_delta_loss,
+    masked_l1,
+    masked_mse,
+    masked_wav_l1,
+    pad_1d,
+    pad_mels,
+    pad_wavs,
+    prepare_row_features,
+    save_checkpoint,
+    set_trainable_by_mode,
+    token_mse,
+    token_mse_nd,
+)
+
+
+# ---------------------------------------------------------------------------
+# 1.  Deeper residual postnet
+# ---------------------------------------------------------------------------
+class ResidualPostnet(nn.Module):
+    """Stacked conv blocks with *per-block* residual connections."""
+
+    def __init__(self, n_mels: int, hidden: int, layers: int = 5, kernel: int = 5):
+        super().__init__()
+        self.blocks = nn.ModuleList()
+        self.entry = nn.Conv1d(n_mels, hidden, kernel, padding=kernel // 2)
+        for _ in range(layers):
+            self.blocks.append(
+                nn.Sequential(
+                    nn.Conv1d(hidden, hidden, kernel, padding=kernel // 2),
+                    nn.BatchNorm1d(hidden),
+                    nn.Tanh(),
+                    nn.Conv1d(hidden, hidden, kernel, padding=kernel // 2),
+                    nn.BatchNorm1d(hidden),
+                )
+            )
+        self.exit = nn.Conv1d(hidden, n_mels, kernel, padding=kernel // 2)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        residual = x
+        h = self.entry(x)
+        h = torch.tanh(h)
+        for block in self.blocks:
+            h = h + block(h)
+        return residual + self.exit(h)
+
+
+# ---------------------------------------------------------------------------
+# 2.  Mel discriminator (simple conv2d stack)
+# ---------------------------------------------------------------------------
+class MelDiscriminator(nn.Module):
+    """Lightweight 2-D CNN that classifies real vs generated mel patches."""
+
+    def __init__(self, n_mels: int = 80):
+        super().__init__()
+        self.convs = nn.ModuleList(
+            [
+                nn.utils.parametrizations.weight_norm(nn.Conv2d(1, 32, (3, 5), stride=(1, 2), padding=(1, 2))),
+                nn.utils.parametrizations.weight_norm(nn.Conv2d(32, 64, (3, 5), stride=(1, 2), padding=(1, 2))),
+                nn.utils.parametrizations.weight_norm(nn.Conv2d(64, 128, (3, 5), stride=(1, 2), padding=(1, 2))),
+                nn.utils.parametrizations.weight_norm(nn.Conv2d(128, 256, (3, 5), stride=(1, 2), padding=(1, 2))),
+                nn.utils.parametrizations.weight_norm(nn.Conv2d(256, 1, (3, 5), padding=(1, 2))),
+            ]
+        )
+
+    def forward(self, mel: torch.Tensor) -> list[torch.Tensor]:
+        # mel: [B, n_mels, T]
+        fmap: list[torch.Tensor] = []
+        x = mel.unsqueeze(1)  # [B, 1, n_mels, T]
+        for conv in self.convs:
+            x = conv(x)
+            x = F.leaky_relu(x, 0.2)
+            fmap.append(x)
+        return fmap
+
+
+# ---------------------------------------------------------------------------
+# 3.  Enhanced acoustic model (drop-in replacement)
+# ---------------------------------------------------------------------------
+@dataclass
+class SmoothAcousticConfig(MicroFastSpeechConfig):
+    postnet_layers: int = 5      # depth of residual postnet
+    postnet_kernel: int = 5      # kernel size for postnet convs
+
+
+class SmoothMicroFastSpeech(MicroFastSpeech):
+    """MicroFastSpeech with a deeper residual postnet."""
+
+    def __init__(self, cfg: SmoothAcousticConfig):
+        super().__init__(cfg)
+        # Replace the original shallow postnet
+        self.postnet = ResidualPostnet(
+            n_mels=cfg.n_mels,
+            hidden=cfg.hidden,
+            layers=cfg.postnet_layers,
+            kernel=cfg.postnet_kernel,
+        )
+
+
+# ---------------------------------------------------------------------------
+# 4.  Multi-resolution STFT loss wrapper
+# ---------------------------------------------------------------------------
+def multi_resolution_stft_loss(
+    pred_mel: torch.Tensor,
+    target_mel: torch.Tensor,
+    frame_mask: torch.Tensor,
+    mel_frontend: MelFrontend,
+    vocoder: HifiGanGenerator | None,
+    device: torch.device,
+    fft_sizes: tuple[int, ...] = (512, 1024, 2048),
+    hop_sizes: tuple[int, ...] = (128, 256, 512),
+    win_lengths: tuple[int, ...] = (512, 1024, 2048),
+) -> torch.Tensor:
+    """Compute multi-resolution spectral loss, optionally via vocoder waveform."""
+    B = pred_mel.shape[0]
+    common = min(pred_mel.shape[-1], target_mel.shape[-1], frame_mask.shape[-1])
+    pm = pred_mel[..., :common]
+    tm = target_mel[..., :common]
+    fm = frame_mask[:, :common]
+
+    if vocoder is not None:
+        with torch.no_grad():
+            pred_wav = vocoder(pm).squeeze(1)
+            target_wav = vocoder(tm).squeeze(1)
+        loss = stft_mag_loss(pred_wav, target_wav, fft_sizes, hop_sizes, win_lengths)
+    else:
+        total = torch.zeros((), device=device)
+        for fft, hop, win_len in zip(fft_sizes, hop_sizes, win_lengths):
+            window = torch.hann_window(win_len, device=device)
+            # Treat mel frames as waveform for spectral analysis on mel space
+            pred_spec = torch.stft(
+                pm.reshape(B * pm.shape[1], -1).T.reshape(B, pm.shape[1], -1)[:, :1, :].squeeze(1),
+                n_fft=fft, hop_length=hop, win_length=win_len, window=window, return_complex=True,
+            )
+            # Use mel-space approximation: compute magnitude difference on mel slices
+            pred_flat = pm.transpose(1, 2)  # [B, T, 80]
+            targ_flat = tm.transpose(1, 2)  # [B, T, 80]
+            mask_flat = fm.unsqueeze(-1)     # [B, T, 1]
+            spec_loss = (F.l1_loss(pred_flat * mask_flat, targ_flat * mask_flat) /
+                         mask_flat.sum().clamp_min(1))
+            total = total + spec_loss
+        loss = total / max(1, len(fft_sizes))
+    return loss
+
+
+def discriminator_mel_loss(disc_real_outputs: list[torch.Tensor], disc_generated_outputs: list[torch.Tensor]) -> torch.Tensor:
+    loss = torch.zeros((), device=disc_real_outputs[0][0].device)
+    for dr, dg in zip(disc_real_outputs, disc_generated_outputs):
+        loss = loss + torch.mean((1 - dr[-1]) ** 2) + torch.mean(dg[-1] ** 2)
+    return loss
+
+
+def generator_mel_loss(disc_outputs: list[torch.Tensor]) -> torch.Tensor:
+    loss = torch.zeros((), device=disc_outputs[0][0].device)
+    for dg in disc_outputs:
+        loss = loss + torch.mean((1 - dg[-1]) ** 2)
+    return loss
+
+
+# ---------------------------------------------------------------------------
+# 5.  Enhanced training loop
+# ---------------------------------------------------------------------------
+def train_smooth(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)
+
+    acoustic_cfg = SmoothAcousticConfig(
+        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,
+        postnet_layers=args.postnet_layers,
+        postnet_kernel=args.postnet_kernel,
+    )
+
+    for row in rows:
+        row["speaker_id"] = speakers[str(row.get("voice_id") or "mark")]
+    random.Random(args.seed).shuffle(rows)
+
+    model = SmoothMicroFastSpeech(acoustic_cfg).to(device)
+    start_step = 0
+
+    # Initialise from existing checkpoint or resume
+    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"Initialised from {args.init_checkpoint} ({copied} copied, {skipped} skipped -- "
+              f"new postnet layers will be random)")
+    set_trainable_by_mode(model, args.trainable)
+    trainable_params = [p for p in model.parameters() if p.requires_grad]
+    optim_g = torch.optim.AdamW(trainable_params, lr=args.lr, betas=(0.9, 0.98), weight_decay=args.weight_decay)
+
+    # Cosine LR schedule with linear warmup
+    warmup_steps = args.warmup_steps
+    total_steps = args.steps
+
+    if args.resume:
+        ckpt_path = None
+        for p in args.out_dir.glob("inflect-smooth-*.pt"):
+            if p.stem.endswith("-latest"):
+                ckpt_path = p
+                break
+        if ckpt_path:
+            ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
+            model.load_state_dict(ckpt["model"])
+            optim_g.load_state_dict(ckpt["optim_g"])
+            start_step = int(ckpt.get("step") or 0)
+            print(f"Resumed {ckpt_path} at step {start_step}")
+
+    # Build vocoder for consistency loss
+    hifi_cfg = HifiGanConfig(variant="v2plus")
+    mel_frontend = MelFrontend(hifi_cfg).to(device)
+
+    consistency_vocoder = None
+    if args.vocoder_checkpoint:
+        consistency_vocoder, consistency_cfg = load_frozen_vocoder(args.vocoder_checkpoint, device)
+        print(f"Loaded frozen vocoder: {args.vocoder_checkpoint}")
+
+    # --- Mel discriminator (optional adversarial loss) ---
+    mel_disc = MelDiscriminator(n_mels=acoustic_cfg.n_mels).to(device) if args.adv_mel_weight > 0 else None
+    if mel_disc is not None:
+        optim_d = torch.optim.AdamW(mel_disc.parameters(), lr=args.lr, betas=(0.5, 0.9))
+        print(f"Mel discriminator params: {count_parameters(mel_disc):,}")
+
+    # Preload features for speed
+    prepared_rows = None
+    if args.preload_features:
+        print("Preloading audio/mel/pitch features...", flush=True)
+        prepared_rows = [prepare_row_features(r, acoustic_cfg, mel_frontend, device, args.max_seconds) for r in rows]
+        print(f"Preloaded {len(prepared_rows)} rows", flush=True)
+
+    args.out_dir.mkdir(parents=True, exist_ok=True)
+    (args.out_dir / "smooth_config.json").write_text(
+        json.dumps({
+            "acoustic_config": acoustic_cfg.__dict__ if hasattr(acoustic_cfg, '__dict__') else {},
+            "speakers": speakers,
+            "rows": len(rows),
+            "params": count_parameters(model),
+        }, indent=2, default=str),
+        encoding="utf-8",
+    )
+
+    print(f"Rows: {len(rows)}  Speakers: {speakers}")
+    print(f"Acoustic params: {count_parameters(model):,} ({count_parameters(model)/1e6:.3f}M)")
+    print(f"Trainable: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}")
+    print(f"Postnet: {args.postnet_layers} layers, kernel={args.postnet_kernel}, scale={args.postnet_scale}")
+    print(f"STFT weight: {args.stft_weight}  Adv mel weight: {args.adv_mel_weight}")
+    print(f"Vocoder consistency: wav={args.vocoder_wav_weight} mel={args.vocoder_mel_weight}")
+
+    rng = random.Random(args.seed + start_step)
+    step = start_step
+    started = time.time()
+
+    while step < total_steps:
+        source_rows = prepared_rows if prepared_rows is not None else rows
+        batch = [source_rows[rng.randrange(len(source_rows))] 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, acoustic_cfg, mel_frontend, device, args.max_seconds, hifi_cfg.hop_size
+            )
+
+        # ---- LR schedule (cosine with warmup) ----
+        if step < warmup_steps:
+            lr_scale = (step + 1) / max(1, warmup_steps)
+        else:
+            progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)
+            lr_scale = 0.5 * (1.0 + math.cos(math.pi * progress))
+        for pg in optim_g.param_groups:
+            pg["lr"] = args.lr * lr_scale
+        if mel_disc is not None:
+            for pg in optim_d.param_groups:
+                pg["lr"] = args.lr * lr_scale
+
+        # ---- Forward pass ----
+        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)
+
+        # Base mel losses
+        mel_l1 = masked_l1(out["mel"], target_mel, frame_mask)
+        mel_mse = masked_mse(out["mel"], target_mel, frame_mask)
+        delta = masked_delta_loss(out["mel"], target_mel, frame_mask)
+        accel = masked_accel_loss(out["mel"], target_mel, frame_mask)
+
+        # Token-level losses
+        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)
+
+        # ---- Multi-resolution STFT loss ----
+        stft_loss = torch.zeros((), device=device)
+        if args.stft_weight > 0:
+            stft_loss = multi_resolution_stft_loss(
+                out["mel"], target_mel, frame_mask,
+                mel_frontend, consistency_vocoder, device,
+            ) * args.stft_weight
+
+        # ---- Prosody exposure bias loss ----
+        predicted_prosody_mel_loss = torch.zeros((), device=device)
+        predicted_prosody_delta_loss = torch.zeros((), device=device)
+        if args.predicted_prosody_mel_weight > 0 or args.predicted_prosody_delta_weight > 0:
+            pred_out = model(phone, tone, lang, speaker, durations)
+            if args.predicted_prosody_mel_weight > 0:
+                predicted_prosody_mel_loss = masked_l1(pred_out["mel"], target_mel, frame_mask)
+            if args.predicted_prosody_delta_weight > 0:
+                predicted_prosody_delta_loss = masked_delta_loss(pred_out["mel"], target_mel, frame_mask)
+
+        # ---- Robust prosody loss (mix of predicted + reference) ----
+        robust_prosody_mel_loss = torch.zeros((), device=device)
+        robust_prosody_delta_loss = torch.zeros((), device=device)
+        if args.robust_prosody_mel_weight > 0 or args.robust_prosody_delta_weight > 0:
+            robust_out = 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:
+                robust_prosody_mel_loss = masked_l1(robust_out["mel"], target_mel, frame_mask)
+            if args.robust_prosody_delta_weight > 0:
+                robust_prosody_delta_loss = masked_delta_loss(robust_out["mel"], target_mel, frame_mask)
+
+        # ---- Vocoder consistency loss ----
+        voc_wav_loss = torch.zeros((), device=device)
+        voc_mel_loss = torch.zeros((), device=device)
+        if consistency_vocoder is not None:
+            if args.vocoder_wav_weight > 0:
+                pred_wav = consistency_vocoder(out["mel"].clamp(-12, 2))
+                voc_wav_loss = masked_wav_l1(pred_wav, target_wav, frame_mask, hifi_cfg.hop_size)
+            if args.vocoder_mel_weight > 0:
+                pred_wav = consistency_vocoder(out["mel"].clamp(-12, 2))
+                pred_recon_mel = mel_frontend(pred_wav.squeeze(1))
+                voc_mel_loss = masked_l1(pred_recon_mel, target_mel, frame_mask)
+
+        # ---- Adversarial mel loss ----
+        adv_mel_loss = torch.zeros((), device=device)
+        disc_loss = torch.zeros((), device=device)
+        if mel_disc is not None and args.adv_mel_weight > 0:
+            # Train discriminator
+            common = min(out["mel"].shape[-1], target_mel.shape[-1])
+            real_mel = target_mel[..., :common]
+            fake_mel = out["mel"][..., :common].detach()
+
+            optim_d.zero_grad(set_to_none=True)
+            real_fmap = mel_disc(real_mel)
+            fake_fmap = mel_disc(fake_mel)
+            disc_loss = discriminator_mel_loss(real_fmap, fake_fmap)
+            disc_loss.backward()
+            torch.nn.utils.clip_grad_norm_(mel_disc.parameters(), args.grad_clip)
+            optim_d.step()
+
+            # Generator adversarial loss
+            adv_fake_fmap = mel_disc(out["mel"][..., :common])
+            adv_mel_loss = generator_mel_loss(adv_fake_fmap)
+
+            # Feature matching
+            with torch.no_grad():
+                real_fmap_detached = mel_disc(real_mel)
+            fm_loss = torch.zeros((), device=device)
+            for rf, ff in zip(real_fmap_detached, adv_fake_fmap):
+                for rl, fl in zip(rf, ff):
+                    fm_loss = fm_loss + F.l1_loss(rl, fl)
+            adv_mel_loss = adv_mel_loss + args.fm_weight * fm_loss
+
+        # ---- Total generator loss ----
+        loss_g = (
+            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
+            + stft_loss
+            + args.adv_mel_weight * adv_mel_loss
+        )
+
+        optim_g.zero_grad(set_to_none=True)
+        loss_g.backward()
+        grad = torch.nn.utils.clip_grad_norm_(trainable_params, args.grad_clip)
+        optim_g.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 = (total_steps - step) / max(1e-6, speed)
+            print(
+                f"step={step}/{total_steps} loss={loss_g.item():.4f} mel={mel_l1.item():.4f} "
+                f"mse={mel_mse.item():.4f} delta={delta.item():.4f} accel={accel.item():.4f} "
+                f"dur={dur_loss.item():.4f} gdur={group_dur_loss.item():.4f} "
+                f"energy={energy_loss.item():.4f} bright={bright_loss.item():.4f} "
+                f"pitch={pitch_loss.item():.4f} stft={stft_loss.item():.4f} "
+                f"pmel={predicted_prosody_mel_loss.item():.4f} "
+                f"pdelta={predicted_prosody_delta_loss.item():.4f} "
+                f"rmel={robust_prosody_mel_loss.item():.4f} "
+                f"rdelta={robust_prosody_delta_loss.item():.4f} "
+                f"vwav={voc_wav_loss.item():.4f} vmel={voc_mel_loss.item():.4f} "
+                f"adv={adv_mel_loss.item():.4f} disc={disc_loss.item():.4f} "
+                f"lr={lr_scale*args.lr:.2g} 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 >= total_steps:
+            payload = {
+                "model": model.state_dict(),
+                "optim_g": optim_g.state_dict(),
+                "step": step,
+                "speakers": speakers,
+                "params": count_parameters(model),
+            }
+            if mel_disc is not None:
+                payload["mel_disc"] = mel_disc.state_dict()
+                payload["optim_d"] = optim_d.state_dict()
+            tmp = args.out_dir / f"inflect-smooth-{step}.pt.tmp"
+            torch.save(payload, tmp)
+            tmp.replace(args.out_dir / f"inflect-smooth-{step}.pt")
+            torch.save(payload, args.out_dir / "inflect-smooth-latest.pt")
+
+    print(f"Done. {args.out_dir}")
+
+
+# ---------------------------------------------------------------------------
+# CLI
+# ---------------------------------------------------------------------------
+def main() -> None:
+    ap = argparse.ArgumentParser(
+        description="Train Inflect-Nano acoustic model with enhanced smoothness losses.",
+    )
+    # Data
+    ap.add_argument("--durations-jsonl", type=Path, required=True)
+    ap.add_argument("--out-dir", type=Path, required=True)
+    ap.add_argument("--max-rows", type=int, default=0)
+
+    # Architecture
+    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("--postnet-scale", type=float, default=0.35,
+                    help="Postnet refinement scale (higher = postnet has more influence)")
+    ap.add_argument("--postnet-layers", type=int, default=5,
+                    help="Depth of residual postnet")
+    ap.add_argument("--postnet-kernel", type=int, default=5,
+                    help="Kernel size for postnet convolutions")
+    ap.add_argument("--abs-frame-bins", type=int, default=512)
+
+    # Training
+    ap.add_argument("--steps", type=int, default=20000)
+    ap.add_argument("--batch-size", type=int, default=6)
+    ap.add_argument("--lr", type=float, default=2e-4)
+    ap.add_argument("--weight-decay", type=float, default=1e-4)
+    ap.add_argument("--warmup-steps", type=int, default=1000,
+                    help="Linear warmup steps for cosine LR schedule")
+    ap.add_argument("--grad-clip", type=float, default=5.0)
+
+    # Loss weights (tuned for smoothness)
+    ap.add_argument("--mse-weight", type=float, default=0.25)
+    ap.add_argument("--delta-weight", type=float, default=0.25,
+                    help="Spectral delta smoothness (higher = smoother transitions)")
+    ap.add_argument("--accel-weight", type=float, default=0.08,
+                    help="Spectral acceleration penalty (higher = less jitter)")
+    ap.add_argument("--duration-weight", type=float, default=0.08)
+    ap.add_argument("--group-duration-weight", type=float, default=0.02)
+    ap.add_argument("--energy-weight", type=float, default=0.06)
+    ap.add_argument("--bright-weight", type=float, default=0.06)
+    ap.add_argument("--pitch-weight", type=float, default=0.06)
+
+    # Spectral / perceptual losses
+    ap.add_argument("--stft-weight", type=float, default=0.15,
+                    help="Multi-resolution STFT loss (improves spectral smoothness)")
+    ap.add_argument("--predicted-prosody-mel-weight", type=float, default=0.05,
+                    help="Exposure bias: use only predicted prosody for mel")
+    ap.add_argument("--predicted-prosody-delta-weight", type=float, default=0.03)
+    ap.add_argument("--robust-prosody-mix", type=float, default=0.5,
+                    help="Mix ratio for robust prosody training")
+    ap.add_argument("--robust-prosody-mel-weight", type=float, default=0.05)
+    ap.add_argument("--robust-prosody-delta-weight", type=float, default=0.03)
+    ap.add_argument("--adv-mel-weight", type=float, default=0.08,
+                    help="Adversarial mel loss (pushes toward realistic spectrograms)")
+    ap.add_argument("--fm-weight", type=float, default=2.0,
+                    help="Feature matching weight for adversarial loss")
+
+    # Vocoder consistency
+    ap.add_argument("--vocoder-checkpoint", type=Path, default=None,
+                    help="Path to vocoder for consistency loss")
+    ap.add_argument("--vocoder-wav-weight", type=float, default=0.12,
+                    help="Vocoder waveform consistency loss")
+    ap.add_argument("--vocoder-mel-weight", type=float, default=0.08,
+                    help="Vocoder mel-reconstruction consistency loss")
+
+    # Checkpointing
+    ap.add_argument("--init-checkpoint", type=Path,
+                    help="Start from an existing acoustic checkpoint for fine-tuning")
+    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")
+
+    # Misc
+    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_smooth(args)
+
+
+if __name__ == "__main__":
+    main()

preprocess_dataset.py

@@ -0,0 +1,463 @@
+"""
+Preprocess an audio+text dataset into the durations.jsonl format required
+by Inflect-Nano training scripts.
+
+Supports:
+  - HuggingFace datasets (CMU ARCTIC, LJSpeech, etc.)
+  - Local directory of .wav files with matching .txt transcriptions
+  - LJSpeech-format metadata.csv
+
+Output: a .jsonl file with one JSON object per line containing:
+  phone_ids, tone_ids, lang_ids, hifigan_durations, target_audio, speaker_id, voice_id
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import sys
+from pathlib import Path
+
+import numpy as np
+import soundfile as sf
+import torch
+import torchaudio
+
+REPO_ROOT = Path(__file__).resolve().parent
+VENDORED_FRONTEND = REPO_ROOT / "third_party" / "tiny_tts_frontend"
+sys.path.insert(0, str(REPO_ROOT))
+sys.path.insert(0, str(VENDORED_FRONTEND))
+
+from tiny_tts.nn import commons
+from tiny_tts.text import phonemes_to_ids
+from tiny_tts.text.english import grapheme_to_phoneme, normalize_text
+from tiny_tts.utils import ADD_BLANK
+
+from inflect_nano.text_cleaning import clean_tinytts_text
+from inflect_nano.vocoder import HifiGanConfig, MelFrontend
+
+
+def text_to_ids(text: str) -> tuple[list[int], list[int], list[int]]:
+    """Convert English text to TinyTTS phone/tone/lang ID lists."""
+    cleaned = clean_tinytts_text(text)
+    normalized = normalize_text(cleaned)
+    phones, tones, _ = grapheme_to_phoneme(normalized)
+    phone_ids, tone_ids, lang_ids = phonemes_to_ids(phones, tones, "EN")
+    if ADD_BLANK:
+        phone_ids = commons.insert_blanks(phone_ids, 0)
+        tone_ids = commons.insert_blanks(tone_ids, 0)
+        lang_ids = commons.insert_blanks(lang_ids, 0)
+    return phone_ids, tone_ids, lang_ids
+
+
+def estimate_durations_uniform(
+    phone_ids: list[int],
+    audio_path: str,
+    mel_frontend: MelFrontend,
+    sample_rate: int = 24000,
+) -> list[int]:
+    """Estimate mel-frame durations uniformly across all phones.
+
+    Loads the audio, extracts mel spectrogram, then distributes the total
+    frame count evenly across phones.  This is a coarse estimate suitable
+    for decoder-only fine-tuning where the duration predictor is frozen.
+    """
+    wav, sr = torchaudio.load(audio_path)
+    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.clamp(-1, 1)
+
+    with torch.no_grad():
+        mel = mel_frontend(wav)
+    total_frames = mel.shape[-1]
+
+    # Count non-blank phones (id != 0)
+    visible_phones = [i for i, pid in enumerate(phone_ids) if pid != 0]
+
+    if not visible_phones:
+        return [0] * len(phone_ids)
+
+    # Distribute frames: each blank gets 1 frame, rest split evenly
+    blank_count = len(phone_ids) - len(visible_phones)
+    remaining = max(0, total_frames - blank_count)
+    base_dur = max(1, remaining // len(visible_phones))
+    remainder = remaining - base_dur * len(visible_phones)
+
+    durations = []
+    for pid in phone_ids:
+        if pid == 0:
+            durations.append(1)
+        else:
+            extra = 1 if remainder > 0 else 0
+            remainder = max(0, remainder - 1)
+            durations.append(base_dur + extra)
+
+    # Scale to match total_frames
+    current_sum = sum(durations)
+    if current_sum > 0 and current_sum != total_frames:
+        scale = total_frames / current_sum
+        scaled = [max(1, round(d * scale)) for d in durations]
+        # Fix rounding errors
+        diff = total_frames - sum(scaled)
+        for i in range(abs(diff)):
+            if diff > 0:
+                scaled[i % len(scaled)] += 1
+            else:
+                if scaled[i % len(scaled)] > 1:
+                    scaled[i % len(scaled)] -= 1
+        durations = scaled
+
+    # Ensure no zero durations for non-blank phones
+    for i, pid in enumerate(phone_ids):
+        if pid != 0 and durations[i] < 1:
+            durations[i] = 1
+
+    return durations
+
+
+def process_hf_dataset(
+    dataset_path: str,
+    output_jsonl: Path,
+    audio_dir: Path | None,
+    subset: str = "default",
+    split: str | None = None,
+    text_key: str = "text",
+    speaker_key: str = "speaker",
+    max_rows: int = 0,
+    voice_id: str = "speaker",
+) -> int:
+    """Process a HuggingFace dataset (loaded via `datasets` library)."""
+    from datasets import load_dataset
+
+    print(f"Loading HF dataset: {dataset_path}  subset={subset}  split={split}")
+    if split:
+        ds = load_dataset(dataset_path, subset, split=split)
+    else:
+        ds_dict = load_dataset(dataset_path, subset)
+        splits = list(ds_dict.keys())
+        print(f"Available splits: {splits}")
+        preferred = [s for s in splits if "train" in s.lower()]
+        ds = ds_dict[preferred[0] if preferred else splits[0]]
+
+    # Debug: print first row to see available columns
+    first = next(iter(ds))
+    print(f"Dataset columns: {list(first.keys())}")
+    for k, v in first.items():
+        if isinstance(v, dict):
+            print(f"  {k}: dict with keys {list(v.keys())}")
+        elif isinstance(v, str):
+            print(f"  {k}: str ({len(v)} chars) = {v[:80]!r}")
+        else:
+            print(f"  {k}: {type(v).__name__}")
+
+    import tempfile
+
+    mel_frontend = MelFrontend(HifiGanConfig(variant="v2plus"))
+    output_jsonl.parent.mkdir(parents=True, exist_ok=True)
+
+    # Temp dir for audio extracted from HF arrays (cleaned up on exit)
+    audio_tmp = Path(tempfile.mkdtemp(prefix="inflect_audio_"))
+    print(f"Audio tmp dir: {audio_tmp}")
+
+    count = 0
+    skip_no_text = 0
+    skip_no_audio = 0
+    skip_bad_ids = 0
+    skip_bad_audio = 0
+    try:
+        with output_jsonl.open("w", encoding="utf-8") as f:
+            for i, row in enumerate(ds):
+                text = str(row.get(text_key, "")).strip()
+                if not text:
+                    skip_no_text += 1
+                    if skip_no_text <= 3:
+                        print(f"  Row {i}: no text (keys={list(row.keys())[:5]})")
+                    continue
+
+                # Get audio path or array
+                audio_info = row.get("audio", row.get("file", None))
+                if audio_info is None:
+                    skip_no_audio += 1
+                    if skip_no_audio <= 3:
+                        print(f"  Row {i}: no audio/file key (keys={list(row.keys())[:5]})")
+                    continue
+
+                # HF Audio feature returns an object with .array / .sampling_rate / .path
+                if hasattr(audio_info, "array"):
+                    audio_array = audio_info["array"] if isinstance(audio_info, dict) else audio_info.array
+                    sr = audio_info["sampling_rate"] if isinstance(audio_info, dict) else audio_info.sampling_rate
+                    audio_path = str(audio_tmp / f"utt_{i:06d}.wav")
+                    sf.write(audio_path, audio_array, sr, subtype="PCM_16")
+                elif isinstance(audio_info, dict):
+                    audio_array = audio_info.get("array")
+                    sr = audio_info.get("sampling_rate", 24000)
+                    if audio_array is None:
+                        skip_bad_audio += 1
+                        continue
+                    audio_path = str(audio_tmp / f"utt_{i:06d}.wav")
+                    sf.write(audio_path, audio_array, sr, subtype="PCM_16")
+                elif isinstance(audio_info, str):
+                    audio_path = audio_info
+                    if audio_dir:
+                        audio_path = str(audio_dir / Path(audio_info).name)
+                    if not Path(audio_path).is_file():
+                        skip_bad_audio += 1
+                        continue
+                else:
+                    skip_bad_audio += 1
+                    if skip_bad_audio <= 3:
+                        print(f"  Row {i}: unexpected audio type: {type(audio_info).__name__}")
+                    continue
+
+                try:
+                    phone_ids, tone_ids, lang_ids = text_to_ids(text)
+                except Exception as e:
+                    skip_bad_ids += 1
+                    if skip_bad_ids <= 3:
+                        print(f"  Row {i}: text-to-ids failed: {e}  text={text[:60]!r}")
+                    continue
+
+                if not phone_ids:
+                    continue
+
+                durations = estimate_durations_uniform(phone_ids, audio_path, mel_frontend)
+
+                speaker = str(row.get(speaker_key, voice_id))
+
+                row_out = {
+                    "phone_ids": phone_ids,
+                    "tone_ids": tone_ids,
+                    "lang_ids": lang_ids,
+                    "hifigan_durations": durations,
+                    "target_audio": str(Path(audio_path).resolve()),
+                    "speaker_id": hash(speaker) % 256,
+                    "voice_id": speaker,
+                }
+                f.write(json.dumps(row_out, ensure_ascii=False) + "\n")
+                count += 1
+
+                if max_rows > 0 and count >= max_rows:
+                    break
+
+                if count % 100 == 0:
+                    print(f"  Processed {count} rows...")
+    finally:
+        # Audio tmp files must survive until training is done, so we keep them
+        pass
+
+    print(f"Wrote {count} rows to {output_jsonl}")
+    print(f"Skipped: no_text={skip_no_text}  no_audio={skip_no_audio}  bad_ids={skip_bad_ids}  bad_audio={skip_bad_audio}")
+    return count
+
+
+def process_local_dir(
+    audio_dir: Path,
+    output_jsonl: Path,
+    ext: str = ".wav",
+    text_ext: str = ".txt",
+    speaker: str = "speaker",
+    max_rows: int = 0,
+) -> int:
+    """Process a local directory of .wav files with matching .txt files."""
+    audio_files = sorted(audio_dir.glob(f"*{ext}"))
+    if not audio_files:
+        audio_files = sorted(audio_dir.rglob(f"*{ext}"))
+
+    mel_frontend = MelFrontend(HifiGanConfig(variant="v2plus"))
+    output_jsonl.parent.mkdir(parents=True, exist_ok=True)
+
+    count = 0
+    with output_jsonl.open("w", encoding="utf-8") as f:
+        for wav_path in audio_files:
+            txt_path = wav_path.with_suffix(text_ext)
+            if not txt_path.is_file():
+                continue
+
+            text = txt_path.read_text(encoding="utf-8").strip()
+            if not text:
+                continue
+
+            try:
+                phone_ids, tone_ids, lang_ids = text_to_ids(text)
+            except Exception as e:
+                print(f"  Skipping {wav_path.name}: {e}")
+                continue
+
+            durations = estimate_durations_uniform(phone_ids, str(wav_path), mel_frontend)
+
+            row_out = {
+                "phone_ids": phone_ids,
+                "tone_ids": tone_ids,
+                "lang_ids": lang_ids,
+                "hifigan_durations": durations,
+                "target_audio": str(wav_path.resolve()),
+                "speaker_id": hash(speaker) % 256,
+                "voice_id": speaker,
+            }
+            f.write(json.dumps(row_out, ensure_ascii=False) + "\n")
+            count += 1
+
+            if max_rows > 0 and count >= max_rows:
+                break
+
+            if count % 100 == 0:
+                print(f"  Processed {count} rows...")
+
+    print(f"Wrote {count} rows to {output_jsonl}")
+    return count
+
+
+def process_metadata_csv(
+    csv_path: Path,
+    audio_dir: Path,
+    output_jsonl: Path,
+    sep: str = "|",
+    header: bool = False,
+    text_col: int = 1,
+    file_col: int = 0,
+    speaker: str = "speaker",
+    max_rows: int = 0,
+) -> int:
+    """Process an LJSpeech-style metadata.csv."""
+    mel_frontend = MelFrontend(HifiGanConfig(variant="v2plus"))
+    output_jsonl.parent.mkdir(parents=True, exist_ok=True)
+
+    count = 0
+    with csv_path.open("r", encoding="utf-8") as csv_file:
+        lines = csv_file.readlines()
+        if header:
+            lines = lines[1:]
+
+        with output_jsonl.open("w", encoding="utf-8") as out_f:
+            for line in lines:
+                line = line.strip()
+                if not line:
+                    continue
+
+                parts = line.split(sep)
+                if len(parts) < max(text_col, file_col) + 1:
+                    continue
+
+                filename = parts[file_col].strip()
+                if not filename.endswith(".wav"):
+                    filename += ".wav"
+                wav_path = audio_dir / filename
+                if not wav_path.is_file():
+                    print(f"  Missing: {wav_path}")
+                    continue
+
+                text = parts[text_col].strip()
+                if not text:
+                    continue
+
+                try:
+                    phone_ids, tone_ids, lang_ids = text_to_ids(text)
+                except Exception as e:
+                    print(f"  Skipping {filename}: {e}")
+                    continue
+
+                durations = estimate_durations_uniform(phone_ids, str(wav_path), mel_frontend)
+
+                row_out = {
+                    "phone_ids": phone_ids,
+                    "tone_ids": tone_ids,
+                    "lang_ids": lang_ids,
+                    "hifigan_durations": durations,
+                    "target_audio": str(wav_path.resolve()),
+                    "speaker_id": hash(speaker) % 256,
+                    "voice_id": speaker,
+                }
+                out_f.write(json.dumps(row_out, ensure_ascii=False) + "\n")
+                count += 1
+
+                if max_rows > 0 and count >= max_rows:
+                    break
+
+                if count % 100 == 0:
+                    print(f"  Processed {count} rows...")
+
+    print(f"Wrote {count} rows to {output_jsonl}")
+    return count
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser(
+        description="Preprocess a speech dataset for Inflect-Nano training.",
+    )
+
+    sub = ap.add_subparsers(dest="mode", required=True)
+
+    # HuggingFace mode
+    p_hf = sub.add_parser("hf", help="Process a HuggingFace dataset")
+    p_hf.add_argument("--dataset", required=True, help="HF dataset path, e.g. MikhailT/cmu-arctic")
+    p_hf.add_argument("--subset", default="default")
+    p_hf.add_argument("--split", default=None, help="e.g. 'rms' for CMU ARCTIC single speaker")
+    p_hf.add_argument("--text-key", default="text")
+    p_hf.add_argument("--speaker-key", default="speaker")
+    p_hf.add_argument("--audio-dir", type=Path, default=None, help="Local dir for audio if HF paths are remote")
+    p_hf.add_argument("--out", type=Path, required=True, help="Output .jsonl path")
+    p_hf.add_argument("--max-rows", type=int, default=0)
+    p_hf.add_argument("--voice-id", default="speaker")
+
+    # Local dir mode
+    p_dir = sub.add_parser("local", help="Process local directory of .wav+.txt pairs")
+    p_dir.add_argument("--audio-dir", type=Path, required=True)
+    p_dir.add_argument("--ext", default=".wav")
+    p_dir.add_argument("--text-ext", default=".txt")
+    p_dir.add_argument("--speaker", default="speaker")
+    p_dir.add_argument("--out", type=Path, required=True)
+    p_dir.add_argument("--max-rows", type=int, default=0)
+
+    # Metadata CSV mode (LJSpeech-style)
+    p_csv = sub.add_parser("csv", help="Process a metadata.csv file")
+    p_csv.add_argument("--csv", type=Path, required=True)
+    p_csv.add_argument("--audio-dir", type=Path, required=True)
+    p_csv.add_argument("--sep", default="|")
+    p_csv.add_argument("--header", action="store_true")
+    p_csv.add_argument("--file-col", type=int, default=0)
+    p_csv.add_argument("--text-col", type=int, default=1)
+    p_csv.add_argument("--speaker", default="speaker")
+    p_csv.add_argument("--out", type=Path, required=True)
+    p_csv.add_argument("--max-rows", type=int, default=0)
+
+    args = ap.parse_args()
+
+    if args.mode == "hf":
+        process_hf_dataset(
+            dataset_path=args.dataset,
+            output_jsonl=args.out,
+            audio_dir=args.audio_dir,
+            subset=args.subset,
+            split=args.split,
+            text_key=args.text_key,
+            speaker_key=args.speaker_key,
+            max_rows=args.max_rows,
+            voice_id=args.voice_id,
+        )
+    elif args.mode == "local":
+        process_local_dir(
+            audio_dir=args.audio_dir,
+            output_jsonl=args.out,
+            ext=args.ext,
+            text_ext=args.text_ext,
+            speaker=args.speaker,
+            max_rows=args.max_rows,
+        )
+    elif args.mode == "csv":
+        process_metadata_csv(
+            csv_path=args.csv,
+            audio_dir=args.audio_dir,
+            output_jsonl=args.out,
+            sep=args.sep,
+            header=args.header,
+            file_col=args.file_col,
+            text_col=args.text_col,
+            speaker=args.speaker,
+            max_rows=args.max_rows,
+        )
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -5,3 +5,6 @@ numpy
 g2p_en
 transformers
 gradio
+numba
+scipy
+datasets

test_inference.py

@@ -0,0 +1,126 @@
+"""Basic smoke test for Inflect-Nano-v1: model loading, text-to-tokens, and synthesis."""
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+# Ensure repo root and vendored frontend are on sys.path
+REPO_ROOT = Path(__file__).resolve().parent
+VENDORED_FRONTEND = REPO_ROOT / "third_party" / "tiny_tts_frontend"
+sys.path.insert(0, str(REPO_ROOT))
+sys.path.insert(0, str(VENDORED_FRONTEND))
+
+import numpy as np
+import torch
+
+from tiny_tts.nn import commons
+from tiny_tts.text import phonemes_to_ids
+from tiny_tts.text.english import grapheme_to_phoneme, normalize_text
+from tiny_tts.utils import ADD_BLANK
+
+from inflect_nano.text_cleaning import clean_tinytts_text
+from inflect_nano.vocoder import HifiGanGenerator, make_config
+from inflect_nano.acoustic import MicroFastSpeech, MicroFastSpeechConfig
+
+
+def test_text_to_tokens():
+    """Verify the text -> token pipeline runs end-to-end."""
+    text = "Hello world, this is a test."
+    cleaned = clean_tinytts_text(text)
+    normalized = normalize_text(cleaned)
+    phones, tones, _ = grapheme_to_phoneme(normalized)
+    phone_ids, tone_ids, lang_ids = phonemes_to_ids(phones, tones, "EN")
+    if ADD_BLANK:
+        phone_ids = commons.insert_blanks(phone_ids, 0)
+        tone_ids = commons.insert_blanks(tone_ids, 0)
+        lang_ids = commons.insert_blanks(lang_ids, 0)
+
+    assert len(phone_ids) > 0, "Should produce non-empty phone IDs"
+    assert len(phone_ids) == len(tone_ids) == len(lang_ids), "All token sequences should have same length"
+    print(f"  [PASS] text_to_tokens: {len(phone_ids)} tokens from '{text}'")
+
+
+def test_load_acoustic():
+    """Verify the acoustic model loads without error."""
+    device = torch.device("cpu")
+    ckpt_path = REPO_ROOT / "weights" / "inflect_nano_v1_acoustic.pt"
+    ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
+    cfg = MicroFastSpeechConfig(**ckpt["config"])
+    model = MicroFastSpeech(cfg).to(device)
+    model.load_state_dict(ckpt["model"])
+    model.eval()
+
+    params = sum(p.numel() for p in model.parameters())
+    print(f"  [PASS] load_acoustic: {params:,} parameters loaded")
+
+
+def test_load_vocoder():
+    """Verify the vocoder loads without error."""
+    device = torch.device("cpu")
+    ckpt_path = REPO_ROOT / "weights" / "inflect_nano_v1_vocoder.pt"
+    ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
+    cfg = make_config((ckpt.get("config") or {}).get("variant", "snake_v2mid"))
+    model = HifiGanGenerator(cfg).to(device)
+    model.load_state_dict(ckpt["generator"])
+    model.remove_weight_norm()
+    model.eval()
+
+    params = sum(p.numel() for p in model.parameters())
+    print(f"  [PASS] load_vocoder: {params:,} parameters loaded")
+
+
+def test_synthesize():
+    """End-to-end: produce a waveform from text."""
+    device = torch.device("cpu")
+
+    # Load acoustic
+    acoustic_ckpt = torch.load(REPO_ROOT / "weights" / "inflect_nano_v1_acoustic.pt", map_location=device, weights_only=False)
+    acoustic_cfg = MicroFastSpeechConfig(**acoustic_ckpt["config"])
+    acoustic = MicroFastSpeech(acoustic_cfg).to(device)
+    acoustic.load_state_dict(acoustic_ckpt["model"])
+    acoustic.eval()
+    speakers = acoustic_ckpt.get("speakers") or {"mark": 0}
+
+    # Load vocoder
+    vocoder_ckpt = torch.load(REPO_ROOT / "weights" / "inflect_nano_v1_vocoder.pt", map_location=device, weights_only=False)
+    vocoder_cfg = make_config((vocoder_ckpt.get("config") or {}).get("variant", "snake_v2mid"))
+    vocoder = HifiGanGenerator(vocoder_cfg).to(device)
+    vocoder.load_state_dict(vocoder_ckpt["generator"])
+    vocoder.remove_weight_norm()
+    vocoder.eval()
+
+    # Tokenize
+    text = "This is a quick test."
+    cleaned = clean_tinytts_text(text)
+    normalized = normalize_text(cleaned)
+    phones, tones, _ = grapheme_to_phoneme(normalized)
+    phone_ids, tone_ids, lang_ids = phonemes_to_ids(phones, tones, "EN")
+    if ADD_BLANK:
+        phone_ids = commons.insert_blanks(phone_ids, 0)
+        tone_ids = commons.insert_blanks(tone_ids, 0)
+        lang_ids = commons.insert_blanks(lang_ids, 0)
+
+    phone = torch.LongTensor(phone_ids).unsqueeze(0).to(device)
+    tone = torch.LongTensor(tone_ids).unsqueeze(0).to(device)
+    lang = torch.LongTensor(lang_ids).unsqueeze(0).to(device)
+    speaker = torch.LongTensor([int(speakers.get("mark", 0))]).to(device)
+
+    # Synthesize
+    with torch.inference_mode():
+        mel = acoustic.infer(phone, tone, lang, speaker)
+        wav = vocoder(mel).squeeze().cpu().numpy()
+
+    assert isinstance(wav, np.ndarray), "Output should be a numpy array"
+    assert wav.size > 0, "Waveform should not be empty"
+    assert np.abs(wav).max() <= 1.0, "Waveform should be normalized to [-1, 1]"
+
+    print(f"  [PASS] synthesize: {wav.size} samples ({wav.size / 24000:.2f}s) from '{text}'")
+
+
+if __name__ == "__main__":
+    print("Inflect-Nano-v1 smoke tests\n")
+    test_text_to_tokens()
+    test_load_acoustic()
+    test_load_vocoder()
+    test_synthesize()
+    print("\nAll tests passed!")

third_party/tiny_tts_frontend/LICENSE

@@ -0,0 +1,152 @@
+                                 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 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, as submitted to the 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.
+
+      "Contributor" shall mean Licensor and any Legal Entity on behalf of
+      whom a Contribution has been received by the Licensor and included
+      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 the 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 any
+      Contribution embodied 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; and
+
+      (d) If the Work includes a "NOTICE" text file, You must include a
+          readable copy of the attribution notices contained within such
+          NOTICE file.
+
+   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.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor.
+
+   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 reproducing 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 exemplary 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 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
+
+   Copyright 2025 tronghieuit
+
+   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.

third_party/tiny_tts_frontend/tiny_tts/__init__.py

@@ -0,0 +1,90 @@
+import os
+import torch
+import soundfile as sf
+from tiny_tts.text.english import normalize_text, grapheme_to_phoneme
+from tiny_tts.text import phonemes_to_ids
+from tiny_tts.nn import commons
+from tiny_tts.models.synthesizer import VoiceSynthesizer
+from tiny_tts.text.symbols import symbols
+from tiny_tts.utils.config import (
+    SAMPLING_RATE, SEGMENT_FRAMES, ADD_BLANK, SPEC_CHANNELS,
+    N_SPEAKERS, SPK2ID, MODEL_PARAMS,
+)
+from tiny_tts.infer import load_engine
+
+class TinyTTS:
+    def __init__(self, checkpoint_path=None, device=None):
+        if device is None:
+            self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        else:
+            self.device = device
+            
+        if checkpoint_path is None:
+            # Look for default checkpoint in pacakage
+            pkg_dir = os.path.dirname(os.path.abspath(__file__))
+            default_ckpt = os.path.join(os.path.dirname(pkg_dir), "checkpoints", "G.pth")
+            # 2. Check HuggingFace Cache / Download
+            if not os.path.exists(default_ckpt):
+                try:
+                    from huggingface_hub import hf_hub_download
+                    print("Downloading/Loading checkpoint from Hugging Face Hub (backtracking/tiny-tts)...")
+                    default_ckpt = hf_hub_download(repo_id="backtracking/tiny-tts", filename="G.pth")
+                except ImportError:
+                    raise ImportError("huggingface_hub is required to auto-download the model. Run: pip install huggingface_hub")
+                except Exception as e:
+                    raise ValueError(f"Failed to download checkpoint from Hugging Face: {e}")
+
+            checkpoint_path = default_ckpt
+                
+        self.model = load_engine(checkpoint_path, self.device)
+
+    def speak(self, text, output_path="output.wav", speaker="MALE", speed=1.0):
+        """Synthesize text to speech and save to output_path."""
+        print(f"Synthesizing: {text}")
+
+        # Normalize text
+        normalized = normalize_text(text)
+
+        # Phonemize
+        phones, tones, word2ph = grapheme_to_phoneme(normalized)
+
+        # Convert to sequence
+        phone_ids, tone_ids, lang_ids = phonemes_to_ids(phones, tones, "EN")
+
+        # Add blanks
+        if ADD_BLANK:
+            phone_ids = commons.insert_blanks(phone_ids, 0)
+            tone_ids = commons.insert_blanks(tone_ids, 0)
+            lang_ids = commons.insert_blanks(lang_ids, 0)
+
+        x = torch.LongTensor(phone_ids).unsqueeze(0).to(self.device)
+        x_lengths = torch.LongTensor([len(phone_ids)]).to(self.device)
+        tone = torch.LongTensor(tone_ids).unsqueeze(0).to(self.device)
+        language = torch.LongTensor(lang_ids).unsqueeze(0).to(self.device)
+
+        # Speaker ID
+        if speaker not in SPK2ID:
+            print(f"Warning: Speaker '{speaker}' not found, using ID 0. Available: {list(SPK2ID.keys())}")
+            sid = torch.LongTensor([0]).to(self.device)
+        else:
+            sid = torch.LongTensor([SPK2ID[speaker]]).to(self.device)
+
+        # BERT features (disabled - using zero tensors)
+        bert = torch.zeros(1024, len(phone_ids)).to(self.device).unsqueeze(0)
+        ja_bert = torch.zeros(768, len(phone_ids)).to(self.device).unsqueeze(0)
+
+        # speed > 1.0 = faster speech, < 1.0 = slower speech
+        length_scale = 1.0 / speed
+
+        with torch.no_grad():
+            audio, *_ = self.model.infer(
+                x, x_lengths, sid, tone, language, bert, ja_bert,
+                noise_scale=0.667,
+                noise_scale_w=0.8,
+                length_scale=length_scale
+            )
+
+        audio_np = audio[0, 0].cpu().numpy()
+        sf.write(output_path, audio_np, SAMPLING_RATE)
+        print(f"Saved audio to {output_path}")
+        return audio_np

third_party/tiny_tts_frontend/tiny_tts/alignment/__init__.py

@@ -0,0 +1,16 @@
+from numpy import zeros, int32, float32
+from torch import from_numpy
+
+from .core import viterbi_decode_kernel
+
+
+def viterbi_decode(neg_cent, mask):
+    device = neg_cent.device
+    dtype = neg_cent.dtype
+    neg_cent = neg_cent.data.cpu().numpy().astype(float32)
+    path = zeros(neg_cent.shape, dtype=int32)
+
+    t_t_max = mask.sum(1)[:, 0].data.cpu().numpy().astype(int32)
+    t_s_max = mask.sum(2)[:, 0].data.cpu().numpy().astype(int32)
+    viterbi_decode_kernel(path, neg_cent, t_t_max, t_s_max)
+    return from_numpy(path).to(device=device, dtype=dtype)

third_party/tiny_tts_frontend/tiny_tts/alignment/core.py

@@ -0,0 +1,46 @@
+import numba
+
+
+@numba.jit(
+    numba.void(
+        numba.int32[:, :, ::1],
+        numba.float32[:, :, ::1],
+        numba.int32[::1],
+        numba.int32[::1],
+    ),
+    nopython=True,
+    nogil=True,
+)
+def viterbi_decode_kernel(paths, values, t_ys, t_xs):
+    b = paths.shape[0]
+    max_neg_val = -1e9
+    for i in range(int(b)):
+        path = paths[i]
+        value = values[i]
+        t_y = t_ys[i]
+        t_x = t_xs[i]
+
+        v_prev = v_cur = 0.0
+        index = t_x - 1
+
+        for y in range(t_y):
+            for x in range(max(0, t_x + y - t_y), min(t_x, y + 1)):
+                if x == y:
+                    v_cur = max_neg_val
+                else:
+                    v_cur = value[y - 1, x]
+                if x == 0:
+                    if y == 0:
+                        v_prev = 0.0
+                    else:
+                        v_prev = max_neg_val
+                else:
+                    v_prev = value[y - 1, x - 1]
+                value[y, x] += max(v_prev, v_cur)
+
+        for y in range(t_y - 1, -1, -1):
+            path[y, index] = 1
+            if index != 0 and (
+                index == y or value[y - 1, index] < value[y - 1, index - 1]
+            ):
+                index = index - 1

third_party/tiny_tts_frontend/tiny_tts/infer.py

@@ -0,0 +1,191 @@
+import os
+import sys
+import re
+import torch
+import soundfile as sf
+import argparse
+from tiny_tts.text.english import normalize_text, grapheme_to_phoneme
+from tiny_tts.text import phonemes_to_ids
+from tiny_tts.nn import commons
+from tiny_tts.models import VoiceSynthesizer
+from tiny_tts.text.symbols import symbols
+from tiny_tts.utils import (
+    SAMPLING_RATE, SEGMENT_FRAMES, ADD_BLANK, SPEC_CHANNELS,
+    N_SPEAKERS, SPK2ID, MODEL_PARAMS,
+)
+
+
+def load_engine(checkpoint_path, device='cuda'):
+    print(f"Loading model from {checkpoint_path}")
+    net_g = VoiceSynthesizer(
+        len(symbols),
+        SPEC_CHANNELS,
+        SEGMENT_FRAMES,
+        n_speakers=N_SPEAKERS,
+        **MODEL_PARAMS
+    ).to(device)
+
+    # Count model parameters
+    total_params = sum(p.numel() for p in net_g.parameters())
+    trainable_params = sum(p.numel() for p in net_g.parameters() if p.requires_grad)
+    print(f"Model parameters: {total_params/1e6:.2f}M total, {trainable_params/1e6:.2f}M trainable")
+
+    checkpoint = torch.load(checkpoint_path, map_location=device)
+    state_dict = checkpoint['model']
+
+    # Remove module. prefix and filter shape mismatches
+    model_state = net_g.state_dict()
+    new_state_dict = {}
+    skipped = []
+    for k, v in state_dict.items():
+        key = k[7:] if k.startswith('module.') else k
+        if key in model_state:
+            if v.shape == model_state[key].shape:
+                new_state_dict[key] = v
+            else:
+                skipped.append(f"{key}: ckpt{v.shape} vs model{model_state[key].shape}")
+        else:
+            new_state_dict[key] = v
+
+    if skipped:
+        print(f"Skipped {len(skipped)} mismatched keys:")
+        for s in skipped[:5]:
+            print(f"  {s}")
+        if len(skipped) > 5:
+            print(f"  ... and {len(skipped)-5} more")
+
+    net_g.load_state_dict(new_state_dict, strict=False)
+    net_g.eval()
+
+    # Fold weight_norm into weight tensors for faster inference (~18% speedup)
+    net_g.dec.remove_weight_norm()
+
+    return net_g
+
+
+def synthesize(text, output_path, model, speaker="MALE", device='cuda', speed=1.0):
+    print(f"Synthesizing: {text}")
+
+    # Normalize text
+    normalized = normalize_text(text)
+
+    # Phonemize
+    phones, tones, word2ph = grapheme_to_phoneme(normalized)
+
+    # Convert to sequence
+    phone_ids, tone_ids, lang_ids = phonemes_to_ids(phones, tones, "EN")
+
+    # Add blanks
+    if ADD_BLANK:
+        phone_ids = commons.insert_blanks(phone_ids, 0)
+        tone_ids = commons.insert_blanks(tone_ids, 0)
+        lang_ids = commons.insert_blanks(lang_ids, 0)
+
+    x = torch.LongTensor(phone_ids).unsqueeze(0).to(device)
+    x_lengths = torch.LongTensor([len(phone_ids)]).to(device)
+    tone = torch.LongTensor(tone_ids).unsqueeze(0).to(device)
+    language = torch.LongTensor(lang_ids).unsqueeze(0).to(device)
+
+    # Speaker ID
+    if speaker not in SPK2ID:
+        print(f"Warning: Speaker {speaker} not found, using ID 0")
+        sid = torch.LongTensor([0]).to(device)
+    else:
+        sid = torch.LongTensor([SPK2ID[speaker]]).to(device)
+
+    # BERT features (disabled - using zero tensors)
+    bert = torch.zeros(1024, len(phone_ids)).to(device).unsqueeze(0)
+    ja_bert = torch.zeros(768, len(phone_ids)).to(device).unsqueeze(0)
+
+    # speed > 1.0 = faster speech, < 1.0 = slower speech
+    length_scale = 1.0 / speed
+
+    with torch.no_grad():
+        audio, *_ = model.infer(
+            x, x_lengths, sid, tone, language, bert, ja_bert,
+            noise_scale=0.667,
+            noise_scale_w=0.8,
+            length_scale=length_scale
+        )
+
+    audio = audio[0, 0].cpu().numpy()
+    sf.write(output_path, audio, SAMPLING_RATE)
+    print(f"Saved audio to {output_path}")
+
+
+def get_latest_checkpoint(checkpoint_dir):
+    """Finds the latest G_*.pth checkpoint in the given directory."""
+    checkpoints = [f for f in os.listdir(checkpoint_dir) if f.startswith('G_') and f.endswith('.pth')]
+    if not checkpoints:
+        return None
+
+    def get_step(filename):
+        match = re.search(r'_(\d+)\.pth', filename)
+        return int(match.group(1)) if match else -1
+
+    latest_ckpt = max(checkpoints, key=get_step)
+    return os.path.join(checkpoint_dir, latest_ckpt)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="TinyTTS — English Text-to-Speech Inference")
+    parser.add_argument("--text", "-t", type=str, default="The weather is nice today, and I feel very relaxed.", help="Text to synthesize")
+    parser.add_argument("--checkpoint", "-c", type=str, default=None, help="Path to checkpoint. Auto-downloads if not provided.")
+    parser.add_argument("--output", "-o", type=str, default="output.wav", help="Output audio file path")
+    parser.add_argument("--speaker", "-s", type=str, default="MALE", help="Speaker ID")
+    parser.add_argument("--speed", type=float, default=1.0, help="Speech speed (1.0=normal, 1.5=faster, 0.7=slower)")
+    parser.add_argument("--device", type=str, default="cuda", help="Device to use (cuda or cpu)")
+
+    args = parser.parse_args()
+
+    if args.checkpoint is None:
+        try:
+            from huggingface_hub import hf_hub_download
+            print("Downloading/Loading checkpoint from Hugging Face Hub (backtracking/tiny-tts)...")
+            args.checkpoint = hf_hub_download(repo_id="backtracking/tiny-tts", filename="G.pth")
+        except ImportError:
+            print("Error: huggingface_hub is required for auto-download. Run: pip install huggingface_hub")
+            sys.exit(1)
+        except Exception as e:
+            print(f"Error downloading checkpoint: {e}")
+            sys.exit(1)
+
+    if not os.path.exists(args.checkpoint):
+        print(f"Error: Checkpoint or directory not found at {args.checkpoint}")
+        sys.exit(1)
+
+    if os.path.isdir(args.checkpoint):
+        latest_ckpt = get_latest_checkpoint(args.checkpoint)
+        if not latest_ckpt:
+            print(f"Error: No G_*.pth checkpoints found in directory {args.checkpoint}")
+            sys.exit(1)
+        args.checkpoint = latest_ckpt
+        print(f"Auto-detected latest checkpoint: {args.checkpoint}")
+
+    # Extract step from checkpoint filename
+    ckpt_basename = os.path.basename(args.checkpoint)
+    match = re.search(r'_(\d+)\.pth', ckpt_basename)
+    step_str = match.group(1) if match else "unknown"
+
+    # Save to output folder
+    out_dir = "infer_outputs"
+    os.makedirs(out_dir, exist_ok=True)
+
+    out_name = os.path.basename(args.output)
+    name, ext = os.path.splitext(out_name)
+    model = load_engine(args.checkpoint, args.device)
+
+    if args.speaker.lower() == "all":
+        if not SPK2ID:
+            print("Error: No speakers found")
+            sys.exit(1)
+        print(f"Synthesizing for all {len(SPK2ID)} speakers...")
+        for spk in SPK2ID.keys():
+            final_output = os.path.join(out_dir, f"{name}_step{step_str}_spk{spk}{ext}")
+            synthesize(args.text, final_output, model, speaker=spk, device=args.device, speed=args.speed)
+    else:
+        final_output = os.path.join(out_dir, f"{name}_step{step_str}_spk{args.speaker}{ext}")
+        synthesize(args.text, final_output, model, speaker=args.speaker, device=args.device, speed=args.speed)
+
+if __name__ == "__main__":
+    main()

third_party/tiny_tts_frontend/tiny_tts/infer_onnx.py

@@ -0,0 +1,199 @@
+"""
+ONNX Runtime inference engine for TinyTTS.
+
+Replaces the PyTorch VoiceSynthesizer.infer() with equivalent
+ONNX Runtime sessions + NumPy ops for the non-exported parts
+(alignment path computation).
+"""
+import os
+import numpy as np
+import soundfile as sf
+
+from tiny_tts.text.english import normalize_text, grapheme_to_phoneme
+from tiny_tts.text import phonemes_to_ids
+from tiny_tts.nn import commons
+from tiny_tts.utils.config import (
+    SAMPLING_RATE, ADD_BLANK, SPK2ID,
+)
+
+try:
+    import onnxruntime as ort
+except ImportError:
+    raise ImportError("onnxruntime is required. Run: pip install onnxruntime")
+
+
+def _build_session(path: str, use_gpu: bool = False):
+    """Create an ORT InferenceSession with optional GPU support."""
+    providers = (
+        ["CUDAExecutionProvider", "CPUExecutionProvider"]
+        if use_gpu else
+        ["CPUExecutionProvider"]
+    )
+    opts = ort.SessionOptions()
+    opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+    opts.intra_op_num_threads = os.cpu_count() or 4
+    return ort.InferenceSession(path, sess_options=opts, providers=providers)
+
+
+def _create_length_mask_np(lengths, max_len=None):
+    """NumPy equivalent of commons.create_length_mask."""
+    if max_len is None:
+        max_len = int(lengths.max())
+    ids = np.arange(max_len, dtype=np.float32)     # [T]
+    mask = (ids[None, :] < lengths[:, None]).astype(np.float32)  # [B, T]
+    return mask
+
+
+def _compute_alignment_path_np(w_ceil, attn_mask):
+    """
+    Monotonic alignment path - vectorized via cumsum (much faster than Python loops).
+    w_ceil:    [B, 1, T_x] — integer duration per phone
+    attn_mask: [B, 1, T_y, T_x] — joint mask
+    Returns attn: [B, 1, T_y, T_x]
+    """
+    B, _, T_x = w_ceil.shape
+    T_y = attn_mask.shape[2]
+
+    # Build duration matrix: for each phone column expand the duration
+    # cumulative sum of durations gives us the end frame index for each phone
+    dur = w_ceil[:, 0, :]                            # [B, T_x]
+    cum_dur = np.cumsum(dur, axis=1)                 # [B, T_x]  — end frame (1-indexed)
+    cum_dur_prev = np.pad(cum_dur[:, :-1], ((0,0),(1,0)))  # [B, T_x] — start frame
+
+    # Frame indices: [1, T_y, 1]
+    frame_idx = np.arange(T_y, dtype=np.float32)[None, :, None]   # [1, T_y, 1]
+    # For each phone, mark frames [start, end)
+    # cum_dur_prev: [B,1,T_x], cum_dur: [B,1,T_x]
+    start = cum_dur_prev[:, None, :]                 # [B, 1, T_x]
+    end   = cum_dur[:, None, :]                      # [B, 1, T_x]
+    attn  = ((frame_idx >= start) & (frame_idx < end)).astype(np.float32)  # [B, T_y, T_x]
+    attn  = attn[:, None, :, :]                      # [B, 1, T_y, T_x]
+    return attn * attn_mask
+
+
+class OnnxTinyTTS:
+    """
+    Inference using ONNX Runtime.
+
+    Args:
+        onnx_dir: directory containing the 4 .onnx files
+        use_gpu:  if True, try CUDAExecutionProvider
+    """
+
+    def __init__(self, onnx_dir: str = "onnx", use_gpu: bool = False):
+        onnx_dir = os.path.abspath(onnx_dir)
+        print(f"Loading ONNX sessions from: {onnx_dir}")
+
+        self._enc  = _build_session(os.path.join(onnx_dir, "text_encoder.onnx"), use_gpu)
+        self._dp   = _build_session(os.path.join(onnx_dir, "duration_predictor.onnx"), use_gpu)
+        self._flow = _build_session(os.path.join(onnx_dir, "flow.onnx"), use_gpu)
+        self._dec  = _build_session(os.path.join(onnx_dir, "decoder.onnx"), use_gpu)
+
+        print("ONNX sessions ready ✅")
+
+    def _text_to_ids(self, text: str):
+        normalized = normalize_text(text)
+        phones, tones, _ = grapheme_to_phoneme(normalized)
+        phone_ids, tone_ids, lang_ids = phonemes_to_ids(phones, tones, "EN")
+
+        if ADD_BLANK:
+            phone_ids = commons.insert_blanks(phone_ids, 0)
+            tone_ids  = commons.insert_blanks(tone_ids, 0)
+            lang_ids  = commons.insert_blanks(lang_ids, 0)
+
+        return phone_ids, tone_ids, lang_ids
+
+    def speak(
+        self,
+        text: str,
+        output_path: str = "onnx_output.wav",
+        speaker: str = "female",
+        noise_scale: float = 0.667,
+        noise_scale_w: float = 0.8,
+        length_scale: float = 1.0,
+        output_sr: int = None,
+    ) -> np.ndarray:
+        """Synthesize speech and save to output_path.
+        
+        Args:
+            output_sr: If set (e.g. 22050), resample the output from 44100 Hz.
+                       Useful to reduce file size while keeping quality.
+        """
+        print(f"[ONNX] Synthesizing: {text}")
+
+        phone_ids, tone_ids, lang_ids = self._text_to_ids(text)
+        T = len(phone_ids)
+
+        # Prepare inputs as float32 / int64 arrays
+        x       = np.array(phone_ids, dtype=np.int64)[None, :]        # [1, T]
+        x_len   = np.array([T], dtype=np.int64)                        # [1]
+        tone    = np.array(tone_ids, dtype=np.int64)[None, :]          # [1, T]
+        lang    = np.array(lang_ids, dtype=np.int64)[None, :]          # [1, T]
+        bert    = np.zeros((1, 1024, T), dtype=np.float32)
+        ja_bert = np.zeros((1, 768,  T), dtype=np.float32)
+        sid_val = SPK2ID.get(speaker, 0)
+        sid     = np.array([sid_val], dtype=np.int64)                  # [1]
+
+        # ── 1. Text Encoder ──────────────────────────────────────────────
+        x_enc, m_p, logs_p, x_mask, g = self._enc.run(
+            None,
+            {
+                "phone_ids":    x,
+                "phone_lengths":x_len,
+                "tone_ids":     tone,
+                "language_ids": lang,
+                "bert":         bert,
+                "ja_bert":      ja_bert,
+                "speaker_id":   sid,
+            },
+        )
+
+        # ── 2. Duration Predictor ─────────────────────────────────────────
+        logw = self._dp.run(None, {"x": x_enc, "x_mask": x_mask, "g": g})[0]
+
+        # ── 3. Alignment Path (NumPy) ─────────────────────────────────────
+        w        = np.exp(logw) * x_mask * length_scale          # [1, 1, T]
+        w_ceil   = np.ceil(w)                                    # [1, 1, T]
+        y_len    = max(1, int(w_ceil.sum()))
+        y_lens   = np.array([y_len], dtype=np.int64)
+
+        y_mask   = _create_length_mask_np(y_lens, y_len)         # [1, T_y]
+        y_mask   = y_mask[:, None, :]                            # [1, 1, T_y]
+        # attn_mask: [1, 1, T_y, T_x]  (outer product of frame mask and phone mask)
+        attn_mask = y_mask[:, :, :, None] * x_mask[:, :, None, :]   # [1,1,T_y,T_x]
+        attn     = _compute_alignment_path_np(w_ceil, attn_mask) # [1, 1, T_y, T_x]
+
+        # Expand prior stats via alignment
+        m_p_exp    = np.matmul(attn[:, 0], m_p.transpose(0, 2, 1)).transpose(0, 2, 1)
+        logs_p_exp = np.matmul(attn[:, 0], logs_p.transpose(0, 2, 1)).transpose(0, 2, 1)
+
+        # ── 4. Sample z_p ─────────────────────────────────────────────────
+        z_p = m_p_exp + np.random.randn(*m_p_exp.shape).astype(np.float32) * \
+              np.exp(logs_p_exp) * noise_scale
+
+        # ── 5. Flow (reverse) ─────────────────────────────────────────────
+        z = self._flow.run(
+            None,
+            {"z_p": z_p, "y_mask": y_mask.astype(np.float32), "g": g},
+        )[0]
+
+        # ── 6. Decoder ────────────────────────────────────────────────────
+        z_masked = (z * y_mask).astype(np.float32)
+        audio = self._dec.run(None, {"z": z_masked, "g": g})[0]  # [1, 1, samples]
+
+        audio_np = audio[0, 0]
+        save_sr = SAMPLING_RATE
+        if output_sr is not None and output_sr != SAMPLING_RATE:
+            try:
+                import torchaudio
+                import torch
+                wav_t = torch.from_numpy(audio_np).unsqueeze(0)
+                resampler = torchaudio.transforms.Resample(SAMPLING_RATE, output_sr)
+                audio_np = resampler(wav_t).squeeze(0).numpy()
+                save_sr = output_sr
+            except Exception as e:
+                print(f"[ONNX] Resampling failed ({e}), saving at {SAMPLING_RATE}Hz")
+
+        sf.write(output_path, audio_np, save_sr)
+        print(f"[ONNX] Saved: {output_path} ({save_sr}Hz)")
+        return audio_np

third_party/tiny_tts_frontend/tiny_tts/models/__init__.py

@@ -0,0 +1 @@
+from .synthesizer import VoiceSynthesizer

third_party/tiny_tts_frontend/tiny_tts/models/synthesizer.py

@@ -0,0 +1,718 @@
+import math
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from tiny_tts.nn import commons
+from tiny_tts.nn import modules
+from tiny_tts.nn import attentions
+
+from torch.nn import Conv1d, ConvTranspose1d
+from torch.nn.utils import weight_norm, remove_weight_norm
+
+from tiny_tts.nn.commons import initialize_weights, compute_padding
+import tiny_tts.alignment as alignment
+
+
+class AttentionFlowBlock(nn.Module):
+    def __init__(
+        self,
+        channels,
+        hidden_channels,
+        filter_channels,
+        n_heads,
+        n_layers,
+        kernel_size,
+        p_dropout,
+        n_flows=4,
+        gin_channels=0,
+        share_parameter=False,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.hidden_channels = hidden_channels
+        self.kernel_size = kernel_size
+        self.n_layers = n_layers
+        self.n_flows = n_flows
+        self.gin_channels = gin_channels
+
+        self.flows = nn.ModuleList()
+
+        self.wn = (
+            attentions.FeedForward(
+                hidden_channels,
+                filter_channels,
+                n_heads,
+                n_layers,
+                kernel_size,
+                p_dropout,
+                isflow=True,
+                gin_channels=self.gin_channels,
+            )
+            if share_parameter
+            else None
+        )
+
+        for i in range(n_flows):
+            self.flows.append(
+                modules.TransformerCouplingLayer(
+                    channels,
+                    hidden_channels,
+                    kernel_size,
+                    n_layers,
+                    n_heads,
+                    p_dropout,
+                    filter_channels,
+                    mean_only=True,
+                    wn_sharing_parameter=self.wn,
+                    gin_channels=self.gin_channels,
+                )
+            )
+            self.flows.append(modules.FlipTransform())
+
+    def forward(self, x, x_mask, g=None, reverse=False):
+        if not reverse:
+            for flow in self.flows:
+                x, _ = flow(x, x_mask, g=g, reverse=reverse)
+        else:
+            for flow in reversed(self.flows):
+                x = flow(x, x_mask, g=g, reverse=reverse)
+        return x
+
+
+class VariationalDurationModel(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        filter_channels,
+        kernel_size,
+        p_dropout,
+        n_flows=4,
+        gin_channels=0,
+    ):
+        super().__init__()
+        filter_channels = in_channels
+        self.in_channels = in_channels
+        self.filter_channels = filter_channels
+        self.kernel_size = kernel_size
+        self.p_dropout = p_dropout
+        self.n_flows = n_flows
+        self.gin_channels = gin_channels
+
+        self.log_flow = modules.LogTransform()
+        self.flows = nn.ModuleList()
+        self.flows.append(modules.AffineCoupling(2))
+        for i in range(n_flows):
+            self.flows.append(
+                modules.ConvolutionalFlow(2, filter_channels, kernel_size, n_layers=3)
+            )
+            self.flows.append(modules.FlipTransform())
+
+        self.post_pre = nn.Conv1d(1, filter_channels, 1)
+        self.post_proj = nn.Conv1d(filter_channels, filter_channels, 1)
+        self.post_convs = modules.DepthwiseSepConv(
+            filter_channels, kernel_size, n_layers=3, p_dropout=p_dropout
+        )
+        self.post_flows = nn.ModuleList()
+        self.post_flows.append(modules.AffineCoupling(2))
+        for i in range(4):
+            self.post_flows.append(
+                modules.ConvolutionalFlow(2, filter_channels, kernel_size, n_layers=3)
+            )
+            self.post_flows.append(modules.FlipTransform())
+
+        self.pre = nn.Conv1d(in_channels, filter_channels, 1)
+        self.proj = nn.Conv1d(filter_channels, filter_channels, 1)
+        self.convs = modules.DepthwiseSepConv(
+            filter_channels, kernel_size, n_layers=3, p_dropout=p_dropout
+        )
+        if gin_channels != 0:
+            self.cond = nn.Conv1d(gin_channels, filter_channels, 1)
+
+    def forward(self, x, x_mask, w=None, g=None, reverse=False, noise_scale=1.0):
+        x = torch.detach(x)
+        x = self.pre(x)
+        if g is not None:
+            g = torch.detach(g)
+            x = x + self.cond(g)
+        x = self.convs(x, x_mask)
+        x = self.proj(x) * x_mask
+
+        if not reverse:
+            flows = self.flows
+            assert w is not None
+
+            logdet_tot_q = 0
+            h_w = self.post_pre(w)
+            h_w = self.post_convs(h_w, x_mask)
+            h_w = self.post_proj(h_w) * x_mask
+            e_q = (
+                torch.randn(w.size(0), 2, w.size(2)).to(device=x.device, dtype=x.dtype)
+                * x_mask
+            )
+            z_q = e_q
+            for flow in self.post_flows:
+                z_q, logdet_q = flow(z_q, x_mask, g=(x + h_w))
+                logdet_tot_q += logdet_q
+            z_u, z1 = torch.split(z_q, [1, 1], 1)
+            u = torch.sigmoid(z_u) * x_mask
+            z0 = (w - u) * x_mask
+            logdet_tot_q += torch.sum(
+                (F.logsigmoid(z_u) + F.logsigmoid(-z_u)) * x_mask, [1, 2]
+            )
+            logq = (
+                torch.sum(-0.5 * (math.log(2 * math.pi) + (e_q**2)) * x_mask, [1, 2])
+                - logdet_tot_q
+            )
+
+            logdet_tot = 0
+            z0, logdet = self.log_flow(z0, x_mask)
+            logdet_tot += logdet
+            z = torch.cat([z0, z1], 1)
+            for flow in flows:
+                z, logdet = flow(z, x_mask, g=x, reverse=reverse)
+                logdet_tot = logdet_tot + logdet
+            nll = (
+                torch.sum(0.5 * (math.log(2 * math.pi) + (z**2)) * x_mask, [1, 2])
+                - logdet_tot
+            )
+            return nll + logq
+        else:
+            flows = list(reversed(self.flows))
+            flows = flows[:-2] + [flows[-1]]
+            z = (
+                torch.randn(x.size(0), 2, x.size(2)).to(device=x.device, dtype=x.dtype)
+                * noise_scale
+            )
+            for flow in flows:
+                z = flow(z, x_mask, g=x, reverse=reverse)
+            z0, z1 = torch.split(z, [1, 1], 1)
+            logw = z0
+            return logw
+
+
+class DurationEstimator(nn.Module):
+    def __init__(
+        self, in_channels, filter_channels, kernel_size, p_dropout, gin_channels=0
+    ):
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.filter_channels = filter_channels
+        self.kernel_size = kernel_size
+        self.p_dropout = p_dropout
+        self.gin_channels = gin_channels
+
+        self.drop = nn.Dropout(p_dropout)
+        self.conv_1 = nn.Conv1d(
+            in_channels, filter_channels, kernel_size, padding=kernel_size // 2
+        )
+        self.norm_1 = modules.ChannelNorm(filter_channels)
+        self.conv_2 = nn.Conv1d(
+            filter_channels, filter_channels, kernel_size, padding=kernel_size // 2
+        )
+        self.norm_2 = modules.ChannelNorm(filter_channels)
+        self.proj = nn.Conv1d(filter_channels, 1, 1)
+
+        if gin_channels != 0:
+            self.cond = nn.Conv1d(gin_channels, in_channels, 1)
+
+    def forward(self, x, x_mask, g=None):
+        x = torch.detach(x)
+        if g is not None:
+            g = torch.detach(g)
+            x = x + self.cond(g)
+        x = self.conv_1(x * x_mask)
+        x = torch.relu(x)
+        x = self.norm_1(x)
+        x = self.drop(x)
+        x = self.conv_2(x * x_mask)
+        x = torch.relu(x)
+        x = self.norm_2(x)
+        x = self.drop(x)
+        x = self.proj(x * x_mask)
+        return x * x_mask
+
+
+class PhonemeEncoder(nn.Module):
+    def __init__(
+        self,
+        n_vocab,
+        out_channels,
+        hidden_channels,
+        filter_channels,
+        n_heads,
+        n_layers,
+        kernel_size,
+        p_dropout,
+        gin_channels=0,
+        num_languages=None,
+        num_tones=None,
+    ):
+        super().__init__()
+        if num_languages is None:
+            from tiny_tts.text import num_languages
+        if num_tones is None:
+            from tiny_tts.text import num_tones
+        self.n_vocab = n_vocab
+        self.out_channels = out_channels
+        self.hidden_channels = hidden_channels
+        self.filter_channels = filter_channels
+        self.n_heads = n_heads
+        self.n_layers = n_layers
+        self.kernel_size = kernel_size
+        self.p_dropout = p_dropout
+        self.gin_channels = gin_channels
+        self.emb = nn.Embedding(n_vocab, hidden_channels)
+        nn.init.normal_(self.emb.weight, 0.0, hidden_channels**-0.5)
+        self.tone_emb = nn.Embedding(num_tones, hidden_channels)
+        nn.init.normal_(self.tone_emb.weight, 0.0, hidden_channels**-0.5)
+        self.language_emb = nn.Embedding(num_languages, hidden_channels)
+        nn.init.normal_(self.language_emb.weight, 0.0, hidden_channels**-0.5)
+        self.bert_proj = nn.Conv1d(1024, hidden_channels, 1)
+        self.ja_bert_proj = nn.Conv1d(768, hidden_channels, 1)
+
+        self.encoder = attentions.TransformerBlock(
+            hidden_channels,
+            filter_channels,
+            n_heads,
+            n_layers,
+            kernel_size,
+            p_dropout,
+            gin_channels=self.gin_channels,
+        )
+        self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
+
+    def forward(self, x, x_lengths, tone, language, bert, ja_bert, g=None):
+        bert_emb = self.bert_proj(bert).transpose(1, 2)
+        ja_bert_emb = self.ja_bert_proj(ja_bert).transpose(1, 2)
+        x = (
+            self.emb(x)
+            + self.tone_emb(tone)
+            + self.language_emb(language)
+            + bert_emb
+            + ja_bert_emb
+        ) * math.sqrt(
+            self.hidden_channels
+        )
+        x = torch.transpose(x, 1, -1)
+        x_mask = torch.unsqueeze(commons.create_length_mask(x_lengths, x.size(2)), 1).to(
+            x.dtype
+        )
+
+        x = self.encoder(x * x_mask, x_mask, g=g)
+        stats = self.proj(x) * x_mask
+
+        m, logs = torch.split(stats, self.out_channels, dim=1)
+        return x, m, logs, x_mask
+
+
+class FlowBlock(nn.Module):
+    def __init__(
+        self,
+        channels,
+        hidden_channels,
+        kernel_size,
+        dilation_rate,
+        n_layers,
+        n_flows=4,
+        gin_channels=0,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.hidden_channels = hidden_channels
+        self.kernel_size = kernel_size
+        self.dilation_rate = dilation_rate
+        self.n_layers = n_layers
+        self.n_flows = n_flows
+        self.gin_channels = gin_channels
+
+        self.flows = nn.ModuleList()
+        for i in range(n_flows):
+            self.flows.append(
+                modules.FlowCouplingLayer(
+                    channels,
+                    hidden_channels,
+                    kernel_size,
+                    dilation_rate,
+                    n_layers,
+                    gin_channels=gin_channels,
+                    mean_only=True,
+                )
+            )
+            self.flows.append(modules.FlipTransform())
+
+    def forward(self, x, x_mask, g=None, reverse=False):
+        if not reverse:
+            for flow in self.flows:
+                x, _ = flow(x, x_mask, g=g, reverse=reverse)
+        else:
+            for flow in reversed(self.flows):
+                x = flow(x, x_mask, g=g, reverse=reverse)
+        return x
+
+
+class LatentEncoder(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        hidden_channels,
+        kernel_size,
+        dilation_rate,
+        n_layers,
+        gin_channels=0,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.hidden_channels = hidden_channels
+        self.kernel_size = kernel_size
+        self.dilation_rate = dilation_rate
+        self.n_layers = n_layers
+        self.gin_channels = gin_channels
+
+        self.pre = nn.Conv1d(in_channels, hidden_channels, 1)
+        self.enc = modules.WaveNet(
+            hidden_channels,
+            kernel_size,
+            dilation_rate,
+            n_layers,
+            gin_channels=gin_channels,
+        )
+        self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
+
+    def forward(self, x, x_lengths, g=None, tau=1.0):
+        x_mask = torch.unsqueeze(commons.create_length_mask(x_lengths, x.size(2)), 1).to(
+            x.dtype
+        )
+        x = self.pre(x) * x_mask
+        x = self.enc(x, x_mask, g=g)
+        stats = self.proj(x) * x_mask
+        m, logs = torch.split(stats, self.out_channels, dim=1)
+        z = (m + torch.randn_like(m) * tau * torch.exp(logs)) * x_mask
+        return z, m, logs, x_mask
+
+
+class WaveformDecoder(torch.nn.Module):
+    def __init__(
+        self,
+        initial_channel,
+        resblock,
+        resblock_kernel_sizes,
+        resblock_dilation_sizes,
+        upsample_rates,
+        upsample_initial_channel,
+        upsample_kernel_sizes,
+        gin_channels=0,
+    ):
+        super(WaveformDecoder, self).__init__()
+        self.num_kernels = len(resblock_kernel_sizes)
+        self.num_upsamples = len(upsample_rates)
+        self.conv_pre = Conv1d(
+            initial_channel, upsample_initial_channel, 7, 1, padding=3
+        )
+        resblock = modules.ConvResBlock if resblock == "1" else modules.ConvResBlockLight
+
+        self.ups = nn.ModuleList()
+        for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
+            self.ups.append(
+                weight_norm(
+                    ConvTranspose1d(
+                        upsample_initial_channel // (2**i),
+                        upsample_initial_channel // (2 ** (i + 1)),
+                        k,
+                        u,
+                        padding=(k - u) // 2,
+                    )
+                )
+            )
+
+        self.resblocks = nn.ModuleList()
+        for i in range(len(self.ups)):
+            ch = upsample_initial_channel // (2 ** (i + 1))
+            for j, (k, d) in enumerate(
+                zip(resblock_kernel_sizes, resblock_dilation_sizes)
+            ):
+                self.resblocks.append(resblock(ch, k, d))
+
+        self.conv_post = Conv1d(ch, 1, 7, 1, padding=3, bias=False)
+        self.ups.apply(initialize_weights)
+
+        if gin_channels != 0:
+            self.cond = nn.Conv1d(gin_channels, upsample_initial_channel, 1)
+
+    def forward(self, x, g=None):
+        x = self.conv_pre(x)
+        if g is not None:
+            x = x + self.cond(g)
+
+        for i in range(self.num_upsamples):
+            x = F.leaky_relu(x, modules.LRELU_SLOPE)
+            x = self.ups[i](x)
+            xs = None
+            for j in range(self.num_kernels):
+                if xs is None:
+                    xs = self.resblocks[i * self.num_kernels + j](x)
+                else:
+                    xs += self.resblocks[i * self.num_kernels + j](x)
+            x = xs / self.num_kernels
+        x = F.leaky_relu(x)
+        x = self.conv_post(x)
+        x = torch.tanh(x)
+
+        return x
+
+    def remove_weight_norm(self):
+        for layer in self.ups:
+            remove_weight_norm(layer)
+        for layer in self.resblocks:
+            layer.remove_weight_norm()
+
+
+class StyleEncoder(nn.Module):
+    def __init__(self, spec_channels, gin_channels=0, layernorm=False):
+        super().__init__()
+        self.spec_channels = spec_channels
+        ref_enc_filters = [32, 32, 64, 64, 128, 128]
+        K = len(ref_enc_filters)
+        filters = [1] + ref_enc_filters
+        convs = [
+            weight_norm(
+                nn.Conv2d(
+                    in_channels=filters[i],
+                    out_channels=filters[i + 1],
+                    kernel_size=(3, 3),
+                    stride=(2, 2),
+                    padding=(1, 1),
+                )
+            )
+            for i in range(K)
+        ]
+        self.convs = nn.ModuleList(convs)
+
+        out_channels = self.calculate_channels(spec_channels, 3, 2, 1, K)
+        self.gru = nn.GRU(
+            input_size=ref_enc_filters[-1] * out_channels,
+            hidden_size=256 // 2,
+            batch_first=True,
+        )
+        self.proj = nn.Linear(128, gin_channels)
+        if layernorm:
+            self.layernorm = nn.LayerNorm(self.spec_channels)
+        else:
+            self.layernorm = None
+
+    def forward(self, inputs, mask=None):
+        N = inputs.size(0)
+
+        out = inputs.view(N, 1, -1, self.spec_channels)
+        if self.layernorm is not None:
+            out = self.layernorm(out)
+
+        for conv in self.convs:
+            out = conv(out)
+            out = F.relu(out)
+
+        out = out.transpose(1, 2)
+        T = out.size(1)
+        N = out.size(0)
+        out = out.contiguous().view(N, T, -1)
+
+        self.gru.flatten_parameters()
+        memory, out = self.gru(out)
+
+        return self.proj(out.squeeze(0))
+
+    def calculate_channels(self, L, kernel_size, stride, pad, n_convs):
+        for i in range(n_convs):
+            L = (L - kernel_size + 2 * pad) // stride + 1
+        return L
+
+
+class VoiceSynthesizer(nn.Module):
+    """Voice synthesis model for inference."""
+
+    def __init__(
+        self,
+        n_vocab,
+        spec_channels,
+        segment_size,
+        inter_channels,
+        hidden_channels,
+        filter_channels,
+        n_heads,
+        n_layers,
+        kernel_size,
+        p_dropout,
+        resblock,
+        resblock_kernel_sizes,
+        resblock_dilation_sizes,
+        upsample_rates,
+        upsample_initial_channel,
+        upsample_kernel_sizes,
+        n_speakers=256,
+        gin_channels=256,
+        use_sdp=True,
+        n_flow_layer=4,
+        n_layers_trans_flow=6,
+        flow_share_parameter=False,
+        use_transformer_flow=True,
+        use_vc=False,
+        num_languages=None,
+        num_tones=None,
+        norm_refenc=False,
+        **kwargs
+    ):
+        super().__init__()
+        self.n_vocab = n_vocab
+        self.spec_channels = spec_channels
+        self.inter_channels = inter_channels
+        self.hidden_channels = hidden_channels
+        self.filter_channels = filter_channels
+        self.n_heads = n_heads
+        self.n_layers = n_layers
+        self.kernel_size = kernel_size
+        self.p_dropout = p_dropout
+        self.resblock = resblock
+        self.resblock_kernel_sizes = resblock_kernel_sizes
+        self.resblock_dilation_sizes = resblock_dilation_sizes
+        self.upsample_rates = upsample_rates
+        self.upsample_initial_channel = upsample_initial_channel
+        self.upsample_kernel_sizes = upsample_kernel_sizes
+        self.segment_size = segment_size
+        self.n_speakers = n_speakers
+        self.gin_channels = gin_channels
+        self.n_layers_trans_flow = n_layers_trans_flow
+        self.use_spk_conditioned_encoder = kwargs.get(
+            "use_spk_conditioned_encoder", True
+        )
+        self.use_sdp = use_sdp
+        self.use_noise_scaled_mas = kwargs.get("use_noise_scaled_mas", False)
+        self.mas_noise_scale_initial = kwargs.get("mas_noise_scale_initial", 0.01)
+        self.noise_scale_delta = kwargs.get("noise_scale_delta", 2e-6)
+        self.current_mas_noise_scale = self.mas_noise_scale_initial
+        if self.use_spk_conditioned_encoder and gin_channels > 0:
+            self.enc_gin_channels = gin_channels
+        else:
+            self.enc_gin_channels = 0
+        self.enc_p = PhonemeEncoder(
+            n_vocab,
+            inter_channels,
+            hidden_channels,
+            filter_channels,
+            n_heads,
+            n_layers,
+            kernel_size,
+            p_dropout,
+            gin_channels=self.enc_gin_channels,
+            num_languages=num_languages,
+            num_tones=num_tones,
+        )
+        self.dec = WaveformDecoder(
+            inter_channels,
+            resblock,
+            resblock_kernel_sizes,
+            resblock_dilation_sizes,
+            upsample_rates,
+            upsample_initial_channel,
+            upsample_kernel_sizes,
+            gin_channels=gin_channels,
+        )
+        self.enc_q = LatentEncoder(
+            spec_channels,
+            inter_channels,
+            hidden_channels,
+            5,
+            1,
+            16,
+            gin_channels=gin_channels,
+        )
+        if use_transformer_flow:
+            self.flow = AttentionFlowBlock(
+                inter_channels,
+                hidden_channels,
+                filter_channels,
+                n_heads,
+                n_layers_trans_flow,
+                5,
+                p_dropout,
+                n_flow_layer,
+                gin_channels=gin_channels,
+                share_parameter=flow_share_parameter,
+            )
+        else:
+            self.flow = FlowBlock(
+                inter_channels,
+                hidden_channels,
+                5,
+                1,
+                n_flow_layer,
+                gin_channels=gin_channels,
+            )
+        self.sdp = VariationalDurationModel(
+            hidden_channels, 192, 3, 0.5, 4, gin_channels=gin_channels
+        )
+        self.dp = DurationEstimator(
+            hidden_channels, 256, 3, 0.5, gin_channels=gin_channels
+        )
+
+        if n_speakers > 0:
+            self.emb_g = nn.Embedding(n_speakers, gin_channels)
+        else:
+            self.ref_enc = StyleEncoder(spec_channels, gin_channels, layernorm=norm_refenc)
+        self.use_vc = use_vc
+
+    def infer(
+        self,
+        x,
+        x_lengths,
+        sid,
+        tone,
+        language,
+        bert,
+        ja_bert,
+        noise_scale=0.667,
+        length_scale=1,
+        noise_scale_w=0.8,
+        max_len=None,
+        sdp_ratio=0,
+        y=None,
+        g=None,
+    ):
+        if g is None:
+            if self.n_speakers > 0:
+                g = self.emb_g(sid).unsqueeze(-1)
+            else:
+                g = self.ref_enc(y.transpose(1, 2)).unsqueeze(-1)
+        if self.use_vc:
+            g_p = None
+        else:
+            g_p = g
+        x, m_p, logs_p, x_mask = self.enc_p(
+            x, x_lengths, tone, language, bert, ja_bert, g=g_p
+        )
+        logw = self.sdp(x, x_mask, g=g, reverse=True, noise_scale=noise_scale_w) * (
+            sdp_ratio
+        ) + self.dp(x, x_mask, g=g) * (1 - sdp_ratio)
+        w = torch.exp(logw) * x_mask * length_scale
+
+        w_ceil = torch.ceil(w)
+        y_lengths = torch.clamp_min(torch.sum(w_ceil, [1, 2]), 1).long()
+        y_mask = torch.unsqueeze(commons.create_length_mask(y_lengths, None), 1).to(
+            x_mask.dtype
+        )
+        attn_mask = torch.unsqueeze(x_mask, 2) * torch.unsqueeze(y_mask, -1)
+        attn = commons.compute_alignment_path(w_ceil, attn_mask)
+
+        m_p = torch.matmul(attn.squeeze(1), m_p.transpose(1, 2)).transpose(
+            1, 2
+        )
+        logs_p = torch.matmul(attn.squeeze(1), logs_p.transpose(1, 2)).transpose(
+            1, 2
+        )
+
+        z_p = m_p + torch.randn_like(m_p) * torch.exp(logs_p) * noise_scale
+        z = self.flow(z_p, y_mask, g=g, reverse=True)
+        o = self.dec((z * y_mask)[:, :, :max_len], g=g)
+        return o, attn, y_mask, (z, z_p, m_p, logs_p)

third_party/tiny_tts_frontend/tiny_tts/nn/__init__.py

@@ -0,0 +1 @@
+# Neural network building blocks

third_party/tiny_tts_frontend/tiny_tts/nn/attentions.py

@@ -0,0 +1,424 @@
+import math
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from . import commons
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+class ChannelLayerNorm(nn.Module):
+    def __init__(self, channels, eps=1e-5):
+        super().__init__()
+        self.channels = channels
+        self.eps = eps
+
+        self.gamma = nn.Parameter(torch.ones(channels))
+        self.beta = nn.Parameter(torch.zeros(channels))
+
+    def forward(self, x):
+        x = x.transpose(1, -1)
+        x = F.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps)
+        return x.transpose(1, -1)
+
+
+@torch.jit.script
+def gated_activation(input_a, input_b, n_channels):
+    n_channels_int = n_channels[0]
+    in_act = input_a + input_b
+    t_act = torch.tanh(in_act[:, :n_channels_int, :])
+    s_act = torch.sigmoid(in_act[:, n_channels_int:, :])
+    acts = t_act * s_act
+    return acts
+
+
+class TransformerBlock(nn.Module):
+    def __init__(
+        self,
+        hidden_channels,
+        filter_channels,
+        n_heads,
+        n_layers,
+        kernel_size=1,
+        p_dropout=0.0,
+        window_size=4,
+        isflow=True,
+        **kwargs
+    ):
+        super().__init__()
+        self.hidden_channels = hidden_channels
+        self.filter_channels = filter_channels
+        self.n_heads = n_heads
+        self.n_layers = n_layers
+        self.kernel_size = kernel_size
+        self.p_dropout = p_dropout
+        self.window_size = window_size
+
+        self.cond_layer_idx = self.n_layers
+        if "gin_channels" in kwargs:
+            self.gin_channels = kwargs["gin_channels"]
+            if self.gin_channels != 0:
+                self.spk_emb_linear = nn.Linear(self.gin_channels, self.hidden_channels)
+                self.cond_layer_idx = (
+                    kwargs["cond_layer_idx"] if "cond_layer_idx" in kwargs else 2
+                )
+                assert (
+                    self.cond_layer_idx < self.n_layers
+                ), "cond_layer_idx should be less than n_layers"
+        self.drop = nn.Dropout(p_dropout)
+        self.attn_layers = nn.ModuleList()
+        self.norm_layers_1 = nn.ModuleList()
+        self.ffn_layers = nn.ModuleList()
+        self.norm_layers_2 = nn.ModuleList()
+
+        for i in range(self.n_layers):
+            self.attn_layers.append(
+                MultiHeadSelfAttention(
+                    hidden_channels,
+                    hidden_channels,
+                    n_heads,
+                    p_dropout=p_dropout,
+                    window_size=window_size,
+                )
+            )
+            self.norm_layers_1.append(ChannelLayerNorm(hidden_channels))
+            self.ffn_layers.append(
+                FeedForward(
+                    hidden_channels,
+                    hidden_channels,
+                    filter_channels,
+                    kernel_size,
+                    p_dropout=p_dropout,
+                )
+            )
+            self.norm_layers_2.append(ChannelLayerNorm(hidden_channels))
+
+    def forward(self, x, x_mask, g=None):
+        attn_mask = x_mask.unsqueeze(2) * x_mask.unsqueeze(-1)
+        x = x * x_mask
+        for i in range(self.n_layers):
+            if i == self.cond_layer_idx and g is not None:
+                g = self.spk_emb_linear(g.transpose(1, 2))
+                g = g.transpose(1, 2)
+                x = x + g
+                x = x * x_mask
+            y = self.attn_layers[i](x, x, attn_mask)
+            y = self.drop(y)
+            x = self.norm_layers_1[i](x + y)
+
+            y = self.ffn_layers[i](x, x_mask)
+            y = self.drop(y)
+            x = self.norm_layers_2[i](x + y)
+        x = x * x_mask
+        return x
+
+
+class TransformerDecoder(nn.Module):
+    def __init__(
+        self,
+        hidden_channels,
+        filter_channels,
+        n_heads,
+        n_layers,
+        kernel_size=1,
+        p_dropout=0.0,
+        proximal_bias=False,
+        proximal_init=True,
+        **kwargs
+    ):
+        super().__init__()
+        self.hidden_channels = hidden_channels
+        self.filter_channels = filter_channels
+        self.n_heads = n_heads
+        self.n_layers = n_layers
+        self.kernel_size = kernel_size
+        self.p_dropout = p_dropout
+        self.proximal_bias = proximal_bias
+        self.proximal_init = proximal_init
+
+        self.drop = nn.Dropout(p_dropout)
+        self.self_attn_layers = nn.ModuleList()
+        self.norm_layers_0 = nn.ModuleList()
+        self.encdec_attn_layers = nn.ModuleList()
+        self.norm_layers_1 = nn.ModuleList()
+        self.ffn_layers = nn.ModuleList()
+        self.norm_layers_2 = nn.ModuleList()
+        for i in range(self.n_layers):
+            self.self_attn_layers.append(
+                MultiHeadSelfAttention(
+                    hidden_channels,
+                    hidden_channels,
+                    n_heads,
+                    p_dropout=p_dropout,
+                    proximal_bias=proximal_bias,
+                    proximal_init=proximal_init,
+                )
+            )
+            self.norm_layers_0.append(ChannelLayerNorm(hidden_channels))
+            self.encdec_attn_layers.append(
+                MultiHeadSelfAttention(
+                    hidden_channels, hidden_channels, n_heads, p_dropout=p_dropout
+                )
+            )
+            self.norm_layers_1.append(ChannelLayerNorm(hidden_channels))
+            self.ffn_layers.append(
+                FeedForward(
+                    hidden_channels,
+                    hidden_channels,
+                    filter_channels,
+                    kernel_size,
+                    p_dropout=p_dropout,
+                    causal=True,
+                )
+            )
+            self.norm_layers_2.append(ChannelLayerNorm(hidden_channels))
+
+    def forward(self, x, x_mask, h, h_mask):
+        self_attn_mask = commons.subsequent_mask(x_mask.size(2)).to(
+            device=x.device, dtype=x.dtype
+        )
+        encdec_attn_mask = h_mask.unsqueeze(2) * x_mask.unsqueeze(-1)
+        x = x * x_mask
+        for i in range(self.n_layers):
+            y = self.self_attn_layers[i](x, x, self_attn_mask)
+            y = self.drop(y)
+            x = self.norm_layers_0[i](x + y)
+
+            y = self.encdec_attn_layers[i](x, h, encdec_attn_mask)
+            y = self.drop(y)
+            x = self.norm_layers_1[i](x + y)
+
+            y = self.ffn_layers[i](x, x_mask)
+            y = self.drop(y)
+            x = self.norm_layers_2[i](x + y)
+        x = x * x_mask
+        return x
+
+
+class MultiHeadSelfAttention(nn.Module):
+    def __init__(
+        self,
+        channels,
+        out_channels,
+        n_heads,
+        p_dropout=0.0,
+        window_size=None,
+        heads_share=True,
+        block_length=None,
+        proximal_bias=False,
+        proximal_init=False,
+    ):
+        super().__init__()
+        assert channels % n_heads == 0
+
+        self.channels = channels
+        self.out_channels = out_channels
+        self.n_heads = n_heads
+        self.p_dropout = p_dropout
+        self.window_size = window_size
+        self.heads_share = heads_share
+        self.block_length = block_length
+        self.proximal_bias = proximal_bias
+        self.proximal_init = proximal_init
+        self.attn = None
+
+        self.k_channels = channels // n_heads
+        self.conv_q = nn.Conv1d(channels, channels, 1)
+        self.conv_k = nn.Conv1d(channels, channels, 1)
+        self.conv_v = nn.Conv1d(channels, channels, 1)
+        self.conv_o = nn.Conv1d(channels, out_channels, 1)
+        self.drop = nn.Dropout(p_dropout)
+
+        if window_size is not None:
+            n_heads_rel = 1 if heads_share else n_heads
+            rel_stddev = self.k_channels**-0.5
+            self.emb_rel_k = nn.Parameter(
+                torch.randn(n_heads_rel, window_size * 2 + 1, self.k_channels)
+                * rel_stddev
+            )
+            self.emb_rel_v = nn.Parameter(
+                torch.randn(n_heads_rel, window_size * 2 + 1, self.k_channels)
+                * rel_stddev
+            )
+
+        nn.init.xavier_uniform_(self.conv_q.weight)
+        nn.init.xavier_uniform_(self.conv_k.weight)
+        nn.init.xavier_uniform_(self.conv_v.weight)
+        if proximal_init:
+            with torch.no_grad():
+                self.conv_k.weight.copy_(self.conv_q.weight)
+                self.conv_k.bias.copy_(self.conv_q.bias)
+
+    def forward(self, x, c, attn_mask=None):
+        q = self.conv_q(x)
+        k = self.conv_k(c)
+        v = self.conv_v(c)
+
+        x, self.attn = self.attention(q, k, v, mask=attn_mask)
+
+        x = self.conv_o(x)
+        return x
+
+    def attention(self, query, key, value, mask=None):
+        b, d, t_s, t_t = (*key.size(), query.size(2))
+        query = query.view(b, self.n_heads, self.k_channels, t_t).transpose(2, 3)
+        key = key.view(b, self.n_heads, self.k_channels, t_s).transpose(2, 3)
+        value = value.view(b, self.n_heads, self.k_channels, t_s).transpose(2, 3)
+
+        scores = torch.matmul(query / math.sqrt(self.k_channels), key.transpose(-2, -1))
+        if self.window_size is not None:
+            assert (
+                t_s == t_t
+            ), "Relative attention is only available for self-attention."
+            key_relative_embeddings = self._get_relative_embeddings(self.emb_rel_k, t_s)
+            rel_logits = self._matmul_with_relative_keys(
+                query / math.sqrt(self.k_channels), key_relative_embeddings
+            )
+            scores_local = self._relative_position_to_absolute_position(rel_logits)
+            scores = scores + scores_local
+        if self.proximal_bias:
+            assert t_s == t_t, "Proximal bias is only available for self-attention."
+            scores = scores + self._attention_bias_proximal(t_s).to(
+                device=scores.device, dtype=scores.dtype
+            )
+        if mask is not None:
+            scores = scores.masked_fill(mask == 0, -1e4)
+            if self.block_length is not None:
+                assert (
+                    t_s == t_t
+                ), "Local attention is only available for self-attention."
+                block_mask = (
+                    torch.ones_like(scores)
+                    .triu(-self.block_length)
+                    .tril(self.block_length)
+                )
+                scores = scores.masked_fill(block_mask == 0, -1e4)
+        p_attn = F.softmax(scores, dim=-1)
+        p_attn = self.drop(p_attn)
+        output = torch.matmul(p_attn, value)
+        if self.window_size is not None:
+            relative_weights = self._absolute_position_to_relative_position(p_attn)
+            value_relative_embeddings = self._get_relative_embeddings(
+                self.emb_rel_v, t_s
+            )
+            output = output + self._matmul_with_relative_values(
+                relative_weights, value_relative_embeddings
+            )
+        output = (
+            output.transpose(2, 3).contiguous().view(b, d, t_t)
+        )
+        return output, p_attn
+
+    def _matmul_with_relative_values(self, x, y):
+        ret = torch.matmul(x, y.unsqueeze(0))
+        return ret
+
+    def _matmul_with_relative_keys(self, x, y):
+        ret = torch.matmul(x, y.unsqueeze(0).transpose(-2, -1))
+        return ret
+
+    def _get_relative_embeddings(self, relative_embeddings, length):
+        2 * self.window_size + 1
+        pad_length = max(length - (self.window_size + 1), 0)
+        slice_start_position = max((self.window_size + 1) - length, 0)
+        slice_end_position = slice_start_position + 2 * length - 1
+        if pad_length > 0:
+            padded_relative_embeddings = F.pad(
+                relative_embeddings,
+                commons.flatten_pad_shape([[0, 0], [pad_length, pad_length], [0, 0]]),
+            )
+        else:
+            padded_relative_embeddings = relative_embeddings
+        used_relative_embeddings = padded_relative_embeddings[
+            :, slice_start_position:slice_end_position
+        ]
+        return used_relative_embeddings
+
+    def _relative_position_to_absolute_position(self, x):
+        batch, heads, length, _ = x.size()
+        x = F.pad(x, commons.flatten_pad_shape([[0, 0], [0, 0], [0, 0], [0, 1]]))
+
+        x_flat = x.view([batch, heads, length * 2 * length])
+        x_flat = F.pad(
+            x_flat, commons.flatten_pad_shape([[0, 0], [0, 0], [0, length - 1]])
+        )
+
+        x_final = x_flat.view([batch, heads, length + 1, 2 * length - 1])[
+            :, :, :length, length - 1 :
+        ]
+        return x_final
+
+    def _absolute_position_to_relative_position(self, x):
+        batch, heads, length, _ = x.size()
+        x = F.pad(
+            x, commons.flatten_pad_shape([[0, 0], [0, 0], [0, 0], [0, length - 1]])
+        )
+        x_flat = x.view([batch, heads, length**2 + length * (length - 1)])
+        x_flat = F.pad(x_flat, commons.flatten_pad_shape([[0, 0], [0, 0], [length, 0]]))
+        x_final = x_flat.view([batch, heads, length, 2 * length])[:, :, :, 1:]
+        return x_final
+
+    def _attention_bias_proximal(self, length):
+        r = torch.arange(length, dtype=torch.float32)
+        diff = torch.unsqueeze(r, 0) - torch.unsqueeze(r, 1)
+        return torch.unsqueeze(torch.unsqueeze(-torch.log1p(torch.abs(diff)), 0), 0)
+
+
+class FeedForward(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        filter_channels,
+        kernel_size,
+        p_dropout=0.0,
+        activation=None,
+        causal=False,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.filter_channels = filter_channels
+        self.kernel_size = kernel_size
+        self.p_dropout = p_dropout
+        self.activation = activation
+        self.causal = causal
+
+        if causal:
+            self.padding = self._causal_padding
+        else:
+            self.padding = self._same_padding
+
+        self.conv_1 = nn.Conv1d(in_channels, filter_channels, kernel_size)
+        self.conv_2 = nn.Conv1d(filter_channels, out_channels, kernel_size)
+        self.drop = nn.Dropout(p_dropout)
+
+    def forward(self, x, x_mask):
+        x = self.conv_1(self.padding(x * x_mask))
+        if self.activation == "gelu":
+            x = x * torch.sigmoid(1.702 * x)
+        else:
+            x = torch.relu(x)
+        x = self.drop(x)
+        x = self.conv_2(self.padding(x * x_mask))
+        return x * x_mask
+
+    def _causal_padding(self, x):
+        if self.kernel_size == 1:
+            return x
+        pad_l = self.kernel_size - 1
+        pad_r = 0
+        padding = [[0, 0], [0, 0], [pad_l, pad_r]]
+        x = F.pad(x, commons.flatten_pad_shape(padding))
+        return x
+
+    def _same_padding(self, x):
+        if self.kernel_size == 1:
+            return x
+        pad_l = (self.kernel_size - 1) // 2
+        pad_r = self.kernel_size // 2
+        padding = [[0, 0], [0, 0], [pad_l, pad_r]]
+        x = F.pad(x, commons.flatten_pad_shape(padding))
+        return x

third_party/tiny_tts_frontend/tiny_tts/nn/commons.py

@@ -0,0 +1,151 @@
+import math
+import torch
+from torch.nn import functional as F
+
+
+def initialize_weights(m, mean=0.0, std=0.01):
+    classname = m.__class__.__name__
+    if classname.find("Conv") != -1:
+        m.weight.data.normal_(mean, std)
+
+
+def compute_padding(kernel_size, dilation=1):
+    return int((kernel_size * dilation - dilation) / 2)
+
+
+def flatten_pad_shape(pad_shape):
+    layer = pad_shape[::-1]
+    pad_shape = [item for sublist in layer for item in sublist]
+    return pad_shape
+
+
+def insert_blanks(lst, item):
+    result = [item] * (len(lst) * 2 + 1)
+    result[1::2] = lst
+    return result
+
+
+def kl_divergence(m_p, logs_p, m_q, logs_q):
+    kl = (logs_q - logs_p) - 0.5
+    kl += (
+        0.5 * (torch.exp(2.0 * logs_p) + ((m_p - m_q) ** 2)) * torch.exp(-2.0 * logs_q)
+    )
+    return kl
+
+
+def rand_gumbel(shape):
+    uniform_samples = torch.rand(shape) * 0.99998 + 0.00001
+    return -torch.log(-torch.log(uniform_samples))
+
+
+def rand_gumbel_like(x):
+    g = rand_gumbel(x.size()).to(dtype=x.dtype, device=x.device)
+    return g
+
+
+def extract_segments(x, ids_str, segment_size=4):
+    ret = torch.zeros_like(x[:, :, :segment_size])
+    for i in range(x.size(0)):
+        idx_str = max(0, ids_str[i].item())
+        idx_end = idx_str + segment_size
+        available = x.size(2) - idx_str
+        if available >= segment_size:
+            ret[i] = x[i, :, idx_str:idx_end]
+        elif available > 0:
+            ret[i, :, :available] = x[i, :, idx_str:idx_str + available]
+    return ret
+
+
+def random_segments(x, x_lengths=None, segment_size=4):
+    b, d, t = x.size()
+    if x_lengths is None:
+        x_lengths = t
+    ids_str_max = torch.clamp(x_lengths - segment_size + 1, min=0)
+    ids_str = (torch.rand([b]).to(device=x.device) * ids_str_max).to(dtype=torch.long)
+    ret = extract_segments(x, ids_str, segment_size)
+    return ret, ids_str
+
+
+def get_timing_signal_1d(length, channels, min_timescale=1.0, max_timescale=1.0e4):
+    position = torch.arange(length, dtype=torch.float)
+    num_timescales = channels // 2
+    log_timescale_increment = math.log(float(max_timescale) / float(min_timescale)) / (
+        num_timescales - 1
+    )
+    inv_timescales = min_timescale * torch.exp(
+        torch.arange(num_timescales, dtype=torch.float) * -log_timescale_increment
+    )
+    scaled_time = position.unsqueeze(0) * inv_timescales.unsqueeze(1)
+    signal = torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], 0)
+    signal = F.pad(signal, [0, 0, 0, channels % 2])
+    signal = signal.view(1, channels, length)
+    return signal
+
+
+def add_timing_signal_1d(x, min_timescale=1.0, max_timescale=1.0e4):
+    b, channels, length = x.size()
+    signal = get_timing_signal_1d(length, channels, min_timescale, max_timescale)
+    return x + signal.to(dtype=x.dtype, device=x.device)
+
+
+def cat_timing_signal_1d(x, min_timescale=1.0, max_timescale=1.0e4, axis=1):
+    b, channels, length = x.size()
+    signal = get_timing_signal_1d(length, channels, min_timescale, max_timescale)
+    return torch.cat([x, signal.to(dtype=x.dtype, device=x.device)], axis)
+
+
+def subsequent_mask(length):
+    mask = torch.tril(torch.ones(length, length)).unsqueeze(0).unsqueeze(0)
+    return mask
+
+
+@torch.jit.script
+def gated_activation(input_a, input_b, n_channels):
+    n_channels_int = n_channels[0]
+    in_act = input_a + input_b
+    t_act = torch.tanh(in_act[:, :n_channels_int, :])
+    s_act = torch.sigmoid(in_act[:, n_channels_int:, :])
+    acts = t_act * s_act
+    return acts
+
+
+def shift_1d(x):
+    x = F.pad(x, flatten_pad_shape([[0, 0], [0, 0], [1, 0]]))[:, :, :-1]
+    return x
+
+
+def create_length_mask(length, max_length=None):
+    if max_length is None:
+        max_length = length.max()
+    x = torch.arange(max_length, dtype=length.dtype, device=length.device)
+    return x.unsqueeze(0) < length.unsqueeze(1)
+
+
+def compute_alignment_path(duration, mask):
+    b, _, t_y, t_x = mask.shape
+    cum_duration = torch.cumsum(duration, -1)
+
+    cum_duration_flat = cum_duration.view(b * t_x)
+    path = create_length_mask(cum_duration_flat, t_y).to(mask.dtype)
+    path = path.view(b, t_x, t_y)
+    path = path - F.pad(path, flatten_pad_shape([[0, 0], [1, 0], [0, 0]]))[:, :-1]
+    path = path.unsqueeze(1).transpose(2, 3) * mask
+    return path
+
+
+def clip_grad_value_(parameters, clip_value, norm_type=2):
+    if isinstance(parameters, torch.Tensor):
+        parameters = [parameters]
+    parameters = list(filter(lambda p: p.grad is not None, parameters))
+    norm_type = float(norm_type)
+    if clip_value is not None:
+        clip_value = float(clip_value)
+
+    total_norm = 0
+    for p in parameters:
+        param_norm = p.grad.data.norm(norm_type)
+        total_norm += param_norm.item() ** norm_type
+        if clip_value is not None:
+            p.grad.data.clamp_(min=-clip_value, max=clip_value)
+    total_norm = total_norm ** (1.0 / norm_type)
+    return total_norm

third_party/tiny_tts_frontend/tiny_tts/nn/modules.py

@@ -0,0 +1,578 @@
+import math
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from torch.nn import Conv1d
+from torch.nn.utils import weight_norm, remove_weight_norm
+
+from . import commons
+from .commons import initialize_weights, compute_padding
+from .transforms import spline_transform
+from .attentions import TransformerBlock
+
+LRELU_SLOPE = 0.1
+
+
+class ChannelNorm(nn.Module):
+    def __init__(self, channels, eps=1e-5):
+        super().__init__()
+        self.channels = channels
+        self.eps = eps
+
+        self.gamma = nn.Parameter(torch.ones(channels))
+        self.beta = nn.Parameter(torch.zeros(channels))
+
+    def forward(self, x):
+        x = x.transpose(1, -1)
+        x = F.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps)
+        return x.transpose(1, -1)
+
+
+class ConvReluNorm(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        hidden_channels,
+        out_channels,
+        kernel_size,
+        n_layers,
+        p_dropout,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.hidden_channels = hidden_channels
+        self.out_channels = out_channels
+        self.kernel_size = kernel_size
+        self.n_layers = n_layers
+        self.p_dropout = p_dropout
+        assert n_layers > 1, "Number of layers should be larger than 0."
+
+        self.conv_layers = nn.ModuleList()
+        self.norm_layers = nn.ModuleList()
+        self.conv_layers.append(
+            nn.Conv1d(
+                in_channels, hidden_channels, kernel_size, padding=kernel_size // 2
+            )
+        )
+        self.norm_layers.append(ChannelNorm(hidden_channels))
+        self.relu_drop = nn.Sequential(nn.ReLU(), nn.Dropout(p_dropout))
+        for _ in range(n_layers - 1):
+            self.conv_layers.append(
+                nn.Conv1d(
+                    hidden_channels,
+                    hidden_channels,
+                    kernel_size,
+                    padding=kernel_size // 2,
+                )
+            )
+            self.norm_layers.append(ChannelNorm(hidden_channels))
+        self.proj = nn.Conv1d(hidden_channels, out_channels, 1)
+        self.proj.weight.data.zero_()
+        self.proj.bias.data.zero_()
+
+    def forward(self, x, x_mask):
+        x_org = x
+        for i in range(self.n_layers):
+            x = self.conv_layers[i](x * x_mask)
+            x = self.norm_layers[i](x)
+            x = self.relu_drop(x)
+        x = x_org + self.proj(x)
+        return x * x_mask
+
+
+class DepthwiseSepConv(nn.Module):
+    """Dilated and Depth-Separable Convolution"""
+
+    def __init__(self, channels, kernel_size, n_layers, p_dropout=0.0):
+        super().__init__()
+        self.channels = channels
+        self.kernel_size = kernel_size
+        self.n_layers = n_layers
+        self.p_dropout = p_dropout
+
+        self.drop = nn.Dropout(p_dropout)
+        self.convs_sep = nn.ModuleList()
+        self.convs_1x1 = nn.ModuleList()
+        self.norms_1 = nn.ModuleList()
+        self.norms_2 = nn.ModuleList()
+        for i in range(n_layers):
+            dilation = kernel_size**i
+            padding = (kernel_size * dilation - dilation) // 2
+            self.convs_sep.append(
+                nn.Conv1d(
+                    channels,
+                    channels,
+                    kernel_size,
+                    groups=channels,
+                    dilation=dilation,
+                    padding=padding,
+                )
+            )
+            self.convs_1x1.append(nn.Conv1d(channels, channels, 1))
+            self.norms_1.append(ChannelNorm(channels))
+            self.norms_2.append(ChannelNorm(channels))
+
+    def forward(self, x, x_mask, g=None):
+        if g is not None:
+            x = x + g
+        for i in range(self.n_layers):
+            y = self.convs_sep[i](x * x_mask)
+            y = self.norms_1[i](y)
+            y = F.gelu(y)
+            y = self.convs_1x1[i](y)
+            y = self.norms_2[i](y)
+            y = F.gelu(y)
+            y = self.drop(y)
+            x = x + y
+        return x * x_mask
+
+
+class WaveNet(torch.nn.Module):
+    def __init__(
+        self,
+        hidden_channels,
+        kernel_size,
+        dilation_rate,
+        n_layers,
+        gin_channels=0,
+        p_dropout=0,
+    ):
+        super(WaveNet, self).__init__()
+        assert kernel_size % 2 == 1
+        self.hidden_channels = hidden_channels
+        self.kernel_size = (kernel_size,)
+        self.dilation_rate = dilation_rate
+        self.n_layers = n_layers
+        self.gin_channels = gin_channels
+        self.p_dropout = p_dropout
+
+        self.in_layers = torch.nn.ModuleList()
+        self.res_skip_layers = torch.nn.ModuleList()
+        self.drop = nn.Dropout(p_dropout)
+
+        if gin_channels != 0:
+            cond_layer = torch.nn.Conv1d(
+                gin_channels, 2 * hidden_channels * n_layers, 1
+            )
+            self.cond_layer = torch.nn.utils.weight_norm(cond_layer, name="weight")
+
+        for i in range(n_layers):
+            dilation = dilation_rate**i
+            padding = int((kernel_size * dilation - dilation) / 2)
+            in_layer = torch.nn.Conv1d(
+                hidden_channels,
+                2 * hidden_channels,
+                kernel_size,
+                dilation=dilation,
+                padding=padding,
+            )
+            in_layer = torch.nn.utils.weight_norm(in_layer, name="weight")
+            self.in_layers.append(in_layer)
+
+            if i < n_layers - 1:
+                res_skip_channels = 2 * hidden_channels
+            else:
+                res_skip_channels = hidden_channels
+
+            res_skip_layer = torch.nn.Conv1d(hidden_channels, res_skip_channels, 1)
+            res_skip_layer = torch.nn.utils.weight_norm(res_skip_layer, name="weight")
+            self.res_skip_layers.append(res_skip_layer)
+
+    def forward(self, x, x_mask, g=None, **kwargs):
+        output = torch.zeros_like(x)
+        n_channels_tensor = torch.IntTensor([self.hidden_channels])
+
+        if g is not None:
+            g = self.cond_layer(g)
+
+        for i in range(self.n_layers):
+            x_in = self.in_layers[i](x)
+            if g is not None:
+                cond_offset = i * 2 * self.hidden_channels
+                g_l = g[:, cond_offset : cond_offset + 2 * self.hidden_channels, :]
+            else:
+                g_l = torch.zeros_like(x_in)
+
+            acts = commons.gated_activation(x_in, g_l, n_channels_tensor)
+            acts = self.drop(acts)
+
+            res_skip_acts = self.res_skip_layers[i](acts)
+            if i < self.n_layers - 1:
+                res_acts = res_skip_acts[:, : self.hidden_channels, :]
+                x = (x + res_acts) * x_mask
+                output = output + res_skip_acts[:, self.hidden_channels :, :]
+            else:
+                output = output + res_skip_acts
+        return output * x_mask
+
+    def remove_weight_norm(self):
+        if self.gin_channels != 0:
+            torch.nn.utils.remove_weight_norm(self.cond_layer)
+        for l in self.in_layers:
+            torch.nn.utils.remove_weight_norm(l)
+        for l in self.res_skip_layers:
+            torch.nn.utils.remove_weight_norm(l)
+
+
+class ConvResBlock(torch.nn.Module):
+    def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5)):
+        super(ConvResBlock, self).__init__()
+        self.convs1 = nn.ModuleList(
+            [
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[0],
+                        padding=compute_padding(kernel_size, dilation[0]),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[1],
+                        padding=compute_padding(kernel_size, dilation[1]),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[2],
+                        padding=compute_padding(kernel_size, dilation[2]),
+                    )
+                ),
+            ]
+        )
+        self.convs1.apply(initialize_weights)
+
+        self.convs2 = nn.ModuleList(
+            [
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=1,
+                        padding=compute_padding(kernel_size, 1),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=1,
+                        padding=compute_padding(kernel_size, 1),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=1,
+                        padding=compute_padding(kernel_size, 1),
+                    )
+                ),
+            ]
+        )
+        self.convs2.apply(initialize_weights)
+
+    def forward(self, x, x_mask=None):
+        for c1, c2 in zip(self.convs1, self.convs2):
+            xt = F.leaky_relu(x, LRELU_SLOPE)
+            if x_mask is not None:
+                xt = xt * x_mask
+            xt = c1(xt)
+            xt = F.leaky_relu(xt, LRELU_SLOPE)
+            if x_mask is not None:
+                xt = xt * x_mask
+            xt = c2(xt)
+            x = xt + x
+        if x_mask is not None:
+            x = x * x_mask
+        return x
+
+    def remove_weight_norm(self):
+        for l in self.convs1:
+            remove_weight_norm(l)
+        for l in self.convs2:
+            remove_weight_norm(l)
+
+
+class ConvResBlockLight(torch.nn.Module):
+    def __init__(self, channels, kernel_size=3, dilation=(1, 3)):
+        super(ConvResBlockLight, self).__init__()
+        self.convs = nn.ModuleList(
+            [
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[0],
+                        padding=compute_padding(kernel_size, dilation[0]),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[1],
+                        padding=compute_padding(kernel_size, dilation[1]),
+                    )
+                ),
+            ]
+        )
+        self.convs.apply(initialize_weights)
+
+    def forward(self, x, x_mask=None):
+        for c in self.convs:
+            xt = F.leaky_relu(x, LRELU_SLOPE)
+            if x_mask is not None:
+                xt = xt * x_mask
+            xt = c(xt)
+            x = xt + x
+        if x_mask is not None:
+            x = x * x_mask
+        return x
+
+    def remove_weight_norm(self):
+        for l in self.convs:
+            remove_weight_norm(l)
+
+
+class LogTransform(nn.Module):
+    def forward(self, x, x_mask, reverse=False, **kwargs):
+        if not reverse:
+            y = torch.log(torch.clamp_min(x, 1e-5)) * x_mask
+            logdet = torch.sum(-y, [1, 2])
+            return y, logdet
+        else:
+            x = torch.exp(x) * x_mask
+            return x
+
+
+class FlipTransform(nn.Module):
+    def forward(self, x, *args, reverse=False, **kwargs):
+        x = torch.flip(x, [1])
+        if not reverse:
+            logdet = torch.zeros(x.size(0)).to(dtype=x.dtype, device=x.device)
+            return x, logdet
+        else:
+            return x
+
+
+class AffineCoupling(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.channels = channels
+        self.m = nn.Parameter(torch.zeros(channels, 1))
+        self.logs = nn.Parameter(torch.zeros(channels, 1))
+
+    def forward(self, x, x_mask, reverse=False, **kwargs):
+        if not reverse:
+            y = self.m + torch.exp(self.logs) * x
+            y = y * x_mask
+            logdet = torch.sum(self.logs * x_mask, [1, 2])
+            return y, logdet
+        else:
+            x = (x - self.m) * torch.exp(-self.logs) * x_mask
+            return x
+
+
+class FlowCouplingLayer(nn.Module):
+    def __init__(
+        self,
+        channels,
+        hidden_channels,
+        kernel_size,
+        dilation_rate,
+        n_layers,
+        p_dropout=0,
+        gin_channels=0,
+        mean_only=False,
+    ):
+        assert channels % 2 == 0, "channels should be divisible by 2"
+        super().__init__()
+        self.channels = channels
+        self.hidden_channels = hidden_channels
+        self.kernel_size = kernel_size
+        self.dilation_rate = dilation_rate
+        self.n_layers = n_layers
+        self.half_channels = channels // 2
+        self.mean_only = mean_only
+
+        self.pre = nn.Conv1d(self.half_channels, hidden_channels, 1)
+        self.enc = WaveNet(
+            hidden_channels,
+            kernel_size,
+            dilation_rate,
+            n_layers,
+            p_dropout=p_dropout,
+            gin_channels=gin_channels,
+        )
+        self.post = nn.Conv1d(hidden_channels, self.half_channels * (2 - mean_only), 1)
+        self.post.weight.data.zero_()
+        self.post.bias.data.zero_()
+
+    def forward(self, x, x_mask, g=None, reverse=False):
+        x0, x1 = torch.split(x, [self.half_channels] * 2, 1)
+        h = self.pre(x0) * x_mask
+        h = self.enc(h, x_mask, g=g)
+        stats = self.post(h) * x_mask
+        if not self.mean_only:
+            m, logs = torch.split(stats, [self.half_channels] * 2, 1)
+        else:
+            m = stats
+            logs = torch.zeros_like(m)
+
+        if not reverse:
+            x1 = m + x1 * torch.exp(logs) * x_mask
+            x = torch.cat([x0, x1], 1)
+            logdet = torch.sum(logs, [1, 2])
+            return x, logdet
+        else:
+            x1 = (x1 - m) * torch.exp(-logs) * x_mask
+            x = torch.cat([x0, x1], 1)
+            return x
+
+
+class ConvolutionalFlow(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        filter_channels,
+        kernel_size,
+        n_layers,
+        num_bins=10,
+        tail_bound=5.0,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.filter_channels = filter_channels
+        self.kernel_size = kernel_size
+        self.n_layers = n_layers
+        self.num_bins = num_bins
+        self.tail_bound = tail_bound
+        self.half_channels = in_channels // 2
+
+        self.pre = nn.Conv1d(self.half_channels, filter_channels, 1)
+        self.convs = DepthwiseSepConv(filter_channels, kernel_size, n_layers, p_dropout=0.0)
+        self.proj = nn.Conv1d(
+            filter_channels, self.half_channels * (num_bins * 3 - 1), 1
+        )
+        self.proj.weight.data.zero_()
+        self.proj.bias.data.zero_()
+
+    def forward(self, x, x_mask, g=None, reverse=False):
+        x0, x1 = torch.split(x, [self.half_channels] * 2, 1)
+        h = self.pre(x0)
+        h = self.convs(h, x_mask, g=g)
+        h = self.proj(h) * x_mask
+
+        b, c, t = x0.shape
+        h = h.reshape(b, c, -1, t).permute(0, 1, 3, 2)
+
+        unnormalized_widths = h[..., : self.num_bins] / math.sqrt(self.filter_channels)
+        unnormalized_heights = h[..., self.num_bins : 2 * self.num_bins] / math.sqrt(
+            self.filter_channels
+        )
+        unnormalized_derivatives = h[..., 2 * self.num_bins :]
+
+        x1, logabsdet = spline_transform(
+            x1,
+            unnormalized_widths,
+            unnormalized_heights,
+            unnormalized_derivatives,
+            inverse=reverse,
+            tails="linear",
+            tail_bound=self.tail_bound,
+        )
+
+        x = torch.cat([x0, x1], 1) * x_mask
+        logdet = torch.sum(logabsdet * x_mask, [1, 2])
+        if not reverse:
+            return x, logdet
+        else:
+            return x
+
+
+class TransformerCouplingLayer(nn.Module):
+    def __init__(
+        self,
+        channels,
+        hidden_channels,
+        kernel_size,
+        n_layers,
+        n_heads,
+        p_dropout=0,
+        filter_channels=0,
+        mean_only=False,
+        wn_sharing_parameter=None,
+        gin_channels=0,
+    ):
+        assert n_layers == 3, n_layers
+        assert channels % 2 == 0, "channels should be divisible by 2"
+        super().__init__()
+        self.channels = channels
+        self.hidden_channels = hidden_channels
+        self.kernel_size = kernel_size
+        self.n_layers = n_layers
+        self.half_channels = channels // 2
+        self.mean_only = mean_only
+
+        self.pre = nn.Conv1d(self.half_channels, hidden_channels, 1)
+        self.enc = (
+            TransformerBlock(
+                hidden_channels,
+                filter_channels,
+                n_heads,
+                n_layers,
+                kernel_size,
+                p_dropout,
+                isflow=True,
+                gin_channels=gin_channels,
+            )
+            if wn_sharing_parameter is None
+            else wn_sharing_parameter
+        )
+        self.post = nn.Conv1d(hidden_channels, self.half_channels * (2 - mean_only), 1)
+        self.post.weight.data.zero_()
+        self.post.bias.data.zero_()
+
+    def forward(self, x, x_mask, g=None, reverse=False):
+        x0, x1 = torch.split(x, [self.half_channels] * 2, 1)
+        h = self.pre(x0) * x_mask
+        h = self.enc(h, x_mask, g=g)
+        stats = self.post(h) * x_mask
+        if not self.mean_only:
+            m, logs = torch.split(stats, [self.half_channels] * 2, 1)
+        else:
+            m = stats
+            logs = torch.zeros_like(m)
+
+        if not reverse:
+            x1 = m + x1 * torch.exp(logs) * x_mask
+            x = torch.cat([x0, x1], 1)
+            logdet = torch.sum(logs, [1, 2])
+            return x, logdet
+        else:
+            x1 = (x1 - m) * torch.exp(-logs) * x_mask
+            x = torch.cat([x0, x1], 1)
+            return x

third_party/tiny_tts_frontend/tiny_tts/nn/transforms.py

@@ -0,0 +1,209 @@
+import torch
+from torch.nn import functional as F
+
+import numpy as np
+
+
+DEFAULT_MIN_BIN_WIDTH = 1e-3
+DEFAULT_MIN_BIN_HEIGHT = 1e-3
+DEFAULT_MIN_DERIVATIVE = 1e-3
+
+
+def spline_transform(
+    inputs,
+    unnormalized_widths,
+    unnormalized_heights,
+    unnormalized_derivatives,
+    inverse=False,
+    tails=None,
+    tail_bound=1.0,
+    min_bin_width=DEFAULT_MIN_BIN_WIDTH,
+    min_bin_height=DEFAULT_MIN_BIN_HEIGHT,
+    min_derivative=DEFAULT_MIN_DERIVATIVE,
+):
+    if tails is None:
+        spline_fn = quadratic_spline
+        spline_kwargs = {}
+    else:
+        spline_fn = unbounded_spline
+        spline_kwargs = {"tails": tails, "tail_bound": tail_bound}
+
+    outputs, logabsdet = spline_fn(
+        inputs=inputs,
+        unnormalized_widths=unnormalized_widths,
+        unnormalized_heights=unnormalized_heights,
+        unnormalized_derivatives=unnormalized_derivatives,
+        inverse=inverse,
+        min_bin_width=min_bin_width,
+        min_bin_height=min_bin_height,
+        min_derivative=min_derivative,
+        **spline_kwargs
+    )
+    return outputs, logabsdet
+
+
+def searchsorted(bin_locations, inputs, eps=1e-6):
+    bin_locations[..., -1] += eps
+    return torch.sum(inputs[..., None] >= bin_locations, dim=-1) - 1
+
+
+def unbounded_spline(
+    inputs,
+    unnormalized_widths,
+    unnormalized_heights,
+    unnormalized_derivatives,
+    inverse=False,
+    tails="linear",
+    tail_bound=1.0,
+    min_bin_width=DEFAULT_MIN_BIN_WIDTH,
+    min_bin_height=DEFAULT_MIN_BIN_HEIGHT,
+    min_derivative=DEFAULT_MIN_DERIVATIVE,
+):
+    inside_interval_mask = (inputs >= -tail_bound) & (inputs <= tail_bound)
+    outside_interval_mask = ~inside_interval_mask
+
+    outputs = torch.zeros_like(inputs)
+    logabsdet = torch.zeros_like(inputs)
+
+    if tails == "linear":
+        unnormalized_derivatives = F.pad(unnormalized_derivatives, pad=(1, 1))
+        constant = np.log(np.exp(1 - min_derivative) - 1)
+        unnormalized_derivatives[..., 0] = constant
+        unnormalized_derivatives[..., -1] = constant
+
+        outputs[outside_interval_mask] = inputs[outside_interval_mask]
+        logabsdet[outside_interval_mask] = 0
+    else:
+        raise RuntimeError("{} tails are not implemented.".format(tails))
+
+    (
+        outputs[inside_interval_mask],
+        logabsdet[inside_interval_mask],
+    ) = quadratic_spline(
+        inputs=inputs[inside_interval_mask],
+        unnormalized_widths=unnormalized_widths[inside_interval_mask, :],
+        unnormalized_heights=unnormalized_heights[inside_interval_mask, :],
+        unnormalized_derivatives=unnormalized_derivatives[inside_interval_mask, :],
+        inverse=inverse,
+        left=-tail_bound,
+        right=tail_bound,
+        bottom=-tail_bound,
+        top=tail_bound,
+        min_bin_width=min_bin_width,
+        min_bin_height=min_bin_height,
+        min_derivative=min_derivative,
+    )
+
+    return outputs, logabsdet
+
+
+def quadratic_spline(
+    inputs,
+    unnormalized_widths,
+    unnormalized_heights,
+    unnormalized_derivatives,
+    inverse=False,
+    left=0.0,
+    right=1.0,
+    bottom=0.0,
+    top=1.0,
+    min_bin_width=DEFAULT_MIN_BIN_WIDTH,
+    min_bin_height=DEFAULT_MIN_BIN_HEIGHT,
+    min_derivative=DEFAULT_MIN_DERIVATIVE,
+):
+    if torch.min(inputs) < left or torch.max(inputs) > right:
+        raise ValueError("Input to a transform is not within its domain")
+
+    num_bins = unnormalized_widths.shape[-1]
+
+    if min_bin_width * num_bins > 1.0:
+        raise ValueError("Minimal bin width too large for the number of bins")
+    if min_bin_height * num_bins > 1.0:
+        raise ValueError("Minimal bin height too large for the number of bins")
+
+    widths = F.softmax(unnormalized_widths, dim=-1)
+    widths = min_bin_width + (1 - min_bin_width * num_bins) * widths
+    cumwidths = torch.cumsum(widths, dim=-1)
+    cumwidths = F.pad(cumwidths, pad=(1, 0), mode="constant", value=0.0)
+    cumwidths = (right - left) * cumwidths + left
+    cumwidths[..., 0] = left
+    cumwidths[..., -1] = right
+    widths = cumwidths[..., 1:] - cumwidths[..., :-1]
+
+    derivatives = min_derivative + F.softplus(unnormalized_derivatives)
+
+    heights = F.softmax(unnormalized_heights, dim=-1)
+    heights = min_bin_height + (1 - min_bin_height * num_bins) * heights
+    cumheights = torch.cumsum(heights, dim=-1)
+    cumheights = F.pad(cumheights, pad=(1, 0), mode="constant", value=0.0)
+    cumheights = (top - bottom) * cumheights + bottom
+    cumheights[..., 0] = bottom
+    cumheights[..., -1] = top
+    heights = cumheights[..., 1:] - cumheights[..., :-1]
+
+    if inverse:
+        bin_idx = searchsorted(cumheights, inputs)[..., None]
+    else:
+        bin_idx = searchsorted(cumwidths, inputs)[..., None]
+
+    input_cumwidths = cumwidths.gather(-1, bin_idx)[..., 0]
+    input_bin_widths = widths.gather(-1, bin_idx)[..., 0]
+
+    input_cumheights = cumheights.gather(-1, bin_idx)[..., 0]
+    delta = heights / widths
+    input_delta = delta.gather(-1, bin_idx)[..., 0]
+
+    input_derivatives = derivatives.gather(-1, bin_idx)[..., 0]
+    input_derivatives_plus_one = derivatives[..., 1:].gather(-1, bin_idx)[..., 0]
+
+    input_heights = heights.gather(-1, bin_idx)[..., 0]
+
+    if inverse:
+        a = (inputs - input_cumheights) * (
+            input_derivatives + input_derivatives_plus_one - 2 * input_delta
+        ) + input_heights * (input_delta - input_derivatives)
+        b = input_heights * input_derivatives - (inputs - input_cumheights) * (
+            input_derivatives + input_derivatives_plus_one - 2 * input_delta
+        )
+        c = -input_delta * (inputs - input_cumheights)
+
+        discriminant = b.pow(2) - 4 * a * c
+        assert (discriminant >= 0).all()
+
+        root = (2 * c) / (-b - torch.sqrt(discriminant))
+        outputs = root * input_bin_widths + input_cumwidths
+
+        theta_one_minus_theta = root * (1 - root)
+        denominator = input_delta + (
+            (input_derivatives + input_derivatives_plus_one - 2 * input_delta)
+            * theta_one_minus_theta
+        )
+        derivative_numerator = input_delta.pow(2) * (
+            input_derivatives_plus_one * root.pow(2)
+            + 2 * input_delta * theta_one_minus_theta
+            + input_derivatives * (1 - root).pow(2)
+        )
+        logabsdet = torch.log(derivative_numerator) - 2 * torch.log(denominator)
+
+        return outputs, -logabsdet
+    else:
+        theta = (inputs - input_cumwidths) / input_bin_widths
+        theta_one_minus_theta = theta * (1 - theta)
+
+        numerator = input_heights * (
+            input_delta * theta.pow(2) + input_derivatives * theta_one_minus_theta
+        )
+        denominator = input_delta + (
+            (input_derivatives + input_derivatives_plus_one - 2 * input_delta)
+            * theta_one_minus_theta
+        )
+        outputs = input_cumheights + numerator / denominator
+
+        derivative_numerator = input_delta.pow(2) * (
+            input_derivatives_plus_one * theta.pow(2)
+            + 2 * input_delta * theta_one_minus_theta
+            + input_derivatives * (1 - theta).pow(2)
+        )
+        logabsdet = torch.log(derivative_numerator) - 2 * torch.log(denominator)
+
+        return outputs, logabsdet

third_party/tiny_tts_frontend/tiny_tts/text/__init__.py

@@ -0,0 +1,19 @@
+from .symbols import *
+
+
+_symbol_to_id = {s: i for i, s in enumerate(symbols)}
+
+
+def phonemes_to_ids(cleaned_text, tones, language, symbol_to_id=None):
+    """Converts a list of phoneme symbols to a sequence of integer IDs."""
+    symbol_to_id_map = symbol_to_id if symbol_to_id else _symbol_to_id
+    unk_id = symbol_to_id_map.get("UNK")
+    if unk_id is None:
+        phones = [symbol_to_id_map[symbol] for symbol in cleaned_text]
+    else:
+        phones = [symbol_to_id_map.get(symbol, unk_id) for symbol in cleaned_text]
+    tone_start = language_tone_start_map[language]
+    tones = [i + tone_start for i in tones]
+    lang_id = language_id_map[language]
+    lang_ids = [lang_id for _ in phones]
+    return phones, tones, lang_ids

third_party/tiny_tts_frontend/tiny_tts/text/english.py

@@ -0,0 +1,173 @@
+import pickle
+import os
+import re
+from g2p_en import G2p
+
+from . import symbols
+
+from .english_utils.abbreviations import expand_abbreviations
+from .english_utils.time_norm import expand_time_english
+from .english_utils.number_norm import normalize_numbers
+
+
+def distribute_phone(n_phone, n_word):
+    phones_per_word = [0] * n_word
+    for task in range(n_phone):
+        min_tasks = min(phones_per_word)
+        min_indices = [
+            i for i, x in enumerate(phones_per_word) if x == min_tasks
+        ]
+        chosen_index = min_indices[len(min_indices) // 2]
+        phones_per_word[chosen_index] += 1
+    return phones_per_word
+
+
+from transformers import AutoTokenizer
+
+current_file_path = os.path.dirname(__file__)
+CMU_DICT_PATH = os.path.join(current_file_path, "cmudict.rep")
+CACHE_PATH = os.path.join(current_file_path, "cmudict_cache.pickle")
+_g2p = G2p()
+
+arpa = {
+    "AH0", "S", "AH1", "EY2", "AE2", "EH0", "OW2", "UH0", "NG", "B",
+    "G", "AY0", "M", "AA0", "F", "AO0", "ER2", "UH1", "IY1", "AH2",
+    "DH", "IY0", "EY1", "IH0", "K", "N", "W", "IY2", "T", "AA1",
+    "ER1", "EH2", "OY0", "UH2", "UW1", "Z", "AW2", "AW1", "V", "UW2",
+    "AA2", "ER", "AW0", "UW0", "R", "OW1", "EH1", "ZH", "AE0", "IH2",
+    "IH", "Y", "JH", "P", "AY1", "EY0", "OY2", "TH", "HH", "D",
+    "ER0", "CH", "AO1", "AE1", "AO2", "OY1", "AY2", "IH1", "OW0", "L", "SH",
+}
+
+
+def map_phoneme(ph):
+    rep_map = {
+        "：": ",", "；": ",", "，": ",", "。": ".", "！": "!",
+        "？": "?", "\n": ".", "·": ",", "、": ",", "...": "…", "v": "V",
+    }
+    if ph in rep_map.keys():
+        ph = rep_map[ph]
+    if ph in symbols:
+        return ph
+    if ph not in symbols:
+        ph = "UNK"
+    return ph
+
+
+def read_dict():
+    g2p_dict = {}
+    start_line = 49
+    with open(CMU_DICT_PATH) as f:
+        line = f.readline()
+        line_index = 1
+        while line:
+            if line_index >= start_line:
+                line = line.strip()
+                word_split = line.split("  ")
+                word = word_split[0]
+
+                syllable_split = word_split[1].split(" - ")
+                g2p_dict[word] = []
+                for syllable in syllable_split:
+                    phone_split = syllable.split(" ")
+                    g2p_dict[word].append(phone_split)
+
+            line_index = line_index + 1
+            line = f.readline()
+
+    return g2p_dict
+
+
+def cache_dict(g2p_dict, file_path):
+    with open(file_path, "wb") as pickle_file:
+        pickle.dump(g2p_dict, pickle_file)
+
+
+def get_dict():
+    if os.path.exists(CACHE_PATH):
+        with open(CACHE_PATH, "rb") as pickle_file:
+            g2p_dict = pickle.load(pickle_file)
+    else:
+        g2p_dict = read_dict()
+        cache_dict(g2p_dict, CACHE_PATH)
+
+    return g2p_dict
+
+
+eng_dict = get_dict()
+
+
+def parse_phoneme(phn):
+    tone = 0
+    if re.search(r"\d$", phn):
+        tone = int(phn[-1]) + 1
+        phn = phn[:-1]
+    return phn.lower(), tone
+
+
+def parse_syllables(syllables):
+    tones = []
+    phonemes = []
+    for phn_list in syllables:
+        for i in range(len(phn_list)):
+            phn = phn_list[i]
+            phn, tone = parse_phoneme(phn)
+            phonemes.append(phn)
+            tones.append(tone)
+    return phonemes, tones
+
+
+def normalize_text(text):
+    text = text.lower()
+    text = expand_time_english(text)
+    text = normalize_numbers(text)
+    text = expand_abbreviations(text)
+    return text
+
+
+model_id = 'bert-base-uncased'
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+
+
+def grapheme_to_phoneme(text, pad_start_end=True, tokenized=None):
+    if tokenized is None:
+        tokenized = tokenizer.tokenize(text)
+    ph_groups = []
+    for t in tokenized:
+        if not t.startswith("#"):
+            ph_groups.append([t])
+        else:
+            ph_groups[-1].append(t.replace("#", ""))
+
+    phones = []
+    tones = []
+    word2ph = []
+    for group in ph_groups:
+        w = "".join(group)
+        phone_len = 0
+        word_len = len(group)
+        if w.upper() in eng_dict:
+            phns, tns = parse_syllables(eng_dict[w.upper()])
+            phones += phns
+            tones += tns
+            phone_len += len(phns)
+        else:
+            phone_list = list(filter(lambda p: p != " ", _g2p(w)))
+            for ph in phone_list:
+                if ph in arpa:
+                    ph, tn = parse_phoneme(ph)
+                    phones.append(ph)
+                    tones.append(tn)
+                else:
+                    phones.append(ph)
+                    tones.append(0)
+                phone_len += 1
+        aaa = distribute_phone(phone_len, word_len)
+        word2ph += aaa
+    phones = [map_phoneme(i) for i in phones]
+
+    if pad_start_end:
+        phones = ["_"] + phones + ["_"]
+        tones = [0] + tones + [0]
+        word2ph = [1] + word2ph + [1]
+    return phones, tones, word2ph

third_party/tiny_tts_frontend/tiny_tts/text/english_utils/abbreviations.py

@@ -0,0 +1,35 @@
+import re
+
+# List of (regular expression, replacement) pairs for abbreviations in english:
+abbreviations_en = [
+    (re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1])
+    for x in [
+        ("mrs", "misess"),
+        ("mr", "mister"),
+        ("dr", "doctor"),
+        ("st", "saint"),
+        ("co", "company"),
+        ("jr", "junior"),
+        ("maj", "major"),
+        ("gen", "general"),
+        ("drs", "doctors"),
+        ("rev", "reverend"),
+        ("lt", "lieutenant"),
+        ("hon", "honorable"),
+        ("sgt", "sergeant"),
+        ("capt", "captain"),
+        ("esq", "esquire"),
+        ("ltd", "limited"),
+        ("col", "colonel"),
+        ("ft", "fort"),
+    ]
+]
+
+def expand_abbreviations(text, lang="en"):
+    if lang == "en":
+        _abbreviations = abbreviations_en
+    else:
+        raise NotImplementedError()
+    for regex, replacement in _abbreviations:
+        text = re.sub(regex, replacement, text)
+    return text

third_party/tiny_tts_frontend/tiny_tts/text/english_utils/number_norm.py

@@ -0,0 +1,97 @@
+""" from https://github.com/keithito/tacotron """
+
+import re
+from typing import Dict
+
+import inflect
+
+_inflect = inflect.engine()
+_comma_number_re = re.compile(r"([0-9][0-9\,]+[0-9])")
+_decimal_number_re = re.compile(r"([0-9]+\.[0-9]+)")
+_currency_re = re.compile(r"(£|\$|¥)([0-9\,\.]*[0-9]+)")
+_ordinal_re = re.compile(r"[0-9]+(st|nd|rd|th)")
+_number_re = re.compile(r"-?[0-9]+")
+
+
+def _remove_commas(m):
+    return m.group(1).replace(",", "")
+
+
+def _expand_decimal_point(m):
+    return m.group(1).replace(".", " point ")
+
+
+def __expand_currency(value: str, inflection: Dict[float, str]) -> str:
+    parts = value.replace(",", "").split(".")
+    if len(parts) > 2:
+        return f"{value} {inflection[2]}"  # Unexpected format
+    text = []
+    integer = int(parts[0]) if parts[0] else 0
+    if integer > 0:
+        integer_unit = inflection.get(integer, inflection[2])
+        text.append(f"{integer} {integer_unit}")
+    fraction = int(parts[1]) if len(parts) > 1 and parts[1] else 0
+    if fraction > 0:
+        fraction_unit = inflection.get(fraction / 100, inflection[0.02])
+        text.append(f"{fraction} {fraction_unit}")
+    if len(text) == 0:
+        return f"zero {inflection[2]}"
+    return " ".join(text)
+
+
+def _expand_currency(m: "re.Match") -> str:
+    currencies = {
+        "$": {
+            0.01: "cent",
+            0.02: "cents",
+            1: "dollar",
+            2: "dollars",
+        },
+        "€": {
+            0.01: "cent",
+            0.02: "cents",
+            1: "euro",
+            2: "euros",
+        },
+        "£": {
+            0.01: "penny",
+            0.02: "pence",
+            1: "pound sterling",
+            2: "pounds sterling",
+        },
+        "¥": {
+            # TODO rin
+            0.02: "sen",
+            2: "yen",
+        },
+    }
+    unit = m.group(1)
+    currency = currencies[unit]
+    value = m.group(2)
+    return __expand_currency(value, currency)
+
+
+def _expand_ordinal(m):
+    return _inflect.number_to_words(m.group(0))
+
+
+def _expand_number(m):
+    num = int(m.group(0))
+    if 1000 < num < 3000:
+        if num == 2000:
+            return "two thousand"
+        if 2000 < num < 2010:
+            return "two thousand " + _inflect.number_to_words(num % 100)
+        if num % 100 == 0:
+            return _inflect.number_to_words(num // 100) + " hundred"
+        return _inflect.number_to_words(num, andword="", zero="oh", group=2).replace(", ", " ")
+    return _inflect.number_to_words(num, andword="")
+
+
+def normalize_numbers(text):
+    text = re.sub(_comma_number_re, _remove_commas, text)
+    text = re.sub(_currency_re, _expand_currency, text)
+    text = re.sub(_decimal_number_re, _expand_decimal_point, text)
+    text = re.sub(_ordinal_re, _expand_ordinal, text)
+    text = re.sub(_number_re, _expand_number, text)
+    return text