Duplicate from arisha07/stem-separator

Co-authored-by: Arisha Kumar <arisha07@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,36 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+onnx/htdemucs_fwd.onnx.data filter=lfs diff=lfs merge=lfs -text

LICENSE.md

@@ -0,0 +1,31 @@
+# License
+
+This model is derived from [facebookresearch/demucs](https://github.com/facebookresearch/demucs),
+licensed under the MIT License. No training or fine-tuning was performed — this repository
+contains an ONNX format conversion of the original pretrained `htdemucs.th` checkpoint.
+
+## MIT License (original demucs project)
+
+Copyright (c) Facebook, Inc. and its affiliates.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+
+> Note: verify this text against the exact `LICENSE` file in the upstream
+> `facebookresearch/demucs` repository before publishing, and update the copyright
+> line if it differs.

README.md

@@ -0,0 +1,82 @@
+---
+license: mit
+tags:
+  - onnx
+  - webnn
+  - audio-source-separation
+  - htdemucs
+pipeline_tag: audio-to-audio
+---
+
+# Stem Separator (HTDemucs v4, fwd-only ONNX export) - WebNN
+
+ONNX export of Meta's pretrained **HTDemucs v4** (`htdemucs.th`, base/non-fine-tuned checkpoint)
+used by the [WebNN Developer Preview](https://github.com/microsoft/webnn-developer-preview) Stem
+Separator demo. Splits a stereo mix into **drums / bass / other / vocals**, running entirely
+client-side via ONNX Runtime Web with WebNN (NPU/GPU) acceleration, cascading to WebGPU/WASM.
+
+## Files
+
+| File | Size | Description |
+|---|---|---|
+| `onnx/htdemucs_fwd.onnx` | ~2.3 MB | Model graph |
+| `onnx/htdemucs_fwd.onnx.data` | ~168 MB | External weights |
+
+## Provenance
+
+- **Source checkpoint**: Meta Research's pretrained `htdemucs.th` base checkpoint (signature
+  `955717e8`, filename `955717e8-8726e21a.th`) from the
+  [facebookresearch/demucs](https://github.com/facebookresearch/demucs) project (MIT license),
+  [checkpoint link](https://dl.fbaipublicfiles.com/demucs/hybrid_transformer/955717e8-8726e21a.th).
+  No training or fine-tuning was performed. This is a format conversion of an already-trained model.
+- **Conversion**: exported via `convert_htdemucs_fwd_only.py`, a script modeled on the HTDemucs
+  conversion utilities from the
+  [RyanMetcalfeInt8/Music-Source-Separation-Training](https://github.com/RyanMetcalfeInt8/Music-Source-Separation-Training)
+  fork, purpose-built to produce a clean, opset-pinned ONNX file for web/WebNN consumption rather
+  than an OpenVINO IR intermediate.
+- **Why fwd-only**: `torch.onnx.export` cannot trace the complex-tensor ops used by HTDemucs's
+  STFT/iSTFT stages. Those stages are reimplemented client-side in JavaScript instead, matching the
+  original PyTorch implementation within floating-point tolerance.
+- **Conversion command**:
+  ```
+  python convert_htdemucs_fwd_only.py \
+    --config configs/config_musdb18_htdemucs.yaml \
+    --checkpoint 955717e8-8726e21a.th \
+    --output htdemucs_fwd.onnx
+  ```
+  The exact script used is included in this repository at
+  [`tools/convert_htdemucs_fwd_only.py`](./tools/convert_htdemucs_fwd_only.py) as reproducibility
+  collateral (adapted from the fork above, MIT License).
+
+## Input / Output shapes
+
+| Tensor | Shape | Description |
+|---|---|---|
+| `x` (input) | `[1, 4, 2048, 336]` | Normalized CAC spectrogram (freq-domain branch) |
+| `xt` (input) | `[1, 2, 343980]` | Normalized stereo waveform (time-domain branch), ~7.8s @ 44.1kHz |
+| `x_out` (output) | `[1, 16, 2048, 336]` | 4 stems × 4-channel CAC spectrogram |
+| `xt_out` (output) | `[1, 8, 343980]` | 4 stems × stereo waveform |
+
+Segments are processed with 50% overlap (`SEGMENT_LENGTH=343980`, `OVERLAP=171990`) and
+reassembled with triangular-weighted overlap-add.
+
+## Validation
+
+Compared against a Python fp32 reference pipeline (sine-wave input, every intermediate stage dumped):
+
+| Stage | Max diff |
+|---|---|
+| STFT | 9.5e-7 |
+| pre_forward | 7.6e-6 |
+| ONNX `fwd` (fp16) | MAE 9e-4 |
+| post_forward | 1.8e-7 |
+
+## Usage
+
+See the [Stem Separator demo](https://github.com/microsoft/webnn-developer-preview/tree/main/demos/stem-separator)
+in the WebNN Developer Preview repository for the full client-side pre/post-processing pipeline.
+
+## License
+
+MIT - see [LICENSE.md](./LICENSE.md). Derived from
+[facebookresearch/demucs](https://github.com/facebookresearch/demucs) (MIT).

onnx/htdemucs_fwd.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:555f511fb653348a1c9495ab6b6091767d1920798f3fba68629a1fe5b135a45f
+size 2385507

onnx/htdemucs_fwd.onnx.data

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:23a7c68669c041363caaef94641e961dc8c4e16b9891cb87f58f0e199a5e8820
+size 168361984

tools/convert_htdemucs_fwd_only.py

@@ -0,0 +1,176 @@
+"""Export ONLY the fwd(x, xt) branch of HTDemucs to ONNX.
+
+This matches what Intel's OpenVINO plugin does: the ONNX contains only the
+transformer/convolution inner network. Pre-processing (STFT, magnitude reshape,
+normalization) and post-processing (iSTFT, mask, de-norm, sum, trim) remain
+outside the ONNX and must be done at inference time.
+
+Usage:
+    python convert_htdemucs_fwd_only.py \
+        --config <repo>/configs/config_musdb18_htdemucs.yaml \
+        --checkpoint 955717e8-8726e21a.th \
+        --output htdemucs_fwd.onnx
+
+Source: modeled on the HTDemucs conversion utilities from the `RyanMetcalfeInt8`
+fork of Music-Source-Separation-Training (https://github.com/RyanMetcalfeInt8/Music-Source-Separation-Training),
+itself forked from https://github.com/ZFTurbo/Music-Source-Separation-Training
+(MIT License, Copyright (c) 2024 Roman Solovyev (ZFTurbo)). Purpose-built to
+produce a clean, opset-pinned ONNX file for web/WebNN consumption. Included here
+as reproducibility collateral for the htdemucs_fwd.onnx model in this repository.
+"""
+import argparse
+import sys
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+import yaml
+from omegaconf import OmegaConf
+
+THIS_DIR = Path(__file__).resolve().parent
+REPO_ROOT = THIS_DIR.parent
+sys.path.insert(0, str(REPO_ROOT))
+
+from models.demucs4ht import HTDemucs  # noqa: E402
+
+
+def patch_demucs_for_export() -> None:
+    from demucs import hdemucs as hd
+    import torch.nn.functional as F
+    import typing as tp
+
+    def pad1d_no_assert(x, paddings: tp.Tuple[int, int], mode: str = 'constant', value: float = 0.):
+        length = x.shape[-1]
+        padding_left, padding_right = paddings
+        if mode == 'reflect':
+            max_pad = max(padding_left, padding_right)
+            if length <= max_pad:
+                extra_pad = max_pad - length + 1
+                extra_pad_right = min(padding_right, extra_pad)
+                extra_pad_left = extra_pad - extra_pad_right
+                paddings = (padding_left - extra_pad_left, padding_right - extra_pad_right)
+                x = F.pad(x, (extra_pad_left, extra_pad_right))
+        return F.pad(x, paddings, mode, value)
+
+    hd.pad1d = pad1d_no_assert
+    import models.demucs4ht as dmht
+    dmht.pad1d = pad1d_no_assert
+
+
+def build_model_from_config(config_path: Path) -> HTDemucs:
+    with open(config_path, 'r', encoding='utf-8') as handle:
+        raw = yaml.safe_load(handle)
+    numeric_fields = ['dconv_init', 'freq_emb', 'emb_scale', 'rescale',
+                      't_dropout', 't_weight_decay', 't_sparsity', 't_hidden_scale',
+                      't_weight_pos_embed']
+    if 'htdemucs' in raw:
+        for key in numeric_fields:
+            if key in raw['htdemucs'] and isinstance(raw['htdemucs'][key], str):
+                try:
+                    raw['htdemucs'][key] = float(raw['htdemucs'][key])
+                except ValueError:
+                    pass
+    args = OmegaConf.create(raw)
+    extra = {
+        'sources': list(args.training.instruments),
+        'audio_channels': args.training.channels,
+        'samplerate': args.training.samplerate,
+        'segment': args.training.segment,
+    }
+    kw = OmegaConf.to_container(args.htdemucs, resolve=True)
+    return HTDemucs(**extra, **kw)
+
+
+def load_checkpoint(model: nn.Module, ckpt_path: Path) -> None:
+    state_dict = torch.load(str(ckpt_path), map_location='cpu', weights_only=False)
+    if isinstance(state_dict, dict) and 'state' in state_dict:
+        state_dict = state_dict['state']
+    missing, unexpected = model.load_state_dict(state_dict, strict=False)
+    if missing:
+        print(f"  [load] missing keys ({len(missing)}): {missing[:5]}")
+    if unexpected:
+        print(f"  [load] unexpected keys ({len(unexpected)}): {unexpected[:5]}")
+
+
+class FwdWrapper(nn.Module):
+    def __init__(self, model: HTDemucs) -> None:
+        super().__init__()
+        self.model = model
+
+    def forward(self, x: torch.Tensor, xt: torch.Tensor):
+        return self.model.fwd(x, xt)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument('--config', type=Path, required=True)
+    parser.add_argument('--checkpoint', type=Path, required=True)
+    parser.add_argument('--output', type=Path, required=True)
+    parser.add_argument('--segment', type=int, default=343980)
+    parser.add_argument('--no-dynamo', action='store_true',
+                        help='Use the legacy (scripting) exporter instead of dynamo')
+    parser.add_argument('--opset', type=int, default=17)
+    args = parser.parse_args()
+
+    print(f"== HTDemucs fwd-only ONNX export ==")
+    print(f"  config:     {args.config}")
+    print(f"  checkpoint: {args.checkpoint}")
+    print(f"  output:     {args.output}")
+    print(f"  segment:    {args.segment}")
+
+    patch_demucs_for_export()
+
+    print("[1/4] Building model + loading checkpoint...")
+    model = build_model_from_config(args.config)
+    model.eval()
+    load_checkpoint(model, args.checkpoint)
+
+    print("[2/4] Running pre_forward on dummy input to get (x, xt) example tensors...")
+    dummy_mix = torch.zeros([1, 2, args.segment], dtype=torch.float32)
+    with torch.inference_mode():
+        x, xt, std, mean, meant, stdt, z, length, length_pre_pad = model.pre_forward(dummy_mix)
+    print(f"  [shapes] x={tuple(x.shape)} dtype={x.dtype}")
+    print(f"  [shapes] xt={tuple(xt.shape)} dtype={xt.dtype}")
+
+    wrapped = FwdWrapper(model)
+    wrapped.eval()
+
+    print("[3/4] Python forward check...")
+    with torch.inference_mode():
+        x_out, xt_out = wrapped(x, xt)
+    print(f"  [shapes] x_out={tuple(x_out.shape)} dtype={x_out.dtype}")
+    print(f"  [shapes] xt_out={tuple(xt_out.shape)} dtype={xt_out.dtype}")
+
+    use_dynamo = not args.no_dynamo
+    print(f"[4/4] Exporting to ONNX (dynamo={use_dynamo}, opset={args.opset})...")
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    with torch.inference_mode():
+        if use_dynamo:
+            torch.onnx.export(
+                wrapped,
+                (x, xt),
+                str(args.output),
+                input_names=['x', 'xt'],
+                output_names=['x_out', 'xt_out'],
+                dynamo=True,
+                report=False,
+                opset_version=args.opset,
+            )
+        else:
+            torch.onnx.export(
+                wrapped,
+                (x, xt),
+                str(args.output),
+                input_names=['x', 'xt'],
+                output_names=['x_out', 'xt_out'],
+                opset_version=args.opset,
+                do_constant_folding=True,
+            )
+    size_mb = args.output.stat().st_size / (1024 * 1024)
+    print(f"  [export] done: {args.output} ({size_mb:.1f} MB)")
+
+    print("DONE.")
+
+
+if __name__ == '__main__':
+    main()