import gradio as gr
from infer_onnx import TTS

natasha = TTS("TeraTTS/natasha-vits")
glados = TTS("TeraTTS/glados-vits")


def text_to_speech(model_name, text):
    
    if model_name == "TeraTTS/natasha-vits":
        audio = natasha(text)
        natasha.save_wav(audio, 'temp.wav')

    if model_name == "TeraTTS/glados-vits":
        audio = glados(text)
        glados.save_wav(audio, 'temp.wav')
    
    return 'temp.wav'
    
input_text = gr.Textbox(label="Введите текст для синтеза речи")
model_choice = gr.Dropdown(choices=["TeraTTS/natasha-vits", "TeraTTS/glados-vits"], value="TeraTTS/natasha-vits", label="Выберите модель")
output_audio = gr.Audio(label="Аудио", type="numpy")

iface = gr.Interface(fn=text_to_speech, inputs=[model_choice, input_text], outputs=output_audio,
                    description="## [Ссылка](https://github.com/Tera2Space/RUTTS/) на наш гитхаб с кодом для локального запуска")
iface.launch()