app.py

# app.py
import os
import sys
import subprocess
import logging

# каб лагі ішлі адразу ў stdout
os.environ.setdefault("PYTHONUNBUFFERED", "1")

import spaces
import gradio as gr
import torch
import numpy as np
from huggingface_hub import hf_hub_download

# -----------------------------
# Лагаванне толькі ў stdout
# -----------------------------
def _setup_logging():
    handler = logging.StreamHandler(sys.stdout)
    handler.setFormatter(logging.Formatter("%(asctime)s %(levelname)s %(message)s"))
    logging.basicConfig(level=logging.INFO, handlers=[handler], force=True)
    logging.captureWarnings(True)
    logger = logging.getLogger("be-xtts-app")
    logger.propagate = False
    logger.setLevel(logging.INFO)
    return logger

logger = _setup_logging()

# ---------------------------------------------------------
# 1) Клануем fork coqui-ai-TTS з падтрымкай беларускай
# ---------------------------------------------------------
REPO_URL = "https://github.com/tuteishygpt/coqui-ai-TTS.git"
REPO_DIR = "coqui-ai-TTS"

if not os.path.exists(REPO_DIR):
    subprocess.run(["git", "clone", REPO_URL, REPO_DIR], check=True)

repo_root = os.path.abspath(REPO_DIR)
if repo_root not in sys.path:
    sys.path.insert(0, repo_root)

# ---------------------------------------------------------
# 2) Імпарты з TTS
# ---------------------------------------------------------
from TTS.tts.configs.xtts_config import XttsConfig
from TTS.tts.models.xtts import Xtts
from TTS.tts.layers.xtts.tokenizer import (
    split_sentence,
    VoiceBpeTokenizer,
)

# ---------------------------------------------------------
# 3) Шляхі да файлаў мадэлі
# ---------------------------------------------------------
repo_id = "archivartaunik/BE_XTTS_V2_10ep250k"
model_dir = "./model"
os.makedirs(model_dir, exist_ok=True)

checkpoint_file = os.path.join(model_dir, "model.pth")
config_file = os.path.join(model_dir, "config.json")
vocab_file = os.path.join(model_dir, "vocab.json")
default_voice_file = os.path.join(model_dir, "voice.wav")
speakers_file = os.path.join(model_dir, "speakers_xtts.pth")

if not os.path.exists(checkpoint_file):
    hf_hub_download(repo_id, filename="model.pth", local_dir=model_dir)
if not os.path.exists(config_file):
    hf_hub_download(repo_id, filename="config.json", local_dir=model_dir)
if not os.path.exists(vocab_file):
    hf_hub_download(repo_id, filename="vocab.json", local_dir=model_dir)
if not os.path.exists(default_voice_file):
    hf_hub_download(repo_id, filename="voice.wav", local_dir=model_dir)
# новае: падцягваем speakers_xtts.pth
if not os.path.exists(speakers_file):
    try:
        hf_hub_download(repo_id, filename="speakers_xtts.pth", local_dir=model_dir)
    except Exception as e:
        logger.warning("Не атрымалася спампаваць speakers_xtts.pth: %s", e)

# ---------------------------------------------------------
# 4) Загрузка мадэлі і токенайзера
# ---------------------------------------------------------
config = XttsConfig()
config.load_json(config_file)

XTTS_MODEL = Xtts.init_from_config(config)
XTTS_MODEL.load_checkpoint(
    config,
    checkpoint_path=checkpoint_file,
    vocab_path=vocab_file,
    use_deepspeed=False,
)

device = "cuda:0" if torch.cuda.is_available() else "cpu"
XTTS_MODEL.to(device)
sampling_rate = int(XTTS_MODEL.config.audio["sample_rate"])

# Ініцыялізацыя VoiceBpeTokenizer і падкладанне ў мадэль
tokenizer = VoiceBpeTokenizer(vocab_file=vocab_file)
XTTS_MODEL.tokenizer = tokenizer

