import gradio as gr
import torch
from transformers import AutoProcessor, VitsModel
from scipy.io.wavfile import write

model_id = "facebook/mms-tts-nko"
processor = AutoProcessor.from_pretrained(model_id)
model = VitsModel.from_pretrained(model_id)
model.to("cpu")

# Bảng chuyển N'Ko sang Latin đơn giản (mở rộng dần)
nko_to_latin = {
    "ߞ": "k", "ߐ": "a", "ߣ": "n", "ߌ": "i",
    "ߔ": "b", "ߘ": "d"
}

def transliterate_nko(text):
    return ''.join(nko_to_latin.get(ch, ch) for ch in text)

def tts_nko(text):
    try:
        latin_text = transliterate_nko(text)
        inputs = processor(text=latin_text, return_tensors="pt")
        inputs["input_ids"] = inputs["input_ids"].long()
        with torch.no_grad():
            output = model(**inputs)
        audio = output.waveform[0].numpy()
        path = "nko.wav"
        write(path, rate=model.config.sampling_rate, data=audio)
        return f"✅ Phát âm: {latin_text}", path
    except Exception as e:
        return f"❌ Lỗi: {str(e)}", None

with gr.Blocks() as demo:
    gr.Markdown("# 🔊 N'Ko TTS (ߞߐߣߌ ➝ /kani/)")
    text_input = gr.Textbox(label="Nhập văn bản bằng chữ N'Ko (VD: ߞߐߣߌ)")
    output_text = gr.Textbox(label="Ghi chú")
    output_audio = gr.Audio(label="Phát âm", type="filepath")
    button = gr.Button("Đọc")

    button.click(fn=tts_nko, inputs=[text_input], outputs=[output_text, output_audio])

demo.launch()