97 lines
3.1 KiB
Plaintext

import json
import os
import asyncio
import time
import traceback
from ahserver.filestorage import FileStorage
from urllib.parse import quote
try:
DEMUCS_WRAPPER = '/tmp/demucs_wrapper.py'
PYTHON_BIN = '/share/vllm-0.8.5/bin/python'
OUTPUT_BASE_DIR = '/tmp/demucs_output'
os.makedirs(OUTPUT_BASE_DIR, exist_ok=True)
fs = FileStorage()
filesroot = os.path.abspath(fs.root) # e.g. /tmp
input_file = None
web_path = None
# 方式1: 文件上传 (multipart/form-data)
if 'audio_file' in params_kw:
web_path = params_kw['audio_file']
if isinstance(web_path, str) and len(web_path) > 0:
input_file = fs.realPath(web_path)
# 方式2: 文件路径
elif 'filepath' in params_kw or 'audio_path' in params_kw:
input_file = params_kw.get('filepath') or params_kw.get('audio_path')
if not input_file:
return json.dumps({
'status': 'error',
'message': 'Missing parameter: audio_file (upload) or filepath (path)'
}, ensure_ascii=False)
if not os.path.isfile(input_file):
return json.dumps({
'status': 'error',
'message': f'File not found: {input_file}',
'web_path': web_path
}, ensure_ascii=False)
run_id = f'demucs_{int(time.time() * 1000)}'
output_dir = os.path.join(OUTPUT_BASE_DIR, run_id)
os.makedirs(output_dir, exist_ok=True)
cmd = [PYTHON_BIN, DEMUCS_WRAPPER, '--two-stems=vocals', '-o', output_dir, input_file]
proc = await asyncio.create_subprocess_exec(
*cmd,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE
)
stdout, stderr = await proc.communicate()
if proc.returncode != 0:
error_msg = stderr.decode('utf-8', errors='replace').strip()
return json.dumps({
'status': 'error',
'message': f'Demucs failed: {error_msg}'
}, ensure_ascii=False)
filename_no_ext = os.path.splitext(os.path.basename(input_file))[0]
stems_dir = os.path.join(output_dir, 'htdemucs', filename_no_ext)
vocals_path = os.path.join(stems_dir, 'vocals.wav')
no_vocals_path = os.path.join(stems_dir, 'no_vocals.wav')
if not os.path.isfile(vocals_path):
return json.dumps({
'status': 'error',
'message': f'Output not found: {vocals_path}'
}, ensure_ascii=False)
# 计算相对于 filesroot 的 web path
vocals_webpath = fs.webpath(vocals_path)
no_vocals_webpath = fs.webpath(no_vocals_path)
# 构建下载 URL (通过 idfile 端点)
vocals_url = f'/idfile?path={quote(vocals_webpath)}'
no_vocals_url = f'/idfile?path={quote(no_vocals_webpath)}'
return json.dumps({
'status': 'success',
'vocals_url': vocals_url,
'no_vocals_url': no_vocals_url,
'input_file': input_file
}, ensure_ascii=False)
except Exception as e:
return json.dumps({
'status': 'error',
'message': str(e),
'traceback': traceback.format_exc()
}, ensure_ascii=False)