bandit-infer

Inference-only Bandit checkpoint registry and session contract for cinematic audio source separation. It preserves separate v1 and v2 inference graphs, with verified public execution for the official v1 Music64 and v2 Multi CASS checkpoints.

Bandit and Bandit v2 are research implementations for the Cinematic Audio Source Separation (CASS) task. Their public repositories include training and experiment systems; this standalone package is intended to preserve only a small, installable inference surface with explicit checkpoint provenance and model lifetime.

• Karn Watcharasupat and collaborators published Bandit v1 and Bandit v2, both Apache-2.0.
• The official v1 weights are hosted by Zenodo record 10160698 under CC-BY-NC-4.0; the official v2 weights are hosted by Zenodo record 12701995 under CC-BY-SA-4.0. Those licenses apply to weights, not the code.

Please cite the original work, not this packaging effort:

@article{watcharasupat2023bandit,
  title={A Generalized Bandsplit Neural Network for Cinematic Audio Source Separation},
  author={Watcharasupat, Karn and Wu, Chih-Wei and Ding, Yiwei and Orife, Iroro and Hipple, Aaron J. and Williams, Phillip A. and Kramer, Scott and Lerch, Alexander and Wolcott, William},
  journal={IEEE Open Journal of Signal Processing}, year={2023},
  doi={10.1109/OJSP.2023.3339428}
}

• A package-owned 28-entry catalog: 21 v1 Bandit artifacts and seven v2 artifacts. dnr-demucs.ckpt and dnr-umxhq.ckpt are explicitly excluded.
• Independent v1 and v2 dispatch namespaces, selected by family=. No checkpoint conversion and no shared model-layer abstraction is permitted.
• BanditSession provides idempotent load, ready-only infer, reloadable release, terminal close, status, cache_info, and context management.
• auto, cpu, cuda, cuda:N, and mps device validation. Explicit unavailable CUDA or MPS requests fail instead of silently falling back. auto deliberately still means cuda-else-cpu: it never promotes a Mac caller onto MPS, since that would move their numerical output without their asking. Request device="mps" explicitly to use it.
• backend= selects the compute framework ("torch" default, optional "mlx", or "auto") -- a separate axis from family= and device=. See "MLX backend" below.

v1-mus64-l1snr and v2-multi are strict-loaded from official downloads and checked stem-by-stem against fixtures generated by their untouched upstream implementations. Both return speech, music, and effects at their native sample rates. v1 direct inference processes each input channel independently; v2 follows its upstream chunked tensor handler.

The other 26 catalog entries remain provenance records, not supported models: their official bytes have not yet been SHA-256 recorded and parity-tested. A catalog entry never implies that an unverified checkpoint can load.

pip install bandit-infer

For development, use uv sync --all-extras --dev.

bandit-infer --list-models

from bandit_infer import BanditSession

with BanditSession("v1-mus64-l1snr", device="auto") as session:
    stems = session.infer(audio, sample_rate=44100)

v1-mus64-l1snr returns speech, music, and effects from 44.1 kHz input. v2-multi returns the same stems from 48 kHz input. The package deliberately rejects other sample rates: resampling is not a hidden convenience layer and has no recorded parity evidence. Multi-channel arrays are shaped (channels, samples); v1 runs channels independently, while v2 retains its native tensor-handler behavior.

Two independent choices, plus family= (see "Catalog and usage" above):

backend= selects the compute framework -- a separate axis from family= (the v1/v2 model architecture, renamed from an earlier backend= to free this name for its org-wide meaning -- see CHANGELOG.md) and from device= (where a Torch backend runs). This is the same two-axis shape six sibling OpenMIRLab packages use. backend="mlx" owns its own Apple Silicon execution and accepts only device="auto" or device="mps", refusing anything else rather than ignoring it (a Torch-only value like "cuda" raises rather than being silently ignored); requesting backend="mlx" without the [mlx] extra installed also raises rather than falling back to Torch.

from bandit_infer import BanditSession

with BanditSession("v1-mus64-l1snr", backend="mlx") as session:
    stems = session.infer(audio, sample_rate=44100)

There is no --backend CLI flag yet (bandit-infer's CLI is catalog-only: --model, --list-models) -- the compute-framework switch is Python-API-only for now.

An optional, additive MLX compute path -- native execution on Apple Silicon, alongside the default Torch path. There is no upstream MLX implementation of Bandit, so both _v1 and _v2 graphs are ported from scratch directly from this package's own Torch source (src/bandit_infer/mlx/v1/, src/bandit_infer/mlx/v2/), not vendored from a third-party MLX project the way this org's other MLX backends are. Install it with the extra, which is never part of the core install:

pip install 'bandit-infer[mlx]'

Support covers exactly what the Torch v1/v2 runtimes construct for every registry entry: the RNN-based SeqBandModellingModule core (every one of the 28 registry variant strings resolves to it -- v1's TransformerTimeFreqModule/ConvolutionalTimeFreqModule and v2's Wiener/ non-musical-band paths are never reached by any checkpoint this package lists, so they are out of scope for the MLX port too). As with the Torch path, only v1-mus64-l1snr and v2-multi have independently verified SHA-256 values and can actually load -- the checkpoint gate in checkpoints.py applies identically regardless of backend.

Measured Torch-vs-MLX parity through the public BanditSession API on real audio, including a zero-padded and a near-silent tail (Apple Silicon, 2026-07-31): see CHANGELOG.md's MLX entry for the recorded numbers.

MPS and MLX both need an arm64 Python interpreter. Under Rosetta/x86_64 they report as unavailable rather than failing loudly -- an x86_64 interpreter makes torch.backends.mps.is_available() return False, and MLX ships no macOS x86_64 wheel at all, so it cannot even be installed there (published wheels are macosx arm64, manylinux aarch64/x86_64, and win); an accelerated path just looks absent rather than misconfigured. This is easy to hit without noticing: an x86_64 uv resolves x86_64 interpreters, so uv sync can silently produce an environment where the accelerated paths structurally cannot exist. Check with python -c "import platform; print(platform.machine())" -- it must print arm64.

Weights will never be bundled, committed, rehosted, or converted. The default folder is ~/.cache/bandit-infer/, overridable by weights_dir or BANDIT_INFER_WEIGHTS; direct-path loading accepts checkpoint_path plus a caller-provided SHA-256. Official URLs and sizes are in the packaged checkpoints.toml.

Automatic cache/download is enabled only for v1-mus64-l1snr and v2-multi, whose SHA-256 values were computed from official downloads. The remaining entries fail closed until they receive the same evidence. Zenodo's published MD5 is retained for provenance, but never substitutes for SHA-256.

• model weights or a weight mirror;
• model training, evaluation metrics, datasets/dataloaders, experiments, Hydra, Lightning trainer surfaces, or checkpoint conversion;
• the baseline Demucs/UMXHQ artifacts in v1's Zenodo record.

Run the commands recorded in CLAUDE.md. The two flagship parity gates require caller-local official checkpoints and GPU memory; fixtures are small deterministic inputs, not model weights.

Code is Apache-2.0. Model weights have their separate licenses stated above.