scripts/run_mac_tts.py

#!/usr/bin/env python3
"""
Run Supertone/supertonic TTS on macOS using ONNX Runtime with Core ML EP when available.

This matches the app's contract (fixed ``text_ids`` / ``text_mask`` length from ONNX,
padded latent grid like ``SupertonicTTSService``), not the variable-length batch path in
upstream ``helper.py`` alone.

Usage::

    uv sync && uv run python convert_supertonic_coreml.py --output ./export
    uv run python run_mac_tts.py --models-dir ./export --text "Hello world." --out ./out.wav
"""

from __future__ import annotations

import argparse
import json
import math
import re
import struct
import sys
import wave
from pathlib import Path
from typing import Any
from unicodedata import normalize

import numpy as np

try:
    import onnxruntime as ort
except ImportError as e:  # pragma: no cover
    raise SystemExit("Install onnxruntime: pip install -r requirements-export.txt") from e


AVAILABLE_LANGS = ("en", "ko", "es", "pt", "fr")


def _providers(prefer_coreml: bool) -> list[str]:
    if not prefer_coreml:
        return ["CPUExecutionProvider"]
    avail = set(ort.get_available_providers())
    if "CoreMLExecutionProvider" in avail:
        return ["CoreMLExecutionProvider", "CPUExecutionProvider"]
    print("CoreMLExecutionProvider not available; using CPU.", file=sys.stderr)
    return ["CPUExecutionProvider"]


def _session(path: Path, providers: list[str]) -> ort.InferenceSession:
    so = ort.SessionOptions()
    so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
    return ort.InferenceSession(str(path), sess_options=so, providers=providers)


def _input_dim(sess: ort.InferenceSession, name: str, axis: int) -> int | None:
    for i in sess.get_inputs():
        if i.name == name:
            shape = i.shape
            if axis < len(shape):
                d = shape[axis]
                if isinstance(d, int) and d > 0:
                    return d
            break
    return None


def preprocess_text(text: str, lang: str) -> str:
    text = normalize("NFKD", text)
    emoji_pattern = re.compile(
        "[\U0001f600-\U0001f64f"
        "\U0001f300-\U0001f5ff"
        "\U0001f680-\U0001f6ff"
        "\U0001f700-\U0001f77f"
        "\U0001f780-\U0001f7ff"
        "\U0001f800-\U0001f8ff"
        "\U0001f900-\U0001f9ff"
        "\U0001fa00-\U0001fa6f"
        "\U0001fa70-\U0001faff"
        "\u2600-\u26ff"
        "\u2700-\u27bf"
        "\U0001f1e6-\U0001f1ff]+",
        flags=re.UNICODE,
    )
    text = emoji_pattern.sub("", text)
    replacements = {
        "–": "-",
        "‑": "-",
        "—": "-",
        "_": " ",
        "\u201c": '"',
        "\u201d": '"',
        "\u2018": "'",
        "\u2019": "'",
        "´": "'",
        "`": "'",
        "[": " ",
        "]": " ",
        "|": " ",
        "/": " ",
        "#": " ",
        "→": " ",
        "←": " ",
    }
    for k, v in replacements.items():
        text = text.replace(k, v)
    text = re.sub(r"[♥☆♡©\\]", "", text)
    for k, v in (("@", " at "), ("e.g.,", "for example, "), ("i.e.,", "that is, ")):
        text = text.replace(k, v)
    text = re.sub(r" ,", ",", text)
    text = re.sub(r" \.", ".", text)
    text = re.sub(r" !", "!", text)
    text = re.sub(r" \?", "?", text)
    text = re.sub(r" ;", ";", text)
    text = re.sub(r" :", ":", text)
    text = re.sub(r" '", "'", text)
    while '""' in text:
        text = text.replace('""', '"')
    while "''" in text:
        text = text.replace("''", "'")
    while "``" in text:
        text = text.replace("``", "`")
    text = re.sub(r"\s+", " ", text).strip()
    if not re.search(r"[.!?;:,'\"')\]}…。」』】〉》›»]$", text):
        text += "."
    if lang not in AVAILABLE_LANGS:
        raise ValueError(f"Invalid language: {lang}")
    # HF ``unicode_indexer.json`` maps ASCII ``<`` / ``>`` to -1 (invalid). Use
    # parentheses so language markup is indexable; matches ``SupertonicTTSService``.
    return f"({lang})" + text + " "


