Harden ace step transcription engine to detect and break out of loops. Add melbanroformer to seperate out the vocals and music for better transcriptions. Much better results.

toolkit/audio/melbandroformer/HF_README.md

@@ -0,0 +1,52 @@
+---
+license: mit
+pipeline_tag: audio-to-audio
+tags:
+  - music-source-separation
+  - vocals
+  - mel-band-roformer
+  - ai-toolkit
+---
+
+# Mel-Band RoFormer (vocals) for ai-toolkit
+
+Safetensors repack of Kimberley Jensen's Mel-Band RoFormer vocal separation model, used by
+[ai-toolkit](https://github.com/ostris/ai-toolkit) to split songs into a vocals track and an
+instrumental track (`instrumental = mix - vocals`).
+
+**These weights are not ours.** All credit goes to the original authors:
+
+- Weights: [KimberleyJSN/melbandroformer](https://huggingface.co/KimberleyJSN/melbandroformer) (`MelBandRoformer.ckpt`, MIT) by Kimberley Jensen
+- Training framework and config: [ZFTurbo/Music-Source-Separation-Training](https://github.com/ZFTurbo/Music-Source-Separation-Training)
+  ([`config_vocals_mel_band_roformer_kj.yaml`](https://github.com/ZFTurbo/Music-Source-Separation-Training/blob/main/configs/KimberleyJensen/config_vocals_mel_band_roformer_kj.yaml)) by Roman Solovyev
+- Architecture implementation: [lucidrains/BS-RoFormer](https://github.com/lucidrains/BS-RoFormer) by Phil Wang
+- Paper: [Mel-Band RoFormer for Music Source Separation](https://arxiv.org/abs/2310.01809), Wang et al. 2023
+
+## Files
+
+| file | notes |
+|---|---|
+| `melbandroformer_vocals_kj.safetensors` | fp32, byte-identical tensors to the original `.ckpt`. Model kwargs and inference defaults (`chunk_size`, `num_overlap`) are stored in the safetensors metadata, so no separate config is needed. |
+
+Conversion script: [`scripts/convert_melbandroformer.py`](https://github.com/ostris/ai-toolkit/blob/main/scripts/convert_melbandroformer.py).
+
+## Usage
+
+ai-toolkit downloads this file automatically on first use:
+
+```bash
+python -m toolkit.audio.melbandroformer song.flac
+# -> song_vocals.flac, song_instrumental.flac
+```
+
+```python
+from toolkit.audio.melbandroformer import load_melbandroformer, separate
+model = load_melbandroformer(device="cuda", compile=True)
+vocals, instrumental = separate(model, wav, sample_rate)  # wav: [channels, samples]
+```
+
+Input: mono or stereo at any sample rate (resampled to 44.1 kHz internally, output at the input rate).
+
+## License
+
+MIT, same as the original weights and code.

toolkit/audio/melbandroformer/LICENSE

@@ -0,0 +1,28 @@
+The code in this directory is derived from:
+
+  https://github.com/ZFTurbo/Music-Source-Separation-Training  (MIT, Copyright (c) 2024 Roman Solovyev)
+  https://github.com/lucidrains/BS-RoFormer                     (MIT, Copyright (c) 2023 Phil Wang)
+  https://github.com/lucidrains/rotary-embedding-torch          (MIT, Copyright (c) 2021 Phil Wang)
+
+The default weights (melbandroformer_vocals_kj.safetensors) are converted from
+https://huggingface.co/KimberleyJSN/melbandroformer (MIT).
+
+MIT License
+
+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.

toolkit/audio/melbandroformer/__init__.py

@@ -0,0 +1,12 @@
+from .model import MelBandRoformer
+from .separate import load_melbandroformer, separate, separate_stems, get_weights_path, HF_REPO, DEFAULT_WEIGHTS
+
+__all__ = [
+    "MelBandRoformer",
+    "load_melbandroformer",
+    "separate",
+    "separate_stems",
+    "get_weights_path",
+    "HF_REPO",
+    "DEFAULT_WEIGHTS",
+]

toolkit/audio/melbandroformer/__main__.py

@@ -0,0 +1,120 @@
+"""python -m toolkit.audio.melbandroformer song.flac [more files or dirs] [--out_dir DIR] [--format flac]
+Writes <name>_vocals.<format> and <name>_instrumental.<format> next to each input (or in --out_dir).
+Decode and encode run in worker threads so the GPU is never waiting on file IO."""
+import argparse
+import os
+import time
+from concurrent.futures import ThreadPoolExecutor
+
+from dotenv import load_dotenv
+
+load_dotenv(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "..", ".env"))  # MODELS_PATH, before toolkit.paths
+
+import torch
+
+from . import load_melbandroformer, separate, DEFAULT_WEIGHTS
+
+AUDIO_EXTS = {".flac", ".wav", ".mp3", ".m4a", ".ogg", ".opus", ".aac", ".wma", ".aif", ".aiff"}
+
+
+def collect_inputs(paths):
+    files = []
+    for p in paths:
+        if os.path.isdir(p):
+            for root, _, names in os.walk(p):
+                files += [os.path.join(root, n) for n in sorted(names) if os.path.splitext(n)[1].lower() in AUDIO_EXTS]
+        else:
+            files.append(p)
+    return files
+
+
+def output_paths(path, out_dir, fmt):
+    out_dir = out_dir or os.path.dirname(path)
+    name = os.path.splitext(os.path.basename(path))[0]
+    return {stem: os.path.join(out_dir, f"{name}_{stem}.{fmt}") for stem in ("vocals", "instrumental")}
+
+
+def decode(path):
+    from torchcodec.decoders import AudioDecoder
+    samples = AudioDecoder(path).get_all_samples()
+    return samples.data, samples.sample_rate
+
+
+def encode(audio, sample_rate, out_path):
+    from torchcodec.encoders import AudioEncoder
+    os.makedirs(os.path.dirname(out_path) or ".", exist_ok=True)
+    AudioEncoder(audio, sample_rate=sample_rate).to_file(out_path)
+    return out_path
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    parser.add_argument("inputs", nargs="+", help="audio files or directories")
+    parser.add_argument("--out_dir", default=None)
+    parser.add_argument("--format", default="flac", help="output container/codec by extension (flac is lossless and ~6x faster to encode than mp3)")
+    parser.add_argument("--weights", default=DEFAULT_WEIGHTS, help="filename under MODELS_PATH/checkpoints")
+    parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
+    parser.add_argument("--batch_size", type=int, default=8)
+    parser.add_argument("--fp32", action="store_true", help="disable fp16 autocast")
+    parser.add_argument("--no_compile", action="store_true", help="skip torch.compile of the trunk (faster startup, slower per file)")
+    parser.add_argument("--overwrite", action="store_true", help="redo files whose outputs already exist")
+    parser.add_argument("--io_workers", type=int, default=4)
+    args = parser.parse_args()
+
+    files = collect_inputs(args.inputs)
+    todo = []
+    for path in files:
+        outs = output_paths(path, args.out_dir, args.format)
+        if args.overwrite or not all(os.path.exists(p) for p in outs.values()):
+            todo.append((path, outs))
+    print(f"{len(todo)} of {len(files)} files to separate")
+    if not todo:
+        return
+
+    t_start = time.perf_counter()
+    model = load_melbandroformer(args.weights, device=args.device, compile=not args.no_compile)
+    dtype = None if args.fp32 else torch.float16
+    total_audio = 0.0
+
+    decoder_pool = ThreadPoolExecutor(max_workers=2)
+    encoder_pool = ThreadPoolExecutor(max_workers=args.io_workers)
+    pending_encodes = []
+
+    # prefetch two files ahead of the GPU
+    futures = [decoder_pool.submit(decode, path) for path, _ in todo[:2]]
+    for idx, (path, outs) in enumerate(todo):
+        if idx + 2 < len(todo):
+            futures.append(decoder_pool.submit(decode, todo[idx + 2][0]))
+        t_wait = time.perf_counter()
+        try:
+            wav, sample_rate = futures.pop(0).result()
+        except Exception as e:
+            print(f"skip {path}: {type(e).__name__}: {e}")
+            continue
+
+        t0 = time.perf_counter()
+        wav = wav.to(args.device, non_blocking=True)
+        vocals, instrumental = separate(model, wav, sample_rate, batch_size=args.batch_size, dtype=dtype)
+        for stem, audio in (("vocals", vocals), ("instrumental", instrumental)):
+            pending_encodes.append(encoder_pool.submit(encode, audio.cpu(), sample_rate, outs[stem]))
+        seconds = wav.shape[-1] / sample_rate
+        total_audio += seconds
+        print(f"[{idx + 1}/{len(todo)}] {path}  {seconds:.0f}s audio in {time.perf_counter() - t0:.2f}s"
+              f" (waited {t0 - t_wait:.2f}s for decode)", flush=True)
+
+        # keep the encode queue bounded so memory does not grow with the file count
+        while len(pending_encodes) > 2 * args.io_workers:
+            pending_encodes.pop(0).result()
+
+    t_drain = time.perf_counter()
+    for f in pending_encodes:
+        f.result()
+    decoder_pool.shutdown()
+    encoder_pool.shutdown()
+    wall = time.perf_counter() - t_start
+    print(f"done: {len(todo)} files, {total_audio:.0f}s audio in {wall:.1f}s ({total_audio / wall:.0f}x realtime, "
+          f"final encode drain {time.perf_counter() - t_drain:.2f}s)", flush=True)
+
+
+if __name__ == "__main__":
+    main()

toolkit/audio/melbandroformer/model.py

@@ -0,0 +1,470 @@
+# Vendored from ZFTurbo/Music-Source-Separation-Training (MIT, Roman Solovyev 2024),
+# derived from lucidrains/BS-RoFormer (MIT, Phil Wang 2023). See LICENSE in this directory.
+# Inference-only: training loss, PoPE, linear attention and checkpointing removed;
+# rotary_embedding_torch inlined; attention is plain torch SDPA; the 60 per-band
+# linears are stacked bmm's (upstream per-band checkpoint keys are repacked on load).
+# core() is torch.compile-friendly: no einops, no complex ops, dynamic batch dim.
+from functools import partial
+
+import torch
+from torch import nn
+from torch.nn import Module, ModuleList
+import torch.nn.functional as F
+from einops import rearrange, pack, unpack, reduce, repeat
+
+
+def exists(val):
+    return val is not None
+
+
+def pack_one(t, pattern):
+    return pack([t], pattern)
+
+
+def unpack_one(t, ps, pattern):
+    return unpack(t, ps, pattern)[0]
+
+
+class RMSNorm(Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.scale = dim ** 0.5
+        self.gamma = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        return F.normalize(x, dim=-1) * self.scale * self.gamma
+
+
+# rotary (matches rotary_embedding_torch.RotaryEmbedding(dim) defaults: freqs_for='lang', theta=1e4)
+
+def rotate_half(x):
+    x1, x2 = x.unflatten(-1, (-1, 2)).unbind(dim=-1)
+    return torch.stack((-x2, x1), dim=-1).flatten(-2)
+
+
+class RotaryEmbedding(Module):
+    def __init__(self, dim):
+        super().__init__()
+        freqs = 1.0 / (10000 ** (torch.arange(0, dim, 2)[: dim // 2].float() / dim))
+        # nn.Parameter (not buffer) so the key matches upstream checkpoints
+        self.freqs = nn.Parameter(freqs, requires_grad=False)
+
+    def rotate_queries_or_keys(self, t):
+        # t: (b h n d); rotation is computed in fp32 regardless of autocast
+        n = t.shape[-2]
+        with torch.autocast(device_type=t.device.type, enabled=False):
+            seq = torch.arange(n, device=t.device, dtype=self.freqs.dtype)
+            freqs = torch.einsum('n,f->nf', seq, self.freqs).repeat_interleave(2, dim=-1)
+            tf = t.float()
+            out = tf * freqs.cos() + rotate_half(tf) * freqs.sin()
+        return out.type(t.dtype)
+
+
+# attention
+
+class FeedForward(Module):
+    def __init__(self, dim, mult=4, dropout=0.):
+        super().__init__()
+        dim_inner = int(dim * mult)
+        self.net = nn.Sequential(
+            RMSNorm(dim),
+            nn.Linear(dim, dim_inner),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(dim_inner, dim),
+            nn.Dropout(dropout),
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+class Attention(Module):
+    def __init__(self, dim, heads=8, dim_head=64, dropout=0., rotary_embed=None):
+        super().__init__()
+        self.heads = heads
+        self.dropout = dropout
+        dim_inner = heads * dim_head
+
+        self.rotary_embed = rotary_embed
+
+        self.norm = RMSNorm(dim)
+        self.to_qkv = nn.Linear(dim, dim_inner * 3, bias=False)
+        self.to_gates = nn.Linear(dim, heads)
+        self.to_out = nn.Sequential(
+            nn.Linear(dim_inner, dim, bias=False),
+            nn.Dropout(dropout),
+        )
+
+    def forward(self, x):
+        x = self.norm(x)
+
+        b, n, _ = x.shape
+        q, k, v = self.to_qkv(x).view(b, n, 3, self.heads, -1).permute(2, 0, 3, 1, 4).unbind(0)  # (b h n d)
+
+        if exists(self.rotary_embed):
+            q = self.rotary_embed.rotate_queries_or_keys(q)
+            k = self.rotary_embed.rotate_queries_or_keys(k)
+
+        out = F.scaled_dot_product_attention(q, k, v, dropout_p=self.dropout if self.training else 0.)
+
+        gates = self.to_gates(x)
+        out = out * gates.transpose(1, 2).unsqueeze(-1).sigmoid()
+
+        out = out.transpose(1, 2).reshape(b, n, -1)
+        return self.to_out(out)
+
+
+class Transformer(Module):
+    def __init__(
+        self,
+        *,
+        dim,
+        depth,
+        dim_head=64,
+        heads=8,
+        attn_dropout=0.,
+        ff_dropout=0.,
+        ff_mult=4,
+        norm_output=True,
+        rotary_embed=None,
+    ):
+        super().__init__()
+        self.layers = ModuleList([])
+
+        for _ in range(depth):
+            self.layers.append(ModuleList([
+                Attention(dim=dim, dim_head=dim_head, heads=heads, dropout=attn_dropout, rotary_embed=rotary_embed),
+                FeedForward(dim=dim, mult=ff_mult, dropout=ff_dropout),
+            ]))
+
+        self.norm = RMSNorm(dim) if norm_output else nn.Identity()
+
+    def forward(self, x):
+        for attn, ff in self.layers:
+            x = attn(x) + x
+            x = ff(x) + x
+        return self.norm(x)
+
+
+# band split / mask estimator: all bands in stacked bmm's instead of 60 per-band linears.
+# Parameters are stored stacked and zero-padded to the widest band; upstream per-band
+# checkpoints are repacked on load (see _load_from_state_dict), so keys stay compatible.
+
+class BandSplit(Module):
+    def __init__(self, dim, dim_inputs):
+        super().__init__()
+        self.dim_inputs = tuple(dim_inputs)
+        n, max_in = len(dim_inputs), max(dim_inputs)
+        self.n, self.max_in = n, max_in
+
+        self.gamma = nn.Parameter(torch.ones(n, max_in))
+        self.weight = nn.Parameter(torch.empty(n, max_in, dim))
+        self.bias = nn.Parameter(torch.zeros(n, dim))
+        for i, d in enumerate(dim_inputs):
+            bound = d ** -0.5
+            nn.init.uniform_(self.weight[i, :d], -bound, bound)
+            nn.init.uniform_(self.bias[i], -bound, bound)
+
+        # gather map from the concatenated band features into the padded (n, max_in) layout;
+        # pads read the extra zero column appended in forward. Built on cpu for meta-device init.
+        pad_src = torch.full((n, max_in), sum(dim_inputs), dtype=torch.long, device='cpu')
+        offset = 0
+        for i, d in enumerate(dim_inputs):
+            pad_src[i, :d] = torch.arange(offset, offset + d, device='cpu')
+            offset += d
+        self.register_buffer('pad_src', pad_src.flatten(), persistent=False)
+        self.register_buffer('scale', torch.tensor([d ** 0.5 for d in dim_inputs], device='cpu')[:, None], persistent=False)
+
+    def _load_from_state_dict(self, state_dict, prefix, *args):
+        if f"{prefix}to_features.0.1.weight" in state_dict:
+            ws, bs, gs = [], [], []
+            for i, d in enumerate(self.dim_inputs):
+                ws.append(F.pad(state_dict.pop(f"{prefix}to_features.{i}.1.weight").t(), (0, 0, 0, self.max_in - d)))
+                bs.append(state_dict.pop(f"{prefix}to_features.{i}.1.bias"))
+                gs.append(F.pad(state_dict.pop(f"{prefix}to_features.{i}.0.gamma"), (0, self.max_in - d)))
+            state_dict[f"{prefix}weight"] = torch.stack(ws)
+            state_dict[f"{prefix}bias"] = torch.stack(bs)
+            state_dict[f"{prefix}gamma"] = torch.stack(gs)
+        super()._load_from_state_dict(state_dict, prefix, *args)
+
+    def forward(self, x):
+        b, t, _ = x.shape
+        x = F.pad(x.reshape(b * t, -1), (0, 1))[:, self.pad_src].view(b * t, self.n, self.max_in)
+        # per-band RMSNorm; zero pads do not change the norm
+        norm = x.float().norm(dim=-1, keepdim=True).clamp(min=1e-12)
+        x = (x / norm * self.scale * self.gamma).transpose(0, 1)  # (n, M, max_in)
+        out = torch.baddbmm(self.bias[:, None, :], x, self.weight)  # (n, M, dim)
+        return out.transpose(0, 1).reshape(b, t, self.n, -1)
+
+
+class MaskEstimator(Module):
+    def __init__(self, dim, dim_inputs, depth, mlp_expansion_factor=4):
+        super().__init__()
+        self.dim_inputs = tuple(dim_inputs)
+        n, hidden = len(dim_inputs), dim * mlp_expansion_factor
+        dim_outs = [2 * d for d in dim_inputs]  # GLU halves it back to d
+        max_out = max(dim_outs)
+        self.n, self.depth, self.max_out = n, depth, max_out
+
+        dims = (dim, *((hidden,) * depth), max_out)
+        self.weights = nn.ParameterList([nn.Parameter(torch.empty(n, i, o)) for i, o in zip(dims[:-1], dims[1:])])
+        self.biases = nn.ParameterList([nn.Parameter(torch.zeros(n, o)) for o in dims[1:]])
+        for w, bias, i in zip(self.weights, self.biases, dims[:-1]):
+            nn.init.uniform_(w, -i ** -0.5, i ** -0.5)
+            nn.init.uniform_(bias, -i ** -0.5, i ** -0.5)
+
+        # GLU gather: a * sigmoid(b) with a/b the first/second half of each band's valid output columns
+        glu_a, glu_b = [], []
+        for i, d in enumerate(dim_inputs):
+            base = i * max_out
+            glu_a.append(torch.arange(base, base + d, device='cpu'))
+            glu_b.append(torch.arange(base + d, base + 2 * d, device='cpu'))
+        self.register_buffer('glu_a', torch.cat(glu_a), persistent=False)
+        self.register_buffer('glu_b', torch.cat(glu_b), persistent=False)
+
+    def _load_from_state_dict(self, state_dict, prefix, *args):
+        if f"{prefix}to_freqs.0.0.0.weight" in state_dict:
+            for j in range(self.depth + 1):
+                ws, bs = [], []
+                for i in range(self.n):
+                    w = state_dict.pop(f"{prefix}to_freqs.{i}.0.{2 * j}.weight").t()
+                    bias = state_dict.pop(f"{prefix}to_freqs.{i}.0.{2 * j}.bias")
+                    if j == self.depth:
+                        w = F.pad(w, (0, self.max_out - w.shape[1]))
+                        bias = F.pad(bias, (0, self.max_out - bias.shape[0]))
+                    ws.append(w)
+                    bs.append(bias)
+                state_dict[f"{prefix}weights.{j}"] = torch.stack(ws)
+                state_dict[f"{prefix}biases.{j}"] = torch.stack(bs)
+        super()._load_from_state_dict(state_dict, prefix, *args)
+
+    def forward(self, x):
+        b, t, n, d = x.shape
+        h = x.permute(2, 0, 1, 3).reshape(n, b * t, d)
+        for j in range(self.depth):
+            h = torch.tanh(torch.baddbmm(self.biases[j][:, None, :], h, self.weights[j]))
+        out = torch.baddbmm(self.biases[-1][:, None, :], h, self.weights[-1])  # (n, M, max_out)
+        out = out.transpose(0, 1).reshape(b * t, n * self.max_out)
+        out = out[:, self.glu_a] * torch.sigmoid(out[:, self.glu_b])
+        return out.view(b, t, -1)
+
+
+# main class
+
+class MelBandRoformer(Module):
+    def __init__(
+        self,
+        dim,
+        *,
+        depth,
+        stereo=False,
+        num_stems=1,
+        time_transformer_depth=2,
+        freq_transformer_depth=2,
+        num_bands=60,
+        dim_head=64,
+        heads=8,
+        attn_dropout=0.1,
+        ff_dropout=0.1,
+        sample_rate=44100,  # for the librosa mel filter bank
+        stft_n_fft=2048,
+        stft_hop_length=512,
+        stft_win_length=2048,
+        stft_normalized=False,
+        zero_dc=True,
+        mask_estimator_depth=1,
+        mlp_expansion_factor=4,
+    ):
+        super().__init__()
+
+        self.stereo = stereo
+        self.audio_channels = 2 if stereo else 1
+        self.num_stems = num_stems
+
+        self.layers = ModuleList([])
+
+        transformer_kwargs = dict(
+            dim=dim,
+            heads=heads,
+            dim_head=dim_head,
+            attn_dropout=attn_dropout,
+            ff_dropout=ff_dropout,
+        )
+
+        # one rotary module shared across all time (resp. freq) transformers, as upstream
+        time_rotary_embed = RotaryEmbedding(dim=dim_head)
+        freq_rotary_embed = RotaryEmbedding(dim=dim_head)
+
+        for _ in range(depth):
+            self.layers.append(nn.ModuleList([
+                Transformer(depth=time_transformer_depth, rotary_embed=time_rotary_embed, **transformer_kwargs),
+                Transformer(depth=freq_transformer_depth, rotary_embed=freq_rotary_embed, **transformer_kwargs),
+            ]))
+
+        self.stft_window_fn = partial(torch.hann_window, stft_win_length)
+
+        self.stft_kwargs = dict(
+            n_fft=stft_n_fft,
+            hop_length=stft_hop_length,
+            win_length=stft_win_length,
+            normalized=stft_normalized,
+        )
+
+        freqs = stft_n_fft // 2 + 1
+
+        # binary mel filter bank as in section 2 of the paper; overlapping bands' masks are averaged
+        from librosa import filters  # lazy: ~1s import, only needed here
+        mel_filter_bank = torch.from_numpy(filters.mel(sr=sample_rate, n_fft=stft_n_fft, n_mels=num_bands))
+        # upstream forces the first/last bins in so every frequency is covered by some band
+        mel_filter_bank[0][0] = 1.
+        mel_filter_bank[-1, -1] = 1.
+
+        freqs_per_band = mel_filter_bank > 0
+        assert freqs_per_band.any(dim=0).all(), 'all frequencies need to be covered by all bands for now'
+
+        # explicit cpu so the loader can build the module under torch.device('meta')
+        repeated_freq_indices = repeat(torch.arange(freqs, device='cpu'), 'f -> b f', b=num_bands)
+        freq_indices = repeated_freq_indices[freqs_per_band]
+
+        if stereo:
+            freq_indices = repeat(freq_indices, 'f -> f s', s=2)
+            freq_indices = freq_indices * 2 + torch.arange(2, device='cpu')
+            freq_indices = rearrange(freq_indices, 'f s -> (f s)')
+
+        self.register_buffer('freq_indices', freq_indices, persistent=False)
+        self.register_buffer('freqs_per_band', freqs_per_band, persistent=False)
+
+        num_freqs_per_band = reduce(freqs_per_band, 'b f -> b', 'sum')
+        num_bands_per_freq = reduce(freqs_per_band, 'b f -> f', 'sum')
+
+        self.register_buffer('num_freqs_per_band', num_freqs_per_band, persistent=False)
+        self.register_buffer('num_bands_per_freq', num_bands_per_freq, persistent=False)
+
+        freqs_per_bands_with_complex = tuple(2 * f * self.audio_channels for f in num_freqs_per_band.tolist())
+
+        self.band_split = BandSplit(dim=dim, dim_inputs=freqs_per_bands_with_complex)
+
+        self.mask_estimators = nn.ModuleList([
+            MaskEstimator(
+                dim=dim,
+                dim_inputs=freqs_per_bands_with_complex,
+                depth=mask_estimator_depth,
+                mlp_expansion_factor=mlp_expansion_factor,
+            )
+            for _ in range(num_stems)
+        ])
+
+        self.zero_dc = zero_dc
+
+    def stft(self, raw_audio):
+        """raw_audio (b s t) -> band-gathered real features (b t (f c)), and stft_repr (b (f s) t c)."""
+        device = raw_audio.device
+        batch = raw_audio.shape[0]
+
+        raw_audio, batch_audio_channel_packed_shape = pack_one(raw_audio, '* t')
+
+        stft_window = self.stft_window_fn(device=device)
+
+        stft_repr = torch.stft(raw_audio, **self.stft_kwargs, window=stft_window, return_complex=True)
+        stft_repr = torch.view_as_real(stft_repr)
+
+        stft_repr = unpack_one(stft_repr, batch_audio_channel_packed_shape, '* f t c')
+
+        # merge stereo / mono into the frequency dim, frequency leading, for band splitting
+        stft_repr = rearrange(stft_repr, 'b s f t c -> b (f s) t c')
+
+        batch_arange = torch.arange(batch, device=device)[..., None]
+
+        x = stft_repr[batch_arange, self.freq_indices]
+
+        # fold complex into the frequency dim
+        x = rearrange(x, 'b f t c -> b t (f c)')
+        return x, stft_repr
+
+    def core(self, x):
+        """Real-valued trunk (b t (f c)) -> masks (b n t (f c)); the torch.compile target."""
+        x = self.band_split(x)
+
+        # axial attention: time then freq
+
+        b, t, f, d = x.shape
+        # native views instead of einops: same math, ~3x faster dynamo tracing
+        for time_transformer, freq_transformer in self.layers:
+            x = time_transformer(x.transpose(1, 2).reshape(b * f, t, d)).view(b, f, t, d)
+            x = freq_transformer(x.transpose(1, 2).reshape(b * t, f, d)).view(b, t, f, d)
+
+        return torch.stack([fn(x) for fn in self.mask_estimators], dim=1)
+
+    def apply_masks(self, stft_repr, masks, channels):
+        batch = stft_repr.shape[0]
+        num_stems = masks.shape[1]
+
+        masks = rearrange(masks, 'b n t (f c) -> b n f t c', c=2)
+
+        # modulate frequency representation (complex multiply)
+
+        stft_repr = rearrange(stft_repr, 'b f t c -> b 1 f t c')
+
+        stft_repr = torch.view_as_complex(stft_repr)
+        masks = torch.view_as_complex(masks)
+
+        masks = masks.type(stft_repr.dtype)
+
+        # average the estimated mask over overlapping bands
+
+        scatter_indices = repeat(self.freq_indices, 'f -> b n f t', b=batch, n=num_stems, t=stft_repr.shape[-1])
+
+        stft_repr_expanded_stems = repeat(stft_repr, 'b 1 ... -> b n ...', n=num_stems)
+        masks_summed = torch.zeros_like(stft_repr_expanded_stems).scatter_add_(2, scatter_indices, masks)
+
+        denom = repeat(self.num_bands_per_freq, 'f -> (f r) 1', r=channels)
+
+        masks_averaged = masks_summed / denom.clamp(min=1e-8)
+
+        stft_repr = stft_repr * masks_averaged
+
+        # istft
+
+        stft_repr = rearrange(stft_repr, 'b n (f s) t -> (b n s) f t', s=self.audio_channels)
+
+        if self.zero_dc:
+            stft_repr[:, 0] = 0.
+
+        stft_window = self.stft_window_fn(device=stft_repr.device)
+        recon_audio = torch.istft(stft_repr, **self.stft_kwargs, window=stft_window, return_complex=False)
+
+        recon_audio = rearrange(recon_audio, '(b n s) t -> b n s t', b=batch, s=self.audio_channels, n=num_stems)
+
+        if num_stems == 1:
+            recon_audio = rearrange(recon_audio, 'b 1 s t -> b s t')
+
+        return recon_audio
+
+    def compile_core(self, **kwargs):
+        """Fuse the trunk's elementwise chains; batch dim is marked dynamic so any batch size hits one graph."""
+        compiled = torch.compile(self.core, **kwargs)
+
+        def core(x):
+            torch._dynamo.mark_dynamic(x, 0)
+            return compiled(x)
+
+        self.core = core
+        return self
+
+    def forward(self, raw_audio):
+        """
+        raw_audio: (b t) or (b s t). Returns (b s t) for a single stem, else (b n s t).
+        b batch, f freq, t time, s audio channel, n stem, c complex (2), d feature
+        """
+        if raw_audio.ndim == 2:
+            raw_audio = rearrange(raw_audio, 'b t -> b 1 t')
+
+        channels = raw_audio.shape[1]
+
+        assert (not self.stereo and channels == 1) or (self.stereo and channels == 2), \
+            'stereo model needs 2-channel audio, mono model needs 1-channel audio'
+
+        x, stft_repr = self.stft(raw_audio)
+        masks = self.core(x)
+        return self.apply_masks(stft_repr, masks, channels)

toolkit/audio/melbandroformer/separate.py

@@ -0,0 +1,158 @@
+import json
+import os
+
+import torch
+import torch.nn.functional as F
+from safetensors import safe_open
+
+from .model import MelBandRoformer
+
+HF_REPO = "ai-toolkit/melbandroformer"
+DEFAULT_WEIGHTS = "melbandroformer_vocals_kj.safetensors"
+MODEL_SAMPLE_RATE = 44100
+
+
+def get_weights_path(filename: str = DEFAULT_WEIGHTS) -> str:
+    """MODELS_PATH/checkpoints/<filename>, pulled from the HF repo root if missing."""
+    from toolkit.paths import MODELS_PATH  # lazy: read after the CLI has loaded .env
+    ckpt_dir = os.path.join(MODELS_PATH, "checkpoints")
+    path = os.path.join(ckpt_dir, filename)
+    if os.path.exists(path):
+        return path
+    from huggingface_hub import hf_hub_download
+    os.makedirs(ckpt_dir, exist_ok=True)
+    print(f"Downloading {HF_REPO}/{filename} to {ckpt_dir}")
+    return hf_hub_download(repo_id=HF_REPO, filename=filename, local_dir=ckpt_dir)
+
+
+def load_melbandroformer(filename: str = DEFAULT_WEIGHTS, device=None, compile: bool = False) -> MelBandRoformer:
+    path = get_weights_path(filename)
+    device = torch.device(device if device is not None else "cpu")
+    with safe_open(path, framework="pt", device=str(device)) as f:
+        metadata = f.metadata() or {}
+        state_dict = {k: f.get_tensor(k) for k in f.keys()}
+    if "config" not in metadata:
+        raise ValueError(
+            f"{path} has no model config in its metadata; convert it with scripts/convert_melbandroformer.py"
+        )
+    # meta init + assign: no cpu random init of 228M params, tensors go file -> device directly
+    with torch.device("meta"):
+        model = MelBandRoformer(**json.loads(metadata["config"]))
+    model.load_state_dict(state_dict, strict=True, assign=True)
+    model.to(device)  # index buffers are built on cpu
+    # chunk_size / num_overlap the weights were validated with
+    model.inference_config = json.loads(metadata.get("inference", "{}"))
+    model.eval()
+    if compile:
+        model.compile_core()
+    return model
+
+
+def _fade_window(chunk_size: int, fade_size: int) -> torch.Tensor:
+    window = torch.ones(chunk_size)
+    window[:fade_size] = torch.linspace(0, 1, fade_size)
+    window[-fade_size:] = torch.linspace(1, 0, fade_size)
+    return window
+
+
+@torch.no_grad()
+def separate_stems(
+    model: MelBandRoformer,
+    mix: torch.Tensor,
+    chunk_size: int = None,
+    num_overlap: int = None,
+    batch_size: int = 4,
+    dtype=torch.float16,
+    progress: bool = False,
+) -> torch.Tensor:
+    """Overlap-add chunked inference (MSST `demix`). mix: [C, T] at MODEL_SAMPLE_RATE. Returns [num_stems, C, T] on cpu."""
+    cfg = getattr(model, "inference_config", {})
+    chunk_size = chunk_size or cfg.get("chunk_size", 352800)
+    num_overlap = num_overlap or cfg.get("num_overlap", 2)
+
+    device = next(model.parameters()).device
+    mix = mix.to(device=device, dtype=torch.float32)
+    length = mix.shape[-1]
+
+    step = chunk_size // num_overlap
+    fade_size = chunk_size // 10
+    border = chunk_size - step
+
+    # reflect-pad the edges so the first/last chunks are not fading in/out on real audio
+    padded = length > 2 * border and border > 0
+    if padded:
+        mix = F.pad(mix, (border, border), mode="reflect")
+    total = mix.shape[-1]
+
+    window = _fade_window(chunk_size, fade_size).to(device)
+    result = torch.zeros((model.num_stems, *mix.shape), device=device)
+    counter = torch.zeros_like(result)
+
+    use_autocast = dtype is not None and device.type == "cuda"
+    positions = list(range(0, total, step))
+    if progress:
+        from tqdm import tqdm
+        batches = tqdm(range(0, len(positions), batch_size), desc="separating", leave=False)
+    else:
+        batches = range(0, len(positions), batch_size)
+
+    for b in batches:
+        batch_pos = positions[b:b + batch_size]
+        parts, lens = [], []
+        for start in batch_pos:
+            part = mix[:, start:start + chunk_size]
+            chunk_len = part.shape[-1]
+            pad_mode = "reflect" if chunk_len > chunk_size // 2 else "constant"
+            parts.append(F.pad(part, (0, chunk_size - chunk_len), mode=pad_mode))
+            lens.append(chunk_len)
+        if len(parts) == 1:
+            parts.append(parts[0])  # a batch of 1 would make dynamo specialize (0/1) and recompile
+
+        with torch.autocast(device_type="cuda", dtype=dtype, enabled=use_autocast):
+            out = model(torch.stack(parts))
+        if out.ndim == 3:
+            out = out[:, None]  # [b, n, c, t]
+        out = out.float()
+
+        for j, (start, seg_len) in enumerate(zip(batch_pos, lens)):
+            w = window.clone()
+            if start == 0:
+                w[:fade_size] = 1
+            if start + step >= total:
+                w[-fade_size:] = 1
+            result[..., start:start + seg_len] += out[j, ..., :seg_len] * w[:seg_len]
+            counter[..., start:start + seg_len] += w[:seg_len]
+
+    stems = torch.nan_to_num(result / counter)
+    if padded:
+        stems = stems[..., border:-border]
+    return stems.cpu()
+
+
+@torch.no_grad()
+def separate(model: MelBandRoformer, wav: torch.Tensor, sample_rate: int, **kwargs):
+    """wav: [C, T] or [T] at any sample rate. Returns (vocals, instrumental), same shape/sr/device as wav,
+    with vocals + instrumental == wav exactly."""
+    if wav.ndim == 1:
+        wav = wav[None]
+    wav = wav.float()
+    channels, length = wav.shape
+    if channels > 2:
+        raise ValueError(f"expected mono or stereo audio, got {channels} channels")
+
+    mix = wav.repeat(2, 1) if channels == 1 else wav
+    if sample_rate != MODEL_SAMPLE_RATE:
+        import torchaudio
+        mix = torchaudio.functional.resample(mix, sample_rate, MODEL_SAMPLE_RATE)
+
+    vocals = separate_stems(model, mix, **kwargs)[0]
+
+    if sample_rate != MODEL_SAMPLE_RATE:
+        vocals = torchaudio.functional.resample(vocals, MODEL_SAMPLE_RATE, sample_rate)
+    vocals = F.pad(vocals[..., :length], (0, max(0, length - vocals.shape[-1])))
+    if channels == 1:
+        vocals = vocals.mean(0, keepdim=True)
+    vocals = vocals.to(wav.device)
+
+    instrumental = wav - vocals
+    return vocals, instrumental