teratts.py

#!/usr/bin/env python3
"""Standalone ONNX Runtime command-line inference for the encoder-free release."""

from __future__ import annotations

import json
import math
import os
import re
import unicodedata
import warnings
import wave
from collections.abc import Iterator
from dataclasses import dataclass
from functools import lru_cache
from pathlib import Path

import numpy as np
import onnxruntime as ort

try:  # Local CLI import and Hugging Face remote-code package import.
    from .teratts_ruaccent import RUAccent
except ImportError:  # pragma: no cover - exercised by ``python teratts.py``.
    from teratts_ruaccent import RUAccent

SAMPLE_RATE = 44_100
SAMPLES_PER_COMPRESSED_FRAME = 3_072
VOCODER_CONTEXT_FRAMES = 20
DEFAULT_STREAM_CHUNK_FRAMES = 16
SPEED = 1.05
SEED = 1234
RUSSIAN_TAG = re.compile(r"<ru>(.*?)</ru>", flags=re.DOTALL)
LANGUAGE_TAG = re.compile(r"<(ru|en)>(.*?)</\1>", flags=re.DOTALL)
LANGUAGE_TAG_TOKEN = re.compile(r"<(/?)([a-z]{2})>")
TAGGED_NUMBER = re.compile(r"(?<![\w.])[-−]?\d+(?:[.,]\d+)?(?![\w.])")
PUNCTUATION_NEEDS_SPACE = re.compile(r"[,.!?;:…](?=[^\s<])")
NUMBER_NEEDS_SPACE = re.compile(r"(?<=\d)(?=[A-Za-zА-Яа-яЁё])")


def prepare_raw_text(raw_text: str) -> tuple[str, str]:
    model_text = unicodedata.normalize("NFKD", raw_text)
    return model_text, model_text.replace("+", "")


def _add_punctuation_spaces(text: str) -> str:
    """Separate punctuation without splitting decimal literals or closing tags."""
    def space_after(match: re.Match[str]) -> str:
        punctuation = match.group(0)
        index = match.start()
        previous = text[index - 1] if index else ""
        following = text[index + 1] if index + 1 < len(text) else ""
        if punctuation in ".," and previous.isdigit() and following.isdigit():
            return punctuation
        return punctuation + " "

    return PUNCTUATION_NEEDS_SPACE.sub(space_after, text)


def validate_language_tags(text: str) -> None:
    """Require balanced ``<ru>`` / ``<en>`` spans for all public synthesis."""
    tokens = list(LANGUAGE_TAG_TOKEN.finditer(text))
    if not tokens or not LANGUAGE_TAG.search(text):
        raise ValueError(
            "text must contain a language tag: wrap text in <ru>...</ru> or <en>...</en>"
        )
    stack: list[str] = []
    for token in tokens:
        closing, language = token.groups()
        if language not in {"ru", "en"}:
            raise ValueError(f"unsupported language tag <{language}>; use <ru> or <en>")
        if not closing:
            stack.append(language)
        elif not stack or stack.pop() != language:
            raise ValueError("language tags must be balanced: use <ru>...</ru> or <en>...</en>")
    if stack:
        raise ValueError("language tags must be balanced: use <ru>...</ru> or <en>...</en>")
    # Angle brackets that did not form a valid tag would be accepted by the
    # character vocabulary but are not meaningful model input.
    if "<" in LANGUAGE_TAG_TOKEN.sub("", text) or ">" in LANGUAGE_TAG_TOKEN.sub("", text):
        raise ValueError("invalid language tags; use only <ru>...</ru> or <en>...</en>")


def _skip_unsupported_characters(
    text: str,
    indexer: "UnicodeIndexer",
    *,
    preserve_digits: bool = False,
) -> str:
    """Return supported text and issue one clear warning for skipped characters."""
    kept: list[str] = []
    skipped: list[str] = []
    for character in text:
        # The released table was trained on NFKD text.  Keep the human-readable
        # NFC spelling here (especially ``й`` and ``ё``) as long as all of its
        # decomposed codepoints exist in the table.  RUAccent must receive this
        # spelling: passing ``и`` + COMBINING BREVE makes its text cleaner drop
        # the breve and turn ``й`` into ``и``.
        encoded = unicodedata.normalize("NFKD", character)
        supported = bool(encoded) and all(
            (indexer.table[ord(item)] if ord(item) < 65_536 else -1) >= 0
            for item in encoded
        )
        if not supported and not (preserve_digits and character.isdigit()):
            skipped.append(character)
        else:
            kept.append(character)
    if skipped:
        labels = ", ".join(
            f"{character!r} (U+{ord(character):04X})" for character in sorted(set(skipped))
        )
        warnings.warn(
            f"skipped unsupported characters not present in the TeraTTS vocabulary: {labels}",
            RuntimeWarning,
            stacklevel=2,
        )
    return "".join(kept)


