Update README.md

README.md

@@ -11,136 +11,180 @@ tags:
 
 # AudioSeparatorONNX
 
-Collection of popular audio separation and vocal removal models converted to **ONNX** format, optimized for high-performance CPU/GPU inference using `onnxruntime`.
+Popular audio source separation models — UVR5, MDX-Net, VR Architecture, BSRoformer — converted to **ONNX** format and packaged as single self-contained files.
 
-## 📌 Overview
+## Why ONNX?
 
-This repository provides pre-converted ONNX versions of popular music source separation (MSS) and vocal removal models (such as UVR5 / Ultimate Vocal Remover models, MDX-Net, VR Architecture, Demucs, etc.). 
+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:
 
-Using ONNX models allows for:
-- **Cross-platform compatibility** (Python, C++, C#, Rust, JS, etc.).
-- **Faster execution** with `onnxruntime` using CPU, CUDA, DirectML, CoreML, or TensorRT.
-- **Lightweight deployments** without requiring heavy frameworks like PyTorch or TensorFlow in production.
+- 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.
 
 ---
 
-## 💻 Quickstart
+## ⚠️ 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`.
 
-## 📦 Embedded Metadata (`sep_meta`)
+### Roformer models (BSRoformer)
 
-Every `.onnx` file in this repository is **self-contained** — no separate sidecar `.json` file is needed. All inference parameters are embedded directly inside the ONNX model as a `metadata_props` entry with key `sep_meta`.
+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.
 
-### Reading metadata with Python (`onnxruntime`)
+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("Reverb_HQ_By_FoxJoy.onnx", providers=["CPUExecutionProvider"])
-meta_map = sess.get_modelmeta().custom_metadata_map   # dict[str, str]
-sep_meta = json.loads(meta_map["sep_meta"])
+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"])   # "No Reverb"
-print(sep_meta["secondary_stem"]) # "Reverb"
+print(sep_meta["arch"])          # "MDX" | "VR" | "ROFORMER"
+print(sep_meta["primary_stem"])  # e.g. "No Reverb"
 ```
 
-### Reading metadata with Python (`onnx` library, no ORT session)
+### Python — `onnx` library (no inference session)
 
 ```python
 import json
 import onnx
 
-model = onnx.load("UVR-MDX-NET-Inst_HQ_1.onnx")
-meta_map = {p.key: p.value for p in model.metadata_props}
-sep_meta = json.loads(meta_map["sep_meta"])
+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"])
 ```
 
-### Reading metadata with C (fast byte-scan, no ORT dependency)
+### C — byte-scan (no ORT, no Python)
 
-The `sep_meta` key is stored near the **end** of the ONNX protobuf binary, after all weight initializers. You can extract it with a simple byte scan — no need to load the entire model into memory:
+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 JSON string, or NULL on failure.
-   Caller must free() the result. */
-char* read_onnx_sep_meta(const char* onnx_path) {
-    FILE* f = fopen(onnx_path, "rb");
+/*
+ * 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;
 
-    /* Read the last 64 KB — sep_meta is always at the end */
-    const size_t SCAN_SIZE = 64 * 1024;
+    const size_t SCAN_SIZE = 65536;
     fseek(f, 0, SEEK_END);
-    long file_size = ftell(f);
-    size_t read_size = (file_size < (long)SCAN_SIZE) ? (size_t)file_size : SCAN_SIZE;
-    long offset = file_size - (long)read_size;
-    fseek(f, offset, SEEK_SET);
+    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_size + 1);
+    char* buf = (char*)malloc(read_sz + 1);
     if (!buf) { fclose(f); return NULL; }
-    size_t n = fread(buf, 1, read_size, f);
+    size_t n = fread(buf, 1, read_sz, f);
     fclose(f);
     buf[n] = '\0';
 
-    /* Find the marker string "sep_meta" */
-    const char* marker = "sep_meta";
-    char* pos = buf;
+    size_t klen = strlen(key);
     char* found = NULL;
-    while ((pos = (char*)memchr(pos, 's', buf + n - pos)) != NULL) {
-        if ((size_t)(buf + n - pos) >= strlen(marker) &&
-            memcmp(pos, marker, strlen(marker)) == 0) {
-            found = pos;  /* keep the last occurrence */
-        }
-        pos++;
-    }
+    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; }
 
-    /* After "sep_meta" the protobuf stores the value string.
-       Skip past the key and any protobuf length bytes until '{' */
-    char* json_start = found + strlen(marker);
-    while (json_start < buf + n && *json_start != '{') json_start++;
-    if (json_start >= buf + n) { 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;
 
-    /* Find matching closing '}' */
     int depth = 0;
-    char* p = json_start;
     while (p < buf + n) {
-        if (*p == '{') depth++;
-        else if (*p == '}') { depth--; if (depth == 0) break; }
+        if      (*p == '{') depth++;
+        else if (*p == '}') { if (--depth == 0) break; }
         p++;
     }
     if (depth != 0) { free(buf); return NULL; }
 
-    size_t json_len = (size_t)(p - json_start) + 1;
-    char* result = (char*)malloc(json_len + 1);
-    memcpy(result, json_start, json_len);
-    result[json_len] = '\0';
-
+    size_t len = (size_t)(p - start) + 1;
+    char* out  = (char*)malloc(len + 1);
+    memcpy(out, start, len);
+    out[len] = '\0';
     free(buf);
-    return result;
+    return out;
 }
 
-/* Example usage */
 int main(void) {
-    char* meta = read_onnx_sep_meta("Reverb_HQ_By_FoxJoy.onnx");
-    if (meta) {
-        printf("%s\n", meta);
-        free(meta);
-    }
+    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;
 }
 ```
 
 ---
 
-## 🔬 `sep_meta` JSON Schema
+## 🔧 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"]
+```
+
+---
 
-The `sep_meta` value is a compact JSON object. Fields differ by architecture:
+## 🔬 `sep_meta` Reference
 
-### MDX-Net models
+### MDX-Net
 
 ```json
 {
@@ -157,20 +201,9 @@ The `sep_meta` value is a compact JSON object. Fields differ by architecture:
 }
 ```
 
-| Field | Type | Description |
-|-------|------|-------------|
-| `n_fft` | int | STFT window size |
-| `hop_length` | int | STFT hop length |
-| `dim_f` | int | Frequency bins fed to the network |
-| `dim_t` | int | Time frames per chunk |
-| `compensate` | float | Amplitude gain applied after iSTFT |
-| `overlap` | float | Overlap ratio between consecutive chunks (0–0.75) |
+**ONNX I/O** — Input `input (1, 4, dim_f, dim_t)`: `[real_L, imag_L, real_R, imag_R]` · Output `output` same shape
 
-**ONNX I/O:**
-- Input: `input` — shape `(1, 4, dim_f, dim_t)` — four channels: `[real_L, imag_L, real_R, imag_R]`
-- Output: `output` — same shape as input
-
-### VR Architecture models
+### VR Architecture
 
 ```json
 {
@@ -182,217 +215,121 @@ The `sep_meta` value is a compact JSON object. Fields differ by architecture:
   "bins":           672,
   "window_size":    512,
   "is_vr51":        true,
-  "nn_arch_size":   218,
+  "nn_arch_size":   218409,
   "model_capacity": [32, 128],
   "band_params": {
-    "1": {"sr": 11025, "n_fft": 2048, "trim": 0, "window": "hann", "agg": 10},
-    "2": {"sr": 22050, "n_fft": 2048, "trim": 48, "window": "hann", "agg": 10},
-    "3": {"sr": 44100, "n_fft": 2048, "trim": 48, "window": "hann", "agg": 10},
-    "4": {"sr": 44100, "n_fft": 8192, "trim": 0, "window": "hann", "agg": 10}
+    "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}
   }
 }
 ```
 
-| Field | Type | Description |
-|-------|------|-------------|
-| `vr_model_param` | string | Band config preset name |
-| `bins` | int | Total frequency bins after multi-band STFT |
-| `window_size` | int | Window size for the primary STFT pass |
-| `is_vr51` | bool | `true` = CascadedNet (UVR v5.1), `false` = v4 |
-| `nn_arch_size` | int | Network depth parameter |
-| `band_params` | object | Per-band STFT config used to build the input spectrogram |
+**ONNX I/O** — Input `input (1, 2, bins+1, window_size)`: stereo multi-band magnitude · Output `output` same shape (source mask)
 
-**ONNX I/O:**
-- Input: `input` — shape `(batch, 2, bins, frames)` — stereo multi-band magnitude spectrogram
-- Output: `output` — shape `(batch, 2, bins, frames)` — predicted source mask
+### BSRoformer
 
-### Roformer models (MelBandRoformer / BSRoformer)
+The ONNX graph covers `band_split + transformer + mask_estimators`. STFT and iSTFT are handled by the caller.
 
-> **Note:** Roformer ONNX exports contain only the **transformer core**. STFT and iSTFT are not part of the ONNX graph (complex ops are not exportable). The inference wrapper must perform STFT → band-split → ONNX → scatter → iSTFT around the model.
+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":      "MelBandRoformer",
+  "roformer_type":      "BSRoformer",
   "primary_stem":       "No Reverb",
   "secondary_stem":     "Reverb",
   "sample_rate":        44100,
-  "chunk_size":         264600,
-  "hop_length":         1024,
+  "num_channels":       2,
+  "chunk_size":         112455,
+  "native_chunk_size":  352800,
+  "hop_length":         441,
   "n_fft":              2048,
-  "dim_t":              256,
-  "num_bands":          64,
-  "dim":                128,
   "frames":             256,
-  "n_freq_indices":     384,
-  "num_bins":           1025,
-  "n_full_freqs":       1025,
+  "n_freq_indices":     2050,
+  "n_full_freqs":       2050,
   "overlap":            2,
-  "freq_indices":       [...],
-  "num_bands_per_freq": [...]
+  "freq_indices":       [0, 1, 2, "..."],
+  "num_bands_per_freq": [1, 1, 1, "..."]
 }
 ```
 
-| Field | Type | Description |
-|-------|------|-------------|
-| `chunk_size` | int | Audio samples per inference chunk |
-| `hop_length` / `n_fft` | int | STFT params applied **outside** the ONNX graph |
-| `num_bands` | int | Number of mel bands |
-| `frames` | int | Time frames (= `dim_t`) |
-| `n_freq_indices` | int | Number of selected frequency bins after mel gather |
-| `freq_indices` | int[] | Frequency bin indices to gather from full STFT before passing to network |
-| `num_bands_per_freq` | int[] | Band assignment per frequency bin (for scatter-add in iSTFT) |
-| `overlap` | int | Number of overlapping inference passes per chunk |
-
-**ONNX I/O:**
-- Input: `x_bands` — shape `(1, frames, num_bands, dim)` — band-split STFT features
-- Output: `masks` — shape `(1, 1, n_freq_indices, frames, 2)` — complex mask
-
----
-
-## 🚀 Quickstart with `audio-separator`
+| 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 |
 
-The recommended way to run these models is [`audio-separator`](https://github.com/nomadkaraoke/python-audio-separator) — the same inference engine that powers [Ultimate Vocal Remover](https://github.com/Anjok07/ultimatevocalremovergui).
+**ONNX I/O** — Input `x_flat (1, frames, n_freq_indices*2)` · Output `masks (1, 1, n_freq_indices, frames, 2)`
 
-All `.onnx` and `.pth` files in this repo are fully compatible with `audio-separator` out of the box.
+---
 
-### Installation
+## 🚀 Quickstart — MDX and VR with `audio-separator`
 
 ```bash
-# CPU only (macOS / Linux / Windows)
-pip install "audio-separator[cpu]"
-
-# Nvidia GPU (CUDA)
-pip install "audio-separator[gpu]"
-
-# Apple Silicon (CoreML acceleration)
-pip install "audio-separator[cpu]"   # CoreML is auto-detected on macOS
+pip install "audio-separator[cpu]"   # CPU / Apple Silicon
+pip install "audio-separator[gpu]"   # Nvidia CUDA
 ```
 
-### CLI usage
-
-Download the model file from the **Files and versions** tab, then point `--model_file_dir` at the folder containing it:
-
 ```bash
-# Remove reverb from a mix
-audio-separator mix.wav \
-  --model_filename Reverb_HQ_By_FoxJoy.onnx \
-  --model_file_dir /path/to/downloaded/models \
-  --output_dir ./output \
-  --output_format WAV
-
-# Vocal / instrumental separation
+# CLI
 audio-separator mix.wav \
-  --model_filename UVR-MDX-NET-Inst_HQ_3.onnx \
-  --model_file_dir /path/to/downloaded/models \
+  --model_filename UVR-DeEcho-DeReverb.onnx \
+  --model_file_dir /path/to/models \
   --output_dir ./output
 ```
 
-### Python API
-
 ```python
+# Python API
 from audio_separator.separator import Separator
 
-separator = Separator(
-    model_file_dir="/path/to/downloaded/models",  # folder with the .onnx files
-    output_dir="./output",
-    output_format="WAV",
-)
-
-# Load and run — arch is auto-detected from the file extension
-separator.load_model("Reverb_HQ_By_FoxJoy.onnx")
-output_files = separator.separate("mix.wav")
-print("Output:", output_files)
+sep = Separator(model_file_dir="/path/to/models", output_dir="./output")
+sep.load_model("UVR-DeEcho-DeReverb.onnx")
+sep.separate("mix.wav")
 ```
 
-### Adjusting inference parameters
-
-`audio-separator` exposes the most important per-arch parameters:
+If hash auto-detection fails (see compatibility note above), pass the parameters manually:
 
 ```python
-# MDX — tune overlap and segment size
-separator = Separator(
-    model_file_dir="/path/to/downloaded/models",
-    mdx_params={
-        "hop_length":    1024,
-        "segment_size":  256,    # dim_t from sep_meta
-        "overlap":       0.25,   # matches sep_meta default
-        "batch_size":    1,
-        "enable_denoise": False,
-    }
+sep = Separator(
+    model_file_dir="/path/to/models",
+    mdx_params={"hop_length": 1024, "segment_size": 256, "overlap": 0.25},
 )
-
-# VR — tune aggression and window size
-separator = Separator(
-    model_file_dir="/path/to/downloaded/models",
-    vr_params={
-        "batch_size":    1,
-        "window_size":   512,    # 320 = slower but better
-        "aggression":    5,      # 0-100, higher = more aggressive separation
-        "enable_tta":    False,
-        "high_end_process": False,
-    }
-)
-
-separator.load_model("UVR-DeEcho-DeReverb.onnx")
-output_files = separator.separate("mix.wav")
-```
-
-> **Note on `sep_meta`:** `audio-separator` reads model parameters from its own internal `mdx_model_data.json` lookup table — it does **not** read the embedded `sep_meta` metadata. The `sep_meta` field is provided for custom inference pipelines and other runtimes (see sections below). embed_sep_meta(onnx_path: str, meta: dict) -> None:
-```python
-    model = onnx.load(onnx_path)
-    # Remove stale entry if present
-    existing = [p for p in model.metadata_props if p.key != "sep_meta"]
-    del model.metadata_props[:]
-    model.metadata_props.extend(existing)
-    # Add new entry
-    entry = model.metadata_props.add()
-    entry.key   = "sep_meta"
-    entry.value = json.dumps(meta, separators=(",", ":"))
-    onnx.save_model(model, onnx_path)
+sep.load_model("UVR-MDX-NET-Inst_HQ_5.onnx")
 ```
-# Example — MDX model
 
-```python
-embed_sep_meta("my_model.onnx", {
-    "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,
-})
-```
 ---
 
 ## 📋 Requirements
 
 ```
 onnxruntime >= 1.16
-numpy >= 1.24
-soundfile >= 0.12
 ```
 
-For metadata-only access (no inference):
+For reading metadata without running inference:
 ```
 onnx >= 1.14
 ```
 
+---
+
 ## 📄 License
 
-The repository structure and conversion scripts are licensed under the [MIT License](https://www.google.com/search?q=LICENSE).
-*Please check the original model licenses before using them for commercial purposes.*
+[MIT License](LICENSE). Check individual model licenses before commercial use.
 
 ---
 
-## 🙏 Acknowledgments & Credits
-
-Special thanks to the authors and maintainers of:
+## 🙏 Acknowledgments
 
-* [Ultimate Vocal Remover (UVR5)](https://www.google.com/search?q=https://github.com/Anemone95/Ultimate-Vocal-Remover-GUI)
-* [audio-separator](https://www.google.com/search?q=https://github.com/beverlis/audio-separator)
-* [ONNX Runtime](https://onnxruntime.ai/)
+- [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)