all-in-one-mlx

GPU-accelerated music structure analysis on Apple Silicon (MLX).

all-in-one-mlx is an Apple Silicon–optimized port of the original All‑In‑One Music Structure Analyzer (upstream: mir-aidj/all-in-one). It runs end‑to‑end inference locally using Apple MLX, with an integrated pipeline designed for real songs (including demixing + fast spectrograms).

Given one or more audio tracks, it produces:

• Tempo (BPM)
• Beat times
• Downbeat times
• Section boundaries
• Section labels (intro / verse / chorus / bridge / outro, etc.)

---

Why this repo exists

The upstream project is a strong reference implementation, but macOS Apple Silicon users historically lacked a first‑class GPU‑accelerated inference path. This repository provides that acceleration via MLX, with an emphasis on:

• High performance on M‑series GPUs
• Practical CLI defaults for song inference (demix → spectrogram → model → outputs)
• Faithful behavior to the upstream model + method

I’m releasing all-in-one-mlx alongside all-in-one-mps (PyTorch/MPS acceleration) so Apple Silicon users can choose the stack that fits their environment.

---

Performance

Benchmark on a single file — Apple M4 Max, 128 GB RAM, macOS 26.3:

Project	Time	vs upstream
mir-aidj/all-in-one	75.25s	baseline
mir-aidj/all-in-one	24.63s	~3.1x faster (patched to use MPS)
all-in-one-mps	13.43s	~5.6x faster
all-in-one-mlx (this repo)	5.96s	~12.6x faster

One run, one file — results will vary by hardware and content.

---

Related projects & attribution

Project	Purpose
mir-aidj/all-in-one	Original reference implementation and training code
all-in-one-mlx	This repo: MLX inference + packaging for Apple Silicon
all-in-one-mps	Companion repo: PyTorch/MPS inference for Apple Silicon

This repository began as a fork/port of the upstream project. The original method/model is described in:

• Taejun Kim & Juhan Nam, All‑In‑One Metrical and Functional Structure Analysis with Neighborhood Attentions on Demixed Audio (arXiv:2307.16425)

If you use this in academic work, please cite the paper and the upstream repository.

---

Requirements

Component	Requirement
Hardware	Apple Silicon (M-series)
OS	macOS 14+ (required by MLX wheels)
Python	3.10+

Need CUDA / Linux / Windows? Use the upstream project.

---

Installation

pip

pip install all-in-one-mlx

uv (recommended)

uv pip install all-in-one-mlx

---

Quickstart

Analyze one or more tracks:

allin1-mlx path/to/song.wav
# or multiple:
allin1-mlx path/to/a.wav path/to/b.wav

By default, results are written under:

• ./struct (set with --out-dir)

Common options

• Choose output directory:

allin1-mlx song.wav --out-dir ./struct

• Save visualizations / sonifications:

allin1-mlx song.wav --visualize --viz-dir ./viz
allin1-mlx song.wav --sonify --sonif-dir ./sonif

• Keep intermediate byproducts (demixed audio + spectrograms):

allin1-mlx song.wav --keep-byproducts
# demix files: ./demix (override with --demix-dir)
# specs:       ./spec  (override with --spec-dir)

• Fast spectrogram backend (default) vs reference backend:

allin1-mlx song.wav --spec-backend mlx_fast   # default
allin1-mlx song.wav --spec-backend mlx        # reference path

• Parity guard for mlx_fast (enabled by default):

allin1-mlx song.wav --spec-fast-guard
allin1-mlx song.wav --no-spec-fast-guard
allin1-mlx song.wav --spec-fast-guard-max-abs 1e-3 --spec-fast-guard-mean-abs 1e-4

When --spec-backend mlx_fast is used, the guard runs a one-time comparison against mlx. If the diff exceeds thresholds, the run automatically falls back to mlx for the remaining tracks.

• One-time spectrogram correctness check (reports max/mean diff):

allin1-mlx song.wav --spec-check