def normalize_input_text(raw_text: str, indexer: "UnicodeIndexer") -> str:
    """Normalize spacing and skip unsupported vocabulary characters with a warning."""
    if not isinstance(raw_text, str) or not raw_text.strip():
        raise ValueError("text must not be empty; use <ru>...</ru> or <en>...</en>")
    # Retain composed characters through RUAccent.  ``prepare_raw_text``
    # performs the required NFKD conversion immediately before ONNX encoding.
    text = unicodedata.normalize("NFC", raw_text)
    text = _add_punctuation_spaces(text)
    text = NUMBER_NEEDS_SPACE.sub(" ", text)
    # Digits are retained only long enough for tagged ``num2words`` expansion;
    # any remaining unsupported digits are skipped after that expansion.
    text = _skip_unsupported_characters(text, indexer, preserve_digits=True)
    validate_language_tags(text)
    return text


def load_ruaccent(
    *,
    model_size: str = "turbo3.1",
    device: str = "CPU",
    workdir: Path | None = None,
    mode: str = "full",
) -> object:
    """Load the bundled RUAccent-derived ONNX models without downloading."""
    if workdir is None:
        raise ValueError("load_ruaccent requires the release's ruaccent asset directory")
    return RUAccent(workdir, model_size=model_size, device=device, mode=mode)


def add_russian_stress(text: str, accentizer: object | None) -> str:
    """Fill stress marks in ``<ru>`` spans while preserving manual markers."""
    if accentizer is None:
        return text

    def accent(match: re.Match[str]) -> str:
        content = match.group(1)
        # Explicit stress from the caller is authoritative. RUAccent is only
        # used for spans that have not already been annotated.
        if "+" in content:
            return match.group(0)
        process_all = getattr(accentizer, "process_all")
        return f"<ru>{process_all(content)}</ru>"

    return RUSSIAN_TAG.sub(accent, text)


def expand_tagged_numbers(text: str) -> str:
    """Spell out numeric literals inside ``<ru>`` and ``<en>`` text spans.

    Language tags are intentionally required: this avoids guessing a language
    for bare text or for identifiers such as versions and file names.
    """
    spans = list(LANGUAGE_TAG.finditer(text))
    if not any(TAGGED_NUMBER.search(match.group(2)) for match in spans):
        return text
    try:
        from num2words import num2words
    except ImportError as error:
        raise RuntimeError(
            "number expansion requires num2words; install the model requirements"
        ) from error

    def expand_span(match: re.Match[str]) -> str:
        language, content = match.groups()

        def expand_number(number: re.Match[str]) -> str:
            literal = number.group(0).replace("−", "-")
            value: int | float
            if "." in literal or "," in literal:
                value = float(literal.replace(",", "."))
            else:
                value = int(literal)
            return str(num2words(value, lang=language))

        return f"<{language}>{TAGGED_NUMBER.sub(expand_number, content)}</{language}>"

    return LANGUAGE_TAG.sub(expand_span, text)


def normalize_text(loaded: "LoadedTTS", text: str) -> str:
    """Return the exact text tensorized by the text encoder for an utterance."""
    normalized_input = normalize_input_text(text, loaded.indexer)
    expanded_text = _skip_unsupported_characters(
        expand_tagged_numbers(normalized_input), loaded.indexer
    )
    model_text, _ = prepare_raw_text(add_russian_stress(expanded_text, loaded.accentizer))
    return model_text


class UnicodeIndexer:
    def __init__(self, indexer_path: Path):
        self.table = json.loads(indexer_path.read_text())
        if len(self.table) != 65_536:
            raise ValueError("unicode_indexer.json must have 65,536 entries")

    def batch(self, text: str) -> tuple[np.ndarray, np.ndarray]:
        ids = []
        for character in text:
            token = self.table[ord(character)] if ord(character) < 65_536 else -1
            if token < 0:
                raise ValueError(
                    f"unsupported character {character!r} (U+{ord(character):04X})"
                )
            ids.append(token)
        if not ids:
            raise ValueError("text produced no tokens")
        values = np.asarray(ids, dtype=np.int64)[None, :]
        return values, np.ones((1, 1, values.shape[1]), dtype=np.float32)


