StemRoller: desktop stem separator, Android stems player, and server daemons

- stemroller/: PySide6/QML desktop app and player for separating music into stems
- android/: Plex music player that plays per-stem mixes, with Android Auto support,
  a graphic + parametric (Equalizer APO-style) EQ, and offline downloads
- server/: stemserve (serves stems to the app) and stemd (separates the library on
  the server's GPU, with "Stem Now" priority requests)

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

stemroller/worker/__main__.py

@@ -0,0 +1,289 @@
+"""Separation worker process.
+
+Connects to the UI's local socket and speaks newline-delimited JSON:
+
+  UI -> worker   start {id, file, model, device, two_stems, out_dir, segment, shifts, overlap, fp16}
+                 cancel {id}
+                 release {}   drop cached models and free VRAM
+                 quit {}
+  worker -> UI   hello {devices[], torch, hip}
+                 vram {device, free_mb, total_mb}
+                 progress {id, pct, stage}
+                 done {id, stems[], peaks_path, samplerate, duration, out_dir}
+                 error {id, msg}
+                 cancelled {id}
+
+Runs one job at a time; the UI owns the queue.
+"""
+
+import argparse
+import json
+import os
+import queue
+import re
+import shutil
+import socket
+import subprocess
+import sys
+import threading
+import time
+import traceback
+
+# Fast MIOpen kernel selection: avoids multi-minute first-run tuning on RDNA4.
+os.environ.setdefault("MIOPEN_FIND_MODE", "FAST")
+os.environ.setdefault("TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL", "1")
+os.environ.setdefault("MIOPEN_LOG_LEVEL", "2")  # silence workspace warnings
+
+PROGRESS_INTERVAL = 1 / 30  # doc: at most 30 progress updates per second
+MODEL_DIR = os.path.join(os.path.expanduser("~/.cache/stemroller"), "models")
+
+
+class Conn:
+    def __init__(self, path):
+        self.sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+        self.sock.connect(path)
+        self.lock = threading.Lock()
+        self.rfile = self.sock.makefile("r", encoding="utf-8")
+
+    def send(self, type_, **payload):
+        line = json.dumps({"type": type_, **payload}) + "\n"
+        with self.lock:
+            self.sock.sendall(line.encode())
+
+    def __iter__(self):
+        for line in self.rfile:
+            line = line.strip()
+            if line:
+                yield json.loads(line)
+
+
+class Job:
+    def __init__(self, msg, conn):
+        from stemroller.catalog import MODELS_BY_ID
+        self.id = msg["id"]
+        self.file = msg["file"]
+        self.model = MODELS_BY_ID[msg["model"]]
+        self.device = msg.get("device", "auto")
+        self.two_stems = bool(msg.get("two_stems"))
+        self.out_dir = msg["out_dir"]
+        self.segment = float(msg.get("segment") or 0)
+        self.shifts = int(msg.get("shifts") or 0)
+        self.overlap = float(msg.get("overlap", 0.25))
+        self.fp16 = bool(msg.get("fp16", True))
+        # audio-separator expresses overlap as a divisor of the chunk (step = chunk / overlap);
+        # RoFormer's windowed overlap-add wants at least 4 (75%).
+        self.overlap_int = max(4, round(1 / max(0.05, 1 - self.overlap)))
+        self.cancel_event = threading.Event()
+        self._conn = conn
+        self._last_sent = 0.0
+
+    def check(self):
+        if self.cancel_event.is_set():
+            from stemroller.worker.backends import Cancelled
+            raise Cancelled()
+
+    def report(self, stage, pct, force=False):
+        now = time.monotonic()
+        if force or now - self._last_sent >= PROGRESS_INTERVAL:
+            self._last_sent = now
+            self._conn.send("progress", id=self.id, pct=round(float(pct), 4), stage=stage)
+
+
+def gpu_marketing_names():
+    """gfx arch -> marketing name from rocminfo (torch reports a generic name on ROCm)."""
+    names = {}
+    if not shutil.which("rocminfo"):
+        return names
+    try:
+        out = subprocess.run(["rocminfo"], capture_output=True, text=True, timeout=5).stdout
+    except Exception:
+        return names
+    arch = None
+    for line in out.splitlines():
+        m = re.match(r"\s*Name:\s+(gfx\w+)\s*$", line)
+        if m:
+            arch = m.group(1)
+        m = re.match(r"\s*Marketing Name:\s+(.+?)\s*$", line)
+        if m and arch:
+            names.setdefault(arch, m.group(1))
+            arch = None
+    return names
+
+
+def pci_address(props):
+    bus = getattr(props, "pci_bus_id", None)
+    if bus is None:
+        return None
+    return f"{props.pci_domain_id:04x}:{bus:02x}:{props.pci_device_id:02x}.0"
+
+
+def sysfs_vram_paths(props):
+    pci = pci_address(props)
+    if pci is None:
+        return None
+    base = f"/sys/bus/pci/devices/{pci}"
+    total, used = base + "/mem_info_vram_total", base + "/mem_info_vram_used"
+    return (total, used) if os.path.exists(used) else None
+
+
+def enumerate_devices():
+    import torch
+    devices = []
+    if torch.cuda.is_available():
+        names = gpu_marketing_names() if torch.version.hip else {}
+        for i in range(torch.cuda.device_count()):
+            p = torch.cuda.get_device_properties(i)
+            if getattr(p, "is_integrated", False):
+                continue  # skip APU iGPUs (e.g. the 9800X3D's gfx1036)
+            arch = getattr(p, "gcnArchName", "").split(":")[0]
+            devices.append({
+                "id": f"cuda:{i}",
+                "name": names.get(arch, p.name),
+                "arch": arch,
+                "total_mb": p.total_memory // 2**20,
+                "pci": pci_address(p),
+                "_sysfs": sysfs_vram_paths(p),
+            })
+    return devices
+
+
+def vram_loop(conn, devices):
+    import torch
+    while True:
+        for d in devices:
+            try:
+                if d["_sysfs"]:
+                    total = int(open(d["_sysfs"][0]).read()) // 2**20
+                    used = int(open(d["_sysfs"][1]).read()) // 2**20
+                    free = total - used
+                else:
+                    free, total = (v // 2**20 for v in torch.cuda.mem_get_info(int(d["id"].split(":")[1])))
+                conn.send("vram", device=d["id"], free_mb=free, total_mb=total)
+            except Exception:
+                pass
+        time.sleep(1.0)
+
+
+def resolve_device(choice, devices):
+    import torch
+    if choice == "cpu" or not devices:
+        return torch.device("cpu")
+    if choice != "auto" and any(d["id"] == choice for d in devices):
+        return torch.device(choice)
+    best = max(devices, key=lambda d: d["total_mb"])
+    return torch.device(best["id"])
+
+
+def run_job(job, backends, devices, conn):
+    import numpy as np
+    import torch
+    from stemroller.worker.audio_io import save_peaks, write_stem
+
+    device = resolve_device(job.device, devices)
+    backend = backends[job.model["family"]]
+    os.makedirs(job.out_dir, exist_ok=True)
+    try:
+        stems, sr = backend.run(job, device)
+    except torch.OutOfMemoryError:
+        torch.cuda.empty_cache()
+        raise RuntimeError("Out of GPU memory. Lower the segment size in Settings or pick CPU.")
+
+    if job.two_stems and "vocals" in stems and len(stems) > 2:
+        rest = [a for k, a in stems.items() if k != "vocals"]
+        stems = {"vocals": stems["vocals"], "instrumental": np.sum(rest, axis=0)}
+
+    job.report("Writing stems", 0.94, force=True)
+    out = []
+    for name, audio in stems.items():
+        job.check()
+        path = os.path.join(job.out_dir, f"{name}.wav")
+        write_stem(path, audio, sr)
+        out.append({"name": name, "path": path})
+    job.report("Building waveforms", 0.98, force=True)
+    peaks_path = os.path.join(job.out_dir, "peaks.npz")
+    save_peaks(peaks_path, stems, sr)
+    duration = next(iter(stems.values())).shape[1] / sr
+    meta = {"source": job.file, "model": job.model["id"], "two_stems": job.two_stems,
+            "stems": out, "peaks_path": peaks_path, "samplerate": sr, "duration": duration,
+            "device": str(device), "finished": time.time()}
+    with open(os.path.join(job.out_dir, "meta.json"), "w") as f:
+        json.dump(meta, f)
+    if device.type == "cuda":
+        torch.cuda.empty_cache()
+    return {"stems": out, "peaks_path": peaks_path, "samplerate": sr, "duration": duration,
+            "out_dir": job.out_dir, "device": str(device)}
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--socket", required=True)
+    args = ap.parse_args()
+
+    try:
+        os.nice(5)  # UI first; the GPU does the heavy lifting anyway
+    except OSError:
+        pass
+    conn = Conn(args.socket)
+    conn.send("booting")
+
+    import torch
+    torch.set_num_threads(max(1, (os.cpu_count() or 4) - 2))
+    from stemroller.worker.backends import Cancelled, DemucsBackend, RoformerBackend
+
+    devices = enumerate_devices()
+    conn.send("hello", devices=[{k: v for k, v in d.items() if not k.startswith("_")} for d in devices],
+              torch=torch.__version__, hip=torch.version.hip)
+    threading.Thread(target=vram_loop, args=(conn, devices), daemon=True).start()
+
+    jobs: queue.Queue = queue.Queue()
+    current: dict = {}
+
+    def reader():
+        try:
+            for msg in conn:
+                t = msg.get("type")
+                if t in ("start", "release"):
+                    jobs.put(msg)
+                elif t == "cancel":
+                    job = current.get("job")
+                    if job and job.id == msg.get("id"):
+                        job.cancel_event.set()
+                elif t == "quit":
+                    break
+        except (OSError, ValueError):
+            pass
+        # UI went away or asked us to quit: no point finishing work nobody will read.
+        os._exit(0)
+
+    threading.Thread(target=reader, daemon=True).start()
+
+    backends = {"demucs": DemucsBackend(), "roformer": RoformerBackend(MODEL_DIR)}
+    while True:
+        msg = jobs.get()
+        if msg.get("type") == "release":
+            for b in backends.values():
+                b.release()
+            continue
+        try:
+            job = Job(msg, conn)
+        except Exception as e:
+            conn.send("error", id=msg.get("id"), msg=f"Bad job: {e}")
+            continue
+        current["job"] = job
+        try:
+            result = run_job(job, backends, devices, conn)
+            conn.send("done", id=job.id, **result)
+        except Cancelled:
+            shutil.rmtree(job.out_dir, ignore_errors=True)
+            conn.send("cancelled", id=job.id)
+        except Exception as e:
+            traceback.print_exc()
+            shutil.rmtree(job.out_dir, ignore_errors=True)
+            conn.send("error", id=job.id, msg=str(e) or e.__class__.__name__)
+        finally:
+            current.pop("job", None)
+            sys.stdout.flush()
+
+
+if __name__ == "__main__":
+    main()

stemroller/worker/audio_io.py

@@ -0,0 +1,79 @@
+"""Audio decode/encode and waveform peak generation for the worker."""
+
+import subprocess
+
+import numpy as np
+import soundfile as sf
+
+PEAK_BASE = 256  # samples per bucket at mip level 0
+
+
+def _source_channels(path: str) -> int:
+    try:
+        out = subprocess.run(["ffprobe", "-v", "error", "-select_streams", "a:0", "-show_entries",
+                              "stream=channels", "-of", "csv=p=0", path],
+                             capture_output=True, text=True, timeout=10).stdout
+        return int(out.strip().splitlines()[0])
+    except (OSError, ValueError, IndexError, subprocess.SubprocessError):
+        return 0
+
+
+def load_audio(path: str, samplerate: int, channels: int = 2) -> np.ndarray:
+    """Decode any ffmpeg-readable file to float32 (channels, samples)."""
+    # ffmpeg's -ac upmix treats mono as centre and spreads it at -3 dB per side;
+    # decode mono as mono and duplicate it so stems keep the source level.
+    decode_ch = 1 if _source_channels(path) == 1 else channels
+    cmd = ["ffmpeg", "-nostdin", "-v", "error", "-i", path, "-map", "0:a:0",
+           "-f", "f32le", "-acodec", "pcm_f32le", "-ac", str(decode_ch), "-ar", str(samplerate), "-"]
+    proc = subprocess.run(cmd, capture_output=True)
+    if proc.returncode != 0:
+        lines = [l.strip() for l in proc.stderr.decode(errors="replace").splitlines() if l.strip()]
+        reason = lines[-1].split(": ", 1)[-1] if lines else f"ffmpeg exited with {proc.returncode}"
+        raise RuntimeError(f"Could not decode the file ({reason[:160]})")
+    data = np.frombuffer(proc.stdout, dtype=np.float32)
+    if data.size == 0:
+        raise RuntimeError("file contains no audio")
+    data = data.reshape(-1, decode_ch).T
+    if decode_ch != channels:
+        data = np.repeat(data, channels, axis=0)
+    return data.copy()
+
+
+def write_stem(path: str, audio: np.ndarray, samplerate: int):
+    """Write (channels, samples) audio. 24-bit unless it would clip, then float."""
+    peak = float(np.abs(audio).max()) if audio.size else 0.0
+    subtype = "FLOAT" if peak > 0.999 else "PCM_24"
+    sf.write(path, audio.T, samplerate, subtype=subtype, format="WAV")
+
+
+def compute_peaks(audio: np.ndarray) -> list[np.ndarray]:
+    """Min/max mip pyramid of a (channels, samples) signal, mono-folded.
+
+    Level k has PEAK_BASE * 2**k samples per bucket; each level is (n, 2) float16.
+    """
+    lo = audio.min(axis=0)
+    hi = audio.max(axis=0)
+    n = -(-lo.size // PEAK_BASE)
+    pad = n * PEAK_BASE - lo.size
+    if pad:
+        lo = np.concatenate([lo, np.zeros(pad, lo.dtype)])
+        hi = np.concatenate([hi, np.zeros(pad, hi.dtype)])
+    lvl_lo = lo.reshape(n, PEAK_BASE).min(axis=1)
+    lvl_hi = hi.reshape(n, PEAK_BASE).max(axis=1)
+    levels = [np.stack([lvl_lo, lvl_hi], axis=1).astype(np.float16)]
+    while lvl_lo.size > 64:
+        if lvl_lo.size % 2:
+            lvl_lo = np.append(lvl_lo, 0)
+            lvl_hi = np.append(lvl_hi, 0)
+        lvl_lo = lvl_lo.reshape(-1, 2).min(axis=1)
+        lvl_hi = lvl_hi.reshape(-1, 2).max(axis=1)
+        levels.append(np.stack([lvl_lo, lvl_hi], axis=1).astype(np.float16))
+    return levels
+
+
+def save_peaks(path: str, stems: dict[str, np.ndarray], samplerate: int):
+    arrays = {"samplerate": np.array(samplerate), "base": np.array(PEAK_BASE)}
+    for name, audio in stems.items():
+        for k, level in enumerate(compute_peaks(audio)):
+            arrays[f"{name}/{k}"] = level
+    np.savez(path, **arrays)

stemroller/worker/backends.py

@@ -0,0 +1,251 @@
+"""Separation backends. Each returns ({stem: float32 (channels, samples)}, samplerate)."""
+
+import gc
+import logging
+import os
+import random
+
+import numpy as np
+import soundfile as sf
+import torch
+
+from .audio_io import load_audio
+
+
+class Cancelled(BaseException):
+    """Raised from progress hooks to abort a running separation."""
+
+
+class DemucsBackend:
+    """Demucs with batched segments.
+
+    demucs.apply.apply_model runs one segment per forward pass, which leaves an
+    RX 9070 XT mostly idle waiting on MIOpen host overhead. Batching 8 segments
+    under fp16 autocast measured ~110x realtime vs ~31x for the stock loop.
+    """
+
+    BATCH = 8
+
+    def __init__(self):
+        self._model = None
+        self._model_name = None
+
+    def _get(self, name):
+        if self._model_name != name:
+            from demucs.pretrained import get_model
+            self._model = None
+            torch.cuda.empty_cache()
+            self._model = get_model(name).eval()
+            self._model_name = name
+        return self._model
+
+    def release(self):
+        self._model = self._model_name = None
+        gc.collect()
+        torch.cuda.empty_cache()
+
+    def run(self, job, device):
+        from demucs.apply import BagOfModels
+
+        job.report("Loading model", 0.01)
+        bag = self._get(job.model["id"])
+        sub_models = bag.models if isinstance(bag, BagOfModels) else [bag]
+        weights = bag.weights if isinstance(bag, BagOfModels) else [[1.0] * len(bag.sources)]
+        job.check()
+        job.report("Decoding audio", 0.03)
+        wav = torch.from_numpy(load_audio(job.file, bag.samplerate, bag.audio_channels))
+        job.check()
+
+        ref = wav.mean(0)
+        mean, std = ref.mean(), ref.std() + 1e-8
+        mix = ((wav - mean) / std).to(device)
+        length = mix.shape[-1]
+
+        passes = max(1, job.shifts)
+        max_shift = int(0.5 * bag.samplerate) if job.shifts else 0
+        padded = torch.nn.functional.pad(mix, (max_shift, max_shift)) if max_shift else mix
+
+        # Work units for progress: one per segment per model per shift.
+        plans = []
+        for sub in sub_models:
+            seg_len = int(min(job.segment, float(sub.segment)) * bag.samplerate) if job.segment else int(float(sub.segment) * bag.samplerate)
+            stride = max(1, int((1 - job.overlap) * seg_len))
+            plans.append((sub, seg_len, stride))
+        total_units = sum(passes * len(range(0, length + max_shift, stride)) for _, _, stride in plans)
+        done_units = 0
+
+        def tick(n):
+            nonlocal done_units
+            done_units += n
+            job.check()
+            job.report("Separating", 0.05 + 0.88 * min(1.0, done_units / total_units))
+
+        self.fp16 = job.fp16
+        job.report("Separating", 0.05)
+        estimates = torch.zeros(len(bag.sources), mix.shape[0], length, device=device)
+        totals = [0.0] * len(bag.sources)
+        for (sub, seg_len, stride), w in zip(plans, weights):
+            sub.to(device)
+            out = torch.zeros_like(estimates)
+            for _ in range(passes):
+                offset = random.randint(0, max_shift) if max_shift else 0
+                piece = padded[..., offset: offset + length + max_shift - offset] if max_shift else padded
+                res = self._split(sub, piece, seg_len, stride, job.overlap, device, tick)
+                out += res[..., max_shift - offset: max_shift - offset + length]
+            out /= passes
+            for k, wk in enumerate(w):
+                estimates[k] += out[k] * wk
+                totals[k] += wk
+            sub.cpu()
+        for k in range(len(totals)):
+            estimates[k] /= totals[k]
+
+        out = (estimates * std.to(device) + mean.to(device)).cpu().numpy()
+        return dict(zip(bag.sources, out)), bag.samplerate
+
+    def _split(self, model, mix, seg_len, stride, overlap, device, tick):
+        """Overlap-add over fixed-length segments, BATCH segments per forward pass."""
+        channels, length = mix.shape
+        n_src = len(model.sources)
+        weight = torch.cat([torch.arange(1, seg_len // 2 + 1, device=device),
+                            torch.arange(seg_len - seg_len // 2, 0, -1, device=device)]).float()
+        weight /= weight.max()
+        out = torch.zeros(n_src, channels, length, device=device)
+        sum_weight = torch.zeros(length, device=device)
+        offsets = list(range(0, length, stride))
+        use_amp = device.type == "cuda" and self.fp16
+        batch = self.BATCH
+        i = 0
+        while i < len(offsets):
+            group = offsets[i:i + batch]
+            chunks = torch.zeros(len(group), channels, seg_len, device=device)
+            for j, off in enumerate(group):
+                piece = mix[:, off:off + seg_len]
+                chunks[j, :, :piece.shape[-1]] = piece
+            try:
+                with torch.no_grad(), torch.autocast("cuda", dtype=torch.float16, enabled=use_amp):
+                    res = model(chunks).float()
+            except torch.OutOfMemoryError:
+                if batch == 1:
+                    raise
+                torch.cuda.empty_cache()
+                batch = max(1, batch // 2)
+                continue
+            for j, off in enumerate(group):
+                n = min(seg_len, length - off)
+                out[..., off:off + n] += weight[:n] * res[j, ..., :n]
+                sum_weight[off:off + n] += weight[:n]
+            i += len(group)
+            tick(len(group))
+        return out / sum_weight
+
+
+def _progress_iter_factory(job, lo, hi):
+    """Stand-in for tqdm inside audio-separator: reports progress and checks cancel."""
+
+    class ProgressIter:
+        def __init__(self, iterable=None, *args, total=None, **kwargs):
+            self.iterable = iterable
+            self.total = total if total is not None else (len(iterable) if hasattr(iterable, "__len__") else None)
+            self.n = 0
+
+        def __iter__(self):
+            for item in self.iterable:
+                self.update(0)
+                yield item
+                self.n += 1
+
+        def update(self, n=1):
+            self.n += n
+            job.check()
+            if self.total:
+                job.report("Separating", lo + (hi - lo) * min(1.0, self.n / self.total))
+
+        def close(self):
+            pass
+
+        def __enter__(self):
+            return self
+
+        def __exit__(self, *exc):
+            return False
+
+    return ProgressIter
+
+
+class RoformerBackend:
+    def __init__(self, model_dir):
+        self.model_dir = model_dir
+        self._sep = None
+        self._key = None
+
+    def _get(self, job, device):
+        key = (job.model["file"], str(device), job.overlap_int, job.fp16)
+        if self._key != key:
+            from audio_separator.separator import Separator
+            self._sep = None
+            torch.cuda.empty_cache()
+            if not os.path.exists(os.path.join(self.model_dir, job.model["file"])):
+                job.report("Downloading model", 0.01)
+            sep = Separator(
+                log_level=logging.WARNING,
+                model_file_dir=self.model_dir,
+                output_dir=job.out_dir,
+                normalization_threshold=1.0,
+                amplification_threshold=0.0,
+                use_autocast=device.type == "cuda" and job.fp16,
+                mdxc_params={"segment_size": 256, "override_model_segment_size": False,
+                             "batch_size": 1, "overlap": job.overlap_int, "pitch_shift": 0},
+            )
+            # Separator picks "cuda" (current device) on its own; honour an explicit choice.
+            if device.type == "cpu":
+                sep.torch_device = torch.device("cpu")
+            elif device.index is not None:
+                torch.cuda.set_device(device.index)
+            sep.load_model(model_filename=job.model["file"])
+            self._sep, self._key = sep, key
+        return self._sep
+
+    def release(self):
+        self._sep = self._key = None
+        gc.collect()
+        torch.cuda.empty_cache()
+
+    def run(self, job, device):
+        from audio_separator.separator.architectures import mdxc_separator
+
+        job.report("Loading model", 0.02)
+        sep = self._get(job, device)
+        job.check()
+
+        captured = {}
+
+        def capture(stem_path, source, stem_name):
+            captured[stem_name] = source
+            return {stem_name: source}
+
+        inst = sep.model_instance
+        inst.final_process = capture
+        inst.output_dir = job.out_dir
+        # audio-separator reads with libsndfile, which rejects many MP3s (big ID3 tags) and
+        # FLACs with one damaged frame; ffmpeg decodes both, so hand it a clean WAV instead.
+        job.report("Decoding audio", 0.03)
+        os.makedirs(job.out_dir, exist_ok=True)
+        mix_path = os.path.join(job.out_dir, "mix.wav")
+        sf.write(mix_path, load_audio(job.file, inst.sample_rate, 2).T, inst.sample_rate, subtype="FLOAT")
+        job.check()
+        mdxc_separator.tqdm = _progress_iter_factory(job, 0.05, 0.93)
+        job.report("Separating", 0.05)
+        try:
+            sep.separate(mix_path)
+        finally:
+            os.remove(mix_path)
+
+        two_stem_model = len(job.model["stems"]) == 2
+        stems = {}
+        for name, source in captured.items():
+            key = name.lower()
+            if two_stem_model and key != "vocals":
+                key = "instrumental"
+            stems[key] = np.ascontiguousarray(source.T, dtype=np.float32)
+        return stems, inst.sample_rate

stemroller/worker/cli.py

@@ -0,0 +1,106 @@
+"""CLI harness for the worker: spawns it, speaks the IPC protocol, prints progress.
+
+  stemroller-worker-cli song.mp3 --model htdemucs [--device cuda:0] [--two-stems] [--cancel-after 3]
+"""
+
+import argparse
+import hashlib
+import json
+import os
+import socket
+import subprocess
+import sys
+import tempfile
+import time
+
+from stemroller.catalog import MODELS_BY_ID
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("files", nargs="+")
+    ap.add_argument("--model", default="htdemucs", choices=list(MODELS_BY_ID))
+    ap.add_argument("--device", default="auto")
+    ap.add_argument("--two-stems", action="store_true")
+    ap.add_argument("--shifts", type=int, default=0)
+    ap.add_argument("--segment", type=float, default=0)
+    ap.add_argument("--overlap", type=float, default=0.25)
+    ap.add_argument("--fp32", action="store_true")
+    ap.add_argument("--out", default=None, help="output root (default: a temp dir)")
+    ap.add_argument("--cancel-after", type=float, default=None, help="send cancel after N seconds")
+    args = ap.parse_args()
+
+    out_root = args.out or tempfile.mkdtemp(prefix="stemroller-cli-")
+    sock_path = os.path.join(tempfile.mkdtemp(), "w.sock")
+    server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+    server.bind(sock_path)
+    server.listen(1)
+
+    t0 = time.monotonic()
+    proc = subprocess.Popen([sys.executable, "-m", "stemroller.worker", "--socket", sock_path])
+    conn, _ = server.accept()
+    rfile = conn.makefile("r")
+
+    def send(**msg):
+        conn.sendall((json.dumps(msg) + "\n").encode())
+
+    pending = []
+    for i, f in enumerate(args.files):
+        jid = f"cli{i}"
+        out_dir = os.path.join(out_root, hashlib.sha1(f.encode()).hexdigest()[:10] + "_" + args.model)
+        pending.append(jid)
+        send(type="start", id=jid, file=os.path.abspath(f), model=args.model, device=args.device,
+             two_stems=args.two_stems, out_dir=out_dir, segment=args.segment, shifts=args.shifts,
+             overlap=args.overlap, fp16=not args.fp32)
+
+    job_started = None
+    n_progress = 0
+    cancel_sent = False
+    rc = 0
+    conn.settimeout(0.2)
+    buf = ""
+    while pending:
+        if args.cancel_after and job_started and not cancel_sent and time.monotonic() - job_started > args.cancel_after:
+            send(type="cancel", id=pending[0])
+            cancel_sent = True
+            print("\n-> cancel sent")
+        try:
+            chunk = conn.recv(65536).decode()
+        except socket.timeout:
+            continue
+        if not chunk:
+            print("worker exited unexpectedly")
+            rc = 1
+            break
+        buf += chunk
+        while "\n" in buf:
+            line, buf = buf.split("\n", 1)
+            msg = json.loads(line)
+            t = msg["type"]
+            el = time.monotonic() - t0
+            if t == "hello":
+                print(f"[{el:6.1f}s] hello torch={msg['torch']} hip={msg['hip']} devices={msg['devices']}")
+            elif t == "progress":
+                n_progress += 1
+                job_started = job_started or time.monotonic()
+                print(f"\r[{el:6.1f}s] {msg['id']} {msg['stage']:<20} {msg['pct'] * 100:5.1f}%", end="", flush=True)
+            elif t == "done":
+                dur = time.monotonic() - job_started
+                print(f"\n[{el:6.1f}s] done {msg['id']} in {dur:.1f}s ({msg['duration'] / dur:.1f}x realtime) "
+                      f"on {msg['device']}, {n_progress} progress msgs ({n_progress / dur:.1f}/s)")
+                for s in msg["stems"]:
+                    print("   ", s["name"], s["path"])
+                pending.remove(msg["id"])
+                job_started, n_progress = None, 0
+            elif t in ("error", "cancelled"):
+                print(f"\n[{el:6.1f}s] {t} {msg}")
+                pending.remove(msg["id"])
+                job_started, n_progress = None, 0
+                rc = rc or (1 if t == "error" else 0)
+    send(type="quit")
+    proc.wait(timeout=10)
+    sys.exit(rc)
+
+
+if __name__ == "__main__":
+    main()