Xiaohan-Tian
121 lines
4.2 KiB
Python
121 lines
4.2 KiB
Python
"""Stem separator runner — invokes audio-separator CLI per-request.
|
|
|
|
Uses a single-worker ThreadPoolExecutor to serialize GPU use.
|
|
Coordinates with model_manager: callers should check model_manager.active_model
|
|
before submitting, and model load routes should check separator_runner.active.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import subprocess
|
|
import threading
|
|
import time
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
ROOT_DIR = Path(__file__).parent.parent
|
|
SEPARATOR_DIR = ROOT_DIR / "separator"
|
|
OUTPUT_DIR = ROOT_DIR / "outputs" / "separator"
|
|
UPLOAD_DIR = ROOT_DIR / "uploads" / "separator"
|
|
|
|
ALLOWED_MODELS = frozenset([
|
|
"UVR-MDX-NET-Inst_HQ_3.onnx",
|
|
"MDX23C-8KFFT-InstVoc_HQ.ckpt",
|
|
"htdemucs_6s.yaml",
|
|
])
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class SeparatorRunner:
|
|
def __init__(self) -> None:
|
|
self._tasks: dict[str, dict] = {}
|
|
self._executor = ThreadPoolExecutor(max_workers=1)
|
|
self._running_count = 0
|
|
self._lock = threading.Lock()
|
|
|
|
@property
|
|
def active(self) -> bool:
|
|
with self._lock:
|
|
return self._running_count > 0
|
|
|
|
def submit(self, task_id: str, file_path: Path, model_filename: str) -> None:
|
|
"""Schedule a separation task. Returns immediately; poll get_task() for status."""
|
|
self._tasks[task_id] = {"status": "pending", "created_at": time.time()}
|
|
with self._lock:
|
|
self._running_count += 1
|
|
self._executor.submit(self._run, task_id, file_path, model_filename)
|
|
|
|
def get_task(self, task_id: str) -> Optional[dict]:
|
|
return self._tasks.get(task_id)
|
|
|
|
def _run(self, task_id: str, file_path: Path, model_filename: str) -> None:
|
|
self._tasks[task_id]["status"] = "running"
|
|
try:
|
|
cmd = [
|
|
"uv", "run", "audio-separator",
|
|
str(file_path),
|
|
"--model_filename", model_filename,
|
|
"--output_dir", str(OUTPUT_DIR),
|
|
"--output_format", "MP3",
|
|
]
|
|
logger.info("Running separator: %s", " ".join(cmd))
|
|
proc = subprocess.Popen(
|
|
cmd,
|
|
cwd=str(SEPARATOR_DIR),
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.STDOUT, # merge stderr into stdout
|
|
text=True,
|
|
bufsize=1,
|
|
)
|
|
output_lines: list[str] = []
|
|
|
|
def _read_output() -> None:
|
|
for line in proc.stdout:
|
|
stripped = line.rstrip()
|
|
logger.info("[separator] %s", stripped)
|
|
output_lines.append(stripped)
|
|
|
|
reader = threading.Thread(target=_read_output, daemon=True)
|
|
reader.start()
|
|
try:
|
|
proc.wait(timeout=600)
|
|
except subprocess.TimeoutExpired:
|
|
proc.kill()
|
|
proc.wait()
|
|
raise RuntimeError("audio-separator timed out after 600s")
|
|
finally:
|
|
reader.join(timeout=5)
|
|
|
|
if proc.returncode != 0:
|
|
detail = "\n".join(output_lines[-20:]) or "unknown error"
|
|
raise RuntimeError(f"audio-separator exited {proc.returncode}: {detail}")
|
|
|
|
# Collect output files — named {input_stem}_{stem_type}_{model_base}.mp3
|
|
stem = file_path.stem
|
|
files = sorted(p.name for p in OUTPUT_DIR.glob(f"{stem}_*"))
|
|
if not files:
|
|
raise RuntimeError("audio-separator completed but produced no output files")
|
|
|
|
self._tasks[task_id].update({"status": "complete", "files": files})
|
|
logger.info("Separation complete for task %s: %s", task_id, files)
|
|
except Exception:
|
|
logger.exception("Separation failed for task %s", task_id)
|
|
self._tasks[task_id].update({
|
|
"status": "error",
|
|
"error": "Separation failed — check server logs.",
|
|
})
|
|
finally:
|
|
with self._lock:
|
|
self._running_count -= 1
|
|
# Clean up upload file after processing
|
|
try:
|
|
file_path.unlink(missing_ok=True)
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
separator_runner = SeparatorRunner()
|