# Базавыя значэнні для кандыцыянавання
CFG_GPT_COND = int(getattr(XTTS_MODEL.config, "gpt_cond_len", 6))
CFG_MAX_REF  = int(getattr(XTTS_MODEL.config, "max_ref_len", 20))
CFG_NORM     = bool(getattr(XTTS_MODEL.config, "sound_norm_refs", True))

# ---------------------------------------------------------
# 4.1) Загрузка speakers_xtts.pth
# ---------------------------------------------------------
SPEAKERS_DB: dict[str, dict] = {}
SPEAKER_CHOICES: list[str] = ["— з аўдыё (reference) —"]

if os.path.exists(speakers_file):
    try:
        raw = torch.load(speakers_file, map_location="cpu")
        # магчымыя фарматы:
        # 1) {"speakers": {name: {...}}}
        # 2) {name: {...}}
        if isinstance(raw, dict) and "speakers" in raw and isinstance(raw["speakers"], dict):
            speakers_dict = raw["speakers"]
        else:
            speakers_dict = raw

        valid_count = 0
        if isinstance(speakers_dict, dict):
            for name, val in speakers_dict.items():
                if (
                    isinstance(val, dict)
                    and "gpt_cond_latent" in val
                    and "speaker_embedding" in val
                ):
                    SPEAKERS_DB[str(name)] = {
                        "gpt_cond_latent": val["gpt_cond_latent"],
                        "speaker_embedding": val["speaker_embedding"],
                    }
                    valid_count += 1

        if valid_count > 0:
            S_NAMES = sorted(SPEAKERS_DB.keys())
            SPEAKER_CHOICES.extend(S_NAMES)
            logger.info("Загружана %d галасоў з speakers_xtts.pth", valid_count)
        else:
            logger.warning(
                "speakers_xtts.pth загружаны, але не знойдзена ніводнага "
                "галасу з ключамі 'gpt_cond_latent' і 'speaker_embedding'."
            )
    except Exception as e:
        logger.exception("Памылка пры загрузцы speakers_xtts.pth: %s", e)
else:
    logger.warning("speakers_xtts.pth не знойдзены па шляху: %s", speakers_file)

# ---------------------------------------------------------
# Утыліты
# ---------------------------------------------------------
def clip_for_log(s: str, limit: int = 600):
    s = (s or "").replace("\n", " ").strip()
    return s if len(s) <= limit else s[:limit] + " ... [clipped]"

def log_after_chunk(idx: int, text: str, ui_logs: list):
    line = f"[TEXT] chunk {idx}: AFTER  :: {clip_for_log(text)}"
    logger.info(line)
    print(line, flush=True)  # дублюем у stdout, гарантуем бачнасць
    ui_logs.append(line)