def build_text_inputs(processed: str, indexer: list[int], max_len: int) -> tuple[np.ndarray, np.ndarray]:
    ids: list[int] = []
    for ch in processed:
        v = ord(ch)
        if v >= len(indexer):
            raise ValueError("Unsupported character")
        idx = indexer[v]
        if idx < 0:
            raise ValueError("Unsupported character")
        ids.append(idx)
    if len(ids) > max_len:
        raise ValueError("Text too long")
    text_ids = np.zeros((1, max_len), dtype=np.int64)
    text_ids[0, : len(ids)] = np.array(ids, dtype=np.int64)
    mask = np.zeros((1, 1, max_len), dtype=np.float32)
    mask[0, 0, : len(ids)] = 1.0
    return text_ids, mask


def load_voice(path: Path) -> tuple[np.ndarray, np.ndarray]:
    with path.open() as f:
        j = json.load(f)
    ttl = np.array(j["style_ttl"]["data"], dtype=np.float32).reshape(j["style_ttl"]["dims"])
    dp = np.array(j["style_dp"]["data"], dtype=np.float32).reshape(j["style_dp"]["dims"])
    return ttl, dp


def gaussian_like_swift(rng: np.random.Generator, shape: tuple[int, ...]) -> np.ndarray:
    u1 = rng.random(shape, dtype=np.float64)
    u1 = np.maximum(u1, 1e-6)
    u2 = rng.random(shape, dtype=np.float64)
    return (np.sqrt(-2 * np.log(u1)) * np.cos(2 * math.pi * u2)).astype(np.float32)