def write_wav(path: Path, samples: np.ndarray) -> None:
    pcm16 = np.clip(samples, -1.0, 1.0)
    pcm16 = np.rint(pcm16 * 32767.0).astype("<i2")
    with wave.open(str(path), "wb") as output:
        output.setnchannels(1)
        output.setsampwidth(2)
        output.setframerate(SAMPLE_RATE)
        output.writeframes(pcm16.tobytes())


def play_stream(
    chunks: Iterator[np.ndarray], *, device: int | str | None = None
) -> None:
    """Play streamed mono float32 chunks through a Windows/Linux/macOS device.

    This requires the optional ``sounddevice`` dependency. ``device`` accepts
    a PortAudio device ID or name; ``None`` uses the operating-system default.
    The supplied iterator is consumed exactly once.
    """
    try:
        import sounddevice as sd
    except ImportError as error:
        raise RuntimeError(
            "stream playback requires sounddevice; install the project's audio extra"
        ) from error

    # RawOutputStream accepts buffer objects, avoiding any extra float32 copy
    # after the unavoidable PCM conversion for the audio device.
    with sd.RawOutputStream(
        samplerate=SAMPLE_RATE,
        channels=1,
        dtype="int16",
        device=device,
    ) as output:
        for chunk in chunks:
            pcm16 = np.rint(np.clip(chunk, -1.0, 1.0) * 32767.0).astype("<i2")
            output.write(pcm16.tobytes())