# ---------------------------------------------------------
# 5) Функцыя TTS (лагі толькі AFTER)
#    + падтрымка speakers_xtts.pth
#    + асобнае аўдыё "Прыклад прасодыі" толькі для gpt_cond_latent
# ---------------------------------------------------------
@spaces.GPU(duration=60)
def text_to_speech(
    belarusian_story: str,
    speaker_audio_file: str | None,
    prosody_audio_file: str | None,   # НОВЫ ПАРАМЕТР
    preset_speaker: str = "— з аўдыё (reference) —",
    language: str = "be",
    preprocess_text_flag: bool = True,
    gpt_cond_len: int = CFG_GPT_COND,
    max_ref_len: int = CFG_MAX_REF,
    sound_norm_refs: bool = CFG_NORM,
    temperature: float = 0.2,
    length_penalty: float = 1.0,
    repetition_penalty: float = 7.0,
    top_k: int = 30,
    top_p: float = 0.8,
):
    """
    Вяртае: (sr, waveform), LOG_TEXT

    Лагі ўключаюць ТОЛЬКІ радкі '[TEXT] chunk N: AFTER  :: ...'

    Параметр preset_speaker:
      - калі значэнне з SPEAKERS_DB — бярэцца голас з speakers_xtts.pth
      - калі '— з аўдыё (reference) —' — выкарыстоўваецца speaker_audio_file / voice.wav

    Дадаткова:
      - калі prosody_audio_file загружаны, ён выкарыстоўваецца ТОЛЬКІ для gpt_cond_latent
      - speaker_embedding вызначаецца так, як рэалізавана раней:
        * альбо з SPEAKERS_DB
        * альбо з reference-аудыё (speaker_audio_file / voice.wav)
    """
    if not belarusian_story or belarusian_story.strip() == "":
        raise gr.Error("Увядзі хоць нейкі тэкст 🙂")

    lang_short = (language or "be").split("-")[0]
    chunk_limit = tokenizer.char_limits.get(lang_short, 250)

    # 1) падзел на чанкі
    try:
        tts_texts = split_sentence(
            belarusian_story.strip(),
            lang=lang_short,
            text_split_length=chunk_limit,
        )
        tts_texts = [s.strip() for s in tts_texts if s and s.strip()]
        if not tts_texts:
            raise gr.Error("Не атрымалася падзяліць тэкст на сказы/чанкі.")
    except Exception as e:
        logger.exception("Памылка пры падзеле тэксту")
        raise gr.Error(f"Памылка пры падзеле тэксту: {e}")

    # 2) поўная апрацоўка (лагі толькі AFTER)
    ui_logs = []
    if preprocess_text_flag:
        processed = []
        for idx, s in enumerate(tts_texts, start=1):
            tokenizer.check_input_length(s, lang_short)  # можа вывесці WARN у stdout
            s_proc = tokenizer.preprocess_text(s, lang_short)
            log_after_chunk(idx, s_proc, ui_logs)        # ЛОГ ТОЛЬКІ AFTER
            processed.append(s_proc)
        tts_texts = processed

    # 3) атрыманне латэнтаў голасу:
    #    - speaker_embedding: "як зараз рэалізавана"
    #    - gpt_cond_latent: альбо з таго ж месца, альбо з prosody_audio_file
    use_preset = (
        isinstance(preset_speaker, str)
        and preset_speaker in SPEAKERS_DB
    )

    gpt_cond_latent = None
    speaker_embedding = None

    # 3a) speaker_embedding (і, калі няма prosody_audio_file, gpt_cond_latent)
    if use_preset:
        # галасавы прэсэт з speakers_xtts.pth
        try:
            sp = SPEAKERS_DB[preset_speaker]
            speaker_embedding = sp["speaker_embedding"].to(device)
            if not prosody_audio_file:
                # калі асобны прыклад прасодыі НЕ зададзены,
                # то gpt_cond_latent таксама бярэм з прэсэта (як раней)
                gpt_cond_latent = sp["gpt_cond_latent"].to(device)
        except Exception as e:
            logger.exception(
                "Памылка пры выкарыстанні галасу '%s' з speakers_xtts.pth", preset_speaker
            )
            raise gr.Error(
                f"Памылка пры выкарыстанні падрыхтаванага галасу '{preset_speaker}': {e}"
            )
    else:
        # reference-аудыё (Прыклад голасу / voice.wav)
        ref_path = speaker_audio_file
        if not ref_path or (
            not isinstance(ref_path, str)
            and getattr(ref_path, "name", "") == ""
        ):
            ref_path = default_voice_file

        try:
            ref_gpt_cond_latent, ref_speaker_embedding = XTTS_MODEL.get_conditioning_latents(
                audio_path=ref_path,
                gpt_cond_len=int(gpt_cond_len),
                max_ref_length=int(max_ref_len),
                sound_norm_refs=bool(sound_norm_refs),
            )
            speaker_embedding = ref_speaker_embedding.to(device)
            if not prosody_audio_file:
                # калі асобны прыклад прасодыі НЕ зададзены,
                # то gpt_cond_latent таксама бярэм з reference-аудыё (як раней)
                gpt_cond_latent = ref_gpt_cond_latent.to(device)
        except Exception as e:
            logger.exception("Памылка пры атрыманні латэнтаў голасу з reference-аудыё")
            raise gr.Error(f"Памылка пры атрыманні латэнтаў голасу: {e}")

    # 3b) Калі загружаны "Прыклад прасодыі" — выкарыстоўваем яго ТОЛЬКІ для gpt_cond_latent
    if prosody_audio_file:
        prosody_path = prosody_audio_file
        if not isinstance(prosody_path, str) and getattr(prosody_path, "name", ""):
            prosody_path = prosody_path.name

        try:
            prosody_gpt_cond_latent, _ = XTTS_MODEL.get_conditioning_latents(
                audio_path=prosody_path,
                gpt_cond_len=int(gpt_cond_len),
                max_ref_length=int(max_ref_len),
                sound_norm_refs=bool(sound_norm_refs),
            )
            gpt_cond_latent = prosody_gpt_cond_latent.to(device)
        except Exception as e:
            logger.exception("Памылка пры атрыманні прасодыі з 'Прыклад прасодыі'")
            raise gr.Error(f"Памылка пры атрыманні прасодыі з аўдыё: {e}")

    # праверка, што абодва латэнты ёсць
    if gpt_cond_latent is None or speaker_embedding is None:
        raise gr.Error(
            "Не атрымалася вызначыць gpt_cond_latent або speaker_embedding. "
            "Праверце налады галасу і файлы аўдыё."
        )

    # 4) генерацыя
    all_wavs = []
    for text in tts_texts:
        try:
            with torch.no_grad():
                wav_chunk = XTTS_MODEL.inference(
                    text=text,
                    language=lang_short,
                    gpt_cond_latent=gpt_cond_latent,
                    speaker_embedding=speaker_embedding,
                    temperature=float(temperature),
                    length_penalty=float(length_penalty),
                    repetition_penalty=float(repetition_penalty),
                    top_k=int(top_k),
                    top_p=float(top_p),
                )
            all_wavs.append(wav_chunk["wav"])
        except Exception as e:
            logger.exception("Памылка пры генерырацыі аўдыя")
            raise gr.Error(f"Памылка пры генерырацыі аўдыя: {e}")

    if not all_wavs:
        raise gr.Error("Нічога не згенеравалася — праверце ўваходныя даныя.")

    try:
        out_wav = np.concatenate(all_wavs).astype(np.float32)
    except Exception as e:
        logger.exception("Памылка пры аб'яднанні аўдыя")
        raise gr.Error(f"Памылка пры аб'яднанні аўдыя: {e}")

    return (sampling_rate, out_wav), "\n".join(ui_logs)

