Blame for app.py

"""tiny-voca — a 4.5M-parameter English text-to-speech model.

This Space runs the ONNX export of [n4ze3m/tiny-voca](https://huggingface.co/n4ze3m/tiny-voca)
(24 kHz mono, two voices: F2 female, M2 male) with the author's own
[dizione](https://github.com/n4ze3m/dizione) phonemizer: a CMUdict-derived
lexicon, a small neural G2P model for unknown words, and a homograph
classifier ("I read the book you read yesterday").

The utterance pipeline is a port of the author's browser demo
(`web/public/worker.js`): text is split into sentences (without breaking after
titles or initials), long sentences are recursively fitted to the model's
512-token limit at clause boundaries, then word halves, each utterance is
synthesized through the ONNX graph, and the chunks are concatenated.
"""

import json
import math
import os
import re
import sys
import tempfile
from pathlib import Path

import gradio as gr
import numpy as np
import onnxruntime as ort
from huggingface_hub import hf_hub_download, snapshot_download

REPO_ID = "n4ze3m/tiny-voca"
TOKEN_HEADROOM = 12         # the reference demo fits sentences to max_tokens - 12

# ---------------------------------------------------------------- the model

# The phonemizer's built data (lexicon, aligner, G2P weights, homograph model)
# lives in the model repo under phonemizer/; dizione finds it via $DIZIONE_DATA.
_PHONEME_DIR = Path(tempfile.mkdtemp(prefix="dizione-"))
snapshot_download(REPO_ID, allow_patterns=["phonemizer/*"], local_dir=str(_PHONEME_DIR))
os.environ["DIZIONE_DATA"] = str(_PHONEME_DIR / "phonemizer")

sys.path.insert(0, str(Path(__file__).parent))
import dizione  # noqa: E402  (vendored; needs $DIZIONE_DATA set before import)

_CONFIG_PATH = hf_hub_download(REPO_ID, "config.json")
_MODEL_PATH = hf_hub_download(REPO_ID, "onnx/model.onnx")
with open(_CONFIG_PATH, encoding="utf-8") as file:
    CONFIG = json.load(file)

_SESSION = ort.InferenceSession(_MODEL_PATH, providers=["CPUExecutionProvider"])
_MAX_TOKENS = int(CONFIG["model"]["max_tokens"])
SAMPLE_RATE = int(CONFIG["model"]["sample_rate"])
VOICES = tuple(CONFIG["voices"])

dizione.phonemize("Warm up.")   # build the lexicon and G2P tables once, at startup

# ------------------------------------------------ text -> sentences (worker.js)

# Split after . ! ? … (optionally followed by a closing quote or bracket), at
# whitespace, before a new sentence that starts with a capital or a digit.
# Python's re needs fixed-width lookbehinds, so the optional trailing mark
# becomes an alternation of one- and two-character lookbehinds.
_SENTENCE_END = re.compile(
    r'''(?:(?<=[.!?…])|(?<=[.!?…]["')\]]))\s+(?=["'(\[]?[A-Z0-9])''')
# Sentences ending like this are titles, initials or numbers, not sentence ends.
_ABBREVIATION = re.compile(
    r"(?:\b(?:mr|mrs|ms|dr|prof|sr|jr|st|vs|etc|no|inc|ltd|co|e\.g|i\.e|a\.m|p\.m)|\b[A-Z])\.$",
    re.IGNORECASE)
_CLAUSE_END = re.compile(r"(?<=[,;:—])\s+")


def _pieces(text):
    parts, last = [], 0
    for match in _SENTENCE_END.finditer(text):
        parts.append(text[last:match.start()])
        last = match.end()
    parts.append(text[last:])
    return parts


def sentences(text):
    """Sentences, without breaking after titles, initials or inside numbers."""
    text = re.sub(r"\s+", " ", text).strip()
    parts, current = [], ""
    for piece in _pieces(text):
        current = f"{current} {piece}" if current else piece
        if not _ABBREVIATION.search(current):
            parts.append(current)
            current = ""
    if current:
        parts.append(current)
    return [part for part in parts if any(c.isalpha() or c.isdigit() for c in part)]


def fit(sentence):
    """Phonemized utterances of one sentence, each within the token limit:
    cut at clause boundaries, then at word halves, like the reference demo."""
    try:
        result = dizione.phonemize(sentence)
    except ValueError:      # nothing pronounceable in this piece
        return []
    if len(result.ids) <= _MAX_TOKENS - TOKEN_HEADROOM:
        return [result]
    clauses = _CLAUSE_END.split(sentence)
    if len(clauses) > 1:
        return [u for clause in clauses for u in fit(clause)]
    words = sentence.split(" ")
    if len(words) < 2:
        return []           # a single word too long for one utterance: drop it
    half = math.ceil(len(words) / 2)
    return fit(" ".join(words[:half])) + fit(" ".join(words[half:]))


