import torch
from transformers import SpeechT5Processor, SpeechT5ForTextToSpeech
import gradio as gr
import soundfile as sf
import io

# Load processor và model MMS TTS cho tiếng Tây Tạng (bod)
processor = SpeechT5Processor.from_pretrained("facebook/mms-tts-bod")
model = SpeechT5ForTextToSpeech.from_pretrained("facebook/mms-tts-bod")

device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)

def tts(text):
    # Xử lý text input
    inputs = processor(text=text, return_tensors="pt")
    input_ids = inputs.input_ids.to(device)

    # Sinh âm thanh
    with torch.no_grad():
        speech = model.generate_speech(input_ids, speaker_embeddings=None, vocoder=None)

    # Lưu ra buffer WAV để phát
    buf = io.BytesIO()
    sf.write(buf, speech.cpu().numpy(), samplerate=16000, format="WAV")
    buf.seek(0)
    return buf

# Tạo giao diện Gradio
iface = gr.Interface(fn=tts, inputs="text", outputs="audio", title="TTS tiếng Tây Tạng - Facebook MMS")

if __name__ == "__main__":
    iface.launch()