# ==== F5-TTS_RUSSIAN — свой минимальный Gradio Space, по образу GigaAM ====
#
# Деплой (как GigaAM):
# 1. huggingface.co -> New Space -> SDK: Gradio -> Hardware: ZeroGPU
#    (бесплатно, но с очередью и лимитом времени на вызов — для одной
#    реплики этого достаточно; если не подойдёт, платный T4 small).
# 2. Положить рядом с этим файлом:
#    - requirements.txt (см. ниже)
#    - ref_voice.wav — 5-15 секунд ЧИСТОЙ речи голосом, которым должен
#      говорить бот (без музыки/шума на фоне, лучше строгая, ровная
#      дикция — F5-TTS клонирует именно тембр и манеру произношения)
#    - переписать REF_TEXT ниже на ТОЧНУЮ дословную расшифровку того, что
#      реально сказано на ref_voice.wav (лишний/неточный текст — главная
#      причина кривых результатов у F5-TTS)
# 3. Закоммитить (через веб-интерфейс Space или git push) — Space сам
#    соберётся и на первый вызов подтянет веса модели с HF (holodный
#    старт может занять пару минут — это нормально).
# 4. В переменных окружения бота на Render задать:
#      TTS_SPACE=твой-юзернейм/имя-space
#
# requirements.txt:
#   f5-tts
#   gradio
#   soundfile
#   spaces

import os
import subprocess
import tempfile

# ZeroGPU требует, чтобы `spaces` импортировался ДО любого CUDA-пакета
# (torch тянется транзитивно через f5_tts) — иначе падает с
# "CUDA has been initialized before importing the `spaces` package".
# На этом Space хардвер всё ещё ZeroGPU (HF ставит spaces сам,
# независимо от requirements.txt), поэтому просто держим декоратор
# рабочим, а не пытаемся его обойти.
try:
    import spaces
    GPU_DECORATOR = spaces.GPU(duration=60)
except Exception:
    def GPU_DECORATOR(func):
        return func

import gradio as gr
import soundfile as sf
from f5_tts.api import F5TTS
from ruaccent import RUAccent

# ---- Ударения и произношение ----
# F5TTS_v1_Base_v2 обучен с полной разметкой ударений (символ "+" перед
# ударной гласной, см. карточку модели) — без неё модель периодически
# путает омографы и часть слов звучит "смазанно". RUAccent расставляет
# ударения автоматически в этом же формате, поэтому просто прогоняем
# через него весь текст перед синтезом.
#
# Отдельная история — заимствования типа "тест", где буква "е" на самом
# деле читается как твёрдое "э" (тэст, а не тест как в слове "тесто").
# Обычные правила русского языка это не знают, RUAccent тоже (это не
# вопрос ударения, а вопрос твёрдости согласного), поэтому модель
# смягчает согласный и на слух "тест" путается с "тесто". Чинится
# точечно через custom_dict — сюда же дописывать новые проблемные
# слова по мере обнаружения (формат: 'слово': 'правильное произношение
# с ударением через +').
_CUSTOM_PRONUNCIATION = {
    "тест": "тэст",
    "теста": "тэста",
    "тесту": "тэсту",
    "тесты": "тэсты",
    "тестов": "тэстов",
    "тестом": "тэстом",
    "тестами": "тэстами",
    "тестах": "тэстах",
    "тестирование": "тэстир+ование",
    "тестировать": "тэстир+овать",
    "тестирую": "тэстир+ую",
    "тестируем": "тэстир+уем",
    "тестируется": "тэстир+уется",
}

# Загружается один раз при старте контейнера — это лёгкая CPU-модель,
# GPU (ZeroGPU) ей не нужен, поэтому не прячем её внутрь @spaces.GPU
# и не тратим на неё выделенное GPU-время.
_accentizer = RUAccent()
_accentizer.load(
    omograph_model_size="turbo3.1",
    use_dictionary=True,
    custom_dict=_CUSTOM_PRONUNCIATION,
)


def _prepare_text(text: str) -> str:
    text = text.strip()
    # Если текст не заканчивается знаком завершения предложения, F5-TTS
    # иногда финиширует раньше времени и обрезает последнее слово/слог —
    # модель ориентируется на пунктуацию, чтобы понять, что фраза
    # закончена. LLM-ответы бота не всегда кончаются точкой (обрывки,
    # смайлики без знака и т.п.), поэтому подстраховываемся тут.
    if text and text[-1] not in ".!?…":
        text += "."
    return _accentizer.process_all(text)