def _synthesize(phones, voice, speed, pitch_shift, pitch_range):
    """One utterance through the ONNX graph -> float32 [samples] at 24 kHz."""
    count = len(phones.ids)
    audio, _durations = _SESSION.run(None, {
        "tokens": np.array([phones.ids], dtype=np.int64),
        "stress": np.array([phones.stress], dtype=np.int64),
        "voice": np.array([voice], dtype=np.int64),
        "speed": np.array([speed], dtype=np.float32),
        "pitch_shift": np.array([pitch_shift], dtype=np.float32),
        "pitch_range": np.array([pitch_range], dtype=np.float32),
    })
    return audio[0]


# ------------------------------------------------------------- the demo

CSS = """
#col-container { max-width: 880px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""


def speak(text, voice="F2", speed=1.0, pitch_shift=0.0, pitch_range=1.0):
    """Synthesize speech from English text.

    Args:
        text: English text to speak. Long text is split into sentences; each
            sentence into utterances that fit the model's 512-token limit.
        voice: "F2" (female) or "M2" (male).
        speed: speaking rate, 0.7 (slow) to 1.4 (fast).
        pitch_shift: added to log F0; 0.058 is about one semitone.
        pitch_range: scale of pitch movement around the utterance mean
            (1 = as predicted, <1 flatter, >1 more lively).
    """
    if not isinstance(text, str) or not text.strip():
        raise gr.Error("Enter some text to speak.")
    chunks, transcript = [], []
    for sentence in sentences(text):
        for phones in fit(sentence):
            chunks.append(_synthesize(phones, VOICES.index(voice), speed, pitch_shift, pitch_range))
            transcript.append(str(phones))
    if not chunks:
        raise gr.Error("That text has nothing to pronounce — try some English words.")
    audio = np.concatenate(chunks)
    if not np.isfinite(audio).all():
        raise gr.Error("The model produced non-finite audio; try different text or settings.")
    audio = np.clip(audio, -1.0, 1.0)
    return (SAMPLE_RATE, audio.astype(np.float32)), " ".join(transcript)


with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            "# 🗣️ tiny-voca\n"
            "A 4.5M-parameter English text-to-speech model — two voices, 24 kHz audio, "
            "~70× real-time on CPU. Weights: [n4ze3m/tiny-voca](https://huggingface.co/n4ze3m/tiny-voca) "
            "(Apache-2.0) · phonemized by the author's "
            "[dizione](https://github.com/n4ze3m/dizione) dictionary + neural G2P."
        )
        with gr.Row():
            text = gr.Textbox(
                label="Text",
                placeholder="Type anything in English — try a homograph sentence…",
                lines=4,
                autofocus=True,
                scale=4,
            )
            run = gr.Button("Speak", variant="primary", scale=1)
        audio = gr.Audio(label="Speech", type="numpy")
        phones = gr.Textbox(label="Phonemes (dizione)", interactive=False, lines=2)
        with gr.Row():
            voice = gr.Radio(choices=list(VOICES), value="F2", label="Voice")
            speed = gr.Slider(minimum=0.7, maximum=1.4, value=1.0, step=0.05, label="Speed")
        with gr.Accordion("Pitch controls", open=False):
            pitch_shift = gr.Slider(minimum=-0.174, maximum=0.174, value=0.0, step=0.058,
                                    label="Pitch shift", info="0.058 ≈ one semitone; 0 = as trained")
            pitch_range = gr.Slider(minimum=0.5, maximum=1.5, value=1.0, step=0.1,
                                    label="Pitch range", info="Scale of pitch movement around the utterance mean")
        gr.Examples(
            examples=[
                ["Hello! This is tiny voca, a small text to speech model running entirely in your browser."],
                ["The meeting moved to 3:30 pm on March 14th, so we have about 45 minutes left. It costs $12.99."],
                ["Did you really think I would forget your birthday? Of course not!"],
                ["She had never seen the sea before, and for a long while she stood at the edge of the water without saying a word."],
                ["I read the book you read yesterday. They live near the live music hall."],
            ],
            inputs=[text],
            outputs=[audio, phones],
            fn=speak,
            cache_examples=True,
            cache_mode="lazy",
        )
    run.click(speak, inputs=[text, voice, speed, pitch_shift, pitch_range],
              outputs=[audio, phones], api_name="speak")

demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)