import gradio as gr
import tempfile

from huggingface_hub import hf_hub_download
from torch import no_grad, package
import ctypes
import gc

from accentor import accentification, stress_replace_and_shift, accentors


config = {
    "mykyta": "theodotus/tts-vits-mykyta-uk",
    "mykyta-lite": "theodotus/tts-vits-mykyta-low-uk",
    "olena": "theodotus/tts-vits-olena-uk",
    "lada": "theodotus/tts-vits-lada-uk",
    "dmytro": "theodotus/tts-vits-dmytro-uk",
    "harakternyk": "theodotus/tts-vits-harakternyk-uk",
}

voices = list(config.keys())

tts_kwargs = {
    "speaker_name": "uk",
    "language_name": "uk",
}


def trim_memory():
    libc = ctypes.CDLL("libc.so.6")
    libc.malloc_trim(0)
    gc.collect()


def init_models():
    models = {}
    for name, model_name in config.items():
        model_path = hf_hub_download(model_name, "model.pt")
        importer = package.PackageImporter(model_path)
        synt = importer.load_pickle("tts_models", "model")
        models[name] = synt
    return models


def tts(text: str, voice: str, mode: str):
    # accentor
    accented_text = accentification(text, mode)
    if (mode != "none"):
        plussed_text = stress_replace_and_shift(accented_text)
    else:
        plussed_text = accented_text
    # TTS
    synt = models[voice]
    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as fp:
        with no_grad():
            wav_data = synt.tts(plussed_text, **tts_kwargs)
            synt.save_wav(wav_data, fp)
            trim_memory()
        return fp.name, accented_text



models = init_models()



iface = gr.Interface(
    fn=tts,
    inputs=[
        gr.Textbox(
            label="Input",
            value="Кам'янець-Подільський - місто в Хмельницькій області України, центр Кам'янець-Подільської міської об'єднаної територіальної громади і Кам'янець-Подільського району.",
        ),
        gr.Radio(
            label="Voice",
            choices=voices,
            value=voices[0],
        ),
        gr.Radio(
            label="Accentor",
            choices=accentors,
            value=accentors[0],
        ),
    ],
    outputs=[
        gr.Audio(label="Output"),
        gr.Textbox(label="Stressed")
    ],
    title="🇺🇦 - Ukrainian Voices",
    article="[Harakternyk](https://huggingface.co/theodotus/tts-vits-harakternyk-uk) model is licensed under [CC BY-NC 4.0](https://creativecommons.org/licenses/by-nc/4.0/)",
)

iface.launch()