schism-audio/mini-bs-roformer-v2-mlx

Mini-BS-RoFormer V2 (46.8M) — MLX

HiDolen/Mini-BS-RoFormer-V2-46.8M — a compact 2025 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.

License: CC-BY-NC-4.0 (non-commercial), inherited from the source weights. For commercial use, see the MIT-licensed HTDemucs ports in this org instead.

Runs with the open schism-mlx package:

from schism_mlx.separate.bs_roformer import BSRoformer

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

Architecture highlights: dual-resolution STFT (analysis at n_fft 4096, mask application and synthesis at 2048), 80-band split, 4× temporal token compression, 9 axial transformer layer pairs (time/freq) with RoPE, grouped KV attention and gated MLPs, register tokens. 46.8M params, ~0.7 s for 15 s of stereo on an M5 Max (RTF 0.05).

Upstream reports 10.03 SDR on the MUSDB18HQ validation set, but notes the model was trained on MUSDB18HQ train + test, so treat that number as optimistic.

Conversion & verification

• Weights are the original bf16 checkpoint, module paths renamed only — no transposes, no quantization, no finetuning.
• Verified against the repo's own reference implementation (run in fp32): single-chunk forward residual −132.4 dB; full separated tracks, including the reference's fade-window overlap-add, match to −134.0 dB.
• Parity subtlety worth knowing: the source checkpoint stores rotary_emb.inv_freq in bfloat16, and the reference computes its RoPE angles from those quantized values. This port quantizes its recomputed inv_freq through the checkpoint dtype to match. (Ports before 2026-07-22 recomputed inv_freq in fp32 — a ~0.4 rad angle error at long time positions that degraded parity to −90 dB and diverged further on strongly tonal low-frequency content; update schism-mlx if you installed earlier.)
• A drums-only recording from schism-audio/e-gmd lands 100.0% of output energy in the drums stem.
• The dual-resolution 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.57 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

CC-BY-NC-4.0, inherited from the source weights. Model by HiDolen; MLX conversion by schism-audio.