python/inflect_ax_tts/frontend.py

"""VITS text frontend: text -> phoneme token ids.

Faithful port of origin/inflect_vits_frontend.py + the token mapping from
origin/runtime/text/__init__.py + commons.intersperse. Text normalization and
eSpeak configuration live in text_normalize.py (verbatim copy of
origin/inflect_nano_v2_frontend.py).

Pipeline (mirrors origin/inference.py:InflectTTS._tokens):
    text -> normalize_text -> eSpeak phonemize (en-us, with_stress,
    preserve_punctuation) -> phoneme overrides -> per-character symbol ids
    -> intersperse(0)  (add_blank)
"""

from __future__ import annotations

from .symbols import symbols
from .text_normalize import _configure_espeak, normalize_text

__all__ = ["normalize_text", "phonemize_text", "text_to_token_ids", "intersperse"]

# eSpeak is the general fallback. This table contains verified exceptions only
# (verbatim from origin/inflect_vits_frontend.py).
PHONEME_OVERRIDES = {
    "sˈæskɐtʃˌuːən": "sɐskˈætʃəwən",
    "flʊɹɹˈɛsənt": "flʊˈɹɛsənt",
}

_SYMBOL_TO_ID = {s: i for i, s in enumerate(symbols)}
_BLANK_ID = 0  # symbols[0] == '_'


def intersperse(seq: list[int], item: int = _BLANK_ID) -> list[int]:
    """commons.intersperse: insert `item` around every element -> 2N+1."""
    result = [item] * (len(seq) * 2 + 1)
    result[1::2] = seq
    return result


def phonemize_text(text: str) -> str:
    """normalize + eSpeak phonemize + overrides (run_vits_frontend equivalent)."""
    import re

    normalized = normalize_text(text)
    _configure_espeak()
    from phonemizer import phonemize

    phoneme_text = phonemize(
        [normalized],
        language="en-us",
        backend="espeak",
        strip=True,
        preserve_punctuation=True,
        with_stress=True,
        njobs=1,
    )[0]
    for source, replacement in PHONEME_OVERRIDES.items():
        phoneme_text = phoneme_text.replace(source, replacement)
    return re.sub(r"\s+", " ", phoneme_text).strip()


def text_to_token_ids(text: str, add_blank: bool = True) -> list[int]:
    """Full frontend: text -> interspersed phoneme ids (2N+1 with add_blank)."""
    phoneme_text = phonemize_text(text)
    try:
        sequence = [_SYMBOL_TO_ID[ch] for ch in phoneme_text]
    except KeyError as exc:
        raise ValueError(
            f"phoneme {exc.args[0]!r} from {phoneme_text!r} is not in the "
            "model symbol table"
        ) from exc
    if add_blank:
        sequence = intersperse(sequence, _BLANK_ID)
    if not sequence:
        raise ValueError("The text frontend produced no speakable tokens.")
    return sequence