schism-audio/mini-bs-roformer-18m-mlx

Mini-BS-RoFormer 18M — MLX

HiDolen/Mini-BS-RoFormer-18M — a compact Band-Split RoPE Transformer (arXiv:2309.02612) for 4-stem music source separation (bass / drums / other / vocals, note the stem order differs from Demucs) — converted to MLX for Apple Silicon.

MIT licensed — unlike the higher-quality but CC-BY-NC V2 port, this one is commercial-safe. At 17.9M params (72 MB fp32) it is the smallest separation model in this org.

Runs with the open schism-mlx package:

from schism_mlx.separate.bs_roformer import BSRoformer

model = BSRoformer.from_pretrained("schism-audio/mini-bs-roformer-18m-mlx")
stems = model.separate("song.wav")   # {"bass": (2, L), "drums": ..., ...}

Architecture: single-resolution STFT (n_fft 2048, hop 512), 62-band split, 8 axial transformer layer pairs (time/freq) with RoPE, grouped KV attention and gated MLPs, register tokens, plain per-band linear mask heads. This is the generation before V2 — no dual-resolution STFT, no temporal token compression. ~1.7 s for 15 s of stereo on an M5 Max (RTF 0.12).

Upstream reports 9.0 mean SDR (bass 8.31 / drums 9.55 / other 8.14 / vocals 10.03) on its own MUSDB18HQ validation carve-out; the repo's training notebook shows the model was trained on MUSDB18HQ train + test, so the figure is not comparable to published test-set SDRs.

Conversion & verification

• Weights are the original fp32 checkpoint, module paths renamed only — no transposes, no quantization, no finetuning.
• Verified against the repo's own reference implementation: single-chunk forward residual −137.6 dB (max sample diff 2.4e-7 — the fp32 noise floor); full separated tracks, including the reference's fade-window overlap-add, match to −138 dB.
• The reference separate() defaults to a 1 s zeroed gap at each chunk edge, which silences the first and last second of every track; the upstream model card itself passes gap_size=0. This port defaults to gap_size=0 and implements both — parity holds in both modes (−138.2 dB at gap 0, −137.6 dB at the reference default).
• A drums-only recording from schism-audio/e-gmd lands 100.0% of output energy in the drums stem.
• The STFT/iSTFT are reimplemented in numpy and verified against torch.stft/torch.istft in isolation. No torch at inference.
• Note: the source repo's custom code targets transformers 4.55 and fails under transformers 5.x from_pretrained; this port has no transformers dependency at all.

Precision note: on M5-class GPUs, MLX silently runs float32 matmuls at TF32-class precision (ml-explore/mlx#3860). schism_mlx disables this on import; set MLX_ENABLE_TF32=0 if using these weights with your own MLX code.

License

MIT, inherited from the source weights. Model by HiDolen; MLX conversion by schism-audio.