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.

extensions_built_in/captioner/AceStepCaptioner.py

@@ -4,10 +4,13 @@ try:
     import librosa
 except ImportError:
     librosa = None
+import threading
+
 import numpy as np
 import torch
 import torchaudio
 from transformers import Qwen2_5OmniForConditionalGeneration, Qwen2_5OmniProcessor
+from transformers.generation import LogitsProcessor, LogitsProcessorList
 from collections import OrderedDict
 import traceback
 
@@ -33,10 +36,39 @@ CAPTION_PROMPT = "*Task* Describe this music in detail. Include genre, mood, ins
 
 # same budget the Qwen2.5-Omni top-level generate used (thinker_max_new_tokens)
 MAX_NEW_TOKENS = 1024
+# lyric transcription loops on repeated musical phrases more than captioning
+TRANSCRIBER_REPETITION_PENALTY = 1.05
+# whisper-style temperature fallback when greedy decode degenerates into a
+# token loop; the penalty alone cannot break a confident single-token loop
+TRANSCRIBER_RETRIES = (
+    {"do_sample": True, "temperature": 0.4, "repetition_penalty": 1.20},
+    {"do_sample": True, "temperature": 0.7, "repetition_penalty": 1.30},
+    {"do_sample": True, "temperature": 1.0, "repetition_penalty": 1.50},
+)
+# a loop is a unit of up to LOOP_MAX_PERIOD tokens repeated at least
+# LOOP_REPEATS times over a span of at least LOOP_MIN_SPAN tokens; sampled
+# decode jitters a token now and then, so the match is a fraction, not exact.
+# real outros repeat a line up to ~12 times (~128 tokens); loops never stop
+LOOP_MAX_PERIOD = 32
+LOOP_REPEATS = 4
+LOOP_MIN_SPAN = 160
+LOOP_MATCH_FRACTION = 0.9
+# on a detected loop the unit's tokens are banned for max(2*period, this) steps
+LOOP_BAN_STEPS = 8
+# loops broken before decode is forced to stop
+LOOP_BREAKS_BEFORE_STOP = 2
+# repeats of the looping unit kept when truncating
+LOOP_KEEP_REPEATS = 2
+CAPTIONER_REPETITION_PENALTY = 1.05
 # fixed cache length under compiled decode: the processor caps audio at 300s
 # (7500 audio tokens) so prompt + MAX_NEW_TOKENS always fits; a constant keeps
 # the static kv cache (and the compiled decode graph) at one shape for every file
 STATIC_MAX_LENGTH = 8704
