from TTS.utils.synthesizer import Synthesizer
from huggingface_hub import hf_hub_download
import gradio as gr
import tempfile

REPO_ID = "jhlfrfufyfn/bel-tts"

my_title = "Беларускі тэкст-у-маўленне"
my_description = "Беларускамоўны мадэль для агучвання тэкста. "

be_text = "Гепарды жывуць у адкрытых і прасторных месцах, дзе ёсць шмат здабычы."

my_inputs = [
  gr.inputs.Textbox(lines=5, label="Input Text", default=be_text),
]

my_outputs = gr.outputs.Audio(type="file", label="Output Audio")

def tts(text: str):
    best_model_path = hf_hub_download(repo_id=REPO_ID, filename="model.pth") 
    config_path = hf_hub_download(repo_id=REPO_ID, filename="config.json")
    vocoder_path = hf_hub_download(repo_id=REPO_ID, filename="vocoder.pth")
    scale_stats_path = hf_hub_download(repo_id=REPO_ID, filename="scale_stats.npy")
    vocoder_config_path = hf_hub_download(repo_id=REPO_ID, filename="vocoder_config.json")
    
    # init synthesizer  
    synthesizer = Synthesizer(
        best_model_path,
        config_path,
        None,
        None,
        vocoder_path,
        vocoder_config_path,
        None,
        None,
        False
    )

    # create audio file
    wavs = synthesizer.tts(text)
    with tempfile.NamedTemporaryFile(suffix = ".wav", delete = False) as fp:
        synthesizer.save_wav(wavs, fp)                      
    return fp.name 


iface = gr.Interface(
    fn=tts, 
    inputs=my_inputs, 
    outputs=my_outputs, 
    title=my_title, 
    description = my_description, 
    article = "",
    examples = "", 
    allow_flagging=False
)
iface.launch()