demucs-service/workers/separate.py

88 lines
2.9 KiB
Python

import os
import json
import asyncio
import time
from appPublic.log import debug, error
async def run_separate(task_obj, payload):
"""
Run demucs vocal/accompaniment separation.
payload:
audio_path (str, required): Path to input audio file
output_dir (str, optional): Output directory, default /tmp/demucs_{task_id}
"""
audio_path = payload.get('audio_path')
if not audio_path:
raise ValueError('audio_path is required')
if not os.path.isfile(audio_path):
raise FileNotFoundError(f'Audio file not found: {audio_path}')
task_id = payload.get('task_id', str(int(time.time())))
output_dir = payload.get('output_dir', f'/tmp/demucs_{task_id}')
gpu_id = task_obj.gpu_id
basename = os.path.splitext(os.path.basename(audio_path))[0]
# Expected output paths from demucs
result_dir = os.path.join(output_dir, 'htdemucs', basename)
vocals_path = os.path.join(result_dir, 'vocals.wav')
no_vocals_path = os.path.join(result_dir, 'no_vocals.wav')
# Build the command
env = os.environ.copy()
env['CUDA_VISIBLE_DEVICES'] = str(gpu_id)
cmd = [
'/data/ymq/demucs_venv/bin/python', '-m', 'demucs',
'--two-stems', 'vocals',
audio_path,
'-o', output_dir
]
debug(f'[demucs] Running separation: audio={audio_path}, output={output_dir}, gpu={gpu_id}')
debug(f'[demucs] Command: {" ".join(cmd)}')
start_time = time.time()
proc = await asyncio.create_subprocess_exec(
*cmd,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
env=env
)
stdout, stderr = await proc.communicate()
elapsed = time.time() - start_time
if proc.returncode != 0:
stderr_text = stderr.decode('utf-8', errors='replace')
stdout_text = stdout.decode('utf-8', errors='replace')
error(f'[demucs] Process failed (rc={proc.returncode})')
error(f'[demucs] stdout: {stdout_text[-2000:]}')
error(f'[demucs] stderr: {stderr_text[-2000:]}')
raise RuntimeError(
f'Demucs separation failed (rc={proc.returncode}): {stderr_text[-500:]}'
)
# Verify output files exist
if not os.path.isfile(vocals_path):
raise FileNotFoundError(f'Expected vocals output not found: {vocals_path}')
if not os.path.isfile(no_vocals_path):
raise FileNotFoundError(f'Expected no_vocals output not found: {no_vocals_path}')
vocals_size = os.path.getsize(vocals_path)
no_vocals_size = os.path.getsize(no_vocals_path)
debug(f'[demucs] Separation complete in {elapsed:.1f}s')
debug(f'[demucs] vocals.wav: {vocals_size} bytes, no_vocals.wav: {no_vocals_size} bytes')
return {
'vocals_path': vocals_path,
'no_vocals_path': no_vocals_path,
'duration': round(elapsed, 2),
'output_dir': output_dir,
'model': 'htdemucs'
}