import gradio as gr
from fairseq.checkpoint_utils import load_model_ensemble_and_task_from_hf_hub
from fairseq.models.text_to_speech.hub_interface import TTSHubInterface
import numpy as np

def load_tts_model():
    models, cfg, task = load_model_ensemble_and_task_from_hf_hub(
        "facebook/tts_transformer-zh-cv7_css10",  # Considere usar um modelo para inglês
        arg_overrides={"vocoder": "hifigan", "fp16": False}
    )
    model = models[0]
    TTSHubInterface.update_cfg_with_data_cfg(cfg, task.data_cfg)
    generator = task.build_generator(model, cfg)
    return task, model, generator

task, model, generator = load_tts_model()

def synthesize_text(text):
    sample = TTSHubInterface.get_model_input(task, text)
    wav, rate = TTSHubInterface.get_prediction(task, model, generator, sample)
    return np.array(wav), rate

# Exemplos pré-carregados em inglês
examples = [
    ["Hello, how are you today?"],
    ["What's the weather like?"],
    ["Learning new languages is fun."],
    # Adicione mais exemplos aqui
]

iface = gr.Interface(
    fn=synthesize_text,
    inputs=gr.inputs.Textbox(lines=2, placeholder="Enter English text here..."),
    outputs=gr.outputs.Audio(label="Synthesized Speech"),
    title="Text to Speech Synthesis",
    description="A simple text-to-speech app. Note: The model is trained for Chinese, results may vary for English.",
    examples=examples,
    theme="huggingface"
)