import torch
import gradio as gr
import soundfile as sf
from transformers import AutoTokenizer, VitsForTextToSpeech

# Tải mô hình và tokenizer
model_name = "facebook/mms-tts-bod"
device = "cuda" if torch.cuda.is_available() else "cpu"

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = VitsForTextToSpeech.from_pretrained(model_name).to(device)

# Hàm chuyển văn bản thành giọng nói
def tts_bod(text):
    inputs = tokenizer(text, return_tensors="pt").to(device)
    with torch.no_grad():
        output = model(**inputs)
        audio = output.waveform.squeeze().cpu().numpy()

    # Lưu tạm âm thanh ra file
    sf.write("output.wav", audio, 16000)
    return "output.wav"

# Giao diện Gradio
interface = gr.Interface(
    fn=tts_bod,
    inputs=gr.Textbox(label="Nhập văn bản tiếng Tây Tạng"),
    outputs=gr.Audio(label="Âm thanh đầu ra", type="filepath"),
    title="TTS Tiếng Tây Tạng (facebook/mms-tts-bod)"
)

interface.launch()