import gradio as gr
import spaces
import torch
from huggingface_hub import snapshot_download

from frysk.fy_num_engine import num_filter
from omnivoice import OmniVoice

gr.close_all()

title = "Prototype Frisian TTS"
description = """
OmniVoice model, fine-tuned on:<br>- A part of the dataset from Project FRIS (a collaboration between Provinsje Fryslân, Taalstudio Prins, and Juvoly)<br>- Audiobooks read by Geartsje de Vries of books written by Koos Tiemersma and published Audiofrysk<br>
"""

css = """
h1 {
    text-align: center;
    display:block;
}
"""

snapshot_download(repo_id="phatdo/Frysk-TTS-test-model", local_dir='./checkpoints')

def load_omnivoice(path):
    model = OmniVoice.from_pretrained(
        path,
        device_map="cpu",
        dtype=torch.float16
    )
    return model

juvoly_model = load_omnivoice("checkpoints/juvoly_speaker102/checkpoint-5000")
geartsje_model = load_omnivoice("checkpoints/geartsje/checkpoint-5000")

@spaces.GPU
def synthesize(model, text):
    model.to('cuda:0')
    text = num_filter(text)
    audio = model.generate(text=text,
                       language="fy")
    return (24000, audio[0])

@spaces.GPU
def infer_juvoly(text):
    output = synthesize(juvoly_model, text)
    return output

@spaces.GPU
def infer_geartsje(text):
    output = synthesize(geartsje_model, text)
    return output

iface_juvoly = gr.Interface(fn=infer_juvoly,
                     inputs=[
                            gr.components.Textbox(lines=3, placeholder="Please input Frisian text to synthesize", label='Text to synthesize'),
                            ],
                     outputs=gr.components.Audio(type="numpy", label='Synthesized speech'),
                     title=title,
                     description=description,
                     css=css,
                     theme='huggingface',
                     )

iface_geartsje = gr.Interface(fn=infer_geartsje,
                     inputs=[
                            gr.components.Textbox(lines=3, placeholder="Please input Frisian text to synthesize", label='Text to synthesize'),
                            ],
                     outputs=gr.components.Audio(type="numpy", label='Synthesized speech'),
                     title=title,
                     description=description,
                     css=css,
                     theme='huggingface',
                     )

demo = gr.TabbedInterface([iface_juvoly, iface_geartsje], ["Juvoly voice (speaker 102)", "Geartsje de Vries' voice "])

demo.launch(share=False, inline=False, server_name="0.0.0.0", server_port=7860)