import gradio as gr
from transformers import AutoTokenizer, VitsModel
import torch

# Load model and tokenizer
model = VitsModel.from_pretrained("facebook/mms-tts-che")
tokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-che")

def tts_chechen(text):
    inputs = tokenizer(text, return_tensors="pt")
    with torch.no_grad():
        output = model(**inputs).waveform
    return (model.config.sampling_rate, output.squeeze().numpy())

# Gradio UI
iface = gr.Interface(
    fn=tts_chechen,
    inputs=gr.Textbox(label="Nhập tiếng Chechen"),
    outputs=gr.Audio(label="Giọng đọc Chechen"),
    title="TTS Chechen - MMS by Meta",
    description="Chuyển văn bản tiếng Chechen thành giọng nói"
)

iface.launch()