import torch
import torchaudio
from transformers import AutoProcessor, SpeechT5ForTextToSpeech, SpeechT5HifiGan
import gradio as gr
import soundfile as sf

# Tải model và processor
print("Đang tải mô hình TTS Tây Tạng từ Facebook MMS...")
processor = AutoProcessor.from_pretrained("facebook/mms-tts-bod")
model = SpeechT5ForTextToSpeech.from_pretrained("facebook/mms-tts-bod")
vocoder = SpeechT5HifiGan.from_pretrained("microsoft/speecht5_hifigan")

# Tải speaker embedding mẫu
speaker_embedding = torch.randn(1, 512)  # MMS không có mẫu sẵn nên dùng random embedding

# Hàm xử lý TTS
def tts(text):
    inputs = processor(text=text, return_tensors="pt")
    with torch.no_grad():
        speech = model.generate_speech(
            inputs["input_ids"],
            speaker_embedding=speaker_embedding,
            vocoder=vocoder
        )
    sf.write("output.wav", speech.numpy(), samplerate=16000)
    return "output.wav"

# Giao diện Gradio
interface = gr.Interface(
    fn=tts,
    inputs=gr.Textbox(label="Nhập văn bản tiếng Tây Tạng"),
    outputs=gr.Audio(label="Giọng đọc TTS (Tây Tạng)"),
    title="TTS Tiếng Tây Tạng (Meta MMS)"
)

interface.launch()