import os
os.system("pip install --upgrade pip")
import gradio as gr
import torch
from tqdm import tqdm
from underthesea import sent_tokenize
from TTS.tts.configs.xtts_config import XttsConfig
from TTS.tts.models.xtts import Xtts
from huggingface_hub import hf_hub_download

# Вызначэнне прылады (выкарыстоўваецца GPU, калі даступна)
device = "cuda:0" if torch.cuda.is_available() else "cpu"

# Шлях да мадэлі ў Hugging Face
repo_id = "archivartaunik/GPT_XTTS_V2_40EP"  # Імя рэпазіторыя Hugging Face
checkpoint_file = hf_hub_download(repo_id, filename="model.pth")  # Загружаем файл праверачнай кропкі
config_file = hf_hub_download(repo_id, filename="config.json")    # Загружаем канфігурацыю
vocab_file = hf_hub_download(repo_id, filename="vocab.json")      # Загружаем слоўнік
default_voice_file = hf_hub_download(repo_id, filename="voice.wav")  # Файл голасу па змаўчанні

# Загрузка канфігурацыі мадэлі
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)
XTTS_MODEL.to(device)

# Функцыя для пераўтварэння тэксту ў гукавы фармат
def text_to_speech(belarusian_story, lang="be", speaker_audio_file=default_voice_file):
    # Атрыманне латэнтных умоў і эмацый
    gpt_cond_latent, speaker_embedding = XTTS_MODEL.get_conditioning_latents(
        audio_path=speaker_audio_file,
        gpt_cond_len=XTTS_MODEL.config.gpt_cond_len,
        max_ref_length=XTTS_MODEL.config.max_ref_len,
        sound_norm_refs=XTTS_MODEL.config.sound_norm_refs,
    )

    # Токенізацыя тэксту на асобныя сказы
    tts_texts = sent_tokenize(belarusian_story)

    # Генерацыя аўдыё для кожнага сказы
    wav_chunks = []
    for text in tqdm(tts_texts):
        wav_chunk = XTTS_MODEL.inference(
            text=text,
            language=lang,
            gpt_cond_latent=gpt_cond_latent,
            speaker_embedding=speaker_embedding,
            temperature=0.1,
            length_penalty=1.0,
            repetition_penalty=10.0,
            top_k=10,
            top_p=0.3,
        )
        wav_chunks.append(torch.tensor(wav_chunk["wav"]))

    # Аб'ядноўваем усе часткі аўдыё ў адзін
    out_wav = torch.cat(wav_chunks, dim=0).unsqueeze(0).cpu()

    return out_wav.numpy()

# Інтэрфейс Gradio
demo = gr.Interface(
    fn=text_to_speech,            # Функцыя пераўтварэння
    inputs=[
        gr.Textbox(lines=5, label="Тэкст на беларускай мове"),
        gr.Textbox(value="be", label="Мова (па змаўчанні BE)", visible=False),  # Па змаўчанні BE
        gr.File(label="Файл голасу (па змаўчанні voice.wav)", file_types=[".wav"], optional=True),
    ],
    outputs="audio",              # Выхадныя дадзеныя (аўдыя)
    title="XTTS TTS Demo",        # Назва праекта
    description="Увядзіце тэкст, і мадэль пераўтворыць яго ў аўдыя. Вы можаце выкарыстоўваць голас па змаўчанні або загрузіць уласны.",
)

# Запуск прыкладання
if __name__ == "__main__":
    demo.launch()