import gradio as gr
import librosa
import numpy as np
import torch

from transformers import SpeechT5Processor, SpeechT5ForTextToSpeech, SpeechT5HifiGan

from datasets import load_dataset


checkpoint = "microsoft/speecht5_tts"
processor = SpeechT5Processor.from_pretrained(checkpoint)
model = SpeechT5ForTextToSpeech.from_pretrained(checkpoint)
vocoder = SpeechT5HifiGan.from_pretrained("microsoft/speecht5_hifigan")

default_voice = "CLB (female)"

embeddings_dataset = load_dataset("Matthijs/cmu-arctic-xvectors", split="validation")
speaker_embeddings = torch.tensor(embeddings_dataset[7306]["xvector"]).unsqueeze(0)

speaker_embedding = {
    "BDL": "spkemb/cmu_us_bdl_arctic-wav-arctic_a0009.npy",
    "CLB": "spkemb/cmu_us_clb_arctic-wav-arctic_a0144.npy",
    "KSP": "spkemb/cmu_us_ksp_arctic-wav-arctic_b0087.npy",
    "RMS": "spkemb/cmu_us_rms_arctic-wav-arctic_b0353.npy",
    "SLT": "spkemb/cmu_us_slt_arctic-wav-arctic_a0508.npy",
}


def predict(text):
    if len(text.strip()) == 0:
        return (16000, np.zeros(0).astype(np.int16))

    inputs = processor(text=text, return_tensors="pt")

    # limit input length
    input_ids = inputs["input_ids"]
    input_ids = input_ids[..., :model.config.max_text_positions]

    
    

    speech = model.generate_speech(input_ids, speaker_embeddings, vocoder=vocoder)

    speech = (speech.numpy() * 32767).astype(np.int16)
    return (16000, speech)


title = "Prosody Project"

description = """
This is the Prosody Project for DT2112 Speech Technology
"""


# examples = [
#     ["Hi, my name is Santiago", "CLB (female)"],
#     ["Two bros, chilling in a hot tub, five feet apart cause they are not gay.", "CLB (female)"]
# ]

examples = [
    ["Hi, my name is Santiago"],
    ["I am becoming a vampire, so I would like no garlic, please."]
]

gr.Interface(
    fn=predict,
    inputs=[
        gr.Text(label="Input Text"),
        #gr.Radio(label="Speaker", choices=[
        #    "CLB (female)"
        #],
        # value="CLB (female)"),
    ],
    outputs=[
        gr.Audio(label="Generated Speech", type="numpy"),
    ],
    title=title,
    description=description,
    article=None,
    examples=examples,
).launch()