# ---------------------------------------------------------
# 5.1) Прэв'ю выбранага галасу (кароткая фраза)
# ---------------------------------------------------------
@spaces.GPU(duration=30)
def preview_speaker(
    preset_speaker: str = "— з аўдыё (reference) —",
    language: str = "be",
):
    """
    Генеруе кароткі прыклад для выбранага прэсэта/рефэрэнса.
    Вяртае толькі (sr, waveform) без лагаў.
    """
    lang_short = (language or "be").split("-")[0]

    sample_texts = {
        "be": "Гэта прыклад беларускага голасу.",
        "ru": "Это пример голоса.",
        "uk": "Це приклад голосу.",
        "pl": "To jest przykładowy głos.",
        "en": "This is a sample voice.",
        "de": "Dies ist eine Beispielstimme.",
        "fr": "Ceci est une voix d'exemple.",
        "es": "Esta es una voz de ejemplo.",
    }
    sample_text = sample_texts.get(lang_short, "This is a sample voice.")

    # Выкарыстоўваем тыя ж налады, што і ў асноўнай функцыі
    (audio, _logs) = text_to_speech(
        belarusian_story=sample_text,
        speaker_audio_file=None,         # для прэсэта не трэба
        prosody_audio_file=None,         # асобны прыклад прасодыі для прэв'ю не патрабуецца
        preset_speaker=preset_speaker,
        language=language,
        preprocess_text_flag=True,
        gpt_cond_len=CFG_GPT_COND,
        max_ref_len=CFG_MAX_REF,
        sound_norm_refs=CFG_NORM,
        temperature=0.2,
        length_penalty=1.0,
        repetition_penalty=7.0,
        top_k=30,
        top_p=0.8,
    )
    return audio