def _cpu_thread_count(
    execution_providers: tuple[str, ...], threads: int | None
) -> int | None:
    """Choose CPU inference parallelism, defaulting to physical-core scale."""
    if execution_providers != ("CPUExecutionProvider",):
        return None
    if threads is not None:
        if threads < 1:
            raise ValueError("threads must be positive")
        return threads

    available = os.cpu_count() or 1
    # Most desktop CPUs expose two logical threads per physical core. Limiting
    # one inference to that physical-core count avoids the oversubscription
    # measured on the target Ryzen 5 5600X; pass ``threads`` when another CPU
    # topology needs a different choice.
    return max(1, available // 2)


@lru_cache(maxsize=12)
def _cached_session(
    model_path: str,
    execution_providers: tuple[str, ...],
    cpu_threads: int | None,
    revision: tuple[int, int],
) -> ort.InferenceSession:
    del revision  # It is part of the cache key so replaced model files reload.
    options = ort.SessionOptions()
    if cpu_threads is not None:
        options.intra_op_num_threads = cpu_threads
        options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
    return ort.InferenceSession(model_path, sess_options=options, providers=list(execution_providers))


def session(
    models: Path,
    name: str,
    execution_providers: list[str],
    *,
    threads: int | None = None,
) -> ort.InferenceSession:
    """Load a reusable session, reloading automatically after a file swap."""
    model_path = (models / name).resolve()
    stamp = model_path.stat()
    providers = tuple(execution_providers)
    return _cached_session(
        str(model_path),
        providers,
        _cpu_thread_count(providers, threads),
        (stamp.st_mtime_ns, stamp.st_size),
    )


def clear_session_cache() -> None:
    """Release cached ONNX sessions, useful before an in-place model swap."""
    _cached_session.cache_clear()


@dataclass(frozen=True)
class LoadedTTS:
    """Reusable encoder-free TTS runtime loaded from one release directory."""

    release: Path
    model: str
    text_encoder: ort.InferenceSession
    duration_predictor: ort.InferenceSession
    sampler: ort.InferenceSession
    vocoder: ort.InferenceSession
    indexer: UnicodeIndexer
    accentizer: object | None


def load_model(
    release: Path,
    *,
    model: str = "distilled",
    provider: str = "CPUExecutionProvider",
    threads: int | None = None,
    russian_stress: bool = True,
    ruaccent_model_size: str = "turbo3.1",
    ruaccent_device: str = "CPU",
    ruaccent_workdir: Path | None = None,
    ruaccent_mode: str = "full",
) -> LoadedTTS:
    """Load reusable ONNX sessions; call once before generating many utterances.

    ``threads`` controls CPU intra-op parallelism. ``None`` chooses a
    physical-core-scale default; CUDA ignores this value. Russian ``<ru>``
    spans receive automatic ``+`` stress markers when ``russian_stress`` is
    enabled; manually supplied markers are preserved. ``ruaccent_mode`` is
    ``"full"`` (neural ONNX models plus dictionaries) or ``"dictionary"``
    (dictionaries only, with no accentuation-model ONNX sessions).
    """
    if model not in {"teacher", "distilled"}:
        raise ValueError("model must be 'teacher' or 'distilled'")
    release = release.resolve()
    models = release / "models"
    providers = [provider]
    sampler_name = (
        "sampler_teacher_8step.onnx"
        if model == "teacher"
        else "sampler_distilled_cfg3_8step.onnx"
    )
    return LoadedTTS(
        release=release,
        model=model,
        text_encoder=session(models, "text_encoder.onnx", providers, threads=threads),
        duration_predictor=session(
            models, "duration_predictor.onnx", providers, threads=threads
        ),
        sampler=session(models, sampler_name, providers, threads=threads),
        vocoder=session(models, "vocoder.onnx", providers, threads=threads),
        indexer=UnicodeIndexer(release / "unicode_indexer.json"),
        accentizer=(
            load_ruaccent(
                model_size=ruaccent_model_size,
                device=ruaccent_device,
                workdir=ruaccent_workdir or release / "ruaccent",
                mode=ruaccent_mode,
            )
            if russian_stress
            else None
        ),
    )


def iter_vocoder_audio(
    vocoder: ort.InferenceSession,
    latent: np.ndarray,
    *,
    chunk_frames: int = DEFAULT_STREAM_CHUNK_FRAMES,
    maximum_samples: int | None = None,
) -> Iterator[np.ndarray]:
    """Decode a latent in causal overlap-save chunks.

    Each yielded array is mono float32 audio at 44,100 Hz.  A consumer can send
    it directly to a playback or network sink.  Concatenating the chunks
    matches a full vocoder decode up to normal floating-point kernel variation.
    """
    if latent.ndim != 3 or latent.shape[:2] != (1, 144):
        raise ValueError("streaming expects latent shape [1, 144, frames]")
    if chunk_frames < 1:
        raise ValueError("chunk_frames must be positive")

    total_frames = latent.shape[-1]
    full_samples = total_frames * SAMPLES_PER_COMPRESSED_FRAME
    if maximum_samples is None:
        maximum_samples = full_samples
    maximum_samples = max(0, min(int(maximum_samples), full_samples))

    emitted = 0
    for start in range(0, total_frames, chunk_frames):
        end = min(start + chunk_frames, total_frames)
        input_start = max(0, start - VOCODER_CONTEXT_FRAMES)
        decoded = vocoder.run(None, {"latent": latent[..., input_start:end]})[0]
        if decoded.ndim != 2 or decoded.shape[0] != 1:
            raise ValueError("vocoder returned an unexpected waveform shape")
        decoded = decoded[0]
        discard = (start - input_start) * SAMPLES_PER_COMPRESSED_FRAME
        new_samples = (end - start) * SAMPLES_PER_COMPRESSED_FRAME
        chunk = decoded[discard : discard + new_samples]
        if chunk.shape[0] != new_samples:
            raise ValueError("vocoder returned fewer samples than its latent input requires")

        remaining = maximum_samples - emitted
        if remaining <= 0:
            break
        chunk = chunk[:remaining]
        if chunk.size:
            emitted += chunk.size
            yield chunk


def _generate_latent(
    loaded: LoadedTTS,
    text: str,
    voice: str,
    duration_scale: float,
    *,
    guidance: float,
    seed: int,
) -> tuple[ort.InferenceSession, np.ndarray, int]:
    if not math.isfinite(guidance) or guidance < 0:
        raise ValueError("guidance must be finite and non-negative")
    if not math.isfinite(duration_scale) or duration_scale <= 0:
        raise ValueError("duration_scale must be finite and positive")
    voice_dir = loaded.release / "styles" / voice
    if not voice_dir.is_dir():
        choices = ", ".join(
            path.name
            for path in sorted((loaded.release / "styles").glob("*"))
            if path.is_dir()
        )
        raise ValueError(f"unknown voice {voice!r}; choices: {choices or '(none)'}")
    style_ttl = np.load(voice_dir / "style_ttl.npy").astype(np.float32, copy=False)
    style_dp = np.load(voice_dir / "style_dp.npy").astype(np.float32, copy=False)
    if style_ttl.shape != (1, 50, 256) or style_dp.shape != (1, 8, 16):
        raise ValueError("style assets have unexpected shapes")

    model_text = normalize_text(loaded, text)
    duration_text = model_text.replace("+", "")
    text_ids, text_mask = loaded.indexer.batch(model_text)
    duration_ids, duration_mask = loaded.indexer.batch(duration_text)
    text_emb = loaded.text_encoder.run(
        None,
        {
            "text_ids": text_ids,
            "style_ttl": style_ttl,
            "text_mask": text_mask,
        },
    )[0]
    raw_duration = loaded.duration_predictor.run(
        None,
        {
            "text_ids": duration_ids,
            "style_dp": style_dp,
            "text_mask": duration_mask,
        },
    )[0]
    duration_seconds = float(raw_duration[0]) * duration_scale / SPEED
    if not math.isfinite(duration_seconds) or duration_seconds <= 0:
        raise ValueError("duration predictor returned a non-positive duration")
    latent_length = max(
        1, math.ceil(duration_seconds * SAMPLE_RATE / SAMPLES_PER_COMPRESSED_FRAME)
    )
    latent_mask = np.ones((1, 1, latent_length), dtype=np.float32)
    latent = (
        np.random.default_rng(seed)
        .standard_normal((1, 144, latent_length))
        .astype(np.float32)
    )
    # The selected sampler graph owns its diffusion architecture and complete
    # Euler schedule.  Replacing it with another graph that keeps this input
    # contract changes the diffusion model without changing host code.
    latent = loaded.sampler.run(
        None,
        {
            "initial_latent": latent,
            "text_emb": text_emb,
            "style_ttl": style_ttl,
            "latent_mask": latent_mask,
            "text_mask": text_mask,
            "guidance": np.asarray([guidance], dtype=np.float32),
        },
    )[0]
    maximum_samples = round(duration_seconds * SAMPLE_RATE)
    return loaded.vocoder, latent, maximum_samples


def generate_speech_stream(
    loaded: LoadedTTS,
    text: str,
    voice: str,
    *,
    duration_scale: float = 1.0,
    guidance: float = 3.0,
    seed: int = SEED,
    chunk_frames: int = DEFAULT_STREAM_CHUNK_FRAMES,
) -> Iterator[np.ndarray]:
    """Generate the latent, then yield vocoder audio as it becomes available.

    The flow-model sampling phase necessarily completes before this iterator
    emits its first audio chunk.  The default chunk has 16 compressed frames
    (49,152 samples); use a smaller positive value to reduce playback latency.
    ``guidance`` is used only by the teacher model and is ignored by distilled.
    """
    vocoder, latent, maximum_samples = _generate_latent(
        loaded,
        text,
        voice,
        duration_scale,
        guidance=guidance,
        seed=seed,
    )
    yield from iter_vocoder_audio(
        vocoder,
        latent,
        chunk_frames=chunk_frames,
        maximum_samples=maximum_samples,
    )


def generate_speech(
    loaded: LoadedTTS,
    text: str,
    voice: str,
    *,
    duration_scale: float = 1.0,
    guidance: float = 3.0,
    seed: int = SEED,
) -> np.ndarray:
    """Generate and fully decode one utterance, trimmed to audible duration."""
    vocoder, latent, maximum_samples = _generate_latent(
        loaded,
        text,
        voice,
        duration_scale,
        guidance=guidance,
        seed=seed,
    )
    waveform = vocoder.run(None, {"latent": latent})[0]
    if waveform.ndim != 2 or waveform.shape[0] != 1:
        raise ValueError("vocoder returned an unexpected waveform shape")
    return waveform[0, :maximum_samples]


def synthesize_stream(
    release: Path,
    text: str,
    voice: str,
    model: str,
    duration_scale: float,
    provider: str,
    *,
    guidance: float = 3.0,
    seed: int = SEED,
    chunk_frames: int = DEFAULT_STREAM_CHUNK_FRAMES,
    threads: int | None = None,
) -> Iterator[np.ndarray]:
    """Compatibility wrapper around :func:`load_model` and streaming generation."""
    loaded = load_model(release, model=model, provider=provider, threads=threads)
    yield from generate_speech_stream(
        loaded,
        text,
        voice,
        duration_scale=duration_scale,
        guidance=guidance,
        seed=seed,
        chunk_frames=chunk_frames,
    )


def synthesize(
    release: Path,
    text: str,
    voice: str,
    model: str,
    duration_scale: float,
    provider: str,
    *,
    guidance: float = 3.0,
    seed: int = SEED,
    threads: int | None = None,
) -> np.ndarray:
    """Compatibility wrapper around :func:`load_model` and full generation."""
    loaded = load_model(release, model=model, provider=provider, threads=threads)
    return generate_speech(
        loaded,
        text,
        voice,
        duration_scale=duration_scale,
        guidance=guidance,
        seed=seed,
    )