• Overwrite specific stages (demix,spec,json,viz,sonify) or everything:

allin1-mlx song.wav --overwrite all
allin1-mlx song.wav --overwrite demix,spec,json

• Timing / performance instrumentation:

allin1-mlx song.wav --timings-path timings.jsonl
allin1-mlx song.wav --timings-path timings.jsonl --timings-viz-path timings.png

---

MLX inference controls

• Select model (pretrained name):

allin1-mlx song.wav --model harmonix-all

• Batch size:

allin1-mlx song.wav --mlx-batch-size 1

• mx.compile for model forward (enabled by default):

allin1-mlx song.wav --no-mlx-compile

• In-memory pipeline for demix + spectrograms (enabled by default):

allin1-mlx song.wav --no-mlx-in-memory

• Ensemble inference parallelism (enabled by default):

allin1-mlx song.wav --no-ensemble-parallel

• Disable multiprocessing (debug / determinism / constrained envs):

allin1-mlx song.wav --no-multiprocess

---

Outputs

The CLI writes analysis artifacts under --out-dir (default ./struct). Exact filenames may vary by model/pipeline version, but outputs include tempo, beats, downbeats, and segment boundaries/labels in machine-readable form.

Optional outputs:

Artifact	Enable with
Visualizations	--visualize and --viz-dir
Sonifications	--sonify and --sonif-dir
Frame-level activations	--activ and optionally --array-dir
Frame-level embeddings	--embed and optionally --array-dir
JSONL timings	--timings-path

Activation and embedding arrays are saved beside the JSON results by default. Use --array-dir to keep these much larger files in a separate directory:

allin1-mlx song.wav --activ --embed \
  --out-dir ./struct --array-dir ./arrays

Demix reproducibility

Demucs applies a random time shift before separation by default. Supply a seed to reproduce that shift selection across runs:

allin1-mlx song.wav --demix-seed 1234

The seed only applies when Demucs actually runs. External stems and reusable cached stems bypass it; use --overwrite demix,spec,json when replacing a disk-backed cached analysis with a seeded run. Seeding Demucs does not guarantee bitwise-stable end-to-end results because later MLX inference can still vary.

Model weights

Item	Behavior
Source	The optional mlx-weights integration is preferred when installed; otherwise verified checkpoints are downloaded on demand
Packaging	Release wheels/sdists do not bundle model weights
Standalone cache	~/.cache/all-in-one-mlx/weights
Custom paths	Use --mlx-weights-dir or explicit --mlx-weights-path and --mlx-config-path

Standalone downloads are pinned to an immutable source revision and checked against the published SHA-256 digest before use. Cached files are rechecked, and a damaged entry is discarded and downloaded again. Install mlx-weights separately if you want to use its shared conversion cache instead.

Known limitations

• Artifact naming uses input basename/stem for intermediate and output files.
• If multiple inputs share the same basename (for example a/song.mp3 and b/song.wav), artifacts may overwrite each other or be reused unexpectedly.
• Workaround: process those files separately or rename files so basenames are unique.

Beat parity workflow

Use the reproducible parity harness to compare upstream CPU, all-in-one-mps, and all-in-one-mlx:

python scripts/compare_beat_parity.py \
  "/path/to/song.wav" \
  --upstream-python /path/to/all-in-one/.venv/bin/python \
  --mps-python /path/to/all-in-one-mps/.venv/bin/python \
  --mlx-python /path/to/all-in-one-mlx/.venv/bin/python \
  --mlx-weights-dir ~/.cache/all-in-one-mlx/weights \
  --output /tmp/allin1_beat_parity.json

The report includes beat/downbeat counts and timing delta summaries (median/mean signed, p90/max abs).

---

License

This project retains the upstream license (MIT). See LICENSE.

---

Issues

Please include:

Include	Example
macOS version + Apple Silicon model	macOS 26.3, M4 Max
Python + MLX versions	Python 3.12.7, mlx x.y
Exact command and logs/traceback	Full allin1-mlx ... command + stack trace