import gradio as gr
import torch
import soundfile as sf  # 使用官方文档推荐的 soundfile 库
from voxcpm import VoxCPM     # 使用官方文档指定的正确导入方式

# 检查是否有可用的 GPU，否则使用 CPU
device = "cuda" if torch.cuda.is_available() else "cpu"

# --- 模型加载 ---
# 根据官方文档，我们只需要这两行
print("正在加载模型...")
model = VoxCPM.from_pretrained("openbmb/VoxCPM-0.5B")
print("模型加载完成！")

def generate_speech(text):
    """
    使用 VoxCPM 模型生成语音的函数。
    """
    if not text or text.strip() == "":
        print("警告：输入的文本为空。")
        return None

    print(f"开始为文本生成语音: '{text[:30]}...'")
    
    # --- 语音生成 ---
    # 使用官方文档中的 generate 函数
    wav = model.generate(
        text=text,
        prompt_wav_path=None,
        prompt_text=None,
        cfg_value=2.0,
        inference_timesteps=10,
        normalize=True,
        denoise=True
    )
    
    # --- 保存音频文件 ---
    # 使用 soundfile.write，并指定官方推荐的 16000 采样率
    output_filename = "output.wav"
    sampling_rate = 16000
    sf.write(output_filename, wav, sampling_rate)
    
    print(f"语音生成成功，已保存为 {output_filename}")
    return output_filename

# --- Gradio 界面 ---
iface = gr.Interface(
    fn=generate_speech,
    inputs=gr.Textbox(lines=5, label="输入文本 (Input Text)", placeholder="在这里输入你想要转换为语音的中文或英文文本..."),
    outputs=gr.Audio(label="生成的语音 (Generated Speech)"),
    title="🎙️ VoxCPM-0.5B 文本到语音转换",
    description="这是一个严格按照官方文档构建的 VoxCPM 模型演示。输入文本后点击 'Submit' 即可生成语音。",
    examples=[
        ["VoxCPM is an innovative end-to-end TTS model designed to generate highly expressive speech."],
        ["今天天气真不错，我们一起去散步吧！"],
        ["通过在连续空间中对语音进行建模，它克服了离散标记化的局限性。"]
    ],
    allow_flagging="never"
)

# 启动 Gradio 应用
iface.launch()