import os
import gradio as gr
from scipy.io.wavfile import write
import tempfile
import shutil

def inference(audio_file):
    """处理上传的音频文件并分离人声和伴奏"""
    # 创建输出目录
    os.makedirs("out", exist_ok=True)
    
    # 使用demucs分离音频
    output_dir = "out"
    os.system(f"python -m demucs.separate -n htdemucs --two-stems=vocals '{audio_file}' -o {output_dir}")
    
    # 获取分离后的文件路径
    base_name = os.path.basename(audio_file)
    name_without_ext = os.path.splitext(base_name)[0]
    
    vocals_path = os.path.join(output_dir, "htdemucs", name_without_ext, "vocals.wav")
    no_vocals_path = os.path.join(output_dir, "htdemucs", name_without_ext, "no_vocals.wav")
    
    return vocals_path, no_vocals_path

# 创建API接口
title = "Suno 音乐分离工具"
description = """
### 使用说明
1. 上传音频文件（支持mp3、wav等格式）
2. 点击"Submit"按钮
3. 等待处理完成后下载分离后的人声和伴奏

### 技术说明
- 本工具使用使用独家AI模型进行音频分离
- 分离质量取决于原始音频的质量和特性
"""

# 创建应用界面
demo = gr.Interface(
    fn=inference,
    inputs=gr.Audio(type="filepath", label="上传音频文件"),
    outputs=[
        gr.Audio(type="filepath", label="人声"),
        gr.Audio(type="filepath", label="伴奏")
    ],
    title=title,
    description=description,
    theme="huggingface",
    examples=[["test.mp3"]]
)

if __name__ == "__main__":
    # 启动服务器，API 在新版 Gradio 中自动启用
    demo.launch(share=True, server_name="0.0.0.0")