432 lines
17 KiB
Python
432 lines
17 KiB
Python
#!/usr/bin/env python3
|
|
"""stemd: separates every track in the music library that doesn't have stems yet.
|
|
|
|
Runs on the server next to stemserve. Walks MUSIC_DIR, and for each track whose folder
|
|
under STEMS_DIR has no marker it runs the separation on the GPU and writes the stems with
|
|
the same layout and marker as the desktop app (stemify/library_core.py), so stemserve
|
|
and the phone pick them up as they land. Once caught up it looks again every RESCAN_MIN.
|
|
|
|
"Stem Now" in the app goes through stemserve, which drops a request (a JSON list of Plex
|
|
file paths) into QUEUE_DIR. Those tracks are separated before the rest: between tracks
|
|
during a pass, and within a few seconds while idle.
|
|
|
|
Settings (environment):
|
|
|
|
MUSIC_DIR /music the music library (read-only is fine)
|
|
STEMS_DIR /stems where stems go
|
|
STATE_DIR /state failed-track list
|
|
QUEUE_DIR STATE_DIR/queue "Stem Now" requests from stemserve
|
|
MODELS_DIR /models model weights, downloaded on first use
|
|
MODEL roformer_bs_sw any id from stemify/catalog.py
|
|
DEVICE cuda cuda, cuda:N or cpu
|
|
RESCAN_MIN 30 minutes between scans once caught up
|
|
AUTO_SCAN 0 1 = scan the library on a timer; 0 = only separate Stem Now requests
|
|
FP16 1 half precision on the GPU
|
|
OVERLAP / SEGMENT / SHIFTS as in the desktop app's settings
|
|
"""
|
|
|
|
import glob
|
|
import json
|
|
import os
|
|
import shutil
|
|
import signal
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
import time
|
|
|
|
_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
sys.path.insert(0, os.path.join(_ROOT, "desktop")) # in the image the package sits at /app/stemify
|
|
sys.path.insert(0, _ROOT)
|
|
|
|
from stemify.catalog import MODELS_BY_ID # noqa: E402
|
|
from stemify.library_core import (MARKER, MAX_SECONDS, SOURCE_RANK, duration, is_done, is_share_error, # noqa: E402
|
|
iter_tracks, signature, write_track)
|
|
|
|
MUSIC_DIR = os.environ.get("MUSIC_DIR", "/music")
|
|
STEMS_DIR = os.environ.get("STEMS_DIR", "/stems")
|
|
STATE_DIR = os.environ.get("STATE_DIR", "/state")
|
|
MODELS_DIR = os.environ.get("MODELS_DIR", "/models")
|
|
MODEL = os.environ.get("MODEL", "roformer_bs_sw")
|
|
DEVICE = os.environ.get("DEVICE", "cuda")
|
|
RESCAN_S = float(os.environ.get("RESCAN_MIN", "30")) * 60
|
|
AUTO_SCAN = os.environ.get("AUTO_SCAN", "0") in ("1", "true", "yes")
|
|
FP16 = os.environ.get("FP16", "1") not in ("0", "false", "no", "")
|
|
OVERLAP = float(os.environ.get("OVERLAP", "0.25"))
|
|
SEGMENT = float(os.environ.get("SEGMENT", "0"))
|
|
SHIFTS = int(os.environ.get("SHIFTS", "0"))
|
|
QUEUE_DIR = os.environ.get("QUEUE_DIR", os.path.join(STATE_DIR, "queue"))
|
|
FAILED_PATH = os.path.join(STATE_DIR, "failed.json")
|
|
QUEUE_POLL_S = 5 # how often an idle stemd looks for Stem Now requests
|
|
SHARE_RETRY_S = 5 * 60 # stems folder unwritable: try again this soon rather than at RESCAN_S
|
|
MAX_CRASHES = 2 # GPU crashes on one track before it's given up on
|
|
|
|
stop = threading.Event()
|
|
|
|
|
|
def log(*args):
|
|
print(time.strftime("%Y-%m-%d %H:%M:%S"), *args, flush=True)
|
|
|
|
|
|
class Failed:
|
|
"""Tracks that didn't separate, keyed by path, retried once the source file changes."""
|
|
|
|
def __init__(self):
|
|
try:
|
|
with open(FAILED_PATH) as f:
|
|
self.data = json.load(f)
|
|
except (OSError, ValueError):
|
|
self.data = {}
|
|
|
|
def skip(self, path, sig):
|
|
e = self.data.get(path)
|
|
return bool(e) and e["sig"] == sig and e.get("tries", MAX_CRASHES) >= MAX_CRASHES
|
|
|
|
def add(self, path, sig, msg, crash=False):
|
|
tries = MAX_CRASHES
|
|
if crash:
|
|
e = self.data.get(path)
|
|
tries = (e.get("tries", 0) if e and e["sig"] == sig else 0) + 1
|
|
self.data[path] = {"sig": sig, "msg": msg, "tries": tries, "at": time.time()}
|
|
os.makedirs(STATE_DIR, exist_ok=True)
|
|
tmp = FAILED_PATH + ".tmp"
|
|
with open(tmp, "w") as f:
|
|
json.dump(self.data, f, indent=1)
|
|
os.replace(tmp, FAILED_PATH)
|
|
|
|
|
|
class _NoProgress:
|
|
"""Stands in for the worker's socket; per-track progress isn't worth a log line."""
|
|
|
|
def send(self, *_, **__):
|
|
pass
|
|
|
|
|
|
class Separator:
|
|
def __init__(self):
|
|
import torch
|
|
from stemify.worker.backends import DemucsBackend, RoformerBackend
|
|
|
|
self.torch = torch
|
|
self.model = MODELS_BY_ID[MODEL]
|
|
self.device = torch.device(DEVICE if DEVICE == "cpu" or torch.cuda.is_available() else "cpu")
|
|
if DEVICE != "cpu" and self.device.type == "cpu":
|
|
log("WARNING: no CUDA device visible; separating on the CPU (slow). Is the container "
|
|
"running with the nvidia runtime?")
|
|
self.backend = (DemucsBackend() if self.model["family"] == "demucs" else RoformerBackend(MODELS_DIR))
|
|
name = torch.cuda.get_device_name(self.device) if self.device.type == "cuda" else "CPU"
|
|
log(f"model {self.model['label']} on {name}, torch {torch.__version__}")
|
|
|
|
def run(self, path, tmp):
|
|
# Imported late: worker/__main__ sets ROCm env defaults, harmless here, and owns Job's
|
|
# parameter rules (overlap, fp16, ...), which must match the desktop's.
|
|
from stemify.worker.__main__ import Job
|
|
job = Job({"id": "stemd", "file": path, "model": MODEL, "out_dir": tmp, "segment": SEGMENT,
|
|
"shifts": SHIFTS, "overlap": OVERLAP, "fp16": FP16}, _NoProgress())
|
|
if "file" in self.model and not os.path.exists(os.path.join(MODELS_DIR, self.model["file"])):
|
|
log(" downloading the model (first run only)")
|
|
job.cancel_event = stop
|
|
stems, sr = self.backend.run(job, self.device)
|
|
# Backends return (channels, samples); the writer wants (samples, channels).
|
|
return {name: audio.T for name, audio in stems.items()}, sr
|
|
|
|
def release(self):
|
|
self.backend.release()
|
|
|
|
def is_fatal(self, e):
|
|
"""CUDA errors other than OOM poison the context; only a fresh process recovers."""
|
|
if isinstance(e, self.torch.OutOfMemoryError):
|
|
return False
|
|
return isinstance(e, getattr(self.torch, "AcceleratorError", ())) or "CUDA error" in str(e)
|
|
|
|
|
|
def separate(sep, failed, path, dest, sig, label):
|
|
"""Separate one track into dest. Returns True if it was written.
|
|
|
|
Raises on a stop request, an unwritable stems share, or a GPU error that needs a restart.
|
|
"""
|
|
rel = os.path.relpath(path, MUSIC_DIR)
|
|
dur = duration(path)
|
|
if dur is not None and dur > MAX_SECONDS:
|
|
log(f"{label} skip {rel}: {dur / 60:.0f} min long")
|
|
failed.add(path, sig, f"longer than {MAX_SECONDS // 60} minutes")
|
|
return False
|
|
log(f"{label} {rel}" + (f" ({int(dur) // 60}:{int(dur) % 60:02d})" if dur else ""))
|
|
t1 = time.monotonic()
|
|
tmp = tempfile.mkdtemp(prefix="stemd-")
|
|
try:
|
|
stems, sr = sep.run(path, tmp)
|
|
meta = {"source": path, "size": sig[0], "mtime_ns": sig[1], "model": MODEL, "samplerate": sr,
|
|
"duration": len(next(iter(stems.values()))) / sr}
|
|
write_track(dest, stems, meta)
|
|
except BaseException as e:
|
|
if stop.is_set():
|
|
raise # Cancelled; the track is redone next start
|
|
if is_share_error(e):
|
|
log(f" can't write to {STEMS_DIR}: {e}; pausing")
|
|
raise
|
|
if not isinstance(e, Exception):
|
|
raise
|
|
msg = str(e).strip().splitlines()[0] if str(e).strip() else e.__class__.__name__
|
|
if isinstance(e, sep.torch.OutOfMemoryError):
|
|
# Something else (Plex, a game server) may be holding VRAM; try again next pass.
|
|
log(" out of GPU memory; will retry next pass")
|
|
sep.release()
|
|
return False
|
|
if sep.is_fatal(e):
|
|
failed.add(path, sig, msg, crash=True)
|
|
log(f" GPU error: {msg}; restarting")
|
|
sys.exit(75)
|
|
failed.add(path, sig, msg)
|
|
log(f" failed: {msg}")
|
|
return False
|
|
finally:
|
|
shutil.rmtree(tmp, ignore_errors=True)
|
|
took = time.monotonic() - t1
|
|
log(f" done in {took:.0f}s" + (f" ({meta['duration'] / took:.1f}x realtime)" if took else ""))
|
|
return True
|
|
|
|
|
|
# The project was renamed to Stemify; songs separated before that carry the old marker name.
|
|
# migrate_markers() renames them once so they aren't separated again. Once stemd has run it
|
|
# against your library, this constant, the function and its call in main() can be deleted.
|
|
LEGACY_MARKER = ".stemroller.json"
|
|
|
|
|
|
def migrate_markers():
|
|
"""Rename every legacy done-marker under STEMS_DIR to MARKER (a few seconds per 10k songs)."""
|
|
n = 0
|
|
for root, dirs, files in os.walk(STEMS_DIR):
|
|
dirs[:] = [d for d in dirs if not d.startswith(".")]
|
|
if LEGACY_MARKER not in files:
|
|
continue
|
|
old, new = os.path.join(root, LEGACY_MARKER), os.path.join(root, MARKER)
|
|
try:
|
|
if MARKER in files:
|
|
os.remove(old)
|
|
else:
|
|
os.replace(old, new)
|
|
n += 1
|
|
except OSError as e:
|
|
log(f"can't rename {old}: {e}")
|
|
if n:
|
|
log(f"renamed {n:,} {LEGACY_MARKER} markers to {MARKER}")
|
|
|
|
|
|
def resolve(plex_path):
|
|
"""The music file a Plex path names, or None. Plex sees the library under its own mount
|
|
point, so this takes the longest tail of the path (folder and file at least) that exists
|
|
under MUSIC_DIR, and prefers the best format when a folder holds the song in several."""
|
|
parts = [p for p in plex_path.replace("\\", "/").split("/") if p not in ("", ".", "..")]
|
|
root = os.path.realpath(MUSIC_DIR)
|
|
for i in range(len(parts) - 1):
|
|
path = os.path.join(root, *parts[i:])
|
|
if os.path.isfile(path) and os.path.realpath(path).startswith(root + os.sep):
|
|
break
|
|
else:
|
|
return None
|
|
folder, name = os.path.split(path)
|
|
base = os.path.splitext(name)[0].casefold()
|
|
try:
|
|
same = [f for f in os.listdir(folder) if os.path.splitext(f)[0].casefold() == base
|
|
and os.path.splitext(f)[1].lower() in SOURCE_RANK]
|
|
except OSError:
|
|
same = []
|
|
if same:
|
|
path = os.path.join(folder, min(same, key=lambda f: (SOURCE_RANK[os.path.splitext(f)[1].lower()], f != name)))
|
|
return path if os.path.splitext(path)[1].lower() in SOURCE_RANK else None
|
|
|
|
|
|
class Status:
|
|
"""What stemd is up to, for stemserve's /api/queue: the track in progress (the Plex path for
|
|
a Stem Now, else the path under MUSIC_DIR), how many tracks were finished since start, and the
|
|
last Stem Now paths that couldn't be done. Kept in QUEUE_DIR/.status.json, which the "*.json"
|
|
request glob skips because it starts with a dot."""
|
|
|
|
def __init__(self):
|
|
self.path = os.path.join(QUEUE_DIR, ".status.json")
|
|
self.current, self.done, self.failed = None, 0, []
|
|
self.write()
|
|
|
|
def start(self, current):
|
|
self.current = current
|
|
self.write()
|
|
|
|
def finish(self, ok, failed_path=None):
|
|
self.current = None
|
|
if ok:
|
|
self.done += 1
|
|
elif failed_path:
|
|
self.failed = [p for p in self.failed if p != failed_path][-49:] + [failed_path]
|
|
self.write()
|
|
|
|
def write(self):
|
|
tmp = self.path + ".tmp"
|
|
try:
|
|
with open(tmp, "w") as f:
|
|
json.dump({"current": self.current, "done": self.done, "failed": self.failed,
|
|
"at": time.time()}, f)
|
|
os.replace(tmp, self.path)
|
|
except OSError:
|
|
pass
|
|
|
|
|
|
status = None # set in main()
|
|
|
|
|
|
def serve_queue(sep, failed):
|
|
"""Separate every Stem Now request, oldest first. Returns whether any track was attempted.
|
|
|
|
A request file is deleted once all its tracks are done, so one cut short by a restart
|
|
picks up where it left off (finished tracks are skipped)."""
|
|
used = False
|
|
while not stop.is_set():
|
|
reqs = sorted(glob.glob(os.path.join(QUEUE_DIR, "*.json")))
|
|
if not reqs:
|
|
break
|
|
req = reqs[0]
|
|
try:
|
|
with open(req) as f:
|
|
paths = json.load(f).get("paths", [])
|
|
except (OSError, ValueError, AttributeError) as e:
|
|
log(f"stem now: unreadable request {os.path.basename(req)}: {e}")
|
|
paths = []
|
|
todo = []
|
|
for p in paths if isinstance(paths, list) else []:
|
|
src = resolve(p) if isinstance(p, str) else None
|
|
if src is None:
|
|
log(f"stem now: not in {MUSIC_DIR}: {p}")
|
|
if isinstance(p, str):
|
|
status.finish(False, p)
|
|
continue
|
|
dest = os.path.join(STEMS_DIR, os.path.splitext(os.path.relpath(src, MUSIC_DIR))[0])
|
|
if is_done(dest):
|
|
continue
|
|
try:
|
|
sig = signature(src)
|
|
except OSError:
|
|
continue
|
|
if failed.skip(src, sig):
|
|
log(f"stem now: skip {os.path.relpath(src, MUSIC_DIR)}: failed before ({failed.data[src]['msg']})")
|
|
status.finish(False, p)
|
|
continue
|
|
todo.append((p, src, dest, sig))
|
|
if todo:
|
|
log(f"stem now: {len(todo)} track{'' if len(todo) == 1 else 's'} ahead of the rest")
|
|
for i, (p, src, dest, sig) in enumerate(todo, 1):
|
|
if stop.is_set():
|
|
return used
|
|
if is_done(dest):
|
|
continue
|
|
used = True
|
|
status.start(p)
|
|
ok = separate(sep, failed, src, dest, sig, f"[now {i}/{len(todo)}]")
|
|
status.finish(ok, None if ok else p)
|
|
if stop.is_set():
|
|
break
|
|
try:
|
|
os.remove(req)
|
|
except FileNotFoundError:
|
|
pass
|
|
return used
|
|
|
|
|
|
def run_pass(sep, failed):
|
|
"""Separate everything that needs it. Returns (tracks done, whether the GPU was used)."""
|
|
t0 = time.monotonic()
|
|
todo, total = [], 0
|
|
for path, dest in iter_tracks(MUSIC_DIR, STEMS_DIR):
|
|
if stop.is_set():
|
|
return 0, False
|
|
total += 1
|
|
if is_done(dest):
|
|
continue
|
|
try:
|
|
sig = signature(path)
|
|
except OSError:
|
|
continue
|
|
if not failed.skip(path, sig):
|
|
todo.append((path, dest, sig))
|
|
log(f"scan: {total:,} tracks, {total - len(todo):,} done or skipped, {len(todo):,} to do "
|
|
f"({time.monotonic() - t0:.0f}s)")
|
|
if not total:
|
|
log(f"WARNING: no music found in {MUSIC_DIR}; is the share mounted and readable?")
|
|
|
|
done, used = 0, False
|
|
for i, (path, dest, sig) in enumerate(todo, 1):
|
|
used = serve_queue(sep, failed) or used
|
|
if stop.is_set():
|
|
break
|
|
if is_done(dest): # a full pass takes days; the desktop or a Stem Now may have got there first
|
|
continue
|
|
used = True
|
|
status.start(os.path.relpath(path, MUSIC_DIR))
|
|
ok = separate(sep, failed, path, dest, sig, f"[{i}/{len(todo)}]")
|
|
status.finish(ok)
|
|
if ok:
|
|
done += 1
|
|
return done, used
|
|
|
|
|
|
def idle(sep, failed, seconds):
|
|
"""Wait for the next scan, separating Stem Now requests as they arrive."""
|
|
end = time.monotonic() + seconds
|
|
while not stop.is_set() and time.monotonic() < end:
|
|
if serve_queue(sep, failed):
|
|
sep.release()
|
|
stop.wait(min(QUEUE_POLL_S, max(0.0, end - time.monotonic())))
|
|
|
|
|
|
def main():
|
|
for name, path in (("MUSIC_DIR", MUSIC_DIR), ("STEMS_DIR", STEMS_DIR)):
|
|
if not os.path.isdir(path):
|
|
sys.exit(f"{name} not found: {path}")
|
|
if MODEL not in MODELS_BY_ID:
|
|
sys.exit(f"unknown MODEL {MODEL!r}; pick one of: {', '.join(MODELS_BY_ID)}")
|
|
# docker stop sends SIGTERM; stop between (or inside) tracks instead of dying mid-write.
|
|
signal.signal(signal.SIGTERM, lambda *_: stop.set())
|
|
signal.signal(signal.SIGINT, lambda *_: stop.set())
|
|
try:
|
|
os.nice(10)
|
|
except OSError:
|
|
pass
|
|
try:
|
|
os.makedirs(QUEUE_DIR, exist_ok=True)
|
|
except OSError as e:
|
|
log(f"WARNING: can't create {QUEUE_DIR} ({e}); Stem Now won't work")
|
|
global status
|
|
status = Status()
|
|
log(f"stemd: {MUSIC_DIR} -> {STEMS_DIR}")
|
|
migrate_markers()
|
|
sep = Separator()
|
|
failed = Failed()
|
|
if not AUTO_SCAN:
|
|
log("AUTO_SCAN off: only Stem Now requests will be separated")
|
|
while not stop.is_set():
|
|
try:
|
|
if not AUTO_SCAN:
|
|
idle(sep, failed, RESCAN_S)
|
|
continue
|
|
done, used = run_pass(sep, failed)
|
|
if done:
|
|
log(f"pass finished: {done:,} tracks separated")
|
|
if used:
|
|
sep.release() # hand the VRAM back while idle
|
|
idle(sep, failed, RESCAN_S)
|
|
except OSError as e:
|
|
if stop.is_set():
|
|
break
|
|
if not is_share_error(e):
|
|
raise
|
|
sep.release()
|
|
stop.wait(SHARE_RETRY_S)
|
|
except BaseException:
|
|
if stop.is_set():
|
|
break
|
|
raise
|
|
log("stopped")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|