import torch
import gradio as gr
from transformers import AutoProcessor, BarkModel
import soundfile as sf
import traceback

from transformers import AutoProcessor, AutoModelForSpeechSeq2Seq

# Tải model và processor
processor = AutoProcessor.from_pretrained("facebook/mms-tts-bod")
model = AutoModelForSpeechSeq2Seq.from_pretrained("facebook/mms-tts-bod")

# Kiểm tra thiết bị
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)

def tts(text):
    try:
        print(f"Đầu vào: {text}")
        inputs = processor(text=text, return_tensors="pt").to(device)
        with torch.no_grad():
            output = model.generate(**inputs)
        audio_arr = processor.batch_decode(output, return_tensors="pt")[0].numpy()
        sf.write("output.wav", audio_arr, 16000)
        return "output.wav"
    except Exception as e:
        print("LỖI:\n", traceback.format_exc())
        return f"Lỗi: {str(e)}"

gr.Interface(
    fn=tts,
    inputs=gr.Textbox(label="Nhập văn bản tiếng Tây Tạng"),
    outputs=gr.Audio(label="Kết quả TTS"),
    title="TTS Tiếng Tây Tạng (MMS)"
).launch()