# ---------------------------------------------------------
# 6) UI (Gradio Blocks)
# ---------------------------------------------------------
with gr.Blocks() as demo:
    gr.Markdown("# Belarusian TTS Demo (XTTSv2 + tokenizer.py) — лагі толькі AFTER")

    # НОВЫ РАДОК: тэкст + прыклад голасу + прыклад прасодыі
    with gr.Row():
        txt = gr.Textbox(lines=8, label="Тэкст")
        ref = gr.Audio(type="filepath", label="Прыклад голасу (≥7 с)")
        prosody = gr.Audio(
            type="filepath",
            label="Прыклад прасодыі (≥7 с, неабавязкова)",
        )

    # выбар галасу з speakers_xtts.pth
    speaker_dropdown = gr.Dropdown(
        label="Падрыхтаваныя галасы (speakers_xtts.pth)",
        choices=SPEAKER_CHOICES,
        value=SPEAKER_CHOICES[0],
    )

    # Аўдыё для прэв'ю галасу
    preview_audio = gr.Audio(type="numpy", label="Прэв'ю выбранага галасу")

    # Кнопка для праслухоўвання выбранага прэсэта
    preview_btn = gr.Button("▶️ Прайграць выбраны голас")

    with gr.Row():
        language = gr.Dropdown(
            label="Мова (language)",
            choices=[
                "be","ru","uk","pl","cs","en","de","fr","es",
                "it","pt","tr","vi","zh","ja","ko","nl","hu","ar","hi"
            ],
            value="be",
        )
        preprocess_text_flag = gr.Checkbox(
            value=True,
            label="Апрацоўваць тэкст праз tokenizer.py"
        )

    with gr.Accordion("Параметры кандыцыянавання (для reference-аудыё)", open=False):
        with gr.Row():
            gpt_cond_len = gr.Slider(
                1, max(1, CFG_GPT_COND*3),
                step=1,
                value=CFG_GPT_COND,
                label="gpt_cond_len (сек.)"
            )
            max_ref_len = gr.Slider(
                1, max(1, CFG_MAX_REF*3),
                step=1,
                value=CFG_MAX_REF,
                label="max_ref_len (сек.)"
            )
            sound_norm_refs = gr.Checkbox(
                value=CFG_NORM,
                label="sound_norm_refs"
            )

    with gr.Accordion("Параметры генерацыі", open=True):
        with gr.Row():
            temperature = gr.Slider(
                0.0, 2.0,
                value=0.2,
                step=0.01,
                label="temperature"
            )
            top_k = gr.Slider(
                1, 100,
                value=30,
                step=1,
                label="top_k"
            )
        with gr.Row():
            top_p = gr.Slider(
                0.0, 1.0,
                value=0.8,
                step=0.01,
                label="top_p"
            )
            length_penalty = gr.Slider(
                0.5, 3.5,
                value=1.0,
                step=0.05,
                label="length_penalty"
            )
            repetition_penalty = gr.Slider(
                0.5, 20.0,
                value=7.0,
                step=0.1,
                label="repetition_penalty"
            )

    out_audio = gr.Audio(type="numpy", label="Згенераванае аўдыя")
    out_logs  = gr.Textbox(lines=16, label="Лагі (толькі AFTER)")

    btn = gr.Button("🔊 Генераваць")

    # асноўная генерацыя
    btn.click(
        fn=text_to_speech,
        inputs=[
            txt,
            ref,
            prosody,          # НОВЫ INPUT
            speaker_dropdown,
            language,
            preprocess_text_flag,
            gpt_cond_len,
            max_ref_len,
            sound_norm_refs,
            temperature,
            length_penalty,
            repetition_penalty,
            top_k,
            top_p,
        ],
        outputs=[out_audio, out_logs],
    )

    # прэв'ю выбранага галасу
    preview_btn.click(
        fn=preview_speaker,
        inputs=[speaker_dropdown, language],
        outputs=preview_audio,
    )

# ---------------------------------------------------------
# 7) Запуск
# ---------------------------------------------------------
if __name__ == "__main__":
    demo.launch()