import os

# PyTorch Dynamo နဲ့ Asyncio Warning တွေကို System level ကနေ ပိတ်ခြင်း
os.environ["TORCHDYNAMO_DISABLE"] = "1"

import asyncio
import warnings

warnings.filterwarnings("ignore")

import gradio as gr
import numpy as np
import spaces
import torch
from voxcpm import VoxCPM

model = None


def get_model():
    global model
    if model is None:
        print("Loading VoxCPM2 Model...")
        model = VoxCPM.from_pretrained("openbmb/VoxCPM2")
    return model


@spaces.GPU
def generate_speech(target_text, reference_audio):
    cpm_model = get_model()

    # Reference Audio သို့မဟုတ် Target Text ကို စစ်ဆေးခြင်း
    if reference_audio:
        wav = cpm_model.generate(
            text=target_text,
            prompt_wav_path=reference_audio,
            prompt_text=target_text,
            cfg_value=2.0,
            inference_timesteps=10,
        )
    else:
        wav = cpm_model.generate(
            text=target_text, cfg_value=2.0, inference_timesteps=10
        )

    # wav ထွက်လာမှု စစ်ဆေးခြင်း
    if wav is None:
        raise ValueError("Audio generation failed and returned None.")

    # Audio Data ကို Gradio ဖတ်နိုင်သော numpy float32 / int16 သို့ သေချာပြောင်းပေးခြင်း
    sample_rate = getattr(cpm_model.tts_model, "sample_rate", 24000)

    if torch.is_tensor(wav):
        wav = wav.cpu().numpy()

    wav = np.squeeze(wav)

    return (sample_rate, wav)


# UI Layout
demo = gr.Interface(
    fn=generate_speech,
    inputs=[
        gr.Textbox(
            label="Target Text (မြန်မာစာ)",
            value="မင်္ဂလာပါ၊ ဒီမှာမြန်မာလို စာသားရိုက်ထည်ပြီး အသံထုတ်နိုင်ပါတယ်။",
            lines=3,
        ),
        gr.Audio(
            label="Reference Audio (မြန်မာလို ပြောထားသည့် အသံဖိုင် ၃-၅ စက္ကန့်)",
            type="filepath",
            sources=["upload", "microphone"],
        ),
    ],
    outputs=gr.Audio(label="Generated Audio"),
    title="VoxCPM2 myanmar voice",
)

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