frontend.py

"""Self-contained text->token frontend for the ONNX demo.

Mirrors `inference.text_to_tokens` but imports ONLY the light text path, so it
avoids `inference.py`'s torch-model imports (inflect_nano.vocoder pulls torchaudio,
which the ONNX demo never needs).
"""
from __future__ import annotations

import sys
from pathlib import Path

import torch

# g2p_en triggers nltk loads at import (G2p() is constructed at module level in
# tiny_tts.text.english). Newer nltk renamed the tagger to
# 'averaged_perceptron_tagger_eng'; download every name g2p_en might want first.
import nltk
for _pkg in ("averaged_perceptron_tagger", "averaged_perceptron_tagger_eng", "cmudict"):
    try:
        nltk.download(_pkg, quiet=True)
    except Exception:
        pass

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))
sys.path.insert(0, str(HERE / "third_party" / "tiny_tts_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


def text_to_tokens(text: str):
    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 torch.LongTensor(phone_ids), torch.LongTensor(tone_ids), torch.LongTensor(lang_ids)