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

# ១. ទាញយកម៉ូដែល TTS ភាសាខ្មែររបស់ Meta
model_name = "facebook/mms-tts-khm"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = VitsModel.from_pretrained(model_name)

def tts_khmer(text):
    # បំលែងអក្សរខ្មែរទៅជា Inputs សម្រាប់ម៉ូដែល
    inputs = tokenizer(text, return_tensors="pt")
    
    # បង្កើតសំឡេង
    with torch.no_grad():
        output = model(**inputs).waveform
    
    # ផ្លាស់ប្តូរទម្រង់ទៅជា File សំឡេង WAV
    sampling_rate = model.config.sampling_rate
    audio_data = output.numpy().squeeze()
    
    output_path = "khmer_output.wav"
    wavfile.write(output_path, rate=sampling_rate, data=audio_data)
    
    return output_path

# ២. បង្កើតផ្ទាំង WebUI សម្រាប់ភាសាខ្មែរ
interface = gr.Interface(
    fn=tts_khmer,
    inputs=gr.Textbox(
        label="វាយអត្ថបទជាភាសាខ្មែរនៅទីនេះ", 
        value="សួស្តី! ស្វាគមន៍មកកាន់ប្រព័ន្ធបំប្លែងអត្ថបទទៅជាសំឡេងនិយាយ។"
    ),
    outputs=gr.Audio(label="លទ្ធផលសំឡេងខ្មែរ"),
    title="Khmer TTS Demo (Meta MMS)",
    description="កម្មវិធីតេស្តសាកល្បងបំប្លែងអក្សរខ្មែរទៅជាសំឡេង AI ដោយប្រើប្រាស់ Meta MMS-TTS"
)

if __name__ == "__main__":
    interface.launch(server_name="0.0.0.0")