# ---- Модель ----
# F5TTS_v1_Base_v2 — самый свежий из выложенных чекпоинтов Misha24-10 на
# момент написания; если на HF появится более новый (проверить
# huggingface.co/Misha24-10/F5-TTS_RUSSIAN) — можно смело поменять здесь.
REPO = "Misha24-10/F5-TTS_RUSSIAN"
CKPT_SUBPATH = "F5TTS_v1_Base_v2/model_last_inference.safetensors"
VOCAB_SUBPATH = "F5TTS_v1_Base/vocab.txt"

from huggingface_hub import hf_hub_download

# Сами файлы весов можно тянуть на старте — это обычное скачивание,
# GPU для этого не нужен.
ckpt_path = hf_hub_download(REPO, CKPT_SUBPATH)
vocab_path = hf_hub_download(REPO, VOCAB_SUBPATH)

# ---- Референсный голос ----
REF_AUDIO = os.path.join(os.path.dirname(__file__), "ref_voice.wav")
# !!! ОБЯЗАТЕЛЬНО заменить на точный текст, реально сказанный в ref_voice.wav
REF_TEXT = "Обычно я говорю спокойно и ровно, стараясь не торопиться и чётко проговаривать каждое слово. Мне нравится, когда речь звучит естественно, без лишней суеты."

# А вот сам F5TTS(...) грузит чекпоинт СРАЗУ на CUDA-устройство — на
# ZeroGPU видеокарта физически подключена к контейнеру только внутри
# вызова функции, помеченной @spaces.GPU. Если создать модель на
# верхнем уровне модуля (т.е. при старте контейнера), CUDA ещё не
# подключена -> "RuntimeError: No CUDA GPUs are available". Поэтому
# создаём модель лениво, при первом реальном запросе, уже внутри
# GPU-функции.
_tts = None


# Telegram рисует нормальный voice-плеер (waveform, длительность, кнопка
# расшифровки речи) только если реально получает OGG-контейнер с Opus
# внутри. soundfile пишет только WAV (или OGG/Vorbis, что тоже не
# подходит), поэтому финальную перекодировку в OGG/Opus делаем через
# ffmpeg — он уже есть в образе HF Spaces, доп. системные зависимости
# не нужны.
def _to_ogg_opus(wav_path: str) -> str:
    ogg_path = tempfile.mktemp(suffix=".ogg")
    subprocess.run(
        [
            "ffmpeg", "-y",
            "-i", wav_path,
            "-c:a", "libopus",
            "-b:a", "32k",
            "-ar", "48000",
            "-ac", "1",
            ogg_path,
        ],
        check=True,
        capture_output=True,
    )
    return ogg_path


def _get_tts():
    global _tts
    if _tts is None:
        _tts = F5TTS(
            model="F5TTS_v1_Base",
            ckpt_file=ckpt_path,
            vocab_file=vocab_path,
        )
    return _tts


@GPU_DECORATOR
def tts_infer(text: str) -> str:
    if not text or not text.strip():
        raise gr.Error("пустой текст")

    tts = _get_tts()
    gen_text = _prepare_text(text)
    wav, sr, _ = tts.infer(
        ref_file=REF_AUDIO,
        ref_text=REF_TEXT,
        gen_text=gen_text,
        # nfe_step повыше — чище результат, меньше заиканий/кривляний
        # (дефолт 32 маловат для русской модели с ударениями RUAccent).
        nfe_step=64,
        # speed чуть ниже 1.0 — модель закладывает больше времени на
        # символ при оценке длительности, снижает риск обрезки конца
        # фразы (см. _prepare_text про пунктуацию — работает в паре).
        speed=0.93,
    )

    wav_path = tempfile.mktemp(suffix=".wav")
    sf.write(wav_path, wav, sr)
    return _to_ogg_opus(wav_path)


demo = gr.Interface(
    fn=tts_infer,
    inputs=gr.Textbox(label="Текст"),
    outputs=gr.Audio(type="filepath", label="Озвучка"),
    api_name="tts",  # именно этот путь бот дёргает как "/tts" через @gradio/client
    title="F5-TTS_RUSSIAN — озвучка для бота",
)

demo.queue().launch()