83 lines
3.2 KiB
Python
83 lines
3.2 KiB
Python
"""Optional vocal isolation (Demucs CLI). Off unless VOCAL_ISOLATION=1; never downloads weights."""
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import shutil
|
|
import subprocess
|
|
import tempfile
|
|
import threading
|
|
import time
|
|
import uuid
|
|
from pathlib import Path
|
|
|
|
from core.audio import _duration, _to_wav_24k
|
|
from core.registry import TEMP_DIR
|
|
|
|
_WEIGHTS_GLOB = "955717e8*.th" # htdemucs
|
|
_busy = threading.Lock() # ponytail: one job at a time; a queue only if several users isolate concurrently
|
|
|
|
|
|
class VocalIsolationError(Exception):
|
|
def __init__(self, message: str, status: int = 500):
|
|
super().__init__(message)
|
|
self.status = status
|
|
|
|
|
|
def _cmd() -> str | None:
|
|
return os.environ.get("VOCAL_ISOLATION_CMD") or shutil.which("demucs")
|
|
|
|
|
|
def _weights_present() -> bool:
|
|
home = Path(os.environ.get("TORCH_HOME") or Path.home() / ".cache" / "torch")
|
|
return any((home / "hub" / "checkpoints").glob(_WEIGHTS_GLOB))
|
|
|
|
|
|
def enabled() -> bool:
|
|
return os.environ.get("VOCAL_ISOLATION", "").strip().lower() in {"1", "true", "yes", "on"}
|
|
|
|
|
|
def available() -> bool:
|
|
return enabled() and _cmd() is not None and _weights_present()
|
|
|
|
|
|
def isolate_vocals(path: Path, cancel: threading.Event | None = None) -> Path:
|
|
"""Return a new 24 kHz mono WAV with vocals only; `path` is left untouched."""
|
|
if not enabled():
|
|
raise VocalIsolationError("Vocal isolation is turned off (set VOCAL_ISOLATION=1).", 503)
|
|
cmd = _cmd()
|
|
if cmd is None or not _weights_present():
|
|
raise VocalIsolationError("Vocal isolation is not installed (Demucs or its weights are missing).", 503)
|
|
max_s = float(os.environ.get("VOCAL_ISOLATION_MAX_SECONDS", "900"))
|
|
timeout = float(os.environ.get("VOCAL_ISOLATION_TIMEOUT", "900"))
|
|
if _duration(path) > max_s:
|
|
raise VocalIsolationError(f"Audio is longer than {int(max_s // 60)} minutes; trim it first.", 413)
|
|
if not _busy.acquire(blocking=False):
|
|
raise VocalIsolationError("Another vocal isolation is running; try again shortly.", 409)
|
|
work = Path(tempfile.mkdtemp(prefix="vocal_", dir=TEMP_DIR))
|
|
try:
|
|
proc = subprocess.Popen(
|
|
[cmd, "--two-stems=vocals", "-n", "htdemucs", "-o", str(work), str(path)],
|
|
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
|
|
)
|
|
deadline = time.monotonic() + timeout
|
|
while proc.poll() is None:
|
|
if cancel is not None and cancel.is_set():
|
|
proc.kill(); proc.wait()
|
|
raise VocalIsolationError("Vocal isolation was cancelled.", 499)
|
|
if time.monotonic() > deadline:
|
|
proc.kill(); proc.wait()
|
|
raise VocalIsolationError("Vocal isolation timed out.", 504)
|
|
time.sleep(0.2)
|
|
found = next(work.glob("*/*/vocals.*"), None)
|
|
if proc.returncode != 0 or found is None:
|
|
raise VocalIsolationError("Vocal isolation failed.", 500)
|
|
staged = TEMP_DIR / f"{uuid.uuid4().hex}_vocals{found.suffix}"
|
|
shutil.move(str(found), staged)
|
|
try:
|
|
return _to_wav_24k(staged)
|
|
finally:
|
|
staged.unlink(missing_ok=True)
|
|
finally:
|
|
_busy.release()
|
|
shutil.rmtree(work, ignore_errors=True)
|