Prune stems whose source song was removed or replaced

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>

server/README.md

@@ -130,6 +130,9 @@ the app would briefly see no stems.
 | `DEVICE`     | `cuda`           | `cuda`, `cuda:N` or `cpu`                            |
 | `RESCAN_MIN` | `30`             | minutes between scans once caught up                 |
 | `AUTO_SCAN`  | `0`              | `1` = scan the library every `RESCAN_MIN`; `0` = only do Stem Now requests |
+| `PRUNE`      | `1`              | Remove stems whose song is gone from the library, at start and every `RESCAN_MIN` |
+| `PRUNE_CHANGED` | `0` | Also remove stems of songs replaced in place (size or mtime changed) |
+| `PRUNE_MAX` | `0.2` | Skip a prune that would remove more than this share of the stems. `server/prune.py` is the manual dry run, `--apply` to remove |
 | `FP16`       | `1`              | half precision on the GPU (`0` for fp32)             |
 | `OVERLAP`, `SEGMENT`, `SHIFTS` | `0.25`, `0`, `0` | as in the desktop app's settings     |
 

server/prune.py

@@ -0,0 +1,49 @@
+#!/usr/bin/env python3
+"""prune: list or remove stems whose song is no longer in the music library.
+
+  prune.py                dry run: print orphans (song gone) and changed (song replaced)
+  prune.py --apply        remove the orphans
+  prune.py --apply --changed   also remove stems of replaced songs, so stemd redoes them
+
+MUSIC_DIR and STEMS_DIR as for stemd. Refuses when the music folder holds no tracks.
+Without --force it also refuses to remove more than 20% of the stems in one go.
+"""
+
+import argparse
+import os
+import sys
+
+_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+sys.path.insert(0, os.path.join(_ROOT, "desktop"))
+sys.path.insert(0, _ROOT)
+
+from stemify.library_core import find_stale, iter_tracks, remove_stems  # noqa: E402
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--music", default=os.environ.get("MUSIC_DIR", "/music"))
+    ap.add_argument("--stems", default=os.environ.get("STEMS_DIR", "/stems"))
+    ap.add_argument("--apply", action="store_true")
+    ap.add_argument("--changed", action="store_true")
+    ap.add_argument("--force", action="store_true")
+    a = ap.parse_args()
+    total = sum(1 for _ in iter_tracks(a.music, a.stems))
+    if not total:
+        sys.exit(f"no tracks in {a.music}; is the share mounted?")
+    orphans, changed = find_stale(a.music, a.stems)
+    drop = orphans + (changed if a.changed else [])
+    for label, items in (("orphan", orphans), ("changed", changed)):
+        for p in items:
+            print(f"{label}\t{os.path.relpath(p, a.stems)}")
+    print(f"{len(orphans)} orphans, {len(changed)} changed, {total} tracks in the library")
+    if not a.apply or not drop:
+        return
+    if len(drop) > 0.2 * (total + len(orphans)) and not a.force:
+        sys.exit("that is over 20% of the stems; check the paths or pass --force")
+    remove_stems(a.stems, drop)
+    print(f"removed {len(drop)}")
+
+
+if __name__ == "__main__":
+    main()

server/stemd.Dockerfile

@@ -15,7 +15,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends gcc libc6-dev \
 COPY desktop/stemify/__init__.py desktop/stemify/catalog.py desktop/stemify/library_core.py /app/stemify/
 COPY desktop/stemify/worker/__init__.py desktop/stemify/worker/__main__.py desktop/stemify/worker/audio_io.py \
      desktop/stemify/worker/backends.py /app/stemify/worker/
-COPY server/stemd.py /app/server/
+COPY server/stemd.py server/prune.py /app/server/
 ENV MUSIC_DIR=/music STEMS_DIR=/stems STATE_DIR=/state MODELS_DIR=/models \
     TORCH_HOME=/models/torch HF_HOME=/models/hf NUMBA_CACHE_DIR=/tmp/numba HOME=/tmp \
     PYTHONUNBUFFERED=1 PYTHONDONTWRITEBYTECODE=1 PYTHONWARNINGS=ignore::FutureWarning

server/stemd.py

@@ -21,6 +21,9 @@ Settings (environment):
   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
+  PRUNE       1                remove stems whose song left the library; 0 = off
+  PRUNE_CHANGED 0              1 = also remove stems of songs replaced in place (redone by the next scan)
+  PRUNE_MAX   0.2              skip a prune that would remove more than this share of the stems
   FP16        1                half precision on the GPU
   OVERLAP / SEGMENT / SHIFTS   as in the desktop app's settings
 """
@@ -40,8 +43,8 @@ sys.path.insert(0, os.path.join(_ROOT, "desktop"))  # in the image the package s
 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)
+from stemify.library_core import (MARKER, MAX_SECONDS, SOURCE_RANK, duration, find_stale, is_done,  # noqa: E402
+                                     is_share_error, iter_tracks, remove_stems, signature, write_track)
 
 MUSIC_DIR = os.environ.get("MUSIC_DIR", "/music")
 STEMS_DIR = os.environ.get("STEMS_DIR", "/stems")
@@ -51,6 +54,9 @@ 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")
+PRUNE = os.environ.get("PRUNE", "1") not in ("0", "false", "no", "")
+PRUNE_CHANGED = os.environ.get("PRUNE_CHANGED", "0") in ("1", "true", "yes")
+PRUNE_MAX = float(os.environ.get("PRUNE_MAX", "0.2"))
 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"))
@@ -331,6 +337,26 @@ def serve_queue(sep, failed):
     return used
 
 
+def prune():
+    """Remove stems whose song is gone from MUSIC_DIR (and optionally replaced ones)."""
+    if not PRUNE:
+        return
+    orphans, changed = find_stale(MUSIC_DIR, STEMS_DIR)
+    drop = orphans + (changed if PRUNE_CHANGED else [])
+    if changed and not PRUNE_CHANGED:
+        log(f"prune: {len(changed):,} stems are of a replaced song (PRUNE_CHANGED=1 removes them)")
+    if not drop:
+        return
+    total = sum(1 for _ in iter_tracks(MUSIC_DIR, STEMS_DIR))
+    kept = total - len(orphans)
+    if kept <= 0 or len(drop) > PRUNE_MAX * (kept + len(drop)):
+        log(f"prune: refusing to remove {len(drop):,} of {kept + len(drop):,} stems "
+            f"(over PRUNE_MAX {PRUNE_MAX:.0%}, or the music share looks empty); run server/prune.py --apply to force")
+        return
+    remove_stems(STEMS_DIR, drop)
+    log(f"prune: removed {len(drop):,} stems folders no longer in the library")
+
+
 def run_pass(sep, failed):
     """Separate everything that needs it. Returns (tracks done, whether the GPU was used)."""
     t0 = time.monotonic()
@@ -404,6 +430,7 @@ def main():
         log("AUTO_SCAN off: only Stem Now requests will be separated")
     while not stop.is_set():
         try:
+            prune()
             if not AUTO_SCAN:
                 idle(sep, failed, RESCAN_S)
                 continue