+# silence trim on the extracted vocal stem: frames below this fraction of the
+# peak frame RMS (-40 dB) at the head and tail are cut, keeping a short margin
+SILENCE_FRAME_SECONDS = 0.02
+SILENCE_RMS_FRACTION = 0.01
+SILENCE_MARGIN_SECONDS = 0.25
 
 # Key profiles for Krumhansl-Schmuckler key detection
 MAJOR_PROFILE = np.array(
@@ -132,6 +164,88 @@ from extensions_built_in.llm_models.src.thinker import (  # noqa: E402  shared f
 CAPTION_FORMATS = ("ace_step", "yue2")
 
 
+def find_loop_period(tokens) -> int:
+    """Period of a unit repeated across the tail of the token ids, 0 if none."""
+    arr = np.asarray(tokens)
+    for period in range(1, LOOP_MAX_PERIOD + 1):
+        span = max(LOOP_REPEATS * period, LOOP_MIN_SPAN)
+        if span > arr.shape[0]:
+            continue
+        tail = arr[-span:]
+        if np.mean(tail[period:] == tail[:-period]) >= LOOP_MATCH_FRACTION:
+            return period
+    return 0
+
+
+def find_loop_start(tokens, period: int) -> int:
+    """Index where the periodic run at the tail of the token ids begins."""
+    # a mismatch run longer than this is real content, shorter is jitter
+    tolerance = max(period, 4)
+    start = len(tokens) - 1
+    mismatch_run = 0
+    while start - period >= 0:
+        if tokens[start] == tokens[start - period]:
+            mismatch_run = 0
+        else:
+            mismatch_run += 1
+            if mismatch_run > tolerance:
+                start += mismatch_run  # back out of the real content
+                break
+        start -= 1
+    return start + 1
+
+
+class LoopBreaker(LogitsProcessor):
+    """Breaks token loops inside a single decode: bans the repeating unit's
+    tokens for a stretch so the model has to move on, and forces EOS once it
+    has looped LOOP_BREAKS_BEFORE_STOP times. Batch size 1 only."""
+
+    def __init__(self, prompt_len: int, eos_ids):
+        self.prompt_len = prompt_len
+        self.eos_ids = list(eos_ids)
+        self.generated: list = []
+        self.banned: list = []
+        self.ban_until = 0
+        self.window_start = 0  # detection only looks at tokens after the last break
+        self.breaks = 0
+        self.stopped = False
+        # first loop: where it began, its period, and where decode resumed after the ban
+        self.loop_start = 0
+        self.loop_period = 0
+        self.loop_resume = 0
+
+    def __call__(self, input_ids, scores):
+        n = input_ids.shape[1] - self.prompt_len
+        if n > len(self.generated):
+            self.generated.extend(input_ids[0, self.prompt_len + len(self.generated):].tolist())
+        if self.stopped:
+            return self._force_eos(scores)
+        if n < self.ban_until:
+            scores[:, self.banned] = float("-inf")
+            return scores
+        period = find_loop_period(self.generated[self.window_start:])
+        if not period:
+            return scores
+        self.breaks += 1
+        if self.breaks >= LOOP_BREAKS_BEFORE_STOP:
+            self.stopped = True
+            return self._force_eos(scores)
+        self.banned = sorted(set(self.generated[-period:]))
+        self.ban_until = n + max(2 * period, LOOP_BAN_STEPS)
+        self.window_start = self.ban_until
+        if self.breaks == 1:
+            self.loop_start = find_loop_start(self.generated, period)
+            self.loop_period = period
+            self.loop_resume = self.ban_until
+        scores[:, self.banned] = float("-inf")
+        return scores
+
+    def _force_eos(self, scores):
+        scores[:] = float("-inf")
+        scores[:, self.eos_ids] = 0.0
+        return scores
+
+
 def clean_lyrics(text: str) -> str:
     """Strip the transcriber's '# Languages ... # Lyrics' header, keep the sectioned lyrics."""
     if "# Lyrics" in text:
@@ -146,6 +260,11 @@ class AceStepCaptionConfig(CaptionConfig):
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
         self.fixed_caption: Optional[str] = kwargs.get("fixed_caption", None)
+        # transcribe a MelBandRoformer vocal stem instead of the full mix; the
+        # captioner still hears the full mix
+        self.extract_vocals_before_transcribe: bool = bool(
+            kwargs.get("extract_vocals_before_transcribe", False)
+        )
         self.caption_format: str = kwargs.get("caption_format", "ace_step") or "ace_step"
         if self.caption_format not in CAPTION_FORMATS:
             raise ValueError(
@@ -159,6 +278,9 @@ class AceStepCaptioner(BaseCaptioner):
 
     def __init__(self, process_id: int, job, config: OrderedDict, **kwargs):
         super(AceStepCaptioner, self).__init__(process_id, job, config, **kwargs)
+        self.separator = None
+        # prep runs in worker threads; one separation on the GPU at a time
+        self.separator_lock = threading.Lock()
 
     def _load_thinker(self, name_or_path: str, label: str):
         """Load a Qwen2.5-Omni checkpoint and keep only its thinker."""
@@ -190,6 +312,11 @@ class AceStepCaptioner(BaseCaptioner):
             self.model2, self.processor2 = self._load_thinker(
                 self.caption_config.model_name_or_path2, "captioner"
             )
+        if self.caption_config.extract_vocals_before_transcribe:
+            from toolkit.audio.melbandroformer import load_melbandroformer
+
+            self.print_and_status_update("Loading vocal separator")
+            self.separator = load_melbandroformer(device=self.device_torch)
         flush()
 
     def maybe_compile_models(self):
@@ -238,25 +365,59 @@ class AceStepCaptioner(BaseCaptioner):
         )
 
     @staticmethod
-    def _load_audio(file_path: str):
-        waveform, sr = torchaudio.load(file_path)
+    def _to_model_audio(waveform: torch.Tensor, sr: int) -> np.ndarray:
+        """[C, T] at any rate -> mono float array at TARGET_SAMPLE_RATE."""
         if waveform.shape[0] > 1:
             waveform = waveform.mean(dim=0, keepdim=True)
         if sr != TARGET_SAMPLE_RATE:
             waveform = torchaudio.functional.resample(waveform, sr, TARGET_SAMPLE_RATE)
         return waveform.squeeze(0).numpy()
 
+    @staticmethod
+    def _load_audio(file_path: str):
+        waveform, sr = torchaudio.load(file_path)
+        return AceStepCaptioner._to_model_audio(waveform, sr)
+
+    @staticmethod
+    def trim_silence(audio: np.ndarray, sample_rate: int = TARGET_SAMPLE_RATE) -> np.ndarray:
+        """Cut leading and trailing silence, judged by frame RMS against the loudest frame."""
+        frame = max(1, int(sample_rate * SILENCE_FRAME_SECONDS))
+        n_frames = len(audio) // frame
+        if n_frames < 2:
+            return audio
+        rms = np.sqrt(np.mean(audio[: n_frames * frame].reshape(n_frames, frame) ** 2, axis=1))
+        loud = np.flatnonzero(rms >= rms.max() * SILENCE_RMS_FRACTION)
+        if len(loud) == 0:
+            return audio
+        margin = int(sample_rate * SILENCE_MARGIN_SECONDS)
+        start = max(0, loud[0] * frame - margin)
+        end = min(len(audio), (loud[-1] + 1) * frame + margin)
+        return audio[start:end]
+
+    def _extract_vocals(self, waveform: torch.Tensor, sr: int) -> np.ndarray:
+        """Vocal stem of [C, T] audio as mono TARGET_SAMPLE_RATE, silence-trimmed."""
+        from toolkit.audio.melbandroformer import separate
+
+        with self.separator_lock:
+            vocals, _ = separate(self.separator, waveform, sr)
+        return self.trim_silence(self._to_model_audio(vocals.cpu(), sr))
+
     def _prep_file(self, file_path: str) -> dict:
         """CPU side of one file (runs in a worker thread): librosa analysis,
-        decode + resample, and the mel/tokenize pass for each model."""
-        audio_data = self._load_audio(file_path)
+        decode + resample, and the mel/tokenize pass for each model. Vocal
+        separation is the one GPU step here, serialized by separator_lock."""
+        waveform, sr = torchaudio.load(file_path)
+        audio_data = self._to_model_audio(waveform, sr)
+        lyrics_audio = (
+            self._extract_vocals(waveform, sr) if self.separator is not None else audio_data
+        )
         item = {
             "file": file_path,
             # BPM/key/time signature only appear in the ace_step layout
             "analysis": analyze_audio(file_path)
             if self.caption_config.caption_format == "ace_step"
             else None,
-            "lyrics_inputs": self._make_inputs(self.processor, audio_data, LYRICS_PROMPT),
+            "lyrics_inputs": self._make_inputs(self.processor, lyrics_audio, LYRICS_PROMPT),
         }
         if self.model2 is not None:
             item["caption_inputs"] = self._make_inputs(
@@ -264,7 +425,16 @@ class AceStepCaptioner(BaseCaptioner):
             )
         return item
 
-    def _generate(self, model, processor, inputs) -> str:
+    def _generate(
+        self,
+        model,
+        processor,
+        inputs,
+        repetition_penalty: float,
+        retries: tuple = (),
+        break_loops: bool = False,
+        file_path: str = "",
+    ) -> str:
         inputs = inputs.to(model.device).to(model.dtype)
         # a generate that dies between static-cache creation and its first
         # forward leaves a half-built cache that breaks every later call
@@ -273,12 +443,14 @@ class AceStepCaptioner(BaseCaptioner):
             del model._cache
         model._pad_mask_2d = inputs.get("attention_mask", None)
         input_len = inputs["input_ids"].shape[1]
-        gen_kwargs = {"max_new_tokens": MAX_NEW_TOKENS}
+        # greedy decode loops on repeated phrases without a penalty
+        gen_kwargs = {"max_new_tokens": MAX_NEW_TOKENS, "repetition_penalty": repetition_penalty}
         if model.generation_config.cache_implementation == "static":
             if input_len + 16 < STATIC_MAX_LENGTH:
                 from transformers.generation import MaxLengthCriteria, StoppingCriteriaList
 
                 gen_kwargs = {
+                    "repetition_penalty": repetition_penalty,
                     "max_length": STATIC_MAX_LENGTH,
                     "stopping_criteria": StoppingCriteriaList(
                         [MaxLengthCriteria(max_length=min(input_len + MAX_NEW_TOKENS, STATIC_MAX_LENGTH))]
@@ -287,15 +459,62 @@ class AceStepCaptioner(BaseCaptioner):
             else:
                 # prompt would not fit the fixed cache; eager for this file
                 gen_kwargs["cache_implementation"] = "dynamic"
-        generated = model.generate(**inputs, **gen_kwargs)
+        eos_ids = model.generation_config.eos_token_id
+        eos_ids = set(eos_ids if isinstance(eos_ids, (list, tuple)) else [eos_ids])
+
+        def strip_eos(tokens: torch.Tensor) -> torch.Tensor:
+            end = tokens.shape[0]
+            while end > 0 and int(tokens[end - 1]) in eos_ids:
+                end -= 1
+            return tokens[:end]
+
+        def run_pass():
+            breaker = LoopBreaker(input_len, eos_ids) if break_loops else None
+            kwargs = dict(gen_kwargs)
+            if breaker is not None:
+                kwargs["logits_processor"] = LogitsProcessorList([breaker])
+            model._pad_mask_2d = inputs.get("attention_mask", None)
+            generated = model.generate(**inputs, **kwargs)
+            return strip_eos(generated[0, input_len:]), breaker
+
+        new_tokens, breaker = run_pass()
+        first_tokens, first_breaker = new_tokens, breaker
+        stopped = breaker is not None and breaker.stopped
+        for attempt, retry in enumerate(retries, 1):
+            if not stopped:
+                break
+            print(
+                f"\n[AITK] {file_path}: transcription looped {breaker.breaks} times and was stopped, "
+                f"retry {attempt}/{len(retries)} with temperature={retry['temperature']} "
+                f"repetition_penalty={retry['repetition_penalty']}"
+            )
+            gen_kwargs.update(retry)
+            new_tokens, breaker = run_pass()
+            stopped = breaker.stopped
+        if stopped:
+            # every pass re-looped after a break; keep the greedy pass up to its first loop
+            keep = first_breaker.loop_start + LOOP_KEEP_REPEATS * first_breaker.loop_period
+            new_tokens = first_tokens[:keep]
+            print(
+                f"\n[AITK] {file_path}: transcription looped on every pass; keeping the greedy "
+                f"pass truncated at its first loop (period {first_breaker.loop_period} tokens)"
+            )
+        elif breaker is not None and breaker.breaks:
+            # drop the excess repeats and the tokens forced out under the ban
+            keep = breaker.loop_start + LOOP_KEEP_REPEATS * breaker.loop_period
+            new_tokens = torch.cat([new_tokens[:keep], new_tokens[breaker.loop_resume :]])
+            print(
+                f"\n[AITK] {file_path}: broke a transcription loop "
+                f"(period {breaker.loop_period} tokens) and continued"
+            )
         output = processor.batch_decode(
-            generated[:, input_len:],
+            new_tokens[None],
             skip_special_tokens=True,
             clean_up_tokenization_spaces=False,
         )[0]
         return output.strip()
 
-    def get_audio_lyrics(self, inputs) -> str:
+    def get_audio_lyrics(self, inputs, file_path: str = "") -> str:
         if (
             self.caption_config.low_vram
             and self.model2 is not None
@@ -304,18 +523,26 @@ class AceStepCaptioner(BaseCaptioner):
             self.model2.to("cpu")
         if self.model.device == torch.device("cpu"):
             self.model.to(self.device_torch)
-        return self._generate(self.model, self.processor, inputs)
+        return self._generate(
+            self.model,
+            self.processor,
+            inputs,
+            TRANSCRIBER_REPETITION_PENALTY,
+            retries=TRANSCRIBER_RETRIES,
+            break_loops=True,
+            file_path=file_path,
+        )
 
     def get_audio_caption(self, inputs) -> str:
         if self.caption_config.low_vram and self.model.device != torch.device("cpu"):
             self.model.to("cpu")
         if self.model2.device == torch.device("cpu"):
             self.model2.to(self.device_torch)
-        return self._generate(self.model2, self.processor2, inputs)
+        return self._generate(self.model2, self.processor2, inputs, CAPTIONER_REPETITION_PENALTY)
 
     def _caption_item(self, item: dict) -> str:
         if self.caption_config.caption_format == "yue2":
-            lyrics = clean_lyrics(self.get_audio_lyrics(item["lyrics_inputs"]))
+            lyrics = clean_lyrics(self.get_audio_lyrics(item["lyrics_inputs"], item["file"]))
             if self.caption_config.fixed_caption is not None:
                 caption = self.caption_config.fixed_caption
             else:
@@ -324,7 +551,7 @@ class AceStepCaptioner(BaseCaptioner):
             return f"{caption}\n[Lyrics]\n{lyrics}"
 
         analysis = item["analysis"]
-        lyrics = self.get_audio_lyrics(item["lyrics_inputs"])
+        lyrics = self.get_audio_lyrics(item["lyrics_inputs"], item["file"])
 
         language = "en"
         if "# Languages" in lyrics and "# Lyrics" in lyrics:

scripts/convert_melbandroformer.py

@@ -0,0 +1,127 @@
+"""Convert an MSST / lucidrains Mel-Band RoFormer .ckpt into a self-describing safetensors file
+(model kwargs + inference defaults in the metadata) that toolkit.audio.melbandroformer loads.
+
+    python scripts/convert_melbandroformer.py                       # KimberleyJSN vocals model -> MODELS_PATH/checkpoints
+    python scripts/convert_melbandroformer.py --ckpt x.ckpt --config msst_config.yaml --out x.safetensors
+"""
+import argparse
+import inspect
+import json
+import os
+import sys
+
+TOOLKIT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+sys.path.insert(0, TOOLKIT_ROOT)
+
+from dotenv import load_dotenv
+
+load_dotenv(os.path.join(TOOLKIT_ROOT, ".env"))  # MODELS_PATH etc., before toolkit.paths is imported
+
+import torch
+import yaml
+from safetensors.torch import save_file
+
+from toolkit.paths import MODELS_PATH
+from toolkit.audio.melbandroformer.model import MelBandRoformer
+from toolkit.audio.melbandroformer.separate import DEFAULT_WEIGHTS
+
+# https://huggingface.co/KimberleyJSN/melbandroformer (MIT), config from
+# MSST configs/KimberleyJensen/config_vocals_mel_band_roformer_kj.yaml
+KJ_VOCALS = dict(
+    repo_id="KimberleyJSN/melbandroformer",
+    filename="MelBandRoformer.ckpt",
+    license="MIT",
+    model=dict(
+        dim=384,
+        depth=6,
+        stereo=True,
+        num_stems=1,
+        time_transformer_depth=1,
+        freq_transformer_depth=1,
+        num_bands=60,
+        dim_head=64,
+        heads=8,
+        attn_dropout=0,
+        ff_dropout=0,
+        sample_rate=44100,
+        stft_n_fft=2048,
+        stft_hop_length=441,
+        stft_win_length=2048,
+        stft_normalized=False,
+        mask_estimator_depth=2,
+    ),
+    inference=dict(chunk_size=352800, num_overlap=2),
+    stems=["vocals"],
+)
+
+
+def config_from_msst_yaml(path):
+    with open(path) as f:
+        cfg = yaml.load(f, Loader=yaml.FullLoader)  # MSST configs carry !!python/tuple
+    allowed = set(inspect.signature(MelBandRoformer.__init__).parameters) - {"self"}
+    model = {k: v for k, v in cfg["model"].items() if k in allowed}
+    dropped = sorted(set(cfg["model"]) - allowed)
+    if dropped:
+        print(f"ignoring non-inference model keys: {dropped}")
+    assert cfg["model"].get("linear_transformer_depth", 0) == 0, "linear transformer layers are not vendored"
+    inference = dict(
+        chunk_size=cfg.get("inference", {}).get("chunk_size", cfg["audio"]["chunk_size"]),
+        num_overlap=cfg["inference"]["num_overlap"],
+    )
+    training = cfg.get("training", {})
+    stems = [training["target_instrument"]] if training.get("target_instrument") else training.get("instruments", [])
+    return model, inference, stems
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    parser.add_argument("--ckpt", default=None, help="local .ckpt (default: download KimberleyJSN/melbandroformer)")
+    parser.add_argument("--config", default=None, help="MSST yaml config (default: KJ vocals config)")
+    parser.add_argument("--out", default=os.path.join(MODELS_PATH, "checkpoints", DEFAULT_WEIGHTS))
+    parser.add_argument("--dtype", default="fp32", choices=["fp32", "fp16", "bf16"])
+    args = parser.parse_args()
+
+    if args.config:
+        model_cfg, inference_cfg, stems = config_from_msst_yaml(args.config)
+        source = args.ckpt
+        license_ = ""
+    else:
+        model_cfg, inference_cfg, stems = KJ_VOCALS["model"], KJ_VOCALS["inference"], KJ_VOCALS["stems"]
+        source = f"{KJ_VOCALS['repo_id']}/{KJ_VOCALS['filename']}"
+        license_ = KJ_VOCALS["license"]
+
+    ckpt = args.ckpt
+    if ckpt is None:
+        from huggingface_hub import hf_hub_download
+        ckpt = hf_hub_download(KJ_VOCALS["repo_id"], KJ_VOCALS["filename"])
+
+    print(f"loading {ckpt}")
+    state_dict = torch.load(ckpt, map_location="cpu", weights_only=True)
+    if "state_dict" in state_dict:
+        state_dict = state_dict["state_dict"]
+    state_dict = {k[len("module."):] if k.startswith("module.") else k: v for k, v in state_dict.items()}
+
+    # validate against the vendored module before writing anything
+    model = MelBandRoformer(**model_cfg)
+    model.load_state_dict(state_dict, strict=True)
+
+    dtype = {"fp32": torch.float32, "fp16": torch.float16, "bf16": torch.bfloat16}[args.dtype]
+    # shared rotary freqs alias one storage in the ckpt; safetensors needs distinct tensors
+    tensors = {k: v.detach().to(dtype).clone().contiguous() for k, v in state_dict.items()}
+
+    metadata = dict(
+        klass="MelBandRoformer",
+        config=json.dumps(model_cfg),
+        inference=json.dumps(inference_cfg),
+        stems=json.dumps(stems),
+        source=str(source),
+        license=license_,
+    )
+    os.makedirs(os.path.dirname(args.out), exist_ok=True)
+    save_file(tensors, args.out, metadata=metadata)
+    n_params = sum(v.numel() for v in tensors.values())
+    print(f"wrote {args.out} ({n_params / 1e6:.1f}M params, {args.dtype}, {os.path.getsize(args.out) / 1e6:.0f} MB)")
+
+
+if __name__ == "__main__":
+    main()

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

ui/src/components/CaptionSimpleJob.tsx

@@ -276,6 +276,13 @@ const CaptionSimpleJob: React.FC<Props> = ({ jobConfig, setJobConfig, gpuIDs, se
               checked={jobConfig.config.process[0].caption.compile || false}
               onChange={value => setJobConfig(value, 'config.process[0].caption.compile')}
             />
+            {additionalSections.includes('caption.extract_vocals_before_transcribe') && (
+              <Checkbox
+                label="Extract Vocals Before Transcribing"
+                checked={jobConfig.config.process[0].caption.extract_vocals_before_transcribe || false}
+                onChange={value => setJobConfig(value, 'config.process[0].caption.extract_vocals_before_transcribe')}
+              />
+            )}
             {additionalSections.includes('caption.thinking') && (
               <Checkbox
                 label="Thinking"

ui/src/helpers/captionOptions.ts

@@ -1,7 +1,7 @@
 import { CloudLora, GroupedSelectOption, SelectOption } from "@/types";
 
 type CaptionGroup = 'image' | 'music' | 'video' | 'image/video/sound';
-type AdditionalSections = 'caption.model_name_or_path2' | 'caption.caption_prompt' | 'caption.max_res' | 'caption.max_new_tokens' | 'caption.fixed_caption' | 'caption.caption_format' | 'caption.thinking' | 'caption.batch_size' | 'caption.layer_offloading';
+type AdditionalSections = 'caption.model_name_or_path2' | 'caption.caption_prompt' | 'caption.max_res' | 'caption.max_new_tokens' | 'caption.fixed_caption' | 'caption.caption_format' | 'caption.extract_vocals_before_transcribe' | 'caption.thinking' | 'caption.batch_size' | 'caption.layer_offloading';
 
 export interface CaptionOption {
     name: string;
@@ -111,6 +111,7 @@ export const captionerTypes: CaptionOption[] = [
             'caption.model_name_or_path2',
             'caption.fixed_caption',
             'caption.caption_format',
+            'caption.extract_vocals_before_transcribe',
         ],
     },
     {

ui/src/types.ts

@@ -338,6 +338,7 @@ export interface CaptionProcessConfig {
     max_new_tokens?: number;
     fixed_caption?: string;
     caption_format?: string;
+    extract_vocals_before_transcribe?: boolean;
     caption_extension?: string;
     thinking?: boolean;
     batch_size?: number;