import spaces
import gradio as gr
import torch
from huggingface_hub import hf_hub_download 
import os
import sys
import tempfile
from scipy.io.wavfile import write

# Клонуем рэпазіторый, калі ён яшчэ не загружаны
if not os.path.exists("XTTSv2-Finetuning-for-New-Languages"):
    os.system("git clone https://github.com/hellcatmon/XTTSv2-Finetuning-for-New-Languages.git")

# Перамяшчаем тэчку TTS у асноўную дырэкторыю
if os.path.exists("XTTSv2-Finetuning-for-New-Languages/TTS"):
    os.system("mv XTTSv2-Finetuning-for-New-Languages/TTS ./")

# Дадаем тэчку TTS у PYTHONPATH
sys.path.append("./TTS")
from tqdm import tqdm
from underthesea import sent_tokenize
from TTS.tts.configs.xtts_config import XttsConfig
from TTS.tts.models.xtts import Xtts

# Шлях да мадэлі ў Hugging Face
repo_id = "archivartaunik/BE_XTTS_V2_60epoch3Dataset"
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)

@spaces.GPU(duration=60)
def text_to_speech(belarusian_story, lang="be", speaker_audio_file=None):
    device = "cuda:0" if torch.cuda.is_available() else "cpu"
    XTTS_MODEL.to(device)
    # Калі файл не пададзены, выкарыстоўваем голас па змаўчанні
    if not speaker_audio_file or (not isinstance(speaker_audio_file, str) and speaker_audio_file.name == ""):
        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).squeeze().cpu().numpy()

    # Захоўваем аўдыё ў часовы файл
    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
    write(temp_file.name, 24000, out_wav)

    return temp_file.name


demo = gr.Interface(
    fn=text_to_speech,
    inputs=[
        gr.Textbox(lines=5, label="Тэкст на беларускай мове"),
        gr.Textbox(value="be", label="Мова (па змаўчанні BE)", visible=False),
        gr.Audio(type="filepath", label="Запішыце або загрузіце файл голасу (без іншых гукаў) не карацей 7 секунд", interactive=True),
    ],
    outputs="audio",
    title="XTTS Belarusian TTS Demo",
    description="Увядзіце тэкст, і мадэль пераўтворыць яго ў аўдыя. Вы можаце выкарыстоўваць голас па змаўчанні, загрузіць уласны файл або запісаць аўдыё.",
)

if __name__ == "__main__":
    demo.launch()