import gradio as gr
from transformers import VitsModel, AutoTokenizer
import torch
import soundfile as sf
import os

# Tải mô hình và tokenizer
model_id = "facebook/mms-tts-vie"
model = VitsModel.from_pretrained(model_id)
tokenizer = AutoTokenizer.from_pretrained(model_id)

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

    # Tạo thư mục lưu nếu chưa có
    os.makedirs("outputs", exist_ok=True)
    output_path = "outputs/output.wav"
    sf.write(output_path, audio, samplerate=model.config.sampling_rate, format='WAV', subtype='PCM_16')
    return output_path

# Giao diện Gradio
gr.Interface(
    fn=tts,
    inputs=gr.Textbox(label="Nhập văn bản tiếng Việt"),
    outputs=gr.Audio(label="Kết quả âm thanh"),
    title="Vietnamese TTS - Hugging Face",
    description="Ứng dụng chuyển văn bản tiếng Việt thành giọng nói bằng mô hình MMS-TTS"
).launch()