allinone-onnx
This repo runs allin1 to generate stem separation and a compact
analysis JSON for a given audio file.
This project includes ONNX + DLL exports based on:
Setup (uv)
- Install Python 3.10 (matches
.python-version):uv python install 3.10 - Sync dependencies:
uv sync
Run
Use the uv environment to execute the script:
uv run python src_python/run_allinone.py --input "path\\to\\audio.wav"
Outputs:
- Stems under
output/htdemucs/<stem>/ - Compact analysis JSON at
output/analysis.json
Select a specific model and keep intermediate outputs (for parity checks):
uv run python src_python/run_allinone.py --input "path\\to\\audio.wav" --model harmonix-fold0 --keep-byproducts
Deterministic Demucs pre-demix (seed + shifts), keeping stems/specs:
uv run python src_python/run_allinone.py --input "path\\to\\audio.wav" --demucs-shifts 0 --demucs-seed 1234 --keep-byproducts
If the allin1 JSON lands somewhere else, pass --result-json:
uv run python src_python/run_allinone.py --input "path\\to\\audio.wav" --result-json "path\\to\\MASTER.json"
Notes
- This project is configured to pull PyTorch CUDA 12.8 wheels via the custom
index in
pyproject.toml. If you need CPU-only wheels, remove the custom torch/torchaudio/torchvision index entries and re-runuv sync. - On Windows,
allin1expects NATTEN CUDA kernels. This repo includes a PyTorch-only fallback insrc_python/nattenso the pipeline runs, but it is slower than the CUDA implementation.
ONNX export
Export a model and config (requires PyTorch):
uv run python src_python/convert_onnx.py --model harmonix-fold0 --output onnx\\harmonix-fold0.onnx --config-json onnx\\harmonix-fold0.json
For the higher-accuracy ensemble, export harmonix-all with a longer dummy
frame length to avoid trace-time shape issues:
uv run python src_python/convert_onnx.py --model harmonix-all --frames 10240 --output onnx\\harmonix-all.onnx --config-json onnx\\harmonix-all.json
Export fold0-7 separately (recommended for stable ensemble):
uv run python src_python/export_onnx_folds.py --out-dir onnx\\folds --frames 10240
ONNX inference (no torch)
This path runs ONNXRuntime only. You must already have demixed stems
(bass.wav, drums.wav, other.wav, vocals.wav) or a precomputed spec.
Stems must be WAV (PCM or IEEE float).
From stems:
uv run python src_python/run_allinone_onnx.py --stems-dir "demix\\htdemucs\\2_23_AM" --onnx-model onnx\\harmonix-fold0.onnx --config-json onnx\\harmonix-fold0.json
Ensemble from multiple fold ONNX files (Python):
uv run python src_python/run_allinone_onnx.py --stems-dir "demix\\htdemucs\\2_23_AM" --onnx-model onnx\\folds\\harmonix-fold0.onnx --onnx-model onnx\\folds\\harmonix-fold1.onnx --onnx-model onnx\\folds\\harmonix-fold2.onnx --onnx-model onnx\\folds\\harmonix-fold3.onnx --onnx-model onnx\\folds\\harmonix-fold4.onnx --onnx-model onnx\\folds\\harmonix-fold5.onnx --onnx-model onnx\\folds\\harmonix-fold6.onnx --onnx-model onnx\\folds\\harmonix-fold7.onnx --config-json onnx\\folds\\harmonix-fold0.json --config-json onnx\\folds\\harmonix-fold1.json --config-json onnx\\folds\\harmonix-fold2.json --config-json onnx\\folds\\harmonix-fold3.json --config-json onnx\\folds\\harmonix-fold4.json --config-json onnx\\folds\\harmonix-fold5.json --config-json onnx\\folds\\harmonix-fold6.json --config-json onnx\\folds\\harmonix-fold7.json
From an existing spec:
uv run python src_python/run_allinone_onnx.py --spec-path "spec_onnx\\2_23_AM.npy" --onnx-model onnx\\harmonix-fold0.onnx --config-json onnx\\harmonix-fold0.json
Deterministic PyTorch demix (disable random shifts):
uv run python src_python/run_allinone_onnx.py --input samples\\2_23_AM.mp3 --demix-backend pytorch --demucs-shifts 0 --demucs-seed 1234 --onnx-model onnx\\harmonix-fold0.onnx --config-json onnx\\harmonix-fold0.json
Force ONNX Runtime to CPU (to match Rust CPU parity):
uv run python src_python/run_allinone_onnx.py --stems-dir "demix\\htdemucs\\2_23_AM" --onnx-model onnx\\harmonix-fold0.onnx --config-json onnx\\harmonix-fold0.json --onnx-provider cpu
Debug dumps for parity checks:
uv run python src_python/run_allinone_onnx.py --stems-dir "demix_onnx\\htdemucs\\2_23_AM" --onnx-model onnx\\harmonix-fold0.onnx --config-json onnx\\harmonix-fold0.json --debug-dir debug_py --debug-preproc
Rust inference
Run the Rust pipeline with ONNX Runtime (Demucs + allin1):
cargo run -- --input samples\\2_23_AM.mp3 --onnx-model onnx\\harmonix-fold0.onnx --config-json onnx\\harmonix-fold0.json --output-dir output_rust
For the ensemble model:
cargo run -- --input samples\\2_23_AM.mp3 --onnx-model onnx\\harmonix-all.onnx --config-json onnx\\harmonix-all.json --output-dir output_rust_all
Ensemble from multiple fold ONNX files:
cargo run -- --input samples\\2_23_AM.mp3 --onnx-model onnx\\folds\\harmonix-fold0.onnx --onnx-model onnx\\folds\\harmonix-fold1.onnx --onnx-model onnx\\folds\\harmonix-fold2.onnx --onnx-model onnx\\folds\\harmonix-fold3.onnx --onnx-model onnx\\folds\\harmonix-fold4.onnx --onnx-model onnx\\folds\\harmonix-fold5.onnx --onnx-model onnx\\folds\\harmonix-fold6.onnx --onnx-model onnx\\folds\\harmonix-fold7.onnx --config-json onnx\\folds\\harmonix-fold0.json --config-json onnx\\folds\\harmonix-fold1.json --config-json onnx\\folds\\harmonix-fold2.json --config-json onnx\\folds\\harmonix-fold3.json --config-json onnx\\folds\\harmonix-fold4.json --config-json onnx\\folds\\harmonix-fold5.json --config-json onnx\\folds\\harmonix-fold6.json --config-json onnx\\folds\\harmonix-fold7.json --output-dir output_rust_all
Debug dumps for parity checks:
cargo run -- --stems-dir demix_onnx\\htdemucs\\2_23_AM --onnx-model onnx\\harmonix-fold0.onnx --config-json onnx\\harmonix-fold0.json --debug-dir debug_rust --debug-preproc
Parity checks
Compare preprocessing + logits between Python ONNX and Rust:
uv run python src_python/compare_parity.py --stems-dir demix_onnx\\htdemucs\\2_23_AM --debug-preproc
Run PyTorch model logits vs ONNX logits (requires allin1/torch):
uv run python src_python/compare_parity.py --stems-dir demix_onnx\\htdemucs\\2_23_AM --run-pytorch-model
Demucs diff visualization
Compare two stem folders (RMS + correlation):
uv run python src_python/demucs_diff.py --a demix_onnx\\htdemucs\\2_23_AM --b debug_py_check\\demix\\htdemucs\\2_23_AM --out-json documents\\demucs_diff.json --out-svg viz\\demucs_diff_2_23_AM.svg
Analysis comparison + sonification
Compare analysis JSONs:
uv run python src_python/compare_analysis.py --a output_compare\\pytorch\\analysis.json --b output_compare\\onnx\\analysis.json --c output_compare\\rust\\analysis.json --out-json documents\\analysis_compare.json
Generate demo audio with tempo/section cues:
uv run python src_python/sonify_analysis.py --input samples\\2_23_AM.mp3 --analysis-json output_compare\\pytorch\\analysis.json --output output_compare\\demos\\pytorch_demo.wav
uv run python src_python/sonify_analysis.py --input samples\\2_23_AM.mp3 --analysis-json output_compare\\onnx\\analysis.json --output output_compare\\demos\\onnx_demo.wav
uv run python src_python/sonify_analysis.py --input samples\\2_23_AM.mp3 --analysis-json output_compare\\rust\\analysis.json --output output_compare\\demos\\rust_demo.wav
Note: To align PyTorch/ONNX/Rust demo timing, ensure the same model and the same stems.
For example, run PyTorch with --model harmonix-fold0 and --demix-dir matching the
stems used by ONNX/Rust, or pass PyTorch-generated stems to ONNX via --stems-dir.
Pipeline mapping (Rust port notes)
allin1.analyze runs these steps; the ONNX path mirrors them with
NumPy-only replacements for Rust-friendly behavior:
- Demix (external):
demucs.separateto produce stems. - Spectrogram:
madmom.audio.signalโsrc_python/native_dsp.py(read_wav_mono,frame_signal)madmom.audio.stftโsrc_python/native_dsp.py(stft,fft_frequencies)madmom.audio.spectrogramfilterbank/log โsrc_python/native_dsp.py(logarithmic_filterbank,log_spectrogram)
- Model inference: ONNX Runtime (
src_python/run_allinone_onnx.py). - Postprocessing:
madmom.features.downbeats.DBNDownBeatTrackingProcessorโsrc_python/native_dbn.pylibrosa.frames_to_timeโ direct hop/sample-rate math insrc_python/run_allinone_onnx.py- peak picking/local maxima โ NumPy in
src_python/run_allinone_onnx.py
Licenses
License references for bundled/derived assets under onnx/.
If you redistribute binaries/models, include the corresponding MIT license text
from upstream.
- allin1 (harmonix models, including
onnx/harmonix-fold0.onnx): MIT, https://github.com/mir-aidj/all-in-one - Demucs models (including
onnx/htdemucs.onnxandonnx/htdemucs.onnx.data): MIT, https://github.com/facebookresearch/demucs - demucs.onnx conversion code (used to export Demucs to ONNX): MIT, https://github.com/sevagh/demucs.onnx
- ONNX Runtime: MIT, https://github.com/microsoft/onnxruntime