From 080873bec9df0bef663897aac4cbd67a4d19f60b Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Thu, 17 Sep 2026 09:49:19 -0600 Subject: [PATCH] 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. --- .../captioner/AceStepCaptioner.py | 255 +++++++++- scripts/convert_melbandroformer.py | 127 +++++ toolkit/audio/melbandroformer/HF_README.md | 52 ++ toolkit/audio/melbandroformer/LICENSE | 28 ++ toolkit/audio/melbandroformer/__init__.py | 12 + toolkit/audio/melbandroformer/__main__.py | 120 +++++ toolkit/audio/melbandroformer/model.py | 470 ++++++++++++++++++ toolkit/audio/melbandroformer/separate.py | 158 ++++++ ui/src/components/CaptionSimpleJob.tsx | 7 + ui/src/helpers/captionOptions.ts | 3 +- ui/src/types.ts | 1 + 11 files changed, 1218 insertions(+), 15 deletions(-) create mode 100644 scripts/convert_melbandroformer.py create mode 100644 toolkit/audio/melbandroformer/HF_README.md create mode 100644 toolkit/audio/melbandroformer/LICENSE create mode 100644 toolkit/audio/melbandroformer/__init__.py create mode 100644 toolkit/audio/melbandroformer/__main__.py create mode 100644 toolkit/audio/melbandroformer/model.py create mode 100644 toolkit/audio/melbandroformer/separate.py diff --git a/extensions_built_in/captioner/AceStepCaptioner.py b/extensions_built_in/captioner/AceStepCaptioner.py index cac935d..1f7f78f 100644 --- a/extensions_built_in/captioner/AceStepCaptioner.py +++ b/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: diff --git a/scripts/convert_melbandroformer.py b/scripts/convert_melbandroformer.py new file mode 100644 index 0000000..26c5356 --- /dev/null +++ b/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() diff --git a/toolkit/audio/melbandroformer/HF_README.md b/toolkit/audio/melbandroformer/HF_README.md new file mode 100644 index 0000000..3e4832c --- /dev/null +++ b/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. diff --git a/toolkit/audio/melbandroformer/LICENSE b/toolkit/audio/melbandroformer/LICENSE new file mode 100644 index 0000000..064c15c --- /dev/null +++ b/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. diff --git a/toolkit/audio/melbandroformer/__init__.py b/toolkit/audio/melbandroformer/__init__.py new file mode 100644 index 0000000..33fb50b --- /dev/null +++ b/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", +] diff --git a/toolkit/audio/melbandroformer/__main__.py b/toolkit/audio/melbandroformer/__main__.py new file mode 100644 index 0000000..0a38e7e --- /dev/null +++ b/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 _vocals. and _instrumental. 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() diff --git a/toolkit/audio/melbandroformer/model.py b/toolkit/audio/melbandroformer/model.py new file mode 100644 index 0000000..17338dd --- /dev/null +++ b/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) diff --git a/toolkit/audio/melbandroformer/separate.py b/toolkit/audio/melbandroformer/separate.py new file mode 100644 index 0000000..d65a972 --- /dev/null +++ b/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/, 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 diff --git a/ui/src/components/CaptionSimpleJob.tsx b/ui/src/components/CaptionSimpleJob.tsx index a3941df..ef05ec1 100644 --- a/ui/src/components/CaptionSimpleJob.tsx +++ b/ui/src/components/CaptionSimpleJob.tsx @@ -276,6 +276,13 @@ const CaptionSimpleJob: React.FC = ({ 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') && ( + setJobConfig(value, 'config.process[0].caption.extract_vocals_before_transcribe')} + /> + )} {additionalSections.includes('caption.thinking') && (