import gradio as gr
from teeteeass.infer import TTS

def generate_tts(text, language, speaker):
    if language == "MS":
        language = "EN"
    tts = TTS(speaker=speaker)
    audio = tts.generate(text, language)
    print(audio.shape)
    return 44100,audio

# Fetch the list of speakers
speakers = TTS.list_speakers()

# Define the Gradio interface
iface = gr.Interface(
    fn=generate_tts,
    inputs=[
        gr.Textbox(lines=5, placeholder="Type your text here..."), 
        gr.Radio(choices=["MS", "ZH", "JP"], label="Language"),
        gr.Dropdown(choices=speakers, label="Speaker")  # Dropdown for speaker selection
    ],
    outputs=[gr.Audio(type="numpy", label="Generated Speech")],
    title="TTS Generation",
    description="Beta version for https://github.com/khursani8/teeteeass\nStill need to add more phonemes, hope can finish before Godam Sahur end.\n\nTrained 50+ speakers for now, will add around 100+ more male and female with Malaysian tongue. First inference it will download the model which might slow at first\n\nAppreciate some feedback and bug report especially pronounciation bug to https://github.com/khursani8/teeteeass/issues/new"
)

# Run the Gradio app
iface.launch()