def sample_noisy_latent(
    duration_s: float,
    sample_rate: int,
    base_chunk_size: int,
    chunk_compress: int,
    latent_dim: int,
    latent_len_max: int,
    rng: np.random.Generator,
) -> tuple[np.ndarray, np.ndarray]:
    wav_len = int(duration_s * sample_rate)
    chunk_size = base_chunk_size * chunk_compress
    latent_len = min((wav_len + chunk_size - 1) // chunk_size, latent_len_max)
    latent = np.zeros((1, latent_dim, latent_len_max), dtype=np.float32)
    noise = gaussian_like_swift(rng, (1, latent_dim, latent_len))
    latent[:, :, :latent_len] = noise[:, :, :latent_len]
    latent_mask = np.zeros((1, 1, latent_len_max), dtype=np.float32)
    latent_mask[0, 0, :latent_len] = 1.0
    return latent, latent_mask


def write_wav_f32(path: Path, samples: np.ndarray, sample_rate: int) -> None:
    s = np.clip(samples.flatten(), -1.0, 1.0)
    pcm = (s * 32767.0).astype(np.int16)
    with wave.open(str(path), "wb") as w:
        w.setnchannels(1)
        w.setsampwidth(2)
        w.setframerate(sample_rate)
        w.writeframes(pcm.tobytes())


def main() -> int:
    ap = argparse.ArgumentParser(description="Supertonic TTS on Mac (ORT + Core ML EP)")
    ap.add_argument("--models-dir", type=Path, required=True, help="Root with onnx/ and voice_styles/")
    ap.add_argument("--text", default="Hello from TranslateBlue.")
    ap.add_argument("--lang", default="en")
    ap.add_argument("--voice", default="", help="voice_styles basename without .json; default: first .json")
    ap.add_argument("--out", type=Path, default=Path("supertonic_mac.wav"))
    ap.add_argument("--steps", type=int, default=20)
    ap.add_argument("--speed", type=float, default=1.0)
    ap.add_argument("--no-coreml", action="store_true", help="Force CPU only")
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument(
        "--max-text",
        type=int,
        default=0,
        help="Pad width T for text_ids/text_mask (default: manifest or 300)",
    )
    ap.add_argument(
        "--latent-len-max",
        type=int,
        default=0,
        help="Max latent time bins (default: 512, matches app fallback)",
    )
    args = ap.parse_args()

    onnx_dir = args.models_dir / "onnx"
    voice_dir = args.models_dir / "voice_styles"
    if not onnx_dir.is_dir():
        raise SystemExit(f"Missing {onnx_dir} (run convert_supertonic_coreml.py first)")

    providers = _providers(not args.no_coreml)
    print("Providers:", providers)

    dp_p = onnx_dir / "duration_predictor.onnx"
    te_p = onnx_dir / "text_encoder.onnx"
    ve_p = onnx_dir / "vector_estimator.onnx"
    voc_p = onnx_dir / "vocoder.onnx"
    for p in (dp_p, te_p, ve_p, voc_p):
        if not p.is_file():
            raise SystemExit(f"Missing ONNX: {p}")

    dp = _session(dp_p, providers)
    te = _session(te_p, providers)
    ve = _session(ve_p, providers)
    voc = _session(voc_p, providers)

    manifest = args.models_dir / "manifest.json"
    manifest_max = 300
    if manifest.is_file():
        manifest_max = int(json.loads(manifest.read_text()).get("max_text_len", 300))

    max_text = args.max_text or _input_dim(dp, "text_mask", 2) or manifest_max
    latent_len_max = args.latent_len_max or _input_dim(ve, "noisy_latent", 2) or 512
    latent_dim_ve = _input_dim(ve, "noisy_latent", 1)
    if latent_dim_ve is None:
        raise SystemExit("Cannot read latent dim from vector_estimator.onnx")
    with (onnx_dir / "tts.json").open() as f:
        cfgs: dict[str, Any] = json.load(f)
    sample_rate = int(cfgs["ae"]["sample_rate"])
    base_chunk = int(cfgs["ae"]["base_chunk_size"])
    chunk_c = int(cfgs["ttl"]["chunk_compress_factor"])
    # Prefer ONNX channel (e.g. 144); config may omit latent_dim in some revisions.
    latent_dim = latent_dim_ve

    with (onnx_dir / "unicode_indexer.json").open() as f:
        indexer: list[int] = json.load(f)

    voice_name = args.voice
    if not voice_name:
        jsons = sorted(voice_dir.glob("*.json"))
        if not jsons:
            raise SystemExit(f"No voice JSON in {voice_dir}")
        voice_name = jsons[0].stem
    voice_path = voice_dir / f"{voice_name}.json"
    if not voice_path.is_file():
        raise SystemExit(f"Voice not found: {voice_path}")

    style_ttl, style_dp = load_voice(voice_path)
    proc = preprocess_text(args.text, args.lang)
    # Swift pads to ONNX ``text_mask`` width (``max_text``), not variable batch max.
    text_ids, text_mask = build_text_inputs(proc, indexer, max_text)

    rng = np.random.default_rng(args.seed)

    dur = dp.run(None, {"text_ids": text_ids, "style_dp": style_dp, "text_mask": text_mask})[0]
    dur_f = float(dur.reshape(-1)[0]) / max(args.speed, 0.01)
    max_dur = latent_len_max * base_chunk * chunk_c / float(sample_rate)
    dur_f = max(0.05, min(dur_f, max_dur))

    te_out = te.run(
        None,
        {"text_ids": text_ids, "style_ttl": style_ttl, "text_mask": text_mask},
    )[0]
    if te_out.shape[1] != 256 or te_out.shape[2] != max_text:
        print("Warning: text_emb shape", te_out.shape, "expected (1,256,T)", file=sys.stderr)

    xt, latent_mask = sample_noisy_latent(
        dur_f, sample_rate, base_chunk, chunk_c, latent_dim, latent_len_max, rng
    )
    total_step = np.array([float(args.steps)], dtype=np.float32)
    for step in range(args.steps):
        cur = np.array([float(step)], dtype=np.float32)
        xt = ve.run(
            None,
            {
                "noisy_latent": xt,
                "text_emb": te_out,
                "style_ttl": style_ttl,
                "text_mask": text_mask,
                "latent_mask": latent_mask,
                "current_step": cur,
                "total_step": total_step,
            },
        )[0]

    outs = voc.run(None, {"latent": xt})
    wav = outs[0]
    # Prefer name wav_tts if present
    for i, o in enumerate(voc.get_outputs()):
        if o.name == "wav_tts" and i < len(outs):
            wav = outs[i]
            break

    trim = min(int(dur_f * sample_rate), int(wav.size))
    wav1 = wav.reshape(-1)[:trim].astype(np.float32)
    peak = float(np.max(np.abs(wav1))) or 1.0
    wav1 *= 0.95 / peak

    args.out.parent.mkdir(parents=True, exist_ok=True)
    write_wav_f32(args.out, wav1, sample_rate)
    print(
        f"Wrote {args.out.resolve()} ({wav1.size / sample_rate:.2f}s, voice={voice_name}, "
        f"max_text_len={max_text})"
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())