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",
        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

iface = gr.Interface(
    fn=synthesize_text,
    inputs=gr.inputs.Textbox(lines=2, placeholder="Enter Chinese text here..."),
    outputs=gr.outputs.Audio(label="Synthesized Speech"),
    title="Text to Speech Synthesis",
    description="A simple text-to-speech app using Fairseq TTS model."
)

iface.launch()