tts.py

"""Inference engine for the Dialogs-RU expressive VITS2 model.

This module is intentionally free of any UI dependency so it can be imported
and tested on its own. ``app.py`` builds the Gradio interface on top of it.

The model is a multi-band iSTFT VITS2 conditioned on a speaker id and an
emotion id. Russian text is used as-is (Cyrillic, lowercased); an optional
``+`` marks the stressed vowel, which the model was trained with. Stress is
placed automatically with ``ruaccent`` when available.
"""

import json
import os
import re

import numpy as np
import torch

import commons
from models import SynthesizerTrn
from text import text_to_sequence
from text.symbols import symbols

# Russian TTS text normalizer (numbers, dates, money, units, abbreviations…).
# Vendored single file: github.com/shigabeev/russian_tts_normalization (MIT).
try:
    from text.rutextnorm import normalize_russian
except Exception:  # pragma: no cover - fall back to the installed package
    try:
        from rutextnorm import normalize_russian
    except Exception:
        normalize_russian = None

DEVICE = "cpu"
MODEL_REPO = os.environ.get("DIALOGS_MODEL_REPO", "frappuccino/dialogs-ru-vits2")
WEIGHTS_FILE = "averaged_G_615000.pth"

# --- speaker / emotion catalogues -------------------------------------------
# speaker id -> (English label, Russian label, gender)
SPEAKERS = {
    0: ("Masha", "Маша", "female"),
    1: ("Sveta", "Света", "female"),
    2: ("Dima", "Дима", "male"),
}

# emotion id -> (English label, Russian label)  (training order, 13 styles)
EMOTIONS = {
    0: ("neutral", "нейтральный"),
    1: ("happy", "радостный"),
    2: ("surprise", "удивление"),
    3: ("arrogance", "высокомерие"),
    4: ("yawn", "зевота"),
    5: ("fear", "страх"),
    6: ("laughing", "смех"),
    7: ("whisper", "шёпот"),
    8: ("disgust", "отвращение"),
    9: ("angry", "злость"),
    10: ("sad", "грусть"),
    11: ("tongue-twister", "скороговорка"),
    12: ("poem", "стихи"),
}


def _resolve_assets():
    """Return (config_path, weights_path).

    Resolution order:
      1. ``DIALOGS_LOCAL_CKPT`` env var (a folder with config.json + weights),
      2. files sitting next to this module (model-repo clone),
      3. download from the Hugging Face Hub.
    """
    local = os.environ.get("DIALOGS_LOCAL_CKPT")
    if local:
        return os.path.join(local, "config.json"), os.path.join(local, WEIGHTS_FILE)

    here = os.path.dirname(os.path.abspath(__file__))
    cfg_local = os.path.join(here, "config.json")
    wts_local = os.path.join(here, WEIGHTS_FILE)
    if os.path.isfile(cfg_local) and os.path.isfile(wts_local):
        return cfg_local, wts_local

    from huggingface_hub import hf_hub_download

    cfg = hf_hub_download(MODEL_REPO, "config.json")
    wts = hf_hub_download(MODEL_REPO, WEIGHTS_FILE)
    return cfg, wts


def _load_checkpoint(path, model):
    """Load generator weights, tolerating missing/extra keys."""
    ckpt = torch.load(path, map_location="cpu")
    saved = ckpt["model"] if "model" in ckpt else ckpt
    state = model.state_dict()
    new_state = {}
    for k in state:
        new_state[k] = saved[k] if k in saved else state[k]
    model.load_state_dict(new_state)
    return ckpt.get("iteration")


class DialogsTTS:
    def __init__(self):
        config_path, weights_path = _resolve_assets()
        with open(config_path, "r", encoding="utf-8") as f:
            cfg = json.load(f)
        self.cfg = cfg
        self.sampling_rate = int(cfg["data"]["sampling_rate"])
        self.add_blank = bool(cfg["data"].get("add_blank", True))
        self.cleaners = cfg["data"].get("text_cleaners", ["basic_cleaners"])

        net_g = SynthesizerTrn(
            len(symbols),
            80,  # mel posterior channels (VITS2)
            cfg["train"]["segment_size"] // cfg["data"]["hop_length"],
            n_speakers=cfg["data"]["n_speakers"],
            n_emotions=cfg["data"]["n_emotions"],
            **cfg["model"],
        ).to(DEVICE)
        net_g.eval()
        self.iteration = _load_checkpoint(weights_path, net_g)
        self.net_g = net_g

        # Optional automatic stress placement.
        self.accentizer = None
        try:
            from ruaccent import RUAccent

            acc = RUAccent()
            acc.load(omograph_model_size="tiny", use_dictionary=True)
            self.accentizer = acc
        except Exception as e:  # pragma: no cover - best effort
            print("ruaccent unavailable, stress will not be auto-placed:", repr(e))

    # -- text helpers --------------------------------------------------------
    def normalize_text(self, text):
        if normalize_russian is None:
            return text
        try:
            return normalize_russian(text)
        except Exception as e:  # pragma: no cover
            print("rutextnorm failed on input:", repr(e))
            return text

    def add_stress(self, text):
        if self.accentizer is None:
            return text
        try:
            return self.accentizer.process_all(text)
        except Exception as e:  # pragma: no cover
            print("ruaccent failed on input:", repr(e))
            return text

    def _to_tokens(self, text):
        seq = text_to_sequence(text, self.cleaners)
        if self.add_blank:
            seq = commons.intersperse(seq, 0)
        return torch.LongTensor(seq)

    # -- synthesis -----------------------------------------------------------
    @torch.no_grad()
    def synthesize(
        self,
        text,
        speaker_id=0,
        emotion_id=0,
        auto_stress=True,
        normalize=True,
        noise_scale=0.667,
        noise_scale_w=0.8,
        length_scale=1.0,
    ):
        if not text or not text.strip():
            raise ValueError("Empty text")
        used = self.normalize_text(text) if normalize else text
        used = self.add_stress(used) if auto_stress else used
        tokens = self._to_tokens(used)
        if tokens.numel() == 0:
            raise ValueError("No recognizable characters in the input text")
        x = tokens.unsqueeze(0).to(DEVICE)
        x_len = torch.LongTensor([tokens.size(0)]).to(DEVICE)
        sid = torch.LongTensor([int(speaker_id)]).to(DEVICE)
        emo = torch.LongTensor([int(emotion_id)]).to(DEVICE)
        audio = self.net_g.infer(
            x,
            x_len,
            sid=sid,
            emotion=emo,
            noise_scale=float(noise_scale),
            noise_scale_w=float(noise_scale_w),
            length_scale=float(length_scale),
        )[0][0, 0].cpu().float().numpy()
        # peak-normalize to avoid clipping, keep some headroom
        peak = float(np.abs(audio).max())
        if peak > 0:
            audio = audio / peak * 0.95
        audio_i16 = (audio * 32767.0).astype(np.int16)
        return self.sampling_rate, audio_i16, used