@audio/neural-separate 0.1.0 → 0.4.0

package/scripts/export-scnet.py

@@ -0,0 +1,274 @@
+#!/usr/bin/env python3
+"""Export SCNet (Tong et al., "SCNet: Sparse Compression Network for Music Source
+Separation", ICASSP 2024) to the ONNX contract @audio/neural-separate expects for
+modelType 'complex': the network without its STFT and iSTFT, which JS computes.
+
+    pip install torch onnx onnxruntime pyyaml
+    python3 export-scnet.py --model scnet-large --verify
+    python3 export-scnet.py --model scnet --verify
+
+The models are SCNet-large and SCNet as their author trained them on MUSDB18-HQ
+(starrytong/SCNet, MIT; the author confirms the weights are MIT too and may be
+redistributed converted, ONNX included: github.com/starrytong/SCNet/issues/35), in
+the layout of ZFTurbo/Music-Source-Separation-Training (MIT), whose releases host
+them with their configs (v1.0.9 SCNet-large_starrytong_fixed.ckpt, 41.2 M
+parameters; v.1.0.6 scnet_checkpoint_musdb18.ckpt, 10.1 M); the model code is that
+repository's models/scnet, fetched at a pinned commit. MUSDB18 test, its own
+measure: SCNet-large vocals 10.94, drums 11.15, bass 9.38, other 7.31 dB SDR; SCNet
+9.90, 10.44, 8.89, 6.89.
+
+One 11 s segment (485100 samples at 44.1 kHz, padded as SCNet.forward pads it to
+486400, 476 frames) per run:
+  input    mix_spec    [1, 4, 2049, 476]     STFT, complex as channels (L re, L im, R re, R im)
+  output   stems_spec  [1, 16, 2049, 476]    per source (drums, bass, other, vocals) and channel, re and im
+The STFT is SCNet's own: n_fft 4096, hop 1024, no window (torch.stft's default,
+ones), normalized (1/sqrt(4096)), centered with reflect padding.
+
+Writes <out-dir>/<model>.onnx (scnet-large 169 MB, scnet 43 MB) and, with --fp16,
+<model>.fp16.onnx (half that: float16 weights, float32 compute, scripts/compact.py's), weights embedded. Default out-dir:
+$AUDIO_NEURAL_CACHE/<model>, else ~/.cache/audiojs/neural/<model>. --verify compares the stems rebuilt from the ONNX
+output with SCNet.forward on the same segment, and writes the reference
+separation test.js compares the JS pipeline against (scripts/reference.py),
+chunked as scripts/reference.py's `chunked` does. Last run: torch 2.14,
+onnxruntime 1.30.
+"""
+import argparse
+import gc
+import math
+import os
+import sys
+import urllib.request
+from pathlib import Path
+
+REPO = "https://raw.githubusercontent.com/ZFTurbo/Music-Source-Separation-Training/84b1eac0887756b4f1a9d7a1ff49105939749ed2/"
+RELEASE = "https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/"
+MODELS = {  # release, checkpoint, config
+    "scnet-large": ("v1.0.9", "SCNet-large_starrytong_fixed.ckpt", "config_musdb18_scnet_large_starrytong.yaml"),
+    "scnet": ("v.1.0.6", "scnet_checkpoint_musdb18.ckpt", "config_musdb18_scnet.yaml"),
+}
+SEGMENT = 485100
+
+
+def fetch(url, path):
+    if not path.exists():
+        print(f"fetching {url}")
+        urllib.request.urlretrieve(url, str(path) + ".part")
+        os.replace(str(path) + ".part", path)
+    return path
+
+
+def load(src, name="scnet-large"):
+    """SCNet or SCNet-large with its weights: the model code and checkpoint fetched into `src` once."""
+    import torch
+    import yaml
+
+    pkg = src / "scnet_code" / "scnet"
+    pkg.mkdir(parents=True, exist_ok=True)
+    (pkg / "__init__.py").touch()
+    for f in ("scnet.py", "separation.py"):
+        fetch(REPO + "models/scnet/" + f, pkg / f)
+    sys.path.insert(0, str(src / "scnet_code"))
+    sys.dont_write_bytecode = True
+    from scnet.scnet import SCNet
+
+    class Loader(yaml.SafeLoader): pass  # the config's augmentation settings hold a !!python/tuple
+    Loader.add_constructor("tag:yaml.org,2002:python/tuple", lambda l, n: tuple(l.construct_sequence(n)))
+    release, ckpt, config = MODELS[name]
+    cfg = yaml.load(fetch(f"{RELEASE}{release}/{config}", src / config).read_text(), Loader=Loader)
+    model = SCNet(**cfg["model"])
+    model.load_state_dict(torch.load(fetch(f"{RELEASE}{release}/{ckpt}", src / ckpt), map_location="cpu", weights_only=True))
+    return model.eval()
+
+
+def stft(model, x):
+    """SCNet.forward's padding and STFT: (B, C, L) -> (B, 2C, F, T), and the padding."""
+    import torch
+    import torch.nn.functional as F
+
+    hop = model.hop_length
+    pad = hop - x.shape[-1] % hop
+    if (x.shape[-1] + pad) // hop % 2 == 0:
+        pad += hop
+    x = F.pad(x, (0, pad))
+    B, C, L = x.shape
+    z = torch.view_as_real(torch.stft(x.reshape(-1, L), **model.stft_config, return_complex=True))
+    return z.permute(0, 3, 1, 2).reshape(B, 2 * C, z.shape[1], z.shape[2]), pad
+
+
+def istft(model, y, pad):
+    """SCNet.forward's tail: (B, 2·S·C, F, T) -> (B, S, C, L)."""
+    import torch
+
+    B, _, Fr, T = y.shape
+    y = torch.view_as_complex(y.reshape(-1, 2, Fr, T).permute(0, 2, 3, 1).contiguous())
+    x = torch.istft(y, **model.stft_config)
+    return x.reshape(B, len(model.sources), model.audio_channels, -1)[..., :-pad]
+
+
+def dft(x, inverse):
+    """FeatureConversion's rfft and irfft over time (norm "ortho") as products with cos and sin matrices: no ONNX
+    exporter of torch 2.14 takes aten::fft_rfft, and a MatMul runs on every backend. x: (B, C, F, T), T even."""
+    import torch
+
+    T = x.shape[3] if not inverse else 2 * (x.shape[3] - 1)
+    t, k = torch.arange(T, dtype=torch.float64), torch.arange(T // 2 + 1, dtype=torch.float64)
+    a = 2 * math.pi * k[:, None] * t[None] / T                     # (K, T)
+    if not inverse:
+        c, s = (torch.cos(a).T / math.sqrt(T)).to(x.dtype), (-torch.sin(a).T / math.sqrt(T)).to(x.dtype)
+        return torch.cat([x @ c, x @ s], dim=1)                     # (re, im) on channels, as torch.cat([x.real, x.imag], 1)
+    w = torch.full((T // 2 + 1, 1), 2.0, dtype=torch.float64)
+    w[0] = w[-1] = 1                                                # irfft: DC and Nyquist once, their imaginary parts unread
+    c, s = (w * torch.cos(a) / math.sqrt(T)).to(x.dtype), (-w * torch.sin(a) / math.sqrt(T)).to(x.dtype)
+    h = x.shape[1] // 2
+    return x[:, :h] @ c + x[:, h:] @ s
+
+
+def group_norm(mod):
+    """GroupNorm(1, C) with its mean and variance reduced one axis at a time, the last first. Exported as it is,
+    onnxruntime 1.30 sums a whole (C, F, T) group in float32 (10 million values in the dual-path layers): the stems
+    came out 52 to 55 dB from torch's; reduced by axis, 110 dB, as torch's own float32 against float64."""
+    import torch
+
+    def mean(x):
+        for d in range(x.dim() - 1, 0, -1):
+            x = x.mean(d, keepdim=True)
+        return x
+
+    def forward(x):
+        mu = mean(x)
+        var = mean((x - mu) ** 2)
+        shape = (1, -1) + (1,) * (x.dim() - 2)
+        return (x - mu) / torch.sqrt(var + mod.eps) * mod.weight.view(shape) + mod.bias.view(shape)
+
+    return forward
+
+
+class Core:
+    """SCNet.forward between its STFT and its iSTFT."""
+
+    @staticmethod
+    def build(model):
+        import torch.nn as nn
+        from collections import deque
+
+        for f in model.separation_net.feature_conversion:
+            f.forward = (lambda inv: lambda x: dft(x, inv))(f.inverse)
+        for m in model.modules():
+            if isinstance(m, nn.GroupNorm) and m.num_groups == 1:
+                m.forward = group_norm(m)
+
+        class C(nn.Module):
+            def __init__(self, m):
+                super().__init__()
+                self.m = m
+
+            def forward(self, x):
+                m, skips = self.m, deque()
+                for sd in m.encoder:
+                    x, skip, lengths, orig = sd(x)
+                    skips.append((skip, lengths, orig))
+                x = m.separation_net(x)
+                for fusion, su in m.decoder:
+                    skip, lengths, orig = skips.pop()
+                    x = su(fusion(x, skip), lengths, orig)
+                return x
+
+        return C(model).eval()
+
+
+def export(model, path, opset):
+    import torch
+
+    spec, _ = stft(model, torch.randn(1, model.audio_channels, SEGMENT))
+    with torch.no_grad():
+        torch.onnx.export(Core.build(model), (spec,), str(path), input_names=["mix_spec"], output_names=["stems_spec"],
+                          opset_version=opset, dynamo=False)
+    print(f"wrote {path} ({path.stat().st_size / 1e6:.1f} MB), segment {SEGMENT} samples, spec {list(spec.shape)}")
+
+
+def fp16(path, out):
+    """float16 weights, float32 compute (scripts/compact.py --as fp16)."""
+    sys.path.insert(0, str(Path(__file__).parent))
+    import compact
+
+    compact.compact(path, out, "fp16")
+    print(f"wrote {out} ({out.stat().st_size / 1e6:.1f} MB)")
+
+
+def session(path):
+    import onnxruntime as ort
+
+    o = ort.SessionOptions()
+    o.enable_cpu_mem_arena = False
+    o.enable_mem_pattern = False
+    return ort.InferenceSession(str(path), o, providers=["CPUExecutionProvider"])
+
+
+def verify(model, path, tol):
+    """Stems rebuilt from onnxruntime's output against SCNet.forward: noise and tones."""
+    import numpy as np
+    import torch
+
+    s = session(path)
+    t = torch.arange(SEGMENT) / 44100
+    tone = sum(torch.sin(2 * math.pi * f * t) / k for k, f in enumerate((110, 220, 440, 880, 1760), 1))
+    for name, mix in (("noise", 0.1 * torch.randn(1, 2, SEGMENT)), ("tones", 0.2 * torch.stack([tone, 0.8 * tone])[None])):
+        with torch.no_grad():
+            want = model(mix).numpy()
+            spec, pad = stft(model, mix)
+            got = istft(model, torch.from_numpy(s.run(None, {"mix_spec": spec.numpy()})[0]), pad).numpy()
+        rel = np.abs(got - want).max() / max(np.abs(want).max(), 1e-12)
+        print(f"  verify {path.name} {name}: max|diff|/max|y| = {rel:.2e} " + ("OK" if rel < tol else "FAIL"))
+        if rel >= tol:
+            raise SystemExit(f"{path}: verification failed ({rel:.2e} >= {tol})")
+    del s
+    gc.collect()
+
+
+def reference(model, out_dir):
+    """The reference separation: SCNet.forward on segments as scripts/reference.py's normalized() and chunked() cut
+    them, over a 20 s mix (longer than two segments less a step: its ends reflected out)."""
+    import torch
+
+    sys.path.insert(0, str(Path(__file__).parent))
+    import reference as ref
+
+    def run(seg):
+        with torch.no_grad():
+            return model(torch.from_numpy(seg)[None])[0].numpy()
+
+    ref.write(out_dir, lambda mix: dict(zip(model.sources, ref.normalized(mix, lambda x: ref.chunked(x, run, SEGMENT)))), 20.0)
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--model", default="scnet-large", choices=list(MODELS))
+    ap.add_argument("--out-dir", default=None, help="default: $AUDIO_NEURAL_CACHE/<model> or ~/.cache/audiojs/neural/<model>")
+    ap.add_argument("--opset", type=int, default=17)
+    ap.add_argument("--fp16", action="store_true", help="also write <model>.fp16.onnx (float16 weights, float32 compute)")
+    ap.add_argument("--verify", action="store_true", help="compare with SCNet.forward; write the reference separation")
+    args = ap.parse_args()
+
+    cache = Path(os.environ.get("AUDIO_NEURAL_CACHE") or Path.home() / ".cache" / "audiojs" / "neural")
+    out_dir = Path(args.out_dir) if args.out_dir else cache / args.model
+    out_dir.mkdir(parents=True, exist_ok=True)
+    src = out_dir / "src"
+    src.mkdir(exist_ok=True)
+    try:
+        model = load(src, args.model)
+    except ImportError as e:
+        raise SystemExit(f"torch and pyyaml are required: pip install torch onnx onnxruntime pyyaml ({e})")
+    path = out_dir / f"{args.model}.onnx"
+    export(model, path, args.opset)
+    if args.verify:
+        verify(model, path, 1e-4)
+        reference(model, out_dir)
+    if args.fp16:
+        fp16(path, out_dir / f"{args.model}.fp16.onnx")
+        if args.verify:
+            verify(model, out_dir / f"{args.model}.fp16.onnx", 2e-2)
+    print(f"done: {args.model} in {out_dir}")
+
+
+if __name__ == "__main__":
+    sys.exit(main())

package/scripts/export-tiger.py

@@ -0,0 +1,304 @@
+#!/usr/bin/env python3
+"""Export TIGER's Divide and Remaster model (Xu, Li, Chen, Hu, "TIGER: Time-frequency Interleaved Gain Extraction and
+Reconstruction for Efficient Speech Separation", ICLR 2025) to the ONNX contract @audio/neural-separate expects for
+modelType 'complex': the networks between their STFT and iSTFT, which JS computes.
+
+    pip install torch onnx onnxruntime safetensors
+    python3 export-tiger.py --verify
+
+TIGERDNR is three band-split TIGER models of 1.4 M parameters, each separating three sources, of which it keeps one:
+the dialogue of the first, the effects of the second, the music of the third. The weights are the authors'
+(huggingface.co/JusperLee/TIGER-DnR, whose model card states license: apache-2.0; the code, JusperLee/TIGER, is MIT),
+fetched at a pinned revision with the model code at a pinned commit. Its paper, DnR test (v1), SI-SDR: music 7.4,
+speech 15.5, effects 6.5 dB.
+
+One mono 12 s segment (529200 samples at 44.1 kHz, 1034 frames) per run, the three models in one graph:
+  input    mix_spec    [1, 2, 1025, 1034]    STFT, re and im as channels
+  output   stems_spec  [1, 6, 1025, 1034]    dialogue, music, effects, each re and im
+The STFT is TIGER's own: n_fft 2048, hop 512, Hann, unnormalized, centered with reflect padding. Three changes to the
+graph, none to its function: F.adaptive_avg_pool1d (the exporter takes only output sizes that divide the input's)
+becomes a cumulative sum read at torch's windows, UConvBlock's global sum starts at its first term instead of a
+zeros tensor (exported, a constant as large as the features), and each GroupNorm(1, C) reduces its statistics one
+axis at a time (export-scnet.py's; exported whole, each was an InstanceNormalization with constants as large as its
+batch).
+
+Writes <out-dir>/tiger.onnx (29 MB), weights embedded. Default out-dir: $AUDIO_NEURAL_CACHE/tiger, else
+~/.cache/audiojs/neural/tiger. --verify compares the graph with the three TIGER.forward on one segment (in PyTorch on a
+CPU, minutes), and writes the reference separation test.js compares the JS pipeline against: upstream's
+TIGERDNR.wav_chunk_inference (12 s segments every 4 s, unweighted, zero-padded) over reference.py's 20 s mix, each
+channel apart, its segments run through the verified graph. Last run: torch 2.14, onnxruntime 1.30.
+"""
+import argparse
+import os
+import sys
+import types
+import importlib.util
+import json
+import urllib.request
+from pathlib import Path
+
+COMMIT = "9f18d4a10a7137e1ce8052cfb62215179f1287b6"
+REPO = f"https://raw.githubusercontent.com/JusperLee/TIGER/{COMMIT}/look2hear/"
+REVISION = "b7a59560bbca10febbcd46fb01600f868e587f57"
+HF = f"https://huggingface.co/JusperLee/TIGER-DnR/resolve/{REVISION}/"
+RATE, N_FFT, HOP, SEGMENT = 44100, 2048, 512, 529200
+# TIGERDNR.forward: each model's kept source, in the order the graph gives them
+KEEP = (("dialog", 2), ("music", 0), ("effect", 1))
+
+
+def fetch(url, path):
+    if not path.exists():
+        print(f"fetching {url}")
+        urllib.request.urlretrieve(url, str(path) + ".part")
+        os.replace(str(path) + ".part", path)
+    return path
+
+
+def module(name, path):
+    s = importlib.util.spec_from_file_location(name, path)
+    m = importlib.util.module_from_spec(s)
+    sys.modules[name] = m
+    s.loader.exec_module(m)
+    return m
+
+
+def load(src):
+    """TIGERDNR with its weights: tiger_dnr.py and the two layer modules it reads, without the package's training stack
+    (its layers/__init__ and base_model import lightning, rich and huggingface_hub)."""
+    import torch
+    from safetensors.torch import load_file
+
+    for f in ("models/tiger_dnr.py", "layers/activations.py", "layers/normalizations.py"):
+        fetch(REPO + f, src / Path(f).name)
+    for f in ("config.json", "model.safetensors"):
+        fetch(HF + f, src / f)
+    for p in ("look2hear", "look2hear.layers", "look2hear.models"):
+        sys.modules[p] = types.ModuleType(p)
+        sys.modules[p].__path__ = []
+    sys.dont_write_bytecode = True
+    layers = sys.modules["look2hear.layers"]
+    layers.activations = module("look2hear.layers.activations", src / "activations.py")
+    layers.normalizations = module("look2hear.layers.normalizations", src / "normalizations.py")
+    base = types.ModuleType("look2hear.models.base_model")
+
+    class BaseModel(torch.nn.Module):
+        def __init__(self, sample_rate, in_chan=1):
+            super().__init__()
+
+    base.BaseModel = BaseModel
+    sys.modules["look2hear.models.base_model"] = base
+    tiger = module("look2hear.models.tiger_dnr", src / "tiger_dnr.py")
+    model = tiger.TIGERDNR(**json.loads((src / "config.json").read_text()))
+    model.load_state_dict(load_file(str(src / "model.safetensors")))
+    return model.eval(), tiger
+
+
+def aap(x, output_size):
+    """F.adaptive_avg_pool1d, windows floor(i L / n) to ceil((i + 1) L / n) as torch's, by a cumulative sum read at the
+    windows' ends"""
+    import torch
+    import torch.nn.functional as F
+
+    L, n = int(x.shape[-1]), int(output_size[0] if isinstance(output_size, (list, tuple)) else output_size)
+    a = torch.tensor([(i * L) // n for i in range(n)])
+    b = torch.tensor([-((-(i + 1) * L) // n) for i in range(n)])
+    c = F.pad(torch.cumsum(x, -1), (1, 0))
+    return (c.index_select(-1, b) - c.index_select(-1, a)) / (b - a).to(x.dtype)
+
+
+def uconv(self, x):
+    """UConvBlock.forward, its global sum started at the first pooled level (0 + a is a)"""
+    residual = x
+    output = [self.spp_dw[0](self.proj_1x1(x))]
+    for k in range(1, self.depth):
+        output.append(self.spp_dw[k](output[-1]))
+    g = None
+    for fea in output:
+        p = aap(fea, output[-1].shape[-1])
+        g = p if g is None else g + p
+    g = self.globalatt(g)
+    fused = [self.loc_glo_fus[i](output[i], g) for i in range(self.depth)]
+    expanded = None
+    for i in range(self.depth - 2, -1, -1):
+        expanded = self.last_layer[i](fused[i], fused[i - 1] if i == self.depth - 2 else expanded)
+    return self.res_conv(expanded) + residual
+
+
+def group_norm(mod):
+    """GroupNorm(1, C), its mean and variance reduced one axis at a time"""
+    import torch
+
+    def mean(x):
+        for d in range(x.dim() - 1, 0, -1):
+            x = x.mean(d, keepdim=True)
+        return x
+
+    def forward(x):
+        mu = mean(x)
+        var = mean((x - mu) ** 2)
+        shape = (1, -1) + (1,) * (x.dim() - 2)
+        return (x - mu) / torch.sqrt(var + mod.eps) * mod.weight.view(shape) + mod.bias.view(shape)
+
+    return forward
+
+
+def stft(x):
+    """TIGER.forward's STFT: (B, L) -> (B, 2, F, T)"""
+    import torch
+
+    z = torch.stft(x, N_FFT, HOP, window=torch.hann_window(N_FFT), return_complex=True)
+    return torch.stack([z.real, z.imag], 1)
+
+
+def istft(y, length):
+    """(B, 2·S, F, T) -> (B, S, length), as TIGER.forward's iSTFT"""
+    import torch
+
+    B, C, F, T = y.shape
+    z = torch.complex(y[:, 0::2], y[:, 1::2]).reshape(-1, F, T)
+    return torch.istft(z, N_FFT, HOP, window=torch.hann_window(N_FFT), length=length).reshape(B, C // 2, length)
+
+
+class Core:
+    """The three TIGER.forward between their STFT and iSTFT, each its kept source."""
+
+    @staticmethod
+    def build(model, tiger):
+        import torch
+        import torch.nn as nn
+        import torch.nn.functional as F
+
+        # the patches, on the module the model's code reads them from
+        tiger.F = types.SimpleNamespace(**{k: getattr(F, k) for k in dir(F) if not k.startswith("_")}, )
+        tiger.F.adaptive_avg_pool1d = aap
+        tiger.UConvBlock.forward = uconv
+        for m in model.modules():
+            if isinstance(m, nn.GroupNorm) and m.num_groups == 1:
+                m.forward = group_norm(m)
+
+        def one(t, src, spec):
+            B, _, _, T = spec.shape
+            re, im = spec[:, 0], spec[:, 1]
+            feats, bands, i0 = [], [], 0
+            for i, bw in enumerate(t.band_width):
+                bw = int(bw)
+                bands.append((i0, bw))
+                feats.append(t.BN[i](spec[:, :, i0:i0 + bw].reshape(B, bw * 2, T)))
+                i0 += bw
+            x = t.separator(torch.stack(feats, 1)).view(B, t.nband, t.feature_dim, T)
+            outr, outi = [], []
+            for i, (i0, bw) in enumerate(bands):
+                o = t.mask[i](x[:, i]).view(B, 2, 2, t.num_output, bw, T)
+                mk = o[:, 0] * torch.sigmoid(o[:, 1])
+                mr, mi = mk[:, 0], mk[:, 1]
+                mr = (mr - (mr.sum(1, keepdim=True) - 1) / t.num_output)[:, src]
+                mi = (mi - mi.sum(1, keepdim=True) / t.num_output)[:, src]
+                r, j = re[:, i0:i0 + bw], im[:, i0:i0 + bw]
+                outr.append(r * mr - j * mi)
+                outi.append(r * mi + j * mr)
+            return torch.stack([torch.cat(outr, 1), torch.cat(outi, 1)], 1)
+
+        class C(nn.Module):
+            def __init__(self, m):
+                super().__init__()
+                self.m = m
+
+            def forward(self, spec):
+                return torch.cat([one(getattr(self.m, name), src, spec) for name, src in KEEP], 1)
+
+        return C(model).eval()
+
+
+def export(core, path, opset):
+    import torch
+
+    spec = stft(0.1 * torch.randn(1, SEGMENT))
+    with torch.no_grad():
+        torch.onnx.export(core, (spec,), str(path), input_names=["mix_spec"], output_names=["stems_spec"],
+                          opset_version=opset, dynamo=False)
+    print(f"wrote {path} ({path.stat().st_size / 1e6:.1f} MB), segment {SEGMENT} samples, spec {list(spec.shape)}")
+
+
+def session(path):
+    import onnxruntime as ort
+
+    o = ort.SessionOptions()
+    o.enable_cpu_mem_arena = False
+    o.enable_mem_pattern = False
+    return ort.InferenceSession(str(path), o, providers=["CPUExecutionProvider"])
+
+
+def verify(model, path, tol):
+    """One segment: the graph's sources through the iSTFT against each TIGER.forward's kept source (unpatched)."""
+    import math
+    import numpy as np
+    import torch
+
+    s = session(path)
+    t = torch.arange(SEGMENT) / RATE
+    x = (0.2 * sum(torch.sin(2 * math.pi * f * t) / k for k, f in enumerate((110, 220, 440, 880), 1)) + 0.05 * torch.randn(SEGMENT))[None]
+    got = istft(torch.from_numpy(s.run(None, {"mix_spec": stft(x).numpy()})[0]), SEGMENT).numpy()
+    with torch.no_grad():
+        want = np.stack([getattr(model, name)(x)[:, src].numpy() for name, src in KEEP], 1)
+    rel = np.abs(got - want).max() / max(np.abs(want).max(), 1e-12)
+    print(f"  verify {path.name}: max|diff|/max|y| = {rel:.2e} " + ("OK" if rel < tol else "FAIL"))
+    if rel >= tol:
+        raise SystemExit(f"{path}: verification failed ({rel:.2e} >= {tol})")
+    return s
+
+
+def reference(model, s, out_dir):
+    """upstream's wav_chunk_inference over reference.py's 20 s mix, each channel apart, its segments through the graph"""
+    import torch
+
+    sys.path.insert(0, str(Path(__file__).parent))
+    import reference as ref
+
+    def graph(k):
+        # TIGER.forward's contract for wav_chunk_inference: (B, L) -> (B, 3, L), the kept source at its index
+        def run(seg):
+            y = istft(torch.from_numpy(s.run(None, {"mix_spec": stft(seg.reshape(-1, seg.shape[-1])).numpy()})[0]), seg.shape[-1])
+            out = torch.zeros(y.shape[0], 3, y.shape[-1])
+            out[:, KEEP[k][1]] = y[:, k]
+            return out
+        return run
+
+    def run(mix):
+        with torch.no_grad():
+            x = torch.from_numpy(mix)[:, None]  # each channel a batch of one
+            out = {}
+            for k, name in enumerate(("dialogue", "music", "effects")):
+                y = torch.stack([model.wav_chunk_inference(graph(k), x[c:c + 1])[KEEP[k][1], 0] for c in range(len(mix))])
+                out[name] = y.numpy()
+            return out
+
+    ref.write(out_dir, run, 20.0)
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--out-dir", default=None, help="default: $AUDIO_NEURAL_CACHE/tiger or ~/.cache/audiojs/neural/tiger")
+    ap.add_argument("--opset", type=int, default=17)
+    ap.add_argument("--verify", action="store_true", help="compare with TIGER.forward; write the reference separation")
+    args = ap.parse_args()
+
+    cache = Path(os.environ.get("AUDIO_NEURAL_CACHE") or Path.home() / ".cache" / "audiojs" / "neural")
+    out_dir = Path(args.out_dir) if args.out_dir else cache / "tiger"
+    out_dir.mkdir(parents=True, exist_ok=True)
+    src = out_dir / "src"
+    src.mkdir(exist_ok=True)
+    try:
+        model, tiger = load(src)
+        ref_model, _ = load(src) if args.verify else (None, None)
+    except ImportError as e:
+        raise SystemExit(f"torch and safetensors are required: pip install torch onnx onnxruntime safetensors ({e})")
+    path = out_dir / "tiger.onnx"
+    export(Core.build(model, tiger), path, args.opset)
+    if args.verify:
+        s = verify(ref_model, path, 1e-4)
+        reference(ref_model, s, out_dir)
+    print(f"done: tiger in {out_dir}")
+
+
+if __name__ == "__main__":
+    sys.exit(main())

package/scripts/reference.py

@@ -25,9 +25,48 @@ def test_mix(seconds=9.0, rate=RATE):
     return np.stack([voice + 1.1 * bass + drums + 0.7 * pad, voice + 0.9 * bass + drums + 1.3 * pad]).astype(np.float32)
 
 
-def write(out_dir, separate):
+def normalized(mix, run):
+    """run() on the mix less its mean over its deviation (of the channels' mean, ddof 1, + 1e-8), its output scaled
+    back, the mean added to no source (separate.js's separateComplex)."""
+    mono = mix.astype(np.float64).mean(0)
+    mean, std = mono.mean(), mono.std(ddof=1) + 1e-8
+    return run(((mix - mean) / std).astype(np.float32)) * std
+
+
+def chunked(mix, run, L, step=0.25):
+    """ZFTurbo/Music-Source-Separation-Training's demix() (utils/model_utils.py, generic mode) over run(segment
+    (C, L) float32) -> (S, C, L): segments every step·L until one reaches the end, linear fades of L/10 (the first
+    segment not in, the last not out), the input reflected L - step out on each side when longer than
+    2 (L - step), a segment past the end reflected out when more than half of it is input, else zero-padded;
+    float64 sums (separate.js's separateComplex). -> (S, C, N)"""
+    C, N = mix.shape
+    hop = max(1, int(L * step))
+    border = L - hop
+    wide = N > 2 * border and border > 0
+    x = np.pad(mix, ((0, 0), (border, border)), mode="reflect") if wide else mix
+    M, fade = x.shape[1], L // 10
+    starts = [k * hop for k in range(max(1, -(-(M - L) // hop) + 1))]  # until one reaches the end
+    acc, wsum = None, np.zeros(M)
+    for k, s in enumerate(starts):
+        part = x[:, s:s + L]
+        n = part.shape[1]
+        part = np.pad(part, ((0, 0), (0, L - n)), mode="reflect" if n > L // 2 else "constant")
+        y = np.asarray(run(part.astype(np.float32)), dtype=np.float64)
+        w = np.ones(L)
+        if k > 0:
+            w[:fade] = np.linspace(0, 1, fade)
+        if k < len(starts) - 1:
+            w[L - fade:] = np.linspace(1, 0, fade)
+        acc = np.zeros((y.shape[0], C, M)) if acc is None else acc
+        acc[..., s:s + n] += y[..., :n] * w[:n]
+        wsum[s:s + n] += w[:n]
+    out = acc / np.where(wsum > 0, wsum, 1)
+    return out[..., border:border + N] if wide else out
+
+
+def write(out_dir, separate, seconds=9.0):
     """separate(mix (2, N) float32) -> { target: (2, N) array }"""
-    mix = test_mix()
+    mix = test_mix(seconds)
     mix.tofile(out_dir / "test.f32")
     for name, x in separate(mix).items():
         np.asarray(x, dtype=np.float32).tofile(out_dir / f"test.{name}.f32")