import gradio as gr
import torch
import torchaudio
from transformers import AutoModelForCausalLM
import tempfile
import os
import numpy as np
import soundfile as sf

# 模型ID（使用官方模型，或你自己复制的）
MODEL_ID = "OpenMOSS/MOSS-TTS-Nano"

def tts_inference(text, prompt_audio_path="assets/audio/zh_1.wav"):
    # 加载模型
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        trust_remote_code=True,
        torch_dtype=torch.float32
    ).eval()
    
    # 加载 prompt 音频
    wav_prompt, sr = torchaudio.load(prompt_audio_path)
    
    # 生成语音
    with torch.no_grad():
        output = model.generate(
            text=text,
            prompt_wav=wav_prompt,
            sampling_rate=sr,
            max_new_tokens=2048,  # 支持长文本
        )
    
    # 保存生成的音频
    wav_out = output["wav"].cpu().numpy().squeeze()
    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp:
        sf.write(tmp.name, wav_out, sr)
        return tmp.name

# 创建 Gradio 界面
with gr.Blocks(title="MOSS-TTS-Nano 长文本语音合成") as demo:
    gr.Markdown("## 🎙️ MOSS-TTS-Nano 文本转语音")
    gr.Markdown("输入文本（最多2000字），选择或上传参考音频（可选），点击合成")
    
    with gr.Row():
        with gr.Column():
            text_input = gr.Textbox(
                label="输入文本",
                placeholder="在这里输入要合成的文本，最长支持2000字...",
                lines=10,
                max_length=2000
            )
            prompt_audio = gr.Audio(
                label="参考音频（可选，用于音色克隆）",
                type="filepath",
                value="assets/audio/zh_1.wav"
            )
            btn = gr.Button("合成语音", variant="primary")
        with gr.Column():
            output_audio = gr.Audio(label="合成结果", type="filepath")
    
    btn.click(
        fn=tts_inference,
        inputs=[text_input, prompt_audio],
        outputs=output_audio
    )
    
    gr.Markdown("### 说明\n- 本模型基于 MOSS-TTS-Nano，针对 CPU 优化\n- 支持最长 2000 个字符的文本\n- 参考音频可选，如果不提供则使用默认音色")

demo.launch()