AudioSeparatorONNX / README.md
raydotac's picture
Update README.md
2aba44b verified
|
Raw
History Blame Contribute Delete
11.1 kB
---
license: mit
pipeline_tag: audio-to-audio
tags:
- audio-separation
- vocal-remover
- stem-separation
- onnx
- onnxruntime
---
# AudioSeparatorONNX
Popular audio source separation models β€” UVR5, MDX-Net, VR Architecture, BSRoformer β€” converted to **ONNX** format and packaged as single self-contained files.
## Why ONNX?
The standard way to run these models is through [audio-separator](https://github.com/nomadkaraoke/python-audio-separator) or [UVR](https://github.com/Anjok07/ultimatevocalremovergui), both of which require Python and PyTorch. That's fine for desktop use, but becomes a problem when you want to:
- Ship a **native application** (C++, Swift, Rust, .NET) without a Python runtime
- Run separation on a **mobile or embedded device**
- Build a **server-side pipeline** where spinning up PyTorch per request is too heavy
- Use **CoreML, TensorRT, DirectML, or other hardware accelerators** via ONNX Runtime
- Load a model in any language that has an ONNX Runtime binding
ONNX Runtime handles all of the above with a single lightweight library and no Python dependency. The models in this repo are ready to drop into any ORT-based pipeline.
## What makes this repo different
Most ONNX model repos ship the file and nothing else. Every model here includes two metadata blobs embedded directly inside the `.onnx` file:
| Key | What it contains |
|-----|-----------------|
| `sep_meta` | All inference parameters: arch, stems, STFT config, chunk size, overlap, and for Roformer β€” `freq_indices` and `num_bands_per_freq` arrays needed for the gather/scatter steps |
| `model_config` | The original training config reconstructed from the model weights β€” lets you recover a YAML for audio-separator or any other PyTorch pipeline without needing the original sidecar file |
No JSON files, no YAML sidecars, no download_checks.json. One file per model.
---
## ⚠️ Compatibility Notes
### MDX and VR models
These work with audio-separator and UVR out of the box. Just point `model_file_dir` at the folder containing the `.onnx` files.
> **Hash detection:** audio-separator identifies models by MD5 of the last 10 MB of the file. Because `sep_meta` and `model_config` are appended at the end, the hash of these files differs from the originals in the UVR database. If auto-detection fails, pass the parameters explicitly via `mdx_params` or `vr_params` β€” all values are available in `sep_meta`.
### Roformer models (BSRoformer)
Roformer `.onnx` files are intended for **ONNX Runtime inference only** β€” audio-separator and UVR run Roformer via PyTorch from the original `.ckpt`, not from ONNX. These files are useful if you are building a custom native pipeline.
The embedded `model_config` key lets you recover the training configuration without the original `.yaml` sidecar β€” see the extraction section below.
---
## πŸ“¦ Reading Embedded Metadata
Both keys are stored as compact JSON strings inside the ONNX `metadata_props` field.
### Python β€” `onnxruntime`
```python
import json
import onnxruntime as ort
sess = ort.InferenceSession("UVR-DeEcho-DeReverb.onnx", providers=["CPUExecutionProvider"])
meta = sess.get_modelmeta().custom_metadata_map # dict[str, str]
sep_meta = json.loads(meta["sep_meta"])
model_config = json.loads(meta["model_config"])
print(sep_meta["arch"]) # "MDX" | "VR" | "ROFORMER"
print(sep_meta["primary_stem"]) # e.g. "No Reverb"
```
### Python β€” `onnx` library (no inference session)
```python
import json
import onnx
model = onnx.load("UVR-DeEcho-DeReverb.onnx")
meta = {p.key: p.value for p in model.metadata_props}
sep_meta = json.loads(meta["sep_meta"])
model_config = json.loads(meta["model_config"])
```
### C β€” byte-scan (no ORT, no Python)
Both keys are stored near the **end** of the ONNX protobuf, after all weight tensors. You can extract either by scanning the last 64 KB without loading any weights:
```c
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
/*
* Returns a heap-allocated null-terminated JSON string for the given key,
* or NULL if not found. Caller must free() the result.
*
* Increase SCAN_SIZE to 524288 for large Roformer models whose
* freq_indices array may exceed 64 KB.
*/
char* read_onnx_meta_key(const char* path, const char* key) {
FILE* f = fopen(path, "rb");
if (!f) return NULL;
const size_t SCAN_SIZE = 65536;
fseek(f, 0, SEEK_END);
long sz = ftell(f);
size_t read_sz = (sz < (long)SCAN_SIZE) ? (size_t)sz : SCAN_SIZE;
fseek(f, sz - (long)read_sz, SEEK_SET);
char* buf = (char*)malloc(read_sz + 1);
if (!buf) { fclose(f); return NULL; }
size_t n = fread(buf, 1, read_sz, f);
fclose(f);
buf[n] = '\0';
size_t klen = strlen(key);
char* found = NULL;
for (size_t i = 0; i + klen < n; i++)
if (memcmp(buf + i, key, klen) == 0) found = buf + i;
if (!found) { free(buf); return NULL; }
char* p = found + klen;
while (p < buf + n && *p != '{') p++;
if (p >= buf + n) { free(buf); return NULL; }
char* start = p;
int depth = 0;
while (p < buf + n) {
if (*p == '{') depth++;
else if (*p == '}') { if (--depth == 0) break; }
p++;
}
if (depth != 0) { free(buf); return NULL; }
size_t len = (size_t)(p - start) + 1;
char* out = (char*)malloc(len + 1);
memcpy(out, start, len);
out[len] = '\0';
free(buf);
return out;
}
int main(void) {
char* sep = read_onnx_meta_key("model.onnx", "sep_meta");
char* cfg = read_onnx_meta_key("model.onnx", "model_config");
if (sep) { printf("sep_meta: %s\n", sep); free(sep); }
if (cfg) { printf("model_config: %s\n", cfg); free(cfg); }
return 0;
}
```
---
## πŸ”§ Using `model_config` in your ONNX pipeline
The `model_config` key contains the original training configuration. If you are building a custom ONNX Runtime pipeline around Roformer inference, you can read it directly at runtime instead of shipping a separate YAML:
```python
import json
import onnxruntime as ort
sess = ort.InferenceSession("deverb_bs_roformer_8_384dim_10depth.onnx",
providers=["CPUExecutionProvider"])
meta = sess.get_modelmeta().custom_metadata_map
sep_meta = json.loads(meta["sep_meta"]) # chunk sizes, freq_indices, etc.
model_config = json.loads(meta["model_config"]) # arch params: dim, depth, n_fft, ...
# Everything you need to run the pipeline is in these two dicts.
# No separate YAML or JSON sidecar required.
n_fft = sep_meta["n_fft"]
hop_length = sep_meta["hop_length"]
chunk_size = sep_meta["chunk_size"]
overlap = sep_meta["overlap"]
```
---
## πŸ”¬ `sep_meta` Reference
### MDX-Net
```json
{
"arch": "MDX",
"primary_stem": "Vocals",
"secondary_stem": "Instrumental",
"sample_rate": 44100,
"n_fft": 7680,
"hop_length": 1024,
"dim_f": 3072,
"dim_t": 256,
"compensate": 1.021,
"overlap": 0.25
}
```
**ONNX I/O** β€” Input `input (1, 4, dim_f, dim_t)`: `[real_L, imag_L, real_R, imag_R]` Β· Output `output` same shape
### VR Architecture
```json
{
"arch": "VR",
"primary_stem": "No Reverb",
"secondary_stem": "Reverb",
"sample_rate": 44100,
"vr_model_param": "4band_v3",
"bins": 672,
"window_size": 512,
"is_vr51": true,
"nn_arch_size": 218409,
"model_capacity": [32, 128],
"band_params": {
"1": {"sr": 11025, "hl": 480, "n_fft": 960, "crop_start": 0, "crop_stop": 245},
"2": {"sr": 22050, "hl": 480, "n_fft": 1920, "crop_start": 245, "crop_stop": 432},
"3": {"sr": 44100, "hl": 480, "n_fft": 3840, "crop_start": 432, "crop_stop": 567},
"4": {"sr": 44100, "hl": 960, "n_fft": 7680, "crop_start": 567, "crop_stop": 673}
}
}
```
**ONNX I/O** β€” Input `input (1, 2, bins+1, window_size)`: stereo multi-band magnitude Β· Output `output` same shape (source mask)
### BSRoformer
The ONNX graph covers `band_split + transformer + mask_estimators`. STFT and iSTFT are handled by the caller.
Pipeline:
1. STFT per channel β†’ interleave channels β†’ `stft_repr` shape `(n_full_freqs, T, 2)`
2. Gather `freq_indices` from `stft_repr` β†’ flatten β†’ `x_flat` shape `(1, frames, n_freq_indices*2)` ← **ONNX input**
3. ONNX forward β†’ `masks` shape `(1, 1, n_freq_indices, frames, 2)`
4. `scatter_add` masks back to `stft_repr` positions, divide by `num_bands_per_freq` β†’ `masks_avg`
5. Multiply `stft_repr * masks_avg` β†’ iSTFT per channel β†’ audio
```json
{
"arch": "ROFORMER",
"roformer_type": "BSRoformer",
"primary_stem": "No Reverb",
"secondary_stem": "Reverb",
"sample_rate": 44100,
"num_channels": 2,
"chunk_size": 112455,
"native_chunk_size": 352800,
"hop_length": 441,
"n_fft": 2048,
"frames": 256,
"n_freq_indices": 2050,
"n_full_freqs": 2050,
"overlap": 2,
"freq_indices": [0, 1, 2, "..."],
"num_bands_per_freq": [1, 1, 1, "..."]
}
```
| Field | Description |
|-------|-------------|
| `chunk_size` | Samples per inference chunk (export size β€” smaller to reduce RAM during export) |
| `native_chunk_size` | Original training chunk size β€” use for best quality if RAM allows |
| `frames` | STFT frame count β€” `x_flat` must have exactly this many time frames |
| `freq_indices` | Indices into `stft_repr` to gather before the ONNX forward pass |
| `num_bands_per_freq` | How many bands cover each frequency bin β€” scatter normalization denominator |
**ONNX I/O** β€” Input `x_flat (1, frames, n_freq_indices*2)` Β· Output `masks (1, 1, n_freq_indices, frames, 2)`
---
## πŸš€ Quickstart β€” MDX and VR with `audio-separator`
```bash
pip install "audio-separator[cpu]" # CPU / Apple Silicon
pip install "audio-separator[gpu]" # Nvidia CUDA
```
```bash
# CLI
audio-separator mix.wav \
--model_filename UVR-DeEcho-DeReverb.onnx \
--model_file_dir /path/to/models \
--output_dir ./output
```
```python
# Python API
from audio_separator.separator import Separator
sep = Separator(model_file_dir="/path/to/models", output_dir="./output")
sep.load_model("UVR-DeEcho-DeReverb.onnx")
sep.separate("mix.wav")
```
If hash auto-detection fails (see compatibility note above), pass the parameters manually:
```python
sep = Separator(
model_file_dir="/path/to/models",
mdx_params={"hop_length": 1024, "segment_size": 256, "overlap": 0.25},
)
sep.load_model("UVR-MDX-NET-Inst_HQ_5.onnx")
```
---
## πŸ“‹ Requirements
```
onnxruntime >= 1.16
```
For reading metadata without running inference:
```
onnx >= 1.14
```
---
## πŸ“„ License
[MIT License](LICENSE). Check individual model licenses before commercial use.
---
## πŸ™ Acknowledgments
- [Ultimate Vocal Remover (UVR5)](https://github.com/Anjok07/ultimatevocalremovergui)
- [audio-separator](https://github.com/nomadkaraoke/python-audio-separator)
- [ONNX Runtime](https://onnxruntime.ai/)
- [BS-RoFormer](https://github.com/lucidrains/BS-RoFormer)