Code actualized

mvsepless/additional_app.py

@@ -279,6 +279,7 @@ class CustomSeparator(Separator, GradioHelper):
         use_spec_invert: bool = False,
         econom_mode: Optional[bool] = None,
         chunk_duration: float = 300,
+        denoise: bool = False,
         progress: gr.Progress = gr.Progress(track_tqdm=True),
     ) -> List:
         """
@@ -295,7 +296,7 @@ class CustomSeparator(Separator, GradioHelper):
             vr_aggr: Агрессивность для VR
             vr_post_process: Постобработка для VR
             vr_high_end_process: Обработка высоких частот для VR
-            mdx_denoise: Шумоподавление для MDX
+            denoise: Шумоподавление для MDX
             use_spec_invert: Использовать инверсию спектрограммы
             econom_mode: Эконом-режим
             chunk_duration: Длительность чанка
@@ -309,7 +310,8 @@ class CustomSeparator(Separator, GradioHelper):
         
         add_settings: Dict[str, Any] = {
             "add_single_sep_text_progress": None,
-            "single_mode": False
+            "single_mode": False,
+            "denoise": denoise
         }
         
         if econom_mode is not None:
@@ -412,6 +414,11 @@ class CustomSeparator(Separator, GradioHelper):
                                     value=False, 
                                     interactive=True
                                 )
+                                denoise = gr.Checkbox(
+                                    label=_i18n("denoise"),
+                                    value=False,
+                                    interactive=True,
+                                )
                                 
                                 gr.Markdown(f"<h4>{_i18n('economy_settings')}</h4>", container=True)
                                 econom_mode = gr.Checkbox(
@@ -497,7 +504,8 @@ class CustomSeparator(Separator, GradioHelper):
                                 stems,
                                 use_spec_for_extract_instrumental,
                                 econom_mode,
-                                chunk_dur_slider
+                                chunk_dur_slider,
+                                denoise
                             ],
                             outputs=[sep_state, status],
                             show_progress="full",
@@ -514,6 +522,7 @@ class CustomSeparator(Separator, GradioHelper):
                             u_spec: bool,
                             ec_mode: bool,
                             ch_dur: float,
+                            den: bool,
                             progress: gr.Progress = gr.Progress(track_tqdm=True),
                         ) -> Tuple[gr.update, gr.update]:
                             results = self._separate_batch(
@@ -527,6 +536,7 @@ class CustomSeparator(Separator, GradioHelper):
                                 u_spec,
                                 ec_mode,
                                 ch_dur * 60,
+                                den,
                                 progress=progress,
                             )
                             return gr.update(value=str(results)), gr.update(visible=False)
@@ -801,10 +811,11 @@ class CustomSeparator(Separator, GradioHelper):
         def refresh_all_models() -> Tuple[gr.update, gr.update, gr.update]:
             models = self.get_mn()
             first_model = models[0] if models else None
+            first_model2 = [models[0]] if models else []
             return (
                 gr.update(choices=models, value=first_model),
                 gr.update(choices=models, value=first_model),
-                gr.update(choices=models, value=[first_model]),
+                gr.update(choices=models, value=first_model2),
             )
 
 class AutoEnsembless(Separator, GradioHelper):

mvsepless/app.py

@@ -278,7 +278,7 @@ class SeparatorGradio(GradioHelper, DownloadModelManager):
         vr_aggr: int = 5,
         vr_post_process: bool = False,
         vr_high_end_process: bool = False,
-        mdx_denoise: bool = False,
+        denoise: bool = False,
         use_spec_invert: bool = False,
         econom_mode: Optional[bool] = None,
         chunk_duration: float = 300,
@@ -298,7 +298,7 @@ class SeparatorGradio(GradioHelper, DownloadModelManager):
             vr_aggr: Агрессивность для VR
             vr_post_process: Постобработка для VR
             vr_high_end_process: Обработка высоких частот для VR
-            mdx_denoise: Шумоподавление для MDX
+            denoise: Шумоподавление для MDX
             use_spec_invert: Использовать инверсию спектрограммы
             econom_mode: Эконом-режим
             chunk_duration: Длительность чанка
@@ -311,7 +311,7 @@ class SeparatorGradio(GradioHelper, DownloadModelManager):
         self.chunk_duration = chunk_duration
         
         add_settings: Dict[str, Any] = {
-            "mdx_denoise": mdx_denoise,
+            "denoise": denoise,
             "vr_aggr": vr_aggr,
             "vr_post_process": vr_post_process,
             "vr_high_end_process": vr_high_end_process,
@@ -487,19 +487,17 @@ class SeparatorGradio(GradioHelper, DownloadModelManager):
                                         interactive=True
                                     )
                                     
-                                    gr.Markdown(f"<h4>{_i18n('mdx_settings')}</h4>", container=True)
-                                    mdx_denoise = gr.Checkbox(
-                                        label=_i18n("mdx_denoise"),
-                                        value=False,
-                                        interactive=True,
-                                    )
-                                    
                                     gr.Markdown(f"<h4>{_i18n('invert_settings')}</h4>", container=True)
                                     use_spec_for_extract_instrumental = gr.Checkbox(
                                         label=_i18n("use_spectrogram_invert"), 
                                         value=False, 
                                         interactive=True
                                     )
+                                    denoise = gr.Checkbox(
+                                        label=_i18n("denoise"),
+                                        value=False,
+                                        interactive=True,
+                                    )
                                     
                                     gr.Markdown(f"<h4>{_i18n('economy_settings')}</h4>", container=True)
                                     econom_mode = gr.Checkbox(
@@ -583,7 +581,7 @@ class SeparatorGradio(GradioHelper, DownloadModelManager):
                                     output_bitrate,
                                     template,
                                     stems,
-                                    mdx_denoise,
+                                    denoise,
                                     vr_aggr,
                                     vr_enable_post_process,
                                     vr_enable_high_end_process,
@@ -603,7 +601,7 @@ class SeparatorGradio(GradioHelper, DownloadModelManager):
                                 output_bitrate: int,
                                 template: str,
                                 stems: List[str],
-                                mdx_denoise: bool,
+                                denoise: bool,
                                 vr_aggr: int,
                                 vr_pp: bool,
                                 vr_hip: bool,
@@ -623,7 +621,7 @@ class SeparatorGradio(GradioHelper, DownloadModelManager):
                                     vr_aggr,
                                     vr_pp,
                                     vr_hip,
-                                    mdx_denoise,
+                                    denoise,
                                     u_spec,
                                     ec_mode,
                                     ch_dur * 60,

mvsepless/i18n.py

@@ -128,7 +128,7 @@ TRANSLATIONS: Dict[Language, Dict[str, str]] = {
         "vr_post_process": "Дополнительная обработка для улучшения качества разделения",
         "vr_high_end_process": "Восстановление недостающих высоких частот",
         "mdx_settings": "MDX-NET",
-        "mdx_denoise": "Шумоподавление",
+        "denoise": "Шумоподавление",
         "invert_settings": "Инвертирование результата",
         "use_spectrogram_invert": "При извлечении инструментала/второго стема/остатка использовать спектрограмму",
         "economy_settings": "Экономия",
@@ -441,7 +441,7 @@ TRANSLATIONS: Dict[Language, Dict[str, str]] = {
         "ext_inst_help": "Извлечь инструментал вычитанием",
         "use_spec_invert_help": "Инверсия спектрограммы для вторичного стема",
         "install_only_help": "Только установка модели",
-        "mdx_denoise_help": "Шумоподавление для MDX-NET моделей",
+        "denoise_help": "Удаление шума из вывода",
         "vr_aggression_help": "Агрессивность для VR моделей (по умолчанию: 5)",
         "vr_high_end_help": "Восстановление недостающих высоких частот на VR моделях",
         "vr_post_process_help": "Дополнительная обработка для улучшения качества разделения VR модели",
@@ -856,7 +856,7 @@ TRANSLATIONS: Dict[Language, Dict[str, str]] = {
         "vr_post_process": "Additional post-processing to improve separation quality",
         "vr_high_end_process": "Restore missing high frequencies",
         "mdx_settings": "MDX-NET",
-        "mdx_denoise": "Denoise",
+        "denoise": "Denoise",
         "invert_settings": "Result Inversion",
         "use_spectrogram_invert": "Use spectrogram for extracting instrumental/second stem/remainder",
         "economy_settings": "Economy",
@@ -1169,7 +1169,7 @@ TRANSLATIONS: Dict[Language, Dict[str, str]] = {
         "ext_inst_help": "Extract instrumental by subtraction",
         "use_spec_invert_help": "Use spectrogram inversion for secondary stem",
         "install_only_help": "Download model only",
-        "mdx_denoise_help": "Denoise for MDX-NET models",
+        "denoise_help": "Remove noise from output",
         "vr_aggression_help": "Aggression for VR models (default: 5)",
         "vr_high_end_help": "Restore missing high frequencies on VR models",
         "vr_post_process_help": "Additional post-processing for VR models",

mvsepless/infer_utils.py

@@ -1,3 +1,4 @@
+import os
 import sys
 sys.stdout.reconfigure(encoding='utf-8')
 sys.stderr.reconfigure(encoding='utf-8')
@@ -70,7 +71,9 @@ def get_model_from_config(model_type: str, config_path: str) -> Tuple[Any, Any]:
         from models.vr_arch import VRNet
         model = VRNet(**dict(config.model))
     elif model_type == "htdemucs":
-        from models.demucs4ht import get_model
+        models_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'models')
+        sys.path.append(models_path)
+        from demucs import get_model
         model = get_model(config)
     elif model_type == "mel_band_roformer":
         if hasattr(config, "windowed"):
@@ -419,6 +422,7 @@ def demix_demucs(
     mix = torch.tensor(mix, dtype=torch.float32)
     chunk_size = config.training.samplerate * config.training.segment
     num_instruments = len(config.training.instruments)
+    denoise = config.inference.denoise
     num_overlap = config.inference.num_overlap
     step = chunk_size // num_overlap
     fade_size = chunk_size // 10
@@ -451,8 +455,12 @@ def demix_demucs(
 
                 if len(batch_data) >= batch_size or i >= mix.shape[1]:
                     arr = torch.stack(batch_data, dim=0)
-                    x = model(arr)
-
+                    if denoise:
+                        x1 = model(arr)
+                        x2 = model(-arr)
+                        x = (x1 + -x2) * 0.5
+                    else:
+                        x = model(arr)
                     window = windowing_array.clone()
                     if i - step == 0:
                         window[:fade_size] = 1
@@ -511,6 +519,7 @@ def demix_generic(
     chunk_size = config.audio.chunk_size
     instruments = prefer_target_instrument(config)
     num_instruments = len(instruments)
+    denoise = config.inference.denoise
     num_overlap = config.inference.num_overlap
 
     fade_size = chunk_size // 10
@@ -550,7 +559,12 @@ def demix_generic(
 
                 if len(batch_data) >= batch_size or i >= mix.shape[1]:
                     arr = torch.stack(batch_data, dim=0)
-                    x = model(arr)
+                    if denoise:
+                        x1 = model(arr)
+                        x2 = model(-arr)
+                        x = (x1 + -x2) * 0.5
+                    else:
+                        x = model(arr)
 
                     window = windowing_array.clone()
                     if i - step == 0:

mvsepless/install.py

@@ -221,33 +221,28 @@ universal_requirements: List[str] = [
     "matplotlib",
     "tqdm",
     "einops",
-    "protobuf",
     "soundfile",
     "pydub",
     "webrtcvad",
     "audiomentations",
     "pedalboard",
     "ml_collections",
-    "timm",
     "wandb",
-    "accelerate",
     "bitsandbytes",
     "tokenizers",
     "huggingface-hub",
     "transformers",
-    "torchseg",
-    "demucs==4.0.0",
+    "diffq>=0.2.1",
+    "julius>=0.2.3",
+    "openunmix",
     "asteroid>=0.6.0",
     "pyloudnorm",
-    "prodigyopt",
-    "torch_log_wmse",
     "rotary_embedding_torch",
     "gradio<6.0",
     "omegaconf",
     "beartype",
     "spafe",
     "torch_audiomentations",
-    "auraloss",
     "onnx>=1.17",
     "onnx2torch>=0.3.0",
     "onnxruntime-gpu>=1.17" if cuda_available else "onnxruntime>=1.17",
@@ -279,32 +274,28 @@ old_requirements: List[str] = [
     "matplotlib==3.10.8",
     "tqdm==4.67.1",
     "einops==0.8.1",
-    "protobuf==6.33.4",
     "soundfile==0.13.1",
     "pydub==0.25.1",
     "webrtcvad==2.0.10",
     "audiomentations==0.43.1",
     "pedalboard==0.8.2",
     "ml_collections==1.1.0",
-    "timm==1.0.24",
     "wandb==0.24.0",
-    "accelerate==1.2.1",
     "bitsandbytes==0.45.0",
     "tokenizers==0.15.2",
     "huggingface-hub==0.34.2",
     "transformers==4.39.3",
-    "torchseg==0.0.1a4",
-    "demucs==4.0.0",
+    "diffq>=0.2.1",
+    "julius>=0.2.3",
+    "openunmix",
     "asteroid==0.6.0",
     "pyloudnorm",
-    "prodigyopt==1.1.2",
     "rotary_embedding_torch==0.3.6",
     "gradio<6.0.0",
     "omegaconf==2.3.0",
     "beartype==0.22.9",
     "spafe==0.3.3",
     "torch_audiomentations==0.12.0",
-    "auraloss==0.4.0",
     "onnx>=1.17",
     "onnx2torch>=0.3.0",
     "onnxruntime-gpu>=1.17" if cuda_available else "onnxruntime>=1.17",

mvsepless/models/bandit/core/__init__.py

@@ -1,669 +1,669 @@
-import os.path
-from collections import defaultdict
-from itertools import chain, combinations
-from typing import Any, Dict, Iterator, Mapping, Optional, Tuple, Type, TypedDict
-
-import pytorch_lightning as pl
-import torch
-import torchaudio as ta
-import torchmetrics as tm
-from asteroid import losses as asteroid_losses
-
-from pytorch_lightning.utilities.types import STEP_OUTPUT
-from torch import nn, optim
-from torch.optim import lr_scheduler
-from torch.optim.lr_scheduler import LRScheduler
-
-from . import loss, metrics as metrics_, model
-from .data._types import BatchedDataDict
-from .data.augmentation import BaseAugmentor, StemAugmentor
-from .utils import audio as audio_
-from .utils.audio import BaseFader
-
-
-ConfigDict = TypedDict("ConfigDict", {"name": str, "kwargs": Dict[str, Any]})
-
-
-class SchedulerConfigDict(ConfigDict):
-    monitor: str
-
-
-OptimizerSchedulerConfigDict = TypedDict(
-    "OptimizerSchedulerConfigDict",
-    {"optimizer": ConfigDict, "scheduler": SchedulerConfigDict},
-    total=False,
-)
-
-
-class LRSchedulerReturnDict(TypedDict, total=False):
-    scheduler: LRScheduler
-    monitor: str
-
-
-class ConfigureOptimizerReturnDict(TypedDict, total=False):
-    optimizer: torch.optim.Optimizer
-    lr_scheduler: LRSchedulerReturnDict
-
-
-OutputType = Dict[str, Any]
-MetricsType = Dict[str, torch.Tensor]
-
-
-def get_optimizer_class(name: str) -> Type[optim.Optimizer]:
-
-    if name == "DeepSpeedCPUAdam":
-        return DeepSpeedCPUAdam
-
-    for module in [optim, gooptim]:
-        if name in module.__dict__:
-            return module.__dict__[name]
-
-    raise NameError
-
-
-def parse_optimizer_config(
-    config: OptimizerSchedulerConfigDict, parameters: Iterator[nn.Parameter]
-) -> ConfigureOptimizerReturnDict:
-    optim_class = get_optimizer_class(config["optimizer"]["name"])
-    optimizer = optim_class(parameters, **config["optimizer"]["kwargs"])
-
-    optim_dict: ConfigureOptimizerReturnDict = {
-        "optimizer": optimizer,
-    }
-
-    if "scheduler" in config:
-
-        lr_scheduler_class_ = config["scheduler"]["name"]
-        lr_scheduler_class = lr_scheduler.__dict__[lr_scheduler_class_]
-        lr_scheduler_dict: LRSchedulerReturnDict = {
-            "scheduler": lr_scheduler_class(optimizer, **config["scheduler"]["kwargs"])
-        }
-
-        if lr_scheduler_class_ == "ReduceLROnPlateau":
-            lr_scheduler_dict["monitor"] = config["scheduler"]["monitor"]
-
-        optim_dict["lr_scheduler"] = lr_scheduler_dict
-
-    return optim_dict
-
-
-def parse_model_config(config: ConfigDict) -> Any:
-    name = config["name"]
-
-    for module in [model]:
-        if name in module.__dict__:
-            return module.__dict__[name](**config["kwargs"])
-
-    raise NameError
-
-
-_LEGACY_LOSS_NAMES = ["HybridL1Loss"]
-
-
-def _parse_legacy_loss_config(config: ConfigDict) -> nn.Module:
-    name = config["name"]
-
-    if name == "HybridL1Loss":
-        return loss.TimeFreqL1Loss(**config["kwargs"])
-
-    raise NameError
-
-
-def parse_loss_config(config: ConfigDict) -> nn.Module:
-    name = config["name"]
-
-    if name in _LEGACY_LOSS_NAMES:
-        return _parse_legacy_loss_config(config)
-
-    for module in [loss, nn.modules.loss, asteroid_losses]:
-        if name in module.__dict__:
-            return module.__dict__[name](**config["kwargs"])
-
-    raise NameError
-
-
-def get_metric(config: ConfigDict) -> tm.Metric:
-    name = config["name"]
-
-    for module in [tm, metrics_]:
-        if name in module.__dict__:
-            return module.__dict__[name](**config["kwargs"])
-    raise NameError
-
-
-def parse_metric_config(config: Dict[str, ConfigDict]) -> tm.MetricCollection:
-    metrics = {}
-
-    for metric in config:
-        metrics[metric] = get_metric(config[metric])
-
-    return tm.MetricCollection(metrics)
-
-
-def parse_fader_config(config: ConfigDict) -> BaseFader:
-    name = config["name"]
-
-    for module in [audio_]:
-        if name in module.__dict__:
-            return module.__dict__[name](**config["kwargs"])
-
-    raise NameError
-
-
-class LightningSystem(pl.LightningModule):
-    _VOX_STEMS = ["speech", "vocals"]
-    _BG_STEMS = ["background", "effects", "mne"]
-
-    def __init__(
-        self, config: Dict, loss_adjustment: float = 1.0, attach_fader: bool = False
-    ) -> None:
-        super().__init__()
-        self.optimizer_config = config["optimizer"]
-        self.model = parse_model_config(config["model"])
-        self.loss = parse_loss_config(config["loss"])
-        self.metrics = nn.ModuleDict(
-            {
-                stem: parse_metric_config(config["metrics"]["dev"])
-                for stem in self.model.stems
-            }
-        )
-
-        self.metrics.disallow_fsdp = True
-
-        self.test_metrics = nn.ModuleDict(
-            {
-                stem: parse_metric_config(config["metrics"]["test"])
-                for stem in self.model.stems
-            }
-        )
-
-        self.test_metrics.disallow_fsdp = True
-
-        self.fs = config["model"]["kwargs"]["fs"]
-
-        self.fader_config = config["inference"]["fader"]
-        if attach_fader:
-            self.fader = parse_fader_config(config["inference"]["fader"])
-        else:
-            self.fader = None
-
-        self.augmentation: Optional[BaseAugmentor]
-        if config.get("augmentation", None) is not None:
-            self.augmentation = StemAugmentor(**config["augmentation"])
-        else:
-            self.augmentation = None
-
-        self.predict_output_path: Optional[str] = None
-        self.loss_adjustment = loss_adjustment
-
-        self.val_prefix = None
-        self.test_prefix = None
-
-    def configure_optimizers(self) -> Any:
-        return parse_optimizer_config(
-            self.optimizer_config, self.trainer.model.parameters()
-        )
-
-    def compute_loss(
-        self, batch: BatchedDataDict, output: OutputType
-    ) -> Dict[str, torch.Tensor]:
-        return {"loss": self.loss(output, batch)}
-
-    def update_metrics(
-        self, batch: BatchedDataDict, output: OutputType, mode: str
-    ) -> None:
-
-        if mode == "test":
-            metrics = self.test_metrics
-        else:
-            metrics = self.metrics
-
-        for stem, metric in metrics.items():
-
-            if stem == "mne:+":
-                stem = "mne"
-
-            if mode == "train":
-                metric.update(
-                    output["audio"][stem],
-                    batch["audio"][stem],
-                )
-            else:
-                if stem not in batch["audio"]:
-                    matched = False
-                    if stem in self._VOX_STEMS:
-                        for bstem in self._VOX_STEMS:
-                            if bstem in batch["audio"]:
-                                batch["audio"][stem] = batch["audio"][bstem]
-                                matched = True
-                                break
-                    elif stem in self._BG_STEMS:
-                        for bstem in self._BG_STEMS:
-                            if bstem in batch["audio"]:
-                                batch["audio"][stem] = batch["audio"][bstem]
-                                matched = True
-                                break
-                else:
-                    matched = True
-
-                if matched:
-                    if stem == "mne" and "mne" not in output["audio"]:
-                        output["audio"]["mne"] = (
-                            output["audio"]["music"] + output["audio"]["effects"]
-                        )
-
-                    metric.update(
-                        output["audio"][stem],
-                        batch["audio"][stem],
-                    )
-
-    def compute_metrics(self, mode: str = "dev") -> Dict[str, torch.Tensor]:
-
-        if mode == "test":
-            metrics = self.test_metrics
-        else:
-            metrics = self.metrics
-
-        metric_dict = {}
-
-        for stem, metric in metrics.items():
-            md = metric.compute()
-            metric_dict.update({f"{stem}/{k}": v for k, v in md.items()})
-
-        self.log_dict(metric_dict, prog_bar=True, logger=False)
-
-        return metric_dict
-
-    def reset_metrics(self, test_mode: bool = False) -> None:
-
-        if test_mode:
-            metrics = self.test_metrics
-        else:
-            metrics = self.metrics
-
-        for _, metric in metrics.items():
-            metric.reset()
-
-    def forward(self, batch: BatchedDataDict) -> Any:
-        batch, output = self.model(batch)
-
-        return batch, output
-
-    def common_step(self, batch: BatchedDataDict, mode: str) -> Any:
-        batch, output = self.forward(batch)
-        loss_dict = self.compute_loss(batch, output)
-
-        with torch.no_grad():
-            self.update_metrics(batch, output, mode=mode)
-
-        if mode == "train":
-            self.log("loss", loss_dict["loss"], prog_bar=True)
-
-        return output, loss_dict
-
-    def training_step(self, batch: BatchedDataDict) -> Dict[str, Any]:
-
-        if self.augmentation is not None:
-            with torch.no_grad():
-                batch = self.augmentation(batch)
-
-        _, loss_dict = self.common_step(batch, mode="train")
-
-        with torch.inference_mode():
-            self.log_dict_with_prefix(
-                loss_dict, "train", batch_size=batch["audio"]["mixture"].shape[0]
-            )
-
-        loss_dict["loss"] *= self.loss_adjustment
-
-        return loss_dict
-
-    def on_train_batch_end(
-        self, outputs: STEP_OUTPUT, batch: BatchedDataDict, batch_idx: int
-    ) -> None:
-
-        metric_dict = self.compute_metrics()
-        self.log_dict_with_prefix(metric_dict, "train")
-        self.reset_metrics()
-
-    def validation_step(
-        self, batch: BatchedDataDict, batch_idx: int, dataloader_idx: int = 0
-    ) -> Dict[str, Any]:
-
-        with torch.inference_mode():
-            curr_val_prefix = f"val{dataloader_idx}" if dataloader_idx > 0 else "val"
-
-            if curr_val_prefix != self.val_prefix:
-                if self.val_prefix is not None:
-                    self._on_validation_epoch_end()
-                self.val_prefix = curr_val_prefix
-            _, loss_dict = self.common_step(batch, mode="val")
-
-            self.log_dict_with_prefix(
-                loss_dict,
-                self.val_prefix,
-                batch_size=batch["audio"]["mixture"].shape[0],
-                prog_bar=True,
-                add_dataloader_idx=False,
-            )
-
-        return loss_dict
-
-    def on_validation_epoch_end(self) -> None:
-        self._on_validation_epoch_end()
-
-    def _on_validation_epoch_end(self) -> None:
-        metric_dict = self.compute_metrics()
-        self.log_dict_with_prefix(
-            metric_dict, self.val_prefix, prog_bar=True, add_dataloader_idx=False
-        )
-        self.reset_metrics()
-
-    def old_predtest_step(
-        self, batch: BatchedDataDict, batch_idx: int, dataloader_idx: int = 0
-    ) -> Tuple[BatchedDataDict, OutputType]:
-
-        audio_batch = batch["audio"]["mixture"]
-        track_batch = batch.get("track", ["" for _ in range(len(audio_batch))])
-
-        output_list_of_dicts = [
-            self.fader(audio[None, ...], lambda a: self.test_forward(a, track))
-            for audio, track in zip(audio_batch, track_batch)
-        ]
-
-        output_dict_of_lists = defaultdict(list)
-
-        for output_dict in output_list_of_dicts:
-            for stem, audio in output_dict.items():
-                output_dict_of_lists[stem].append(audio)
-
-        output = {
-            "audio": {
-                stem: torch.concat(output_list, dim=0)
-                for stem, output_list in output_dict_of_lists.items()
-            }
-        }
-
-        return batch, output
-
-    def predtest_step(
-        self, batch: BatchedDataDict, batch_idx: int = -1, dataloader_idx: int = 0
-    ) -> Tuple[BatchedDataDict, OutputType]:
-
-        if getattr(self.model, "bypass_fader", False):
-            batch, output = self.model(batch)
-        else:
-            audio_batch = batch["audio"]["mixture"]
-            output = self.fader(
-                audio_batch, lambda a: self.test_forward(a, "", batch=batch)
-            )
-
-        return batch, output
-
-    def test_forward(
-        self, audio: torch.Tensor, track: str = "", batch: BatchedDataDict = None
-    ) -> torch.Tensor:
-
-        if self.fader is None:
-            self.attach_fader()
-
-        cond = batch.get("condition", None)
-
-        if cond is not None and cond.shape[0] == 1:
-            cond = cond.repeat(audio.shape[0], 1)
-
-        _, output = self.forward(
-            {
-                "audio": {"mixture": audio},
-                "track": track,
-                "condition": cond,
-            }
-        )
-
-        return output["audio"]
-
-    def on_test_epoch_start(self) -> None:
-        self.attach_fader(force_reattach=True)
-
-    def test_step(
-        self, batch: BatchedDataDict, batch_idx: int, dataloader_idx: int = 0
-    ) -> Any:
-        curr_test_prefix = f"test{dataloader_idx}"
-
-        if curr_test_prefix != self.test_prefix:
-            if self.test_prefix is not None:
-                self._on_test_epoch_end()
-            self.test_prefix = curr_test_prefix
-
-        with torch.inference_mode():
-            _, output = self.predtest_step(batch, batch_idx, dataloader_idx)
-            self.update_metrics(batch, output, mode="test")
-
-        return output
-
-    def on_test_epoch_end(self) -> None:
-        self._on_test_epoch_end()
-
-    def _on_test_epoch_end(self) -> None:
-        metric_dict = self.compute_metrics(mode="test")
-        self.log_dict_with_prefix(
-            metric_dict, self.test_prefix, prog_bar=True, add_dataloader_idx=False
-        )
-        self.reset_metrics()
-
-    def predict_step(
-        self,
-        batch: BatchedDataDict,
-        batch_idx: int = 0,
-        dataloader_idx: int = 0,
-        include_track_name: Optional[bool] = None,
-        get_no_vox_combinations: bool = True,
-        get_residual: bool = False,
-        treat_batch_as_channels: bool = False,
-        fs: Optional[int] = None,
-    ) -> Any:
-        assert self.predict_output_path is not None
-
-        batch_size = batch["audio"]["mixture"].shape[0]
-
-        if include_track_name is None:
-            include_track_name = batch_size > 1
-
-        with torch.inference_mode():
-            batch, output = self.predtest_step(batch, batch_idx, dataloader_idx)
-        print("Pred test finished...")
-        torch.cuda.empty_cache()
-        metric_dict = {}
-
-        if get_residual:
-            mixture = batch["audio"]["mixture"]
-            extracted = sum([output["audio"][stem] for stem in output["audio"]])
-            residual = mixture - extracted
-            print(extracted.shape, mixture.shape, residual.shape)
-
-            output["audio"]["residual"] = residual
-
-        if get_no_vox_combinations:
-            no_vox_stems = [
-                stem for stem in output["audio"] if stem not in self._VOX_STEMS
-            ]
-            no_vox_combinations = chain.from_iterable(
-                combinations(no_vox_stems, r) for r in range(2, len(no_vox_stems) + 1)
-            )
-
-            for combination in no_vox_combinations:
-                combination_ = list(combination)
-                output["audio"]["+".join(combination_)] = sum(
-                    [output["audio"][stem] for stem in combination_]
-                )
-
-        if treat_batch_as_channels:
-            for stem in output["audio"]:
-                output["audio"][stem] = output["audio"][stem].reshape(
-                    1, -1, output["audio"][stem].shape[-1]
-                )
-            batch_size = 1
-
-        for b in range(batch_size):
-            print("!!", b)
-            for stem in output["audio"]:
-                print(f"Saving audio for {stem} to {self.predict_output_path}")
-                track_name = batch["track"][b].split("/")[-1]
-
-                if batch.get("audio", {}).get(stem, None) is not None:
-                    self.test_metrics[stem].reset()
-                    metrics = self.test_metrics[stem](
-                        batch["audio"][stem][[b], ...], output["audio"][stem][[b], ...]
-                    )
-                    snr = metrics["snr"]
-                    sisnr = metrics["sisnr"]
-                    sdr = metrics["sdr"]
-                    metric_dict[stem] = metrics
-                    print(
-                        track_name,
-                        f"snr={snr:2.2f} dB",
-                        f"sisnr={sisnr:2.2f}",
-                        f"sdr={sdr:2.2f} dB",
-                    )
-                    filename = f"{stem} - snr={snr:2.2f}dB - sdr={sdr:2.2f}dB.wav"
-                else:
-                    filename = f"{stem}.wav"
-
-                if include_track_name:
-                    output_dir = os.path.join(self.predict_output_path, track_name)
-                else:
-                    output_dir = self.predict_output_path
-
-                os.makedirs(output_dir, exist_ok=True)
-
-                if fs is None:
-                    fs = self.fs
-
-                ta.save(
-                    os.path.join(output_dir, filename),
-                    output["audio"][stem][b, ...].cpu(),
-                    fs,
-                )
-
-        return metric_dict
-
-    def get_stems(
-        self,
-        batch: BatchedDataDict,
-        batch_idx: int = 0,
-        dataloader_idx: int = 0,
-        include_track_name: Optional[bool] = None,
-        get_no_vox_combinations: bool = True,
-        get_residual: bool = False,
-        treat_batch_as_channels: bool = False,
-        fs: Optional[int] = None,
-    ) -> Any:
-        assert self.predict_output_path is not None
-
-        batch_size = batch["audio"]["mixture"].shape[0]
-
-        if include_track_name is None:
-            include_track_name = batch_size > 1
-
-        with torch.inference_mode():
-            batch, output = self.predtest_step(batch, batch_idx, dataloader_idx)
-        torch.cuda.empty_cache()
-        metric_dict = {}
-
-        if get_residual:
-            mixture = batch["audio"]["mixture"]
-            extracted = sum([output["audio"][stem] for stem in output["audio"]])
-            residual = mixture - extracted
-
-            output["audio"]["residual"] = residual
-
-        if get_no_vox_combinations:
-            no_vox_stems = [
-                stem for stem in output["audio"] if stem not in self._VOX_STEMS
-            ]
-            no_vox_combinations = chain.from_iterable(
-                combinations(no_vox_stems, r) for r in range(2, len(no_vox_stems) + 1)
-            )
-
-            for combination in no_vox_combinations:
-                combination_ = list(combination)
-                output["audio"]["+".join(combination_)] = sum(
-                    [output["audio"][stem] for stem in combination_]
-                )
-
-        if treat_batch_as_channels:
-            for stem in output["audio"]:
-                output["audio"][stem] = output["audio"][stem].reshape(
-                    1, -1, output["audio"][stem].shape[-1]
-                )
-            batch_size = 1
-
-        result = {}
-        for b in range(batch_size):
-            for stem in output["audio"]:
-                track_name = batch["track"][b].split("/")[-1]
-
-                if batch.get("audio", {}).get(stem, None) is not None:
-                    self.test_metrics[stem].reset()
-                    metrics = self.test_metrics[stem](
-                        batch["audio"][stem][[b], ...], output["audio"][stem][[b], ...]
-                    )
-                    snr = metrics["snr"]
-                    sisnr = metrics["sisnr"]
-                    sdr = metrics["sdr"]
-                    metric_dict[stem] = metrics
-                    print(
-                        track_name,
-                        f"snr={snr:2.2f} dB",
-                        f"sisnr={sisnr:2.2f}",
-                        f"sdr={sdr:2.2f} dB",
-                    )
-                    filename = f"{stem} - snr={snr:2.2f}dB - sdr={sdr:2.2f}dB.wav"
-                else:
-                    filename = f"{stem}.wav"
-
-                if include_track_name:
-                    output_dir = os.path.join(self.predict_output_path, track_name)
-                else:
-                    output_dir = self.predict_output_path
-
-                os.makedirs(output_dir, exist_ok=True)
-
-                if fs is None:
-                    fs = self.fs
-
-                result[stem] = output["audio"][stem][b, ...].cpu().numpy()
-
-        return result
-
-    def load_state_dict(
-        self, state_dict: Mapping[str, Any], strict: bool = False
-    ) -> Any:
-
-        return super().load_state_dict(state_dict, strict=False)
-
-    def set_predict_output_path(self, path: str) -> None:
-        self.predict_output_path = path
-        os.makedirs(self.predict_output_path, exist_ok=True)
-
-        self.attach_fader()
-
-    def attach_fader(self, force_reattach=False) -> None:
-        if self.fader is None or force_reattach:
-            self.fader = parse_fader_config(self.fader_config)
-            self.fader.to(self.device)
-
-    def log_dict_with_prefix(
-        self,
-        dict_: Dict[str, torch.Tensor],
-        prefix: str,
-        batch_size: Optional[int] = None,
-        **kwargs: Any,
-    ) -> None:
-        self.log_dict(
-            {f"{prefix}/{k}": v for k, v in dict_.items()},
-            batch_size=batch_size,
-            logger=True,
-            sync_dist=True,
-            **kwargs,
-        )
+import os.path
+from collections import defaultdict
+from itertools import chain, combinations
+from typing import Any, Dict, Iterator, Mapping, Optional, Tuple, Type, TypedDict
+
+import pytorch_lightning as pl
+import torch
+import torchaudio as ta
+import torchmetrics as tm
+from asteroid import losses as asteroid_losses
+
+from pytorch_lightning.utilities.types import STEP_OUTPUT
+from torch import nn, optim
+from torch.optim import lr_scheduler
+from torch.optim.lr_scheduler import LRScheduler
+
+from . import loss, metrics as metrics_, model
+from .data._types import BatchedDataDict
+from .data.augmentation import BaseAugmentor, StemAugmentor
+from .utils import audio as audio_
+from .utils.audio import BaseFader
+
+
+ConfigDict = TypedDict("ConfigDict", {"name": str, "kwargs": Dict[str, Any]})
+
+
+class SchedulerConfigDict(ConfigDict):
+    monitor: str
+
+
+OptimizerSchedulerConfigDict = TypedDict(
+    "OptimizerSchedulerConfigDict",
+    {"optimizer": ConfigDict, "scheduler": SchedulerConfigDict},
+    total=False,
+)
+
+
+class LRSchedulerReturnDict(TypedDict, total=False):
+    scheduler: LRScheduler
+    monitor: str
+
+
+class ConfigureOptimizerReturnDict(TypedDict, total=False):
+    optimizer: torch.optim.Optimizer
+    lr_scheduler: LRSchedulerReturnDict
+
+
+OutputType = Dict[str, Any]
+MetricsType = Dict[str, torch.Tensor]
+
+
+def get_optimizer_class(name: str) -> Type[optim.Optimizer]:
+
+    if name == "DeepSpeedCPUAdam":
+        return DeepSpeedCPUAdam
+
+    for module in [optim, gooptim]:
+        if name in module.__dict__:
+            return module.__dict__[name]
+
+    raise NameError
+
+
+def parse_optimizer_config(
+    config: OptimizerSchedulerConfigDict, parameters: Iterator[nn.Parameter]
+) -> ConfigureOptimizerReturnDict:
+    optim_class = get_optimizer_class(config["optimizer"]["name"])
+    optimizer = optim_class(parameters, **config["optimizer"]["kwargs"])
+
+    optim_dict: ConfigureOptimizerReturnDict = {
+        "optimizer": optimizer,
+    }
+
+    if "scheduler" in config:
+
+        lr_scheduler_class_ = config["scheduler"]["name"]
+        lr_scheduler_class = lr_scheduler.__dict__[lr_scheduler_class_]
+        lr_scheduler_dict: LRSchedulerReturnDict = {
+            "scheduler": lr_scheduler_class(optimizer, **config["scheduler"]["kwargs"])
+        }
+
+        if lr_scheduler_class_ == "ReduceLROnPlateau":
+            lr_scheduler_dict["monitor"] = config["scheduler"]["monitor"]
+
+        optim_dict["lr_scheduler"] = lr_scheduler_dict
+
+    return optim_dict
+
+
+def parse_model_config(config: ConfigDict) -> Any:
+    name = config["name"]
+
+    for module in [model]:
+        if name in module.__dict__:
+            return module.__dict__[name](**config["kwargs"])
+
+    raise NameError
+
+
+_LEGACY_LOSS_NAMES = ["HybridL1Loss"]
+
+
+def _parse_legacy_loss_config(config: ConfigDict) -> nn.Module:
+    name = config["name"]
+
+    if name == "HybridL1Loss":
+        return loss.TimeFreqL1Loss(**config["kwargs"])
+
+    raise NameError
+
+
+def parse_loss_config(config: ConfigDict) -> nn.Module:
+    name = config["name"]
+
+    if name in _LEGACY_LOSS_NAMES:
+        return _parse_legacy_loss_config(config)
+
+    for module in [loss, nn.modules.loss, asteroid_losses]:
+        if name in module.__dict__:
+            return module.__dict__[name](**config["kwargs"])
+
+    raise NameError
+
+
+def get_metric(config: ConfigDict) -> tm.Metric:
+    name = config["name"]
+
+    for module in [tm, metrics_]:
+        if name in module.__dict__:
+            return module.__dict__[name](**config["kwargs"])
+    raise NameError
+
+
+def parse_metric_config(config: Dict[str, ConfigDict]) -> tm.MetricCollection:
+    metrics = {}
+
+    for metric in config:
+        metrics[metric] = get_metric(config[metric])
+
+    return tm.MetricCollection(metrics)
+
+
+def parse_fader_config(config: ConfigDict) -> BaseFader:
+    name = config["name"]
+
+    for module in [audio_]:
+        if name in module.__dict__:
+            return module.__dict__[name](**config["kwargs"])
+
+    raise NameError
+
+
+class LightningSystem(pl.LightningModule):
+    _VOX_STEMS = ["speech", "vocals"]
+    _BG_STEMS = ["background", "effects", "mne"]
+
+    def __init__(
+        self, config: Dict, loss_adjustment: float = 1.0, attach_fader: bool = False
+    ) -> None:
+        super().__init__()
+        self.optimizer_config = config["optimizer"]
+        self.model = parse_model_config(config["model"])
+        self.loss = parse_loss_config(config["loss"])
+        self.metrics = nn.ModuleDict(
+            {
+                stem: parse_metric_config(config["metrics"]["dev"])
+                for stem in self.model.stems
+            }
+        )
+
+        self.metrics.disallow_fsdp = True
+
+        self.test_metrics = nn.ModuleDict(
+            {
+                stem: parse_metric_config(config["metrics"]["test"])
+                for stem in self.model.stems
+            }
+        )
+
+        self.test_metrics.disallow_fsdp = True
+
+        self.fs = config["model"]["kwargs"]["fs"]
+
+        self.fader_config = config["inference"]["fader"]
+        if attach_fader:
+            self.fader = parse_fader_config(config["inference"]["fader"])
+        else:
+            self.fader = None
+
+        self.augmentation: Optional[BaseAugmentor]
+        if config.get("augmentation", None) is not None:
+            self.augmentation = StemAugmentor(**config["augmentation"])
+        else:
+            self.augmentation = None
+
+        self.predict_output_path: Optional[str] = None
+        self.loss_adjustment = loss_adjustment
+
+        self.val_prefix = None
+        self.test_prefix = None
+
+    def configure_optimizers(self) -> Any:
+        return parse_optimizer_config(
+            self.optimizer_config, self.trainer.model.parameters()
+        )
+
+    def compute_loss(
+        self, batch: BatchedDataDict, output: OutputType
+    ) -> Dict[str, torch.Tensor]:
+        return {"loss": self.loss(output, batch)}
+
+    def update_metrics(
+        self, batch: BatchedDataDict, output: OutputType, mode: str
+    ) -> None:
+
+        if mode == "test":
+            metrics = self.test_metrics
+        else:
+            metrics = self.metrics
+
+        for stem, metric in metrics.items():
+
+            if stem == "mne:+":
+                stem = "mne"
+
+            if mode == "train":
+                metric.update(
+                    output["audio"][stem],
+                    batch["audio"][stem],
+                )
+            else:
+                if stem not in batch["audio"]:
+                    matched = False
+                    if stem in self._VOX_STEMS:
+                        for bstem in self._VOX_STEMS:
+                            if bstem in batch["audio"]:
+                                batch["audio"][stem] = batch["audio"][bstem]
+                                matched = True
+                                break
+                    elif stem in self._BG_STEMS:
+                        for bstem in self._BG_STEMS:
+                            if bstem in batch["audio"]:
+                                batch["audio"][stem] = batch["audio"][bstem]
+                                matched = True
+                                break
+                else:
+                    matched = True
+
+                if matched:
+                    if stem == "mne" and "mne" not in output["audio"]:
+                        output["audio"]["mne"] = (
+                            output["audio"]["music"] + output["audio"]["effects"]
+                        )
+
+                    metric.update(
+                        output["audio"][stem],
+                        batch["audio"][stem],
+                    )
+
+    def compute_metrics(self, mode: str = "dev") -> Dict[str, torch.Tensor]:
+
+        if mode == "test":
+            metrics = self.test_metrics
+        else:
+            metrics = self.metrics
+
+        metric_dict = {}
+
+        for stem, metric in metrics.items():
+            md = metric.compute()
+            metric_dict.update({f"{stem}/{k}": v for k, v in md.items()})
+
+        self.log_dict(metric_dict, prog_bar=True, logger=False)
+
+        return metric_dict
+
+    def reset_metrics(self, test_mode: bool = False) -> None:
+
+        if test_mode:
+            metrics = self.test_metrics
+        else:
+            metrics = self.metrics
+
+        for _, metric in metrics.items():
+            metric.reset()
+
+    def forward(self, batch: BatchedDataDict) -> Any:
+        batch, output = self.model(batch)
+
+        return batch, output
+
+    def common_step(self, batch: BatchedDataDict, mode: str) -> Any:
+        batch, output = self.forward(batch)
+        loss_dict = self.compute_loss(batch, output)
+
+        with torch.no_grad():
+            self.update_metrics(batch, output, mode=mode)
+
+        if mode == "train":
+            self.log("loss", loss_dict["loss"], prog_bar=True)
+
+        return output, loss_dict
+
+    def training_step(self, batch: BatchedDataDict) -> Dict[str, Any]:
+
+        if self.augmentation is not None:
+            with torch.no_grad():
+                batch = self.augmentation(batch)
+
+        _, loss_dict = self.common_step(batch, mode="train")
+
+        with torch.inference_mode():
+            self.log_dict_with_prefix(
+                loss_dict, "train", batch_size=batch["audio"]["mixture"].shape[0]
+            )
+
+        loss_dict["loss"] *= self.loss_adjustment
+
+        return loss_dict
+
+    def on_train_batch_end(
+        self, outputs: STEP_OUTPUT, batch: BatchedDataDict, batch_idx: int
+    ) -> None:
+
+        metric_dict = self.compute_metrics()
+        self.log_dict_with_prefix(metric_dict, "train")
+        self.reset_metrics()
+
+    def validation_step(
+        self, batch: BatchedDataDict, batch_idx: int, dataloader_idx: int = 0
+    ) -> Dict[str, Any]:
+
+        with torch.inference_mode():
+            curr_val_prefix = f"val{dataloader_idx}" if dataloader_idx > 0 else "val"
+
+            if curr_val_prefix != self.val_prefix:
+                if self.val_prefix is not None:
+                    self._on_validation_epoch_end()
+                self.val_prefix = curr_val_prefix
+            _, loss_dict = self.common_step(batch, mode="val")
+
+            self.log_dict_with_prefix(
+                loss_dict,
+                self.val_prefix,
+                batch_size=batch["audio"]["mixture"].shape[0],
+                prog_bar=True,
+                add_dataloader_idx=False,
+            )
+
+        return loss_dict
+
+    def on_validation_epoch_end(self) -> None:
+        self._on_validation_epoch_end()
+
+    def _on_validation_epoch_end(self) -> None:
+        metric_dict = self.compute_metrics()
+        self.log_dict_with_prefix(
+            metric_dict, self.val_prefix, prog_bar=True, add_dataloader_idx=False
+        )
+        self.reset_metrics()
+
+    def old_predtest_step(
+        self, batch: BatchedDataDict, batch_idx: int, dataloader_idx: int = 0
+    ) -> Tuple[BatchedDataDict, OutputType]:
+
+        audio_batch = batch["audio"]["mixture"]
+        track_batch = batch.get("track", ["" for _ in range(len(audio_batch))])
+
+        output_list_of_dicts = [
+            self.fader(audio[None, ...], lambda a: self.test_forward(a, track))
+            for audio, track in zip(audio_batch, track_batch)
+        ]
+
+        output_dict_of_lists = defaultdict(list)
+
+        for output_dict in output_list_of_dicts:
+            for stem, audio in output_dict.items():
+                output_dict_of_lists[stem].append(audio)
+
+        output = {
+            "audio": {
+                stem: torch.concat(output_list, dim=0)
+                for stem, output_list in output_dict_of_lists.items()
+            }
+        }
+
+        return batch, output
+
+    def predtest_step(
+        self, batch: BatchedDataDict, batch_idx: int = -1, dataloader_idx: int = 0
+    ) -> Tuple[BatchedDataDict, OutputType]:
+
+        if getattr(self.model, "bypass_fader", False):
+            batch, output = self.model(batch)
+        else:
+            audio_batch = batch["audio"]["mixture"]
+            output = self.fader(
+                audio_batch, lambda a: self.test_forward(a, "", batch=batch)
+            )
+
+        return batch, output
+
+    def test_forward(
+        self, audio: torch.Tensor, track: str = "", batch: BatchedDataDict = None
+    ) -> torch.Tensor:
+
+        if self.fader is None:
+            self.attach_fader()
+
+        cond = batch.get("condition", None)
+
+        if cond is not None and cond.shape[0] == 1:
+            cond = cond.repeat(audio.shape[0], 1)
+
+        _, output = self.forward(
+            {
+                "audio": {"mixture": audio},
+                "track": track,
+                "condition": cond,
+            }
+        )
+
+        return output["audio"]
+
+    def on_test_epoch_start(self) -> None:
+        self.attach_fader(force_reattach=True)
+
+    def test_step(
+        self, batch: BatchedDataDict, batch_idx: int, dataloader_idx: int = 0
+    ) -> Any:
+        curr_test_prefix = f"test{dataloader_idx}"
+
+        if curr_test_prefix != self.test_prefix:
+            if self.test_prefix is not None:
+                self._on_test_epoch_end()
+            self.test_prefix = curr_test_prefix
+
+        with torch.inference_mode():
+            _, output = self.predtest_step(batch, batch_idx, dataloader_idx)
+            self.update_metrics(batch, output, mode="test")
+
+        return output
+
+    def on_test_epoch_end(self) -> None:
+        self._on_test_epoch_end()
+
+    def _on_test_epoch_end(self) -> None:
+        metric_dict = self.compute_metrics(mode="test")
+        self.log_dict_with_prefix(
+            metric_dict, self.test_prefix, prog_bar=True, add_dataloader_idx=False
+        )
+        self.reset_metrics()
+
+    def predict_step(
+        self,
+        batch: BatchedDataDict,
+        batch_idx: int = 0,
+        dataloader_idx: int = 0,
+        include_track_name: Optional[bool] = None,
+        get_no_vox_combinations: bool = True,
+        get_residual: bool = False,
+        treat_batch_as_channels: bool = False,
+        fs: Optional[int] = None,
+    ) -> Any:
+        assert self.predict_output_path is not None
+
+        batch_size = batch["audio"]["mixture"].shape[0]
+
+        if include_track_name is None:
+            include_track_name = batch_size > 1
+
+        with torch.inference_mode():
+            batch, output = self.predtest_step(batch, batch_idx, dataloader_idx)
+        print("Pred test finished...")
+        torch.cuda.empty_cache()
+        metric_dict = {}
+
+        if get_residual:
+            mixture = batch["audio"]["mixture"]
+            extracted = sum([output["audio"][stem] for stem in output["audio"]])
+            residual = mixture - extracted
+            print(extracted.shape, mixture.shape, residual.shape)
+
+            output["audio"]["residual"] = residual
+
+        if get_no_vox_combinations:
+            no_vox_stems = [
+                stem for stem in output["audio"] if stem not in self._VOX_STEMS
+            ]
+            no_vox_combinations = chain.from_iterable(
+                combinations(no_vox_stems, r) for r in range(2, len(no_vox_stems) + 1)
+            )
+
+            for combination in no_vox_combinations:
+                combination_ = list(combination)
+                output["audio"]["+".join(combination_)] = sum(
+                    [output["audio"][stem] for stem in combination_]
+                )
+
+        if treat_batch_as_channels:
+            for stem in output["audio"]:
+                output["audio"][stem] = output["audio"][stem].reshape(
+                    1, -1, output["audio"][stem].shape[-1]
+                )
+            batch_size = 1
+
+        for b in range(batch_size):
+            print("!!", b)
+            for stem in output["audio"]:
+                print(f"Saving audio for {stem} to {self.predict_output_path}")
+                track_name = batch["track"][b].split("/")[-1]
+
+                if batch.get("audio", {}).get(stem, None) is not None:
+                    self.test_metrics[stem].reset()
+                    metrics = self.test_metrics[stem](
+                        batch["audio"][stem][[b], ...], output["audio"][stem][[b], ...]
+                    )
+                    snr = metrics["snr"]
+                    sisnr = metrics["sisnr"]
+                    sdr = metrics["sdr"]
+                    metric_dict[stem] = metrics
+                    print(
+                        track_name,
+                        f"snr={snr:2.2f} dB",
+                        f"sisnr={sisnr:2.2f}",
+                        f"sdr={sdr:2.2f} dB",
+                    )
+                    filename = f"{stem} - snr={snr:2.2f}dB - sdr={sdr:2.2f}dB.wav"
+                else:
+                    filename = f"{stem}.wav"
+
+                if include_track_name:
+                    output_dir = os.path.join(self.predict_output_path, track_name)
+                else:
+                    output_dir = self.predict_output_path
+
+                os.makedirs(output_dir, exist_ok=True)
+
+                if fs is None:
+                    fs = self.fs
+
+                ta.save(
+                    os.path.join(output_dir, filename),
+                    output["audio"][stem][b, ...].cpu(),
+                    fs,
+                )
+
+        return metric_dict
+
+    def get_stems(
+        self,
+        batch: BatchedDataDict,
+        batch_idx: int = 0,
+        dataloader_idx: int = 0,
+        include_track_name: Optional[bool] = None,
+        get_no_vox_combinations: bool = True,
+        get_residual: bool = False,
+        treat_batch_as_channels: bool = False,
+        fs: Optional[int] = None,
+    ) -> Any:
+        assert self.predict_output_path is not None
+
+        batch_size = batch["audio"]["mixture"].shape[0]
+
+        if include_track_name is None:
+            include_track_name = batch_size > 1
+
+        with torch.inference_mode():
+            batch, output = self.predtest_step(batch, batch_idx, dataloader_idx)
+        torch.cuda.empty_cache()
+        metric_dict = {}
+
+        if get_residual:
+            mixture = batch["audio"]["mixture"]
+            extracted = sum([output["audio"][stem] for stem in output["audio"]])
+            residual = mixture - extracted
+
+            output["audio"]["residual"] = residual
+
+        if get_no_vox_combinations:
+            no_vox_stems = [
+                stem for stem in output["audio"] if stem not in self._VOX_STEMS
+            ]
+            no_vox_combinations = chain.from_iterable(
+                combinations(no_vox_stems, r) for r in range(2, len(no_vox_stems) + 1)
+            )
+
+            for combination in no_vox_combinations:
+                combination_ = list(combination)
+                output["audio"]["+".join(combination_)] = sum(
+                    [output["audio"][stem] for stem in combination_]
+                )
+
+        if treat_batch_as_channels:
+            for stem in output["audio"]:
+                output["audio"][stem] = output["audio"][stem].reshape(
+                    1, -1, output["audio"][stem].shape[-1]
+                )
+            batch_size = 1
+
+        result = {}
+        for b in range(batch_size):
+            for stem in output["audio"]:
+                track_name = batch["track"][b].split("/")[-1]
+
+                if batch.get("audio", {}).get(stem, None) is not None:
+                    self.test_metrics[stem].reset()
+                    metrics = self.test_metrics[stem](
+                        batch["audio"][stem][[b], ...], output["audio"][stem][[b], ...]
+                    )
+                    snr = metrics["snr"]
+                    sisnr = metrics["sisnr"]
+                    sdr = metrics["sdr"]
+                    metric_dict[stem] = metrics
+                    print(
+                        track_name,
+                        f"snr={snr:2.2f} dB",
+                        f"sisnr={sisnr:2.2f}",
+                        f"sdr={sdr:2.2f} dB",
+                    )
+                    filename = f"{stem} - snr={snr:2.2f}dB - sdr={sdr:2.2f}dB.wav"
+                else:
+                    filename = f"{stem}.wav"
+
+                if include_track_name:
+                    output_dir = os.path.join(self.predict_output_path, track_name)
+                else:
+                    output_dir = self.predict_output_path
+
+                os.makedirs(output_dir, exist_ok=True)
+
+                if fs is None:
+                    fs = self.fs
+
+                result[stem] = output["audio"][stem][b, ...].cpu().numpy()
+
+        return result
+
+    def load_state_dict(
+        self, state_dict: Mapping[str, Any], strict: bool = False
+    ) -> Any:
+
+        return super().load_state_dict(state_dict, strict=False)
+
+    def set_predict_output_path(self, path: str) -> None:
+        self.predict_output_path = path
+        os.makedirs(self.predict_output_path, exist_ok=True)
+
+        self.attach_fader()
+
+    def attach_fader(self, force_reattach=False) -> None:
+        if self.fader is None or force_reattach:
+            self.fader = parse_fader_config(self.fader_config)
+            self.fader.to(self.device)
+
+    def log_dict_with_prefix(
+        self,
+        dict_: Dict[str, torch.Tensor],
+        prefix: str,
+        batch_size: Optional[int] = None,
+        **kwargs: Any,
+    ) -> None:
+        self.log_dict(
+            {f"{prefix}/{k}": v for k, v in dict_.items()},
+            batch_size=batch_size,
+            logger=True,
+            sync_dist=True,
+            **kwargs,
+        )

mvsepless/models/bandit/core/data/__init__.py

@@ -1,2 +1,2 @@
-from .dnr.datamodule import DivideAndRemasterDataModule
-from .musdb.datamodule import MUSDB18DataModule
+from .dnr.datamodule import DivideAndRemasterDataModule
+from .musdb.datamodule import MUSDB18DataModule

mvsepless/models/bandit/core/data/_types.py

@@ -1,17 +1,17 @@
-from typing import Dict, Sequence, TypedDict
-
-import torch
-
-AudioDict = Dict[str, torch.Tensor]
-
-DataDict = TypedDict("DataDict", {"audio": AudioDict, "track": str})
-
-BatchedDataDict = TypedDict(
-    "BatchedDataDict", {"audio": AudioDict, "track": Sequence[str]}
-)
-
-
-class DataDictWithLanguage(TypedDict):
-    audio: AudioDict
-    track: str
-    language: str
+from typing import Dict, Sequence, TypedDict
+
+import torch
+
+AudioDict = Dict[str, torch.Tensor]
+
+DataDict = TypedDict("DataDict", {"audio": AudioDict, "track": str})
+
+BatchedDataDict = TypedDict(
+    "BatchedDataDict", {"audio": AudioDict, "track": Sequence[str]}
+)
+
+
+class DataDictWithLanguage(TypedDict):
+    audio: AudioDict
+    track: str
+    language: str

mvsepless/models/bandit/core/data/augmentation.py

@@ -1,102 +1,102 @@
-from abc import ABC
-from typing import Any, Dict, Union
-
-import torch
-import torch_audiomentations as tam
-from torch import nn
-
-from ._types import BatchedDataDict, DataDict
-
-
-class BaseAugmentor(nn.Module, ABC):
-    def forward(
-        self, item: Union[DataDict, BatchedDataDict]
-    ) -> Union[DataDict, BatchedDataDict]:
-        raise NotImplementedError
-
-
-class StemAugmentor(BaseAugmentor):
-    def __init__(
-        self,
-        audiomentations: Dict[str, Dict[str, Any]],
-        fix_clipping: bool = True,
-        scaler_margin: float = 0.5,
-        apply_both_default_and_common: bool = False,
-    ) -> None:
-        super().__init__()
-
-        augmentations = {}
-
-        self.has_default = "[default]" in audiomentations
-        self.has_common = "[common]" in audiomentations
-        self.apply_both_default_and_common = apply_both_default_and_common
-
-        for stem in audiomentations:
-            if audiomentations[stem]["name"] == "Compose":
-                augmentations[stem] = getattr(tam, audiomentations[stem]["name"])(
-                    [
-                        getattr(tam, aug["name"])(**aug["kwargs"])
-                        for aug in audiomentations[stem]["kwargs"]["transforms"]
-                    ],
-                    **audiomentations[stem]["kwargs"]["kwargs"],
-                )
-            else:
-                augmentations[stem] = getattr(tam, audiomentations[stem]["name"])(
-                    **audiomentations[stem]["kwargs"]
-                )
-
-        self.augmentations = nn.ModuleDict(augmentations)
-        self.fix_clipping = fix_clipping
-        self.scaler_margin = scaler_margin
-
-    def check_and_fix_clipping(
-        self, item: Union[DataDict, BatchedDataDict]
-    ) -> Union[DataDict, BatchedDataDict]:
-        max_abs = []
-
-        for stem in item["audio"]:
-            max_abs.append(item["audio"][stem].abs().max().item())
-
-        if max(max_abs) > 1.0:
-            scaler = 1.0 / (
-                max(max_abs)
-                + torch.rand((1,), device=item["audio"]["mixture"].device)
-                * self.scaler_margin
-            )
-
-            for stem in item["audio"]:
-                item["audio"][stem] *= scaler
-
-        return item
-
-    def forward(
-        self, item: Union[DataDict, BatchedDataDict]
-    ) -> Union[DataDict, BatchedDataDict]:
-
-        for stem in item["audio"]:
-            if stem == "mixture":
-                continue
-
-            if self.has_common:
-                item["audio"][stem] = self.augmentations["[common]"](
-                    item["audio"][stem]
-                ).samples
-
-            if stem in self.augmentations:
-                item["audio"][stem] = self.augmentations[stem](
-                    item["audio"][stem]
-                ).samples
-            elif self.has_default:
-                if not self.has_common or self.apply_both_default_and_common:
-                    item["audio"][stem] = self.augmentations["[default]"](
-                        item["audio"][stem]
-                    ).samples
-
-        item["audio"]["mixture"] = sum(
-            [item["audio"][stem] for stem in item["audio"] if stem != "mixture"]
-        )  # type: ignore[call-overload, assignment]
-
-        if self.fix_clipping:
-            item = self.check_and_fix_clipping(item)
-
-        return item
+from abc import ABC
+from typing import Any, Dict, Union
+
+import torch
+import torch_audiomentations as tam
+from torch import nn
+
+from ._types import BatchedDataDict, DataDict
+
+
+class BaseAugmentor(nn.Module, ABC):
+    def forward(
+        self, item: Union[DataDict, BatchedDataDict]
+    ) -> Union[DataDict, BatchedDataDict]:
+        raise NotImplementedError
+
+
+class StemAugmentor(BaseAugmentor):
+    def __init__(
+        self,
+        audiomentations: Dict[str, Dict[str, Any]],
+        fix_clipping: bool = True,
+        scaler_margin: float = 0.5,
+        apply_both_default_and_common: bool = False,
+    ) -> None:
+        super().__init__()
+
+        augmentations = {}
+
+        self.has_default = "[default]" in audiomentations
+        self.has_common = "[common]" in audiomentations
+        self.apply_both_default_and_common = apply_both_default_and_common
+
+        for stem in audiomentations:
+            if audiomentations[stem]["name"] == "Compose":
+                augmentations[stem] = getattr(tam, audiomentations[stem]["name"])(
+                    [
+                        getattr(tam, aug["name"])(**aug["kwargs"])
+                        for aug in audiomentations[stem]["kwargs"]["transforms"]
+                    ],
+                    **audiomentations[stem]["kwargs"]["kwargs"],
+                )
+            else:
+                augmentations[stem] = getattr(tam, audiomentations[stem]["name"])(
+                    **audiomentations[stem]["kwargs"]
+                )
+
+        self.augmentations = nn.ModuleDict(augmentations)
+        self.fix_clipping = fix_clipping
+        self.scaler_margin = scaler_margin
+
+    def check_and_fix_clipping(
+        self, item: Union[DataDict, BatchedDataDict]
+    ) -> Union[DataDict, BatchedDataDict]:
+        max_abs = []
+
+        for stem in item["audio"]:
+            max_abs.append(item["audio"][stem].abs().max().item())
+
+        if max(max_abs) > 1.0:
+            scaler = 1.0 / (
+                max(max_abs)
+                + torch.rand((1,), device=item["audio"]["mixture"].device)
+                * self.scaler_margin
+            )
+
+            for stem in item["audio"]:
+                item["audio"][stem] *= scaler
+
+        return item
+
+    def forward(
+        self, item: Union[DataDict, BatchedDataDict]
+    ) -> Union[DataDict, BatchedDataDict]:
+
+        for stem in item["audio"]:
+            if stem == "mixture":
+                continue
+
+            if self.has_common:
+                item["audio"][stem] = self.augmentations["[common]"](
+                    item["audio"][stem]
+                ).samples
+
+            if stem in self.augmentations:
+                item["audio"][stem] = self.augmentations[stem](
+                    item["audio"][stem]
+                ).samples
+            elif self.has_default:
+                if not self.has_common or self.apply_both_default_and_common:
+                    item["audio"][stem] = self.augmentations["[default]"](
+                        item["audio"][stem]
+                    ).samples
+
+        item["audio"]["mixture"] = sum(
+            [item["audio"][stem] for stem in item["audio"] if stem != "mixture"]
+        )  # type: ignore[call-overload, assignment]
+
+        if self.fix_clipping:
+            item = self.check_and_fix_clipping(item)
+
+        return item

mvsepless/models/bandit/core/data/augmented.py

@@ -1,34 +1,34 @@
-import warnings
-from typing import Dict, Optional, Union
-
-import torch
-from torch import nn
-from torch.utils import data
-
-
-class AugmentedDataset(data.Dataset):
-    def __init__(
-        self,
-        dataset: data.Dataset,
-        augmentation: nn.Module = nn.Identity(),
-        target_length: Optional[int] = None,
-    ) -> None:
-        warnings.warn(
-            "This class is no longer used. Attach augmentation to "
-            "the LightningSystem instead.",
-            DeprecationWarning,
-        )
-
-        self.dataset = dataset
-        self.augmentation = augmentation
-
-        self.ds_length: int = len(dataset)  # type: ignore[arg-type]
-        self.length = target_length if target_length is not None else self.ds_length
-
-    def __getitem__(self, index: int) -> Dict[str, Union[str, Dict[str, torch.Tensor]]]:
-        item = self.dataset[index % self.ds_length]
-        item = self.augmentation(item)
-        return item
-
-    def __len__(self) -> int:
-        return self.length
+import warnings
+from typing import Dict, Optional, Union
+
+import torch
+from torch import nn
+from torch.utils import data
+
+
+class AugmentedDataset(data.Dataset):
+    def __init__(
+        self,
+        dataset: data.Dataset,
+        augmentation: nn.Module = nn.Identity(),
+        target_length: Optional[int] = None,
+    ) -> None:
+        warnings.warn(
+            "This class is no longer used. Attach augmentation to "
+            "the LightningSystem instead.",
+            DeprecationWarning,
+        )
+
+        self.dataset = dataset
+        self.augmentation = augmentation
+
+        self.ds_length: int = len(dataset)  # type: ignore[arg-type]
+        self.length = target_length if target_length is not None else self.ds_length
+
+    def __getitem__(self, index: int) -> Dict[str, Union[str, Dict[str, torch.Tensor]]]:
+        item = self.dataset[index % self.ds_length]
+        item = self.augmentation(item)
+        return item
+
+    def __len__(self) -> int:
+        return self.length

mvsepless/models/bandit/core/data/base.py

@@ -1,60 +1,60 @@
-import os
-from abc import ABC, abstractmethod
-from typing import Any, Dict, List, Optional
-
-import numpy as np
-import pedalboard as pb
-import torch
-import torchaudio as ta
-from torch.utils import data
-
-from ._types import AudioDict, DataDict
-
-
-class BaseSourceSeparationDataset(data.Dataset, ABC):
-    def __init__(
-        self,
-        split: str,
-        stems: List[str],
-        files: List[str],
-        data_path: str,
-        fs: int,
-        npy_memmap: bool,
-        recompute_mixture: bool,
-    ):
-        self.split = split
-        self.stems = stems
-        self.stems_no_mixture = [s for s in stems if s != "mixture"]
-        self.files = files
-        self.data_path = data_path
-        self.fs = fs
-        self.npy_memmap = npy_memmap
-        self.recompute_mixture = recompute_mixture
-
-    @abstractmethod
-    def get_stem(self, *, stem: str, identifier: Dict[str, Any]) -> torch.Tensor:
-        raise NotImplementedError
-
-    def _get_audio(self, stems, identifier: Dict[str, Any]):
-        audio = {}
-        for stem in stems:
-            audio[stem] = self.get_stem(stem=stem, identifier=identifier)
-
-        return audio
-
-    def get_audio(self, identifier: Dict[str, Any]) -> AudioDict:
-
-        if self.recompute_mixture:
-            audio = self._get_audio(self.stems_no_mixture, identifier=identifier)
-            audio["mixture"] = self.compute_mixture(audio)
-            return audio
-        else:
-            return self._get_audio(self.stems, identifier=identifier)
-
-    @abstractmethod
-    def get_identifier(self, index: int) -> Dict[str, Any]:
-        pass
-
-    def compute_mixture(self, audio: AudioDict) -> torch.Tensor:
-
-        return sum(audio[stem] for stem in audio if stem != "mixture")
+import os
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List, Optional
+
+import numpy as np
+import pedalboard as pb
+import torch
+import torchaudio as ta
+from torch.utils import data
+
+from ._types import AudioDict, DataDict
+
+
+class BaseSourceSeparationDataset(data.Dataset, ABC):
+    def __init__(
+        self,
+        split: str,
+        stems: List[str],
+        files: List[str],
+        data_path: str,
+        fs: int,
+        npy_memmap: bool,
+        recompute_mixture: bool,
+    ):
+        self.split = split
+        self.stems = stems
+        self.stems_no_mixture = [s for s in stems if s != "mixture"]
+        self.files = files
+        self.data_path = data_path
+        self.fs = fs
+        self.npy_memmap = npy_memmap
+        self.recompute_mixture = recompute_mixture
+
+    @abstractmethod
+    def get_stem(self, *, stem: str, identifier: Dict[str, Any]) -> torch.Tensor:
+        raise NotImplementedError
+
+    def _get_audio(self, stems, identifier: Dict[str, Any]):
+        audio = {}
+        for stem in stems:
+            audio[stem] = self.get_stem(stem=stem, identifier=identifier)
+
+        return audio
+
+    def get_audio(self, identifier: Dict[str, Any]) -> AudioDict:
+
+        if self.recompute_mixture:
+            audio = self._get_audio(self.stems_no_mixture, identifier=identifier)
+            audio["mixture"] = self.compute_mixture(audio)
+            return audio
+        else:
+            return self._get_audio(self.stems, identifier=identifier)
+
+    @abstractmethod
+    def get_identifier(self, index: int) -> Dict[str, Any]:
+        pass
+
+    def compute_mixture(self, audio: AudioDict) -> torch.Tensor:
+
+        return sum(audio[stem] for stem in audio if stem != "mixture")

mvsepless/models/bandit/core/data/dnr/datamodule.py

@@ -1,64 +1,64 @@
-import os
-from typing import Mapping, Optional
-
-import pytorch_lightning as pl
-
-from .dataset import (
-    DivideAndRemasterDataset,
-    DivideAndRemasterDeterministicChunkDataset,
-    DivideAndRemasterRandomChunkDataset,
-    DivideAndRemasterRandomChunkDatasetWithSpeechReverb,
-)
-
-
-def DivideAndRemasterDataModule(
-    data_root: str = "$DATA_ROOT/DnR/v2",
-    batch_size: int = 2,
-    num_workers: int = 8,
-    train_kwargs: Optional[Mapping] = None,
-    val_kwargs: Optional[Mapping] = None,
-    test_kwargs: Optional[Mapping] = None,
-    datamodule_kwargs: Optional[Mapping] = None,
-    use_speech_reverb: bool = False,
-) -> pl.LightningDataModule:
-    if train_kwargs is None:
-        train_kwargs = {}
-
-    if val_kwargs is None:
-        val_kwargs = {}
-
-    if test_kwargs is None:
-        test_kwargs = {}
-
-    if datamodule_kwargs is None:
-        datamodule_kwargs = {}
-
-    if num_workers is None:
-        num_workers = os.cpu_count()
-
-        if num_workers is None:
-            num_workers = 32
-
-        num_workers = min(num_workers, 64)
-
-    if use_speech_reverb:
-        train_cls = DivideAndRemasterRandomChunkDatasetWithSpeechReverb
-    else:
-        train_cls = DivideAndRemasterRandomChunkDataset
-
-    train_dataset = train_cls(data_root, "train", **train_kwargs)
-
-    datamodule = pl.LightningDataModule.from_datasets(
-        train_dataset=train_dataset,
-        val_dataset=DivideAndRemasterDeterministicChunkDataset(
-            data_root, "val", **val_kwargs
-        ),
-        test_dataset=DivideAndRemasterDataset(data_root, "test", **test_kwargs),
-        batch_size=batch_size,
-        num_workers=num_workers,
-        **datamodule_kwargs,
-    )
-
-    datamodule.predict_dataloader = datamodule.test_dataloader  # type: ignore[method-assign]
-
-    return datamodule
+import os
+from typing import Mapping, Optional
+
+import pytorch_lightning as pl
+
+from .dataset import (
+    DivideAndRemasterDataset,
+    DivideAndRemasterDeterministicChunkDataset,
+    DivideAndRemasterRandomChunkDataset,
+    DivideAndRemasterRandomChunkDatasetWithSpeechReverb,
+)
+
+
+def DivideAndRemasterDataModule(
+    data_root: str = "$DATA_ROOT/DnR/v2",
+    batch_size: int = 2,
+    num_workers: int = 8,
+    train_kwargs: Optional[Mapping] = None,
+    val_kwargs: Optional[Mapping] = None,
+    test_kwargs: Optional[Mapping] = None,
+    datamodule_kwargs: Optional[Mapping] = None,
+    use_speech_reverb: bool = False,
+) -> pl.LightningDataModule:
+    if train_kwargs is None:
+        train_kwargs = {}
+
+    if val_kwargs is None:
+        val_kwargs = {}
+
+    if test_kwargs is None:
+        test_kwargs = {}
+
+    if datamodule_kwargs is None:
+        datamodule_kwargs = {}
+
+    if num_workers is None:
+        num_workers = os.cpu_count()
+
+        if num_workers is None:
+            num_workers = 32
+
+        num_workers = min(num_workers, 64)
+
+    if use_speech_reverb:
+        train_cls = DivideAndRemasterRandomChunkDatasetWithSpeechReverb
+    else:
+        train_cls = DivideAndRemasterRandomChunkDataset
+
+    train_dataset = train_cls(data_root, "train", **train_kwargs)
+
+    datamodule = pl.LightningDataModule.from_datasets(
+        train_dataset=train_dataset,
+        val_dataset=DivideAndRemasterDeterministicChunkDataset(
+            data_root, "val", **val_kwargs
+        ),
+        test_dataset=DivideAndRemasterDataset(data_root, "test", **test_kwargs),
+        batch_size=batch_size,
+        num_workers=num_workers,
+        **datamodule_kwargs,
+    )
+
+    datamodule.predict_dataloader = datamodule.test_dataloader  # type: ignore[method-assign]
+
+    return datamodule

mvsepless/models/bandit/core/data/dnr/dataset.py

@@ -1,360 +1,360 @@
-import os
-from abc import ABC
-from typing import Any, Dict, List, Optional
-
-import numpy as np
-import pedalboard as pb
-import torch
-import torchaudio as ta
-from torch.utils import data
-
-from .._types import AudioDict, DataDict
-from ..base import BaseSourceSeparationDataset
-
-
-class DivideAndRemasterBaseDataset(BaseSourceSeparationDataset, ABC):
-    ALLOWED_STEMS = ["mixture", "speech", "music", "effects", "mne"]
-    STEM_NAME_MAP = {
-        "mixture": "mix",
-        "speech": "speech",
-        "music": "music",
-        "effects": "sfx",
-    }
-    SPLIT_NAME_MAP = {"train": "tr", "val": "cv", "test": "tt"}
-
-    FULL_TRACK_LENGTH_SECOND = 60
-    FULL_TRACK_LENGTH_SAMPLES = FULL_TRACK_LENGTH_SECOND * 44100
-
-    def __init__(
-        self,
-        split: str,
-        stems: List[str],
-        files: List[str],
-        data_path: str,
-        fs: int = 44100,
-        npy_memmap: bool = True,
-        recompute_mixture: bool = False,
-    ) -> None:
-        super().__init__(
-            split=split,
-            stems=stems,
-            files=files,
-            data_path=data_path,
-            fs=fs,
-            npy_memmap=npy_memmap,
-            recompute_mixture=recompute_mixture,
-        )
-
-    def get_stem(self, *, stem: str, identifier: Dict[str, Any]) -> torch.Tensor:
-
-        if stem == "mne":
-            return self.get_stem(stem="music", identifier=identifier) + self.get_stem(
-                stem="effects", identifier=identifier
-            )
-
-        track = identifier["track"]
-        path = os.path.join(self.data_path, track)
-
-        if self.npy_memmap:
-            audio = np.load(
-                os.path.join(path, f"{self.STEM_NAME_MAP[stem]}.npy"), mmap_mode="r"
-            )
-        else:
-            audio, _ = ta.load(os.path.join(path, f"{self.STEM_NAME_MAP[stem]}.wav"))
-
-        return audio
-
-    def get_identifier(self, index):
-        return dict(track=self.files[index])
-
-    def __getitem__(self, index: int) -> DataDict:
-        identifier = self.get_identifier(index)
-        audio = self.get_audio(identifier)
-
-        return {"audio": audio, "track": f"{self.split}/{identifier['track']}"}
-
-
-class DivideAndRemasterDataset(DivideAndRemasterBaseDataset):
-    def __init__(
-        self,
-        data_root: str,
-        split: str,
-        stems: Optional[List[str]] = None,
-        fs: int = 44100,
-        npy_memmap: bool = True,
-    ) -> None:
-
-        if stems is None:
-            stems = self.ALLOWED_STEMS
-        self.stems = stems
-
-        data_path = os.path.join(data_root, self.SPLIT_NAME_MAP[split])
-
-        files = sorted(os.listdir(data_path))
-        files = [
-            f
-            for f in files
-            if (not f.startswith(".")) and os.path.isdir(os.path.join(data_path, f))
-        ]
-        if split == "train":
-            assert len(files) == 3406, len(files)
-        elif split == "val":
-            assert len(files) == 487, len(files)
-        elif split == "test":
-            assert len(files) == 973, len(files)
-
-        self.n_tracks = len(files)
-
-        super().__init__(
-            data_path=data_path,
-            split=split,
-            stems=stems,
-            files=files,
-            fs=fs,
-            npy_memmap=npy_memmap,
-        )
-
-    def __len__(self) -> int:
-        return self.n_tracks
-
-
-class DivideAndRemasterRandomChunkDataset(DivideAndRemasterBaseDataset):
-    def __init__(
-        self,
-        data_root: str,
-        split: str,
-        target_length: int,
-        chunk_size_second: float,
-        stems: Optional[List[str]] = None,
-        fs: int = 44100,
-        npy_memmap: bool = True,
-    ) -> None:
-
-        if stems is None:
-            stems = self.ALLOWED_STEMS
-        self.stems = stems
-
-        data_path = os.path.join(data_root, self.SPLIT_NAME_MAP[split])
-
-        files = sorted(os.listdir(data_path))
-        files = [
-            f
-            for f in files
-            if (not f.startswith(".")) and os.path.isdir(os.path.join(data_path, f))
-        ]
-
-        if split == "train":
-            assert len(files) == 3406, len(files)
-        elif split == "val":
-            assert len(files) == 487, len(files)
-        elif split == "test":
-            assert len(files) == 973, len(files)
-
-        self.n_tracks = len(files)
-
-        self.target_length = target_length
-        self.chunk_size = int(chunk_size_second * fs)
-
-        super().__init__(
-            data_path=data_path,
-            split=split,
-            stems=stems,
-            files=files,
-            fs=fs,
-            npy_memmap=npy_memmap,
-        )
-
-    def __len__(self) -> int:
-        return self.target_length
-
-    def get_identifier(self, index):
-        return super().get_identifier(index % self.n_tracks)
-
-    def get_stem(
-        self,
-        *,
-        stem: str,
-        identifier: Dict[str, Any],
-        chunk_here: bool = False,
-    ) -> torch.Tensor:
-
-        stem = super().get_stem(stem=stem, identifier=identifier)
-
-        if chunk_here:
-            start = np.random.randint(
-                0, self.FULL_TRACK_LENGTH_SAMPLES - self.chunk_size
-            )
-            end = start + self.chunk_size
-
-            stem = stem[:, start:end]
-
-        return stem
-
-    def __getitem__(self, index: int) -> DataDict:
-        identifier = self.get_identifier(index)
-        audio = self.get_audio(identifier)
-
-        start = np.random.randint(0, self.FULL_TRACK_LENGTH_SAMPLES - self.chunk_size)
-        end = start + self.chunk_size
-
-        audio = {k: v[:, start:end] for k, v in audio.items()}
-
-        return {"audio": audio, "track": f"{self.split}/{identifier['track']}"}
-
-
-class DivideAndRemasterDeterministicChunkDataset(DivideAndRemasterBaseDataset):
-    def __init__(
-        self,
-        data_root: str,
-        split: str,
-        chunk_size_second: float,
-        hop_size_second: float,
-        stems: Optional[List[str]] = None,
-        fs: int = 44100,
-        npy_memmap: bool = True,
-    ) -> None:
-
-        if stems is None:
-            stems = self.ALLOWED_STEMS
-        self.stems = stems
-
-        data_path = os.path.join(data_root, self.SPLIT_NAME_MAP[split])
-
-        files = sorted(os.listdir(data_path))
-        files = [
-            f
-            for f in files
-            if (not f.startswith(".")) and os.path.isdir(os.path.join(data_path, f))
-        ]
-        if split == "train":
-            assert len(files) == 3406, len(files)
-        elif split == "val":
-            assert len(files) == 487, len(files)
-        elif split == "test":
-            assert len(files) == 973, len(files)
-
-        self.n_tracks = len(files)
-
-        self.chunk_size = int(chunk_size_second * fs)
-        self.hop_size = int(hop_size_second * fs)
-        self.n_chunks_per_track = int(
-            (self.FULL_TRACK_LENGTH_SECOND - chunk_size_second) / hop_size_second
-        )
-
-        self.length = self.n_tracks * self.n_chunks_per_track
-
-        super().__init__(
-            data_path=data_path,
-            split=split,
-            stems=stems,
-            files=files,
-            fs=fs,
-            npy_memmap=npy_memmap,
-        )
-
-    def get_identifier(self, index):
-        return super().get_identifier(index % self.n_tracks)
-
-    def __len__(self) -> int:
-        return self.length
-
-    def __getitem__(self, item: int) -> DataDict:
-
-        index = item % self.n_tracks
-        chunk = item // self.n_tracks
-
-        data_ = super().__getitem__(index)
-
-        audio = data_["audio"]
-
-        start = chunk * self.hop_size
-        end = start + self.chunk_size
-
-        for stem in self.stems:
-            data_["audio"][stem] = audio[stem][:, start:end]
-
-        return data_
-
-
-class DivideAndRemasterRandomChunkDatasetWithSpeechReverb(
-    DivideAndRemasterRandomChunkDataset
-):
-    def __init__(
-        self,
-        data_root: str,
-        split: str,
-        target_length: int,
-        chunk_size_second: float,
-        stems: Optional[List[str]] = None,
-        fs: int = 44100,
-        npy_memmap: bool = True,
-    ) -> None:
-
-        if stems is None:
-            stems = self.ALLOWED_STEMS
-
-        stems_no_mixture = [s for s in stems if s != "mixture"]
-
-        super().__init__(
-            data_root=data_root,
-            split=split,
-            target_length=target_length,
-            chunk_size_second=chunk_size_second,
-            stems=stems_no_mixture,
-            fs=fs,
-            npy_memmap=npy_memmap,
-        )
-
-        self.stems = stems
-        self.stems_no_mixture = stems_no_mixture
-
-    def __getitem__(self, index: int) -> DataDict:
-
-        data_ = super().__getitem__(index)
-
-        dry = data_["audio"]["speech"][:]
-        n_samples = dry.shape[-1]
-
-        wet_level = np.random.rand()
-
-        speech = pb.Reverb(
-            room_size=np.random.rand(),
-            damping=np.random.rand(),
-            wet_level=wet_level,
-            dry_level=(1 - wet_level),
-            width=np.random.rand(),
-        ).process(dry, self.fs, buffer_size=8192 * 4)[..., :n_samples]
-
-        data_["audio"]["speech"] = speech
-
-        data_["audio"]["mixture"] = sum(
-            [data_["audio"][s] for s in self.stems_no_mixture]
-        )
-
-        return data_
-
-    def __len__(self) -> int:
-        return super().__len__()
-
-
-if __name__ == "__main__":
-
-    from pprint import pprint
-    from tqdm.auto import tqdm
-
-    for split_ in ["train", "val", "test"]:
-        ds = DivideAndRemasterRandomChunkDatasetWithSpeechReverb(
-            data_root="$DATA_ROOT/DnR/v2np",
-            split=split_,
-            target_length=100,
-            chunk_size_second=6.0,
-        )
-
-        print(split_, len(ds))
-
-        for track_ in tqdm(ds):  # type: ignore
-            pprint(track_)
-            track_["audio"] = {k: v.shape for k, v in track_["audio"].items()}
-            pprint(track_)
-
-        break
+import os
+from abc import ABC
+from typing import Any, Dict, List, Optional
+
+import numpy as np
+import pedalboard as pb
+import torch
+import torchaudio as ta
+from torch.utils import data
+
+from .._types import AudioDict, DataDict
+from ..base import BaseSourceSeparationDataset
+
+
+class DivideAndRemasterBaseDataset(BaseSourceSeparationDataset, ABC):
+    ALLOWED_STEMS = ["mixture", "speech", "music", "effects", "mne"]
+    STEM_NAME_MAP = {
+        "mixture": "mix",
+        "speech": "speech",
+        "music": "music",
+        "effects": "sfx",
+    }
+    SPLIT_NAME_MAP = {"train": "tr", "val": "cv", "test": "tt"}
+
+    FULL_TRACK_LENGTH_SECOND = 60
+    FULL_TRACK_LENGTH_SAMPLES = FULL_TRACK_LENGTH_SECOND * 44100
+
+    def __init__(
+        self,
+        split: str,
+        stems: List[str],
+        files: List[str],
+        data_path: str,
+        fs: int = 44100,
+        npy_memmap: bool = True,
+        recompute_mixture: bool = False,
+    ) -> None:
+        super().__init__(
+            split=split,
+            stems=stems,
+            files=files,
+            data_path=data_path,
+            fs=fs,
+            npy_memmap=npy_memmap,
+            recompute_mixture=recompute_mixture,
+        )
+
+    def get_stem(self, *, stem: str, identifier: Dict[str, Any]) -> torch.Tensor:
+
+        if stem == "mne":
+            return self.get_stem(stem="music", identifier=identifier) + self.get_stem(
+                stem="effects", identifier=identifier
+            )
+
+        track = identifier["track"]
+        path = os.path.join(self.data_path, track)
+
+        if self.npy_memmap:
+            audio = np.load(
+                os.path.join(path, f"{self.STEM_NAME_MAP[stem]}.npy"), mmap_mode="r"
+            )
+        else:
+            audio, _ = ta.load(os.path.join(path, f"{self.STEM_NAME_MAP[stem]}.wav"))
+
+        return audio
+
+    def get_identifier(self, index):
+        return dict(track=self.files[index])
+
+    def __getitem__(self, index: int) -> DataDict:
+        identifier = self.get_identifier(index)
+        audio = self.get_audio(identifier)
+
+        return {"audio": audio, "track": f"{self.split}/{identifier['track']}"}
+
+
+class DivideAndRemasterDataset(DivideAndRemasterBaseDataset):
+    def __init__(
+        self,
+        data_root: str,
+        split: str,
+        stems: Optional[List[str]] = None,
+        fs: int = 44100,
+        npy_memmap: bool = True,
+    ) -> None:
+
+        if stems is None:
+            stems = self.ALLOWED_STEMS
+        self.stems = stems
+
+        data_path = os.path.join(data_root, self.SPLIT_NAME_MAP[split])
+
+        files = sorted(os.listdir(data_path))
+        files = [
+            f
+            for f in files
+            if (not f.startswith(".")) and os.path.isdir(os.path.join(data_path, f))
+        ]
+        if split == "train":
+            assert len(files) == 3406, len(files)
+        elif split == "val":
+            assert len(files) == 487, len(files)
+        elif split == "test":
+            assert len(files) == 973, len(files)
+
+        self.n_tracks = len(files)
+
+        super().__init__(
+            data_path=data_path,
+            split=split,
+            stems=stems,
+            files=files,
+            fs=fs,
+            npy_memmap=npy_memmap,
+        )
+
+    def __len__(self) -> int:
+        return self.n_tracks
+
+
+class DivideAndRemasterRandomChunkDataset(DivideAndRemasterBaseDataset):
+    def __init__(
+        self,
+        data_root: str,
+        split: str,
+        target_length: int,
+        chunk_size_second: float,
+        stems: Optional[List[str]] = None,
+        fs: int = 44100,
+        npy_memmap: bool = True,
+    ) -> None:
+
+        if stems is None:
+            stems = self.ALLOWED_STEMS
+        self.stems = stems
+
+        data_path = os.path.join(data_root, self.SPLIT_NAME_MAP[split])
+
+        files = sorted(os.listdir(data_path))
+        files = [
+            f
+            for f in files
+            if (not f.startswith(".")) and os.path.isdir(os.path.join(data_path, f))
+        ]
+
+        if split == "train":
+            assert len(files) == 3406, len(files)
+        elif split == "val":
+            assert len(files) == 487, len(files)
+        elif split == "test":
+            assert len(files) == 973, len(files)
+
+        self.n_tracks = len(files)
+
+        self.target_length = target_length
+        self.chunk_size = int(chunk_size_second * fs)
+
+        super().__init__(
+            data_path=data_path,
+            split=split,
+            stems=stems,
+            files=files,
+            fs=fs,
+            npy_memmap=npy_memmap,
+        )
+
+    def __len__(self) -> int:
+        return self.target_length
+
+    def get_identifier(self, index):
+        return super().get_identifier(index % self.n_tracks)
+
+    def get_stem(
+        self,
+        *,
+        stem: str,
+        identifier: Dict[str, Any],
+        chunk_here: bool = False,
+    ) -> torch.Tensor:
+
+        stem = super().get_stem(stem=stem, identifier=identifier)
+
+        if chunk_here:
+            start = np.random.randint(
+                0, self.FULL_TRACK_LENGTH_SAMPLES - self.chunk_size
+            )
+            end = start + self.chunk_size
+
+            stem = stem[:, start:end]
+
+        return stem
+
+    def __getitem__(self, index: int) -> DataDict:
+        identifier = self.get_identifier(index)
+        audio = self.get_audio(identifier)
+
+        start = np.random.randint(0, self.FULL_TRACK_LENGTH_SAMPLES - self.chunk_size)
+        end = start + self.chunk_size
+
+        audio = {k: v[:, start:end] for k, v in audio.items()}
+
+        return {"audio": audio, "track": f"{self.split}/{identifier['track']}"}
+
+
+class DivideAndRemasterDeterministicChunkDataset(DivideAndRemasterBaseDataset):
+    def __init__(
+        self,
+        data_root: str,
+        split: str,
+        chunk_size_second: float,
+        hop_size_second: float,
+        stems: Optional[List[str]] = None,
+        fs: int = 44100,
+        npy_memmap: bool = True,
+    ) -> None:
+
+        if stems is None:
+            stems = self.ALLOWED_STEMS
+        self.stems = stems
+
+        data_path = os.path.join(data_root, self.SPLIT_NAME_MAP[split])
+
+        files = sorted(os.listdir(data_path))
+        files = [
+            f
+            for f in files
+            if (not f.startswith(".")) and os.path.isdir(os.path.join(data_path, f))
+        ]
+        if split == "train":
+            assert len(files) == 3406, len(files)
+        elif split == "val":
+            assert len(files) == 487, len(files)
+        elif split == "test":
+            assert len(files) == 973, len(files)
+
+        self.n_tracks = len(files)
+
+        self.chunk_size = int(chunk_size_second * fs)
+        self.hop_size = int(hop_size_second * fs)
+        self.n_chunks_per_track = int(
+            (self.FULL_TRACK_LENGTH_SECOND - chunk_size_second) / hop_size_second
+        )
+
+        self.length = self.n_tracks * self.n_chunks_per_track
+
+        super().__init__(
+            data_path=data_path,
+            split=split,
+            stems=stems,
+            files=files,
+            fs=fs,
+            npy_memmap=npy_memmap,
+        )
+
+    def get_identifier(self, index):
+        return super().get_identifier(index % self.n_tracks)
+
+    def __len__(self) -> int:
+        return self.length
+
+    def __getitem__(self, item: int) -> DataDict:
+
+        index = item % self.n_tracks
+        chunk = item // self.n_tracks
+
+        data_ = super().__getitem__(index)
+
+        audio = data_["audio"]
+
+        start = chunk * self.hop_size
+        end = start + self.chunk_size
+
+        for stem in self.stems:
+            data_["audio"][stem] = audio[stem][:, start:end]
+
+        return data_
+
+
+class DivideAndRemasterRandomChunkDatasetWithSpeechReverb(
+    DivideAndRemasterRandomChunkDataset
+):
+    def __init__(
+        self,
+        data_root: str,
+        split: str,
+        target_length: int,
+        chunk_size_second: float,
+        stems: Optional[List[str]] = None,
+        fs: int = 44100,
+        npy_memmap: bool = True,
+    ) -> None:
+
+        if stems is None:
+            stems = self.ALLOWED_STEMS
+
+        stems_no_mixture = [s for s in stems if s != "mixture"]
+
+        super().__init__(
+            data_root=data_root,
+            split=split,
+            target_length=target_length,
+            chunk_size_second=chunk_size_second,
+            stems=stems_no_mixture,
+            fs=fs,
+            npy_memmap=npy_memmap,
+        )
+
+        self.stems = stems
+        self.stems_no_mixture = stems_no_mixture
+
+    def __getitem__(self, index: int) -> DataDict:
+
+        data_ = super().__getitem__(index)
+
+        dry = data_["audio"]["speech"][:]
+        n_samples = dry.shape[-1]
+
+        wet_level = np.random.rand()
+
+        speech = pb.Reverb(
+            room_size=np.random.rand(),
+            damping=np.random.rand(),
+            wet_level=wet_level,
+            dry_level=(1 - wet_level),
+            width=np.random.rand(),
+        ).process(dry, self.fs, buffer_size=8192 * 4)[..., :n_samples]
+
+        data_["audio"]["speech"] = speech
+
+        data_["audio"]["mixture"] = sum(
+            [data_["audio"][s] for s in self.stems_no_mixture]
+        )
+
+        return data_
+
+    def __len__(self) -> int:
+        return super().__len__()
+
+
+if __name__ == "__main__":
+
+    from pprint import pprint
+    from tqdm.auto import tqdm
+
+    for split_ in ["train", "val", "test"]:
+        ds = DivideAndRemasterRandomChunkDatasetWithSpeechReverb(
+            data_root="$DATA_ROOT/DnR/v2np",
+            split=split_,
+            target_length=100,
+            chunk_size_second=6.0,
+        )
+
+        print(split_, len(ds))
+
+        for track_ in tqdm(ds):  # type: ignore
+            pprint(track_)
+            track_["audio"] = {k: v.shape for k, v in track_["audio"].items()}
+            pprint(track_)
+
+        break

mvsepless/models/bandit/core/data/dnr/preprocess.py

@@ -1,51 +1,51 @@
-import glob
-import os
-from typing import Tuple
-
-import numpy as np
-import torchaudio as ta
-from tqdm.contrib.concurrent import process_map
-
-
-def process_one(inputs: Tuple[str, str, int]) -> None:
-    infile, outfile, target_fs = inputs
-
-    dir = os.path.dirname(outfile)
-    os.makedirs(dir, exist_ok=True)
-
-    data, fs = ta.load(infile)
-
-    if fs != target_fs:
-        data = ta.functional.resample(
-            data, fs, target_fs, resampling_method="sinc_interp_kaiser"
-        )
-        fs = target_fs
-
-    data = data.numpy()
-    data = data.astype(np.float32)
-
-    if os.path.exists(outfile):
-        data_ = np.load(outfile)
-        if np.allclose(data, data_):
-            return
-
-    np.save(outfile, data)
-
-
-def preprocess(data_path: str, output_path: str, fs: int) -> None:
-    files = glob.glob(os.path.join(data_path, "**", "*.wav"), recursive=True)
-    print(files)
-    outfiles = [
-        f.replace(data_path, output_path).replace(".wav", ".npy") for f in files
-    ]
-
-    os.makedirs(output_path, exist_ok=True)
-    inputs = list(zip(files, outfiles, [fs] * len(files)))
-
-    process_map(process_one, inputs, chunksize=32)
-
-
-if __name__ == "__main__":
-    import fire
-
-    fire.Fire()
+import glob
+import os
+from typing import Tuple
+
+import numpy as np
+import torchaudio as ta
+from tqdm.contrib.concurrent import process_map
+
+
+def process_one(inputs: Tuple[str, str, int]) -> None:
+    infile, outfile, target_fs = inputs
+
+    dir = os.path.dirname(outfile)
+    os.makedirs(dir, exist_ok=True)
+
+    data, fs = ta.load(infile)
+
+    if fs != target_fs:
+        data = ta.functional.resample(
+            data, fs, target_fs, resampling_method="sinc_interp_kaiser"
+        )
+        fs = target_fs
+
+    data = data.numpy()
+    data = data.astype(np.float32)
+
+    if os.path.exists(outfile):
+        data_ = np.load(outfile)
+        if np.allclose(data, data_):
+            return
+
+    np.save(outfile, data)
+
+
+def preprocess(data_path: str, output_path: str, fs: int) -> None:
+    files = glob.glob(os.path.join(data_path, "**", "*.wav"), recursive=True)
+    print(files)
+    outfiles = [
+        f.replace(data_path, output_path).replace(".wav", ".npy") for f in files
+    ]
+
+    os.makedirs(output_path, exist_ok=True)
+    inputs = list(zip(files, outfiles, [fs] * len(files)))
+
+    process_map(process_one, inputs, chunksize=32)
+
+
+if __name__ == "__main__":
+    import fire
+
+    fire.Fire()

mvsepless/models/bandit/core/data/musdb/datamodule.py

@@ -1,75 +1,75 @@
-import os.path
-from typing import Mapping, Optional
-
-import pytorch_lightning as pl
-
-from .dataset import (
-    MUSDB18BaseDataset,
-    MUSDB18FullTrackDataset,
-    MUSDB18SadDataset,
-    MUSDB18SadOnTheFlyAugmentedDataset,
-)
-
-
-def MUSDB18DataModule(
-    data_root: str = "$DATA_ROOT/MUSDB18/HQ",
-    target_stem: str = "vocals",
-    batch_size: int = 2,
-    num_workers: int = 8,
-    train_kwargs: Optional[Mapping] = None,
-    val_kwargs: Optional[Mapping] = None,
-    test_kwargs: Optional[Mapping] = None,
-    datamodule_kwargs: Optional[Mapping] = None,
-    use_on_the_fly: bool = True,
-    npy_memmap: bool = True,
-) -> pl.LightningDataModule:
-    if train_kwargs is None:
-        train_kwargs = {}
-
-    if val_kwargs is None:
-        val_kwargs = {}
-
-    if test_kwargs is None:
-        test_kwargs = {}
-
-    if datamodule_kwargs is None:
-        datamodule_kwargs = {}
-
-    train_dataset: MUSDB18BaseDataset
-
-    if use_on_the_fly:
-        train_dataset = MUSDB18SadOnTheFlyAugmentedDataset(
-            data_root=os.path.join(data_root, "saded-np"),
-            split="train",
-            target_stem=target_stem,
-            **train_kwargs,
-        )
-    else:
-        train_dataset = MUSDB18SadDataset(
-            data_root=os.path.join(data_root, "saded-np"),
-            split="train",
-            target_stem=target_stem,
-            **train_kwargs,
-        )
-
-    datamodule = pl.LightningDataModule.from_datasets(
-        train_dataset=train_dataset,
-        val_dataset=MUSDB18SadDataset(
-            data_root=os.path.join(data_root, "saded-np"),
-            split="val",
-            target_stem=target_stem,
-            **val_kwargs,
-        ),
-        test_dataset=MUSDB18FullTrackDataset(
-            data_root=os.path.join(data_root, "canonical"), split="test", **test_kwargs
-        ),
-        batch_size=batch_size,
-        num_workers=num_workers,
-        **datamodule_kwargs,
-    )
-
-    datamodule.predict_dataloader = (  # type: ignore[method-assign]
-        datamodule.test_dataloader
-    )
-
-    return datamodule
+import os.path
+from typing import Mapping, Optional
+
+import pytorch_lightning as pl
+
+from .dataset import (
+    MUSDB18BaseDataset,
+    MUSDB18FullTrackDataset,
+    MUSDB18SadDataset,
+    MUSDB18SadOnTheFlyAugmentedDataset,
+)
+
+
+def MUSDB18DataModule(
+    data_root: str = "$DATA_ROOT/MUSDB18/HQ",
+    target_stem: str = "vocals",
+    batch_size: int = 2,
+    num_workers: int = 8,
+    train_kwargs: Optional[Mapping] = None,
+    val_kwargs: Optional[Mapping] = None,
+    test_kwargs: Optional[Mapping] = None,
+    datamodule_kwargs: Optional[Mapping] = None,
+    use_on_the_fly: bool = True,
+    npy_memmap: bool = True,
+) -> pl.LightningDataModule:
+    if train_kwargs is None:
+        train_kwargs = {}
+
+    if val_kwargs is None:
+        val_kwargs = {}
+
+    if test_kwargs is None:
+        test_kwargs = {}
+
+    if datamodule_kwargs is None:
+        datamodule_kwargs = {}
+
+    train_dataset: MUSDB18BaseDataset
+
+    if use_on_the_fly:
+        train_dataset = MUSDB18SadOnTheFlyAugmentedDataset(
+            data_root=os.path.join(data_root, "saded-np"),
+            split="train",
+            target_stem=target_stem,
+            **train_kwargs,
+        )
+    else:
+        train_dataset = MUSDB18SadDataset(
+            data_root=os.path.join(data_root, "saded-np"),
+            split="train",
+            target_stem=target_stem,
+            **train_kwargs,
+        )
+
+    datamodule = pl.LightningDataModule.from_datasets(
+        train_dataset=train_dataset,
+        val_dataset=MUSDB18SadDataset(
+            data_root=os.path.join(data_root, "saded-np"),
+            split="val",
+            target_stem=target_stem,
+            **val_kwargs,
+        ),
+        test_dataset=MUSDB18FullTrackDataset(
+            data_root=os.path.join(data_root, "canonical"), split="test", **test_kwargs
+        ),
+        batch_size=batch_size,
+        num_workers=num_workers,
+        **datamodule_kwargs,
+    )
+
+    datamodule.predict_dataloader = (  # type: ignore[method-assign]
+        datamodule.test_dataloader
+    )
+
+    return datamodule

mvsepless/models/bandit/core/data/musdb/dataset.py

@@ -1,241 +1,241 @@
-import os
-from abc import ABC
-from typing import List, Optional, Tuple
-
-import numpy as np
-import torch
-import torchaudio as ta
-from torch.utils import data
-
-from .._types import AudioDict, DataDict
-from ..base import BaseSourceSeparationDataset
-
-
-class MUSDB18BaseDataset(BaseSourceSeparationDataset, ABC):
-
-    ALLOWED_STEMS = ["mixture", "vocals", "bass", "drums", "other"]
-
-    def __init__(
-        self,
-        split: str,
-        stems: List[str],
-        files: List[str],
-        data_path: str,
-        fs: int = 44100,
-        npy_memmap=False,
-    ) -> None:
-        super().__init__(
-            split=split,
-            stems=stems,
-            files=files,
-            data_path=data_path,
-            fs=fs,
-            npy_memmap=npy_memmap,
-            recompute_mixture=False,
-        )
-
-    def get_stem(self, *, stem: str, identifier) -> torch.Tensor:
-        track = identifier["track"]
-        path = os.path.join(self.data_path, track)
-
-        if self.npy_memmap:
-            audio = np.load(os.path.join(path, f"{stem}.wav.npy"), mmap_mode="r")
-        else:
-            audio, _ = ta.load(os.path.join(path, f"{stem}.wav"))
-
-        return audio
-
-    def get_identifier(self, index):
-        return dict(track=self.files[index])
-
-    def __getitem__(self, index: int) -> DataDict:
-        identifier = self.get_identifier(index)
-        audio = self.get_audio(identifier)
-
-        return {"audio": audio, "track": f"{self.split}/{identifier['track']}"}
-
-
-class MUSDB18FullTrackDataset(MUSDB18BaseDataset):
-
-    N_TRAIN_TRACKS = 100
-    N_TEST_TRACKS = 50
-    VALIDATION_FILES = [
-        "Actions - One Minute Smile",
-        "Clara Berry And Wooldog - Waltz For My Victims",
-        "Johnny Lokke - Promises & Lies",
-        "Patrick Talbot - A Reason To Leave",
-        "Triviul - Angelsaint",
-        "Alexander Ross - Goodbye Bolero",
-        "Fergessen - Nos Palpitants",
-        "Leaf - Summerghost",
-        "Skelpolu - Human Mistakes",
-        "Young Griffo - Pennies",
-        "ANiMAL - Rockshow",
-        "James May - On The Line",
-        "Meaxic - Take A Step",
-        "Traffic Experiment - Sirens",
-    ]
-
-    def __init__(
-        self, data_root: str, split: str, stems: Optional[List[str]] = None
-    ) -> None:
-
-        if stems is None:
-            stems = self.ALLOWED_STEMS
-        self.stems = stems
-
-        if split == "test":
-            subset = "test"
-        elif split in ["train", "val"]:
-            subset = "train"
-        else:
-            raise NameError
-
-        data_path = os.path.join(data_root, subset)
-
-        files = sorted(os.listdir(data_path))
-        files = [f for f in files if not f.startswith(".")]
-        if subset == "train":
-            assert len(files) == 100, len(files)
-            if split == "train":
-                files = [f for f in files if f not in self.VALIDATION_FILES]
-                assert len(files) == 100 - len(self.VALIDATION_FILES)
-            else:
-                files = [f for f in files if f in self.VALIDATION_FILES]
-                assert len(files) == len(self.VALIDATION_FILES)
-        else:
-            split = "test"
-            assert len(files) == 50
-
-        self.n_tracks = len(files)
-
-        super().__init__(data_path=data_path, split=split, stems=stems, files=files)
-
-    def __len__(self) -> int:
-        return self.n_tracks
-
-
-class MUSDB18SadDataset(MUSDB18BaseDataset):
-    def __init__(
-        self,
-        data_root: str,
-        split: str,
-        target_stem: str,
-        stems: Optional[List[str]] = None,
-        target_length: Optional[int] = None,
-        npy_memmap=False,
-    ) -> None:
-
-        if stems is None:
-            stems = self.ALLOWED_STEMS
-
-        data_path = os.path.join(data_root, target_stem, split)
-
-        files = sorted(os.listdir(data_path))
-        files = [f for f in files if not f.startswith(".")]
-
-        super().__init__(
-            data_path=data_path,
-            split=split,
-            stems=stems,
-            files=files,
-            npy_memmap=npy_memmap,
-        )
-        self.n_segments = len(files)
-        self.target_stem = target_stem
-        self.target_length = (
-            target_length if target_length is not None else self.n_segments
-        )
-
-    def __len__(self) -> int:
-        return self.target_length
-
-    def __getitem__(self, index: int) -> DataDict:
-
-        index = index % self.n_segments
-
-        return super().__getitem__(index)
-
-    def get_identifier(self, index):
-        return super().get_identifier(index % self.n_segments)
-
-
-class MUSDB18SadOnTheFlyAugmentedDataset(MUSDB18SadDataset):
-    def __init__(
-        self,
-        data_root: str,
-        split: str,
-        target_stem: str,
-        stems: Optional[List[str]] = None,
-        target_length: int = 20000,
-        apply_probability: Optional[float] = None,
-        chunk_size_second: float = 3.0,
-        random_scale_range_db: Tuple[float, float] = (-10, 10),
-        drop_probability: float = 0.1,
-        rescale: bool = True,
-    ) -> None:
-        super().__init__(data_root, split, target_stem, stems)
-
-        if apply_probability is None:
-            apply_probability = (target_length - self.n_segments) / target_length
-
-        self.apply_probability = apply_probability
-        self.drop_probability = drop_probability
-        self.chunk_size_second = chunk_size_second
-        self.random_scale_range_db = random_scale_range_db
-        self.rescale = rescale
-
-        self.chunk_size_sample = int(self.chunk_size_second * self.fs)
-        self.target_length = target_length
-
-    def __len__(self) -> int:
-        return self.target_length
-
-    def __getitem__(self, index: int) -> DataDict:
-
-        index = index % self.n_segments
-
-        audio = {}
-        identifier = self.get_identifier(index)
-
-        for stem in self.stems_no_mixture:
-            if stem == self.target_stem:
-                identifier_ = identifier
-            else:
-                if np.random.rand() < self.apply_probability:
-                    index_ = np.random.randint(self.n_segments)
-                    identifier_ = self.get_identifier(index_)
-                else:
-                    identifier_ = identifier
-
-            audio[stem] = self.get_stem(stem=stem, identifier=identifier_)
-
-            if self.chunk_size_sample < audio[stem].shape[-1]:
-                chunk_start = np.random.randint(
-                    audio[stem].shape[-1] - self.chunk_size_sample
-                )
-            else:
-                chunk_start = 0
-
-            if np.random.rand() < self.drop_probability:
-                linear_scale = 0.0
-            else:
-                db_scale = np.random.uniform(*self.random_scale_range_db)
-                linear_scale = np.power(10, db_scale / 20)
-            audio[stem][..., chunk_start : chunk_start + self.chunk_size_sample] = (
-                linear_scale
-                * audio[stem][..., chunk_start : chunk_start + self.chunk_size_sample]
-            )
-
-        audio["mixture"] = self.compute_mixture(audio)
-
-        if self.rescale:
-            max_abs_val = max(
-                [torch.max(torch.abs(audio[stem])) for stem in self.stems]
-            )  # type: ignore[type-var]
-            if max_abs_val > 1:
-                audio = {k: v / max_abs_val for k, v in audio.items()}
-
-        track = identifier["track"]
-
-        return {"audio": audio, "track": f"{self.split}/{track}"}
+import os
+from abc import ABC
+from typing import List, Optional, Tuple
+
+import numpy as np
+import torch
+import torchaudio as ta
+from torch.utils import data
+
+from .._types import AudioDict, DataDict
+from ..base import BaseSourceSeparationDataset
+
+
+class MUSDB18BaseDataset(BaseSourceSeparationDataset, ABC):
+
+    ALLOWED_STEMS = ["mixture", "vocals", "bass", "drums", "other"]
+
+    def __init__(
+        self,
+        split: str,
+        stems: List[str],
+        files: List[str],
+        data_path: str,
+        fs: int = 44100,
+        npy_memmap=False,
+    ) -> None:
+        super().__init__(
+            split=split,
+            stems=stems,
+            files=files,
+            data_path=data_path,
+            fs=fs,
+            npy_memmap=npy_memmap,
+            recompute_mixture=False,
+        )
+
+    def get_stem(self, *, stem: str, identifier) -> torch.Tensor:
+        track = identifier["track"]
+        path = os.path.join(self.data_path, track)
+
+        if self.npy_memmap:
+            audio = np.load(os.path.join(path, f"{stem}.wav.npy"), mmap_mode="r")
+        else:
+            audio, _ = ta.load(os.path.join(path, f"{stem}.wav"))
+
+        return audio
+
+    def get_identifier(self, index):
+        return dict(track=self.files[index])
+
+    def __getitem__(self, index: int) -> DataDict:
+        identifier = self.get_identifier(index)
+        audio = self.get_audio(identifier)
+
+        return {"audio": audio, "track": f"{self.split}/{identifier['track']}"}
+
+
+class MUSDB18FullTrackDataset(MUSDB18BaseDataset):
+
+    N_TRAIN_TRACKS = 100
+    N_TEST_TRACKS = 50
+    VALIDATION_FILES = [
+        "Actions - One Minute Smile",
+        "Clara Berry And Wooldog - Waltz For My Victims",
+        "Johnny Lokke - Promises & Lies",
+        "Patrick Talbot - A Reason To Leave",
+        "Triviul - Angelsaint",
+        "Alexander Ross - Goodbye Bolero",
+        "Fergessen - Nos Palpitants",
+        "Leaf - Summerghost",
+        "Skelpolu - Human Mistakes",
+        "Young Griffo - Pennies",
+        "ANiMAL - Rockshow",
+        "James May - On The Line",
+        "Meaxic - Take A Step",
+        "Traffic Experiment - Sirens",
+    ]
+
+    def __init__(
+        self, data_root: str, split: str, stems: Optional[List[str]] = None
+    ) -> None:
+
+        if stems is None:
+            stems = self.ALLOWED_STEMS
+        self.stems = stems
+
+        if split == "test":
+            subset = "test"
+        elif split in ["train", "val"]:
+            subset = "train"
+        else:
+            raise NameError
+
+        data_path = os.path.join(data_root, subset)
+
+        files = sorted(os.listdir(data_path))
+        files = [f for f in files if not f.startswith(".")]
+        if subset == "train":
+            assert len(files) == 100, len(files)
+            if split == "train":
+                files = [f for f in files if f not in self.VALIDATION_FILES]
+                assert len(files) == 100 - len(self.VALIDATION_FILES)
+            else:
+                files = [f for f in files if f in self.VALIDATION_FILES]
+                assert len(files) == len(self.VALIDATION_FILES)
+        else:
+            split = "test"
+            assert len(files) == 50
+
+        self.n_tracks = len(files)
+
+        super().__init__(data_path=data_path, split=split, stems=stems, files=files)
+
+    def __len__(self) -> int:
+        return self.n_tracks
+
+
+class MUSDB18SadDataset(MUSDB18BaseDataset):
+    def __init__(
+        self,
+        data_root: str,
+        split: str,
+        target_stem: str,
+        stems: Optional[List[str]] = None,
+        target_length: Optional[int] = None,
+        npy_memmap=False,
+    ) -> None:
+
+        if stems is None:
+            stems = self.ALLOWED_STEMS
+
+        data_path = os.path.join(data_root, target_stem, split)
+
+        files = sorted(os.listdir(data_path))
+        files = [f for f in files if not f.startswith(".")]
+
+        super().__init__(
+            data_path=data_path,
+            split=split,
+            stems=stems,
+            files=files,
+            npy_memmap=npy_memmap,
+        )
+        self.n_segments = len(files)
+        self.target_stem = target_stem
+        self.target_length = (
+            target_length if target_length is not None else self.n_segments
+        )
+
+    def __len__(self) -> int:
+        return self.target_length
+
+    def __getitem__(self, index: int) -> DataDict:
+
+        index = index % self.n_segments
+
+        return super().__getitem__(index)
+
+    def get_identifier(self, index):
+        return super().get_identifier(index % self.n_segments)
+
+
+class MUSDB18SadOnTheFlyAugmentedDataset(MUSDB18SadDataset):
+    def __init__(
+        self,
+        data_root: str,
+        split: str,
+        target_stem: str,
+        stems: Optional[List[str]] = None,
+        target_length: int = 20000,
+        apply_probability: Optional[float] = None,
+        chunk_size_second: float = 3.0,
+        random_scale_range_db: Tuple[float, float] = (-10, 10),
+        drop_probability: float = 0.1,
+        rescale: bool = True,
+    ) -> None:
+        super().__init__(data_root, split, target_stem, stems)
+
+        if apply_probability is None:
+            apply_probability = (target_length - self.n_segments) / target_length
+
+        self.apply_probability = apply_probability
+        self.drop_probability = drop_probability
+        self.chunk_size_second = chunk_size_second
+        self.random_scale_range_db = random_scale_range_db
+        self.rescale = rescale
+
+        self.chunk_size_sample = int(self.chunk_size_second * self.fs)
+        self.target_length = target_length
+
+    def __len__(self) -> int:
+        return self.target_length
+
+    def __getitem__(self, index: int) -> DataDict:
+
+        index = index % self.n_segments
+
+        audio = {}
+        identifier = self.get_identifier(index)
+
+        for stem in self.stems_no_mixture:
+            if stem == self.target_stem:
+                identifier_ = identifier
+            else:
+                if np.random.rand() < self.apply_probability:
+                    index_ = np.random.randint(self.n_segments)
+                    identifier_ = self.get_identifier(index_)
+                else:
+                    identifier_ = identifier
+
+            audio[stem] = self.get_stem(stem=stem, identifier=identifier_)
+
+            if self.chunk_size_sample < audio[stem].shape[-1]:
+                chunk_start = np.random.randint(
+                    audio[stem].shape[-1] - self.chunk_size_sample
+                )
+            else:
+                chunk_start = 0
+
+            if np.random.rand() < self.drop_probability:
+                linear_scale = 0.0
+            else:
+                db_scale = np.random.uniform(*self.random_scale_range_db)
+                linear_scale = np.power(10, db_scale / 20)
+            audio[stem][..., chunk_start : chunk_start + self.chunk_size_sample] = (
+                linear_scale
+                * audio[stem][..., chunk_start : chunk_start + self.chunk_size_sample]
+            )
+
+        audio["mixture"] = self.compute_mixture(audio)
+
+        if self.rescale:
+            max_abs_val = max(
+                [torch.max(torch.abs(audio[stem])) for stem in self.stems]
+            )  # type: ignore[type-var]
+            if max_abs_val > 1:
+                audio = {k: v / max_abs_val for k, v in audio.items()}
+
+        track = identifier["track"]
+
+        return {"audio": audio, "track": f"{self.split}/{track}"}

mvsepless/models/bandit/core/data/musdb/preprocess.py

@@ -1,223 +1,223 @@
-import glob
-import os
-
-import numpy as np
-import torch
-import torchaudio as ta
-from torch import nn
-from torch.nn import functional as F
-from tqdm.contrib.concurrent import process_map
-
-from .._types import DataDict
-from .dataset import MUSDB18FullTrackDataset
-import pyloudnorm as pyln
-
-
-class SourceActivityDetector(nn.Module):
-    def __init__(
-        self,
-        analysis_stem: str,
-        output_path: str,
-        fs: int = 44100,
-        segment_length_second: float = 6.0,
-        hop_length_second: float = 3.0,
-        n_chunks: int = 10,
-        chunk_epsilon: float = 1e-5,
-        energy_threshold_quantile: float = 0.15,
-        segment_epsilon: float = 1e-3,
-        salient_proportion_threshold: float = 0.5,
-        target_lufs: float = -24,
-    ) -> None:
-        super().__init__()
-
-        self.fs = fs
-        self.segment_length = int(segment_length_second * self.fs)
-        self.hop_length = int(hop_length_second * self.fs)
-        self.n_chunks = n_chunks
-        assert self.segment_length % self.n_chunks == 0
-        self.chunk_size = self.segment_length // self.n_chunks
-        self.chunk_epsilon = chunk_epsilon
-        self.energy_threshold_quantile = energy_threshold_quantile
-        self.segment_epsilon = segment_epsilon
-        self.salient_proportion_threshold = salient_proportion_threshold
-        self.analysis_stem = analysis_stem
-
-        self.meter = pyln.Meter(self.fs)
-        self.target_lufs = target_lufs
-
-        self.output_path = output_path
-
-    def forward(self, data: DataDict) -> None:
-
-        stem_ = self.analysis_stem if (self.analysis_stem != "none") else "mixture"
-
-        x = data["audio"][stem_]
-
-        xnp = x.numpy()
-        loudness = self.meter.integrated_loudness(xnp.T)
-
-        for stem in data["audio"]:
-            s = data["audio"][stem]
-            s = pyln.normalize.loudness(s.numpy().T, loudness, self.target_lufs).T
-            s = torch.as_tensor(s)
-            data["audio"][stem] = s
-
-        if x.ndim == 3:
-            assert x.shape[0] == 1
-            x = x[0]
-
-        n_chan, n_samples = x.shape
-
-        n_segments = (
-            int(np.ceil((n_samples - self.segment_length) / self.hop_length)) + 1
-        )
-
-        segments = torch.zeros((n_segments, n_chan, self.segment_length))
-        for i in range(n_segments):
-            start = i * self.hop_length
-            end = start + self.segment_length
-            end = min(end, n_samples)
-
-            xseg = x[:, start:end]
-
-            if end - start < self.segment_length:
-                xseg = F.pad(
-                    xseg, pad=(0, self.segment_length - (end - start)), value=torch.nan
-                )
-
-            segments[i, :, :] = xseg
-
-        chunks = segments.reshape((n_segments, n_chan, self.n_chunks, self.chunk_size))
-
-        if self.analysis_stem != "none":
-            chunk_energies = torch.mean(torch.square(chunks), dim=(1, 3))
-            chunk_energies = torch.nan_to_num(chunk_energies, nan=0)
-            chunk_energies[chunk_energies == 0] = self.chunk_epsilon
-
-            energy_threshold = torch.nanquantile(
-                chunk_energies, q=self.energy_threshold_quantile
-            )
-
-            if energy_threshold < self.segment_epsilon:
-                energy_threshold = self.segment_epsilon  # type: ignore[assignment]
-
-            chunks_above_threshold = chunk_energies > energy_threshold
-            n_chunks_above_threshold = torch.mean(
-                chunks_above_threshold.to(torch.float), dim=-1
-            )
-
-            segment_above_threshold = (
-                n_chunks_above_threshold > self.salient_proportion_threshold
-            )
-
-            if torch.sum(segment_above_threshold) == 0:
-                return
-
-        else:
-            segment_above_threshold = torch.ones((n_segments,))
-
-        for i in range(n_segments):
-            if not segment_above_threshold[i]:
-                continue
-
-            outpath = os.path.join(
-                self.output_path,
-                self.analysis_stem,
-                f"{data['track']} - {self.analysis_stem}{i:03d}",
-            )
-            os.makedirs(outpath, exist_ok=True)
-
-            for stem in data["audio"]:
-                if stem == self.analysis_stem:
-                    segment = torch.nan_to_num(segments[i, :, :], nan=0)
-                else:
-                    start = i * self.hop_length
-                    end = start + self.segment_length
-                    end = min(n_samples, end)
-
-                    segment = data["audio"][stem][:, start:end]
-
-                    if end - start < self.segment_length:
-                        segment = F.pad(
-                            segment, (0, self.segment_length - (end - start))
-                        )
-
-                assert segment.shape[-1] == self.segment_length, segment.shape
-
-                np.save(os.path.join(outpath, f"{stem}.wav"), segment)
-
-
-def preprocess(
-    analysis_stem: str,
-    output_path: str = "/data/MUSDB18/HQ/saded-np",
-    fs: int = 44100,
-    segment_length_second: float = 6.0,
-    hop_length_second: float = 3.0,
-    n_chunks: int = 10,
-    chunk_epsilon: float = 1e-5,
-    energy_threshold_quantile: float = 0.15,
-    segment_epsilon: float = 1e-3,
-    salient_proportion_threshold: float = 0.5,
-) -> None:
-
-    sad = SourceActivityDetector(
-        analysis_stem=analysis_stem,
-        output_path=output_path,
-        fs=fs,
-        segment_length_second=segment_length_second,
-        hop_length_second=hop_length_second,
-        n_chunks=n_chunks,
-        chunk_epsilon=chunk_epsilon,
-        energy_threshold_quantile=energy_threshold_quantile,
-        segment_epsilon=segment_epsilon,
-        salient_proportion_threshold=salient_proportion_threshold,
-    )
-
-    for split in ["train", "val", "test"]:
-        ds = MUSDB18FullTrackDataset(
-            data_root="/data/MUSDB18/HQ/canonical",
-            split=split,
-        )
-
-        tracks = []
-        for i, track in enumerate(tqdm(ds, total=len(ds))):
-            if i % 32 == 0 and tracks:
-                process_map(sad, tracks, max_workers=8)
-                tracks = []
-            tracks.append(track)
-        process_map(sad, tracks, max_workers=8)
-
-
-def loudness_norm_one(inputs):
-    infile, outfile, target_lufs = inputs
-
-    audio, fs = ta.load(infile)
-    audio = audio.mean(dim=0, keepdim=True).numpy().T
-
-    meter = pyln.Meter(fs)
-    loudness = meter.integrated_loudness(audio)
-    audio = pyln.normalize.loudness(audio, loudness, target_lufs)
-
-    os.makedirs(os.path.dirname(outfile), exist_ok=True)
-    np.save(outfile, audio.T)
-
-
-def loudness_norm(
-    data_path: str,
-    target_lufs=-17.0,
-):
-    files = glob.glob(os.path.join(data_path, "**", "*.wav"), recursive=True)
-
-    outfiles = [f.replace(".wav", ".npy").replace("saded", "saded-np") for f in files]
-
-    files = [(f, o, target_lufs) for f, o in zip(files, outfiles)]
-
-    process_map(loudness_norm_one, files, chunksize=2)
-
-
-if __name__ == "__main__":
-
-    from tqdm.auto import tqdm
-    import fire
-
-    fire.Fire()
+import glob
+import os
+
+import numpy as np
+import torch
+import torchaudio as ta
+from torch import nn
+from torch.nn import functional as F
+from tqdm.contrib.concurrent import process_map
+
+from .._types import DataDict
+from .dataset import MUSDB18FullTrackDataset
+import pyloudnorm as pyln
+
+
+class SourceActivityDetector(nn.Module):
+    def __init__(
+        self,
+        analysis_stem: str,
+        output_path: str,
+        fs: int = 44100,
+        segment_length_second: float = 6.0,
+        hop_length_second: float = 3.0,
+        n_chunks: int = 10,
+        chunk_epsilon: float = 1e-5,
+        energy_threshold_quantile: float = 0.15,
+        segment_epsilon: float = 1e-3,
+        salient_proportion_threshold: float = 0.5,
+        target_lufs: float = -24,
+    ) -> None:
+        super().__init__()
+
+        self.fs = fs
+        self.segment_length = int(segment_length_second * self.fs)
+        self.hop_length = int(hop_length_second * self.fs)
+        self.n_chunks = n_chunks
+        assert self.segment_length % self.n_chunks == 0
+        self.chunk_size = self.segment_length // self.n_chunks
+        self.chunk_epsilon = chunk_epsilon
+        self.energy_threshold_quantile = energy_threshold_quantile
+        self.segment_epsilon = segment_epsilon
+        self.salient_proportion_threshold = salient_proportion_threshold
+        self.analysis_stem = analysis_stem
+
+        self.meter = pyln.Meter(self.fs)
+        self.target_lufs = target_lufs
+
+        self.output_path = output_path
+
+    def forward(self, data: DataDict) -> None:
+
+        stem_ = self.analysis_stem if (self.analysis_stem != "none") else "mixture"
+
+        x = data["audio"][stem_]
+
+        xnp = x.numpy()
+        loudness = self.meter.integrated_loudness(xnp.T)
+
+        for stem in data["audio"]:
+            s = data["audio"][stem]
+            s = pyln.normalize.loudness(s.numpy().T, loudness, self.target_lufs).T
+            s = torch.as_tensor(s)
+            data["audio"][stem] = s
+
+        if x.ndim == 3:
+            assert x.shape[0] == 1
+            x = x[0]
+
+        n_chan, n_samples = x.shape
+
+        n_segments = (
+            int(np.ceil((n_samples - self.segment_length) / self.hop_length)) + 1
+        )
+
+        segments = torch.zeros((n_segments, n_chan, self.segment_length))
+        for i in range(n_segments):
+            start = i * self.hop_length
+            end = start + self.segment_length
+            end = min(end, n_samples)
+
+            xseg = x[:, start:end]
+
+            if end - start < self.segment_length:
+                xseg = F.pad(
+                    xseg, pad=(0, self.segment_length - (end - start)), value=torch.nan
+                )
+
+            segments[i, :, :] = xseg
+
+        chunks = segments.reshape((n_segments, n_chan, self.n_chunks, self.chunk_size))
+
+        if self.analysis_stem != "none":
+            chunk_energies = torch.mean(torch.square(chunks), dim=(1, 3))
+            chunk_energies = torch.nan_to_num(chunk_energies, nan=0)
+            chunk_energies[chunk_energies == 0] = self.chunk_epsilon
+
+            energy_threshold = torch.nanquantile(
+                chunk_energies, q=self.energy_threshold_quantile
+            )
+
+            if energy_threshold < self.segment_epsilon:
+                energy_threshold = self.segment_epsilon  # type: ignore[assignment]
+
+            chunks_above_threshold = chunk_energies > energy_threshold
+            n_chunks_above_threshold = torch.mean(
+                chunks_above_threshold.to(torch.float), dim=-1
+            )
+
+            segment_above_threshold = (
+                n_chunks_above_threshold > self.salient_proportion_threshold
+            )
+
+            if torch.sum(segment_above_threshold) == 0:
+                return
+
+        else:
+            segment_above_threshold = torch.ones((n_segments,))
+
+        for i in range(n_segments):
+            if not segment_above_threshold[i]:
+                continue
+
+            outpath = os.path.join(
+                self.output_path,
+                self.analysis_stem,
+                f"{data['track']} - {self.analysis_stem}{i:03d}",
+            )
+            os.makedirs(outpath, exist_ok=True)
+
+            for stem in data["audio"]:
+                if stem == self.analysis_stem:
+                    segment = torch.nan_to_num(segments[i, :, :], nan=0)
+                else:
+                    start = i * self.hop_length
+                    end = start + self.segment_length
+                    end = min(n_samples, end)
+
+                    segment = data["audio"][stem][:, start:end]
+
+                    if end - start < self.segment_length:
+                        segment = F.pad(
+                            segment, (0, self.segment_length - (end - start))
+                        )
+
+                assert segment.shape[-1] == self.segment_length, segment.shape
+
+                np.save(os.path.join(outpath, f"{stem}.wav"), segment)
+
+
+def preprocess(
+    analysis_stem: str,
+    output_path: str = "/data/MUSDB18/HQ/saded-np",
+    fs: int = 44100,
+    segment_length_second: float = 6.0,
+    hop_length_second: float = 3.0,
+    n_chunks: int = 10,
+    chunk_epsilon: float = 1e-5,
+    energy_threshold_quantile: float = 0.15,
+    segment_epsilon: float = 1e-3,
+    salient_proportion_threshold: float = 0.5,
+) -> None:
+
+    sad = SourceActivityDetector(
+        analysis_stem=analysis_stem,
+        output_path=output_path,
+        fs=fs,
+        segment_length_second=segment_length_second,
+        hop_length_second=hop_length_second,
+        n_chunks=n_chunks,
+        chunk_epsilon=chunk_epsilon,
+        energy_threshold_quantile=energy_threshold_quantile,
+        segment_epsilon=segment_epsilon,
+        salient_proportion_threshold=salient_proportion_threshold,
+    )
+
+    for split in ["train", "val", "test"]:
+        ds = MUSDB18FullTrackDataset(
+            data_root="/data/MUSDB18/HQ/canonical",
+            split=split,
+        )
+
+        tracks = []
+        for i, track in enumerate(tqdm(ds, total=len(ds))):
+            if i % 32 == 0 and tracks:
+                process_map(sad, tracks, max_workers=8)
+                tracks = []
+            tracks.append(track)
+        process_map(sad, tracks, max_workers=8)
+
+
+def loudness_norm_one(inputs):
+    infile, outfile, target_lufs = inputs
+
+    audio, fs = ta.load(infile)
+    audio = audio.mean(dim=0, keepdim=True).numpy().T
+
+    meter = pyln.Meter(fs)
+    loudness = meter.integrated_loudness(audio)
+    audio = pyln.normalize.loudness(audio, loudness, target_lufs)
+
+    os.makedirs(os.path.dirname(outfile), exist_ok=True)
+    np.save(outfile, audio.T)
+
+
+def loudness_norm(
+    data_path: str,
+    target_lufs=-17.0,
+):
+    files = glob.glob(os.path.join(data_path, "**", "*.wav"), recursive=True)
+
+    outfiles = [f.replace(".wav", ".npy").replace("saded", "saded-np") for f in files]
+
+    files = [(f, o, target_lufs) for f, o in zip(files, outfiles)]
+
+    process_map(loudness_norm_one, files, chunksize=2)
+
+
+if __name__ == "__main__":
+
+    from tqdm.auto import tqdm
+    import fire
+
+    fire.Fire()

mvsepless/models/bandit/core/data/musdb/validation.yaml

@@ -1,15 +1,15 @@
-validation:
-  - 'Actions - One Minute Smile'
-  - 'Clara Berry And Wooldog - Waltz For My Victims'
-  - 'Johnny Lokke - Promises & Lies'
-  - 'Patrick Talbot - A Reason To Leave'
-  - 'Triviul - Angelsaint'
-  - 'Alexander Ross - Goodbye Bolero'
-  - 'Fergessen - Nos Palpitants'
-  - 'Leaf - Summerghost'
-  - 'Skelpolu - Human Mistakes'
-  - 'Young Griffo - Pennies'
-  - 'ANiMAL - Rockshow'
-  - 'James May - On The Line'
-  - 'Meaxic - Take A Step'
+validation:
+  - 'Actions - One Minute Smile'
+  - 'Clara Berry And Wooldog - Waltz For My Victims'
+  - 'Johnny Lokke - Promises & Lies'
+  - 'Patrick Talbot - A Reason To Leave'
+  - 'Triviul - Angelsaint'
+  - 'Alexander Ross - Goodbye Bolero'
+  - 'Fergessen - Nos Palpitants'
+  - 'Leaf - Summerghost'
+  - 'Skelpolu - Human Mistakes'
+  - 'Young Griffo - Pennies'
+  - 'ANiMAL - Rockshow'
+  - 'James May - On The Line'
+  - 'Meaxic - Take A Step'
   - 'Traffic Experiment - Sirens'

mvsepless/models/bandit/core/loss/__init__.py

@@ -1,8 +1,8 @@
-from ._multistem import MultiStemWrapperFromConfig
-from ._timefreq import (
-    ReImL1Loss,
-    ReImL2Loss,
-    TimeFreqL1Loss,
-    TimeFreqL2Loss,
-    TimeFreqSignalNoisePNormRatioLoss,
-)
+from ._multistem import MultiStemWrapperFromConfig
+from ._timefreq import (
+    ReImL1Loss,
+    ReImL2Loss,
+    TimeFreqL1Loss,
+    TimeFreqL2Loss,
+    TimeFreqSignalNoisePNormRatioLoss,
+)

mvsepless/models/bandit/core/loss/_complex.py

@@ -1,27 +1,27 @@
-from typing import Any
-
-import torch
-from torch import nn
-from torch.nn.modules import loss as _loss
-from torch.nn.modules.loss import _Loss
-
-
-class ReImLossWrapper(_Loss):
-    def __init__(self, module: _Loss) -> None:
-        super().__init__()
-        self.module = module
-
-    def forward(self, preds: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
-        return self.module(torch.view_as_real(preds), torch.view_as_real(target))
-
-
-class ReImL1Loss(ReImLossWrapper):
-    def __init__(self, **kwargs: Any) -> None:
-        l1_loss = _loss.L1Loss(**kwargs)
-        super().__init__(module=(l1_loss))
-
-
-class ReImL2Loss(ReImLossWrapper):
-    def __init__(self, **kwargs: Any) -> None:
-        l2_loss = _loss.MSELoss(**kwargs)
-        super().__init__(module=(l2_loss))
+from typing import Any
+
+import torch
+from torch import nn
+from torch.nn.modules import loss as _loss
+from torch.nn.modules.loss import _Loss
+
+
+class ReImLossWrapper(_Loss):
+    def __init__(self, module: _Loss) -> None:
+        super().__init__()
+        self.module = module
+
+    def forward(self, preds: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+        return self.module(torch.view_as_real(preds), torch.view_as_real(target))
+
+
+class ReImL1Loss(ReImLossWrapper):
+    def __init__(self, **kwargs: Any) -> None:
+        l1_loss = _loss.L1Loss(**kwargs)
+        super().__init__(module=(l1_loss))
+
+
+class ReImL2Loss(ReImLossWrapper):
+    def __init__(self, **kwargs: Any) -> None:
+        l2_loss = _loss.MSELoss(**kwargs)
+        super().__init__(module=(l2_loss))

mvsepless/models/bandit/core/loss/_multistem.py

@@ -1,43 +1,43 @@
-from typing import Any, Dict
-
-import torch
-from asteroid import losses as asteroid_losses
-from torch import nn
-from torch.nn.modules.loss import _Loss
-
-from . import snr
-
-
-def parse_loss(name: str, kwargs: Dict[str, Any]) -> _Loss:
-
-    for module in [nn.modules.loss, snr, asteroid_losses, asteroid_losses.sdr]:
-        if name in module.__dict__:
-            return module.__dict__[name](**kwargs)
-
-    raise NameError
-
-
-class MultiStemWrapper(_Loss):
-    def __init__(self, module: _Loss, modality: str = "audio") -> None:
-        super().__init__()
-        self.loss = module
-        self.modality = modality
-
-    def forward(
-        self,
-        preds: Dict[str, Dict[str, torch.Tensor]],
-        target: Dict[str, Dict[str, torch.Tensor]],
-    ) -> torch.Tensor:
-        loss = {
-            stem: self.loss(preds[self.modality][stem], target[self.modality][stem])
-            for stem in preds[self.modality]
-            if stem in target[self.modality]
-        }
-
-        return sum(list(loss.values()))
-
-
-class MultiStemWrapperFromConfig(MultiStemWrapper):
-    def __init__(self, name: str, kwargs: Any, modality: str = "audio") -> None:
-        loss = parse_loss(name, kwargs)
-        super().__init__(module=loss, modality=modality)
+from typing import Any, Dict
+
+import torch
+from asteroid import losses as asteroid_losses
+from torch import nn
+from torch.nn.modules.loss import _Loss
+
+from . import snr
+
+
+def parse_loss(name: str, kwargs: Dict[str, Any]) -> _Loss:
+
+    for module in [nn.modules.loss, snr, asteroid_losses, asteroid_losses.sdr]:
+        if name in module.__dict__:
+            return module.__dict__[name](**kwargs)
+
+    raise NameError
+
+
+class MultiStemWrapper(_Loss):
+    def __init__(self, module: _Loss, modality: str = "audio") -> None:
+        super().__init__()
+        self.loss = module
+        self.modality = modality
+
+    def forward(
+        self,
+        preds: Dict[str, Dict[str, torch.Tensor]],
+        target: Dict[str, Dict[str, torch.Tensor]],
+    ) -> torch.Tensor:
+        loss = {
+            stem: self.loss(preds[self.modality][stem], target[self.modality][stem])
+            for stem in preds[self.modality]
+            if stem in target[self.modality]
+        }
+
+        return sum(list(loss.values()))
+
+
+class MultiStemWrapperFromConfig(MultiStemWrapper):
+    def __init__(self, name: str, kwargs: Any, modality: str = "audio") -> None:
+        loss = parse_loss(name, kwargs)
+        super().__init__(module=loss, modality=modality)

mvsepless/models/bandit/core/loss/_timefreq.py

@@ -1,94 +1,94 @@
-from typing import Any, Dict, Optional
-
-import torch
-from torch import nn
-from torch.nn.modules.loss import _Loss
-
-from ._multistem import MultiStemWrapper
-from ._complex import ReImL1Loss, ReImL2Loss, ReImLossWrapper
-from .snr import SignalNoisePNormRatio
-
-
-class TimeFreqWrapper(_Loss):
-    def __init__(
-        self,
-        time_module: _Loss,
-        freq_module: Optional[_Loss] = None,
-        time_weight: float = 1.0,
-        freq_weight: float = 1.0,
-        multistem: bool = True,
-    ) -> None:
-        super().__init__()
-
-        if freq_module is None:
-            freq_module = time_module
-
-        if multistem:
-            time_module = MultiStemWrapper(time_module, modality="audio")
-            freq_module = MultiStemWrapper(freq_module, modality="spectrogram")
-
-        self.time_module = time_module
-        self.freq_module = freq_module
-
-        self.time_weight = time_weight
-        self.freq_weight = freq_weight
-
-    def forward(self, preds: Any, target: Any) -> torch.Tensor:
-
-        return self.time_weight * self.time_module(
-            preds, target
-        ) + self.freq_weight * self.freq_module(preds, target)
-
-
-class TimeFreqL1Loss(TimeFreqWrapper):
-    def __init__(
-        self,
-        time_weight: float = 1.0,
-        freq_weight: float = 1.0,
-        tkwargs: Optional[Dict[str, Any]] = None,
-        fkwargs: Optional[Dict[str, Any]] = None,
-        multistem: bool = True,
-    ) -> None:
-        if tkwargs is None:
-            tkwargs = {}
-        if fkwargs is None:
-            fkwargs = {}
-        time_module = nn.L1Loss(**tkwargs)
-        freq_module = ReImL1Loss(**fkwargs)
-        super().__init__(time_module, freq_module, time_weight, freq_weight, multistem)
-
-
-class TimeFreqL2Loss(TimeFreqWrapper):
-    def __init__(
-        self,
-        time_weight: float = 1.0,
-        freq_weight: float = 1.0,
-        tkwargs: Optional[Dict[str, Any]] = None,
-        fkwargs: Optional[Dict[str, Any]] = None,
-        multistem: bool = True,
-    ) -> None:
-        if tkwargs is None:
-            tkwargs = {}
-        if fkwargs is None:
-            fkwargs = {}
-        time_module = nn.MSELoss(**tkwargs)
-        freq_module = ReImL2Loss(**fkwargs)
-        super().__init__(time_module, freq_module, time_weight, freq_weight, multistem)
-
-
-class TimeFreqSignalNoisePNormRatioLoss(TimeFreqWrapper):
-    def __init__(
-        self,
-        time_weight: float = 1.0,
-        freq_weight: float = 1.0,
-        tkwargs: Optional[Dict[str, Any]] = None,
-        fkwargs: Optional[Dict[str, Any]] = None,
-        multistem: bool = True,
-    ) -> None:
-        if tkwargs is None:
-            tkwargs = {}
-        if fkwargs is None:
-            fkwargs = {}
-        time_module = SignalNoisePNormRatio(**tkwargs)
-        freq_module = SignalNoisePNormRatio(**fkwargs)
-        super().__init__(time_module, freq_module, time_weight, freq_weight, multistem)
+from typing import Any, Dict, Optional
+
+import torch
+from torch import nn
+from torch.nn.modules.loss import _Loss
+
+from ._multistem import MultiStemWrapper
+from ._complex import ReImL1Loss, ReImL2Loss, ReImLossWrapper
+from .snr import SignalNoisePNormRatio
+
+
+class TimeFreqWrapper(_Loss):
+    def __init__(
+        self,
+        time_module: _Loss,
+        freq_module: Optional[_Loss] = None,
+        time_weight: float = 1.0,
+        freq_weight: float = 1.0,
+        multistem: bool = True,
+    ) -> None:
+        super().__init__()
+
+        if freq_module is None:
+            freq_module = time_module
+
+        if multistem:
+            time_module = MultiStemWrapper(time_module, modality="audio")
+            freq_module = MultiStemWrapper(freq_module, modality="spectrogram")
+
+        self.time_module = time_module
+        self.freq_module = freq_module
+
+        self.time_weight = time_weight
+        self.freq_weight = freq_weight
+
+    def forward(self, preds: Any, target: Any) -> torch.Tensor:
+
+        return self.time_weight * self.time_module(
+            preds, target
+        ) + self.freq_weight * self.freq_module(preds, target)
+
+
+class TimeFreqL1Loss(TimeFreqWrapper):
+    def __init__(
+        self,
+        time_weight: float = 1.0,
+        freq_weight: float = 1.0,
+        tkwargs: Optional[Dict[str, Any]] = None,
+        fkwargs: Optional[Dict[str, Any]] = None,
+        multistem: bool = True,
+    ) -> None:
+        if tkwargs is None:
+            tkwargs = {}
+        if fkwargs is None:
+            fkwargs = {}
+        time_module = nn.L1Loss(**tkwargs)
+        freq_module = ReImL1Loss(**fkwargs)
+        super().__init__(time_module, freq_module, time_weight, freq_weight, multistem)
+
+
+class TimeFreqL2Loss(TimeFreqWrapper):
+    def __init__(
+        self,
+        time_weight: float = 1.0,
+        freq_weight: float = 1.0,
+        tkwargs: Optional[Dict[str, Any]] = None,
+        fkwargs: Optional[Dict[str, Any]] = None,
+        multistem: bool = True,
+    ) -> None:
+        if tkwargs is None:
+            tkwargs = {}
+        if fkwargs is None:
+            fkwargs = {}
+        time_module = nn.MSELoss(**tkwargs)
+        freq_module = ReImL2Loss(**fkwargs)
+        super().__init__(time_module, freq_module, time_weight, freq_weight, multistem)
+
+
+class TimeFreqSignalNoisePNormRatioLoss(TimeFreqWrapper):
+    def __init__(
+        self,
+        time_weight: float = 1.0,
+        freq_weight: float = 1.0,
+        tkwargs: Optional[Dict[str, Any]] = None,
+        fkwargs: Optional[Dict[str, Any]] = None,
+        multistem: bool = True,
+    ) -> None:
+        if tkwargs is None:
+            tkwargs = {}
+        if fkwargs is None:
+            fkwargs = {}
+        time_module = SignalNoisePNormRatio(**tkwargs)
+        freq_module = SignalNoisePNormRatio(**fkwargs)
+        super().__init__(time_module, freq_module, time_weight, freq_weight, multistem)

mvsepless/models/bandit/core/loss/snr.py

@@ -1,131 +1,131 @@
-import torch
-from torch.nn.modules.loss import _Loss
-from torch.nn import functional as F
-
-
-class SignalNoisePNormRatio(_Loss):
-    def __init__(
-        self,
-        p: float = 1.0,
-        scale_invariant: bool = False,
-        zero_mean: bool = False,
-        take_log: bool = True,
-        reduction: str = "mean",
-        EPS: float = 1e-3,
-    ) -> None:
-        assert reduction != "sum", NotImplementedError
-        super().__init__(reduction=reduction)
-        assert not zero_mean
-
-        self.p = p
-
-        self.EPS = EPS
-        self.take_log = take_log
-
-        self.scale_invariant = scale_invariant
-
-    def forward(self, est_target: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
-
-        target_ = target
-        if self.scale_invariant:
-            ndim = target.ndim
-            dot = torch.sum(est_target * torch.conj(target), dim=-1, keepdim=True)
-            s_target_energy = torch.sum(
-                target * torch.conj(target), dim=-1, keepdim=True
-            )
-
-            if ndim > 2:
-                dot = torch.sum(dot, dim=list(range(1, ndim)), keepdim=True)
-                s_target_energy = torch.sum(
-                    s_target_energy, dim=list(range(1, ndim)), keepdim=True
-                )
-
-            target_scaler = (dot + 1e-8) / (s_target_energy + 1e-8)
-            target = target_ * target_scaler
-
-        if torch.is_complex(est_target):
-            est_target = torch.view_as_real(est_target)
-            target = torch.view_as_real(target)
-
-        batch_size = est_target.shape[0]
-        est_target = est_target.reshape(batch_size, -1)
-        target = target.reshape(batch_size, -1)
-
-        if self.p == 1:
-            e_error = torch.abs(est_target - target).mean(dim=-1)
-            e_target = torch.abs(target).mean(dim=-1)
-        elif self.p == 2:
-            e_error = torch.square(est_target - target).mean(dim=-1)
-            e_target = torch.square(target).mean(dim=-1)
-        else:
-            raise NotImplementedError
-
-        if self.take_log:
-            loss = 10 * (
-                torch.log10(e_error + self.EPS) - torch.log10(e_target + self.EPS)
-            )
-        else:
-            loss = (e_error + self.EPS) / (e_target + self.EPS)
-
-        if self.reduction == "mean":
-            loss = loss.mean()
-        elif self.reduction == "sum":
-            loss = loss.sum()
-
-        return loss
-
-
-class MultichannelSingleSrcNegSDR(_Loss):
-    def __init__(
-        self,
-        sdr_type: str,
-        p: float = 2.0,
-        zero_mean: bool = True,
-        take_log: bool = True,
-        reduction: str = "mean",
-        EPS: float = 1e-8,
-    ) -> None:
-        assert reduction != "sum", NotImplementedError
-        super().__init__(reduction=reduction)
-
-        assert sdr_type in ["snr", "sisdr", "sdsdr"]
-        self.sdr_type = sdr_type
-        self.zero_mean = zero_mean
-        self.take_log = take_log
-        self.EPS = 1e-8
-
-        self.p = p
-
-    def forward(self, est_target: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
-        if target.size() != est_target.size() or target.ndim != 3:
-            raise TypeError(
-                f"Inputs must be of shape [batch, time], got {target.size()} and {est_target.size()} instead"
-            )
-        if self.zero_mean:
-            mean_source = torch.mean(target, dim=[1, 2], keepdim=True)
-            mean_estimate = torch.mean(est_target, dim=[1, 2], keepdim=True)
-            target = target - mean_source
-            est_target = est_target - mean_estimate
-        if self.sdr_type in ["sisdr", "sdsdr"]:
-            dot = torch.sum(est_target * target, dim=[1, 2], keepdim=True)
-            s_target_energy = torch.sum(target**2, dim=[1, 2], keepdim=True) + self.EPS
-            scaled_target = dot * target / s_target_energy
-        else:
-            scaled_target = target
-        if self.sdr_type in ["sdsdr", "snr"]:
-            e_noise = est_target - target
-        else:
-            e_noise = est_target - scaled_target
-
-        if self.p == 2.0:
-            losses = torch.sum(scaled_target**2, dim=[1, 2]) / (
-                torch.sum(e_noise**2, dim=[1, 2]) + self.EPS
-            )
-        else:
-            losses = torch.norm(scaled_target, p=self.p, dim=[1, 2]) / (
-                torch.linalg.vector_norm(e_noise, p=self.p, dim=[1, 2]) + self.EPS
-            )
-        if self.take_log:
-            losses = 10 * torch.log10(losses + self.EPS)
-        losses = losses.mean() if self.reduction == "mean" else losses
-        return -losses
+import torch
+from torch.nn.modules.loss import _Loss
+from torch.nn import functional as F
+
+
+class SignalNoisePNormRatio(_Loss):
+    def __init__(
+        self,
+        p: float = 1.0,
+        scale_invariant: bool = False,
+        zero_mean: bool = False,
+        take_log: bool = True,
+        reduction: str = "mean",
+        EPS: float = 1e-3,
+    ) -> None:
+        assert reduction != "sum", NotImplementedError
+        super().__init__(reduction=reduction)
+        assert not zero_mean
+
+        self.p = p
+
+        self.EPS = EPS
+        self.take_log = take_log
+
+        self.scale_invariant = scale_invariant
+
+    def forward(self, est_target: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+
+        target_ = target
+        if self.scale_invariant:
+            ndim = target.ndim
+            dot = torch.sum(est_target * torch.conj(target), dim=-1, keepdim=True)
+            s_target_energy = torch.sum(
+                target * torch.conj(target), dim=-1, keepdim=True
+            )
+
+            if ndim > 2:
+                dot = torch.sum(dot, dim=list(range(1, ndim)), keepdim=True)
+                s_target_energy = torch.sum(
+                    s_target_energy, dim=list(range(1, ndim)), keepdim=True
+                )
+
+            target_scaler = (dot + 1e-8) / (s_target_energy + 1e-8)
+            target = target_ * target_scaler
+
+        if torch.is_complex(est_target):
+            est_target = torch.view_as_real(est_target)
+            target = torch.view_as_real(target)
+
+        batch_size = est_target.shape[0]
+        est_target = est_target.reshape(batch_size, -1)
+        target = target.reshape(batch_size, -1)
+
+        if self.p == 1:
+            e_error = torch.abs(est_target - target).mean(dim=-1)
+            e_target = torch.abs(target).mean(dim=-1)
+        elif self.p == 2:
+            e_error = torch.square(est_target - target).mean(dim=-1)
+            e_target = torch.square(target).mean(dim=-1)
+        else:
+            raise NotImplementedError
+
+        if self.take_log:
+            loss = 10 * (
+                torch.log10(e_error + self.EPS) - torch.log10(e_target + self.EPS)
+            )
+        else:
+            loss = (e_error + self.EPS) / (e_target + self.EPS)
+
+        if self.reduction == "mean":
+            loss = loss.mean()
+        elif self.reduction == "sum":
+            loss = loss.sum()
+
+        return loss
+
+
+class MultichannelSingleSrcNegSDR(_Loss):
+    def __init__(
+        self,
+        sdr_type: str,
+        p: float = 2.0,
+        zero_mean: bool = True,
+        take_log: bool = True,
+        reduction: str = "mean",
+        EPS: float = 1e-8,
+    ) -> None:
+        assert reduction != "sum", NotImplementedError
+        super().__init__(reduction=reduction)
+
+        assert sdr_type in ["snr", "sisdr", "sdsdr"]
+        self.sdr_type = sdr_type
+        self.zero_mean = zero_mean
+        self.take_log = take_log
+        self.EPS = 1e-8
+
+        self.p = p
+
+    def forward(self, est_target: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+        if target.size() != est_target.size() or target.ndim != 3:
+            raise TypeError(
+                f"Inputs must be of shape [batch, time], got {target.size()} and {est_target.size()} instead"
+            )
+        if self.zero_mean:
+            mean_source = torch.mean(target, dim=[1, 2], keepdim=True)
+            mean_estimate = torch.mean(est_target, dim=[1, 2], keepdim=True)
+            target = target - mean_source
+            est_target = est_target - mean_estimate
+        if self.sdr_type in ["sisdr", "sdsdr"]:
+            dot = torch.sum(est_target * target, dim=[1, 2], keepdim=True)
+            s_target_energy = torch.sum(target**2, dim=[1, 2], keepdim=True) + self.EPS
+            scaled_target = dot * target / s_target_energy
+        else:
+            scaled_target = target
+        if self.sdr_type in ["sdsdr", "snr"]:
+            e_noise = est_target - target
+        else:
+            e_noise = est_target - scaled_target
+
+        if self.p == 2.0:
+            losses = torch.sum(scaled_target**2, dim=[1, 2]) / (
+                torch.sum(e_noise**2, dim=[1, 2]) + self.EPS
+            )
+        else:
+            losses = torch.norm(scaled_target, p=self.p, dim=[1, 2]) / (
+                torch.linalg.vector_norm(e_noise, p=self.p, dim=[1, 2]) + self.EPS
+            )
+        if self.take_log:
+            losses = 10 * torch.log10(losses + self.EPS)
+        losses = losses.mean() if self.reduction == "mean" else losses
+        return -losses

mvsepless/models/bandit/core/metrics/__init__.py

@@ -1,7 +1,7 @@
-from .snr import (
-    ChunkMedianScaleInvariantSignalDistortionRatio,
-    ChunkMedianScaleInvariantSignalNoiseRatio,
-    ChunkMedianSignalDistortionRatio,
-    ChunkMedianSignalNoiseRatio,
-    SafeSignalDistortionRatio,
-)
+from .snr import (
+    ChunkMedianScaleInvariantSignalDistortionRatio,
+    ChunkMedianScaleInvariantSignalNoiseRatio,
+    ChunkMedianSignalDistortionRatio,
+    ChunkMedianSignalNoiseRatio,
+    SafeSignalDistortionRatio,
+)

mvsepless/models/bandit/core/metrics/_squim.py

@@ -1,350 +1,350 @@
-from dataclasses import dataclass
-
-from torchaudio._internal import load_state_dict_from_url
-
-import math
-from typing import List, Optional, Tuple
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-
-def transform_wb_pesq_range(x: float) -> float:
-    return 0.999 + (4.999 - 0.999) / (1 + math.exp(-1.3669 * x + 3.8224))
-
-
-PESQRange: Tuple[float, float] = (
-    1.0,
-    transform_wb_pesq_range(4.5),
-)
-
-
-class RangeSigmoid(nn.Module):
-    def __init__(self, val_range: Tuple[float, float] = (0.0, 1.0)) -> None:
-        super(RangeSigmoid, self).__init__()
-        assert isinstance(val_range, tuple) and len(val_range) == 2
-        self.val_range: Tuple[float, float] = val_range
-        self.sigmoid: nn.modules.Module = nn.Sigmoid()
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        out = (
-            self.sigmoid(x) * (self.val_range[1] - self.val_range[0])
-            + self.val_range[0]
-        )
-        return out
-
-
-class Encoder(nn.Module):
-
-    def __init__(self, feat_dim: int = 512, win_len: int = 32) -> None:
-        super(Encoder, self).__init__()
-
-        self.conv1d = nn.Conv1d(1, feat_dim, win_len, stride=win_len // 2, bias=False)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        out = x.unsqueeze(dim=1)
-        out = F.relu(self.conv1d(out))
-        return out
-
-
-class SingleRNN(nn.Module):
-    def __init__(
-        self, rnn_type: str, input_size: int, hidden_size: int, dropout: float = 0.0
-    ) -> None:
-        super(SingleRNN, self).__init__()
-
-        self.rnn_type = rnn_type
-        self.input_size = input_size
-        self.hidden_size = hidden_size
-
-        self.rnn: nn.modules.Module = getattr(nn, rnn_type)(
-            input_size,
-            hidden_size,
-            1,
-            dropout=dropout,
-            batch_first=True,
-            bidirectional=True,
-        )
-
-        self.proj = nn.Linear(hidden_size * 2, input_size)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        out, _ = self.rnn(x)
-        out = self.proj(out)
-        return out
-
-
-class DPRNN(nn.Module):
-
-    def __init__(
-        self,
-        feat_dim: int = 64,
-        hidden_dim: int = 128,
-        num_blocks: int = 6,
-        rnn_type: str = "LSTM",
-        d_model: int = 256,
-        chunk_size: int = 100,
-        chunk_stride: int = 50,
-    ) -> None:
-        super(DPRNN, self).__init__()
-
-        self.num_blocks = num_blocks
-
-        self.row_rnn = nn.ModuleList([])
-        self.col_rnn = nn.ModuleList([])
-        self.row_norm = nn.ModuleList([])
-        self.col_norm = nn.ModuleList([])
-        for _ in range(num_blocks):
-            self.row_rnn.append(SingleRNN(rnn_type, feat_dim, hidden_dim))
-            self.col_rnn.append(SingleRNN(rnn_type, feat_dim, hidden_dim))
-            self.row_norm.append(nn.GroupNorm(1, feat_dim, eps=1e-8))
-            self.col_norm.append(nn.GroupNorm(1, feat_dim, eps=1e-8))
-        self.conv = nn.Sequential(
-            nn.Conv2d(feat_dim, d_model, 1),
-            nn.PReLU(),
-        )
-        self.chunk_size = chunk_size
-        self.chunk_stride = chunk_stride
-
-    def pad_chunk(self, x: torch.Tensor) -> Tuple[torch.Tensor, int]:
-        seq_len = x.shape[-1]
-
-        rest = (
-            self.chunk_size
-            - (self.chunk_stride + seq_len % self.chunk_size) % self.chunk_size
-        )
-        out = F.pad(x, [self.chunk_stride, rest + self.chunk_stride])
-
-        return out, rest
-
-    def chunking(self, x: torch.Tensor) -> Tuple[torch.Tensor, int]:
-        out, rest = self.pad_chunk(x)
-        batch_size, feat_dim, seq_len = out.shape
-
-        segments1 = (
-            out[:, :, : -self.chunk_stride]
-            .contiguous()
-            .view(batch_size, feat_dim, -1, self.chunk_size)
-        )
-        segments2 = (
-            out[:, :, self.chunk_stride :]
-            .contiguous()
-            .view(batch_size, feat_dim, -1, self.chunk_size)
-        )
-        out = torch.cat([segments1, segments2], dim=3)
-        out = (
-            out.view(batch_size, feat_dim, -1, self.chunk_size)
-            .transpose(2, 3)
-            .contiguous()
-        )
-
-        return out, rest
-
-    def merging(self, x: torch.Tensor, rest: int) -> torch.Tensor:
-        batch_size, dim, _, _ = x.shape
-        out = (
-            x.transpose(2, 3)
-            .contiguous()
-            .view(batch_size, dim, -1, self.chunk_size * 2)
-        )
-        out1 = (
-            out[:, :, :, : self.chunk_size]
-            .contiguous()
-            .view(batch_size, dim, -1)[:, :, self.chunk_stride :]
-        )
-        out2 = (
-            out[:, :, :, self.chunk_size :]
-            .contiguous()
-            .view(batch_size, dim, -1)[:, :, : -self.chunk_stride]
-        )
-        out = out1 + out2
-        if rest > 0:
-            out = out[:, :, :-rest]
-        out = out.contiguous()
-        return out
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x, rest = self.chunking(x)
-        batch_size, _, dim1, dim2 = x.shape
-        out = x
-        for row_rnn, row_norm, col_rnn, col_norm in zip(
-            self.row_rnn, self.row_norm, self.col_rnn, self.col_norm
-        ):
-            row_in = (
-                out.permute(0, 3, 2, 1)
-                .contiguous()
-                .view(batch_size * dim2, dim1, -1)
-                .contiguous()
-            )
-            row_out = row_rnn(row_in)
-            row_out = (
-                row_out.view(batch_size, dim2, dim1, -1)
-                .permute(0, 3, 2, 1)
-                .contiguous()
-            )
-            row_out = row_norm(row_out)
-            out = out + row_out
-
-            col_in = (
-                out.permute(0, 2, 3, 1)
-                .contiguous()
-                .view(batch_size * dim1, dim2, -1)
-                .contiguous()
-            )
-            col_out = col_rnn(col_in)
-            col_out = (
-                col_out.view(batch_size, dim1, dim2, -1)
-                .permute(0, 3, 1, 2)
-                .contiguous()
-            )
-            col_out = col_norm(col_out)
-            out = out + col_out
-        out = self.conv(out)
-        out = self.merging(out, rest)
-        out = out.transpose(1, 2).contiguous()
-        return out
-
-
-class AutoPool(nn.Module):
-    def __init__(self, pool_dim: int = 1) -> None:
-        super(AutoPool, self).__init__()
-        self.pool_dim: int = pool_dim
-        self.softmax: nn.modules.Module = nn.Softmax(dim=pool_dim)
-        self.register_parameter("alpha", nn.Parameter(torch.ones(1)))
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        weight = self.softmax(torch.mul(x, self.alpha))
-        out = torch.sum(torch.mul(x, weight), dim=self.pool_dim)
-        return out
-
-
-class SquimObjective(nn.Module):
-
-    def __init__(
-        self,
-        encoder: nn.Module,
-        dprnn: nn.Module,
-        branches: nn.ModuleList,
-    ):
-        super(SquimObjective, self).__init__()
-        self.encoder = encoder
-        self.dprnn = dprnn
-        self.branches = branches
-
-    def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
-        if x.ndim != 2:
-            raise ValueError(
-                f"The input must be a 2D Tensor. Found dimension {x.ndim}."
-            )
-        x = x / (torch.mean(x**2, dim=1, keepdim=True) ** 0.5 * 20)
-        out = self.encoder(x)
-        out = self.dprnn(out)
-        scores = []
-        for branch in self.branches:
-            scores.append(branch(out).squeeze(dim=1))
-        return scores
-
-
-def _create_branch(d_model: int, nhead: int, metric: str) -> nn.modules.Module:
-    layer1 = nn.TransformerEncoderLayer(
-        d_model, nhead, d_model * 4, dropout=0.0, batch_first=True
-    )
-    layer2 = AutoPool()
-    if metric == "stoi":
-        layer3 = nn.Sequential(
-            nn.Linear(d_model, d_model),
-            nn.PReLU(),
-            nn.Linear(d_model, 1),
-            RangeSigmoid(),
-        )
-    elif metric == "pesq":
-        layer3 = nn.Sequential(
-            nn.Linear(d_model, d_model),
-            nn.PReLU(),
-            nn.Linear(d_model, 1),
-            RangeSigmoid(val_range=PESQRange),
-        )
-    else:
-        layer3: nn.modules.Module = nn.Sequential(
-            nn.Linear(d_model, d_model), nn.PReLU(), nn.Linear(d_model, 1)
-        )
-    return nn.Sequential(layer1, layer2, layer3)
-
-
-def squim_objective_model(
-    feat_dim: int,
-    win_len: int,
-    d_model: int,
-    nhead: int,
-    hidden_dim: int,
-    num_blocks: int,
-    rnn_type: str,
-    chunk_size: int,
-    chunk_stride: Optional[int] = None,
-) -> SquimObjective:
-    if chunk_stride is None:
-        chunk_stride = chunk_size // 2
-    encoder = Encoder(feat_dim, win_len)
-    dprnn = DPRNN(
-        feat_dim, hidden_dim, num_blocks, rnn_type, d_model, chunk_size, chunk_stride
-    )
-    branches = nn.ModuleList(
-        [
-            _create_branch(d_model, nhead, "stoi"),
-            _create_branch(d_model, nhead, "pesq"),
-            _create_branch(d_model, nhead, "sisdr"),
-        ]
-    )
-    return SquimObjective(encoder, dprnn, branches)
-
-
-def squim_objective_base() -> SquimObjective:
-    return squim_objective_model(
-        feat_dim=256,
-        win_len=64,
-        d_model=256,
-        nhead=4,
-        hidden_dim=256,
-        num_blocks=2,
-        rnn_type="LSTM",
-        chunk_size=71,
-    )
-
-
-@dataclass
-class SquimObjectiveBundle:
-
-    _path: str
-    _sample_rate: float
-
-    def _get_state_dict(self, dl_kwargs):
-        url = f"https://download.pytorch.org/torchaudio/models/{self._path}"
-        dl_kwargs = {} if dl_kwargs is None else dl_kwargs
-        state_dict = load_state_dict_from_url(url, **dl_kwargs)
-        return state_dict
-
-    def get_model(self, *, dl_kwargs=None) -> SquimObjective:
-        model = squim_objective_base()
-        model.load_state_dict(self._get_state_dict(dl_kwargs))
-        model.eval()
-        return model
-
-    @property
-    def sample_rate(self):
-        return self._sample_rate
-
-
-SQUIM_OBJECTIVE = SquimObjectiveBundle(
-    "squim_objective_dns2020.pth",
-    _sample_rate=16000,
-)
-SQUIM_OBJECTIVE.__doc__ = """SquimObjective pipeline trained using approach described in
-    :cite:`kumar2023torchaudio` on the *DNS 2020 Dataset* :cite:`reddy2020interspeech`.
-
-    The underlying model is constructed by :py:func:`torchaudio.models.squim_objective_base`.
-    The weights are under `Creative Commons Attribution 4.0 International License
-    <https://github.com/microsoft/DNS-Challenge/blob/interspeech2020/master/LICENSE>`__.
-
-    Please refer to :py:class:`SquimObjectiveBundle` for usage instructions.
-    """
+from dataclasses import dataclass
+
+from torchaudio._internal import load_state_dict_from_url
+
+import math
+from typing import List, Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def transform_wb_pesq_range(x: float) -> float:
+    return 0.999 + (4.999 - 0.999) / (1 + math.exp(-1.3669 * x + 3.8224))
+
+
+PESQRange: Tuple[float, float] = (
+    1.0,
+    transform_wb_pesq_range(4.5),
+)
+
+
+class RangeSigmoid(nn.Module):
+    def __init__(self, val_range: Tuple[float, float] = (0.0, 1.0)) -> None:
+        super(RangeSigmoid, self).__init__()
+        assert isinstance(val_range, tuple) and len(val_range) == 2
+        self.val_range: Tuple[float, float] = val_range
+        self.sigmoid: nn.modules.Module = nn.Sigmoid()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        out = (
+            self.sigmoid(x) * (self.val_range[1] - self.val_range[0])
+            + self.val_range[0]
+        )
+        return out
+
+
+class Encoder(nn.Module):
+
+    def __init__(self, feat_dim: int = 512, win_len: int = 32) -> None:
+        super(Encoder, self).__init__()
+
+        self.conv1d = nn.Conv1d(1, feat_dim, win_len, stride=win_len // 2, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        out = x.unsqueeze(dim=1)
+        out = F.relu(self.conv1d(out))
+        return out
+
+
+class SingleRNN(nn.Module):
+    def __init__(
+        self, rnn_type: str, input_size: int, hidden_size: int, dropout: float = 0.0
+    ) -> None:
+        super(SingleRNN, self).__init__()
+
+        self.rnn_type = rnn_type
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+
+        self.rnn: nn.modules.Module = getattr(nn, rnn_type)(
+            input_size,
+            hidden_size,
+            1,
+            dropout=dropout,
+            batch_first=True,
+            bidirectional=True,
+        )
+
+        self.proj = nn.Linear(hidden_size * 2, input_size)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        out, _ = self.rnn(x)
+        out = self.proj(out)
+        return out
+
+
+class DPRNN(nn.Module):
+
+    def __init__(
+        self,
+        feat_dim: int = 64,
+        hidden_dim: int = 128,
+        num_blocks: int = 6,
+        rnn_type: str = "LSTM",
+        d_model: int = 256,
+        chunk_size: int = 100,
+        chunk_stride: int = 50,
+    ) -> None:
+        super(DPRNN, self).__init__()
+
+        self.num_blocks = num_blocks
+
+        self.row_rnn = nn.ModuleList([])
+        self.col_rnn = nn.ModuleList([])
+        self.row_norm = nn.ModuleList([])
+        self.col_norm = nn.ModuleList([])
+        for _ in range(num_blocks):
+            self.row_rnn.append(SingleRNN(rnn_type, feat_dim, hidden_dim))
+            self.col_rnn.append(SingleRNN(rnn_type, feat_dim, hidden_dim))
+            self.row_norm.append(nn.GroupNorm(1, feat_dim, eps=1e-8))
+            self.col_norm.append(nn.GroupNorm(1, feat_dim, eps=1e-8))
+        self.conv = nn.Sequential(
+            nn.Conv2d(feat_dim, d_model, 1),
+            nn.PReLU(),
+        )
+        self.chunk_size = chunk_size
+        self.chunk_stride = chunk_stride
+
+    def pad_chunk(self, x: torch.Tensor) -> Tuple[torch.Tensor, int]:
+        seq_len = x.shape[-1]
+
+        rest = (
+            self.chunk_size
+            - (self.chunk_stride + seq_len % self.chunk_size) % self.chunk_size
+        )
+        out = F.pad(x, [self.chunk_stride, rest + self.chunk_stride])
+
+        return out, rest
+
+    def chunking(self, x: torch.Tensor) -> Tuple[torch.Tensor, int]:
+        out, rest = self.pad_chunk(x)
+        batch_size, feat_dim, seq_len = out.shape
+
+        segments1 = (
+            out[:, :, : -self.chunk_stride]
+            .contiguous()
+            .view(batch_size, feat_dim, -1, self.chunk_size)
+        )
+        segments2 = (
+            out[:, :, self.chunk_stride :]
+            .contiguous()
+            .view(batch_size, feat_dim, -1, self.chunk_size)
+        )
+        out = torch.cat([segments1, segments2], dim=3)
+        out = (
+            out.view(batch_size, feat_dim, -1, self.chunk_size)
+            .transpose(2, 3)
+            .contiguous()
+        )
+
+        return out, rest
+
+    def merging(self, x: torch.Tensor, rest: int) -> torch.Tensor:
+        batch_size, dim, _, _ = x.shape
+        out = (
+            x.transpose(2, 3)
+            .contiguous()
+            .view(batch_size, dim, -1, self.chunk_size * 2)
+        )
+        out1 = (
+            out[:, :, :, : self.chunk_size]
+            .contiguous()
+            .view(batch_size, dim, -1)[:, :, self.chunk_stride :]
+        )
+        out2 = (
+            out[:, :, :, self.chunk_size :]
+            .contiguous()
+            .view(batch_size, dim, -1)[:, :, : -self.chunk_stride]
+        )
+        out = out1 + out2
+        if rest > 0:
+            out = out[:, :, :-rest]
+        out = out.contiguous()
+        return out
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x, rest = self.chunking(x)
+        batch_size, _, dim1, dim2 = x.shape
+        out = x
+        for row_rnn, row_norm, col_rnn, col_norm in zip(
+            self.row_rnn, self.row_norm, self.col_rnn, self.col_norm
+        ):
+            row_in = (
+                out.permute(0, 3, 2, 1)
+                .contiguous()
+                .view(batch_size * dim2, dim1, -1)
+                .contiguous()
+            )
+            row_out = row_rnn(row_in)
+            row_out = (
+                row_out.view(batch_size, dim2, dim1, -1)
+                .permute(0, 3, 2, 1)
+                .contiguous()
+            )
+            row_out = row_norm(row_out)
+            out = out + row_out
+
+            col_in = (
+                out.permute(0, 2, 3, 1)
+                .contiguous()
+                .view(batch_size * dim1, dim2, -1)
+                .contiguous()
+            )
+            col_out = col_rnn(col_in)
+            col_out = (
+                col_out.view(batch_size, dim1, dim2, -1)
+                .permute(0, 3, 1, 2)
+                .contiguous()
+            )
+            col_out = col_norm(col_out)
+            out = out + col_out
+        out = self.conv(out)
+        out = self.merging(out, rest)
+        out = out.transpose(1, 2).contiguous()
+        return out
+
+
+class AutoPool(nn.Module):
+    def __init__(self, pool_dim: int = 1) -> None:
+        super(AutoPool, self).__init__()
+        self.pool_dim: int = pool_dim
+        self.softmax: nn.modules.Module = nn.Softmax(dim=pool_dim)
+        self.register_parameter("alpha", nn.Parameter(torch.ones(1)))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        weight = self.softmax(torch.mul(x, self.alpha))
+        out = torch.sum(torch.mul(x, weight), dim=self.pool_dim)
+        return out
+
+
+class SquimObjective(nn.Module):
+
+    def __init__(
+        self,
+        encoder: nn.Module,
+        dprnn: nn.Module,
+        branches: nn.ModuleList,
+    ):
+        super(SquimObjective, self).__init__()
+        self.encoder = encoder
+        self.dprnn = dprnn
+        self.branches = branches
+
+    def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
+        if x.ndim != 2:
+            raise ValueError(
+                f"The input must be a 2D Tensor. Found dimension {x.ndim}."
+            )
+        x = x / (torch.mean(x**2, dim=1, keepdim=True) ** 0.5 * 20)
+        out = self.encoder(x)
+        out = self.dprnn(out)
+        scores = []
+        for branch in self.branches:
+            scores.append(branch(out).squeeze(dim=1))
+        return scores
+
+
+def _create_branch(d_model: int, nhead: int, metric: str) -> nn.modules.Module:
+    layer1 = nn.TransformerEncoderLayer(
+        d_model, nhead, d_model * 4, dropout=0.0, batch_first=True
+    )
+    layer2 = AutoPool()
+    if metric == "stoi":
+        layer3 = nn.Sequential(
+            nn.Linear(d_model, d_model),
+            nn.PReLU(),
+            nn.Linear(d_model, 1),
+            RangeSigmoid(),
+        )
+    elif metric == "pesq":
+        layer3 = nn.Sequential(
+            nn.Linear(d_model, d_model),
+            nn.PReLU(),
+            nn.Linear(d_model, 1),
+            RangeSigmoid(val_range=PESQRange),
+        )
+    else:
+        layer3: nn.modules.Module = nn.Sequential(
+            nn.Linear(d_model, d_model), nn.PReLU(), nn.Linear(d_model, 1)
+        )
+    return nn.Sequential(layer1, layer2, layer3)
+
+
+def squim_objective_model(
+    feat_dim: int,
+    win_len: int,
+    d_model: int,
+    nhead: int,
+    hidden_dim: int,
+    num_blocks: int,
+    rnn_type: str,
+    chunk_size: int,
+    chunk_stride: Optional[int] = None,
+) -> SquimObjective:
+    if chunk_stride is None:
+        chunk_stride = chunk_size // 2
+    encoder = Encoder(feat_dim, win_len)
+    dprnn = DPRNN(
+        feat_dim, hidden_dim, num_blocks, rnn_type, d_model, chunk_size, chunk_stride
+    )
+    branches = nn.ModuleList(
+        [
+            _create_branch(d_model, nhead, "stoi"),
+            _create_branch(d_model, nhead, "pesq"),
+            _create_branch(d_model, nhead, "sisdr"),
+        ]
+    )
+    return SquimObjective(encoder, dprnn, branches)
+
+
+def squim_objective_base() -> SquimObjective:
+    return squim_objective_model(
+        feat_dim=256,
+        win_len=64,
+        d_model=256,
+        nhead=4,
+        hidden_dim=256,
+        num_blocks=2,
+        rnn_type="LSTM",
+        chunk_size=71,
+    )
+
+
+@dataclass
+class SquimObjectiveBundle:
+
+    _path: str
+    _sample_rate: float
+
+    def _get_state_dict(self, dl_kwargs):
+        url = f"https://download.pytorch.org/torchaudio/models/{self._path}"
+        dl_kwargs = {} if dl_kwargs is None else dl_kwargs
+        state_dict = load_state_dict_from_url(url, **dl_kwargs)
+        return state_dict
+
+    def get_model(self, *, dl_kwargs=None) -> SquimObjective:
+        model = squim_objective_base()
+        model.load_state_dict(self._get_state_dict(dl_kwargs))
+        model.eval()
+        return model
+
+    @property
+    def sample_rate(self):
+        return self._sample_rate
+
+
+SQUIM_OBJECTIVE = SquimObjectiveBundle(
+    "squim_objective_dns2020.pth",
+    _sample_rate=16000,
+)
+SQUIM_OBJECTIVE.__doc__ = """SquimObjective pipeline trained using approach described in
+    :cite:`kumar2023torchaudio` on the *DNS 2020 Dataset* :cite:`reddy2020interspeech`.
+
+    The underlying model is constructed by :py:func:`torchaudio.models.squim_objective_base`.
+    The weights are under `Creative Commons Attribution 4.0 International License
+    <https://github.com/microsoft/DNS-Challenge/blob/interspeech2020/master/LICENSE>`__.
+
+    Please refer to :py:class:`SquimObjectiveBundle` for usage instructions.
+    """

mvsepless/models/bandit/core/metrics/snr.py

@@ -1,124 +1,124 @@
-from typing import Any, Callable
-
-import numpy as np
-import torch
-import torchmetrics as tm
-from torch._C import _LinAlgError
-from torchmetrics import functional as tmF
-
-
-class SafeSignalDistortionRatio(tm.SignalDistortionRatio):
-    def __init__(self, **kwargs) -> None:
-        super().__init__(**kwargs)
-
-    def update(self, *args, **kwargs) -> Any:
-        try:
-            super().update(*args, **kwargs)
-        except:
-            pass
-
-    def compute(self) -> Any:
-        if self.total == 0:
-            return torch.tensor(torch.nan)
-        return super().compute()
-
-
-class BaseChunkMedianSignalRatio(tm.Metric):
-    def __init__(
-        self,
-        func: Callable,
-        window_size: int,
-        hop_size: int = None,
-        zero_mean: bool = False,
-    ) -> None:
-        super().__init__()
-
-        self.func = func
-        self.window_size = window_size
-        if hop_size is None:
-            hop_size = window_size
-        self.hop_size = hop_size
-
-        self.add_state("sum_snr", default=torch.tensor(0.0), dist_reduce_fx="sum")
-        self.add_state("total", default=torch.tensor(0), dist_reduce_fx="sum")
-
-    def update(self, preds: torch.Tensor, target: torch.Tensor) -> None:
-
-        n_samples = target.shape[-1]
-
-        n_chunks = int(np.ceil((n_samples - self.window_size) / self.hop_size) + 1)
-
-        snr_chunk = []
-
-        for i in range(n_chunks):
-            start = i * self.hop_size
-
-            if n_samples - start < self.window_size:
-                continue
-
-            end = start + self.window_size
-
-            try:
-                chunk_snr = self.func(preds[..., start:end], target[..., start:end])
-
-                if torch.all(torch.isfinite(chunk_snr)):
-                    snr_chunk.append(chunk_snr)
-            except _LinAlgError:
-                pass
-
-        snr_chunk = torch.stack(snr_chunk, dim=-1)
-        snr_batch, _ = torch.nanmedian(snr_chunk, dim=-1)
-
-        self.sum_snr += snr_batch.sum()
-        self.total += snr_batch.numel()
-
-    def compute(self) -> Any:
-        return self.sum_snr / self.total
-
-
-class ChunkMedianSignalNoiseRatio(BaseChunkMedianSignalRatio):
-    def __init__(
-        self, window_size: int, hop_size: int = None, zero_mean: bool = False
-    ) -> None:
-        super().__init__(
-            func=tmF.signal_noise_ratio,
-            window_size=window_size,
-            hop_size=hop_size,
-            zero_mean=zero_mean,
-        )
-
-
-class ChunkMedianScaleInvariantSignalNoiseRatio(BaseChunkMedianSignalRatio):
-    def __init__(
-        self, window_size: int, hop_size: int = None, zero_mean: bool = False
-    ) -> None:
-        super().__init__(
-            func=tmF.scale_invariant_signal_noise_ratio,
-            window_size=window_size,
-            hop_size=hop_size,
-            zero_mean=zero_mean,
-        )
-
-
-class ChunkMedianSignalDistortionRatio(BaseChunkMedianSignalRatio):
-    def __init__(
-        self, window_size: int, hop_size: int = None, zero_mean: bool = False
-    ) -> None:
-        super().__init__(
-            func=tmF.signal_distortion_ratio,
-            window_size=window_size,
-            hop_size=hop_size,
-            zero_mean=zero_mean,
-        )
-
-
-class ChunkMedianScaleInvariantSignalDistortionRatio(BaseChunkMedianSignalRatio):
-    def __init__(
-        self, window_size: int, hop_size: int = None, zero_mean: bool = False
-    ) -> None:
-        super().__init__(
-            func=tmF.scale_invariant_signal_distortion_ratio,
-            window_size=window_size,
-            hop_size=hop_size,
-            zero_mean=zero_mean,
-        )
+from typing import Any, Callable
+
+import numpy as np
+import torch
+import torchmetrics as tm
+from torch._C import _LinAlgError
+from torchmetrics import functional as tmF
+
+
+class SafeSignalDistortionRatio(tm.SignalDistortionRatio):
+    def __init__(self, **kwargs) -> None:
+        super().__init__(**kwargs)
+
+    def update(self, *args, **kwargs) -> Any:
+        try:
+            super().update(*args, **kwargs)
+        except:
+            pass
+
+    def compute(self) -> Any:
+        if self.total == 0:
+            return torch.tensor(torch.nan)
+        return super().compute()
+
+
+class BaseChunkMedianSignalRatio(tm.Metric):
+    def __init__(
+        self,
+        func: Callable,
+        window_size: int,
+        hop_size: int = None,
+        zero_mean: bool = False,
+    ) -> None:
+        super().__init__()
+
+        self.func = func
+        self.window_size = window_size
+        if hop_size is None:
+            hop_size = window_size
+        self.hop_size = hop_size
+
+        self.add_state("sum_snr", default=torch.tensor(0.0), dist_reduce_fx="sum")
+        self.add_state("total", default=torch.tensor(0), dist_reduce_fx="sum")
+
+    def update(self, preds: torch.Tensor, target: torch.Tensor) -> None:
+
+        n_samples = target.shape[-1]
+
+        n_chunks = int(np.ceil((n_samples - self.window_size) / self.hop_size) + 1)
+
+        snr_chunk = []
+
+        for i in range(n_chunks):
+            start = i * self.hop_size
+
+            if n_samples - start < self.window_size:
+                continue
+
+            end = start + self.window_size
+
+            try:
+                chunk_snr = self.func(preds[..., start:end], target[..., start:end])
+
+                if torch.all(torch.isfinite(chunk_snr)):
+                    snr_chunk.append(chunk_snr)
+            except _LinAlgError:
+                pass
+
+        snr_chunk = torch.stack(snr_chunk, dim=-1)
+        snr_batch, _ = torch.nanmedian(snr_chunk, dim=-1)
+
+        self.sum_snr += snr_batch.sum()
+        self.total += snr_batch.numel()
+
+    def compute(self) -> Any:
+        return self.sum_snr / self.total
+
+
+class ChunkMedianSignalNoiseRatio(BaseChunkMedianSignalRatio):
+    def __init__(
+        self, window_size: int, hop_size: int = None, zero_mean: bool = False
+    ) -> None:
+        super().__init__(
+            func=tmF.signal_noise_ratio,
+            window_size=window_size,
+            hop_size=hop_size,
+            zero_mean=zero_mean,
+        )
+
+
+class ChunkMedianScaleInvariantSignalNoiseRatio(BaseChunkMedianSignalRatio):
+    def __init__(
+        self, window_size: int, hop_size: int = None, zero_mean: bool = False
+    ) -> None:
+        super().__init__(
+            func=tmF.scale_invariant_signal_noise_ratio,
+            window_size=window_size,
+            hop_size=hop_size,
+            zero_mean=zero_mean,
+        )
+
+
+class ChunkMedianSignalDistortionRatio(BaseChunkMedianSignalRatio):
+    def __init__(
+        self, window_size: int, hop_size: int = None, zero_mean: bool = False
+    ) -> None:
+        super().__init__(
+            func=tmF.signal_distortion_ratio,
+            window_size=window_size,
+            hop_size=hop_size,
+            zero_mean=zero_mean,
+        )
+
+
+class ChunkMedianScaleInvariantSignalDistortionRatio(BaseChunkMedianSignalRatio):
+    def __init__(
+        self, window_size: int, hop_size: int = None, zero_mean: bool = False
+    ) -> None:
+        super().__init__(
+            func=tmF.scale_invariant_signal_distortion_ratio,
+            window_size=window_size,
+            hop_size=hop_size,
+            zero_mean=zero_mean,
+        )

mvsepless/models/bandit/core/model/__init__.py

@@ -1,3 +1,3 @@
-from .bsrnn.wrapper import (
-    MultiMaskMultiSourceBandSplitRNNSimple,
-)
+from .bsrnn.wrapper import (
+    MultiMaskMultiSourceBandSplitRNNSimple,
+)

mvsepless/models/bandit/core/model/_spectral.py

@@ -1,54 +1,54 @@
-from typing import Dict, Optional
-
-import torch
-import torchaudio as ta
-from torch import nn
-
-
-class _SpectralComponent(nn.Module):
-    def __init__(
-        self,
-        n_fft: int = 2048,
-        win_length: Optional[int] = 2048,
-        hop_length: int = 512,
-        window_fn: str = "hann_window",
-        wkwargs: Optional[Dict] = None,
-        power: Optional[int] = None,
-        center: bool = True,
-        normalized: bool = True,
-        pad_mode: str = "constant",
-        onesided: bool = True,
-        **kwargs,
-    ) -> None:
-        super().__init__()
-
-        assert power is None
-
-        window_fn = torch.__dict__[window_fn]
-
-        self.stft = ta.transforms.Spectrogram(
-            n_fft=n_fft,
-            win_length=win_length,
-            hop_length=hop_length,
-            pad_mode=pad_mode,
-            pad=0,
-            window_fn=window_fn,
-            wkwargs=wkwargs,
-            power=power,
-            normalized=normalized,
-            center=center,
-            onesided=onesided,
-        )
-
-        self.istft = ta.transforms.InverseSpectrogram(
-            n_fft=n_fft,
-            win_length=win_length,
-            hop_length=hop_length,
-            pad_mode=pad_mode,
-            pad=0,
-            window_fn=window_fn,
-            wkwargs=wkwargs,
-            normalized=normalized,
-            center=center,
-            onesided=onesided,
-        )
+from typing import Dict, Optional
+
+import torch
+import torchaudio as ta
+from torch import nn
+
+
+class _SpectralComponent(nn.Module):
+    def __init__(
+        self,
+        n_fft: int = 2048,
+        win_length: Optional[int] = 2048,
+        hop_length: int = 512,
+        window_fn: str = "hann_window",
+        wkwargs: Optional[Dict] = None,
+        power: Optional[int] = None,
+        center: bool = True,
+        normalized: bool = True,
+        pad_mode: str = "constant",
+        onesided: bool = True,
+        **kwargs,
+    ) -> None:
+        super().__init__()
+
+        assert power is None
+
+        window_fn = torch.__dict__[window_fn]
+
+        self.stft = ta.transforms.Spectrogram(
+            n_fft=n_fft,
+            win_length=win_length,
+            hop_length=hop_length,
+            pad_mode=pad_mode,
+            pad=0,
+            window_fn=window_fn,
+            wkwargs=wkwargs,
+            power=power,
+            normalized=normalized,
+            center=center,
+            onesided=onesided,
+        )
+
+        self.istft = ta.transforms.InverseSpectrogram(
+            n_fft=n_fft,
+            win_length=win_length,
+            hop_length=hop_length,
+            pad_mode=pad_mode,
+            pad=0,
+            window_fn=window_fn,
+            wkwargs=wkwargs,
+            normalized=normalized,
+            center=center,
+            onesided=onesided,
+        )

mvsepless/models/bandit/core/model/bsrnn/__init__.py

@@ -1,23 +1,23 @@
-from abc import ABC
-from typing import Iterable, Mapping, Union
-
-from torch import nn
-
-from .bandsplit import BandSplitModule
-from .tfmodel import (
-    SeqBandModellingModule,
-    TransformerTimeFreqModule,
-)
-
-
-class BandsplitCoreBase(nn.Module, ABC):
-    band_split: nn.Module
-    tf_model: nn.Module
-    mask_estim: Union[nn.Module, Mapping[str, nn.Module], Iterable[nn.Module]]
-
-    def __init__(self) -> None:
-        super().__init__()
-
-    @staticmethod
-    def mask(x, m):
-        return x * m
+from abc import ABC
+from typing import Iterable, Mapping, Union
+
+from torch import nn
+
+from .bandsplit import BandSplitModule
+from .tfmodel import (
+    SeqBandModellingModule,
+    TransformerTimeFreqModule,
+)
+
+
+class BandsplitCoreBase(nn.Module, ABC):
+    band_split: nn.Module
+    tf_model: nn.Module
+    mask_estim: Union[nn.Module, Mapping[str, nn.Module], Iterable[nn.Module]]
+
+    def __init__(self) -> None:
+        super().__init__()
+
+    @staticmethod
+    def mask(x, m):
+        return x * m

mvsepless/models/bandit/core/model/bsrnn/bandsplit.py

@@ -1,119 +1,119 @@
-from typing import List, Tuple
-
-import torch
-from torch import nn
-
-from .utils import (
-    band_widths_from_specs,
-    check_no_gap,
-    check_no_overlap,
-    check_nonzero_bandwidth,
-)
-
-
-class NormFC(nn.Module):
-    def __init__(
-        self,
-        emb_dim: int,
-        bandwidth: int,
-        in_channel: int,
-        normalize_channel_independently: bool = False,
-        treat_channel_as_feature: bool = True,
-    ) -> None:
-        super().__init__()
-
-        self.treat_channel_as_feature = treat_channel_as_feature
-
-        if normalize_channel_independently:
-            raise NotImplementedError
-
-        reim = 2
-
-        self.norm = nn.LayerNorm(in_channel * bandwidth * reim)
-
-        fc_in = bandwidth * reim
-
-        if treat_channel_as_feature:
-            fc_in *= in_channel
-        else:
-            assert emb_dim % in_channel == 0
-            emb_dim = emb_dim // in_channel
-
-        self.fc = nn.Linear(fc_in, emb_dim)
-
-    def forward(self, xb):
-
-        batch, n_time, in_chan, ribw = xb.shape
-        xb = self.norm(xb.reshape(batch, n_time, in_chan * ribw))
-
-        if not self.treat_channel_as_feature:
-            xb = xb.reshape(batch, n_time, in_chan, ribw)
-
-        zb = self.fc(xb)
-
-        if not self.treat_channel_as_feature:
-            batch, n_time, in_chan, emb_dim_per_chan = zb.shape
-            zb = zb.reshape((batch, n_time, in_chan * emb_dim_per_chan))
-
-        return zb
-
-
-class BandSplitModule(nn.Module):
-    def __init__(
-        self,
-        band_specs: List[Tuple[float, float]],
-        emb_dim: int,
-        in_channel: int,
-        require_no_overlap: bool = False,
-        require_no_gap: bool = True,
-        normalize_channel_independently: bool = False,
-        treat_channel_as_feature: bool = True,
-    ) -> None:
-        super().__init__()
-
-        check_nonzero_bandwidth(band_specs)
-
-        if require_no_gap:
-            check_no_gap(band_specs)
-
-        if require_no_overlap:
-            check_no_overlap(band_specs)
-
-        self.band_specs = band_specs
-        self.band_widths = band_widths_from_specs(band_specs)
-        self.n_bands = len(band_specs)
-        self.emb_dim = emb_dim
-
-        self.norm_fc_modules = nn.ModuleList(
-            [  # type: ignore
-                (
-                    NormFC(
-                        emb_dim=emb_dim,
-                        bandwidth=bw,
-                        in_channel=in_channel,
-                        normalize_channel_independently=normalize_channel_independently,
-                        treat_channel_as_feature=treat_channel_as_feature,
-                    )
-                )
-                for bw in self.band_widths
-            ]
-        )
-
-    def forward(self, x: torch.Tensor):
-
-        batch, in_chan, _, n_time = x.shape
-
-        z = torch.zeros(
-            size=(batch, self.n_bands, n_time, self.emb_dim), device=x.device
-        )
-
-        xr = torch.view_as_real(x)
-        xr = torch.permute(xr, (0, 3, 1, 4, 2))
-        batch, n_time, in_chan, reim, band_width = xr.shape
-        for i, nfm in enumerate(self.norm_fc_modules):
-            fstart, fend = self.band_specs[i]
-            xb = xr[..., fstart:fend]
-            xb = torch.reshape(xb, (batch, n_time, in_chan, -1))
-            z[:, i, :, :] = nfm(xb.contiguous())
-
-        return z
+from typing import List, Tuple
+
+import torch
+from torch import nn
+
+from .utils import (
+    band_widths_from_specs,
+    check_no_gap,
+    check_no_overlap,
+    check_nonzero_bandwidth,
+)
+
+
+class NormFC(nn.Module):
+    def __init__(
+        self,
+        emb_dim: int,
+        bandwidth: int,
+        in_channel: int,
+        normalize_channel_independently: bool = False,
+        treat_channel_as_feature: bool = True,
+    ) -> None:
+        super().__init__()
+
+        self.treat_channel_as_feature = treat_channel_as_feature
+
+        if normalize_channel_independently:
+            raise NotImplementedError
+
+        reim = 2
+
+        self.norm = nn.LayerNorm(in_channel * bandwidth * reim)
+
+        fc_in = bandwidth * reim
+
+        if treat_channel_as_feature:
+            fc_in *= in_channel
+        else:
+            assert emb_dim % in_channel == 0
+            emb_dim = emb_dim // in_channel
+
+        self.fc = nn.Linear(fc_in, emb_dim)
+
+    def forward(self, xb):
+
+        batch, n_time, in_chan, ribw = xb.shape
+        xb = self.norm(xb.reshape(batch, n_time, in_chan * ribw))
+
+        if not self.treat_channel_as_feature:
+            xb = xb.reshape(batch, n_time, in_chan, ribw)
+
+        zb = self.fc(xb)
+
+        if not self.treat_channel_as_feature:
+            batch, n_time, in_chan, emb_dim_per_chan = zb.shape
+            zb = zb.reshape((batch, n_time, in_chan * emb_dim_per_chan))
+
+        return zb
+
+
+class BandSplitModule(nn.Module):
+    def __init__(
+        self,
+        band_specs: List[Tuple[float, float]],
+        emb_dim: int,
+        in_channel: int,
+        require_no_overlap: bool = False,
+        require_no_gap: bool = True,
+        normalize_channel_independently: bool = False,
+        treat_channel_as_feature: bool = True,
+    ) -> None:
+        super().__init__()
+
+        check_nonzero_bandwidth(band_specs)
+
+        if require_no_gap:
+            check_no_gap(band_specs)
+
+        if require_no_overlap:
+            check_no_overlap(band_specs)
+
+        self.band_specs = band_specs
+        self.band_widths = band_widths_from_specs(band_specs)
+        self.n_bands = len(band_specs)
+        self.emb_dim = emb_dim
+
+        self.norm_fc_modules = nn.ModuleList(
+            [  # type: ignore
+                (
+                    NormFC(
+                        emb_dim=emb_dim,
+                        bandwidth=bw,
+                        in_channel=in_channel,
+                        normalize_channel_independently=normalize_channel_independently,
+                        treat_channel_as_feature=treat_channel_as_feature,
+                    )
+                )
+                for bw in self.band_widths
+            ]
+        )
+
+    def forward(self, x: torch.Tensor):
+
+        batch, in_chan, _, n_time = x.shape
+
+        z = torch.zeros(
+            size=(batch, self.n_bands, n_time, self.emb_dim), device=x.device
+        )
+
+        xr = torch.view_as_real(x)
+        xr = torch.permute(xr, (0, 3, 1, 4, 2))
+        batch, n_time, in_chan, reim, band_width = xr.shape
+        for i, nfm in enumerate(self.norm_fc_modules):
+            fstart, fend = self.band_specs[i]
+            xb = xr[..., fstart:fend]
+            xb = torch.reshape(xb, (batch, n_time, in_chan, -1))
+            z[:, i, :, :] = nfm(xb.contiguous())
+
+        return z

mvsepless/models/bandit/core/model/bsrnn/core.py

@@ -1,619 +1,619 @@
-from typing import Dict, List, Optional, Tuple
-
-import torch
-from torch import nn
-from torch.nn import functional as F
-
-from . import BandsplitCoreBase
-from .bandsplit import BandSplitModule
-from .maskestim import (
-    MaskEstimationModule,
-    OverlappingMaskEstimationModule,
-)
-from .tfmodel import (
-    ConvolutionalTimeFreqModule,
-    SeqBandModellingModule,
-    TransformerTimeFreqModule,
-)
-
-
-class MultiMaskBandSplitCoreBase(BandsplitCoreBase):
-    def __init__(self) -> None:
-        super().__init__()
-
-    def forward(self, x, cond=None, compute_residual: bool = True):
-        batch, in_chan, n_freq, n_time = x.shape
-        x = torch.reshape(x, (-1, 1, n_freq, n_time))
-
-        z = self.band_split(x)
-
-        q = self.tf_model(z)
-
-        out = {}
-
-        for stem, mem in self.mask_estim.items():
-            m = mem(q, cond=cond)
-
-            s = self.mask(x, m)
-            s = torch.reshape(s, (batch, in_chan, n_freq, n_time))
-            out[stem] = s
-
-        return {"spectrogram": out}
-
-    def instantiate_mask_estim(
-        self,
-        in_channel: int,
-        stems: List[str],
-        band_specs: List[Tuple[float, float]],
-        emb_dim: int,
-        mlp_dim: int,
-        cond_dim: int,
-        hidden_activation: str,
-        hidden_activation_kwargs: Optional[Dict] = None,
-        complex_mask: bool = True,
-        overlapping_band: bool = False,
-        freq_weights: Optional[List[torch.Tensor]] = None,
-        n_freq: Optional[int] = None,
-        use_freq_weights: bool = True,
-        mult_add_mask: bool = False,
-    ):
-        if hidden_activation_kwargs is None:
-            hidden_activation_kwargs = {}
-
-        if "mne:+" in stems:
-            stems = [s for s in stems if s != "mne:+"]
-
-        if overlapping_band:
-            assert freq_weights is not None
-            assert n_freq is not None
-
-            if mult_add_mask:
-
-                self.mask_estim = nn.ModuleDict(
-                    {
-                        stem: MultAddMaskEstimationModule(
-                            band_specs=band_specs,
-                            freq_weights=freq_weights,
-                            n_freq=n_freq,
-                            emb_dim=emb_dim,
-                            mlp_dim=mlp_dim,
-                            in_channel=in_channel,
-                            hidden_activation=hidden_activation,
-                            hidden_activation_kwargs=hidden_activation_kwargs,
-                            complex_mask=complex_mask,
-                            use_freq_weights=use_freq_weights,
-                        )
-                        for stem in stems
-                    }
-                )
-            else:
-                self.mask_estim = nn.ModuleDict(
-                    {
-                        stem: OverlappingMaskEstimationModule(
-                            band_specs=band_specs,
-                            freq_weights=freq_weights,
-                            n_freq=n_freq,
-                            emb_dim=emb_dim,
-                            mlp_dim=mlp_dim,
-                            in_channel=in_channel,
-                            hidden_activation=hidden_activation,
-                            hidden_activation_kwargs=hidden_activation_kwargs,
-                            complex_mask=complex_mask,
-                            use_freq_weights=use_freq_weights,
-                        )
-                        for stem in stems
-                    }
-                )
-        else:
-            self.mask_estim = nn.ModuleDict(
-                {
-                    stem: MaskEstimationModule(
-                        band_specs=band_specs,
-                        emb_dim=emb_dim,
-                        mlp_dim=mlp_dim,
-                        cond_dim=cond_dim,
-                        in_channel=in_channel,
-                        hidden_activation=hidden_activation,
-                        hidden_activation_kwargs=hidden_activation_kwargs,
-                        complex_mask=complex_mask,
-                    )
-                    for stem in stems
-                }
-            )
-
-    def instantiate_bandsplit(
-        self,
-        in_channel: int,
-        band_specs: List[Tuple[float, float]],
-        require_no_overlap: bool = False,
-        require_no_gap: bool = True,
-        normalize_channel_independently: bool = False,
-        treat_channel_as_feature: bool = True,
-        emb_dim: int = 128,
-    ):
-        self.band_split = BandSplitModule(
-            in_channel=in_channel,
-            band_specs=band_specs,
-            require_no_overlap=require_no_overlap,
-            require_no_gap=require_no_gap,
-            normalize_channel_independently=normalize_channel_independently,
-            treat_channel_as_feature=treat_channel_as_feature,
-            emb_dim=emb_dim,
-        )
-
-
-class SingleMaskBandsplitCoreBase(BandsplitCoreBase):
-    def __init__(self, **kwargs) -> None:
-        super().__init__()
-
-    def forward(self, x):
-        z = self.band_split(x)
-        q = self.tf_model(z)
-        m = self.mask_estim(q)
-
-        s = self.mask(x, m)
-
-        return s
-
-
-class SingleMaskBandsplitCoreRNN(
-    SingleMaskBandsplitCoreBase,
-):
-    def __init__(
-        self,
-        in_channel: int,
-        band_specs: List[Tuple[float, float]],
-        require_no_overlap: bool = False,
-        require_no_gap: bool = True,
-        normalize_channel_independently: bool = False,
-        treat_channel_as_feature: bool = True,
-        n_sqm_modules: int = 12,
-        emb_dim: int = 128,
-        rnn_dim: int = 256,
-        bidirectional: bool = True,
-        rnn_type: str = "LSTM",
-        mlp_dim: int = 512,
-        hidden_activation: str = "Tanh",
-        hidden_activation_kwargs: Optional[Dict] = None,
-        complex_mask: bool = True,
-    ) -> None:
-        super().__init__()
-        self.band_split = BandSplitModule(
-            in_channel=in_channel,
-            band_specs=band_specs,
-            require_no_overlap=require_no_overlap,
-            require_no_gap=require_no_gap,
-            normalize_channel_independently=normalize_channel_independently,
-            treat_channel_as_feature=treat_channel_as_feature,
-            emb_dim=emb_dim,
-        )
-        self.tf_model = SeqBandModellingModule(
-            n_modules=n_sqm_modules,
-            emb_dim=emb_dim,
-            rnn_dim=rnn_dim,
-            bidirectional=bidirectional,
-            rnn_type=rnn_type,
-        )
-        self.mask_estim = MaskEstimationModule(
-            in_channel=in_channel,
-            band_specs=band_specs,
-            emb_dim=emb_dim,
-            mlp_dim=mlp_dim,
-            hidden_activation=hidden_activation,
-            hidden_activation_kwargs=hidden_activation_kwargs,
-            complex_mask=complex_mask,
-        )
-
-
-class SingleMaskBandsplitCoreTransformer(
-    SingleMaskBandsplitCoreBase,
-):
-    def __init__(
-        self,
-        in_channel: int,
-        band_specs: List[Tuple[float, float]],
-        require_no_overlap: bool = False,
-        require_no_gap: bool = True,
-        normalize_channel_independently: bool = False,
-        treat_channel_as_feature: bool = True,
-        n_sqm_modules: int = 12,
-        emb_dim: int = 128,
-        rnn_dim: int = 256,
-        bidirectional: bool = True,
-        tf_dropout: float = 0.0,
-        mlp_dim: int = 512,
-        hidden_activation: str = "Tanh",
-        hidden_activation_kwargs: Optional[Dict] = None,
-        complex_mask: bool = True,
-    ) -> None:
-        super().__init__()
-        self.band_split = BandSplitModule(
-            in_channel=in_channel,
-            band_specs=band_specs,
-            require_no_overlap=require_no_overlap,
-            require_no_gap=require_no_gap,
-            normalize_channel_independently=normalize_channel_independently,
-            treat_channel_as_feature=treat_channel_as_feature,
-            emb_dim=emb_dim,
-        )
-        self.tf_model = TransformerTimeFreqModule(
-            n_modules=n_sqm_modules,
-            emb_dim=emb_dim,
-            rnn_dim=rnn_dim,
-            bidirectional=bidirectional,
-            dropout=tf_dropout,
-        )
-        self.mask_estim = MaskEstimationModule(
-            in_channel=in_channel,
-            band_specs=band_specs,
-            emb_dim=emb_dim,
-            mlp_dim=mlp_dim,
-            hidden_activation=hidden_activation,
-            hidden_activation_kwargs=hidden_activation_kwargs,
-            complex_mask=complex_mask,
-        )
-
-
-class MultiSourceMultiMaskBandSplitCoreRNN(MultiMaskBandSplitCoreBase):
-    def __init__(
-        self,
-        in_channel: int,
-        stems: List[str],
-        band_specs: List[Tuple[float, float]],
-        require_no_overlap: bool = False,
-        require_no_gap: bool = True,
-        normalize_channel_independently: bool = False,
-        treat_channel_as_feature: bool = True,
-        n_sqm_modules: int = 12,
-        emb_dim: int = 128,
-        rnn_dim: int = 256,
-        bidirectional: bool = True,
-        rnn_type: str = "LSTM",
-        mlp_dim: int = 512,
-        cond_dim: int = 0,
-        hidden_activation: str = "Tanh",
-        hidden_activation_kwargs: Optional[Dict] = None,
-        complex_mask: bool = True,
-        overlapping_band: bool = False,
-        freq_weights: Optional[List[torch.Tensor]] = None,
-        n_freq: Optional[int] = None,
-        use_freq_weights: bool = True,
-        mult_add_mask: bool = False,
-    ) -> None:
-
-        super().__init__()
-        self.instantiate_bandsplit(
-            in_channel=in_channel,
-            band_specs=band_specs,
-            require_no_overlap=require_no_overlap,
-            require_no_gap=require_no_gap,
-            normalize_channel_independently=normalize_channel_independently,
-            treat_channel_as_feature=treat_channel_as_feature,
-            emb_dim=emb_dim,
-        )
-
-        self.tf_model = SeqBandModellingModule(
-            n_modules=n_sqm_modules,
-            emb_dim=emb_dim,
-            rnn_dim=rnn_dim,
-            bidirectional=bidirectional,
-            rnn_type=rnn_type,
-        )
-
-        self.mult_add_mask = mult_add_mask
-
-        self.instantiate_mask_estim(
-            in_channel=in_channel,
-            stems=stems,
-            band_specs=band_specs,
-            emb_dim=emb_dim,
-            mlp_dim=mlp_dim,
-            cond_dim=cond_dim,
-            hidden_activation=hidden_activation,
-            hidden_activation_kwargs=hidden_activation_kwargs,
-            complex_mask=complex_mask,
-            overlapping_band=overlapping_band,
-            freq_weights=freq_weights,
-            n_freq=n_freq,
-            use_freq_weights=use_freq_weights,
-            mult_add_mask=mult_add_mask,
-        )
-
-    @staticmethod
-    def _mult_add_mask(x, m):
-
-        assert m.ndim == 5
-
-        mm = m[..., 0]
-        am = m[..., 1]
-
-        return x * mm + am
-
-    def mask(self, x, m):
-        if self.mult_add_mask:
-
-            return self._mult_add_mask(x, m)
-        else:
-            return super().mask(x, m)
-
-
-class MultiSourceMultiMaskBandSplitCoreTransformer(
-    MultiMaskBandSplitCoreBase,
-):
-    def __init__(
-        self,
-        in_channel: int,
-        stems: List[str],
-        band_specs: List[Tuple[float, float]],
-        require_no_overlap: bool = False,
-        require_no_gap: bool = True,
-        normalize_channel_independently: bool = False,
-        treat_channel_as_feature: bool = True,
-        n_sqm_modules: int = 12,
-        emb_dim: int = 128,
-        rnn_dim: int = 256,
-        bidirectional: bool = True,
-        tf_dropout: float = 0.0,
-        mlp_dim: int = 512,
-        hidden_activation: str = "Tanh",
-        hidden_activation_kwargs: Optional[Dict] = None,
-        complex_mask: bool = True,
-        overlapping_band: bool = False,
-        freq_weights: Optional[List[torch.Tensor]] = None,
-        n_freq: Optional[int] = None,
-        use_freq_weights: bool = True,
-        rnn_type: str = "LSTM",
-        cond_dim: int = 0,
-        mult_add_mask: bool = False,
-    ) -> None:
-        super().__init__()
-        self.instantiate_bandsplit(
-            in_channel=in_channel,
-            band_specs=band_specs,
-            require_no_overlap=require_no_overlap,
-            require_no_gap=require_no_gap,
-            normalize_channel_independently=normalize_channel_independently,
-            treat_channel_as_feature=treat_channel_as_feature,
-            emb_dim=emb_dim,
-        )
-        self.tf_model = TransformerTimeFreqModule(
-            n_modules=n_sqm_modules,
-            emb_dim=emb_dim,
-            rnn_dim=rnn_dim,
-            bidirectional=bidirectional,
-            dropout=tf_dropout,
-        )
-
-        self.instantiate_mask_estim(
-            in_channel=in_channel,
-            stems=stems,
-            band_specs=band_specs,
-            emb_dim=emb_dim,
-            mlp_dim=mlp_dim,
-            cond_dim=cond_dim,
-            hidden_activation=hidden_activation,
-            hidden_activation_kwargs=hidden_activation_kwargs,
-            complex_mask=complex_mask,
-            overlapping_band=overlapping_band,
-            freq_weights=freq_weights,
-            n_freq=n_freq,
-            use_freq_weights=use_freq_weights,
-            mult_add_mask=mult_add_mask,
-        )
-
-
-class MultiSourceMultiMaskBandSplitCoreConv(
-    MultiMaskBandSplitCoreBase,
-):
-    def __init__(
-        self,
-        in_channel: int,
-        stems: List[str],
-        band_specs: List[Tuple[float, float]],
-        require_no_overlap: bool = False,
-        require_no_gap: bool = True,
-        normalize_channel_independently: bool = False,
-        treat_channel_as_feature: bool = True,
-        n_sqm_modules: int = 12,
-        emb_dim: int = 128,
-        rnn_dim: int = 256,
-        bidirectional: bool = True,
-        tf_dropout: float = 0.0,
-        mlp_dim: int = 512,
-        hidden_activation: str = "Tanh",
-        hidden_activation_kwargs: Optional[Dict] = None,
-        complex_mask: bool = True,
-        overlapping_band: bool = False,
-        freq_weights: Optional[List[torch.Tensor]] = None,
-        n_freq: Optional[int] = None,
-        use_freq_weights: bool = True,
-        rnn_type: str = "LSTM",
-        cond_dim: int = 0,
-        mult_add_mask: bool = False,
-    ) -> None:
-        super().__init__()
-        self.instantiate_bandsplit(
-            in_channel=in_channel,
-            band_specs=band_specs,
-            require_no_overlap=require_no_overlap,
-            require_no_gap=require_no_gap,
-            normalize_channel_independently=normalize_channel_independently,
-            treat_channel_as_feature=treat_channel_as_feature,
-            emb_dim=emb_dim,
-        )
-        self.tf_model = ConvolutionalTimeFreqModule(
-            n_modules=n_sqm_modules,
-            emb_dim=emb_dim,
-            rnn_dim=rnn_dim,
-            bidirectional=bidirectional,
-            dropout=tf_dropout,
-        )
-
-        self.instantiate_mask_estim(
-            in_channel=in_channel,
-            stems=stems,
-            band_specs=band_specs,
-            emb_dim=emb_dim,
-            mlp_dim=mlp_dim,
-            cond_dim=cond_dim,
-            hidden_activation=hidden_activation,
-            hidden_activation_kwargs=hidden_activation_kwargs,
-            complex_mask=complex_mask,
-            overlapping_band=overlapping_band,
-            freq_weights=freq_weights,
-            n_freq=n_freq,
-            use_freq_weights=use_freq_weights,
-            mult_add_mask=mult_add_mask,
-        )
-
-
-class PatchingMaskBandsplitCoreBase(MultiMaskBandSplitCoreBase):
-    def __init__(self) -> None:
-        super().__init__()
-
-    def mask(self, x, m):
-
-        _, n_channel, kernel_freq, kernel_time, n_freq, n_time = m.shape
-        padding = ((kernel_freq - 1) // 2, (kernel_time - 1) // 2)
-
-        xf = F.unfold(
-            x,
-            kernel_size=(kernel_freq, kernel_time),
-            padding=padding,
-            stride=(1, 1),
-        )
-
-        xf = xf.view(
-            -1,
-            n_channel,
-            kernel_freq,
-            kernel_time,
-            n_freq,
-            n_time,
-        )
-
-        sf = xf * m
-
-        sf = sf.view(
-            -1,
-            n_channel * kernel_freq * kernel_time,
-            n_freq * n_time,
-        )
-
-        s = F.fold(
-            sf,
-            output_size=(n_freq, n_time),
-            kernel_size=(kernel_freq, kernel_time),
-            padding=padding,
-            stride=(1, 1),
-        ).view(
-            -1,
-            n_channel,
-            n_freq,
-            n_time,
-        )
-
-        return s
-
-    def old_mask(self, x, m):
-
-        s = torch.zeros_like(x)
-
-        _, n_channel, n_freq, n_time = x.shape
-        kernel_freq, kernel_time, _, _, _, _ = m.shape
-
-        kernel_freq_half = (kernel_freq - 1) // 2
-        kernel_time_half = (kernel_time - 1) // 2
-
-        for ifreq in range(kernel_freq):
-            for itime in range(kernel_time):
-                df, dt = kernel_freq_half - ifreq, kernel_time_half - itime
-                x = x.roll(shifts=(df, dt), dims=(2, 3))
-
-                fslice = slice(max(0, df), min(n_freq, n_freq + df))
-                tslice = slice(max(0, dt), min(n_time, n_time + dt))
-
-                s[:, :, fslice, tslice] += (
-                    x[:, :, fslice, tslice] * m[ifreq, itime, :, :, fslice, tslice]
-                )
-
-        return s
-
-
-class MultiSourceMultiPatchingMaskBandSplitCoreRNN(PatchingMaskBandsplitCoreBase):
-    def __init__(
-        self,
-        in_channel: int,
-        stems: List[str],
-        band_specs: List[Tuple[float, float]],
-        mask_kernel_freq: int,
-        mask_kernel_time: int,
-        conv_kernel_freq: int,
-        conv_kernel_time: int,
-        kernel_norm_mlp_version: int,
-        require_no_overlap: bool = False,
-        require_no_gap: bool = True,
-        normalize_channel_independently: bool = False,
-        treat_channel_as_feature: bool = True,
-        n_sqm_modules: int = 12,
-        emb_dim: int = 128,
-        rnn_dim: int = 256,
-        bidirectional: bool = True,
-        rnn_type: str = "LSTM",
-        mlp_dim: int = 512,
-        hidden_activation: str = "Tanh",
-        hidden_activation_kwargs: Optional[Dict] = None,
-        complex_mask: bool = True,
-        overlapping_band: bool = False,
-        freq_weights: Optional[List[torch.Tensor]] = None,
-        n_freq: Optional[int] = None,
-    ) -> None:
-
-        super().__init__()
-        self.band_split = BandSplitModule(
-            in_channel=in_channel,
-            band_specs=band_specs,
-            require_no_overlap=require_no_overlap,
-            require_no_gap=require_no_gap,
-            normalize_channel_independently=normalize_channel_independently,
-            treat_channel_as_feature=treat_channel_as_feature,
-            emb_dim=emb_dim,
-        )
-
-        self.tf_model = SeqBandModellingModule(
-            n_modules=n_sqm_modules,
-            emb_dim=emb_dim,
-            rnn_dim=rnn_dim,
-            bidirectional=bidirectional,
-            rnn_type=rnn_type,
-        )
-
-        if hidden_activation_kwargs is None:
-            hidden_activation_kwargs = {}
-
-        if overlapping_band:
-            assert freq_weights is not None
-            assert n_freq is not None
-            self.mask_estim = nn.ModuleDict(
-                {
-                    stem: PatchingMaskEstimationModule(
-                        band_specs=band_specs,
-                        freq_weights=freq_weights,
-                        n_freq=n_freq,
-                        emb_dim=emb_dim,
-                        mlp_dim=mlp_dim,
-                        in_channel=in_channel,
-                        hidden_activation=hidden_activation,
-                        hidden_activation_kwargs=hidden_activation_kwargs,
-                        complex_mask=complex_mask,
-                        mask_kernel_freq=mask_kernel_freq,
-                        mask_kernel_time=mask_kernel_time,
-                        conv_kernel_freq=conv_kernel_freq,
-                        conv_kernel_time=conv_kernel_time,
-                        kernel_norm_mlp_version=kernel_norm_mlp_version,
-                    )
-                    for stem in stems
-                }
-            )
-        else:
-            raise NotImplementedError
+from typing import Dict, List, Optional, Tuple
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from . import BandsplitCoreBase
+from .bandsplit import BandSplitModule
+from .maskestim import (
+    MaskEstimationModule,
+    OverlappingMaskEstimationModule,
+)
+from .tfmodel import (
+    ConvolutionalTimeFreqModule,
+    SeqBandModellingModule,
+    TransformerTimeFreqModule,
+)
+
+
+class MultiMaskBandSplitCoreBase(BandsplitCoreBase):
+    def __init__(self) -> None:
+        super().__init__()
+
+    def forward(self, x, cond=None, compute_residual: bool = True):
+        batch, in_chan, n_freq, n_time = x.shape
+        x = torch.reshape(x, (-1, 1, n_freq, n_time))
+
+        z = self.band_split(x)
+
+        q = self.tf_model(z)
+
+        out = {}
+
+        for stem, mem in self.mask_estim.items():
+            m = mem(q, cond=cond)
+
+            s = self.mask(x, m)
+            s = torch.reshape(s, (batch, in_chan, n_freq, n_time))
+            out[stem] = s
+
+        return {"spectrogram": out}
+
+    def instantiate_mask_estim(
+        self,
+        in_channel: int,
+        stems: List[str],
+        band_specs: List[Tuple[float, float]],
+        emb_dim: int,
+        mlp_dim: int,
+        cond_dim: int,
+        hidden_activation: str,
+        hidden_activation_kwargs: Optional[Dict] = None,
+        complex_mask: bool = True,
+        overlapping_band: bool = False,
+        freq_weights: Optional[List[torch.Tensor]] = None,
+        n_freq: Optional[int] = None,
+        use_freq_weights: bool = True,
+        mult_add_mask: bool = False,
+    ):
+        if hidden_activation_kwargs is None:
+            hidden_activation_kwargs = {}
+
+        if "mne:+" in stems:
+            stems = [s for s in stems if s != "mne:+"]
+
+        if overlapping_band:
+            assert freq_weights is not None
+            assert n_freq is not None
+
+            if mult_add_mask:
+
+                self.mask_estim = nn.ModuleDict(
+                    {
+                        stem: MultAddMaskEstimationModule(
+                            band_specs=band_specs,
+                            freq_weights=freq_weights,
+                            n_freq=n_freq,
+                            emb_dim=emb_dim,
+                            mlp_dim=mlp_dim,
+                            in_channel=in_channel,
+                            hidden_activation=hidden_activation,
+                            hidden_activation_kwargs=hidden_activation_kwargs,
+                            complex_mask=complex_mask,
+                            use_freq_weights=use_freq_weights,
+                        )
+                        for stem in stems
+                    }
+                )
+            else:
+                self.mask_estim = nn.ModuleDict(
+                    {
+                        stem: OverlappingMaskEstimationModule(
+                            band_specs=band_specs,
+                            freq_weights=freq_weights,
+                            n_freq=n_freq,
+                            emb_dim=emb_dim,
+                            mlp_dim=mlp_dim,
+                            in_channel=in_channel,
+                            hidden_activation=hidden_activation,
+                            hidden_activation_kwargs=hidden_activation_kwargs,
+                            complex_mask=complex_mask,
+                            use_freq_weights=use_freq_weights,
+                        )
+                        for stem in stems
+                    }
+                )
+        else:
+            self.mask_estim = nn.ModuleDict(
+                {
+                    stem: MaskEstimationModule(
+                        band_specs=band_specs,
+                        emb_dim=emb_dim,
+                        mlp_dim=mlp_dim,
+                        cond_dim=cond_dim,
+                        in_channel=in_channel,
+                        hidden_activation=hidden_activation,
+                        hidden_activation_kwargs=hidden_activation_kwargs,
+                        complex_mask=complex_mask,
+                    )
+                    for stem in stems
+                }
+            )
+
+    def instantiate_bandsplit(
+        self,
+        in_channel: int,
+        band_specs: List[Tuple[float, float]],
+        require_no_overlap: bool = False,
+        require_no_gap: bool = True,
+        normalize_channel_independently: bool = False,
+        treat_channel_as_feature: bool = True,
+        emb_dim: int = 128,
+    ):
+        self.band_split = BandSplitModule(
+            in_channel=in_channel,
+            band_specs=band_specs,
+            require_no_overlap=require_no_overlap,
+            require_no_gap=require_no_gap,
+            normalize_channel_independently=normalize_channel_independently,
+            treat_channel_as_feature=treat_channel_as_feature,
+            emb_dim=emb_dim,
+        )
+
+
+class SingleMaskBandsplitCoreBase(BandsplitCoreBase):
+    def __init__(self, **kwargs) -> None:
+        super().__init__()
+
+    def forward(self, x):
+        z = self.band_split(x)
+        q = self.tf_model(z)
+        m = self.mask_estim(q)
+
+        s = self.mask(x, m)
+
+        return s
+
+
+class SingleMaskBandsplitCoreRNN(
+    SingleMaskBandsplitCoreBase,
+):
+    def __init__(
+        self,
+        in_channel: int,
+        band_specs: List[Tuple[float, float]],
+        require_no_overlap: bool = False,
+        require_no_gap: bool = True,
+        normalize_channel_independently: bool = False,
+        treat_channel_as_feature: bool = True,
+        n_sqm_modules: int = 12,
+        emb_dim: int = 128,
+        rnn_dim: int = 256,
+        bidirectional: bool = True,
+        rnn_type: str = "LSTM",
+        mlp_dim: int = 512,
+        hidden_activation: str = "Tanh",
+        hidden_activation_kwargs: Optional[Dict] = None,
+        complex_mask: bool = True,
+    ) -> None:
+        super().__init__()
+        self.band_split = BandSplitModule(
+            in_channel=in_channel,
+            band_specs=band_specs,
+            require_no_overlap=require_no_overlap,
+            require_no_gap=require_no_gap,
+            normalize_channel_independently=normalize_channel_independently,
+            treat_channel_as_feature=treat_channel_as_feature,
+            emb_dim=emb_dim,
+        )
+        self.tf_model = SeqBandModellingModule(
+            n_modules=n_sqm_modules,
+            emb_dim=emb_dim,
+            rnn_dim=rnn_dim,
+            bidirectional=bidirectional,
+            rnn_type=rnn_type,
+        )
+        self.mask_estim = MaskEstimationModule(
+            in_channel=in_channel,
+            band_specs=band_specs,
+            emb_dim=emb_dim,
+            mlp_dim=mlp_dim,
+            hidden_activation=hidden_activation,
+            hidden_activation_kwargs=hidden_activation_kwargs,
+            complex_mask=complex_mask,
+        )
+
+
+class SingleMaskBandsplitCoreTransformer(
+    SingleMaskBandsplitCoreBase,
+):
+    def __init__(
+        self,
+        in_channel: int,
+        band_specs: List[Tuple[float, float]],
+        require_no_overlap: bool = False,
+        require_no_gap: bool = True,
+        normalize_channel_independently: bool = False,
+        treat_channel_as_feature: bool = True,
+        n_sqm_modules: int = 12,
+        emb_dim: int = 128,
+        rnn_dim: int = 256,
+        bidirectional: bool = True,
+        tf_dropout: float = 0.0,
+        mlp_dim: int = 512,
+        hidden_activation: str = "Tanh",
+        hidden_activation_kwargs: Optional[Dict] = None,
+        complex_mask: bool = True,
+    ) -> None:
+        super().__init__()
+        self.band_split = BandSplitModule(
+            in_channel=in_channel,
+            band_specs=band_specs,
+            require_no_overlap=require_no_overlap,
+            require_no_gap=require_no_gap,
+            normalize_channel_independently=normalize_channel_independently,
+            treat_channel_as_feature=treat_channel_as_feature,
+            emb_dim=emb_dim,
+        )
+        self.tf_model = TransformerTimeFreqModule(
+            n_modules=n_sqm_modules,
+            emb_dim=emb_dim,
+            rnn_dim=rnn_dim,
+            bidirectional=bidirectional,
+            dropout=tf_dropout,
+        )
+        self.mask_estim = MaskEstimationModule(
+            in_channel=in_channel,
+            band_specs=band_specs,
+            emb_dim=emb_dim,
+            mlp_dim=mlp_dim,
+            hidden_activation=hidden_activation,
+            hidden_activation_kwargs=hidden_activation_kwargs,
+            complex_mask=complex_mask,
+        )
+
+
+class MultiSourceMultiMaskBandSplitCoreRNN(MultiMaskBandSplitCoreBase):
+    def __init__(
+        self,
+        in_channel: int,
+        stems: List[str],
+        band_specs: List[Tuple[float, float]],
+        require_no_overlap: bool = False,
+        require_no_gap: bool = True,
+        normalize_channel_independently: bool = False,
+        treat_channel_as_feature: bool = True,
+        n_sqm_modules: int = 12,
+        emb_dim: int = 128,
+        rnn_dim: int = 256,
+        bidirectional: bool = True,
+        rnn_type: str = "LSTM",
+        mlp_dim: int = 512,
+        cond_dim: int = 0,
+        hidden_activation: str = "Tanh",
+        hidden_activation_kwargs: Optional[Dict] = None,
+        complex_mask: bool = True,
+        overlapping_band: bool = False,
+        freq_weights: Optional[List[torch.Tensor]] = None,
+        n_freq: Optional[int] = None,
+        use_freq_weights: bool = True,
+        mult_add_mask: bool = False,
+    ) -> None:
+
+        super().__init__()
+        self.instantiate_bandsplit(
+            in_channel=in_channel,
+            band_specs=band_specs,
+            require_no_overlap=require_no_overlap,
+            require_no_gap=require_no_gap,
+            normalize_channel_independently=normalize_channel_independently,
+            treat_channel_as_feature=treat_channel_as_feature,
+            emb_dim=emb_dim,
+        )
+
+        self.tf_model = SeqBandModellingModule(
+            n_modules=n_sqm_modules,
+            emb_dim=emb_dim,
+            rnn_dim=rnn_dim,
+            bidirectional=bidirectional,
+            rnn_type=rnn_type,
+        )
+
+        self.mult_add_mask = mult_add_mask
+
+        self.instantiate_mask_estim(
+            in_channel=in_channel,
+            stems=stems,
+            band_specs=band_specs,
+            emb_dim=emb_dim,
+            mlp_dim=mlp_dim,
+            cond_dim=cond_dim,
+            hidden_activation=hidden_activation,
+            hidden_activation_kwargs=hidden_activation_kwargs,
+            complex_mask=complex_mask,
+            overlapping_band=overlapping_band,
+            freq_weights=freq_weights,
+            n_freq=n_freq,
+            use_freq_weights=use_freq_weights,
+            mult_add_mask=mult_add_mask,
+        )
+
+    @staticmethod
+    def _mult_add_mask(x, m):
+
+        assert m.ndim == 5
+
+        mm = m[..., 0]
+        am = m[..., 1]
+
+        return x * mm + am
+
+    def mask(self, x, m):
+        if self.mult_add_mask:
+
+            return self._mult_add_mask(x, m)
+        else:
+            return super().mask(x, m)
+
+
+class MultiSourceMultiMaskBandSplitCoreTransformer(
+    MultiMaskBandSplitCoreBase,
+):
+    def __init__(
+        self,
+        in_channel: int,
+        stems: List[str],
+        band_specs: List[Tuple[float, float]],
+        require_no_overlap: bool = False,
+        require_no_gap: bool = True,
+        normalize_channel_independently: bool = False,
+        treat_channel_as_feature: bool = True,
+        n_sqm_modules: int = 12,
+        emb_dim: int = 128,
+        rnn_dim: int = 256,
+        bidirectional: bool = True,
+        tf_dropout: float = 0.0,
+        mlp_dim: int = 512,
+        hidden_activation: str = "Tanh",
+        hidden_activation_kwargs: Optional[Dict] = None,
+        complex_mask: bool = True,
+        overlapping_band: bool = False,
+        freq_weights: Optional[List[torch.Tensor]] = None,
+        n_freq: Optional[int] = None,
+        use_freq_weights: bool = True,
+        rnn_type: str = "LSTM",
+        cond_dim: int = 0,
+        mult_add_mask: bool = False,
+    ) -> None:
+        super().__init__()
+        self.instantiate_bandsplit(
+            in_channel=in_channel,
+            band_specs=band_specs,
+            require_no_overlap=require_no_overlap,
+            require_no_gap=require_no_gap,
+            normalize_channel_independently=normalize_channel_independently,
+            treat_channel_as_feature=treat_channel_as_feature,
+            emb_dim=emb_dim,
+        )
+        self.tf_model = TransformerTimeFreqModule(
+            n_modules=n_sqm_modules,
+            emb_dim=emb_dim,
+            rnn_dim=rnn_dim,
+            bidirectional=bidirectional,
+            dropout=tf_dropout,
+        )
+
+        self.instantiate_mask_estim(
+            in_channel=in_channel,
+            stems=stems,
+            band_specs=band_specs,
+            emb_dim=emb_dim,
+            mlp_dim=mlp_dim,
+            cond_dim=cond_dim,
+            hidden_activation=hidden_activation,
+            hidden_activation_kwargs=hidden_activation_kwargs,
+            complex_mask=complex_mask,
+            overlapping_band=overlapping_band,
+            freq_weights=freq_weights,
+            n_freq=n_freq,
+            use_freq_weights=use_freq_weights,
+            mult_add_mask=mult_add_mask,
+        )
+
+
+class MultiSourceMultiMaskBandSplitCoreConv(
+    MultiMaskBandSplitCoreBase,
+):
+    def __init__(
+        self,
+        in_channel: int,
+        stems: List[str],
+        band_specs: List[Tuple[float, float]],
+        require_no_overlap: bool = False,
+        require_no_gap: bool = True,
+        normalize_channel_independently: bool = False,
+        treat_channel_as_feature: bool = True,
+        n_sqm_modules: int = 12,
+        emb_dim: int = 128,
+        rnn_dim: int = 256,
+        bidirectional: bool = True,
+        tf_dropout: float = 0.0,
+        mlp_dim: int = 512,
+        hidden_activation: str = "Tanh",
+        hidden_activation_kwargs: Optional[Dict] = None,
+        complex_mask: bool = True,
+        overlapping_band: bool = False,
+        freq_weights: Optional[List[torch.Tensor]] = None,
+        n_freq: Optional[int] = None,
+        use_freq_weights: bool = True,
+        rnn_type: str = "LSTM",
+        cond_dim: int = 0,
+        mult_add_mask: bool = False,
+    ) -> None:
+        super().__init__()
+        self.instantiate_bandsplit(
+            in_channel=in_channel,
+            band_specs=band_specs,
+            require_no_overlap=require_no_overlap,
+            require_no_gap=require_no_gap,
+            normalize_channel_independently=normalize_channel_independently,
+            treat_channel_as_feature=treat_channel_as_feature,
+            emb_dim=emb_dim,
+        )
+        self.tf_model = ConvolutionalTimeFreqModule(
+            n_modules=n_sqm_modules,
+            emb_dim=emb_dim,
+            rnn_dim=rnn_dim,
+            bidirectional=bidirectional,
+            dropout=tf_dropout,
+        )
+
+        self.instantiate_mask_estim(
+            in_channel=in_channel,
+            stems=stems,
+            band_specs=band_specs,
+            emb_dim=emb_dim,
+            mlp_dim=mlp_dim,
+            cond_dim=cond_dim,
+            hidden_activation=hidden_activation,
+            hidden_activation_kwargs=hidden_activation_kwargs,
+            complex_mask=complex_mask,
+            overlapping_band=overlapping_band,
+            freq_weights=freq_weights,
+            n_freq=n_freq,
+            use_freq_weights=use_freq_weights,
+            mult_add_mask=mult_add_mask,
+        )
+
+
+class PatchingMaskBandsplitCoreBase(MultiMaskBandSplitCoreBase):
+    def __init__(self) -> None:
+        super().__init__()
+
+    def mask(self, x, m):
+
+        _, n_channel, kernel_freq, kernel_time, n_freq, n_time = m.shape
+        padding = ((kernel_freq - 1) // 2, (kernel_time - 1) // 2)
+
+        xf = F.unfold(
+            x,
+            kernel_size=(kernel_freq, kernel_time),
+            padding=padding,
+            stride=(1, 1),
+        )
+
+        xf = xf.view(
+            -1,
+            n_channel,
+            kernel_freq,
+            kernel_time,
+            n_freq,
+            n_time,
+        )
+
+        sf = xf * m
+
+        sf = sf.view(
+            -1,
+            n_channel * kernel_freq * kernel_time,
+            n_freq * n_time,
+        )
+
+        s = F.fold(
+            sf,
+            output_size=(n_freq, n_time),
+            kernel_size=(kernel_freq, kernel_time),
+            padding=padding,
+            stride=(1, 1),
+        ).view(
+            -1,
+            n_channel,
+            n_freq,
+            n_time,
+        )
+
+        return s
+
+    def old_mask(self, x, m):
+
+        s = torch.zeros_like(x)
+
+        _, n_channel, n_freq, n_time = x.shape
+        kernel_freq, kernel_time, _, _, _, _ = m.shape
+
+        kernel_freq_half = (kernel_freq - 1) // 2
+        kernel_time_half = (kernel_time - 1) // 2
+
+        for ifreq in range(kernel_freq):
+            for itime in range(kernel_time):
+                df, dt = kernel_freq_half - ifreq, kernel_time_half - itime
+                x = x.roll(shifts=(df, dt), dims=(2, 3))
+
+                fslice = slice(max(0, df), min(n_freq, n_freq + df))
+                tslice = slice(max(0, dt), min(n_time, n_time + dt))
+
+                s[:, :, fslice, tslice] += (
+                    x[:, :, fslice, tslice] * m[ifreq, itime, :, :, fslice, tslice]
+                )
+
+        return s
+
+
+class MultiSourceMultiPatchingMaskBandSplitCoreRNN(PatchingMaskBandsplitCoreBase):
+    def __init__(
+        self,
+        in_channel: int,
+        stems: List[str],
+        band_specs: List[Tuple[float, float]],
+        mask_kernel_freq: int,
+        mask_kernel_time: int,
+        conv_kernel_freq: int,
+        conv_kernel_time: int,
+        kernel_norm_mlp_version: int,
+        require_no_overlap: bool = False,
+        require_no_gap: bool = True,
+        normalize_channel_independently: bool = False,
+        treat_channel_as_feature: bool = True,
+        n_sqm_modules: int = 12,
+        emb_dim: int = 128,
+        rnn_dim: int = 256,
+        bidirectional: bool = True,
+        rnn_type: str = "LSTM",
+        mlp_dim: int = 512,
+        hidden_activation: str = "Tanh",
+        hidden_activation_kwargs: Optional[Dict] = None,
+        complex_mask: bool = True,
+        overlapping_band: bool = False,
+        freq_weights: Optional[List[torch.Tensor]] = None,
+        n_freq: Optional[int] = None,
+    ) -> None:
+
+        super().__init__()
+        self.band_split = BandSplitModule(
+            in_channel=in_channel,
+            band_specs=band_specs,
+            require_no_overlap=require_no_overlap,
+            require_no_gap=require_no_gap,
+            normalize_channel_independently=normalize_channel_independently,
+            treat_channel_as_feature=treat_channel_as_feature,
+            emb_dim=emb_dim,
+        )
+
+        self.tf_model = SeqBandModellingModule(
+            n_modules=n_sqm_modules,
+            emb_dim=emb_dim,
+            rnn_dim=rnn_dim,
+            bidirectional=bidirectional,
+            rnn_type=rnn_type,
+        )
+
+        if hidden_activation_kwargs is None:
+            hidden_activation_kwargs = {}
+
+        if overlapping_band:
+            assert freq_weights is not None
+            assert n_freq is not None
+            self.mask_estim = nn.ModuleDict(
+                {
+                    stem: PatchingMaskEstimationModule(
+                        band_specs=band_specs,
+                        freq_weights=freq_weights,
+                        n_freq=n_freq,
+                        emb_dim=emb_dim,
+                        mlp_dim=mlp_dim,
+                        in_channel=in_channel,
+                        hidden_activation=hidden_activation,
+                        hidden_activation_kwargs=hidden_activation_kwargs,
+                        complex_mask=complex_mask,
+                        mask_kernel_freq=mask_kernel_freq,
+                        mask_kernel_time=mask_kernel_time,
+                        conv_kernel_freq=conv_kernel_freq,
+                        conv_kernel_time=conv_kernel_time,
+                        kernel_norm_mlp_version=kernel_norm_mlp_version,
+                    )
+                    for stem in stems
+                }
+            )
+        else:
+            raise NotImplementedError

mvsepless/models/bandit/core/model/bsrnn/maskestim.py

@@ -1,327 +1,327 @@
-import warnings
-from typing import Dict, List, Optional, Tuple, Type
-
-import torch
-from torch import nn
-from torch.nn.modules import activation
-
-from .utils import (
-    band_widths_from_specs,
-    check_no_gap,
-    check_no_overlap,
-    check_nonzero_bandwidth,
-)
-
-
-class BaseNormMLP(nn.Module):
-    def __init__(
-        self,
-        emb_dim: int,
-        mlp_dim: int,
-        bandwidth: int,
-        in_channel: Optional[int],
-        hidden_activation: str = "Tanh",
-        hidden_activation_kwargs=None,
-        complex_mask: bool = True,
-    ):
-
-        super().__init__()
-        if hidden_activation_kwargs is None:
-            hidden_activation_kwargs = {}
-        self.hidden_activation_kwargs = hidden_activation_kwargs
-        self.norm = nn.LayerNorm(emb_dim)
-        self.hidden = torch.jit.script(
-            nn.Sequential(
-                nn.Linear(in_features=emb_dim, out_features=mlp_dim),
-                activation.__dict__[hidden_activation](**self.hidden_activation_kwargs),
-            )
-        )
-
-        self.bandwidth = bandwidth
-        self.in_channel = in_channel
-
-        self.complex_mask = complex_mask
-        self.reim = 2 if complex_mask else 1
-        self.glu_mult = 2
-
-
-class NormMLP(BaseNormMLP):
-    def __init__(
-        self,
-        emb_dim: int,
-        mlp_dim: int,
-        bandwidth: int,
-        in_channel: Optional[int],
-        hidden_activation: str = "Tanh",
-        hidden_activation_kwargs=None,
-        complex_mask: bool = True,
-    ) -> None:
-        super().__init__(
-            emb_dim=emb_dim,
-            mlp_dim=mlp_dim,
-            bandwidth=bandwidth,
-            in_channel=in_channel,
-            hidden_activation=hidden_activation,
-            hidden_activation_kwargs=hidden_activation_kwargs,
-            complex_mask=complex_mask,
-        )
-
-        self.output = torch.jit.script(
-            nn.Sequential(
-                nn.Linear(
-                    in_features=mlp_dim,
-                    out_features=bandwidth * in_channel * self.reim * 2,
-                ),
-                nn.GLU(dim=-1),
-            )
-        )
-
-    def reshape_output(self, mb):
-        batch, n_time, _ = mb.shape
-        if self.complex_mask:
-            mb = mb.reshape(
-                batch, n_time, self.in_channel, self.bandwidth, self.reim
-            ).contiguous()
-            mb = torch.view_as_complex(mb)
-        else:
-            mb = mb.reshape(batch, n_time, self.in_channel, self.bandwidth)
-
-        mb = torch.permute(mb, (0, 2, 3, 1))
-
-        return mb
-
-    def forward(self, qb):
-
-        qb = self.norm(qb)
-
-        qb = self.hidden(qb)
-        mb = self.output(qb)
-        mb = self.reshape_output(mb)
-
-        return mb
-
-
-class MultAddNormMLP(NormMLP):
-    def __init__(
-        self,
-        emb_dim: int,
-        mlp_dim: int,
-        bandwidth: int,
-        in_channel: "int | None",
-        hidden_activation: str = "Tanh",
-        hidden_activation_kwargs=None,
-        complex_mask: bool = True,
-    ) -> None:
-        super().__init__(
-            emb_dim,
-            mlp_dim,
-            bandwidth,
-            in_channel,
-            hidden_activation,
-            hidden_activation_kwargs,
-            complex_mask,
-        )
-
-        self.output2 = torch.jit.script(
-            nn.Sequential(
-                nn.Linear(
-                    in_features=mlp_dim,
-                    out_features=bandwidth * in_channel * self.reim * 2,
-                ),
-                nn.GLU(dim=-1),
-            )
-        )
-
-    def forward(self, qb):
-
-        qb = self.norm(qb)
-        qb = self.hidden(qb)
-        mmb = self.output(qb)
-        mmb = self.reshape_output(mmb)
-        amb = self.output2(qb)
-        amb = self.reshape_output(amb)
-
-        return mmb, amb
-
-
-class MaskEstimationModuleSuperBase(nn.Module):
-    pass
-
-
-class MaskEstimationModuleBase(MaskEstimationModuleSuperBase):
-    def __init__(
-        self,
-        band_specs: List[Tuple[float, float]],
-        emb_dim: int,
-        mlp_dim: int,
-        in_channel: Optional[int],
-        hidden_activation: str = "Tanh",
-        hidden_activation_kwargs: Dict = None,
-        complex_mask: bool = True,
-        norm_mlp_cls: Type[nn.Module] = NormMLP,
-        norm_mlp_kwargs: Dict = None,
-    ) -> None:
-        super().__init__()
-
-        self.band_widths = band_widths_from_specs(band_specs)
-        self.n_bands = len(band_specs)
-
-        if hidden_activation_kwargs is None:
-            hidden_activation_kwargs = {}
-
-        if norm_mlp_kwargs is None:
-            norm_mlp_kwargs = {}
-
-        self.norm_mlp = nn.ModuleList(
-            [
-                (
-                    norm_mlp_cls(
-                        bandwidth=self.band_widths[b],
-                        emb_dim=emb_dim,
-                        mlp_dim=mlp_dim,
-                        in_channel=in_channel,
-                        hidden_activation=hidden_activation,
-                        hidden_activation_kwargs=hidden_activation_kwargs,
-                        complex_mask=complex_mask,
-                        **norm_mlp_kwargs,
-                    )
-                )
-                for b in range(self.n_bands)
-            ]
-        )
-
-    def compute_masks(self, q):
-        batch, n_bands, n_time, emb_dim = q.shape
-
-        masks = []
-
-        for b, nmlp in enumerate(self.norm_mlp):
-            qb = q[:, b, :, :]
-            mb = nmlp(qb)
-            masks.append(mb)
-
-        return masks
-
-
-class OverlappingMaskEstimationModule(MaskEstimationModuleBase):
-    def __init__(
-        self,
-        in_channel: int,
-        band_specs: List[Tuple[float, float]],
-        freq_weights: List[torch.Tensor],
-        n_freq: int,
-        emb_dim: int,
-        mlp_dim: int,
-        cond_dim: int = 0,
-        hidden_activation: str = "Tanh",
-        hidden_activation_kwargs: Dict = None,
-        complex_mask: bool = True,
-        norm_mlp_cls: Type[nn.Module] = NormMLP,
-        norm_mlp_kwargs: Dict = None,
-        use_freq_weights: bool = True,
-    ) -> None:
-        check_nonzero_bandwidth(band_specs)
-        check_no_gap(band_specs)
-
-        super().__init__(
-            band_specs=band_specs,
-            emb_dim=emb_dim + cond_dim,
-            mlp_dim=mlp_dim,
-            in_channel=in_channel,
-            hidden_activation=hidden_activation,
-            hidden_activation_kwargs=hidden_activation_kwargs,
-            complex_mask=complex_mask,
-            norm_mlp_cls=norm_mlp_cls,
-            norm_mlp_kwargs=norm_mlp_kwargs,
-        )
-
-        self.n_freq = n_freq
-        self.band_specs = band_specs
-        self.in_channel = in_channel
-
-        if freq_weights is not None:
-            for i, fw in enumerate(freq_weights):
-                self.register_buffer(f"freq_weights/{i}", fw)
-
-                self.use_freq_weights = use_freq_weights
-        else:
-            self.use_freq_weights = False
-
-        self.cond_dim = cond_dim
-
-    def forward(self, q, cond=None):
-
-        batch, n_bands, n_time, emb_dim = q.shape
-
-        if cond is not None:
-            print(cond)
-            if cond.ndim == 2:
-                cond = cond[:, None, None, :].expand(-1, n_bands, n_time, -1)
-            elif cond.ndim == 3:
-                assert cond.shape[1] == n_time
-            else:
-                raise ValueError(f"Invalid cond shape: {cond.shape}")
-
-            q = torch.cat([q, cond], dim=-1)
-        elif self.cond_dim > 0:
-            cond = torch.ones(
-                (batch, n_bands, n_time, self.cond_dim),
-                device=q.device,
-                dtype=q.dtype,
-            )
-            q = torch.cat([q, cond], dim=-1)
-        else:
-            pass
-
-        mask_list = self.compute_masks(q)
-
-        masks = torch.zeros(
-            (batch, self.in_channel, self.n_freq, n_time),
-            device=q.device,
-            dtype=mask_list[0].dtype,
-        )
-
-        for im, mask in enumerate(mask_list):
-            fstart, fend = self.band_specs[im]
-            if self.use_freq_weights:
-                fw = self.get_buffer(f"freq_weights/{im}")[:, None]
-                mask = mask * fw
-            masks[:, :, fstart:fend, :] += mask
-
-        return masks
-
-
-class MaskEstimationModule(OverlappingMaskEstimationModule):
-    def __init__(
-        self,
-        band_specs: List[Tuple[float, float]],
-        emb_dim: int,
-        mlp_dim: int,
-        in_channel: Optional[int],
-        hidden_activation: str = "Tanh",
-        hidden_activation_kwargs: Dict = None,
-        complex_mask: bool = True,
-        **kwargs,
-    ) -> None:
-        check_nonzero_bandwidth(band_specs)
-        check_no_gap(band_specs)
-        check_no_overlap(band_specs)
-        super().__init__(
-            in_channel=in_channel,
-            band_specs=band_specs,
-            freq_weights=None,
-            n_freq=None,
-            emb_dim=emb_dim,
-            mlp_dim=mlp_dim,
-            hidden_activation=hidden_activation,
-            hidden_activation_kwargs=hidden_activation_kwargs,
-            complex_mask=complex_mask,
-        )
-
-    def forward(self, q, cond=None):
-
-        masks = self.compute_masks(q)
-
-        masks = torch.concat(masks, dim=2)
-
-        return masks
+import warnings
+from typing import Dict, List, Optional, Tuple, Type
+
+import torch
+from torch import nn
+from torch.nn.modules import activation
+
+from .utils import (
+    band_widths_from_specs,
+    check_no_gap,
+    check_no_overlap,
+    check_nonzero_bandwidth,
+)
+
+
+class BaseNormMLP(nn.Module):
+    def __init__(
+        self,
+        emb_dim: int,
+        mlp_dim: int,
+        bandwidth: int,
+        in_channel: Optional[int],
+        hidden_activation: str = "Tanh",
+        hidden_activation_kwargs=None,
+        complex_mask: bool = True,
+    ):
+
+        super().__init__()
+        if hidden_activation_kwargs is None:
+            hidden_activation_kwargs = {}
+        self.hidden_activation_kwargs = hidden_activation_kwargs
+        self.norm = nn.LayerNorm(emb_dim)
+        self.hidden = torch.jit.script(
+            nn.Sequential(
+                nn.Linear(in_features=emb_dim, out_features=mlp_dim),
+                activation.__dict__[hidden_activation](**self.hidden_activation_kwargs),
+            )
+        )
+
+        self.bandwidth = bandwidth
+        self.in_channel = in_channel
+
+        self.complex_mask = complex_mask
+        self.reim = 2 if complex_mask else 1
+        self.glu_mult = 2
+
+
+class NormMLP(BaseNormMLP):
+    def __init__(
+        self,
+        emb_dim: int,
+        mlp_dim: int,
+        bandwidth: int,
+        in_channel: Optional[int],
+        hidden_activation: str = "Tanh",
+        hidden_activation_kwargs=None,
+        complex_mask: bool = True,
+    ) -> None:
+        super().__init__(
+            emb_dim=emb_dim,
+            mlp_dim=mlp_dim,
+            bandwidth=bandwidth,
+            in_channel=in_channel,
+            hidden_activation=hidden_activation,
+            hidden_activation_kwargs=hidden_activation_kwargs,
+            complex_mask=complex_mask,
+        )
+
+        self.output = torch.jit.script(
+            nn.Sequential(
+                nn.Linear(
+                    in_features=mlp_dim,
+                    out_features=bandwidth * in_channel * self.reim * 2,
+                ),
+                nn.GLU(dim=-1),
+            )
+        )
+
+    def reshape_output(self, mb):
+        batch, n_time, _ = mb.shape
+        if self.complex_mask:
+            mb = mb.reshape(
+                batch, n_time, self.in_channel, self.bandwidth, self.reim
+            ).contiguous()
+            mb = torch.view_as_complex(mb)
+        else:
+            mb = mb.reshape(batch, n_time, self.in_channel, self.bandwidth)
+
+        mb = torch.permute(mb, (0, 2, 3, 1))
+
+        return mb
+
+    def forward(self, qb):
+
+        qb = self.norm(qb)
+
+        qb = self.hidden(qb)
+        mb = self.output(qb)
+        mb = self.reshape_output(mb)
+
+        return mb
+
+
+class MultAddNormMLP(NormMLP):
+    def __init__(
+        self,
+        emb_dim: int,
+        mlp_dim: int,
+        bandwidth: int,
+        in_channel: "int | None",
+        hidden_activation: str = "Tanh",
+        hidden_activation_kwargs=None,
+        complex_mask: bool = True,
+    ) -> None:
+        super().__init__(
+            emb_dim,
+            mlp_dim,
+            bandwidth,
+            in_channel,
+            hidden_activation,
+            hidden_activation_kwargs,
+            complex_mask,
+        )
+
+        self.output2 = torch.jit.script(
+            nn.Sequential(
+                nn.Linear(
+                    in_features=mlp_dim,
+                    out_features=bandwidth * in_channel * self.reim * 2,
+                ),
+                nn.GLU(dim=-1),
+            )
+        )
+
+    def forward(self, qb):
+
+        qb = self.norm(qb)
+        qb = self.hidden(qb)
+        mmb = self.output(qb)
+        mmb = self.reshape_output(mmb)
+        amb = self.output2(qb)
+        amb = self.reshape_output(amb)
+
+        return mmb, amb
+
+
+class MaskEstimationModuleSuperBase(nn.Module):
+    pass
+
+
+class MaskEstimationModuleBase(MaskEstimationModuleSuperBase):
+    def __init__(
+        self,
+        band_specs: List[Tuple[float, float]],
+        emb_dim: int,
+        mlp_dim: int,
+        in_channel: Optional[int],
+        hidden_activation: str = "Tanh",
+        hidden_activation_kwargs: Dict = None,
+        complex_mask: bool = True,
+        norm_mlp_cls: Type[nn.Module] = NormMLP,
+        norm_mlp_kwargs: Dict = None,
+    ) -> None:
+        super().__init__()
+
+        self.band_widths = band_widths_from_specs(band_specs)
+        self.n_bands = len(band_specs)
+
+        if hidden_activation_kwargs is None:
+            hidden_activation_kwargs = {}
+
+        if norm_mlp_kwargs is None:
+            norm_mlp_kwargs = {}
+
+        self.norm_mlp = nn.ModuleList(
+            [
+                (
+                    norm_mlp_cls(
+                        bandwidth=self.band_widths[b],
+                        emb_dim=emb_dim,
+                        mlp_dim=mlp_dim,
+                        in_channel=in_channel,
+                        hidden_activation=hidden_activation,
+                        hidden_activation_kwargs=hidden_activation_kwargs,
+                        complex_mask=complex_mask,
+                        **norm_mlp_kwargs,
+                    )
+                )
+                for b in range(self.n_bands)
+            ]
+        )
+
+    def compute_masks(self, q):
+        batch, n_bands, n_time, emb_dim = q.shape
+
+        masks = []
+
+        for b, nmlp in enumerate(self.norm_mlp):
+            qb = q[:, b, :, :]
+            mb = nmlp(qb)
+            masks.append(mb)
+
+        return masks
+
+
+class OverlappingMaskEstimationModule(MaskEstimationModuleBase):
+    def __init__(
+        self,
+        in_channel: int,
+        band_specs: List[Tuple[float, float]],
+        freq_weights: List[torch.Tensor],
+        n_freq: int,
+        emb_dim: int,
+        mlp_dim: int,
+        cond_dim: int = 0,
+        hidden_activation: str = "Tanh",
+        hidden_activation_kwargs: Dict = None,
+        complex_mask: bool = True,
+        norm_mlp_cls: Type[nn.Module] = NormMLP,
+        norm_mlp_kwargs: Dict = None,
+        use_freq_weights: bool = True,
+    ) -> None:
+        check_nonzero_bandwidth(band_specs)
+        check_no_gap(band_specs)
+
+        super().__init__(
+            band_specs=band_specs,
+            emb_dim=emb_dim + cond_dim,
+            mlp_dim=mlp_dim,
+            in_channel=in_channel,
+            hidden_activation=hidden_activation,
+            hidden_activation_kwargs=hidden_activation_kwargs,
+            complex_mask=complex_mask,
+            norm_mlp_cls=norm_mlp_cls,
+            norm_mlp_kwargs=norm_mlp_kwargs,
+        )
+
+        self.n_freq = n_freq
+        self.band_specs = band_specs
+        self.in_channel = in_channel
+
+        if freq_weights is not None:
+            for i, fw in enumerate(freq_weights):
+                self.register_buffer(f"freq_weights/{i}", fw)
+
+                self.use_freq_weights = use_freq_weights
+        else:
+            self.use_freq_weights = False
+
+        self.cond_dim = cond_dim
+
+    def forward(self, q, cond=None):
+
+        batch, n_bands, n_time, emb_dim = q.shape
+
+        if cond is not None:
+            print(cond)
+            if cond.ndim == 2:
+                cond = cond[:, None, None, :].expand(-1, n_bands, n_time, -1)
+            elif cond.ndim == 3:
+                assert cond.shape[1] == n_time
+            else:
+                raise ValueError(f"Invalid cond shape: {cond.shape}")
+
+            q = torch.cat([q, cond], dim=-1)
+        elif self.cond_dim > 0:
+            cond = torch.ones(
+                (batch, n_bands, n_time, self.cond_dim),
+                device=q.device,
+                dtype=q.dtype,
+            )
+            q = torch.cat([q, cond], dim=-1)
+        else:
+            pass
+
+        mask_list = self.compute_masks(q)
+
+        masks = torch.zeros(
+            (batch, self.in_channel, self.n_freq, n_time),
+            device=q.device,
+            dtype=mask_list[0].dtype,
+        )
+
+        for im, mask in enumerate(mask_list):
+            fstart, fend = self.band_specs[im]
+            if self.use_freq_weights:
+                fw = self.get_buffer(f"freq_weights/{im}")[:, None]
+                mask = mask * fw
+            masks[:, :, fstart:fend, :] += mask
+
+        return masks
+
+
+class MaskEstimationModule(OverlappingMaskEstimationModule):
+    def __init__(
+        self,
+        band_specs: List[Tuple[float, float]],
+        emb_dim: int,
+        mlp_dim: int,
+        in_channel: Optional[int],
+        hidden_activation: str = "Tanh",
+        hidden_activation_kwargs: Dict = None,
+        complex_mask: bool = True,
+        **kwargs,
+    ) -> None:
+        check_nonzero_bandwidth(band_specs)
+        check_no_gap(band_specs)
+        check_no_overlap(band_specs)
+        super().__init__(
+            in_channel=in_channel,
+            band_specs=band_specs,
+            freq_weights=None,
+            n_freq=None,
+            emb_dim=emb_dim,
+            mlp_dim=mlp_dim,
+            hidden_activation=hidden_activation,
+            hidden_activation_kwargs=hidden_activation_kwargs,
+            complex_mask=complex_mask,
+        )
+
+    def forward(self, q, cond=None):
+
+        masks = self.compute_masks(q)
+
+        masks = torch.concat(masks, dim=2)
+
+        return masks

mvsepless/models/bandit/core/model/bsrnn/tfmodel.py

@@ -1,287 +1,287 @@
-import warnings
-
-import torch
-from torch import nn
-from torch.nn import functional as F
-from torch.nn.modules import rnn
-
-import torch.backends.cuda
-
-
-class TimeFrequencyModellingModule(nn.Module):
-    def __init__(self) -> None:
-        super().__init__()
-
-
-class ResidualRNN(nn.Module):
-    def __init__(
-        self,
-        emb_dim: int,
-        rnn_dim: int,
-        bidirectional: bool = True,
-        rnn_type: str = "LSTM",
-        use_batch_trick: bool = True,
-        use_layer_norm: bool = True,
-    ) -> None:
-        super().__init__()
-
-        self.use_layer_norm = use_layer_norm
-        if use_layer_norm:
-            self.norm = nn.LayerNorm(emb_dim)
-        else:
-            self.norm = nn.GroupNorm(num_groups=emb_dim, num_channels=emb_dim)
-
-        self.rnn = rnn.__dict__[rnn_type](
-            input_size=emb_dim,
-            hidden_size=rnn_dim,
-            num_layers=1,
-            batch_first=True,
-            bidirectional=bidirectional,
-        )
-
-        self.fc = nn.Linear(
-            in_features=rnn_dim * (2 if bidirectional else 1), out_features=emb_dim
-        )
-
-        self.use_batch_trick = use_batch_trick
-        if not self.use_batch_trick:
-            warnings.warn("NOT USING BATCH TRICK IS EXTREMELY SLOW!!")
-
-    def forward(self, z):
-
-        z0 = torch.clone(z)
-
-        if self.use_layer_norm:
-            z = self.norm(z)
-        else:
-            z = torch.permute(z, (0, 3, 1, 2))
-
-            z = self.norm(z)
-
-            z = torch.permute(z, (0, 2, 3, 1))
-
-        batch, n_uncrossed, n_across, emb_dim = z.shape
-
-        if self.use_batch_trick:
-            z = torch.reshape(z, (batch * n_uncrossed, n_across, emb_dim))
-
-            z = self.rnn(z.contiguous())[0]
-
-            z = torch.reshape(z, (batch, n_uncrossed, n_across, -1))
-        else:
-            zlist = []
-            for i in range(n_uncrossed):
-                zi = self.rnn(z[:, i, :, :])[0]
-                zlist.append(zi)
-
-            z = torch.stack(zlist, dim=1)
-
-        z = self.fc(z)
-
-        z = z + z0
-
-        return z
-
-
-class SeqBandModellingModule(TimeFrequencyModellingModule):
-    def __init__(
-        self,
-        n_modules: int = 12,
-        emb_dim: int = 128,
-        rnn_dim: int = 256,
-        bidirectional: bool = True,
-        rnn_type: str = "LSTM",
-        parallel_mode=False,
-    ) -> None:
-        super().__init__()
-        self.seqband = nn.ModuleList([])
-
-        if parallel_mode:
-            for _ in range(n_modules):
-                self.seqband.append(
-                    nn.ModuleList(
-                        [
-                            ResidualRNN(
-                                emb_dim=emb_dim,
-                                rnn_dim=rnn_dim,
-                                bidirectional=bidirectional,
-                                rnn_type=rnn_type,
-                            ),
-                            ResidualRNN(
-                                emb_dim=emb_dim,
-                                rnn_dim=rnn_dim,
-                                bidirectional=bidirectional,
-                                rnn_type=rnn_type,
-                            ),
-                        ]
-                    )
-                )
-        else:
-
-            for _ in range(2 * n_modules):
-                self.seqband.append(
-                    ResidualRNN(
-                        emb_dim=emb_dim,
-                        rnn_dim=rnn_dim,
-                        bidirectional=bidirectional,
-                        rnn_type=rnn_type,
-                    )
-                )
-
-        self.parallel_mode = parallel_mode
-
-    def forward(self, z):
-
-        if self.parallel_mode:
-            for sbm_pair in self.seqband:
-                sbm_t, sbm_f = sbm_pair[0], sbm_pair[1]
-                zt = sbm_t(z)
-                zf = sbm_f(z.transpose(1, 2))
-                z = zt + zf.transpose(1, 2)
-        else:
-            for sbm in self.seqband:
-                z = sbm(z)
-                z = z.transpose(1, 2)
-
-        q = z
-        return q
-
-
-class ResidualTransformer(nn.Module):
-    def __init__(
-        self,
-        emb_dim: int = 128,
-        rnn_dim: int = 256,
-        bidirectional: bool = True,
-        dropout: float = 0.0,
-    ) -> None:
-        super().__init__()
-
-        self.tf = nn.TransformerEncoderLayer(
-            d_model=emb_dim, nhead=4, dim_feedforward=rnn_dim, batch_first=True
-        )
-
-        self.is_causal = not bidirectional
-        self.dropout = dropout
-
-    def forward(self, z):
-        batch, n_uncrossed, n_across, emb_dim = z.shape
-        z = torch.reshape(z, (batch * n_uncrossed, n_across, emb_dim))
-        z = self.tf(z, is_causal=self.is_causal)
-        z = torch.reshape(z, (batch, n_uncrossed, n_across, emb_dim))
-
-        return z
-
-
-class TransformerTimeFreqModule(TimeFrequencyModellingModule):
-    def __init__(
-        self,
-        n_modules: int = 12,
-        emb_dim: int = 128,
-        rnn_dim: int = 256,
-        bidirectional: bool = True,
-        dropout: float = 0.0,
-    ) -> None:
-        super().__init__()
-        self.norm = nn.LayerNorm(emb_dim)
-        self.seqband = nn.ModuleList([])
-
-        for _ in range(2 * n_modules):
-            self.seqband.append(
-                ResidualTransformer(
-                    emb_dim=emb_dim,
-                    rnn_dim=rnn_dim,
-                    bidirectional=bidirectional,
-                    dropout=dropout,
-                )
-            )
-
-    def forward(self, z):
-        z = self.norm(z)
-
-        for sbm in self.seqband:
-            z = sbm(z)
-            z = z.transpose(1, 2)
-
-        q = z
-        return q
-
-
-class ResidualConvolution(nn.Module):
-    def __init__(
-        self,
-        emb_dim: int = 128,
-        rnn_dim: int = 256,
-        bidirectional: bool = True,
-        dropout: float = 0.0,
-    ) -> None:
-        super().__init__()
-        self.norm = nn.InstanceNorm2d(emb_dim, affine=True)
-
-        self.conv = nn.Sequential(
-            nn.Conv2d(
-                in_channels=emb_dim,
-                out_channels=rnn_dim,
-                kernel_size=(3, 3),
-                padding="same",
-                stride=(1, 1),
-            ),
-            nn.Tanhshrink(),
-        )
-
-        self.is_causal = not bidirectional
-        self.dropout = dropout
-
-        self.fc = nn.Conv2d(
-            in_channels=rnn_dim,
-            out_channels=emb_dim,
-            kernel_size=(1, 1),
-            padding="same",
-            stride=(1, 1),
-        )
-
-    def forward(self, z):
-
-        z0 = torch.clone(z)
-
-        z = self.norm(z)
-        z = self.conv(z)
-        z = self.fc(z)
-        z = z + z0
-
-        return z
-
-
-class ConvolutionalTimeFreqModule(TimeFrequencyModellingModule):
-    def __init__(
-        self,
-        n_modules: int = 12,
-        emb_dim: int = 128,
-        rnn_dim: int = 256,
-        bidirectional: bool = True,
-        dropout: float = 0.0,
-    ) -> None:
-        super().__init__()
-        self.seqband = torch.jit.script(
-            nn.Sequential(
-                *[
-                    ResidualConvolution(
-                        emb_dim=emb_dim,
-                        rnn_dim=rnn_dim,
-                        bidirectional=bidirectional,
-                        dropout=dropout,
-                    )
-                    for _ in range(2 * n_modules)
-                ]
-            )
-        )
-
-    def forward(self, z):
-
-        z = torch.permute(z, (0, 3, 1, 2))
-
-        z = self.seqband(z)
-
-        z = torch.permute(z, (0, 2, 3, 1))
-
-        return z
+import warnings
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from torch.nn.modules import rnn
+
+import torch.backends.cuda
+
+
+class TimeFrequencyModellingModule(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+
+
+class ResidualRNN(nn.Module):
+    def __init__(
+        self,
+        emb_dim: int,
+        rnn_dim: int,
+        bidirectional: bool = True,
+        rnn_type: str = "LSTM",
+        use_batch_trick: bool = True,
+        use_layer_norm: bool = True,
+    ) -> None:
+        super().__init__()
+
+        self.use_layer_norm = use_layer_norm
+        if use_layer_norm:
+            self.norm = nn.LayerNorm(emb_dim)
+        else:
+            self.norm = nn.GroupNorm(num_groups=emb_dim, num_channels=emb_dim)
+
+        self.rnn = rnn.__dict__[rnn_type](
+            input_size=emb_dim,
+            hidden_size=rnn_dim,
+            num_layers=1,
+            batch_first=True,
+            bidirectional=bidirectional,
+        )
+
+        self.fc = nn.Linear(
+            in_features=rnn_dim * (2 if bidirectional else 1), out_features=emb_dim
+        )
+
+        self.use_batch_trick = use_batch_trick
+        if not self.use_batch_trick:
+            warnings.warn("NOT USING BATCH TRICK IS EXTREMELY SLOW!!")
+
+    def forward(self, z):
+
+        z0 = torch.clone(z)
+
+        if self.use_layer_norm:
+            z = self.norm(z)
+        else:
+            z = torch.permute(z, (0, 3, 1, 2))
+
+            z = self.norm(z)
+
+            z = torch.permute(z, (0, 2, 3, 1))
+
+        batch, n_uncrossed, n_across, emb_dim = z.shape
+
+        if self.use_batch_trick:
+            z = torch.reshape(z, (batch * n_uncrossed, n_across, emb_dim))
+
+            z = self.rnn(z.contiguous())[0]
+
+            z = torch.reshape(z, (batch, n_uncrossed, n_across, -1))
+        else:
+            zlist = []
+            for i in range(n_uncrossed):
+                zi = self.rnn(z[:, i, :, :])[0]
+                zlist.append(zi)
+
+            z = torch.stack(zlist, dim=1)
+
+        z = self.fc(z)
+
+        z = z + z0
+
+        return z
+
+
+class SeqBandModellingModule(TimeFrequencyModellingModule):
+    def __init__(
+        self,
+        n_modules: int = 12,
+        emb_dim: int = 128,
+        rnn_dim: int = 256,
+        bidirectional: bool = True,
+        rnn_type: str = "LSTM",
+        parallel_mode=False,
+    ) -> None:
+        super().__init__()
+        self.seqband = nn.ModuleList([])
+
+        if parallel_mode:
+            for _ in range(n_modules):
+                self.seqband.append(
+                    nn.ModuleList(
+                        [
+                            ResidualRNN(
+                                emb_dim=emb_dim,
+                                rnn_dim=rnn_dim,
+                                bidirectional=bidirectional,
+                                rnn_type=rnn_type,
+                            ),
+                            ResidualRNN(
+                                emb_dim=emb_dim,
+                                rnn_dim=rnn_dim,
+                                bidirectional=bidirectional,
+                                rnn_type=rnn_type,
+                            ),
+                        ]
+                    )
+                )
+        else:
+
+            for _ in range(2 * n_modules):
+                self.seqband.append(
+                    ResidualRNN(
+                        emb_dim=emb_dim,
+                        rnn_dim=rnn_dim,
+                        bidirectional=bidirectional,
+                        rnn_type=rnn_type,
+                    )
+                )
+
+        self.parallel_mode = parallel_mode
+
+    def forward(self, z):
+
+        if self.parallel_mode:
+            for sbm_pair in self.seqband:
+                sbm_t, sbm_f = sbm_pair[0], sbm_pair[1]
+                zt = sbm_t(z)
+                zf = sbm_f(z.transpose(1, 2))
+                z = zt + zf.transpose(1, 2)
+        else:
+            for sbm in self.seqband:
+                z = sbm(z)
+                z = z.transpose(1, 2)
+
+        q = z
+        return q
+
+
+class ResidualTransformer(nn.Module):
+    def __init__(
+        self,
+        emb_dim: int = 128,
+        rnn_dim: int = 256,
+        bidirectional: bool = True,
+        dropout: float = 0.0,
+    ) -> None:
+        super().__init__()
+
+        self.tf = nn.TransformerEncoderLayer(
+            d_model=emb_dim, nhead=4, dim_feedforward=rnn_dim, batch_first=True
+        )
+
+        self.is_causal = not bidirectional
+        self.dropout = dropout
+
+    def forward(self, z):
+        batch, n_uncrossed, n_across, emb_dim = z.shape
+        z = torch.reshape(z, (batch * n_uncrossed, n_across, emb_dim))
+        z = self.tf(z, is_causal=self.is_causal)
+        z = torch.reshape(z, (batch, n_uncrossed, n_across, emb_dim))
+
+        return z
+
+
+class TransformerTimeFreqModule(TimeFrequencyModellingModule):
+    def __init__(
+        self,
+        n_modules: int = 12,
+        emb_dim: int = 128,
+        rnn_dim: int = 256,
+        bidirectional: bool = True,
+        dropout: float = 0.0,
+    ) -> None:
+        super().__init__()
+        self.norm = nn.LayerNorm(emb_dim)
+        self.seqband = nn.ModuleList([])
+
+        for _ in range(2 * n_modules):
+            self.seqband.append(
+                ResidualTransformer(
+                    emb_dim=emb_dim,
+                    rnn_dim=rnn_dim,
+                    bidirectional=bidirectional,
+                    dropout=dropout,
+                )
+            )
+
+    def forward(self, z):
+        z = self.norm(z)
+
+        for sbm in self.seqband:
+            z = sbm(z)
+            z = z.transpose(1, 2)
+
+        q = z
+        return q
+
+
+class ResidualConvolution(nn.Module):
+    def __init__(
+        self,
+        emb_dim: int = 128,
+        rnn_dim: int = 256,
+        bidirectional: bool = True,
+        dropout: float = 0.0,
+    ) -> None:
+        super().__init__()
+        self.norm = nn.InstanceNorm2d(emb_dim, affine=True)
+
+        self.conv = nn.Sequential(
+            nn.Conv2d(
+                in_channels=emb_dim,
+                out_channels=rnn_dim,
+                kernel_size=(3, 3),
+                padding="same",
+                stride=(1, 1),
+            ),
+            nn.Tanhshrink(),
+        )
+
+        self.is_causal = not bidirectional
+        self.dropout = dropout
+
+        self.fc = nn.Conv2d(
+            in_channels=rnn_dim,
+            out_channels=emb_dim,
+            kernel_size=(1, 1),
+            padding="same",
+            stride=(1, 1),
+        )
+
+    def forward(self, z):
+
+        z0 = torch.clone(z)
+
+        z = self.norm(z)
+        z = self.conv(z)
+        z = self.fc(z)
+        z = z + z0
+
+        return z
+
+
+class ConvolutionalTimeFreqModule(TimeFrequencyModellingModule):
+    def __init__(
+        self,
+        n_modules: int = 12,
+        emb_dim: int = 128,
+        rnn_dim: int = 256,
+        bidirectional: bool = True,
+        dropout: float = 0.0,
+    ) -> None:
+        super().__init__()
+        self.seqband = torch.jit.script(
+            nn.Sequential(
+                *[
+                    ResidualConvolution(
+                        emb_dim=emb_dim,
+                        rnn_dim=rnn_dim,
+                        bidirectional=bidirectional,
+                        dropout=dropout,
+                    )
+                    for _ in range(2 * n_modules)
+                ]
+            )
+        )
+
+    def forward(self, z):
+
+        z = torch.permute(z, (0, 3, 1, 2))
+
+        z = self.seqband(z)
+
+        z = torch.permute(z, (0, 2, 3, 1))
+
+        return z

mvsepless/models/bandit/core/model/bsrnn/utils.py

@@ -1,518 +1,518 @@
-import os
-from abc import abstractmethod
-from typing import Any, Callable
-
-import numpy as np
-import torch
-from librosa import hz_to_midi, midi_to_hz
-from torch import Tensor
-from torchaudio import functional as taF
-from spafe.fbanks import bark_fbanks
-from spafe.utils.converters import erb2hz, hz2bark, hz2erb
-from torchaudio.functional.functional import _create_triangular_filterbank
-
-
-def band_widths_from_specs(band_specs):
-    return [e - i for i, e in band_specs]
-
-
-def check_nonzero_bandwidth(band_specs):
-    for fstart, fend in band_specs:
-        if fend - fstart <= 0:
-            raise ValueError("Bands cannot be zero-width")
-
-
-def check_no_overlap(band_specs):
-    fend_prev = -1
-    for fstart_curr, fend_curr in band_specs:
-        if fstart_curr <= fend_prev:
-            raise ValueError("Bands cannot overlap")
-
-
-def check_no_gap(band_specs):
-    fstart, _ = band_specs[0]
-    assert fstart == 0
-
-    fend_prev = -1
-    for fstart_curr, fend_curr in band_specs:
-        if fstart_curr - fend_prev > 1:
-            raise ValueError("Bands cannot leave gap")
-        fend_prev = fend_curr
-
-
-class BandsplitSpecification:
-    def __init__(self, nfft: int, fs: int) -> None:
-        self.fs = fs
-        self.nfft = nfft
-        self.nyquist = fs / 2
-        self.max_index = nfft // 2 + 1
-
-        self.split500 = self.hertz_to_index(500)
-        self.split1k = self.hertz_to_index(1000)
-        self.split2k = self.hertz_to_index(2000)
-        self.split4k = self.hertz_to_index(4000)
-        self.split8k = self.hertz_to_index(8000)
-        self.split16k = self.hertz_to_index(16000)
-        self.split20k = self.hertz_to_index(20000)
-
-        self.above20k = [(self.split20k, self.max_index)]
-        self.above16k = [(self.split16k, self.split20k)] + self.above20k
-
-    def index_to_hertz(self, index: int):
-        return index * self.fs / self.nfft
-
-    def hertz_to_index(self, hz: float, round: bool = True):
-        index = hz * self.nfft / self.fs
-
-        if round:
-            index = int(np.round(index))
-
-        return index
-
-    def get_band_specs_with_bandwidth(self, start_index, end_index, bandwidth_hz):
-        band_specs = []
-        lower = start_index
-
-        while lower < end_index:
-            upper = int(np.floor(lower + self.hertz_to_index(bandwidth_hz)))
-            upper = min(upper, end_index)
-
-            band_specs.append((lower, upper))
-            lower = upper
-
-        return band_specs
-
-    @abstractmethod
-    def get_band_specs(self):
-        raise NotImplementedError
-
-
-class VocalBandsplitSpecification(BandsplitSpecification):
-    def __init__(self, nfft: int, fs: int, version: str = "7") -> None:
-        super().__init__(nfft=nfft, fs=fs)
-
-        self.version = version
-
-    def get_band_specs(self):
-        return getattr(self, f"version{self.version}")()
-
-    @property
-    def version1(self):
-        return self.get_band_specs_with_bandwidth(
-            start_index=0, end_index=self.max_index, bandwidth_hz=1000
-        )
-
-    def version2(self):
-        below16k = self.get_band_specs_with_bandwidth(
-            start_index=0, end_index=self.split16k, bandwidth_hz=1000
-        )
-        below20k = self.get_band_specs_with_bandwidth(
-            start_index=self.split16k, end_index=self.split20k, bandwidth_hz=2000
-        )
-
-        return below16k + below20k + self.above20k
-
-    def version3(self):
-        below8k = self.get_band_specs_with_bandwidth(
-            start_index=0, end_index=self.split8k, bandwidth_hz=1000
-        )
-        below16k = self.get_band_specs_with_bandwidth(
-            start_index=self.split8k, end_index=self.split16k, bandwidth_hz=2000
-        )
-
-        return below8k + below16k + self.above16k
-
-    def version4(self):
-        below1k = self.get_band_specs_with_bandwidth(
-            start_index=0, end_index=self.split1k, bandwidth_hz=100
-        )
-        below8k = self.get_band_specs_with_bandwidth(
-            start_index=self.split1k, end_index=self.split8k, bandwidth_hz=1000
-        )
-        below16k = self.get_band_specs_with_bandwidth(
-            start_index=self.split8k, end_index=self.split16k, bandwidth_hz=2000
-        )
-
-        return below1k + below8k + below16k + self.above16k
-
-    def version5(self):
-        below1k = self.get_band_specs_with_bandwidth(
-            start_index=0, end_index=self.split1k, bandwidth_hz=100
-        )
-        below16k = self.get_band_specs_with_bandwidth(
-            start_index=self.split1k, end_index=self.split16k, bandwidth_hz=1000
-        )
-        below20k = self.get_band_specs_with_bandwidth(
-            start_index=self.split16k, end_index=self.split20k, bandwidth_hz=2000
-        )
-        return below1k + below16k + below20k + self.above20k
-
-    def version6(self):
-        below1k = self.get_band_specs_with_bandwidth(
-            start_index=0, end_index=self.split1k, bandwidth_hz=100
-        )
-        below4k = self.get_band_specs_with_bandwidth(
-            start_index=self.split1k, end_index=self.split4k, bandwidth_hz=500
-        )
-        below8k = self.get_band_specs_with_bandwidth(
-            start_index=self.split4k, end_index=self.split8k, bandwidth_hz=1000
-        )
-        below16k = self.get_band_specs_with_bandwidth(
-            start_index=self.split8k, end_index=self.split16k, bandwidth_hz=2000
-        )
-        return below1k + below4k + below8k + below16k + self.above16k
-
-    def version7(self):
-        below1k = self.get_band_specs_with_bandwidth(
-            start_index=0, end_index=self.split1k, bandwidth_hz=100
-        )
-        below4k = self.get_band_specs_with_bandwidth(
-            start_index=self.split1k, end_index=self.split4k, bandwidth_hz=250
-        )
-        below8k = self.get_band_specs_with_bandwidth(
-            start_index=self.split4k, end_index=self.split8k, bandwidth_hz=500
-        )
-        below16k = self.get_band_specs_with_bandwidth(
-            start_index=self.split8k, end_index=self.split16k, bandwidth_hz=1000
-        )
-        below20k = self.get_band_specs_with_bandwidth(
-            start_index=self.split16k, end_index=self.split20k, bandwidth_hz=2000
-        )
-        return below1k + below4k + below8k + below16k + below20k + self.above20k
-
-
-class OtherBandsplitSpecification(VocalBandsplitSpecification):
-    def __init__(self, nfft: int, fs: int) -> None:
-        super().__init__(nfft=nfft, fs=fs, version="7")
-
-
-class BassBandsplitSpecification(BandsplitSpecification):
-    def __init__(self, nfft: int, fs: int, version: str = "7") -> None:
-        super().__init__(nfft=nfft, fs=fs)
-
-    def get_band_specs(self):
-        below500 = self.get_band_specs_with_bandwidth(
-            start_index=0, end_index=self.split500, bandwidth_hz=50
-        )
-        below1k = self.get_band_specs_with_bandwidth(
-            start_index=self.split500, end_index=self.split1k, bandwidth_hz=100
-        )
-        below4k = self.get_band_specs_with_bandwidth(
-            start_index=self.split1k, end_index=self.split4k, bandwidth_hz=500
-        )
-        below8k = self.get_band_specs_with_bandwidth(
-            start_index=self.split4k, end_index=self.split8k, bandwidth_hz=1000
-        )
-        below16k = self.get_band_specs_with_bandwidth(
-            start_index=self.split8k, end_index=self.split16k, bandwidth_hz=2000
-        )
-        above16k = [(self.split16k, self.max_index)]
-
-        return below500 + below1k + below4k + below8k + below16k + above16k
-
-
-class DrumBandsplitSpecification(BandsplitSpecification):
-    def __init__(self, nfft: int, fs: int) -> None:
-        super().__init__(nfft=nfft, fs=fs)
-
-    def get_band_specs(self):
-        below1k = self.get_band_specs_with_bandwidth(
-            start_index=0, end_index=self.split1k, bandwidth_hz=50
-        )
-        below2k = self.get_band_specs_with_bandwidth(
-            start_index=self.split1k, end_index=self.split2k, bandwidth_hz=100
-        )
-        below4k = self.get_band_specs_with_bandwidth(
-            start_index=self.split2k, end_index=self.split4k, bandwidth_hz=250
-        )
-        below8k = self.get_band_specs_with_bandwidth(
-            start_index=self.split4k, end_index=self.split8k, bandwidth_hz=500
-        )
-        below16k = self.get_band_specs_with_bandwidth(
-            start_index=self.split8k, end_index=self.split16k, bandwidth_hz=1000
-        )
-        above16k = [(self.split16k, self.max_index)]
-
-        return below1k + below2k + below4k + below8k + below16k + above16k
-
-
-class PerceptualBandsplitSpecification(BandsplitSpecification):
-    def __init__(
-        self,
-        nfft: int,
-        fs: int,
-        fbank_fn: Callable[[int, int, float, float, int], torch.Tensor],
-        n_bands: int,
-        f_min: float = 0.0,
-        f_max: float = None,
-    ) -> None:
-        super().__init__(nfft=nfft, fs=fs)
-        self.n_bands = n_bands
-        if f_max is None:
-            f_max = fs / 2
-
-        self.filterbank = fbank_fn(n_bands, fs, f_min, f_max, self.max_index)
-
-        weight_per_bin = torch.sum(self.filterbank, dim=0, keepdim=True)
-        normalized_mel_fb = self.filterbank / weight_per_bin
-
-        freq_weights = []
-        band_specs = []
-        for i in range(self.n_bands):
-            active_bins = torch.nonzero(self.filterbank[i, :]).squeeze().tolist()
-            if isinstance(active_bins, int):
-                active_bins = (active_bins, active_bins)
-            if len(active_bins) == 0:
-                continue
-            start_index = active_bins[0]
-            end_index = active_bins[-1] + 1
-            band_specs.append((start_index, end_index))
-            freq_weights.append(normalized_mel_fb[i, start_index:end_index])
-
-        self.freq_weights = freq_weights
-        self.band_specs = band_specs
-
-    def get_band_specs(self):
-        return self.band_specs
-
-    def get_freq_weights(self):
-        return self.freq_weights
-
-    def save_to_file(self, dir_path: str) -> None:
-
-        os.makedirs(dir_path, exist_ok=True)
-
-        import pickle
-
-        with open(os.path.join(dir_path, "mel_bandsplit_spec.pkl"), "wb") as f:
-            pickle.dump(
-                {
-                    "band_specs": self.band_specs,
-                    "freq_weights": self.freq_weights,
-                    "filterbank": self.filterbank,
-                },
-                f,
-            )
-
-
-def mel_filterbank(n_bands, fs, f_min, f_max, n_freqs):
-    fb = taF.melscale_fbanks(
-        n_mels=n_bands,
-        sample_rate=fs,
-        f_min=f_min,
-        f_max=f_max,
-        n_freqs=n_freqs,
-    ).T
-
-    fb[0, 0] = 1.0
-
-    return fb
-
-
-class MelBandsplitSpecification(PerceptualBandsplitSpecification):
-    def __init__(
-        self, nfft: int, fs: int, n_bands: int, f_min: float = 0.0, f_max: float = None
-    ) -> None:
-        super().__init__(
-            fbank_fn=mel_filterbank,
-            nfft=nfft,
-            fs=fs,
-            n_bands=n_bands,
-            f_min=f_min,
-            f_max=f_max,
-        )
-
-
-def musical_filterbank(n_bands, fs, f_min, f_max, n_freqs, scale="constant"):
-
-    nfft = 2 * (n_freqs - 1)
-    df = fs / nfft
-    f_max = f_max or fs / 2
-    f_min = f_min or 0
-    f_min = fs / nfft
-
-    n_octaves = np.log2(f_max / f_min)
-    n_octaves_per_band = n_octaves / n_bands
-    bandwidth_mult = np.power(2.0, n_octaves_per_band)
-
-    low_midi = max(0, hz_to_midi(f_min))
-    high_midi = hz_to_midi(f_max)
-    midi_points = np.linspace(low_midi, high_midi, n_bands)
-    hz_pts = midi_to_hz(midi_points)
-
-    low_pts = hz_pts / bandwidth_mult
-    high_pts = hz_pts * bandwidth_mult
-
-    low_bins = np.floor(low_pts / df).astype(int)
-    high_bins = np.ceil(high_pts / df).astype(int)
-
-    fb = np.zeros((n_bands, n_freqs))
-
-    for i in range(n_bands):
-        fb[i, low_bins[i] : high_bins[i] + 1] = 1.0
-
-    fb[0, : low_bins[0]] = 1.0
-    fb[-1, high_bins[-1] + 1 :] = 1.0
-
-    return torch.as_tensor(fb)
-
-
-class MusicalBandsplitSpecification(PerceptualBandsplitSpecification):
-    def __init__(
-        self, nfft: int, fs: int, n_bands: int, f_min: float = 0.0, f_max: float = None
-    ) -> None:
-        super().__init__(
-            fbank_fn=musical_filterbank,
-            nfft=nfft,
-            fs=fs,
-            n_bands=n_bands,
-            f_min=f_min,
-            f_max=f_max,
-        )
-
-
-def bark_filterbank(n_bands, fs, f_min, f_max, n_freqs):
-    nfft = 2 * (n_freqs - 1)
-    fb, _ = bark_fbanks.bark_filter_banks(
-        nfilts=n_bands,
-        nfft=nfft,
-        fs=fs,
-        low_freq=f_min,
-        high_freq=f_max,
-        scale="constant",
-    )
-
-    return torch.as_tensor(fb)
-
-
-class BarkBandsplitSpecification(PerceptualBandsplitSpecification):
-    def __init__(
-        self, nfft: int, fs: int, n_bands: int, f_min: float = 0.0, f_max: float = None
-    ) -> None:
-        super().__init__(
-            fbank_fn=bark_filterbank,
-            nfft=nfft,
-            fs=fs,
-            n_bands=n_bands,
-            f_min=f_min,
-            f_max=f_max,
-        )
-
-
-def triangular_bark_filterbank(n_bands, fs, f_min, f_max, n_freqs):
-
-    all_freqs = torch.linspace(0, fs // 2, n_freqs)
-
-    m_min = hz2bark(f_min)
-    m_max = hz2bark(f_max)
-
-    m_pts = torch.linspace(m_min, m_max, n_bands + 2)
-    f_pts = 600 * torch.sinh(m_pts / 6)
-
-    fb = _create_triangular_filterbank(all_freqs, f_pts)
-
-    fb = fb.T
-
-    first_active_band = torch.nonzero(torch.sum(fb, dim=-1))[0, 0]
-    first_active_bin = torch.nonzero(fb[first_active_band, :])[0, 0]
-
-    fb[first_active_band, :first_active_bin] = 1.0
-
-    return fb
-
-
-class TriangularBarkBandsplitSpecification(PerceptualBandsplitSpecification):
-    def __init__(
-        self, nfft: int, fs: int, n_bands: int, f_min: float = 0.0, f_max: float = None
-    ) -> None:
-        super().__init__(
-            fbank_fn=triangular_bark_filterbank,
-            nfft=nfft,
-            fs=fs,
-            n_bands=n_bands,
-            f_min=f_min,
-            f_max=f_max,
-        )
-
-
-def minibark_filterbank(n_bands, fs, f_min, f_max, n_freqs):
-    fb = bark_filterbank(n_bands, fs, f_min, f_max, n_freqs)
-
-    fb[fb < np.sqrt(0.5)] = 0.0
-
-    return fb
-
-
-class MiniBarkBandsplitSpecification(PerceptualBandsplitSpecification):
-    def __init__(
-        self, nfft: int, fs: int, n_bands: int, f_min: float = 0.0, f_max: float = None
-    ) -> None:
-        super().__init__(
-            fbank_fn=minibark_filterbank,
-            nfft=nfft,
-            fs=fs,
-            n_bands=n_bands,
-            f_min=f_min,
-            f_max=f_max,
-        )
-
-
-def erb_filterbank(
-    n_bands: int,
-    fs: int,
-    f_min: float,
-    f_max: float,
-    n_freqs: int,
-) -> Tensor:
-    A = (1000 * np.log(10)) / (24.7 * 4.37)
-    all_freqs = torch.linspace(0, fs // 2, n_freqs)
-
-    m_min = hz2erb(f_min)
-    m_max = hz2erb(f_max)
-
-    m_pts = torch.linspace(m_min, m_max, n_bands + 2)
-    f_pts = (torch.pow(10, (m_pts / A)) - 1) / 0.00437
-
-    fb = _create_triangular_filterbank(all_freqs, f_pts)
-
-    fb = fb.T
-
-    first_active_band = torch.nonzero(torch.sum(fb, dim=-1))[0, 0]
-    first_active_bin = torch.nonzero(fb[first_active_band, :])[0, 0]
-
-    fb[first_active_band, :first_active_bin] = 1.0
-
-    return fb
-
-
-class EquivalentRectangularBandsplitSpecification(PerceptualBandsplitSpecification):
-    def __init__(
-        self, nfft: int, fs: int, n_bands: int, f_min: float = 0.0, f_max: float = None
-    ) -> None:
-        super().__init__(
-            fbank_fn=erb_filterbank,
-            nfft=nfft,
-            fs=fs,
-            n_bands=n_bands,
-            f_min=f_min,
-            f_max=f_max,
-        )
-
-
-if __name__ == "__main__":
-    import pandas as pd
-
-    band_defs = []
-
-    for bands in [VocalBandsplitSpecification]:
-        band_name = bands.__name__.replace("BandsplitSpecification", "")
-
-        mbs = bands(nfft=2048, fs=44100).get_band_specs()
-
-        for i, (f_min, f_max) in enumerate(mbs):
-            band_defs.append(
-                {"band": band_name, "band_index": i, "f_min": f_min, "f_max": f_max}
-            )
-
-    df = pd.DataFrame(band_defs)
-    df.to_csv("vox7bands.csv", index=False)
+import os
+from abc import abstractmethod
+from typing import Any, Callable
+
+import numpy as np
+import torch
+from librosa import hz_to_midi, midi_to_hz
+from torch import Tensor
+from torchaudio import functional as taF
+from spafe.fbanks import bark_fbanks
+from spafe.utils.converters import erb2hz, hz2bark, hz2erb
+from torchaudio.functional.functional import _create_triangular_filterbank
+
+
+def band_widths_from_specs(band_specs):
+    return [e - i for i, e in band_specs]
+
+
+def check_nonzero_bandwidth(band_specs):
+    for fstart, fend in band_specs:
+        if fend - fstart <= 0:
+            raise ValueError("Bands cannot be zero-width")
+
+
+def check_no_overlap(band_specs):
+    fend_prev = -1
+    for fstart_curr, fend_curr in band_specs:
+        if fstart_curr <= fend_prev:
+            raise ValueError("Bands cannot overlap")
+
+
+def check_no_gap(band_specs):
+    fstart, _ = band_specs[0]
+    assert fstart == 0
+
+    fend_prev = -1
+    for fstart_curr, fend_curr in band_specs:
+        if fstart_curr - fend_prev > 1:
+            raise ValueError("Bands cannot leave gap")
+        fend_prev = fend_curr
+
+
+class BandsplitSpecification:
+    def __init__(self, nfft: int, fs: int) -> None:
+        self.fs = fs
+        self.nfft = nfft
+        self.nyquist = fs / 2
+        self.max_index = nfft // 2 + 1
+
+        self.split500 = self.hertz_to_index(500)
+        self.split1k = self.hertz_to_index(1000)
+        self.split2k = self.hertz_to_index(2000)
+        self.split4k = self.hertz_to_index(4000)
+        self.split8k = self.hertz_to_index(8000)
+        self.split16k = self.hertz_to_index(16000)
+        self.split20k = self.hertz_to_index(20000)
+
+        self.above20k = [(self.split20k, self.max_index)]
+        self.above16k = [(self.split16k, self.split20k)] + self.above20k
+
+    def index_to_hertz(self, index: int):
+        return index * self.fs / self.nfft
+
+    def hertz_to_index(self, hz: float, round: bool = True):
+        index = hz * self.nfft / self.fs
+
+        if round:
+            index = int(np.round(index))
+
+        return index
+
+    def get_band_specs_with_bandwidth(self, start_index, end_index, bandwidth_hz):
+        band_specs = []
+        lower = start_index
+
+        while lower < end_index:
+            upper = int(np.floor(lower + self.hertz_to_index(bandwidth_hz)))
+            upper = min(upper, end_index)
+
+            band_specs.append((lower, upper))
+            lower = upper
+
+        return band_specs
+
+    @abstractmethod
+    def get_band_specs(self):
+        raise NotImplementedError
+
+
+class VocalBandsplitSpecification(BandsplitSpecification):
+    def __init__(self, nfft: int, fs: int, version: str = "7") -> None:
+        super().__init__(nfft=nfft, fs=fs)
+
+        self.version = version
+
+    def get_band_specs(self):
+        return getattr(self, f"version{self.version}")()
+
+    @property
+    def version1(self):
+        return self.get_band_specs_with_bandwidth(
+            start_index=0, end_index=self.max_index, bandwidth_hz=1000
+        )
+
+    def version2(self):
+        below16k = self.get_band_specs_with_bandwidth(
+            start_index=0, end_index=self.split16k, bandwidth_hz=1000
+        )
+        below20k = self.get_band_specs_with_bandwidth(
+            start_index=self.split16k, end_index=self.split20k, bandwidth_hz=2000
+        )
+
+        return below16k + below20k + self.above20k
+
+    def version3(self):
+        below8k = self.get_band_specs_with_bandwidth(
+            start_index=0, end_index=self.split8k, bandwidth_hz=1000
+        )
+        below16k = self.get_band_specs_with_bandwidth(
+            start_index=self.split8k, end_index=self.split16k, bandwidth_hz=2000
+        )
+
+        return below8k + below16k + self.above16k
+
+    def version4(self):
+        below1k = self.get_band_specs_with_bandwidth(
+            start_index=0, end_index=self.split1k, bandwidth_hz=100
+        )
+        below8k = self.get_band_specs_with_bandwidth(
+            start_index=self.split1k, end_index=self.split8k, bandwidth_hz=1000
+        )
+        below16k = self.get_band_specs_with_bandwidth(
+            start_index=self.split8k, end_index=self.split16k, bandwidth_hz=2000
+        )
+
+        return below1k + below8k + below16k + self.above16k
+
+    def version5(self):
+        below1k = self.get_band_specs_with_bandwidth(
+            start_index=0, end_index=self.split1k, bandwidth_hz=100
+        )
+        below16k = self.get_band_specs_with_bandwidth(
+            start_index=self.split1k, end_index=self.split16k, bandwidth_hz=1000
+        )
+        below20k = self.get_band_specs_with_bandwidth(
+            start_index=self.split16k, end_index=self.split20k, bandwidth_hz=2000
+        )
+        return below1k + below16k + below20k + self.above20k
+
+    def version6(self):
+        below1k = self.get_band_specs_with_bandwidth(
+            start_index=0, end_index=self.split1k, bandwidth_hz=100
+        )
+        below4k = self.get_band_specs_with_bandwidth(
+            start_index=self.split1k, end_index=self.split4k, bandwidth_hz=500
+        )
+        below8k = self.get_band_specs_with_bandwidth(
+            start_index=self.split4k, end_index=self.split8k, bandwidth_hz=1000
+        )
+        below16k = self.get_band_specs_with_bandwidth(
+            start_index=self.split8k, end_index=self.split16k, bandwidth_hz=2000
+        )
+        return below1k + below4k + below8k + below16k + self.above16k
+
+    def version7(self):
+        below1k = self.get_band_specs_with_bandwidth(
+            start_index=0, end_index=self.split1k, bandwidth_hz=100
+        )
+        below4k = self.get_band_specs_with_bandwidth(
+            start_index=self.split1k, end_index=self.split4k, bandwidth_hz=250
+        )
+        below8k = self.get_band_specs_with_bandwidth(
+            start_index=self.split4k, end_index=self.split8k, bandwidth_hz=500
+        )
+        below16k = self.get_band_specs_with_bandwidth(
+            start_index=self.split8k, end_index=self.split16k, bandwidth_hz=1000
+        )
+        below20k = self.get_band_specs_with_bandwidth(
+            start_index=self.split16k, end_index=self.split20k, bandwidth_hz=2000
+        )
+        return below1k + below4k + below8k + below16k + below20k + self.above20k
+
+
+class OtherBandsplitSpecification(VocalBandsplitSpecification):
+    def __init__(self, nfft: int, fs: int) -> None:
+        super().__init__(nfft=nfft, fs=fs, version="7")
+
+
+class BassBandsplitSpecification(BandsplitSpecification):
+    def __init__(self, nfft: int, fs: int, version: str = "7") -> None:
+        super().__init__(nfft=nfft, fs=fs)
+
+    def get_band_specs(self):
+        below500 = self.get_band_specs_with_bandwidth(
+            start_index=0, end_index=self.split500, bandwidth_hz=50
+        )
+        below1k = self.get_band_specs_with_bandwidth(
+            start_index=self.split500, end_index=self.split1k, bandwidth_hz=100
+        )
+        below4k = self.get_band_specs_with_bandwidth(
+            start_index=self.split1k, end_index=self.split4k, bandwidth_hz=500
+        )
+        below8k = self.get_band_specs_with_bandwidth(
+            start_index=self.split4k, end_index=self.split8k, bandwidth_hz=1000
+        )
+        below16k = self.get_band_specs_with_bandwidth(
+            start_index=self.split8k, end_index=self.split16k, bandwidth_hz=2000
+        )
+        above16k = [(self.split16k, self.max_index)]
+
+        return below500 + below1k + below4k + below8k + below16k + above16k
+
+
+class DrumBandsplitSpecification(BandsplitSpecification):
+    def __init__(self, nfft: int, fs: int) -> None:
+        super().__init__(nfft=nfft, fs=fs)
+
+    def get_band_specs(self):
+        below1k = self.get_band_specs_with_bandwidth(
+            start_index=0, end_index=self.split1k, bandwidth_hz=50
+        )
+        below2k = self.get_band_specs_with_bandwidth(
+            start_index=self.split1k, end_index=self.split2k, bandwidth_hz=100
+        )
+        below4k = self.get_band_specs_with_bandwidth(
+            start_index=self.split2k, end_index=self.split4k, bandwidth_hz=250
+        )
+        below8k = self.get_band_specs_with_bandwidth(
+            start_index=self.split4k, end_index=self.split8k, bandwidth_hz=500
+        )
+        below16k = self.get_band_specs_with_bandwidth(
+            start_index=self.split8k, end_index=self.split16k, bandwidth_hz=1000
+        )
+        above16k = [(self.split16k, self.max_index)]
+
+        return below1k + below2k + below4k + below8k + below16k + above16k
+
+
+class PerceptualBandsplitSpecification(BandsplitSpecification):
+    def __init__(
+        self,
+        nfft: int,
+        fs: int,
+        fbank_fn: Callable[[int, int, float, float, int], torch.Tensor],
+        n_bands: int,
+        f_min: float = 0.0,
+        f_max: float = None,
+    ) -> None:
+        super().__init__(nfft=nfft, fs=fs)
+        self.n_bands = n_bands
+        if f_max is None:
+            f_max = fs / 2
+
+        self.filterbank = fbank_fn(n_bands, fs, f_min, f_max, self.max_index)
+
+        weight_per_bin = torch.sum(self.filterbank, dim=0, keepdim=True)
+        normalized_mel_fb = self.filterbank / weight_per_bin
+
+        freq_weights = []
+        band_specs = []
+        for i in range(self.n_bands):
+            active_bins = torch.nonzero(self.filterbank[i, :]).squeeze().tolist()
+            if isinstance(active_bins, int):
+                active_bins = (active_bins, active_bins)
+            if len(active_bins) == 0:
+                continue
+            start_index = active_bins[0]
+            end_index = active_bins[-1] + 1
+            band_specs.append((start_index, end_index))
+            freq_weights.append(normalized_mel_fb[i, start_index:end_index])
+
+        self.freq_weights = freq_weights
+        self.band_specs = band_specs
+
+    def get_band_specs(self):
+        return self.band_specs
+
+    def get_freq_weights(self):
+        return self.freq_weights
+
+    def save_to_file(self, dir_path: str) -> None:
+
+        os.makedirs(dir_path, exist_ok=True)
+
+        import pickle
+
+        with open(os.path.join(dir_path, "mel_bandsplit_spec.pkl"), "wb") as f:
+            pickle.dump(
+                {
+                    "band_specs": self.band_specs,
+                    "freq_weights": self.freq_weights,
+                    "filterbank": self.filterbank,
+                },
+                f,
+            )
+
+
+def mel_filterbank(n_bands, fs, f_min, f_max, n_freqs):
+    fb = taF.melscale_fbanks(
+        n_mels=n_bands,
+        sample_rate=fs,
+        f_min=f_min,
+        f_max=f_max,
+        n_freqs=n_freqs,
+    ).T
+
+    fb[0, 0] = 1.0
+
+    return fb
+
+
+class MelBandsplitSpecification(PerceptualBandsplitSpecification):
+    def __init__(
+        self, nfft: int, fs: int, n_bands: int, f_min: float = 0.0, f_max: float = None
+    ) -> None:
+        super().__init__(
+            fbank_fn=mel_filterbank,
+            nfft=nfft,
+            fs=fs,
+            n_bands=n_bands,
+            f_min=f_min,
+            f_max=f_max,
+        )
+
+
+def musical_filterbank(n_bands, fs, f_min, f_max, n_freqs, scale="constant"):
+
+    nfft = 2 * (n_freqs - 1)
+    df = fs / nfft
+    f_max = f_max or fs / 2
+    f_min = f_min or 0
+    f_min = fs / nfft
+
+    n_octaves = np.log2(f_max / f_min)
+    n_octaves_per_band = n_octaves / n_bands
+    bandwidth_mult = np.power(2.0, n_octaves_per_band)
+
+    low_midi = max(0, hz_to_midi(f_min))
+    high_midi = hz_to_midi(f_max)
+    midi_points = np.linspace(low_midi, high_midi, n_bands)
+    hz_pts = midi_to_hz(midi_points)
+
+    low_pts = hz_pts / bandwidth_mult
+    high_pts = hz_pts * bandwidth_mult
+
+    low_bins = np.floor(low_pts / df).astype(int)
+    high_bins = np.ceil(high_pts / df).astype(int)
+
+    fb = np.zeros((n_bands, n_freqs))
+
+    for i in range(n_bands):
+        fb[i, low_bins[i] : high_bins[i] + 1] = 1.0
+
+    fb[0, : low_bins[0]] = 1.0
+    fb[-1, high_bins[-1] + 1 :] = 1.0
+
+    return torch.as_tensor(fb)
+
+
+class MusicalBandsplitSpecification(PerceptualBandsplitSpecification):
+    def __init__(
+        self, nfft: int, fs: int, n_bands: int, f_min: float = 0.0, f_max: float = None
+    ) -> None:
+        super().__init__(
+            fbank_fn=musical_filterbank,
+            nfft=nfft,
+            fs=fs,
+            n_bands=n_bands,
+            f_min=f_min,
+            f_max=f_max,
+        )
+
+
+def bark_filterbank(n_bands, fs, f_min, f_max, n_freqs):
+    nfft = 2 * (n_freqs - 1)
+    fb, _ = bark_fbanks.bark_filter_banks(
+        nfilts=n_bands,
+        nfft=nfft,
+        fs=fs,
+        low_freq=f_min,
+        high_freq=f_max,
+        scale="constant",
+    )
+
+    return torch.as_tensor(fb)
+
+
+class BarkBandsplitSpecification(PerceptualBandsplitSpecification):
+    def __init__(
+        self, nfft: int, fs: int, n_bands: int, f_min: float = 0.0, f_max: float = None
+    ) -> None:
+        super().__init__(
+            fbank_fn=bark_filterbank,
+            nfft=nfft,
+            fs=fs,
+            n_bands=n_bands,
+            f_min=f_min,
+            f_max=f_max,
+        )
+
+
+def triangular_bark_filterbank(n_bands, fs, f_min, f_max, n_freqs):
+
+    all_freqs = torch.linspace(0, fs // 2, n_freqs)
+
+    m_min = hz2bark(f_min)
+    m_max = hz2bark(f_max)
+
+    m_pts = torch.linspace(m_min, m_max, n_bands + 2)
+    f_pts = 600 * torch.sinh(m_pts / 6)
+
+    fb = _create_triangular_filterbank(all_freqs, f_pts)
+
+    fb = fb.T
+
+    first_active_band = torch.nonzero(torch.sum(fb, dim=-1))[0, 0]
+    first_active_bin = torch.nonzero(fb[first_active_band, :])[0, 0]
+
+    fb[first_active_band, :first_active_bin] = 1.0
+
+    return fb
+
+
+class TriangularBarkBandsplitSpecification(PerceptualBandsplitSpecification):
+    def __init__(
+        self, nfft: int, fs: int, n_bands: int, f_min: float = 0.0, f_max: float = None
+    ) -> None:
+        super().__init__(
+            fbank_fn=triangular_bark_filterbank,
+            nfft=nfft,
+            fs=fs,
+            n_bands=n_bands,
+            f_min=f_min,
+            f_max=f_max,
+        )
+
+
+def minibark_filterbank(n_bands, fs, f_min, f_max, n_freqs):
+    fb = bark_filterbank(n_bands, fs, f_min, f_max, n_freqs)
+
+    fb[fb < np.sqrt(0.5)] = 0.0
+
+    return fb
+
+
+class MiniBarkBandsplitSpecification(PerceptualBandsplitSpecification):
+    def __init__(
+        self, nfft: int, fs: int, n_bands: int, f_min: float = 0.0, f_max: float = None
+    ) -> None:
+        super().__init__(
+            fbank_fn=minibark_filterbank,
+            nfft=nfft,
+            fs=fs,
+            n_bands=n_bands,
+            f_min=f_min,
+            f_max=f_max,
+        )
+
+
+def erb_filterbank(
+    n_bands: int,
+    fs: int,
+    f_min: float,
+    f_max: float,
+    n_freqs: int,
+) -> Tensor:
+    A = (1000 * np.log(10)) / (24.7 * 4.37)
+    all_freqs = torch.linspace(0, fs // 2, n_freqs)
+
+    m_min = hz2erb(f_min)
+    m_max = hz2erb(f_max)
+
+    m_pts = torch.linspace(m_min, m_max, n_bands + 2)
+    f_pts = (torch.pow(10, (m_pts / A)) - 1) / 0.00437
+
+    fb = _create_triangular_filterbank(all_freqs, f_pts)
+
+    fb = fb.T
+
+    first_active_band = torch.nonzero(torch.sum(fb, dim=-1))[0, 0]
+    first_active_bin = torch.nonzero(fb[first_active_band, :])[0, 0]
+
+    fb[first_active_band, :first_active_bin] = 1.0
+
+    return fb
+
+
+class EquivalentRectangularBandsplitSpecification(PerceptualBandsplitSpecification):
+    def __init__(
+        self, nfft: int, fs: int, n_bands: int, f_min: float = 0.0, f_max: float = None
+    ) -> None:
+        super().__init__(
+            fbank_fn=erb_filterbank,
+            nfft=nfft,
+            fs=fs,
+            n_bands=n_bands,
+            f_min=f_min,
+            f_max=f_max,
+        )
+
+
+if __name__ == "__main__":
+    import pandas as pd
+
+    band_defs = []
+
+    for bands in [VocalBandsplitSpecification]:
+        band_name = bands.__name__.replace("BandsplitSpecification", "")
+
+        mbs = bands(nfft=2048, fs=44100).get_band_specs()
+
+        for i, (f_min, f_max) in enumerate(mbs):
+            band_defs.append(
+                {"band": band_name, "band_index": i, "f_min": f_min, "f_max": f_max}
+            )
+
+    df = pd.DataFrame(band_defs)
+    df.to_csv("vox7bands.csv", index=False)

mvsepless/models/bandit/core/model/bsrnn/wrapper.py

@@ -1,828 +1,828 @@
-from pprint import pprint
-from typing import Dict, List, Optional, Tuple, Union
-
-import torch
-from torch import nn
-
-from .._spectral import _SpectralComponent
-from .utils import (
-    BarkBandsplitSpecification,
-    BassBandsplitSpecification,
-    DrumBandsplitSpecification,
-    EquivalentRectangularBandsplitSpecification,
-    MelBandsplitSpecification,
-    MusicalBandsplitSpecification,
-    OtherBandsplitSpecification,
-    TriangularBarkBandsplitSpecification,
-    VocalBandsplitSpecification,
-)
-from .core import (
-    MultiSourceMultiMaskBandSplitCoreConv,
-    MultiSourceMultiMaskBandSplitCoreRNN,
-    MultiSourceMultiMaskBandSplitCoreTransformer,
-    MultiSourceMultiPatchingMaskBandSplitCoreRNN,
-    SingleMaskBandsplitCoreRNN,
-    SingleMaskBandsplitCoreTransformer,
-)
-
-import pytorch_lightning as pl
-
-
-def get_band_specs(band_specs, n_fft, fs, n_bands=None):
-    if band_specs in ["dnr:speech", "dnr:vox7", "musdb:vocals", "musdb:vox7"]:
-        bsm = VocalBandsplitSpecification(nfft=n_fft, fs=fs).get_band_specs()
-        freq_weights = None
-        overlapping_band = False
-    elif "tribark" in band_specs:
-        assert n_bands is not None
-        specs = TriangularBarkBandsplitSpecification(nfft=n_fft, fs=fs, n_bands=n_bands)
-        bsm = specs.get_band_specs()
-        freq_weights = specs.get_freq_weights()
-        overlapping_band = True
-    elif "bark" in band_specs:
-        assert n_bands is not None
-        specs = BarkBandsplitSpecification(nfft=n_fft, fs=fs, n_bands=n_bands)
-        bsm = specs.get_band_specs()
-        freq_weights = specs.get_freq_weights()
-        overlapping_band = True
-    elif "erb" in band_specs:
-        assert n_bands is not None
-        specs = EquivalentRectangularBandsplitSpecification(
-            nfft=n_fft, fs=fs, n_bands=n_bands
-        )
-        bsm = specs.get_band_specs()
-        freq_weights = specs.get_freq_weights()
-        overlapping_band = True
-    elif "musical" in band_specs:
-        assert n_bands is not None
-        specs = MusicalBandsplitSpecification(nfft=n_fft, fs=fs, n_bands=n_bands)
-        bsm = specs.get_band_specs()
-        freq_weights = specs.get_freq_weights()
-        overlapping_band = True
-    elif band_specs == "dnr:mel" or "mel" in band_specs:
-        assert n_bands is not None
-        specs = MelBandsplitSpecification(nfft=n_fft, fs=fs, n_bands=n_bands)
-        bsm = specs.get_band_specs()
-        freq_weights = specs.get_freq_weights()
-        overlapping_band = True
-    else:
-        raise NameError
-
-    return bsm, freq_weights, overlapping_band
-
-
-def get_band_specs_map(band_specs_map, n_fft, fs, n_bands=None):
-    if band_specs_map == "musdb:all":
-        bsm = {
-            "vocals": VocalBandsplitSpecification(nfft=n_fft, fs=fs).get_band_specs(),
-            "drums": DrumBandsplitSpecification(nfft=n_fft, fs=fs).get_band_specs(),
-            "bass": BassBandsplitSpecification(nfft=n_fft, fs=fs).get_band_specs(),
-            "other": OtherBandsplitSpecification(nfft=n_fft, fs=fs).get_band_specs(),
-        }
-        freq_weights = None
-        overlapping_band = False
-    elif band_specs_map == "dnr:vox7":
-        bsm_, freq_weights, overlapping_band = get_band_specs(
-            "dnr:speech", n_fft, fs, n_bands
-        )
-        bsm = {"speech": bsm_, "music": bsm_, "effects": bsm_}
-    elif "dnr:vox7:" in band_specs_map:
-        stem = band_specs_map.split(":")[-1]
-        bsm_, freq_weights, overlapping_band = get_band_specs(
-            "dnr:speech", n_fft, fs, n_bands
-        )
-        bsm = {stem: bsm_}
-    else:
-        raise NameError
-
-    return bsm, freq_weights, overlapping_band
-
-
-class BandSplitWrapperBase(pl.LightningModule):
-    bsrnn: nn.Module
-
-    def __init__(self, **kwargs):
-        super().__init__()
-
-
-class SingleMaskMultiSourceBandSplitBase(BandSplitWrapperBase, _SpectralComponent):
-    def __init__(
-        self,
-        band_specs_map: Union[str, Dict[str, List[Tuple[float, float]]]],
-        fs: int = 44100,
-        n_fft: int = 2048,
-        win_length: Optional[int] = 2048,
-        hop_length: int = 512,
-        window_fn: str = "hann_window",
-        wkwargs: Optional[Dict] = None,
-        power: Optional[int] = None,
-        center: bool = True,
-        normalized: bool = True,
-        pad_mode: str = "constant",
-        onesided: bool = True,
-        n_bands: int = None,
-    ) -> None:
-        super().__init__(
-            n_fft=n_fft,
-            win_length=win_length,
-            hop_length=hop_length,
-            window_fn=window_fn,
-            wkwargs=wkwargs,
-            power=power,
-            center=center,
-            normalized=normalized,
-            pad_mode=pad_mode,
-            onesided=onesided,
-        )
-
-        if isinstance(band_specs_map, str):
-            self.band_specs_map, self.freq_weights, self.overlapping_band = (
-                get_band_specs_map(band_specs_map, n_fft, fs, n_bands=n_bands)
-            )
-
-        self.stems = list(self.band_specs_map.keys())
-
-    def forward(self, batch):
-        audio = batch["audio"]
-
-        with torch.no_grad():
-            batch["spectrogram"] = {stem: self.stft(audio[stem]) for stem in audio}
-
-        X = batch["spectrogram"]["mixture"]
-        length = batch["audio"]["mixture"].shape[-1]
-
-        output = {"spectrogram": {}, "audio": {}}
-
-        for stem, bsrnn in self.bsrnn.items():
-            S = bsrnn(X)
-            s = self.istft(S, length)
-            output["spectrogram"][stem] = S
-            output["audio"][stem] = s
-
-        return batch, output
-
-
-class MultiMaskMultiSourceBandSplitBase(BandSplitWrapperBase, _SpectralComponent):
-    def __init__(
-        self,
-        stems: List[str],
-        band_specs: Union[str, List[Tuple[float, float]]],
-        fs: int = 44100,
-        n_fft: int = 2048,
-        win_length: Optional[int] = 2048,
-        hop_length: int = 512,
-        window_fn: str = "hann_window",
-        wkwargs: Optional[Dict] = None,
-        power: Optional[int] = None,
-        center: bool = True,
-        normalized: bool = True,
-        pad_mode: str = "constant",
-        onesided: bool = True,
-        n_bands: int = None,
-    ) -> None:
-        super().__init__(
-            n_fft=n_fft,
-            win_length=win_length,
-            hop_length=hop_length,
-            window_fn=window_fn,
-            wkwargs=wkwargs,
-            power=power,
-            center=center,
-            normalized=normalized,
-            pad_mode=pad_mode,
-            onesided=onesided,
-        )
-
-        if isinstance(band_specs, str):
-            self.band_specs, self.freq_weights, self.overlapping_band = get_band_specs(
-                band_specs, n_fft, fs, n_bands
-            )
-
-        self.stems = stems
-
-    def forward(self, batch):
-        audio = batch["audio"]
-        cond = batch.get("condition", None)
-        with torch.no_grad():
-            batch["spectrogram"] = {stem: self.stft(audio[stem]) for stem in audio}
-
-        X = batch["spectrogram"]["mixture"]
-        length = batch["audio"]["mixture"].shape[-1]
-
-        output = self.bsrnn(X, cond=cond)
-        output["audio"] = {}
-
-        for stem, S in output["spectrogram"].items():
-            s = self.istft(S, length)
-            output["audio"][stem] = s
-
-        return batch, output
-
-
-class MultiMaskMultiSourceBandSplitBaseSimple(BandSplitWrapperBase, _SpectralComponent):
-    def __init__(
-        self,
-        stems: List[str],
-        band_specs: Union[str, List[Tuple[float, float]]],
-        fs: int = 44100,
-        n_fft: int = 2048,
-        win_length: Optional[int] = 2048,
-        hop_length: int = 512,
-        window_fn: str = "hann_window",
-        wkwargs: Optional[Dict] = None,
-        power: Optional[int] = None,
-        center: bool = True,
-        normalized: bool = True,
-        pad_mode: str = "constant",
-        onesided: bool = True,
-        n_bands: int = None,
-    ) -> None:
-        super().__init__(
-            n_fft=n_fft,
-            win_length=win_length,
-            hop_length=hop_length,
-            window_fn=window_fn,
-            wkwargs=wkwargs,
-            power=power,
-            center=center,
-            normalized=normalized,
-            pad_mode=pad_mode,
-            onesided=onesided,
-        )
-
-        if isinstance(band_specs, str):
-            self.band_specs, self.freq_weights, self.overlapping_band = get_band_specs(
-                band_specs, n_fft, fs, n_bands
-            )
-
-        self.stems = stems
-
-    def forward(self, batch):
-        with torch.no_grad():
-            X = self.stft(batch)
-        length = batch.shape[-1]
-        output = self.bsrnn(X, cond=None)
-        res = []
-        for stem, S in output["spectrogram"].items():
-            s = self.istft(S, length)
-            res.append(s)
-        res = torch.stack(res, dim=1)
-        return res
-
-
-class SingleMaskMultiSourceBandSplitRNN(SingleMaskMultiSourceBandSplitBase):
-    def __init__(
-        self,
-        in_channel: int,
-        band_specs_map: Union[str, Dict[str, List[Tuple[float, float]]]],
-        fs: int = 44100,
-        require_no_overlap: bool = False,
-        require_no_gap: bool = True,
-        normalize_channel_independently: bool = False,
-        treat_channel_as_feature: bool = True,
-        n_sqm_modules: int = 12,
-        emb_dim: int = 128,
-        rnn_dim: int = 256,
-        bidirectional: bool = True,
-        rnn_type: str = "LSTM",
-        mlp_dim: int = 512,
-        hidden_activation: str = "Tanh",
-        hidden_activation_kwargs: Optional[Dict] = None,
-        complex_mask: bool = True,
-        n_fft: int = 2048,
-        win_length: Optional[int] = 2048,
-        hop_length: int = 512,
-        window_fn: str = "hann_window",
-        wkwargs: Optional[Dict] = None,
-        power: Optional[int] = None,
-        center: bool = True,
-        normalized: bool = True,
-        pad_mode: str = "constant",
-        onesided: bool = True,
-    ) -> None:
-        super().__init__(
-            band_specs_map=band_specs_map,
-            fs=fs,
-            n_fft=n_fft,
-            win_length=win_length,
-            hop_length=hop_length,
-            window_fn=window_fn,
-            wkwargs=wkwargs,
-            power=power,
-            center=center,
-            normalized=normalized,
-            pad_mode=pad_mode,
-            onesided=onesided,
-        )
-
-        self.bsrnn = nn.ModuleDict(
-            {
-                src: SingleMaskBandsplitCoreRNN(
-                    band_specs=specs,
-                    in_channel=in_channel,
-                    require_no_overlap=require_no_overlap,
-                    require_no_gap=require_no_gap,
-                    normalize_channel_independently=normalize_channel_independently,
-                    treat_channel_as_feature=treat_channel_as_feature,
-                    n_sqm_modules=n_sqm_modules,
-                    emb_dim=emb_dim,
-                    rnn_dim=rnn_dim,
-                    bidirectional=bidirectional,
-                    rnn_type=rnn_type,
-                    mlp_dim=mlp_dim,
-                    hidden_activation=hidden_activation,
-                    hidden_activation_kwargs=hidden_activation_kwargs,
-                    complex_mask=complex_mask,
-                )
-                for src, specs in self.band_specs_map.items()
-            }
-        )
-
-
-class SingleMaskMultiSourceBandSplitTransformer(SingleMaskMultiSourceBandSplitBase):
-    def __init__(
-        self,
-        in_channel: int,
-        band_specs_map: Union[str, Dict[str, List[Tuple[float, float]]]],
-        fs: int = 44100,
-        require_no_overlap: bool = False,
-        require_no_gap: bool = True,
-        normalize_channel_independently: bool = False,
-        treat_channel_as_feature: bool = True,
-        n_sqm_modules: int = 12,
-        emb_dim: int = 128,
-        rnn_dim: int = 256,
-        bidirectional: bool = True,
-        tf_dropout: float = 0.0,
-        mlp_dim: int = 512,
-        hidden_activation: str = "Tanh",
-        hidden_activation_kwargs: Optional[Dict] = None,
-        complex_mask: bool = True,
-        n_fft: int = 2048,
-        win_length: Optional[int] = 2048,
-        hop_length: int = 512,
-        window_fn: str = "hann_window",
-        wkwargs: Optional[Dict] = None,
-        power: Optional[int] = None,
-        center: bool = True,
-        normalized: bool = True,
-        pad_mode: str = "constant",
-        onesided: bool = True,
-    ) -> None:
-        super().__init__(
-            band_specs_map=band_specs_map,
-            fs=fs,
-            n_fft=n_fft,
-            win_length=win_length,
-            hop_length=hop_length,
-            window_fn=window_fn,
-            wkwargs=wkwargs,
-            power=power,
-            center=center,
-            normalized=normalized,
-            pad_mode=pad_mode,
-            onesided=onesided,
-        )
-
-        self.bsrnn = nn.ModuleDict(
-            {
-                src: SingleMaskBandsplitCoreTransformer(
-                    band_specs=specs,
-                    in_channel=in_channel,
-                    require_no_overlap=require_no_overlap,
-                    require_no_gap=require_no_gap,
-                    normalize_channel_independently=normalize_channel_independently,
-                    treat_channel_as_feature=treat_channel_as_feature,
-                    n_sqm_modules=n_sqm_modules,
-                    emb_dim=emb_dim,
-                    rnn_dim=rnn_dim,
-                    bidirectional=bidirectional,
-                    tf_dropout=tf_dropout,
-                    mlp_dim=mlp_dim,
-                    hidden_activation=hidden_activation,
-                    hidden_activation_kwargs=hidden_activation_kwargs,
-                    complex_mask=complex_mask,
-                )
-                for src, specs in self.band_specs_map.items()
-            }
-        )
-
-
-class MultiMaskMultiSourceBandSplitRNN(MultiMaskMultiSourceBandSplitBase):
-    def __init__(
-        self,
-        in_channel: int,
-        stems: List[str],
-        band_specs: Union[str, List[Tuple[float, float]]],
-        fs: int = 44100,
-        require_no_overlap: bool = False,
-        require_no_gap: bool = True,
-        normalize_channel_independently: bool = False,
-        treat_channel_as_feature: bool = True,
-        n_sqm_modules: int = 12,
-        emb_dim: int = 128,
-        rnn_dim: int = 256,
-        cond_dim: int = 0,
-        bidirectional: bool = True,
-        rnn_type: str = "LSTM",
-        mlp_dim: int = 512,
-        hidden_activation: str = "Tanh",
-        hidden_activation_kwargs: Optional[Dict] = None,
-        complex_mask: bool = True,
-        n_fft: int = 2048,
-        win_length: Optional[int] = 2048,
-        hop_length: int = 512,
-        window_fn: str = "hann_window",
-        wkwargs: Optional[Dict] = None,
-        power: Optional[int] = None,
-        center: bool = True,
-        normalized: bool = True,
-        pad_mode: str = "constant",
-        onesided: bool = True,
-        n_bands: int = None,
-        use_freq_weights: bool = True,
-        normalize_input: bool = False,
-        mult_add_mask: bool = False,
-        freeze_encoder: bool = False,
-    ) -> None:
-        super().__init__(
-            stems=stems,
-            band_specs=band_specs,
-            fs=fs,
-            n_fft=n_fft,
-            win_length=win_length,
-            hop_length=hop_length,
-            window_fn=window_fn,
-            wkwargs=wkwargs,
-            power=power,
-            center=center,
-            normalized=normalized,
-            pad_mode=pad_mode,
-            onesided=onesided,
-            n_bands=n_bands,
-        )
-
-        self.bsrnn = MultiSourceMultiMaskBandSplitCoreRNN(
-            stems=stems,
-            band_specs=self.band_specs,
-            in_channel=in_channel,
-            require_no_overlap=require_no_overlap,
-            require_no_gap=require_no_gap,
-            normalize_channel_independently=normalize_channel_independently,
-            treat_channel_as_feature=treat_channel_as_feature,
-            n_sqm_modules=n_sqm_modules,
-            emb_dim=emb_dim,
-            rnn_dim=rnn_dim,
-            bidirectional=bidirectional,
-            rnn_type=rnn_type,
-            mlp_dim=mlp_dim,
-            cond_dim=cond_dim,
-            hidden_activation=hidden_activation,
-            hidden_activation_kwargs=hidden_activation_kwargs,
-            complex_mask=complex_mask,
-            overlapping_band=self.overlapping_band,
-            freq_weights=self.freq_weights,
-            n_freq=n_fft // 2 + 1,
-            use_freq_weights=use_freq_weights,
-            mult_add_mask=mult_add_mask,
-        )
-
-        self.normalize_input = normalize_input
-        self.cond_dim = cond_dim
-
-        if freeze_encoder:
-            for param in self.bsrnn.band_split.parameters():
-                param.requires_grad = False
-
-            for param in self.bsrnn.tf_model.parameters():
-                param.requires_grad = False
-
-
-class MultiMaskMultiSourceBandSplitRNNSimple(MultiMaskMultiSourceBandSplitBaseSimple):
-    def __init__(
-        self,
-        in_channel: int,
-        stems: List[str],
-        band_specs: Union[str, List[Tuple[float, float]]],
-        fs: int = 44100,
-        require_no_overlap: bool = False,
-        require_no_gap: bool = True,
-        normalize_channel_independently: bool = False,
-        treat_channel_as_feature: bool = True,
-        n_sqm_modules: int = 12,
-        emb_dim: int = 128,
-        rnn_dim: int = 256,
-        cond_dim: int = 0,
-        bidirectional: bool = True,
-        rnn_type: str = "LSTM",
-        mlp_dim: int = 512,
-        hidden_activation: str = "Tanh",
-        hidden_activation_kwargs: Optional[Dict] = None,
-        complex_mask: bool = True,
-        n_fft: int = 2048,
-        win_length: Optional[int] = 2048,
-        hop_length: int = 512,
-        window_fn: str = "hann_window",
-        wkwargs: Optional[Dict] = None,
-        power: Optional[int] = None,
-        center: bool = True,
-        normalized: bool = True,
-        pad_mode: str = "constant",
-        onesided: bool = True,
-        n_bands: int = None,
-        use_freq_weights: bool = True,
-        normalize_input: bool = False,
-        mult_add_mask: bool = False,
-        freeze_encoder: bool = False,
-    ) -> None:
-        super().__init__(
-            stems=stems,
-            band_specs=band_specs,
-            fs=fs,
-            n_fft=n_fft,
-            win_length=win_length,
-            hop_length=hop_length,
-            window_fn=window_fn,
-            wkwargs=wkwargs,
-            power=power,
-            center=center,
-            normalized=normalized,
-            pad_mode=pad_mode,
-            onesided=onesided,
-            n_bands=n_bands,
-        )
-
-        self.bsrnn = MultiSourceMultiMaskBandSplitCoreRNN(
-            stems=stems,
-            band_specs=self.band_specs,
-            in_channel=in_channel,
-            require_no_overlap=require_no_overlap,
-            require_no_gap=require_no_gap,
-            normalize_channel_independently=normalize_channel_independently,
-            treat_channel_as_feature=treat_channel_as_feature,
-            n_sqm_modules=n_sqm_modules,
-            emb_dim=emb_dim,
-            rnn_dim=rnn_dim,
-            bidirectional=bidirectional,
-            rnn_type=rnn_type,
-            mlp_dim=mlp_dim,
-            cond_dim=cond_dim,
-            hidden_activation=hidden_activation,
-            hidden_activation_kwargs=hidden_activation_kwargs,
-            complex_mask=complex_mask,
-            overlapping_band=self.overlapping_band,
-            freq_weights=self.freq_weights,
-            n_freq=n_fft // 2 + 1,
-            use_freq_weights=use_freq_weights,
-            mult_add_mask=mult_add_mask,
-        )
-
-        self.normalize_input = normalize_input
-        self.cond_dim = cond_dim
-
-        if freeze_encoder:
-            for param in self.bsrnn.band_split.parameters():
-                param.requires_grad = False
-
-            for param in self.bsrnn.tf_model.parameters():
-                param.requires_grad = False
-
-
-class MultiMaskMultiSourceBandSplitTransformer(MultiMaskMultiSourceBandSplitBase):
-    def __init__(
-        self,
-        in_channel: int,
-        stems: List[str],
-        band_specs: Union[str, List[Tuple[float, float]]],
-        fs: int = 44100,
-        require_no_overlap: bool = False,
-        require_no_gap: bool = True,
-        normalize_channel_independently: bool = False,
-        treat_channel_as_feature: bool = True,
-        n_sqm_modules: int = 12,
-        emb_dim: int = 128,
-        rnn_dim: int = 256,
-        cond_dim: int = 0,
-        bidirectional: bool = True,
-        rnn_type: str = "LSTM",
-        mlp_dim: int = 512,
-        hidden_activation: str = "Tanh",
-        hidden_activation_kwargs: Optional[Dict] = None,
-        complex_mask: bool = True,
-        n_fft: int = 2048,
-        win_length: Optional[int] = 2048,
-        hop_length: int = 512,
-        window_fn: str = "hann_window",
-        wkwargs: Optional[Dict] = None,
-        power: Optional[int] = None,
-        center: bool = True,
-        normalized: bool = True,
-        pad_mode: str = "constant",
-        onesided: bool = True,
-        n_bands: int = None,
-        use_freq_weights: bool = True,
-        normalize_input: bool = False,
-        mult_add_mask: bool = False,
-    ) -> None:
-        super().__init__(
-            stems=stems,
-            band_specs=band_specs,
-            fs=fs,
-            n_fft=n_fft,
-            win_length=win_length,
-            hop_length=hop_length,
-            window_fn=window_fn,
-            wkwargs=wkwargs,
-            power=power,
-            center=center,
-            normalized=normalized,
-            pad_mode=pad_mode,
-            onesided=onesided,
-            n_bands=n_bands,
-        )
-
-        self.bsrnn = MultiSourceMultiMaskBandSplitCoreTransformer(
-            stems=stems,
-            band_specs=self.band_specs,
-            in_channel=in_channel,
-            require_no_overlap=require_no_overlap,
-            require_no_gap=require_no_gap,
-            normalize_channel_independently=normalize_channel_independently,
-            treat_channel_as_feature=treat_channel_as_feature,
-            n_sqm_modules=n_sqm_modules,
-            emb_dim=emb_dim,
-            rnn_dim=rnn_dim,
-            bidirectional=bidirectional,
-            rnn_type=rnn_type,
-            mlp_dim=mlp_dim,
-            cond_dim=cond_dim,
-            hidden_activation=hidden_activation,
-            hidden_activation_kwargs=hidden_activation_kwargs,
-            complex_mask=complex_mask,
-            overlapping_band=self.overlapping_band,
-            freq_weights=self.freq_weights,
-            n_freq=n_fft // 2 + 1,
-            use_freq_weights=use_freq_weights,
-            mult_add_mask=mult_add_mask,
-        )
-
-
-class MultiMaskMultiSourceBandSplitConv(MultiMaskMultiSourceBandSplitBase):
-    def __init__(
-        self,
-        in_channel: int,
-        stems: List[str],
-        band_specs: Union[str, List[Tuple[float, float]]],
-        fs: int = 44100,
-        require_no_overlap: bool = False,
-        require_no_gap: bool = True,
-        normalize_channel_independently: bool = False,
-        treat_channel_as_feature: bool = True,
-        n_sqm_modules: int = 12,
-        emb_dim: int = 128,
-        rnn_dim: int = 256,
-        cond_dim: int = 0,
-        bidirectional: bool = True,
-        rnn_type: str = "LSTM",
-        mlp_dim: int = 512,
-        hidden_activation: str = "Tanh",
-        hidden_activation_kwargs: Optional[Dict] = None,
-        complex_mask: bool = True,
-        n_fft: int = 2048,
-        win_length: Optional[int] = 2048,
-        hop_length: int = 512,
-        window_fn: str = "hann_window",
-        wkwargs: Optional[Dict] = None,
-        power: Optional[int] = None,
-        center: bool = True,
-        normalized: bool = True,
-        pad_mode: str = "constant",
-        onesided: bool = True,
-        n_bands: int = None,
-        use_freq_weights: bool = True,
-        normalize_input: bool = False,
-        mult_add_mask: bool = False,
-    ) -> None:
-        super().__init__(
-            stems=stems,
-            band_specs=band_specs,
-            fs=fs,
-            n_fft=n_fft,
-            win_length=win_length,
-            hop_length=hop_length,
-            window_fn=window_fn,
-            wkwargs=wkwargs,
-            power=power,
-            center=center,
-            normalized=normalized,
-            pad_mode=pad_mode,
-            onesided=onesided,
-            n_bands=n_bands,
-        )
-
-        self.bsrnn = MultiSourceMultiMaskBandSplitCoreConv(
-            stems=stems,
-            band_specs=self.band_specs,
-            in_channel=in_channel,
-            require_no_overlap=require_no_overlap,
-            require_no_gap=require_no_gap,
-            normalize_channel_independently=normalize_channel_independently,
-            treat_channel_as_feature=treat_channel_as_feature,
-            n_sqm_modules=n_sqm_modules,
-            emb_dim=emb_dim,
-            rnn_dim=rnn_dim,
-            bidirectional=bidirectional,
-            rnn_type=rnn_type,
-            mlp_dim=mlp_dim,
-            cond_dim=cond_dim,
-            hidden_activation=hidden_activation,
-            hidden_activation_kwargs=hidden_activation_kwargs,
-            complex_mask=complex_mask,
-            overlapping_band=self.overlapping_band,
-            freq_weights=self.freq_weights,
-            n_freq=n_fft // 2 + 1,
-            use_freq_weights=use_freq_weights,
-            mult_add_mask=mult_add_mask,
-        )
-
-
-class PatchingMaskMultiSourceBandSplitRNN(MultiMaskMultiSourceBandSplitBase):
-    def __init__(
-        self,
-        in_channel: int,
-        stems: List[str],
-        band_specs: Union[str, List[Tuple[float, float]]],
-        kernel_norm_mlp_version: int = 1,
-        mask_kernel_freq: int = 3,
-        mask_kernel_time: int = 3,
-        conv_kernel_freq: int = 1,
-        conv_kernel_time: int = 1,
-        fs: int = 44100,
-        require_no_overlap: bool = False,
-        require_no_gap: bool = True,
-        normalize_channel_independently: bool = False,
-        treat_channel_as_feature: bool = True,
-        n_sqm_modules: int = 12,
-        emb_dim: int = 128,
-        rnn_dim: int = 256,
-        bidirectional: bool = True,
-        rnn_type: str = "LSTM",
-        mlp_dim: int = 512,
-        hidden_activation: str = "Tanh",
-        hidden_activation_kwargs: Optional[Dict] = None,
-        complex_mask: bool = True,
-        n_fft: int = 2048,
-        win_length: Optional[int] = 2048,
-        hop_length: int = 512,
-        window_fn: str = "hann_window",
-        wkwargs: Optional[Dict] = None,
-        power: Optional[int] = None,
-        center: bool = True,
-        normalized: bool = True,
-        pad_mode: str = "constant",
-        onesided: bool = True,
-        n_bands: int = None,
-    ) -> None:
-        super().__init__(
-            stems=stems,
-            band_specs=band_specs,
-            fs=fs,
-            n_fft=n_fft,
-            win_length=win_length,
-            hop_length=hop_length,
-            window_fn=window_fn,
-            wkwargs=wkwargs,
-            power=power,
-            center=center,
-            normalized=normalized,
-            pad_mode=pad_mode,
-            onesided=onesided,
-            n_bands=n_bands,
-        )
-
-        self.bsrnn = MultiSourceMultiPatchingMaskBandSplitCoreRNN(
-            stems=stems,
-            band_specs=self.band_specs,
-            in_channel=in_channel,
-            require_no_overlap=require_no_overlap,
-            require_no_gap=require_no_gap,
-            normalize_channel_independently=normalize_channel_independently,
-            treat_channel_as_feature=treat_channel_as_feature,
-            n_sqm_modules=n_sqm_modules,
-            emb_dim=emb_dim,
-            rnn_dim=rnn_dim,
-            bidirectional=bidirectional,
-            rnn_type=rnn_type,
-            mlp_dim=mlp_dim,
-            hidden_activation=hidden_activation,
-            hidden_activation_kwargs=hidden_activation_kwargs,
-            complex_mask=complex_mask,
-            overlapping_band=self.overlapping_band,
-            freq_weights=self.freq_weights,
-            n_freq=n_fft // 2 + 1,
-            mask_kernel_freq=mask_kernel_freq,
-            mask_kernel_time=mask_kernel_time,
-            conv_kernel_freq=conv_kernel_freq,
-            conv_kernel_time=conv_kernel_time,
-            kernel_norm_mlp_version=kernel_norm_mlp_version,
-        )
+from pprint import pprint
+from typing import Dict, List, Optional, Tuple, Union
+
+import torch
+from torch import nn
+
+from .._spectral import _SpectralComponent
+from .utils import (
+    BarkBandsplitSpecification,
+    BassBandsplitSpecification,
+    DrumBandsplitSpecification,
+    EquivalentRectangularBandsplitSpecification,
+    MelBandsplitSpecification,
+    MusicalBandsplitSpecification,
+    OtherBandsplitSpecification,
+    TriangularBarkBandsplitSpecification,
+    VocalBandsplitSpecification,
+)
+from .core import (
+    MultiSourceMultiMaskBandSplitCoreConv,
+    MultiSourceMultiMaskBandSplitCoreRNN,
+    MultiSourceMultiMaskBandSplitCoreTransformer,
+    MultiSourceMultiPatchingMaskBandSplitCoreRNN,
+    SingleMaskBandsplitCoreRNN,
+    SingleMaskBandsplitCoreTransformer,
+)
+
+import pytorch_lightning as pl
+
+
+def get_band_specs(band_specs, n_fft, fs, n_bands=None):
+    if band_specs in ["dnr:speech", "dnr:vox7", "musdb:vocals", "musdb:vox7"]:
+        bsm = VocalBandsplitSpecification(nfft=n_fft, fs=fs).get_band_specs()
+        freq_weights = None
+        overlapping_band = False
+    elif "tribark" in band_specs:
+        assert n_bands is not None
+        specs = TriangularBarkBandsplitSpecification(nfft=n_fft, fs=fs, n_bands=n_bands)
+        bsm = specs.get_band_specs()
+        freq_weights = specs.get_freq_weights()
+        overlapping_band = True
+    elif "bark" in band_specs:
+        assert n_bands is not None
+        specs = BarkBandsplitSpecification(nfft=n_fft, fs=fs, n_bands=n_bands)
+        bsm = specs.get_band_specs()
+        freq_weights = specs.get_freq_weights()
+        overlapping_band = True
+    elif "erb" in band_specs:
+        assert n_bands is not None
+        specs = EquivalentRectangularBandsplitSpecification(
+            nfft=n_fft, fs=fs, n_bands=n_bands
+        )
+        bsm = specs.get_band_specs()
+        freq_weights = specs.get_freq_weights()
+        overlapping_band = True
+    elif "musical" in band_specs:
+        assert n_bands is not None
+        specs = MusicalBandsplitSpecification(nfft=n_fft, fs=fs, n_bands=n_bands)
+        bsm = specs.get_band_specs()
+        freq_weights = specs.get_freq_weights()
+        overlapping_band = True
+    elif band_specs == "dnr:mel" or "mel" in band_specs:
+        assert n_bands is not None
+        specs = MelBandsplitSpecification(nfft=n_fft, fs=fs, n_bands=n_bands)
+        bsm = specs.get_band_specs()
+        freq_weights = specs.get_freq_weights()
+        overlapping_band = True
+    else:
+        raise NameError
+
+    return bsm, freq_weights, overlapping_band
+
+
+def get_band_specs_map(band_specs_map, n_fft, fs, n_bands=None):
+    if band_specs_map == "musdb:all":
+        bsm = {
+            "vocals": VocalBandsplitSpecification(nfft=n_fft, fs=fs).get_band_specs(),
+            "drums": DrumBandsplitSpecification(nfft=n_fft, fs=fs).get_band_specs(),
+            "bass": BassBandsplitSpecification(nfft=n_fft, fs=fs).get_band_specs(),
+            "other": OtherBandsplitSpecification(nfft=n_fft, fs=fs).get_band_specs(),
+        }
+        freq_weights = None
+        overlapping_band = False
+    elif band_specs_map == "dnr:vox7":
+        bsm_, freq_weights, overlapping_band = get_band_specs(
+            "dnr:speech", n_fft, fs, n_bands
+        )
+        bsm = {"speech": bsm_, "music": bsm_, "effects": bsm_}
+    elif "dnr:vox7:" in band_specs_map:
+        stem = band_specs_map.split(":")[-1]
+        bsm_, freq_weights, overlapping_band = get_band_specs(
+            "dnr:speech", n_fft, fs, n_bands
+        )
+        bsm = {stem: bsm_}
+    else:
+        raise NameError
+
+    return bsm, freq_weights, overlapping_band
+
+
+class BandSplitWrapperBase(pl.LightningModule):
+    bsrnn: nn.Module
+
+    def __init__(self, **kwargs):
+        super().__init__()
+
+
+class SingleMaskMultiSourceBandSplitBase(BandSplitWrapperBase, _SpectralComponent):
+    def __init__(
+        self,
+        band_specs_map: Union[str, Dict[str, List[Tuple[float, float]]]],
+        fs: int = 44100,
+        n_fft: int = 2048,
+        win_length: Optional[int] = 2048,
+        hop_length: int = 512,
+        window_fn: str = "hann_window",
+        wkwargs: Optional[Dict] = None,
+        power: Optional[int] = None,
+        center: bool = True,
+        normalized: bool = True,
+        pad_mode: str = "constant",
+        onesided: bool = True,
+        n_bands: int = None,
+    ) -> None:
+        super().__init__(
+            n_fft=n_fft,
+            win_length=win_length,
+            hop_length=hop_length,
+            window_fn=window_fn,
+            wkwargs=wkwargs,
+            power=power,
+            center=center,
+            normalized=normalized,
+            pad_mode=pad_mode,
+            onesided=onesided,
+        )
+
+        if isinstance(band_specs_map, str):
+            self.band_specs_map, self.freq_weights, self.overlapping_band = (
+                get_band_specs_map(band_specs_map, n_fft, fs, n_bands=n_bands)
+            )
+
+        self.stems = list(self.band_specs_map.keys())
+
+    def forward(self, batch):
+        audio = batch["audio"]
+
+        with torch.no_grad():
+            batch["spectrogram"] = {stem: self.stft(audio[stem]) for stem in audio}
+
+        X = batch["spectrogram"]["mixture"]
+        length = batch["audio"]["mixture"].shape[-1]
+
+        output = {"spectrogram": {}, "audio": {}}
+
+        for stem, bsrnn in self.bsrnn.items():
+            S = bsrnn(X)
+            s = self.istft(S, length)
+            output["spectrogram"][stem] = S
+            output["audio"][stem] = s
+
+        return batch, output
+
+
+class MultiMaskMultiSourceBandSplitBase(BandSplitWrapperBase, _SpectralComponent):
+    def __init__(
+        self,
+        stems: List[str],
+        band_specs: Union[str, List[Tuple[float, float]]],
+        fs: int = 44100,
+        n_fft: int = 2048,
+        win_length: Optional[int] = 2048,
+        hop_length: int = 512,
+        window_fn: str = "hann_window",
+        wkwargs: Optional[Dict] = None,
+        power: Optional[int] = None,
+        center: bool = True,
+        normalized: bool = True,
+        pad_mode: str = "constant",
+        onesided: bool = True,
+        n_bands: int = None,
+    ) -> None:
+        super().__init__(
+            n_fft=n_fft,
+            win_length=win_length,
+            hop_length=hop_length,
+            window_fn=window_fn,
+            wkwargs=wkwargs,
+            power=power,
+            center=center,
+            normalized=normalized,
+            pad_mode=pad_mode,
+            onesided=onesided,
+        )
+
+        if isinstance(band_specs, str):
+            self.band_specs, self.freq_weights, self.overlapping_band = get_band_specs(
+                band_specs, n_fft, fs, n_bands
+            )
+
+        self.stems = stems
+
+    def forward(self, batch):
+        audio = batch["audio"]
+        cond = batch.get("condition", None)
+        with torch.no_grad():
+            batch["spectrogram"] = {stem: self.stft(audio[stem]) for stem in audio}
+
+        X = batch["spectrogram"]["mixture"]
+        length = batch["audio"]["mixture"].shape[-1]
+
+        output = self.bsrnn(X, cond=cond)
+        output["audio"] = {}
+
+        for stem, S in output["spectrogram"].items():
+            s = self.istft(S, length)
+            output["audio"][stem] = s
+
+        return batch, output
+
+
+class MultiMaskMultiSourceBandSplitBaseSimple(BandSplitWrapperBase, _SpectralComponent):
+    def __init__(
+        self,
+        stems: List[str],
+        band_specs: Union[str, List[Tuple[float, float]]],
+        fs: int = 44100,
+        n_fft: int = 2048,
+        win_length: Optional[int] = 2048,
+        hop_length: int = 512,
+        window_fn: str = "hann_window",
+        wkwargs: Optional[Dict] = None,
+        power: Optional[int] = None,
+        center: bool = True,
+        normalized: bool = True,
+        pad_mode: str = "constant",
+        onesided: bool = True,
+        n_bands: int = None,
+    ) -> None:
+        super().__init__(
+            n_fft=n_fft,
+            win_length=win_length,
+            hop_length=hop_length,
+            window_fn=window_fn,
+            wkwargs=wkwargs,
+            power=power,
+            center=center,
+            normalized=normalized,
+            pad_mode=pad_mode,
+            onesided=onesided,
+        )
+
+        if isinstance(band_specs, str):
+            self.band_specs, self.freq_weights, self.overlapping_band = get_band_specs(
+                band_specs, n_fft, fs, n_bands
+            )
+
+        self.stems = stems
+
+    def forward(self, batch):
+        with torch.no_grad():
+            X = self.stft(batch)
+        length = batch.shape[-1]
+        output = self.bsrnn(X, cond=None)
+        res = []
+        for stem, S in output["spectrogram"].items():
+            s = self.istft(S, length)
+            res.append(s)
+        res = torch.stack(res, dim=1)
+        return res
+
+
+class SingleMaskMultiSourceBandSplitRNN(SingleMaskMultiSourceBandSplitBase):
+    def __init__(
+        self,
+        in_channel: int,
+        band_specs_map: Union[str, Dict[str, List[Tuple[float, float]]]],
+        fs: int = 44100,
+        require_no_overlap: bool = False,
+        require_no_gap: bool = True,
+        normalize_channel_independently: bool = False,
+        treat_channel_as_feature: bool = True,
+        n_sqm_modules: int = 12,
+        emb_dim: int = 128,
+        rnn_dim: int = 256,
+        bidirectional: bool = True,
+        rnn_type: str = "LSTM",
+        mlp_dim: int = 512,
+        hidden_activation: str = "Tanh",
+        hidden_activation_kwargs: Optional[Dict] = None,
+        complex_mask: bool = True,
+        n_fft: int = 2048,
+        win_length: Optional[int] = 2048,
+        hop_length: int = 512,
+        window_fn: str = "hann_window",
+        wkwargs: Optional[Dict] = None,
+        power: Optional[int] = None,
+        center: bool = True,
+        normalized: bool = True,
+        pad_mode: str = "constant",
+        onesided: bool = True,
+    ) -> None:
+        super().__init__(
+            band_specs_map=band_specs_map,
+            fs=fs,
+            n_fft=n_fft,
+            win_length=win_length,
+            hop_length=hop_length,
+            window_fn=window_fn,
+            wkwargs=wkwargs,
+            power=power,
+            center=center,
+            normalized=normalized,
+            pad_mode=pad_mode,
+            onesided=onesided,
+        )
+
+        self.bsrnn = nn.ModuleDict(
+            {
+                src: SingleMaskBandsplitCoreRNN(
+                    band_specs=specs,
+                    in_channel=in_channel,
+                    require_no_overlap=require_no_overlap,
+                    require_no_gap=require_no_gap,
+                    normalize_channel_independently=normalize_channel_independently,
+                    treat_channel_as_feature=treat_channel_as_feature,
+                    n_sqm_modules=n_sqm_modules,
+                    emb_dim=emb_dim,
+                    rnn_dim=rnn_dim,
+                    bidirectional=bidirectional,
+                    rnn_type=rnn_type,
+                    mlp_dim=mlp_dim,
+                    hidden_activation=hidden_activation,
+                    hidden_activation_kwargs=hidden_activation_kwargs,
+                    complex_mask=complex_mask,
+                )
+                for src, specs in self.band_specs_map.items()
+            }
+        )
+
+
+class SingleMaskMultiSourceBandSplitTransformer(SingleMaskMultiSourceBandSplitBase):
+    def __init__(
+        self,
+        in_channel: int,
+        band_specs_map: Union[str, Dict[str, List[Tuple[float, float]]]],
+        fs: int = 44100,
+        require_no_overlap: bool = False,
+        require_no_gap: bool = True,
+        normalize_channel_independently: bool = False,
+        treat_channel_as_feature: bool = True,
+        n_sqm_modules: int = 12,
+        emb_dim: int = 128,
+        rnn_dim: int = 256,
+        bidirectional: bool = True,
+        tf_dropout: float = 0.0,
+        mlp_dim: int = 512,
+        hidden_activation: str = "Tanh",
+        hidden_activation_kwargs: Optional[Dict] = None,
+        complex_mask: bool = True,
+        n_fft: int = 2048,
+        win_length: Optional[int] = 2048,
+        hop_length: int = 512,
+        window_fn: str = "hann_window",
+        wkwargs: Optional[Dict] = None,
+        power: Optional[int] = None,
+        center: bool = True,
+        normalized: bool = True,
+        pad_mode: str = "constant",
+        onesided: bool = True,
+    ) -> None:
+        super().__init__(
+            band_specs_map=band_specs_map,
+            fs=fs,
+            n_fft=n_fft,
+            win_length=win_length,
+            hop_length=hop_length,
+            window_fn=window_fn,
+            wkwargs=wkwargs,
+            power=power,
+            center=center,
+            normalized=normalized,
+            pad_mode=pad_mode,
+            onesided=onesided,
+        )
+
+        self.bsrnn = nn.ModuleDict(
+            {
+                src: SingleMaskBandsplitCoreTransformer(
+                    band_specs=specs,
+                    in_channel=in_channel,
+                    require_no_overlap=require_no_overlap,
+                    require_no_gap=require_no_gap,
+                    normalize_channel_independently=normalize_channel_independently,
+                    treat_channel_as_feature=treat_channel_as_feature,
+                    n_sqm_modules=n_sqm_modules,
+                    emb_dim=emb_dim,
+                    rnn_dim=rnn_dim,
+                    bidirectional=bidirectional,
+                    tf_dropout=tf_dropout,
+                    mlp_dim=mlp_dim,
+                    hidden_activation=hidden_activation,
+                    hidden_activation_kwargs=hidden_activation_kwargs,
+                    complex_mask=complex_mask,
+                )
+                for src, specs in self.band_specs_map.items()
+            }
+        )
+
+
+class MultiMaskMultiSourceBandSplitRNN(MultiMaskMultiSourceBandSplitBase):
+    def __init__(
+        self,
+        in_channel: int,
+        stems: List[str],
+        band_specs: Union[str, List[Tuple[float, float]]],
+        fs: int = 44100,
+        require_no_overlap: bool = False,
+        require_no_gap: bool = True,
+        normalize_channel_independently: bool = False,
+        treat_channel_as_feature: bool = True,
+        n_sqm_modules: int = 12,
+        emb_dim: int = 128,
+        rnn_dim: int = 256,
+        cond_dim: int = 0,
+        bidirectional: bool = True,
+        rnn_type: str = "LSTM",
+        mlp_dim: int = 512,
+        hidden_activation: str = "Tanh",
+        hidden_activation_kwargs: Optional[Dict] = None,
+        complex_mask: bool = True,
+        n_fft: int = 2048,
+        win_length: Optional[int] = 2048,
+        hop_length: int = 512,
+        window_fn: str = "hann_window",
+        wkwargs: Optional[Dict] = None,
+        power: Optional[int] = None,
+        center: bool = True,
+        normalized: bool = True,
+        pad_mode: str = "constant",
+        onesided: bool = True,
+        n_bands: int = None,
+        use_freq_weights: bool = True,
+        normalize_input: bool = False,
+        mult_add_mask: bool = False,
+        freeze_encoder: bool = False,
+    ) -> None:
+        super().__init__(
+            stems=stems,
+            band_specs=band_specs,
+            fs=fs,
+            n_fft=n_fft,
+            win_length=win_length,
+            hop_length=hop_length,
+            window_fn=window_fn,
+            wkwargs=wkwargs,
+            power=power,
+            center=center,
+            normalized=normalized,
+            pad_mode=pad_mode,
+            onesided=onesided,
+            n_bands=n_bands,
+        )
+
+        self.bsrnn = MultiSourceMultiMaskBandSplitCoreRNN(
+            stems=stems,
+            band_specs=self.band_specs,
+            in_channel=in_channel,
+            require_no_overlap=require_no_overlap,
+            require_no_gap=require_no_gap,
+            normalize_channel_independently=normalize_channel_independently,
+            treat_channel_as_feature=treat_channel_as_feature,
+            n_sqm_modules=n_sqm_modules,
+            emb_dim=emb_dim,
+            rnn_dim=rnn_dim,
+            bidirectional=bidirectional,
+            rnn_type=rnn_type,
+            mlp_dim=mlp_dim,
+            cond_dim=cond_dim,
+            hidden_activation=hidden_activation,
+            hidden_activation_kwargs=hidden_activation_kwargs,
+            complex_mask=complex_mask,
+            overlapping_band=self.overlapping_band,
+            freq_weights=self.freq_weights,
+            n_freq=n_fft // 2 + 1,
+            use_freq_weights=use_freq_weights,
+            mult_add_mask=mult_add_mask,
+        )
+
+        self.normalize_input = normalize_input
+        self.cond_dim = cond_dim
+
+        if freeze_encoder:
+            for param in self.bsrnn.band_split.parameters():
+                param.requires_grad = False
+
+            for param in self.bsrnn.tf_model.parameters():
+                param.requires_grad = False
+
+
+class MultiMaskMultiSourceBandSplitRNNSimple(MultiMaskMultiSourceBandSplitBaseSimple):
+    def __init__(
+        self,
+        in_channel: int,
+        stems: List[str],
+        band_specs: Union[str, List[Tuple[float, float]]],
+        fs: int = 44100,
+        require_no_overlap: bool = False,
+        require_no_gap: bool = True,
+        normalize_channel_independently: bool = False,
+        treat_channel_as_feature: bool = True,
+        n_sqm_modules: int = 12,
+        emb_dim: int = 128,
+        rnn_dim: int = 256,
+        cond_dim: int = 0,
+        bidirectional: bool = True,
+        rnn_type: str = "LSTM",
+        mlp_dim: int = 512,
+        hidden_activation: str = "Tanh",
+        hidden_activation_kwargs: Optional[Dict] = None,
+        complex_mask: bool = True,
+        n_fft: int = 2048,
+        win_length: Optional[int] = 2048,
+        hop_length: int = 512,
+        window_fn: str = "hann_window",
+        wkwargs: Optional[Dict] = None,
+        power: Optional[int] = None,
+        center: bool = True,
+        normalized: bool = True,
+        pad_mode: str = "constant",
+        onesided: bool = True,
+        n_bands: int = None,
+        use_freq_weights: bool = True,
+        normalize_input: bool = False,
+        mult_add_mask: bool = False,
+        freeze_encoder: bool = False,
+    ) -> None:
+        super().__init__(
+            stems=stems,
+            band_specs=band_specs,
+            fs=fs,
+            n_fft=n_fft,
+            win_length=win_length,
+            hop_length=hop_length,
+            window_fn=window_fn,
+            wkwargs=wkwargs,
+            power=power,
+            center=center,
+            normalized=normalized,
+            pad_mode=pad_mode,
+            onesided=onesided,
+            n_bands=n_bands,
+        )
+
+        self.bsrnn = MultiSourceMultiMaskBandSplitCoreRNN(
+            stems=stems,
+            band_specs=self.band_specs,
+            in_channel=in_channel,
+            require_no_overlap=require_no_overlap,
+            require_no_gap=require_no_gap,
+            normalize_channel_independently=normalize_channel_independently,
+            treat_channel_as_feature=treat_channel_as_feature,
+            n_sqm_modules=n_sqm_modules,
+            emb_dim=emb_dim,
+            rnn_dim=rnn_dim,
+            bidirectional=bidirectional,
+            rnn_type=rnn_type,
+            mlp_dim=mlp_dim,
+            cond_dim=cond_dim,
+            hidden_activation=hidden_activation,
+            hidden_activation_kwargs=hidden_activation_kwargs,
+            complex_mask=complex_mask,
+            overlapping_band=self.overlapping_band,
+            freq_weights=self.freq_weights,
+            n_freq=n_fft // 2 + 1,
+            use_freq_weights=use_freq_weights,
+            mult_add_mask=mult_add_mask,
+        )
+
+        self.normalize_input = normalize_input
+        self.cond_dim = cond_dim
+
+        if freeze_encoder:
+            for param in self.bsrnn.band_split.parameters():
+                param.requires_grad = False
+
+            for param in self.bsrnn.tf_model.parameters():
+                param.requires_grad = False
+
+
+class MultiMaskMultiSourceBandSplitTransformer(MultiMaskMultiSourceBandSplitBase):
+    def __init__(
+        self,
+        in_channel: int,
+        stems: List[str],
+        band_specs: Union[str, List[Tuple[float, float]]],
+        fs: int = 44100,
+        require_no_overlap: bool = False,
+        require_no_gap: bool = True,
+        normalize_channel_independently: bool = False,
+        treat_channel_as_feature: bool = True,
+        n_sqm_modules: int = 12,
+        emb_dim: int = 128,
+        rnn_dim: int = 256,
+        cond_dim: int = 0,
+        bidirectional: bool = True,
+        rnn_type: str = "LSTM",
+        mlp_dim: int = 512,
+        hidden_activation: str = "Tanh",
+        hidden_activation_kwargs: Optional[Dict] = None,
+        complex_mask: bool = True,
+        n_fft: int = 2048,
+        win_length: Optional[int] = 2048,
+        hop_length: int = 512,
+        window_fn: str = "hann_window",
+        wkwargs: Optional[Dict] = None,
+        power: Optional[int] = None,
+        center: bool = True,
+        normalized: bool = True,
+        pad_mode: str = "constant",
+        onesided: bool = True,
+        n_bands: int = None,
+        use_freq_weights: bool = True,
+        normalize_input: bool = False,
+        mult_add_mask: bool = False,
+    ) -> None:
+        super().__init__(
+            stems=stems,
+            band_specs=band_specs,
+            fs=fs,
+            n_fft=n_fft,
+            win_length=win_length,
+            hop_length=hop_length,
+            window_fn=window_fn,
+            wkwargs=wkwargs,
+            power=power,
+            center=center,
+            normalized=normalized,
+            pad_mode=pad_mode,
+            onesided=onesided,
+            n_bands=n_bands,
+        )
+
+        self.bsrnn = MultiSourceMultiMaskBandSplitCoreTransformer(
+            stems=stems,
+            band_specs=self.band_specs,
+            in_channel=in_channel,
+            require_no_overlap=require_no_overlap,
+            require_no_gap=require_no_gap,
+            normalize_channel_independently=normalize_channel_independently,
+            treat_channel_as_feature=treat_channel_as_feature,
+            n_sqm_modules=n_sqm_modules,
+            emb_dim=emb_dim,
+            rnn_dim=rnn_dim,
+            bidirectional=bidirectional,
+            rnn_type=rnn_type,
+            mlp_dim=mlp_dim,
+            cond_dim=cond_dim,
+            hidden_activation=hidden_activation,
+            hidden_activation_kwargs=hidden_activation_kwargs,
+            complex_mask=complex_mask,
+            overlapping_band=self.overlapping_band,
+            freq_weights=self.freq_weights,
+            n_freq=n_fft // 2 + 1,
+            use_freq_weights=use_freq_weights,
+            mult_add_mask=mult_add_mask,
+        )
+
+
+class MultiMaskMultiSourceBandSplitConv(MultiMaskMultiSourceBandSplitBase):
+    def __init__(
+        self,
+        in_channel: int,
+        stems: List[str],
+        band_specs: Union[str, List[Tuple[float, float]]],
+        fs: int = 44100,
+        require_no_overlap: bool = False,
+        require_no_gap: bool = True,
+        normalize_channel_independently: bool = False,
+        treat_channel_as_feature: bool = True,
+        n_sqm_modules: int = 12,
+        emb_dim: int = 128,
+        rnn_dim: int = 256,
+        cond_dim: int = 0,
+        bidirectional: bool = True,
+        rnn_type: str = "LSTM",
+        mlp_dim: int = 512,
+        hidden_activation: str = "Tanh",
+        hidden_activation_kwargs: Optional[Dict] = None,
+        complex_mask: bool = True,
+        n_fft: int = 2048,
+        win_length: Optional[int] = 2048,
+        hop_length: int = 512,
+        window_fn: str = "hann_window",
+        wkwargs: Optional[Dict] = None,
+        power: Optional[int] = None,
+        center: bool = True,
+        normalized: bool = True,
+        pad_mode: str = "constant",
+        onesided: bool = True,
+        n_bands: int = None,
+        use_freq_weights: bool = True,
+        normalize_input: bool = False,
+        mult_add_mask: bool = False,
+    ) -> None:
+        super().__init__(
+            stems=stems,
+            band_specs=band_specs,
+            fs=fs,
+            n_fft=n_fft,
+            win_length=win_length,
+            hop_length=hop_length,
+            window_fn=window_fn,
+            wkwargs=wkwargs,
+            power=power,
+            center=center,
+            normalized=normalized,
+            pad_mode=pad_mode,
+            onesided=onesided,
+            n_bands=n_bands,
+        )
+
+        self.bsrnn = MultiSourceMultiMaskBandSplitCoreConv(
+            stems=stems,
+            band_specs=self.band_specs,
+            in_channel=in_channel,
+            require_no_overlap=require_no_overlap,
+            require_no_gap=require_no_gap,
+            normalize_channel_independently=normalize_channel_independently,
+            treat_channel_as_feature=treat_channel_as_feature,
+            n_sqm_modules=n_sqm_modules,
+            emb_dim=emb_dim,
+            rnn_dim=rnn_dim,
+            bidirectional=bidirectional,
+            rnn_type=rnn_type,
+            mlp_dim=mlp_dim,
+            cond_dim=cond_dim,
+            hidden_activation=hidden_activation,
+            hidden_activation_kwargs=hidden_activation_kwargs,
+            complex_mask=complex_mask,
+            overlapping_band=self.overlapping_band,
+            freq_weights=self.freq_weights,
+            n_freq=n_fft // 2 + 1,
+            use_freq_weights=use_freq_weights,
+            mult_add_mask=mult_add_mask,
+        )
+
+
+class PatchingMaskMultiSourceBandSplitRNN(MultiMaskMultiSourceBandSplitBase):
+    def __init__(
+        self,
+        in_channel: int,
+        stems: List[str],
+        band_specs: Union[str, List[Tuple[float, float]]],
+        kernel_norm_mlp_version: int = 1,
+        mask_kernel_freq: int = 3,
+        mask_kernel_time: int = 3,
+        conv_kernel_freq: int = 1,
+        conv_kernel_time: int = 1,
+        fs: int = 44100,
+        require_no_overlap: bool = False,
+        require_no_gap: bool = True,
+        normalize_channel_independently: bool = False,
+        treat_channel_as_feature: bool = True,
+        n_sqm_modules: int = 12,
+        emb_dim: int = 128,
+        rnn_dim: int = 256,
+        bidirectional: bool = True,
+        rnn_type: str = "LSTM",
+        mlp_dim: int = 512,
+        hidden_activation: str = "Tanh",
+        hidden_activation_kwargs: Optional[Dict] = None,
+        complex_mask: bool = True,
+        n_fft: int = 2048,
+        win_length: Optional[int] = 2048,
+        hop_length: int = 512,
+        window_fn: str = "hann_window",
+        wkwargs: Optional[Dict] = None,
+        power: Optional[int] = None,
+        center: bool = True,
+        normalized: bool = True,
+        pad_mode: str = "constant",
+        onesided: bool = True,
+        n_bands: int = None,
+    ) -> None:
+        super().__init__(
+            stems=stems,
+            band_specs=band_specs,
+            fs=fs,
+            n_fft=n_fft,
+            win_length=win_length,
+            hop_length=hop_length,
+            window_fn=window_fn,
+            wkwargs=wkwargs,
+            power=power,
+            center=center,
+            normalized=normalized,
+            pad_mode=pad_mode,
+            onesided=onesided,
+            n_bands=n_bands,
+        )
+
+        self.bsrnn = MultiSourceMultiPatchingMaskBandSplitCoreRNN(
+            stems=stems,
+            band_specs=self.band_specs,
+            in_channel=in_channel,
+            require_no_overlap=require_no_overlap,
+            require_no_gap=require_no_gap,
+            normalize_channel_independently=normalize_channel_independently,
+            treat_channel_as_feature=treat_channel_as_feature,
+            n_sqm_modules=n_sqm_modules,
+            emb_dim=emb_dim,
+            rnn_dim=rnn_dim,
+            bidirectional=bidirectional,
+            rnn_type=rnn_type,
+            mlp_dim=mlp_dim,
+            hidden_activation=hidden_activation,
+            hidden_activation_kwargs=hidden_activation_kwargs,
+            complex_mask=complex_mask,
+            overlapping_band=self.overlapping_band,
+            freq_weights=self.freq_weights,
+            n_freq=n_fft // 2 + 1,
+            mask_kernel_freq=mask_kernel_freq,
+            mask_kernel_time=mask_kernel_time,
+            conv_kernel_freq=conv_kernel_freq,
+            conv_kernel_time=conv_kernel_time,
+            kernel_norm_mlp_version=kernel_norm_mlp_version,
+        )

mvsepless/models/bandit/core/utils/audio.py

@@ -1,324 +1,324 @@
-from collections import defaultdict
-
-from tqdm.auto import tqdm
-from typing import Callable, Dict, List, Optional, Tuple
-
-import numpy as np
-import torch
-from torch import nn
-from torch.nn import functional as F
-
-
-@torch.jit.script
-def merge(
-    combined: torch.Tensor,
-    original_batch_size: int,
-    n_channel: int,
-    n_chunks: int,
-    chunk_size: int,
-):
-    combined = torch.reshape(
-        combined, (original_batch_size, n_chunks, n_channel, chunk_size)
-    )
-    combined = torch.permute(combined, (0, 2, 3, 1)).reshape(
-        original_batch_size * n_channel, chunk_size, n_chunks
-    )
-
-    return combined
-
-
-@torch.jit.script
-def unfold(
-    padded_audio: torch.Tensor,
-    original_batch_size: int,
-    n_channel: int,
-    chunk_size: int,
-    hop_size: int,
-) -> torch.Tensor:
-
-    unfolded_input = F.unfold(
-        padded_audio[:, :, None, :], kernel_size=(1, chunk_size), stride=(1, hop_size)
-    )
-
-    _, _, n_chunks = unfolded_input.shape
-    unfolded_input = unfolded_input.view(
-        original_batch_size, n_channel, chunk_size, n_chunks
-    )
-    unfolded_input = torch.permute(unfolded_input, (0, 3, 1, 2)).reshape(
-        original_batch_size * n_chunks, n_channel, chunk_size
-    )
-
-    return unfolded_input
-
-
-@torch.jit.script
-def merge_chunks_all(
-    combined: torch.Tensor,
-    original_batch_size: int,
-    n_channel: int,
-    n_samples: int,
-    n_padded_samples: int,
-    n_chunks: int,
-    chunk_size: int,
-    hop_size: int,
-    edge_frame_pad_sizes: Tuple[int, int],
-    standard_window: torch.Tensor,
-    first_window: torch.Tensor,
-    last_window: torch.Tensor,
-):
-    combined = merge(combined, original_batch_size, n_channel, n_chunks, chunk_size)
-
-    combined = combined * standard_window[:, None].to(combined.device)
-
-    combined = F.fold(
-        combined.to(torch.float32),
-        output_size=(1, n_padded_samples),
-        kernel_size=(1, chunk_size),
-        stride=(1, hop_size),
-    )
-
-    combined = combined.view(original_batch_size, n_channel, n_padded_samples)
-
-    pad_front, pad_back = edge_frame_pad_sizes
-    combined = combined[..., pad_front:-pad_back]
-
-    combined = combined[..., :n_samples]
-
-    return combined
-
-
-def merge_chunks_edge(
-    combined: torch.Tensor,
-    original_batch_size: int,
-    n_channel: int,
-    n_samples: int,
-    n_padded_samples: int,
-    n_chunks: int,
-    chunk_size: int,
-    hop_size: int,
-    edge_frame_pad_sizes: Tuple[int, int],
-    standard_window: torch.Tensor,
-    first_window: torch.Tensor,
-    last_window: torch.Tensor,
-):
-    combined = merge(combined, original_batch_size, n_channel, n_chunks, chunk_size)
-
-    combined[..., 0] = combined[..., 0] * first_window
-    combined[..., -1] = combined[..., -1] * last_window
-    combined[..., 1:-1] = combined[..., 1:-1] * standard_window[:, None]
-
-    combined = F.fold(
-        combined,
-        output_size=(1, n_padded_samples),
-        kernel_size=(1, chunk_size),
-        stride=(1, hop_size),
-    )
-
-    combined = combined.view(original_batch_size, n_channel, n_padded_samples)
-
-    combined = combined[..., :n_samples]
-
-    return combined
-
-
-class BaseFader(nn.Module):
-    def __init__(
-        self,
-        chunk_size_second: float,
-        hop_size_second: float,
-        fs: int,
-        fade_edge_frames: bool,
-        batch_size: int,
-    ) -> None:
-        super().__init__()
-
-        self.chunk_size = int(chunk_size_second * fs)
-        self.hop_size = int(hop_size_second * fs)
-        self.overlap_size = self.chunk_size - self.hop_size
-        self.fade_edge_frames = fade_edge_frames
-        self.batch_size = batch_size
-
-    def prepare(self, audio):
-
-        if self.fade_edge_frames:
-            audio = F.pad(audio, self.edge_frame_pad_sizes, mode="reflect")
-
-        n_samples = audio.shape[-1]
-        n_chunks = int(np.ceil((n_samples - self.chunk_size) / self.hop_size) + 1)
-
-        padded_size = (n_chunks - 1) * self.hop_size + self.chunk_size
-        pad_size = padded_size - n_samples
-
-        padded_audio = F.pad(audio, (0, pad_size))
-
-        return padded_audio, n_chunks
-
-    def forward(
-        self,
-        audio: torch.Tensor,
-        model_fn: Callable[[torch.Tensor], Dict[str, torch.Tensor]],
-    ):
-
-        original_dtype = audio.dtype
-        original_device = audio.device
-
-        audio = audio.to("cpu")
-
-        original_batch_size, n_channel, n_samples = audio.shape
-        padded_audio, n_chunks = self.prepare(audio)
-        del audio
-        n_padded_samples = padded_audio.shape[-1]
-
-        if n_channel > 1:
-            padded_audio = padded_audio.view(
-                original_batch_size * n_channel, 1, n_padded_samples
-            )
-
-        unfolded_input = unfold(
-            padded_audio, original_batch_size, n_channel, self.chunk_size, self.hop_size
-        )
-
-        n_total_chunks, n_channel, chunk_size = unfolded_input.shape
-
-        n_batch = np.ceil(n_total_chunks / self.batch_size).astype(int)
-
-        chunks_in = [
-            unfolded_input[b * self.batch_size : (b + 1) * self.batch_size, ...].clone()
-            for b in range(n_batch)
-        ]
-
-        all_chunks_out = defaultdict(
-            lambda: torch.zeros_like(unfolded_input, device="cpu")
-        )
-
-        for b, cin in enumerate(chunks_in):
-            if torch.allclose(cin, torch.tensor(0.0)):
-                del cin
-                continue
-
-            chunks_out = model_fn(cin.to(original_device))
-            del cin
-            for s, c in chunks_out.items():
-                all_chunks_out[s][
-                    b * self.batch_size : (b + 1) * self.batch_size, ...
-                ] = c.cpu()
-            del chunks_out
-
-        del unfolded_input
-        del padded_audio
-
-        if self.fade_edge_frames:
-            fn = merge_chunks_all
-        else:
-            fn = merge_chunks_edge
-        outputs = {}
-
-        torch.cuda.empty_cache()
-
-        for s, c in all_chunks_out.items():
-            combined: torch.Tensor = fn(
-                c,
-                original_batch_size,
-                n_channel,
-                n_samples,
-                n_padded_samples,
-                n_chunks,
-                self.chunk_size,
-                self.hop_size,
-                self.edge_frame_pad_sizes,
-                self.standard_window,
-                self.__dict__.get("first_window", self.standard_window),
-                self.__dict__.get("last_window", self.standard_window),
-            )
-
-            outputs[s] = combined.to(dtype=original_dtype, device=original_device)
-
-        return {"audio": outputs}
-
-
-class LinearFader(BaseFader):
-    def __init__(
-        self,
-        chunk_size_second: float,
-        hop_size_second: float,
-        fs: int,
-        fade_edge_frames: bool = False,
-        batch_size: int = 1,
-    ) -> None:
-
-        assert hop_size_second >= chunk_size_second / 2
-
-        super().__init__(
-            chunk_size_second=chunk_size_second,
-            hop_size_second=hop_size_second,
-            fs=fs,
-            fade_edge_frames=fade_edge_frames,
-            batch_size=batch_size,
-        )
-
-        in_fade = torch.linspace(0.0, 1.0, self.overlap_size + 1)[:-1]
-        out_fade = torch.linspace(1.0, 0.0, self.overlap_size + 1)[1:]
-        center_ones = torch.ones(self.chunk_size - 2 * self.overlap_size)
-        inout_ones = torch.ones(self.overlap_size)
-
-        self.register_buffer(
-            "standard_window", torch.concat([in_fade, center_ones, out_fade])
-        )
-
-        self.fade_edge_frames = fade_edge_frames
-        self.edge_frame_pad_size = (self.overlap_size, self.overlap_size)
-
-        if not self.fade_edge_frames:
-            self.first_window = nn.Parameter(
-                torch.concat([inout_ones, center_ones, out_fade]), requires_grad=False
-            )
-            self.last_window = nn.Parameter(
-                torch.concat([in_fade, center_ones, inout_ones]), requires_grad=False
-            )
-
-
-class OverlapAddFader(BaseFader):
-    def __init__(
-        self,
-        window_type: str,
-        chunk_size_second: float,
-        hop_size_second: float,
-        fs: int,
-        batch_size: int = 1,
-    ) -> None:
-        assert (chunk_size_second / hop_size_second) % 2 == 0
-        assert int(chunk_size_second * fs) % 2 == 0
-
-        super().__init__(
-            chunk_size_second=chunk_size_second,
-            hop_size_second=hop_size_second,
-            fs=fs,
-            fade_edge_frames=True,
-            batch_size=batch_size,
-        )
-
-        self.hop_multiplier = self.chunk_size / (2 * self.hop_size)
-
-        self.edge_frame_pad_sizes = (2 * self.overlap_size, 2 * self.overlap_size)
-
-        self.register_buffer(
-            "standard_window",
-            torch.windows.__dict__[window_type](
-                self.chunk_size,
-                sym=False,
-            )
-            / self.hop_multiplier,
-        )
-
-
-if __name__ == "__main__":
-    import torchaudio as ta
-
-    fs = 44100
-    ola = OverlapAddFader("hann", 6.0, 1.0, fs, batch_size=16)
-    audio_, _ = ta.load(
-        "$DATA_ROOT/MUSDB18/HQ/canonical/test/BKS - Too " "Much/vocals.wav"
-    )
-    audio_ = audio_[None, ...]
-    out = ola(audio_, lambda x: {"stem": x})["audio"]["stem"]
-    print(torch.allclose(out, audio_))
+from collections import defaultdict
+
+from tqdm.auto import tqdm
+from typing import Callable, Dict, List, Optional, Tuple
+
+import numpy as np
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+@torch.jit.script
+def merge(
+    combined: torch.Tensor,
+    original_batch_size: int,
+    n_channel: int,
+    n_chunks: int,
+    chunk_size: int,
+):
+    combined = torch.reshape(
+        combined, (original_batch_size, n_chunks, n_channel, chunk_size)
+    )
+    combined = torch.permute(combined, (0, 2, 3, 1)).reshape(
+        original_batch_size * n_channel, chunk_size, n_chunks
+    )
+
+    return combined
+
+
+@torch.jit.script
+def unfold(
+    padded_audio: torch.Tensor,
+    original_batch_size: int,
+    n_channel: int,
+    chunk_size: int,
+    hop_size: int,
+) -> torch.Tensor:
+
+    unfolded_input = F.unfold(
+        padded_audio[:, :, None, :], kernel_size=(1, chunk_size), stride=(1, hop_size)
+    )
+
+    _, _, n_chunks = unfolded_input.shape
+    unfolded_input = unfolded_input.view(
+        original_batch_size, n_channel, chunk_size, n_chunks
+    )
+    unfolded_input = torch.permute(unfolded_input, (0, 3, 1, 2)).reshape(
+        original_batch_size * n_chunks, n_channel, chunk_size
+    )
+
+    return unfolded_input
+
+
+@torch.jit.script
+def merge_chunks_all(
+    combined: torch.Tensor,
+    original_batch_size: int,
+    n_channel: int,
+    n_samples: int,
+    n_padded_samples: int,
+    n_chunks: int,
+    chunk_size: int,
+    hop_size: int,
+    edge_frame_pad_sizes: Tuple[int, int],
+    standard_window: torch.Tensor,
+    first_window: torch.Tensor,
+    last_window: torch.Tensor,
+):
+    combined = merge(combined, original_batch_size, n_channel, n_chunks, chunk_size)
+
+    combined = combined * standard_window[:, None].to(combined.device)
+
+    combined = F.fold(
+        combined.to(torch.float32),
+        output_size=(1, n_padded_samples),
+        kernel_size=(1, chunk_size),
+        stride=(1, hop_size),
+    )
+
+    combined = combined.view(original_batch_size, n_channel, n_padded_samples)
+
+    pad_front, pad_back = edge_frame_pad_sizes
+    combined = combined[..., pad_front:-pad_back]
+
+    combined = combined[..., :n_samples]
+
+    return combined
+
+
+def merge_chunks_edge(
+    combined: torch.Tensor,
+    original_batch_size: int,
+    n_channel: int,
+    n_samples: int,
+    n_padded_samples: int,
+    n_chunks: int,
+    chunk_size: int,
+    hop_size: int,
+    edge_frame_pad_sizes: Tuple[int, int],
+    standard_window: torch.Tensor,
+    first_window: torch.Tensor,
+    last_window: torch.Tensor,
+):
+    combined = merge(combined, original_batch_size, n_channel, n_chunks, chunk_size)
+
+    combined[..., 0] = combined[..., 0] * first_window
+    combined[..., -1] = combined[..., -1] * last_window
+    combined[..., 1:-1] = combined[..., 1:-1] * standard_window[:, None]
+
+    combined = F.fold(
+        combined,
+        output_size=(1, n_padded_samples),
+        kernel_size=(1, chunk_size),
+        stride=(1, hop_size),
+    )
+
+    combined = combined.view(original_batch_size, n_channel, n_padded_samples)
+
+    combined = combined[..., :n_samples]
+
+    return combined
+
+
+class BaseFader(nn.Module):
+    def __init__(
+        self,
+        chunk_size_second: float,
+        hop_size_second: float,
+        fs: int,
+        fade_edge_frames: bool,
+        batch_size: int,
+    ) -> None:
+        super().__init__()
+
+        self.chunk_size = int(chunk_size_second * fs)
+        self.hop_size = int(hop_size_second * fs)
+        self.overlap_size = self.chunk_size - self.hop_size
+        self.fade_edge_frames = fade_edge_frames
+        self.batch_size = batch_size
+
+    def prepare(self, audio):
+
+        if self.fade_edge_frames:
+            audio = F.pad(audio, self.edge_frame_pad_sizes, mode="reflect")
+
+        n_samples = audio.shape[-1]
+        n_chunks = int(np.ceil((n_samples - self.chunk_size) / self.hop_size) + 1)
+
+        padded_size = (n_chunks - 1) * self.hop_size + self.chunk_size
+        pad_size = padded_size - n_samples
+
+        padded_audio = F.pad(audio, (0, pad_size))
+
+        return padded_audio, n_chunks
+
+    def forward(
+        self,
+        audio: torch.Tensor,
+        model_fn: Callable[[torch.Tensor], Dict[str, torch.Tensor]],
+    ):
+
+        original_dtype = audio.dtype
+        original_device = audio.device
+
+        audio = audio.to("cpu")
+
+        original_batch_size, n_channel, n_samples = audio.shape
+        padded_audio, n_chunks = self.prepare(audio)
+        del audio
+        n_padded_samples = padded_audio.shape[-1]
+
+        if n_channel > 1:
+            padded_audio = padded_audio.view(
+                original_batch_size * n_channel, 1, n_padded_samples
+            )
+
+        unfolded_input = unfold(
+            padded_audio, original_batch_size, n_channel, self.chunk_size, self.hop_size
+        )
+
+        n_total_chunks, n_channel, chunk_size = unfolded_input.shape
+
+        n_batch = np.ceil(n_total_chunks / self.batch_size).astype(int)
+
+        chunks_in = [
+            unfolded_input[b * self.batch_size : (b + 1) * self.batch_size, ...].clone()
+            for b in range(n_batch)
+        ]
+
+        all_chunks_out = defaultdict(
+            lambda: torch.zeros_like(unfolded_input, device="cpu")
+        )
+
+        for b, cin in enumerate(chunks_in):
+            if torch.allclose(cin, torch.tensor(0.0)):
+                del cin
+                continue
+
+            chunks_out = model_fn(cin.to(original_device))
+            del cin
+            for s, c in chunks_out.items():
+                all_chunks_out[s][
+                    b * self.batch_size : (b + 1) * self.batch_size, ...
+                ] = c.cpu()
+            del chunks_out
+
+        del unfolded_input
+        del padded_audio
+
+        if self.fade_edge_frames:
+            fn = merge_chunks_all
+        else:
+            fn = merge_chunks_edge
+        outputs = {}
+
+        torch.cuda.empty_cache()
+
+        for s, c in all_chunks_out.items():
+            combined: torch.Tensor = fn(
+                c,
+                original_batch_size,
+                n_channel,
+                n_samples,
+                n_padded_samples,
+                n_chunks,
+                self.chunk_size,
+                self.hop_size,
+                self.edge_frame_pad_sizes,
+                self.standard_window,
+                self.__dict__.get("first_window", self.standard_window),
+                self.__dict__.get("last_window", self.standard_window),
+            )
+
+            outputs[s] = combined.to(dtype=original_dtype, device=original_device)
+
+        return {"audio": outputs}
+
+
+class LinearFader(BaseFader):
+    def __init__(
+        self,
+        chunk_size_second: float,
+        hop_size_second: float,
+        fs: int,
+        fade_edge_frames: bool = False,
+        batch_size: int = 1,
+    ) -> None:
+
+        assert hop_size_second >= chunk_size_second / 2
+
+        super().__init__(
+            chunk_size_second=chunk_size_second,
+            hop_size_second=hop_size_second,
+            fs=fs,
+            fade_edge_frames=fade_edge_frames,
+            batch_size=batch_size,
+        )
+
+        in_fade = torch.linspace(0.0, 1.0, self.overlap_size + 1)[:-1]
+        out_fade = torch.linspace(1.0, 0.0, self.overlap_size + 1)[1:]
+        center_ones = torch.ones(self.chunk_size - 2 * self.overlap_size)
+        inout_ones = torch.ones(self.overlap_size)
+
+        self.register_buffer(
+            "standard_window", torch.concat([in_fade, center_ones, out_fade])
+        )
+
+        self.fade_edge_frames = fade_edge_frames
+        self.edge_frame_pad_size = (self.overlap_size, self.overlap_size)
+
+        if not self.fade_edge_frames:
+            self.first_window = nn.Parameter(
+                torch.concat([inout_ones, center_ones, out_fade]), requires_grad=False
+            )
+            self.last_window = nn.Parameter(
+                torch.concat([in_fade, center_ones, inout_ones]), requires_grad=False
+            )
+
+
+class OverlapAddFader(BaseFader):
+    def __init__(
+        self,
+        window_type: str,
+        chunk_size_second: float,
+        hop_size_second: float,
+        fs: int,
+        batch_size: int = 1,
+    ) -> None:
+        assert (chunk_size_second / hop_size_second) % 2 == 0
+        assert int(chunk_size_second * fs) % 2 == 0
+
+        super().__init__(
+            chunk_size_second=chunk_size_second,
+            hop_size_second=hop_size_second,
+            fs=fs,
+            fade_edge_frames=True,
+            batch_size=batch_size,
+        )
+
+        self.hop_multiplier = self.chunk_size / (2 * self.hop_size)
+
+        self.edge_frame_pad_sizes = (2 * self.overlap_size, 2 * self.overlap_size)
+
+        self.register_buffer(
+            "standard_window",
+            torch.windows.__dict__[window_type](
+                self.chunk_size,
+                sym=False,
+            )
+            / self.hop_multiplier,
+        )
+
+
+if __name__ == "__main__":
+    import torchaudio as ta
+
+    fs = 44100
+    ola = OverlapAddFader("hann", 6.0, 1.0, fs, batch_size=16)
+    audio_, _ = ta.load(
+        "$DATA_ROOT/MUSDB18/HQ/canonical/test/BKS - Too " "Much/vocals.wav"
+    )
+    audio_ = audio_[None, ...]
+    out = ola(audio_, lambda x: {"stem": x})["audio"]["stem"]
+    print(torch.allclose(out, audio_))

mvsepless/models/bandit/model_from_config.py

@@ -1,26 +1,26 @@
-import sys
-import os.path
-import torch
-
-import yaml
-from ml_collections import ConfigDict
-
-torch.set_float32_matmul_precision("medium")
-
-
-def get_model(
-    config_path,
-    weights_path,
-    device,
-):
-    from .core.model import MultiMaskMultiSourceBandSplitRNNSimple
-
-    f = open(config_path)
-    config = ConfigDict(yaml.load(f, Loader=yaml.FullLoader))
-    f.close()
-
-    model = MultiMaskMultiSourceBandSplitRNNSimple(**config.model)
-    d = torch.load(code_path + "model_bandit_plus_dnr_sdr_11.47.chpt")
-    model.load_state_dict(d)
-    model.to(device)
-    return model, config
+import sys
+import os.path
+import torch
+
+import yaml
+from ml_collections import ConfigDict
+
+torch.set_float32_matmul_precision("medium")
+
+
+def get_model(
+    config_path,
+    weights_path,
+    device,
+):
+    from .core.model import MultiMaskMultiSourceBandSplitRNNSimple
+
+    f = open(config_path)
+    config = ConfigDict(yaml.load(f, Loader=yaml.FullLoader))
+    f.close()
+
+    model = MultiMaskMultiSourceBandSplitRNNSimple(**config.model)
+    d = torch.load(code_path + "model_bandit_plus_dnr_sdr_11.47.chpt")
+    model.load_state_dict(d)
+    model.to(device)
+    return model, config

mvsepless/models/bandit_v2/bandit.py

@@ -1,360 +1,360 @@
-from typing import Dict, List, Optional
-
-import torch
-import torchaudio as ta
-from torch import nn
-import pytorch_lightning as pl
-
-from .bandsplit import BandSplitModule
-from .maskestim import OverlappingMaskEstimationModule
-from .tfmodel import SeqBandModellingModule
-from .utils import MusicalBandsplitSpecification
-
-
-class BaseEndToEndModule(pl.LightningModule):
-    def __init__(
-        self,
-    ) -> None:
-        super().__init__()
-
-
-class BaseBandit(BaseEndToEndModule):
-    def __init__(
-        self,
-        in_channels: int,
-        fs: int,
-        band_type: str = "musical",
-        n_bands: int = 64,
-        require_no_overlap: bool = False,
-        require_no_gap: bool = True,
-        normalize_channel_independently: bool = False,
-        treat_channel_as_feature: bool = True,
-        n_sqm_modules: int = 12,
-        emb_dim: int = 128,
-        rnn_dim: int = 256,
-        bidirectional: bool = True,
-        rnn_type: str = "LSTM",
-        n_fft: int = 2048,
-        win_length: Optional[int] = 2048,
-        hop_length: int = 512,
-        window_fn: str = "hann_window",
-        wkwargs: Optional[Dict] = None,
-        power: Optional[int] = None,
-        center: bool = True,
-        normalized: bool = True,
-        pad_mode: str = "constant",
-        onesided: bool = True,
-    ):
-        super().__init__()
-
-        self.in_channels = in_channels
-
-        self.instantitate_spectral(
-            n_fft=n_fft,
-            win_length=win_length,
-            hop_length=hop_length,
-            window_fn=window_fn,
-            wkwargs=wkwargs,
-            power=power,
-            normalized=normalized,
-            center=center,
-            pad_mode=pad_mode,
-            onesided=onesided,
-        )
-
-        self.instantiate_bandsplit(
-            in_channels=in_channels,
-            band_type=band_type,
-            n_bands=n_bands,
-            require_no_overlap=require_no_overlap,
-            require_no_gap=require_no_gap,
-            normalize_channel_independently=normalize_channel_independently,
-            treat_channel_as_feature=treat_channel_as_feature,
-            emb_dim=emb_dim,
-            n_fft=n_fft,
-            fs=fs,
-        )
-
-        self.instantiate_tf_modelling(
-            n_sqm_modules=n_sqm_modules,
-            emb_dim=emb_dim,
-            rnn_dim=rnn_dim,
-            bidirectional=bidirectional,
-            rnn_type=rnn_type,
-        )
-
-    def instantitate_spectral(
-        self,
-        n_fft: int = 2048,
-        win_length: Optional[int] = 2048,
-        hop_length: int = 512,
-        window_fn: str = "hann_window",
-        wkwargs: Optional[Dict] = None,
-        power: Optional[int] = None,
-        normalized: bool = True,
-        center: bool = True,
-        pad_mode: str = "constant",
-        onesided: bool = True,
-    ):
-        assert power is None
-
-        window_fn = torch.__dict__[window_fn]
-
-        self.stft = ta.transforms.Spectrogram(
-            n_fft=n_fft,
-            win_length=win_length,
-            hop_length=hop_length,
-            pad_mode=pad_mode,
-            pad=0,
-            window_fn=window_fn,
-            wkwargs=wkwargs,
-            power=power,
-            normalized=normalized,
-            center=center,
-            onesided=onesided,
-        )
-
-        self.istft = ta.transforms.InverseSpectrogram(
-            n_fft=n_fft,
-            win_length=win_length,
-            hop_length=hop_length,
-            pad_mode=pad_mode,
-            pad=0,
-            window_fn=window_fn,
-            wkwargs=wkwargs,
-            normalized=normalized,
-            center=center,
-            onesided=onesided,
-        )
-
-    def instantiate_bandsplit(
-        self,
-        in_channels: int,
-        band_type: str = "musical",
-        n_bands: int = 64,
-        require_no_overlap: bool = False,
-        require_no_gap: bool = True,
-        normalize_channel_independently: bool = False,
-        treat_channel_as_feature: bool = True,
-        emb_dim: int = 128,
-        n_fft: int = 2048,
-        fs: int = 44100,
-    ):
-        assert band_type == "musical"
-
-        self.band_specs = MusicalBandsplitSpecification(
-            nfft=n_fft, fs=fs, n_bands=n_bands
-        )
-
-        self.band_split = BandSplitModule(
-            in_channels=in_channels,
-            band_specs=self.band_specs.get_band_specs(),
-            require_no_overlap=require_no_overlap,
-            require_no_gap=require_no_gap,
-            normalize_channel_independently=normalize_channel_independently,
-            treat_channel_as_feature=treat_channel_as_feature,
-            emb_dim=emb_dim,
-        )
-
-    def instantiate_tf_modelling(
-        self,
-        n_sqm_modules: int = 12,
-        emb_dim: int = 128,
-        rnn_dim: int = 256,
-        bidirectional: bool = True,
-        rnn_type: str = "LSTM",
-    ):
-        try:
-            self.tf_model = torch.compile(
-                SeqBandModellingModule(
-                    n_modules=n_sqm_modules,
-                    emb_dim=emb_dim,
-                    rnn_dim=rnn_dim,
-                    bidirectional=bidirectional,
-                    rnn_type=rnn_type,
-                ),
-                disable=True,
-            )
-        except Exception as e:
-            self.tf_model = SeqBandModellingModule(
-                n_modules=n_sqm_modules,
-                emb_dim=emb_dim,
-                rnn_dim=rnn_dim,
-                bidirectional=bidirectional,
-                rnn_type=rnn_type,
-            )
-
-    def mask(self, x, m):
-        return x * m
-
-    def forward(self, batch, mode="train"):
-        init_shape = batch.shape
-        if not isinstance(batch, dict):
-            mono = batch.view(-1, 1, batch.shape[-1])
-            batch = {"mixture": {"audio": mono}}
-
-        with torch.no_grad():
-            mixture = batch["mixture"]["audio"]
-
-            x = self.stft(mixture)
-            batch["mixture"]["spectrogram"] = x
-
-            if "sources" in batch.keys():
-                for stem in batch["sources"].keys():
-                    s = batch["sources"][stem]["audio"]
-                    s = self.stft(s)
-                    batch["sources"][stem]["spectrogram"] = s
-
-        batch = self.separate(batch)
-
-        if 1:
-            b = []
-            for s in self.stems:
-                r = batch["estimates"][s]["audio"].view(
-                    -1, init_shape[1], init_shape[2]
-                )
-                b.append(r)
-            batch = torch.stack(b, dim=1)
-        return batch
-
-    def encode(self, batch):
-        x = batch["mixture"]["spectrogram"]
-        length = batch["mixture"]["audio"].shape[-1]
-
-        z = self.band_split(x)
-        q = self.tf_model(z)
-
-        return x, q, length
-
-    def separate(self, batch):
-        raise NotImplementedError
-
-
-class Bandit(BaseBandit):
-    def __init__(
-        self,
-        in_channels: int,
-        stems: List[str],
-        band_type: str = "musical",
-        n_bands: int = 64,
-        require_no_overlap: bool = False,
-        require_no_gap: bool = True,
-        normalize_channel_independently: bool = False,
-        treat_channel_as_feature: bool = True,
-        n_sqm_modules: int = 12,
-        emb_dim: int = 128,
-        rnn_dim: int = 256,
-        bidirectional: bool = True,
-        rnn_type: str = "LSTM",
-        mlp_dim: int = 512,
-        hidden_activation: str = "Tanh",
-        hidden_activation_kwargs: Dict | None = None,
-        complex_mask: bool = True,
-        use_freq_weights: bool = True,
-        n_fft: int = 2048,
-        win_length: int | None = 2048,
-        hop_length: int = 512,
-        window_fn: str = "hann_window",
-        wkwargs: Dict | None = None,
-        power: int | None = None,
-        center: bool = True,
-        normalized: bool = True,
-        pad_mode: str = "constant",
-        onesided: bool = True,
-        fs: int = 44100,
-        stft_precisions="32",
-        bandsplit_precisions="bf16",
-        tf_model_precisions="bf16",
-        mask_estim_precisions="bf16",
-    ):
-        super().__init__(
-            in_channels=in_channels,
-            band_type=band_type,
-            n_bands=n_bands,
-            require_no_overlap=require_no_overlap,
-            require_no_gap=require_no_gap,
-            normalize_channel_independently=normalize_channel_independently,
-            treat_channel_as_feature=treat_channel_as_feature,
-            n_sqm_modules=n_sqm_modules,
-            emb_dim=emb_dim,
-            rnn_dim=rnn_dim,
-            bidirectional=bidirectional,
-            rnn_type=rnn_type,
-            n_fft=n_fft,
-            win_length=win_length,
-            hop_length=hop_length,
-            window_fn=window_fn,
-            wkwargs=wkwargs,
-            power=power,
-            center=center,
-            normalized=normalized,
-            pad_mode=pad_mode,
-            onesided=onesided,
-            fs=fs,
-        )
-
-        self.stems = stems
-
-        self.instantiate_mask_estim(
-            in_channels=in_channels,
-            stems=stems,
-            emb_dim=emb_dim,
-            mlp_dim=mlp_dim,
-            hidden_activation=hidden_activation,
-            hidden_activation_kwargs=hidden_activation_kwargs,
-            complex_mask=complex_mask,
-            n_freq=n_fft // 2 + 1,
-            use_freq_weights=use_freq_weights,
-        )
-
-    def instantiate_mask_estim(
-        self,
-        in_channels: int,
-        stems: List[str],
-        emb_dim: int,
-        mlp_dim: int,
-        hidden_activation: str,
-        hidden_activation_kwargs: Optional[Dict] = None,
-        complex_mask: bool = True,
-        n_freq: Optional[int] = None,
-        use_freq_weights: bool = False,
-    ):
-        if hidden_activation_kwargs is None:
-            hidden_activation_kwargs = {}
-
-        assert n_freq is not None
-
-        self.mask_estim = nn.ModuleDict(
-            {
-                stem: OverlappingMaskEstimationModule(
-                    band_specs=self.band_specs.get_band_specs(),
-                    freq_weights=self.band_specs.get_freq_weights(),
-                    n_freq=n_freq,
-                    emb_dim=emb_dim,
-                    mlp_dim=mlp_dim,
-                    in_channels=in_channels,
-                    hidden_activation=hidden_activation,
-                    hidden_activation_kwargs=hidden_activation_kwargs,
-                    complex_mask=complex_mask,
-                    use_freq_weights=use_freq_weights,
-                )
-                for stem in stems
-            }
-        )
-
-    def separate(self, batch):
-        batch["estimates"] = {}
-
-        x, q, length = self.encode(batch)
-
-        for stem, mem in self.mask_estim.items():
-            m = mem(q)
-
-            s = self.mask(x, m.to(x.dtype))
-            s = torch.reshape(s, x.shape)
-            batch["estimates"][stem] = {
-                "audio": self.istft(s, length),
-                "spectrogram": s,
-            }
-
-        return batch
+from typing import Dict, List, Optional
+
+import torch
+import torchaudio as ta
+from torch import nn
+import pytorch_lightning as pl
+
+from .bandsplit import BandSplitModule
+from .maskestim import OverlappingMaskEstimationModule
+from .tfmodel import SeqBandModellingModule
+from .utils import MusicalBandsplitSpecification
+
+
+class BaseEndToEndModule(pl.LightningModule):
+    def __init__(
+        self,
+    ) -> None:
+        super().__init__()
+
+
+class BaseBandit(BaseEndToEndModule):
+    def __init__(
+        self,
+        in_channels: int,
+        fs: int,
+        band_type: str = "musical",
+        n_bands: int = 64,
+        require_no_overlap: bool = False,
+        require_no_gap: bool = True,
+        normalize_channel_independently: bool = False,
+        treat_channel_as_feature: bool = True,
+        n_sqm_modules: int = 12,
+        emb_dim: int = 128,
+        rnn_dim: int = 256,
+        bidirectional: bool = True,
+        rnn_type: str = "LSTM",
+        n_fft: int = 2048,
+        win_length: Optional[int] = 2048,
+        hop_length: int = 512,
+        window_fn: str = "hann_window",
+        wkwargs: Optional[Dict] = None,
+        power: Optional[int] = None,
+        center: bool = True,
+        normalized: bool = True,
+        pad_mode: str = "constant",
+        onesided: bool = True,
+    ):
+        super().__init__()
+
+        self.in_channels = in_channels
+
+        self.instantitate_spectral(
+            n_fft=n_fft,
+            win_length=win_length,
+            hop_length=hop_length,
+            window_fn=window_fn,
+            wkwargs=wkwargs,
+            power=power,
+            normalized=normalized,
+            center=center,
+            pad_mode=pad_mode,
+            onesided=onesided,
+        )
+
+        self.instantiate_bandsplit(
+            in_channels=in_channels,
+            band_type=band_type,
+            n_bands=n_bands,
+            require_no_overlap=require_no_overlap,
+            require_no_gap=require_no_gap,
+            normalize_channel_independently=normalize_channel_independently,
+            treat_channel_as_feature=treat_channel_as_feature,
+            emb_dim=emb_dim,
+            n_fft=n_fft,
+            fs=fs,
+        )
+
+        self.instantiate_tf_modelling(
+            n_sqm_modules=n_sqm_modules,
+            emb_dim=emb_dim,
+            rnn_dim=rnn_dim,
+            bidirectional=bidirectional,
+            rnn_type=rnn_type,
+        )
+
+    def instantitate_spectral(
+        self,
+        n_fft: int = 2048,
+        win_length: Optional[int] = 2048,
+        hop_length: int = 512,
+        window_fn: str = "hann_window",
+        wkwargs: Optional[Dict] = None,
+        power: Optional[int] = None,
+        normalized: bool = True,
+        center: bool = True,
+        pad_mode: str = "constant",
+        onesided: bool = True,
+    ):
+        assert power is None
+
+        window_fn = torch.__dict__[window_fn]
+
+        self.stft = ta.transforms.Spectrogram(
+            n_fft=n_fft,
+            win_length=win_length,
+            hop_length=hop_length,
+            pad_mode=pad_mode,
+            pad=0,
+            window_fn=window_fn,
+            wkwargs=wkwargs,
+            power=power,
+            normalized=normalized,
+            center=center,
+            onesided=onesided,
+        )
+
+        self.istft = ta.transforms.InverseSpectrogram(
+            n_fft=n_fft,
+            win_length=win_length,
+            hop_length=hop_length,
+            pad_mode=pad_mode,
+            pad=0,
+            window_fn=window_fn,
+            wkwargs=wkwargs,
+            normalized=normalized,
+            center=center,
+            onesided=onesided,
+        )
+
+    def instantiate_bandsplit(
+        self,
+        in_channels: int,
+        band_type: str = "musical",
+        n_bands: int = 64,
+        require_no_overlap: bool = False,
+        require_no_gap: bool = True,
+        normalize_channel_independently: bool = False,
+        treat_channel_as_feature: bool = True,
+        emb_dim: int = 128,
+        n_fft: int = 2048,
+        fs: int = 44100,
+    ):
+        assert band_type == "musical"
+
+        self.band_specs = MusicalBandsplitSpecification(
+            nfft=n_fft, fs=fs, n_bands=n_bands
+        )
+
+        self.band_split = BandSplitModule(
+            in_channels=in_channels,
+            band_specs=self.band_specs.get_band_specs(),
+            require_no_overlap=require_no_overlap,
+            require_no_gap=require_no_gap,
+            normalize_channel_independently=normalize_channel_independently,
+            treat_channel_as_feature=treat_channel_as_feature,
+            emb_dim=emb_dim,
+        )
+
+    def instantiate_tf_modelling(
+        self,
+        n_sqm_modules: int = 12,
+        emb_dim: int = 128,
+        rnn_dim: int = 256,
+        bidirectional: bool = True,
+        rnn_type: str = "LSTM",
+    ):
+        try:
+            self.tf_model = torch.compile(
+                SeqBandModellingModule(
+                    n_modules=n_sqm_modules,
+                    emb_dim=emb_dim,
+                    rnn_dim=rnn_dim,
+                    bidirectional=bidirectional,
+                    rnn_type=rnn_type,
+                ),
+                disable=True,
+            )
+        except Exception as e:
+            self.tf_model = SeqBandModellingModule(
+                n_modules=n_sqm_modules,
+                emb_dim=emb_dim,
+                rnn_dim=rnn_dim,
+                bidirectional=bidirectional,
+                rnn_type=rnn_type,
+            )
+
+    def mask(self, x, m):
+        return x * m
+
+    def forward(self, batch, mode="train"):
+        init_shape = batch.shape
+        if not isinstance(batch, dict):
+            mono = batch.view(-1, 1, batch.shape[-1])
+            batch = {"mixture": {"audio": mono}}
+
+        with torch.no_grad():
+            mixture = batch["mixture"]["audio"]
+
+            x = self.stft(mixture)
+            batch["mixture"]["spectrogram"] = x
+
+            if "sources" in batch.keys():
+                for stem in batch["sources"].keys():
+                    s = batch["sources"][stem]["audio"]
+                    s = self.stft(s)
+                    batch["sources"][stem]["spectrogram"] = s
+
+        batch = self.separate(batch)
+
+        if 1:
+            b = []
+            for s in self.stems:
+                r = batch["estimates"][s]["audio"].view(
+                    -1, init_shape[1], init_shape[2]
+                )
+                b.append(r)
+            batch = torch.stack(b, dim=1)
+        return batch
+
+    def encode(self, batch):
+        x = batch["mixture"]["spectrogram"]
+        length = batch["mixture"]["audio"].shape[-1]
+
+        z = self.band_split(x)
+        q = self.tf_model(z)
+
+        return x, q, length
+
+    def separate(self, batch):
+        raise NotImplementedError
+
+
+class Bandit(BaseBandit):
+    def __init__(
+        self,
+        in_channels: int,
+        stems: List[str],
+        band_type: str = "musical",
+        n_bands: int = 64,
+        require_no_overlap: bool = False,
+        require_no_gap: bool = True,
+        normalize_channel_independently: bool = False,
+        treat_channel_as_feature: bool = True,
+        n_sqm_modules: int = 12,
+        emb_dim: int = 128,
+        rnn_dim: int = 256,
+        bidirectional: bool = True,
+        rnn_type: str = "LSTM",
+        mlp_dim: int = 512,
+        hidden_activation: str = "Tanh",
+        hidden_activation_kwargs: Dict | None = None,
+        complex_mask: bool = True,
+        use_freq_weights: bool = True,
+        n_fft: int = 2048,
+        win_length: int | None = 2048,
+        hop_length: int = 512,
+        window_fn: str = "hann_window",
+        wkwargs: Dict | None = None,
+        power: int | None = None,
+        center: bool = True,
+        normalized: bool = True,
+        pad_mode: str = "constant",
+        onesided: bool = True,
+        fs: int = 44100,
+        stft_precisions="32",
+        bandsplit_precisions="bf16",
+        tf_model_precisions="bf16",
+        mask_estim_precisions="bf16",
+    ):
+        super().__init__(
+            in_channels=in_channels,
+            band_type=band_type,
+            n_bands=n_bands,
+            require_no_overlap=require_no_overlap,
+            require_no_gap=require_no_gap,
+            normalize_channel_independently=normalize_channel_independently,
+            treat_channel_as_feature=treat_channel_as_feature,
+            n_sqm_modules=n_sqm_modules,
+            emb_dim=emb_dim,
+            rnn_dim=rnn_dim,
+            bidirectional=bidirectional,
+            rnn_type=rnn_type,
+            n_fft=n_fft,
+            win_length=win_length,
+            hop_length=hop_length,
+            window_fn=window_fn,
+            wkwargs=wkwargs,
+            power=power,
+            center=center,
+            normalized=normalized,
+            pad_mode=pad_mode,
+            onesided=onesided,
+            fs=fs,
+        )
+
+        self.stems = stems
+
+        self.instantiate_mask_estim(
+            in_channels=in_channels,
+            stems=stems,
+            emb_dim=emb_dim,
+            mlp_dim=mlp_dim,
+            hidden_activation=hidden_activation,
+            hidden_activation_kwargs=hidden_activation_kwargs,
+            complex_mask=complex_mask,
+            n_freq=n_fft // 2 + 1,
+            use_freq_weights=use_freq_weights,
+        )
+
+    def instantiate_mask_estim(
+        self,
+        in_channels: int,
+        stems: List[str],
+        emb_dim: int,
+        mlp_dim: int,
+        hidden_activation: str,
+        hidden_activation_kwargs: Optional[Dict] = None,
+        complex_mask: bool = True,
+        n_freq: Optional[int] = None,
+        use_freq_weights: bool = False,
+    ):
+        if hidden_activation_kwargs is None:
+            hidden_activation_kwargs = {}
+
+        assert n_freq is not None
+
+        self.mask_estim = nn.ModuleDict(
+            {
+                stem: OverlappingMaskEstimationModule(
+                    band_specs=self.band_specs.get_band_specs(),
+                    freq_weights=self.band_specs.get_freq_weights(),
+                    n_freq=n_freq,
+                    emb_dim=emb_dim,
+                    mlp_dim=mlp_dim,
+                    in_channels=in_channels,
+                    hidden_activation=hidden_activation,
+                    hidden_activation_kwargs=hidden_activation_kwargs,
+                    complex_mask=complex_mask,
+                    use_freq_weights=use_freq_weights,
+                )
+                for stem in stems
+            }
+        )
+
+    def separate(self, batch):
+        batch["estimates"] = {}
+
+        x, q, length = self.encode(batch)
+
+        for stem, mem in self.mask_estim.items():
+            m = mem(q)
+
+            s = self.mask(x, m.to(x.dtype))
+            s = torch.reshape(s, x.shape)
+            batch["estimates"][stem] = {
+                "audio": self.istft(s, length),
+                "spectrogram": s,
+            }
+
+        return batch

mvsepless/models/bandit_v2/bandsplit.py

@@ -1,127 +1,127 @@
-from typing import List, Tuple
-
-import torch
-from torch import nn
-from torch.utils.checkpoint import checkpoint_sequential
-
-from .utils import (
-    band_widths_from_specs,
-    check_no_gap,
-    check_no_overlap,
-    check_nonzero_bandwidth,
-)
-
-
-class NormFC(nn.Module):
-    def __init__(
-        self,
-        emb_dim: int,
-        bandwidth: int,
-        in_channels: int,
-        normalize_channel_independently: bool = False,
-        treat_channel_as_feature: bool = True,
-    ) -> None:
-        super().__init__()
-
-        if not treat_channel_as_feature:
-            raise NotImplementedError
-
-        self.treat_channel_as_feature = treat_channel_as_feature
-
-        if normalize_channel_independently:
-            raise NotImplementedError
-
-        reim = 2
-
-        norm = nn.LayerNorm(in_channels * bandwidth * reim)
-
-        fc_in = bandwidth * reim
-
-        if treat_channel_as_feature:
-            fc_in *= in_channels
-        else:
-            assert emb_dim % in_channels == 0
-            emb_dim = emb_dim // in_channels
-
-        fc = nn.Linear(fc_in, emb_dim)
-
-        self.combined = nn.Sequential(norm, fc)
-
-    def forward(self, xb):
-        return checkpoint_sequential(self.combined, 1, xb, use_reentrant=False)
-
-
-class BandSplitModule(nn.Module):
-    def __init__(
-        self,
-        band_specs: List[Tuple[float, float]],
-        emb_dim: int,
-        in_channels: int,
-        require_no_overlap: bool = False,
-        require_no_gap: bool = True,
-        normalize_channel_independently: bool = False,
-        treat_channel_as_feature: bool = True,
-    ) -> None:
-        super().__init__()
-
-        check_nonzero_bandwidth(band_specs)
-
-        if require_no_gap:
-            check_no_gap(band_specs)
-
-        if require_no_overlap:
-            check_no_overlap(band_specs)
-
-        self.band_specs = band_specs
-        self.band_widths = band_widths_from_specs(band_specs)
-        self.n_bands = len(band_specs)
-        self.emb_dim = emb_dim
-
-        try:
-            self.norm_fc_modules = nn.ModuleList(
-                [  # type: ignore
-                    torch.compile(
-                        NormFC(
-                            emb_dim=emb_dim,
-                            bandwidth=bw,
-                            in_channels=in_channels,
-                            normalize_channel_independently=normalize_channel_independently,
-                            treat_channel_as_feature=treat_channel_as_feature,
-                        ),
-                        disable=True,
-                    )
-                    for bw in self.band_widths
-                ]
-            )
-        except Exception as e:
-            self.norm_fc_modules = nn.ModuleList(
-                [  # type: ignore
-                    NormFC(
-                        emb_dim=emb_dim,
-                        bandwidth=bw,
-                        in_channels=in_channels,
-                        normalize_channel_independently=normalize_channel_independently,
-                        treat_channel_as_feature=treat_channel_as_feature,
-                    )
-                    for bw in self.band_widths
-                ]
-            )
-
-    def forward(self, x: torch.Tensor):
-
-        batch, in_chan, band_width, n_time = x.shape
-
-        z = torch.zeros(
-            size=(batch, self.n_bands, n_time, self.emb_dim), device=x.device
-        )
-
-        x = torch.permute(x, (0, 3, 1, 2)).contiguous()
-
-        for i, nfm in enumerate(self.norm_fc_modules):
-            fstart, fend = self.band_specs[i]
-            xb = x[:, :, :, fstart:fend]
-            xb = torch.view_as_real(xb)
-            xb = torch.reshape(xb, (batch, n_time, -1))
-            z[:, i, :, :] = nfm(xb)
-
-        return z
+from typing import List, Tuple
+
+import torch
+from torch import nn
+from torch.utils.checkpoint import checkpoint_sequential
+
+from .utils import (
+    band_widths_from_specs,
+    check_no_gap,
+    check_no_overlap,
+    check_nonzero_bandwidth,
+)
+
+
+class NormFC(nn.Module):
+    def __init__(
+        self,
+        emb_dim: int,
+        bandwidth: int,
+        in_channels: int,
+        normalize_channel_independently: bool = False,
+        treat_channel_as_feature: bool = True,
+    ) -> None:
+        super().__init__()
+
+        if not treat_channel_as_feature:
+            raise NotImplementedError
+
+        self.treat_channel_as_feature = treat_channel_as_feature
+
+        if normalize_channel_independently:
+            raise NotImplementedError
+
+        reim = 2
+
+        norm = nn.LayerNorm(in_channels * bandwidth * reim)
+
+        fc_in = bandwidth * reim
+
+        if treat_channel_as_feature:
+            fc_in *= in_channels
+        else:
+            assert emb_dim % in_channels == 0
+            emb_dim = emb_dim // in_channels
+
+        fc = nn.Linear(fc_in, emb_dim)
+
+        self.combined = nn.Sequential(norm, fc)
+
+    def forward(self, xb):
+        return checkpoint_sequential(self.combined, 1, xb, use_reentrant=False)
+
+
+class BandSplitModule(nn.Module):
+    def __init__(
+        self,
+        band_specs: List[Tuple[float, float]],
+        emb_dim: int,
+        in_channels: int,
+        require_no_overlap: bool = False,
+        require_no_gap: bool = True,
+        normalize_channel_independently: bool = False,
+        treat_channel_as_feature: bool = True,
+    ) -> None:
+        super().__init__()
+
+        check_nonzero_bandwidth(band_specs)
+
+        if require_no_gap:
+            check_no_gap(band_specs)
+
+        if require_no_overlap:
+            check_no_overlap(band_specs)
+
+        self.band_specs = band_specs
+        self.band_widths = band_widths_from_specs(band_specs)
+        self.n_bands = len(band_specs)
+        self.emb_dim = emb_dim
+
+        try:
+            self.norm_fc_modules = nn.ModuleList(
+                [  # type: ignore
+                    torch.compile(
+                        NormFC(
+                            emb_dim=emb_dim,
+                            bandwidth=bw,
+                            in_channels=in_channels,
+                            normalize_channel_independently=normalize_channel_independently,
+                            treat_channel_as_feature=treat_channel_as_feature,
+                        ),
+                        disable=True,
+                    )
+                    for bw in self.band_widths
+                ]
+            )
+        except Exception as e:
+            self.norm_fc_modules = nn.ModuleList(
+                [  # type: ignore
+                    NormFC(
+                        emb_dim=emb_dim,
+                        bandwidth=bw,
+                        in_channels=in_channels,
+                        normalize_channel_independently=normalize_channel_independently,
+                        treat_channel_as_feature=treat_channel_as_feature,
+                    )
+                    for bw in self.band_widths
+                ]
+            )
+
+    def forward(self, x: torch.Tensor):
+
+        batch, in_chan, band_width, n_time = x.shape
+
+        z = torch.zeros(
+            size=(batch, self.n_bands, n_time, self.emb_dim), device=x.device
+        )
+
+        x = torch.permute(x, (0, 3, 1, 2)).contiguous()
+
+        for i, nfm in enumerate(self.norm_fc_modules):
+            fstart, fend = self.band_specs[i]
+            xb = x[:, :, :, fstart:fend]
+            xb = torch.view_as_real(xb)
+            xb = torch.reshape(xb, (batch, n_time, -1))
+            z[:, i, :, :] = nfm(xb)
+
+        return z

mvsepless/models/bandit_v2/film.py

@@ -1,23 +1,23 @@
-from torch import nn
-import torch
-
-
-class FiLM(nn.Module):
-    def __init__(self):
-        super().__init__()
-
-    def forward(self, x, gamma, beta):
-        return gamma * x + beta
-
-
-class BTFBroadcastedFiLM(nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.film = FiLM()
-
-    def forward(self, x, gamma, beta):
-
-        gamma = gamma[None, None, None, :]
-        beta = beta[None, None, None, :]
-
-        return self.film(x, gamma, beta)
+from torch import nn
+import torch
+
+
+class FiLM(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x, gamma, beta):
+        return gamma * x + beta
+
+
+class BTFBroadcastedFiLM(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.film = FiLM()
+
+    def forward(self, x, gamma, beta):
+
+        gamma = gamma[None, None, None, :]
+        beta = beta[None, None, None, :]
+
+        return self.film(x, gamma, beta)

mvsepless/models/bandit_v2/maskestim.py

@@ -1,269 +1,269 @@
-from typing import Dict, List, Optional, Tuple, Type
-
-import torch
-from torch import nn
-from torch.nn.modules import activation
-from torch.utils.checkpoint import checkpoint_sequential
-
-from .utils import (
-    band_widths_from_specs,
-    check_no_gap,
-    check_no_overlap,
-    check_nonzero_bandwidth,
-)
-
-
-class BaseNormMLP(nn.Module):
-    def __init__(
-        self,
-        emb_dim: int,
-        mlp_dim: int,
-        bandwidth: int,
-        in_channels: Optional[int],
-        hidden_activation: str = "Tanh",
-        hidden_activation_kwargs=None,
-        complex_mask: bool = True,
-    ):
-        super().__init__()
-        if hidden_activation_kwargs is None:
-            hidden_activation_kwargs = {}
-        self.hidden_activation_kwargs = hidden_activation_kwargs
-        self.norm = nn.LayerNorm(emb_dim)
-        self.hidden = nn.Sequential(
-            nn.Linear(in_features=emb_dim, out_features=mlp_dim),
-            activation.__dict__[hidden_activation](**self.hidden_activation_kwargs),
-        )
-
-        self.bandwidth = bandwidth
-        self.in_channels = in_channels
-
-        self.complex_mask = complex_mask
-        self.reim = 2 if complex_mask else 1
-        self.glu_mult = 2
-
-
-class NormMLP(BaseNormMLP):
-    def __init__(
-        self,
-        emb_dim: int,
-        mlp_dim: int,
-        bandwidth: int,
-        in_channels: Optional[int],
-        hidden_activation: str = "Tanh",
-        hidden_activation_kwargs=None,
-        complex_mask: bool = True,
-    ) -> None:
-        super().__init__(
-            emb_dim=emb_dim,
-            mlp_dim=mlp_dim,
-            bandwidth=bandwidth,
-            in_channels=in_channels,
-            hidden_activation=hidden_activation,
-            hidden_activation_kwargs=hidden_activation_kwargs,
-            complex_mask=complex_mask,
-        )
-
-        self.output = nn.Sequential(
-            nn.Linear(
-                in_features=mlp_dim,
-                out_features=bandwidth * in_channels * self.reim * 2,
-            ),
-            nn.GLU(dim=-1),
-        )
-
-        try:
-            self.combined = torch.compile(
-                nn.Sequential(self.norm, self.hidden, self.output), disable=True
-            )
-        except Exception as e:
-            self.combined = nn.Sequential(self.norm, self.hidden, self.output)
-
-    def reshape_output(self, mb):
-        batch, n_time, _ = mb.shape
-        if self.complex_mask:
-            mb = mb.reshape(
-                batch, n_time, self.in_channels, self.bandwidth, self.reim
-            ).contiguous()
-            mb = torch.view_as_complex(mb)
-        else:
-            mb = mb.reshape(batch, n_time, self.in_channels, self.bandwidth)
-
-        mb = torch.permute(mb, (0, 2, 3, 1))
-
-        return mb
-
-    def forward(self, qb):
-
-        mb = checkpoint_sequential(self.combined, 2, qb, use_reentrant=False)
-        mb = self.reshape_output(mb)
-
-        return mb
-
-
-class MaskEstimationModuleSuperBase(nn.Module):
-    pass
-
-
-class MaskEstimationModuleBase(MaskEstimationModuleSuperBase):
-    def __init__(
-        self,
-        band_specs: List[Tuple[float, float]],
-        emb_dim: int,
-        mlp_dim: int,
-        in_channels: Optional[int],
-        hidden_activation: str = "Tanh",
-        hidden_activation_kwargs: Dict = None,
-        complex_mask: bool = True,
-        norm_mlp_cls: Type[nn.Module] = NormMLP,
-        norm_mlp_kwargs: Dict = None,
-    ) -> None:
-        super().__init__()
-
-        self.band_widths = band_widths_from_specs(band_specs)
-        self.n_bands = len(band_specs)
-
-        if hidden_activation_kwargs is None:
-            hidden_activation_kwargs = {}
-
-        if norm_mlp_kwargs is None:
-            norm_mlp_kwargs = {}
-
-        self.norm_mlp = nn.ModuleList(
-            [
-                norm_mlp_cls(
-                    bandwidth=self.band_widths[b],
-                    emb_dim=emb_dim,
-                    mlp_dim=mlp_dim,
-                    in_channels=in_channels,
-                    hidden_activation=hidden_activation,
-                    hidden_activation_kwargs=hidden_activation_kwargs,
-                    complex_mask=complex_mask,
-                    **norm_mlp_kwargs,
-                )
-                for b in range(self.n_bands)
-            ]
-        )
-
-    def compute_masks(self, q):
-        batch, n_bands, n_time, emb_dim = q.shape
-
-        masks = []
-
-        for b, nmlp in enumerate(self.norm_mlp):
-            qb = q[:, b, :, :]
-            mb = nmlp(qb)
-            masks.append(mb)
-
-        return masks
-
-    def compute_mask(self, q, b):
-        batch, n_bands, n_time, emb_dim = q.shape
-        qb = q[:, b, :, :]
-        mb = self.norm_mlp[b](qb)
-        return mb
-
-
-class OverlappingMaskEstimationModule(MaskEstimationModuleBase):
-    def __init__(
-        self,
-        in_channels: int,
-        band_specs: List[Tuple[float, float]],
-        freq_weights: List[torch.Tensor],
-        n_freq: int,
-        emb_dim: int,
-        mlp_dim: int,
-        cond_dim: int = 0,
-        hidden_activation: str = "Tanh",
-        hidden_activation_kwargs: Dict = None,
-        complex_mask: bool = True,
-        norm_mlp_cls: Type[nn.Module] = NormMLP,
-        norm_mlp_kwargs: Dict = None,
-        use_freq_weights: bool = False,
-    ) -> None:
-        check_nonzero_bandwidth(band_specs)
-        check_no_gap(band_specs)
-
-        if cond_dim > 0:
-            raise NotImplementedError
-
-        super().__init__(
-            band_specs=band_specs,
-            emb_dim=emb_dim + cond_dim,
-            mlp_dim=mlp_dim,
-            in_channels=in_channels,
-            hidden_activation=hidden_activation,
-            hidden_activation_kwargs=hidden_activation_kwargs,
-            complex_mask=complex_mask,
-            norm_mlp_cls=norm_mlp_cls,
-            norm_mlp_kwargs=norm_mlp_kwargs,
-        )
-
-        self.n_freq = n_freq
-        self.band_specs = band_specs
-        self.in_channels = in_channels
-
-        if freq_weights is not None and use_freq_weights:
-            for i, fw in enumerate(freq_weights):
-                self.register_buffer(f"freq_weights/{i}", fw)
-
-                self.use_freq_weights = use_freq_weights
-        else:
-            self.use_freq_weights = False
-
-    def forward(self, q):
-
-        batch, n_bands, n_time, emb_dim = q.shape
-
-        masks = torch.zeros(
-            (batch, self.in_channels, self.n_freq, n_time),
-            device=q.device,
-            dtype=torch.complex64,
-        )
-
-        for im in range(n_bands):
-            fstart, fend = self.band_specs[im]
-
-            mask = self.compute_mask(q, im)
-
-            if self.use_freq_weights:
-                fw = self.get_buffer(f"freq_weights/{im}")[:, None]
-                mask = mask * fw
-            masks[:, :, fstart:fend, :] += mask
-
-        return masks
-
-
-class MaskEstimationModule(OverlappingMaskEstimationModule):
-    def __init__(
-        self,
-        band_specs: List[Tuple[float, float]],
-        emb_dim: int,
-        mlp_dim: int,
-        in_channels: Optional[int],
-        hidden_activation: str = "Tanh",
-        hidden_activation_kwargs: Dict = None,
-        complex_mask: bool = True,
-        **kwargs,
-    ) -> None:
-        check_nonzero_bandwidth(band_specs)
-        check_no_gap(band_specs)
-        check_no_overlap(band_specs)
-        super().__init__(
-            in_channels=in_channels,
-            band_specs=band_specs,
-            freq_weights=None,
-            n_freq=None,
-            emb_dim=emb_dim,
-            mlp_dim=mlp_dim,
-            hidden_activation=hidden_activation,
-            hidden_activation_kwargs=hidden_activation_kwargs,
-            complex_mask=complex_mask,
-        )
-
-    def forward(self, q, cond=None):
-
-        masks = self.compute_masks(q)
-
-        masks = torch.concat(masks, dim=2)
-
-        return masks
+from typing import Dict, List, Optional, Tuple, Type
+
+import torch
+from torch import nn
+from torch.nn.modules import activation
+from torch.utils.checkpoint import checkpoint_sequential
+
+from .utils import (
+    band_widths_from_specs,
+    check_no_gap,
+    check_no_overlap,
+    check_nonzero_bandwidth,
+)
+
+
+class BaseNormMLP(nn.Module):
+    def __init__(
+        self,
+        emb_dim: int,
+        mlp_dim: int,
+        bandwidth: int,
+        in_channels: Optional[int],
+        hidden_activation: str = "Tanh",
+        hidden_activation_kwargs=None,
+        complex_mask: bool = True,
+    ):
+        super().__init__()
+        if hidden_activation_kwargs is None:
+            hidden_activation_kwargs = {}
+        self.hidden_activation_kwargs = hidden_activation_kwargs
+        self.norm = nn.LayerNorm(emb_dim)
+        self.hidden = nn.Sequential(
+            nn.Linear(in_features=emb_dim, out_features=mlp_dim),
+            activation.__dict__[hidden_activation](**self.hidden_activation_kwargs),
+        )
+
+        self.bandwidth = bandwidth
+        self.in_channels = in_channels
+
+        self.complex_mask = complex_mask
+        self.reim = 2 if complex_mask else 1
+        self.glu_mult = 2
+
+
+class NormMLP(BaseNormMLP):
+    def __init__(
+        self,
+        emb_dim: int,
+        mlp_dim: int,
+        bandwidth: int,
+        in_channels: Optional[int],
+        hidden_activation: str = "Tanh",
+        hidden_activation_kwargs=None,
+        complex_mask: bool = True,
+    ) -> None:
+        super().__init__(
+            emb_dim=emb_dim,
+            mlp_dim=mlp_dim,
+            bandwidth=bandwidth,
+            in_channels=in_channels,
+            hidden_activation=hidden_activation,
+            hidden_activation_kwargs=hidden_activation_kwargs,
+            complex_mask=complex_mask,
+        )
+
+        self.output = nn.Sequential(
+            nn.Linear(
+                in_features=mlp_dim,
+                out_features=bandwidth * in_channels * self.reim * 2,
+            ),
+            nn.GLU(dim=-1),
+        )
+
+        try:
+            self.combined = torch.compile(
+                nn.Sequential(self.norm, self.hidden, self.output), disable=True
+            )
+        except Exception as e:
+            self.combined = nn.Sequential(self.norm, self.hidden, self.output)
+
+    def reshape_output(self, mb):
+        batch, n_time, _ = mb.shape
+        if self.complex_mask:
+            mb = mb.reshape(
+                batch, n_time, self.in_channels, self.bandwidth, self.reim
+            ).contiguous()
+            mb = torch.view_as_complex(mb)
+        else:
+            mb = mb.reshape(batch, n_time, self.in_channels, self.bandwidth)
+
+        mb = torch.permute(mb, (0, 2, 3, 1))
+
+        return mb
+
+    def forward(self, qb):
+
+        mb = checkpoint_sequential(self.combined, 2, qb, use_reentrant=False)
+        mb = self.reshape_output(mb)
+
+        return mb
+
+
+class MaskEstimationModuleSuperBase(nn.Module):
+    pass
+
+
+class MaskEstimationModuleBase(MaskEstimationModuleSuperBase):
+    def __init__(
+        self,
+        band_specs: List[Tuple[float, float]],
+        emb_dim: int,
+        mlp_dim: int,
+        in_channels: Optional[int],
+        hidden_activation: str = "Tanh",
+        hidden_activation_kwargs: Dict = None,
+        complex_mask: bool = True,
+        norm_mlp_cls: Type[nn.Module] = NormMLP,
+        norm_mlp_kwargs: Dict = None,
+    ) -> None:
+        super().__init__()
+
+        self.band_widths = band_widths_from_specs(band_specs)
+        self.n_bands = len(band_specs)
+
+        if hidden_activation_kwargs is None:
+            hidden_activation_kwargs = {}
+
+        if norm_mlp_kwargs is None:
+            norm_mlp_kwargs = {}
+
+        self.norm_mlp = nn.ModuleList(
+            [
+                norm_mlp_cls(
+                    bandwidth=self.band_widths[b],
+                    emb_dim=emb_dim,
+                    mlp_dim=mlp_dim,
+                    in_channels=in_channels,
+                    hidden_activation=hidden_activation,
+                    hidden_activation_kwargs=hidden_activation_kwargs,
+                    complex_mask=complex_mask,
+                    **norm_mlp_kwargs,
+                )
+                for b in range(self.n_bands)
+            ]
+        )
+
+    def compute_masks(self, q):
+        batch, n_bands, n_time, emb_dim = q.shape
+
+        masks = []
+
+        for b, nmlp in enumerate(self.norm_mlp):
+            qb = q[:, b, :, :]
+            mb = nmlp(qb)
+            masks.append(mb)
+
+        return masks
+
+    def compute_mask(self, q, b):
+        batch, n_bands, n_time, emb_dim = q.shape
+        qb = q[:, b, :, :]
+        mb = self.norm_mlp[b](qb)
+        return mb
+
+
+class OverlappingMaskEstimationModule(MaskEstimationModuleBase):
+    def __init__(
+        self,
+        in_channels: int,
+        band_specs: List[Tuple[float, float]],
+        freq_weights: List[torch.Tensor],
+        n_freq: int,
+        emb_dim: int,
+        mlp_dim: int,
+        cond_dim: int = 0,
+        hidden_activation: str = "Tanh",
+        hidden_activation_kwargs: Dict = None,
+        complex_mask: bool = True,
+        norm_mlp_cls: Type[nn.Module] = NormMLP,
+        norm_mlp_kwargs: Dict = None,
+        use_freq_weights: bool = False,
+    ) -> None:
+        check_nonzero_bandwidth(band_specs)
+        check_no_gap(band_specs)
+
+        if cond_dim > 0:
+            raise NotImplementedError
+
+        super().__init__(
+            band_specs=band_specs,
+            emb_dim=emb_dim + cond_dim,
+            mlp_dim=mlp_dim,
+            in_channels=in_channels,
+            hidden_activation=hidden_activation,
+            hidden_activation_kwargs=hidden_activation_kwargs,
+            complex_mask=complex_mask,
+            norm_mlp_cls=norm_mlp_cls,
+            norm_mlp_kwargs=norm_mlp_kwargs,
+        )
+
+        self.n_freq = n_freq
+        self.band_specs = band_specs
+        self.in_channels = in_channels
+
+        if freq_weights is not None and use_freq_weights:
+            for i, fw in enumerate(freq_weights):
+                self.register_buffer(f"freq_weights/{i}", fw)
+
+                self.use_freq_weights = use_freq_weights
+        else:
+            self.use_freq_weights = False
+
+    def forward(self, q):
+
+        batch, n_bands, n_time, emb_dim = q.shape
+
+        masks = torch.zeros(
+            (batch, self.in_channels, self.n_freq, n_time),
+            device=q.device,
+            dtype=torch.complex64,
+        )
+
+        for im in range(n_bands):
+            fstart, fend = self.band_specs[im]
+
+            mask = self.compute_mask(q, im)
+
+            if self.use_freq_weights:
+                fw = self.get_buffer(f"freq_weights/{im}")[:, None]
+                mask = mask * fw
+            masks[:, :, fstart:fend, :] += mask
+
+        return masks
+
+
+class MaskEstimationModule(OverlappingMaskEstimationModule):
+    def __init__(
+        self,
+        band_specs: List[Tuple[float, float]],
+        emb_dim: int,
+        mlp_dim: int,
+        in_channels: Optional[int],
+        hidden_activation: str = "Tanh",
+        hidden_activation_kwargs: Dict = None,
+        complex_mask: bool = True,
+        **kwargs,
+    ) -> None:
+        check_nonzero_bandwidth(band_specs)
+        check_no_gap(band_specs)
+        check_no_overlap(band_specs)
+        super().__init__(
+            in_channels=in_channels,
+            band_specs=band_specs,
+            freq_weights=None,
+            n_freq=None,
+            emb_dim=emb_dim,
+            mlp_dim=mlp_dim,
+            hidden_activation=hidden_activation,
+            hidden_activation_kwargs=hidden_activation_kwargs,
+            complex_mask=complex_mask,
+        )
+
+    def forward(self, q, cond=None):
+
+        masks = self.compute_masks(q)
+
+        masks = torch.concat(masks, dim=2)
+
+        return masks

mvsepless/models/bandit_v2/tfmodel.py

@@ -1,141 +1,141 @@
-import warnings
-
-import torch
-import torch.backends.cuda
-from torch import nn
-from torch.nn.modules import rnn
-from torch.utils.checkpoint import checkpoint_sequential
-
-
-class TimeFrequencyModellingModule(nn.Module):
-    def __init__(self) -> None:
-        super().__init__()
-
-
-class ResidualRNN(nn.Module):
-    def __init__(
-        self,
-        emb_dim: int,
-        rnn_dim: int,
-        bidirectional: bool = True,
-        rnn_type: str = "LSTM",
-        use_batch_trick: bool = True,
-        use_layer_norm: bool = True,
-    ) -> None:
-        super().__init__()
-
-        assert use_layer_norm
-        assert use_batch_trick
-
-        self.use_layer_norm = use_layer_norm
-        self.norm = nn.LayerNorm(emb_dim)
-        self.rnn = rnn.__dict__[rnn_type](
-            input_size=emb_dim,
-            hidden_size=rnn_dim,
-            num_layers=1,
-            batch_first=True,
-            bidirectional=bidirectional,
-        )
-
-        self.fc = nn.Linear(
-            in_features=rnn_dim * (2 if bidirectional else 1), out_features=emb_dim
-        )
-
-        self.use_batch_trick = use_batch_trick
-        if not self.use_batch_trick:
-            warnings.warn("NOT USING BATCH TRICK IS EXTREMELY SLOW!!")
-
-    def forward(self, z):
-
-        z0 = torch.clone(z)
-        z = self.norm(z)
-
-        batch, n_uncrossed, n_across, emb_dim = z.shape
-        z = torch.reshape(z, (batch * n_uncrossed, n_across, emb_dim))
-        z = self.rnn(z)[0]
-        z = torch.reshape(z, (batch, n_uncrossed, n_across, -1))
-
-        z = self.fc(z)
-
-        z = z + z0
-
-        return z
-
-
-class Transpose(nn.Module):
-    def __init__(self, dim0: int, dim1: int) -> None:
-        super().__init__()
-        self.dim0 = dim0
-        self.dim1 = dim1
-
-    def forward(self, z):
-        return z.transpose(self.dim0, self.dim1)
-
-
-class SeqBandModellingModule(TimeFrequencyModellingModule):
-    def __init__(
-        self,
-        n_modules: int = 12,
-        emb_dim: int = 128,
-        rnn_dim: int = 256,
-        bidirectional: bool = True,
-        rnn_type: str = "LSTM",
-        parallel_mode=False,
-    ) -> None:
-        super().__init__()
-
-        self.n_modules = n_modules
-
-        if parallel_mode:
-            self.seqband = nn.ModuleList([])
-            for _ in range(n_modules):
-                self.seqband.append(
-                    nn.ModuleList(
-                        [
-                            ResidualRNN(
-                                emb_dim=emb_dim,
-                                rnn_dim=rnn_dim,
-                                bidirectional=bidirectional,
-                                rnn_type=rnn_type,
-                            ),
-                            ResidualRNN(
-                                emb_dim=emb_dim,
-                                rnn_dim=rnn_dim,
-                                bidirectional=bidirectional,
-                                rnn_type=rnn_type,
-                            ),
-                        ]
-                    )
-                )
-        else:
-            seqband = []
-            for _ in range(2 * n_modules):
-                seqband += [
-                    ResidualRNN(
-                        emb_dim=emb_dim,
-                        rnn_dim=rnn_dim,
-                        bidirectional=bidirectional,
-                        rnn_type=rnn_type,
-                    ),
-                    Transpose(1, 2),
-                ]
-
-            self.seqband = nn.Sequential(*seqband)
-
-        self.parallel_mode = parallel_mode
-
-    def forward(self, z):
-
-        if self.parallel_mode:
-            for sbm_pair in self.seqband:
-                sbm_t, sbm_f = sbm_pair[0], sbm_pair[1]
-                zt = sbm_t(z)
-                zf = sbm_f(z.transpose(1, 2))
-                z = zt + zf.transpose(1, 2)
-        else:
-            z = checkpoint_sequential(
-                self.seqband, self.n_modules, z, use_reentrant=False
-            )
-
-        q = z
-        return q
+import warnings
+
+import torch
+import torch.backends.cuda
+from torch import nn
+from torch.nn.modules import rnn
+from torch.utils.checkpoint import checkpoint_sequential
+
+
+class TimeFrequencyModellingModule(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+
+
+class ResidualRNN(nn.Module):
+    def __init__(
+        self,
+        emb_dim: int,
+        rnn_dim: int,
+        bidirectional: bool = True,
+        rnn_type: str = "LSTM",
+        use_batch_trick: bool = True,
+        use_layer_norm: bool = True,
+    ) -> None:
+        super().__init__()
+
+        assert use_layer_norm
+        assert use_batch_trick
+
+        self.use_layer_norm = use_layer_norm
+        self.norm = nn.LayerNorm(emb_dim)
+        self.rnn = rnn.__dict__[rnn_type](
+            input_size=emb_dim,
+            hidden_size=rnn_dim,
+            num_layers=1,
+            batch_first=True,
+            bidirectional=bidirectional,
+        )
+
+        self.fc = nn.Linear(
+            in_features=rnn_dim * (2 if bidirectional else 1), out_features=emb_dim
+        )
+
+        self.use_batch_trick = use_batch_trick
+        if not self.use_batch_trick:
+            warnings.warn("NOT USING BATCH TRICK IS EXTREMELY SLOW!!")
+
+    def forward(self, z):
+
+        z0 = torch.clone(z)
+        z = self.norm(z)
+
+        batch, n_uncrossed, n_across, emb_dim = z.shape
+        z = torch.reshape(z, (batch * n_uncrossed, n_across, emb_dim))
+        z = self.rnn(z)[0]
+        z = torch.reshape(z, (batch, n_uncrossed, n_across, -1))
+
+        z = self.fc(z)
+
+        z = z + z0
+
+        return z
+
+
+class Transpose(nn.Module):
+    def __init__(self, dim0: int, dim1: int) -> None:
+        super().__init__()
+        self.dim0 = dim0
+        self.dim1 = dim1
+
+    def forward(self, z):
+        return z.transpose(self.dim0, self.dim1)
+
+
+class SeqBandModellingModule(TimeFrequencyModellingModule):
+    def __init__(
+        self,
+        n_modules: int = 12,
+        emb_dim: int = 128,
+        rnn_dim: int = 256,
+        bidirectional: bool = True,
+        rnn_type: str = "LSTM",
+        parallel_mode=False,
+    ) -> None:
+        super().__init__()
+
+        self.n_modules = n_modules
+
+        if parallel_mode:
+            self.seqband = nn.ModuleList([])
+            for _ in range(n_modules):
+                self.seqband.append(
+                    nn.ModuleList(
+                        [
+                            ResidualRNN(
+                                emb_dim=emb_dim,
+                                rnn_dim=rnn_dim,
+                                bidirectional=bidirectional,
+                                rnn_type=rnn_type,
+                            ),
+                            ResidualRNN(
+                                emb_dim=emb_dim,
+                                rnn_dim=rnn_dim,
+                                bidirectional=bidirectional,
+                                rnn_type=rnn_type,
+                            ),
+                        ]
+                    )
+                )
+        else:
+            seqband = []
+            for _ in range(2 * n_modules):
+                seqband += [
+                    ResidualRNN(
+                        emb_dim=emb_dim,
+                        rnn_dim=rnn_dim,
+                        bidirectional=bidirectional,
+                        rnn_type=rnn_type,
+                    ),
+                    Transpose(1, 2),
+                ]
+
+            self.seqband = nn.Sequential(*seqband)
+
+        self.parallel_mode = parallel_mode
+
+    def forward(self, z):
+
+        if self.parallel_mode:
+            for sbm_pair in self.seqband:
+                sbm_t, sbm_f = sbm_pair[0], sbm_pair[1]
+                zt = sbm_t(z)
+                zf = sbm_f(z.transpose(1, 2))
+                z = zt + zf.transpose(1, 2)
+        else:
+            z = checkpoint_sequential(
+                self.seqband, self.n_modules, z, use_reentrant=False
+            )
+
+        q = z
+        return q

mvsepless/models/bandit_v2/utils.py

@@ -1,384 +1,384 @@
-import os
-from abc import abstractmethod
-from typing import Callable
-
-import numpy as np
-import torch
-from librosa import hz_to_midi, midi_to_hz
-from torchaudio import functional as taF
-
-
-def band_widths_from_specs(band_specs):
-    return [e - i for i, e in band_specs]
-
-
-def check_nonzero_bandwidth(band_specs):
-    for fstart, fend in band_specs:
-        if fend - fstart <= 0:
-            raise ValueError("Bands cannot be zero-width")
-
-
-def check_no_overlap(band_specs):
-    fend_prev = -1
-    for fstart_curr, fend_curr in band_specs:
-        if fstart_curr <= fend_prev:
-            raise ValueError("Bands cannot overlap")
-
-
-def check_no_gap(band_specs):
-    fstart, _ = band_specs[0]
-    assert fstart == 0
-
-    fend_prev = -1
-    for fstart_curr, fend_curr in band_specs:
-        if fstart_curr - fend_prev > 1:
-            raise ValueError("Bands cannot leave gap")
-        fend_prev = fend_curr
-
-
-class BandsplitSpecification:
-    def __init__(self, nfft: int, fs: int) -> None:
-        self.fs = fs
-        self.nfft = nfft
-        self.nyquist = fs / 2
-        self.max_index = nfft // 2 + 1
-
-        self.split500 = self.hertz_to_index(500)
-        self.split1k = self.hertz_to_index(1000)
-        self.split2k = self.hertz_to_index(2000)
-        self.split4k = self.hertz_to_index(4000)
-        self.split8k = self.hertz_to_index(8000)
-        self.split16k = self.hertz_to_index(16000)
-        self.split20k = self.hertz_to_index(20000)
-
-        self.above20k = [(self.split20k, self.max_index)]
-        self.above16k = [(self.split16k, self.split20k)] + self.above20k
-
-    def index_to_hertz(self, index: int):
-        return index * self.fs / self.nfft
-
-    def hertz_to_index(self, hz: float, round: bool = True):
-        index = hz * self.nfft / self.fs
-
-        if round:
-            index = int(np.round(index))
-
-        return index
-
-    def get_band_specs_with_bandwidth(self, start_index, end_index, bandwidth_hz):
-        band_specs = []
-        lower = start_index
-
-        while lower < end_index:
-            upper = int(np.floor(lower + self.hertz_to_index(bandwidth_hz)))
-            upper = min(upper, end_index)
-
-            band_specs.append((lower, upper))
-            lower = upper
-
-        return band_specs
-
-    @abstractmethod
-    def get_band_specs(self):
-        raise NotImplementedError
-
-
-class VocalBandsplitSpecification(BandsplitSpecification):
-    def __init__(self, nfft: int, fs: int, version: str = "7") -> None:
-        super().__init__(nfft=nfft, fs=fs)
-
-        self.version = version
-
-    def get_band_specs(self):
-        return getattr(self, f"version{self.version}")()
-
-    @property
-    def version1(self):
-        return self.get_band_specs_with_bandwidth(
-            start_index=0, end_index=self.max_index, bandwidth_hz=1000
-        )
-
-    def version2(self):
-        below16k = self.get_band_specs_with_bandwidth(
-            start_index=0, end_index=self.split16k, bandwidth_hz=1000
-        )
-        below20k = self.get_band_specs_with_bandwidth(
-            start_index=self.split16k, end_index=self.split20k, bandwidth_hz=2000
-        )
-
-        return below16k + below20k + self.above20k
-
-    def version3(self):
-        below8k = self.get_band_specs_with_bandwidth(
-            start_index=0, end_index=self.split8k, bandwidth_hz=1000
-        )
-        below16k = self.get_band_specs_with_bandwidth(
-            start_index=self.split8k, end_index=self.split16k, bandwidth_hz=2000
-        )
-
-        return below8k + below16k + self.above16k
-
-    def version4(self):
-        below1k = self.get_band_specs_with_bandwidth(
-            start_index=0, end_index=self.split1k, bandwidth_hz=100
-        )
-        below8k = self.get_band_specs_with_bandwidth(
-            start_index=self.split1k, end_index=self.split8k, bandwidth_hz=1000
-        )
-        below16k = self.get_band_specs_with_bandwidth(
-            start_index=self.split8k, end_index=self.split16k, bandwidth_hz=2000
-        )
-
-        return below1k + below8k + below16k + self.above16k
-
-    def version5(self):
-        below1k = self.get_band_specs_with_bandwidth(
-            start_index=0, end_index=self.split1k, bandwidth_hz=100
-        )
-        below16k = self.get_band_specs_with_bandwidth(
-            start_index=self.split1k, end_index=self.split16k, bandwidth_hz=1000
-        )
-        below20k = self.get_band_specs_with_bandwidth(
-            start_index=self.split16k, end_index=self.split20k, bandwidth_hz=2000
-        )
-        return below1k + below16k + below20k + self.above20k
-
-    def version6(self):
-        below1k = self.get_band_specs_with_bandwidth(
-            start_index=0, end_index=self.split1k, bandwidth_hz=100
-        )
-        below4k = self.get_band_specs_with_bandwidth(
-            start_index=self.split1k, end_index=self.split4k, bandwidth_hz=500
-        )
-        below8k = self.get_band_specs_with_bandwidth(
-            start_index=self.split4k, end_index=self.split8k, bandwidth_hz=1000
-        )
-        below16k = self.get_band_specs_with_bandwidth(
-            start_index=self.split8k, end_index=self.split16k, bandwidth_hz=2000
-        )
-        return below1k + below4k + below8k + below16k + self.above16k
-
-    def version7(self):
-        below1k = self.get_band_specs_with_bandwidth(
-            start_index=0, end_index=self.split1k, bandwidth_hz=100
-        )
-        below4k = self.get_band_specs_with_bandwidth(
-            start_index=self.split1k, end_index=self.split4k, bandwidth_hz=250
-        )
-        below8k = self.get_band_specs_with_bandwidth(
-            start_index=self.split4k, end_index=self.split8k, bandwidth_hz=500
-        )
-        below16k = self.get_band_specs_with_bandwidth(
-            start_index=self.split8k, end_index=self.split16k, bandwidth_hz=1000
-        )
-        below20k = self.get_band_specs_with_bandwidth(
-            start_index=self.split16k, end_index=self.split20k, bandwidth_hz=2000
-        )
-        return below1k + below4k + below8k + below16k + below20k + self.above20k
-
-
-class OtherBandsplitSpecification(VocalBandsplitSpecification):
-    def __init__(self, nfft: int, fs: int) -> None:
-        super().__init__(nfft=nfft, fs=fs, version="7")
-
-
-class BassBandsplitSpecification(BandsplitSpecification):
-    def __init__(self, nfft: int, fs: int, version: str = "7") -> None:
-        super().__init__(nfft=nfft, fs=fs)
-
-    def get_band_specs(self):
-        below500 = self.get_band_specs_with_bandwidth(
-            start_index=0, end_index=self.split500, bandwidth_hz=50
-        )
-        below1k = self.get_band_specs_with_bandwidth(
-            start_index=self.split500, end_index=self.split1k, bandwidth_hz=100
-        )
-        below4k = self.get_band_specs_with_bandwidth(
-            start_index=self.split1k, end_index=self.split4k, bandwidth_hz=500
-        )
-        below8k = self.get_band_specs_with_bandwidth(
-            start_index=self.split4k, end_index=self.split8k, bandwidth_hz=1000
-        )
-        below16k = self.get_band_specs_with_bandwidth(
-            start_index=self.split8k, end_index=self.split16k, bandwidth_hz=2000
-        )
-        above16k = [(self.split16k, self.max_index)]
-
-        return below500 + below1k + below4k + below8k + below16k + above16k
-
-
-class DrumBandsplitSpecification(BandsplitSpecification):
-    def __init__(self, nfft: int, fs: int) -> None:
-        super().__init__(nfft=nfft, fs=fs)
-
-    def get_band_specs(self):
-        below1k = self.get_band_specs_with_bandwidth(
-            start_index=0, end_index=self.split1k, bandwidth_hz=50
-        )
-        below2k = self.get_band_specs_with_bandwidth(
-            start_index=self.split1k, end_index=self.split2k, bandwidth_hz=100
-        )
-        below4k = self.get_band_specs_with_bandwidth(
-            start_index=self.split2k, end_index=self.split4k, bandwidth_hz=250
-        )
-        below8k = self.get_band_specs_with_bandwidth(
-            start_index=self.split4k, end_index=self.split8k, bandwidth_hz=500
-        )
-        below16k = self.get_band_specs_with_bandwidth(
-            start_index=self.split8k, end_index=self.split16k, bandwidth_hz=1000
-        )
-        above16k = [(self.split16k, self.max_index)]
-
-        return below1k + below2k + below4k + below8k + below16k + above16k
-
-
-class PerceptualBandsplitSpecification(BandsplitSpecification):
-    def __init__(
-        self,
-        nfft: int,
-        fs: int,
-        fbank_fn: Callable[[int, int, float, float, int], torch.Tensor],
-        n_bands: int,
-        f_min: float = 0.0,
-        f_max: float = None,
-    ) -> None:
-        super().__init__(nfft=nfft, fs=fs)
-        self.n_bands = n_bands
-        if f_max is None:
-            f_max = fs / 2
-
-        self.filterbank = fbank_fn(n_bands, fs, f_min, f_max, self.max_index)
-
-        weight_per_bin = torch.sum(self.filterbank, dim=0, keepdim=True)
-        normalized_mel_fb = self.filterbank / weight_per_bin
-
-        freq_weights = []
-        band_specs = []
-        for i in range(self.n_bands):
-            active_bins = torch.nonzero(self.filterbank[i, :]).squeeze().tolist()
-            if isinstance(active_bins, int):
-                active_bins = (active_bins, active_bins)
-            if len(active_bins) == 0:
-                continue
-            start_index = active_bins[0]
-            end_index = active_bins[-1] + 1
-            band_specs.append((start_index, end_index))
-            freq_weights.append(normalized_mel_fb[i, start_index:end_index])
-
-        self.freq_weights = freq_weights
-        self.band_specs = band_specs
-
-    def get_band_specs(self):
-        return self.band_specs
-
-    def get_freq_weights(self):
-        return self.freq_weights
-
-    def save_to_file(self, dir_path: str) -> None:
-        os.makedirs(dir_path, exist_ok=True)
-
-        import pickle
-
-        with open(os.path.join(dir_path, "mel_bandsplit_spec.pkl"), "wb") as f:
-            pickle.dump(
-                {
-                    "band_specs": self.band_specs,
-                    "freq_weights": self.freq_weights,
-                    "filterbank": self.filterbank,
-                },
-                f,
-            )
-
-
-def mel_filterbank(n_bands, fs, f_min, f_max, n_freqs):
-    fb = taF.melscale_fbanks(
-        n_mels=n_bands,
-        sample_rate=fs,
-        f_min=f_min,
-        f_max=f_max,
-        n_freqs=n_freqs,
-    ).T
-
-    fb[0, 0] = 1.0
-
-    return fb
-
-
-class MelBandsplitSpecification(PerceptualBandsplitSpecification):
-    def __init__(
-        self, nfft: int, fs: int, n_bands: int, f_min: float = 0.0, f_max: float = None
-    ) -> None:
-        super().__init__(
-            fbank_fn=mel_filterbank,
-            nfft=nfft,
-            fs=fs,
-            n_bands=n_bands,
-            f_min=f_min,
-            f_max=f_max,
-        )
-
-
-def musical_filterbank(n_bands, fs, f_min, f_max, n_freqs, scale="constant"):
-    nfft = 2 * (n_freqs - 1)
-    df = fs / nfft
-    f_max = f_max or fs / 2
-    f_min = f_min or 0
-    f_min = fs / nfft
-
-    n_octaves = np.log2(f_max / f_min)
-    n_octaves_per_band = n_octaves / n_bands
-    bandwidth_mult = np.power(2.0, n_octaves_per_band)
-
-    low_midi = max(0, hz_to_midi(f_min))
-    high_midi = hz_to_midi(f_max)
-    midi_points = np.linspace(low_midi, high_midi, n_bands)
-    hz_pts = midi_to_hz(midi_points)
-
-    low_pts = hz_pts / bandwidth_mult
-    high_pts = hz_pts * bandwidth_mult
-
-    low_bins = np.floor(low_pts / df).astype(int)
-    high_bins = np.ceil(high_pts / df).astype(int)
-
-    fb = np.zeros((n_bands, n_freqs))
-
-    for i in range(n_bands):
-        fb[i, low_bins[i] : high_bins[i] + 1] = 1.0
-
-    fb[0, : low_bins[0]] = 1.0
-    fb[-1, high_bins[-1] + 1 :] = 1.0
-
-    return torch.as_tensor(fb)
-
-
-class MusicalBandsplitSpecification(PerceptualBandsplitSpecification):
-    def __init__(
-        self, nfft: int, fs: int, n_bands: int, f_min: float = 0.0, f_max: float = None
-    ) -> None:
-        super().__init__(
-            fbank_fn=musical_filterbank,
-            nfft=nfft,
-            fs=fs,
-            n_bands=n_bands,
-            f_min=f_min,
-            f_max=f_max,
-        )
-
-
-if __name__ == "__main__":
-    import pandas as pd
-
-    band_defs = []
-
-    for bands in [VocalBandsplitSpecification]:
-        band_name = bands.__name__.replace("BandsplitSpecification", "")
-
-        mbs = bands(nfft=2048, fs=44100).get_band_specs()
-
-        for i, (f_min, f_max) in enumerate(mbs):
-            band_defs.append(
-                {"band": band_name, "band_index": i, "f_min": f_min, "f_max": f_max}
-            )
-
-    df = pd.DataFrame(band_defs)
-    df.to_csv("vox7bands.csv", index=False)
+import os
+from abc import abstractmethod
+from typing import Callable
+
+import numpy as np
+import torch
+from librosa import hz_to_midi, midi_to_hz
+from torchaudio import functional as taF
+
+
+def band_widths_from_specs(band_specs):
+    return [e - i for i, e in band_specs]
+
+
+def check_nonzero_bandwidth(band_specs):
+    for fstart, fend in band_specs:
+        if fend - fstart <= 0:
+            raise ValueError("Bands cannot be zero-width")
+
+
+def check_no_overlap(band_specs):
+    fend_prev = -1
+    for fstart_curr, fend_curr in band_specs:
+        if fstart_curr <= fend_prev:
+            raise ValueError("Bands cannot overlap")
+
+
+def check_no_gap(band_specs):
+    fstart, _ = band_specs[0]
+    assert fstart == 0
+
+    fend_prev = -1
+    for fstart_curr, fend_curr in band_specs:
+        if fstart_curr - fend_prev > 1:
+            raise ValueError("Bands cannot leave gap")
+        fend_prev = fend_curr
+
+
+class BandsplitSpecification:
+    def __init__(self, nfft: int, fs: int) -> None:
+        self.fs = fs
+        self.nfft = nfft
+        self.nyquist = fs / 2
+        self.max_index = nfft // 2 + 1
+
+        self.split500 = self.hertz_to_index(500)
+        self.split1k = self.hertz_to_index(1000)
+        self.split2k = self.hertz_to_index(2000)
+        self.split4k = self.hertz_to_index(4000)
+        self.split8k = self.hertz_to_index(8000)
+        self.split16k = self.hertz_to_index(16000)
+        self.split20k = self.hertz_to_index(20000)
+
+        self.above20k = [(self.split20k, self.max_index)]
+        self.above16k = [(self.split16k, self.split20k)] + self.above20k
+
+    def index_to_hertz(self, index: int):
+        return index * self.fs / self.nfft
+
+    def hertz_to_index(self, hz: float, round: bool = True):
+        index = hz * self.nfft / self.fs
+
+        if round:
+            index = int(np.round(index))
+
+        return index
+
+    def get_band_specs_with_bandwidth(self, start_index, end_index, bandwidth_hz):
+        band_specs = []
+        lower = start_index
+
+        while lower < end_index:
+            upper = int(np.floor(lower + self.hertz_to_index(bandwidth_hz)))
+            upper = min(upper, end_index)
+
+            band_specs.append((lower, upper))
+            lower = upper
+
+        return band_specs
+
+    @abstractmethod
+    def get_band_specs(self):
+        raise NotImplementedError
+
+
+class VocalBandsplitSpecification(BandsplitSpecification):
+    def __init__(self, nfft: int, fs: int, version: str = "7") -> None:
+        super().__init__(nfft=nfft, fs=fs)
+
+        self.version = version
+
+    def get_band_specs(self):
+        return getattr(self, f"version{self.version}")()
+
+    @property
+    def version1(self):
+        return self.get_band_specs_with_bandwidth(
+            start_index=0, end_index=self.max_index, bandwidth_hz=1000
+        )
+
+    def version2(self):
+        below16k = self.get_band_specs_with_bandwidth(
+            start_index=0, end_index=self.split16k, bandwidth_hz=1000
+        )
+        below20k = self.get_band_specs_with_bandwidth(
+            start_index=self.split16k, end_index=self.split20k, bandwidth_hz=2000
+        )
+
+        return below16k + below20k + self.above20k
+
+    def version3(self):
+        below8k = self.get_band_specs_with_bandwidth(
+            start_index=0, end_index=self.split8k, bandwidth_hz=1000
+        )
+        below16k = self.get_band_specs_with_bandwidth(
+            start_index=self.split8k, end_index=self.split16k, bandwidth_hz=2000
+        )
+
+        return below8k + below16k + self.above16k
+
+    def version4(self):
+        below1k = self.get_band_specs_with_bandwidth(
+            start_index=0, end_index=self.split1k, bandwidth_hz=100
+        )
+        below8k = self.get_band_specs_with_bandwidth(
+            start_index=self.split1k, end_index=self.split8k, bandwidth_hz=1000
+        )
+        below16k = self.get_band_specs_with_bandwidth(
+            start_index=self.split8k, end_index=self.split16k, bandwidth_hz=2000
+        )
+
+        return below1k + below8k + below16k + self.above16k
+
+    def version5(self):
+        below1k = self.get_band_specs_with_bandwidth(
+            start_index=0, end_index=self.split1k, bandwidth_hz=100
+        )
+        below16k = self.get_band_specs_with_bandwidth(
+            start_index=self.split1k, end_index=self.split16k, bandwidth_hz=1000
+        )
+        below20k = self.get_band_specs_with_bandwidth(
+            start_index=self.split16k, end_index=self.split20k, bandwidth_hz=2000
+        )
+        return below1k + below16k + below20k + self.above20k
+
+    def version6(self):
+        below1k = self.get_band_specs_with_bandwidth(
+            start_index=0, end_index=self.split1k, bandwidth_hz=100
+        )
+        below4k = self.get_band_specs_with_bandwidth(
+            start_index=self.split1k, end_index=self.split4k, bandwidth_hz=500
+        )
+        below8k = self.get_band_specs_with_bandwidth(
+            start_index=self.split4k, end_index=self.split8k, bandwidth_hz=1000
+        )
+        below16k = self.get_band_specs_with_bandwidth(
+            start_index=self.split8k, end_index=self.split16k, bandwidth_hz=2000
+        )
+        return below1k + below4k + below8k + below16k + self.above16k
+
+    def version7(self):
+        below1k = self.get_band_specs_with_bandwidth(
+            start_index=0, end_index=self.split1k, bandwidth_hz=100
+        )
+        below4k = self.get_band_specs_with_bandwidth(
+            start_index=self.split1k, end_index=self.split4k, bandwidth_hz=250
+        )
+        below8k = self.get_band_specs_with_bandwidth(
+            start_index=self.split4k, end_index=self.split8k, bandwidth_hz=500
+        )
+        below16k = self.get_band_specs_with_bandwidth(
+            start_index=self.split8k, end_index=self.split16k, bandwidth_hz=1000
+        )
+        below20k = self.get_band_specs_with_bandwidth(
+            start_index=self.split16k, end_index=self.split20k, bandwidth_hz=2000
+        )
+        return below1k + below4k + below8k + below16k + below20k + self.above20k
+
+
+class OtherBandsplitSpecification(VocalBandsplitSpecification):
+    def __init__(self, nfft: int, fs: int) -> None:
+        super().__init__(nfft=nfft, fs=fs, version="7")
+
+
+class BassBandsplitSpecification(BandsplitSpecification):
+    def __init__(self, nfft: int, fs: int, version: str = "7") -> None:
+        super().__init__(nfft=nfft, fs=fs)
+
+    def get_band_specs(self):
+        below500 = self.get_band_specs_with_bandwidth(
+            start_index=0, end_index=self.split500, bandwidth_hz=50
+        )
+        below1k = self.get_band_specs_with_bandwidth(
+            start_index=self.split500, end_index=self.split1k, bandwidth_hz=100
+        )
+        below4k = self.get_band_specs_with_bandwidth(
+            start_index=self.split1k, end_index=self.split4k, bandwidth_hz=500
+        )
+        below8k = self.get_band_specs_with_bandwidth(
+            start_index=self.split4k, end_index=self.split8k, bandwidth_hz=1000
+        )
+        below16k = self.get_band_specs_with_bandwidth(
+            start_index=self.split8k, end_index=self.split16k, bandwidth_hz=2000
+        )
+        above16k = [(self.split16k, self.max_index)]
+
+        return below500 + below1k + below4k + below8k + below16k + above16k
+
+
+class DrumBandsplitSpecification(BandsplitSpecification):
+    def __init__(self, nfft: int, fs: int) -> None:
+        super().__init__(nfft=nfft, fs=fs)
+
+    def get_band_specs(self):
+        below1k = self.get_band_specs_with_bandwidth(
+            start_index=0, end_index=self.split1k, bandwidth_hz=50
+        )
+        below2k = self.get_band_specs_with_bandwidth(
+            start_index=self.split1k, end_index=self.split2k, bandwidth_hz=100
+        )
+        below4k = self.get_band_specs_with_bandwidth(
+            start_index=self.split2k, end_index=self.split4k, bandwidth_hz=250
+        )
+        below8k = self.get_band_specs_with_bandwidth(
+            start_index=self.split4k, end_index=self.split8k, bandwidth_hz=500
+        )
+        below16k = self.get_band_specs_with_bandwidth(
+            start_index=self.split8k, end_index=self.split16k, bandwidth_hz=1000
+        )
+        above16k = [(self.split16k, self.max_index)]
+
+        return below1k + below2k + below4k + below8k + below16k + above16k
+
+
+class PerceptualBandsplitSpecification(BandsplitSpecification):
+    def __init__(
+        self,
+        nfft: int,
+        fs: int,
+        fbank_fn: Callable[[int, int, float, float, int], torch.Tensor],
+        n_bands: int,
+        f_min: float = 0.0,
+        f_max: float = None,
+    ) -> None:
+        super().__init__(nfft=nfft, fs=fs)
+        self.n_bands = n_bands
+        if f_max is None:
+            f_max = fs / 2
+
+        self.filterbank = fbank_fn(n_bands, fs, f_min, f_max, self.max_index)
+
+        weight_per_bin = torch.sum(self.filterbank, dim=0, keepdim=True)
+        normalized_mel_fb = self.filterbank / weight_per_bin
+
+        freq_weights = []
+        band_specs = []
+        for i in range(self.n_bands):
+            active_bins = torch.nonzero(self.filterbank[i, :]).squeeze().tolist()
+            if isinstance(active_bins, int):
+                active_bins = (active_bins, active_bins)
+            if len(active_bins) == 0:
+                continue
+            start_index = active_bins[0]
+            end_index = active_bins[-1] + 1
+            band_specs.append((start_index, end_index))
+            freq_weights.append(normalized_mel_fb[i, start_index:end_index])
+
+        self.freq_weights = freq_weights
+        self.band_specs = band_specs
+
+    def get_band_specs(self):
+        return self.band_specs
+
+    def get_freq_weights(self):
+        return self.freq_weights
+
+    def save_to_file(self, dir_path: str) -> None:
+        os.makedirs(dir_path, exist_ok=True)
+
+        import pickle
+
+        with open(os.path.join(dir_path, "mel_bandsplit_spec.pkl"), "wb") as f:
+            pickle.dump(
+                {
+                    "band_specs": self.band_specs,
+                    "freq_weights": self.freq_weights,
+                    "filterbank": self.filterbank,
+                },
+                f,
+            )
+
+
+def mel_filterbank(n_bands, fs, f_min, f_max, n_freqs):
+    fb = taF.melscale_fbanks(
+        n_mels=n_bands,
+        sample_rate=fs,
+        f_min=f_min,
+        f_max=f_max,
+        n_freqs=n_freqs,
+    ).T
+
+    fb[0, 0] = 1.0
+
+    return fb
+
+
+class MelBandsplitSpecification(PerceptualBandsplitSpecification):
+    def __init__(
+        self, nfft: int, fs: int, n_bands: int, f_min: float = 0.0, f_max: float = None
+    ) -> None:
+        super().__init__(
+            fbank_fn=mel_filterbank,
+            nfft=nfft,
+            fs=fs,
+            n_bands=n_bands,
+            f_min=f_min,
+            f_max=f_max,
+        )
+
+
+def musical_filterbank(n_bands, fs, f_min, f_max, n_freqs, scale="constant"):
+    nfft = 2 * (n_freqs - 1)
+    df = fs / nfft
+    f_max = f_max or fs / 2
+    f_min = f_min or 0
+    f_min = fs / nfft
+
+    n_octaves = np.log2(f_max / f_min)
+    n_octaves_per_band = n_octaves / n_bands
+    bandwidth_mult = np.power(2.0, n_octaves_per_band)
+
+    low_midi = max(0, hz_to_midi(f_min))
+    high_midi = hz_to_midi(f_max)
+    midi_points = np.linspace(low_midi, high_midi, n_bands)
+    hz_pts = midi_to_hz(midi_points)
+
+    low_pts = hz_pts / bandwidth_mult
+    high_pts = hz_pts * bandwidth_mult
+
+    low_bins = np.floor(low_pts / df).astype(int)
+    high_bins = np.ceil(high_pts / df).astype(int)
+
+    fb = np.zeros((n_bands, n_freqs))
+
+    for i in range(n_bands):
+        fb[i, low_bins[i] : high_bins[i] + 1] = 1.0
+
+    fb[0, : low_bins[0]] = 1.0
+    fb[-1, high_bins[-1] + 1 :] = 1.0
+
+    return torch.as_tensor(fb)
+
+
+class MusicalBandsplitSpecification(PerceptualBandsplitSpecification):
+    def __init__(
+        self, nfft: int, fs: int, n_bands: int, f_min: float = 0.0, f_max: float = None
+    ) -> None:
+        super().__init__(
+            fbank_fn=musical_filterbank,
+            nfft=nfft,
+            fs=fs,
+            n_bands=n_bands,
+            f_min=f_min,
+            f_max=f_max,
+        )
+
+
+if __name__ == "__main__":
+    import pandas as pd
+
+    band_defs = []
+
+    for bands in [VocalBandsplitSpecification]:
+        band_name = bands.__name__.replace("BandsplitSpecification", "")
+
+        mbs = bands(nfft=2048, fs=44100).get_band_specs()
+
+        for i, (f_min, f_max) in enumerate(mbs):
+            band_defs.append(
+                {"band": band_name, "band_index": i, "f_min": f_min, "f_max": f_max}
+            )
+
+    df = pd.DataFrame(band_defs)
+    df.to_csv("vox7bands.csv", index=False)

mvsepless/models/bs_roformer/__init__.py

@@ -1,14 +1,14 @@
-from .bs_roformer import BSRoformer
-from .bs_conformer import BSConformer
-from .bs_roformer_sw import BSRoformer_SW
-try:
-    from neuralop.models import FNO1d
-    from .bs_roformer_fno import BSRoformer_FNO
-except:
-    pass
-from .bs_roformer_hyperace import BSRoformerHyperACE
-from .bs_roformer_hyperace2 import BSRoformerHyperACE_2
-from .bs_roformer_conditional import BSRoformer_Conditional
-from .bs_roformer_unwa_inst_large_2 import BSRoformer_2
-from .mel_band_roformer import MelBandRoformer
-from .mel_band_conformer import MelBandConformer
+from .bs_roformer import BSRoformer
+from .bs_conformer import BSConformer
+from .bs_roformer_sw import BSRoformer_SW
+try:
+    from neuralop.models import FNO1d
+    from .bs_roformer_fno import BSRoformer_FNO
+except:
+    pass
+from .bs_roformer_hyperace import BSRoformerHyperACE
+from .bs_roformer_hyperace2 import BSRoformerHyperACE_2
+from .bs_roformer_conditional import BSRoformer_Conditional
+from .bs_roformer_unwa_inst_large_2 import BSRoformer_2
+from .mel_band_roformer import MelBandRoformer
+from .mel_band_conformer import MelBandConformer

mvsepless/models/bs_roformer/attend.py

@@ -1,128 +1,128 @@
-from functools import wraps
-from packaging import version
-from collections import namedtuple
-
-import os
-import torch
-from torch import nn, einsum
-import torch.nn.functional as F
-
-from einops import rearrange, reduce
-
-
-FlashAttentionConfig = namedtuple(
-    "FlashAttentionConfig", ["enable_flash", "enable_math", "enable_mem_efficient"]
-)
-
-
-def exists(val):
-    return val is not None
-
-
-def default(v, d):
-    return v if exists(v) else d
-
-
-def once(fn):
-    called = False
-
-    @wraps(fn)
-    def inner(x):
-        nonlocal called
-        if called:
-            return
-        called = True
-        return fn(x)
-
-    return inner
-
-
-print_once = once(print)
-
-
-class Attend(nn.Module):
-    def __init__(self, dropout=0.0, flash=False, scale=None):
-        super().__init__()
-        self.scale = scale
-        self.dropout = dropout
-        self.attn_dropout = nn.Dropout(dropout)
-
-        self.flash = flash
-        self.use_torch_2_sdpa = False
-        self._config_checked = False
-        
-        # Проверяем версию PyTorch при первом вызове
-        if flash and not self._config_checked:
-            if version.parse(torch.__version__) >= version.parse("2.0.0"):
-                print_once("PyTorch >= 2.0 detected, will use SDPA if available.")
-                self.use_torch_2_sdpa = True
-                
-                # Настройки для PyTorch >= 2.0
-                self.cpu_config = FlashAttentionConfig(True, True, True)
-                self.cuda_config = None
-                
-                if torch.cuda.is_available():
-                    device_properties = torch.cuda.get_device_properties(torch.device("cuda"))
-                    device_version = version.parse(
-                        f"{device_properties.major}.{device_properties.minor}"
-                    )
-                    
-                    if device_version >= version.parse("8.0"):
-                        if os.name == "nt":
-                            print_once(
-                                "Windows OS detected, using math or mem efficient attention if input tensor is on cuda"
-                            )
-                            self.cuda_config = FlashAttentionConfig(False, True, True)
-                        else:
-                            print_once(
-                                "GPU Compute Capability equal or above 8.0, using flash attention if input tensor is on cuda"
-                            )
-                            self.cuda_config = FlashAttentionConfig(True, False, False)
-                    else:
-                        print_once(
-                            "GPU Compute Capability below 8.0, using math or mem efficient attention if input tensor is on cuda"
-                        )
-                        self.cuda_config = FlashAttentionConfig(False, True, True)
-            else:
-                print_once("PyTorch < 2.0 detected, flash attention will use einsum fallback.")
-                self.use_torch_2_sdpa = False
-            
-            self._config_checked = True
-
-    def flash_attn_torch2(self, q, k, v):
-        """SDPA для PyTorch >= 2.0"""
-        if exists(self.scale):
-            default_scale = q.shape[-1] ** -0.5
-            q = q * (self.scale / default_scale)
-
-        is_cuda = q.is_cuda
-        config = self.cuda_config if is_cuda else self.cpu_config
-
-        with torch.backends.cuda.sdp_kernel(**config._asdict()):
-            out = F.scaled_dot_product_attention(
-                q, k, v, dropout_p=self.dropout if self.training else 0.0
-            )
-
-        return out
-
-    def forward(self, q, k, v):
-        q_len, k_len, device = q.shape[-2], k.shape[-2], q.device
-
-        scale = default(self.scale, q.shape[-1] ** -0.5)
-
-        if self.flash and self.use_torch_2_sdpa:
-            try:
-                return self.flash_attn_torch2(q, k, v)
-            except Exception as e:
-                print(f"Flash attention failed: {e}. Falling back to einsum.")
-                self.use_torch_2_sdpa = False
-
-        # Fallback для PyTorch < 2.0 или если flash отключен
-        sim = einsum(f"b h i d, b h j d -> b h i j", q, k) * scale
-
-        attn = sim.softmax(dim=-1)
-        attn = self.attn_dropout(attn)
-
-        out = einsum(f"b h i j, b h j d -> b h i d", attn, v)
-
+from functools import wraps
+from packaging import version
+from collections import namedtuple
+
+import os
+import torch
+from torch import nn, einsum
+import torch.nn.functional as F
+
+from einops import rearrange, reduce
+
+
+FlashAttentionConfig = namedtuple(
+    "FlashAttentionConfig", ["enable_flash", "enable_math", "enable_mem_efficient"]
+)
+
+
+def exists(val):
+    return val is not None
+
+
+def default(v, d):
+    return v if exists(v) else d
+
+
+def once(fn):
+    called = False
+
+    @wraps(fn)
+    def inner(x):
+        nonlocal called
+        if called:
+            return
+        called = True
+        return fn(x)
+
+    return inner
+
+
+print_once = once(print)
+
+
+class Attend(nn.Module):
+    def __init__(self, dropout=0.0, flash=False, scale=None):
+        super().__init__()
+        self.scale = scale
+        self.dropout = dropout
+        self.attn_dropout = nn.Dropout(dropout)
+
+        self.flash = flash
+        self.use_torch_2_sdpa = False
+        self._config_checked = False
+        
+        # Проверяем версию PyTorch при первом вызове
+        if flash and not self._config_checked:
+            if version.parse(torch.__version__) >= version.parse("2.0.0"):
+                print_once("PyTorch >= 2.0 detected, will use SDPA if available.")
+                self.use_torch_2_sdpa = True
+                
+                # Настройки для PyTorch >= 2.0
+                self.cpu_config = FlashAttentionConfig(True, True, True)
+                self.cuda_config = None
+                
+                if torch.cuda.is_available():
+                    device_properties = torch.cuda.get_device_properties(torch.device("cuda"))
+                    device_version = version.parse(
+                        f"{device_properties.major}.{device_properties.minor}"
+                    )
+                    
+                    if device_version >= version.parse("8.0"):
+                        if os.name == "nt":
+                            print_once(
+                                "Windows OS detected, using math or mem efficient attention if input tensor is on cuda"
+                            )
+                            self.cuda_config = FlashAttentionConfig(False, True, True)
+                        else:
+                            print_once(
+                                "GPU Compute Capability equal or above 8.0, using flash attention if input tensor is on cuda"
+                            )
+                            self.cuda_config = FlashAttentionConfig(True, False, False)
+                    else:
+                        print_once(
+                            "GPU Compute Capability below 8.0, using math or mem efficient attention if input tensor is on cuda"
+                        )
+                        self.cuda_config = FlashAttentionConfig(False, True, True)
+            else:
+                print_once("PyTorch < 2.0 detected, flash attention will use einsum fallback.")
+                self.use_torch_2_sdpa = False
+            
+            self._config_checked = True
+
+    def flash_attn_torch2(self, q, k, v):
+        """SDPA для PyTorch >= 2.0"""
+        if exists(self.scale):
+            default_scale = q.shape[-1] ** -0.5
+            q = q * (self.scale / default_scale)
+
+        is_cuda = q.is_cuda
+        config = self.cuda_config if is_cuda else self.cpu_config
+
+        with torch.backends.cuda.sdp_kernel(**config._asdict()):
+            out = F.scaled_dot_product_attention(
+                q, k, v, dropout_p=self.dropout if self.training else 0.0
+            )
+
+        return out
+
+    def forward(self, q, k, v):
+        q_len, k_len, device = q.shape[-2], k.shape[-2], q.device
+
+        scale = default(self.scale, q.shape[-1] ** -0.5)
+
+        if self.flash and self.use_torch_2_sdpa:
+            try:
+                return self.flash_attn_torch2(q, k, v)
+            except Exception as e:
+                print(f"Flash attention failed: {e}. Falling back to einsum.")
+                self.use_torch_2_sdpa = False
+
+        # Fallback для PyTorch < 2.0 или если flash отключен
+        sim = einsum(f"b h i d, b h j d -> b h i j", q, k) * scale
+
+        attn = sim.softmax(dim=-1)
+        attn = self.attn_dropout(attn)
+
+        out = einsum(f"b h i j, b h j d -> b h i d", attn, v)
+
         return out

mvsepless/models/bs_roformer/attend_sage.py

@@ -1,147 +1,147 @@
-from functools import wraps
-from packaging import version
-from collections import namedtuple
-
-import os
-import torch
-from torch import nn, einsum
-import torch.nn.functional as F
-
-from einops import rearrange, reduce
-
-
-def _print_once(msg):
-    printed = False
-
-    def inner():
-        nonlocal printed
-        if not printed:
-            print(msg)
-            printed = True
-
-    return inner
-
-
-# Проверяем доступность SageAttention
-try:
-    from sageattention import sageattn
-    _has_sage_attention = True
-except ImportError:
-    _has_sage_attention = False
-    _print_sage_not_found = _print_once(
-        "SageAttention not found. Will fall back to PyTorch SDPA (if available) or manual einsum."
-    )
-    _print_sage_not_found()
-
-
-def exists(val):
-    return val is not None
-
-
-def default(v, d):
-    return v if exists(v) else d
-
-
-class Attend(nn.Module):
-    def __init__(self, dropout=0.0, flash=False, scale=None):
-        super().__init__()
-        self.scale = scale
-        self.dropout = dropout
-
-        self.use_sage = flash and _has_sage_attention
-        self.use_pytorch_sdpa = False
-        self._sdpa_checked = False
-        self.flash = flash
-        
-        # Инициализируем сообщения
-        self._init_messages = False
-        
-        if flash and not self.use_sage:
-            if not self._sdpa_checked:
-                if version.parse(torch.__version__) >= version.parse("2.0.0"):
-                    self.use_pytorch_sdpa = True
-                self._sdpa_checked = True
-
-        self.attn_dropout = nn.Dropout(dropout)
-
-    def _print_init_messages(self):
-        """Печатаем сообщения инициализации один раз"""
-        if self._init_messages:
-            return
-            
-        if self.flash:
-            if self.use_sage:
-                print_once = _print_once("Using SageAttention backend.")
-                print_once()
-            elif self.use_pytorch_sdpa:
-                print_once = _print_once(
-                    "Using PyTorch SDPA backend (FlashAttention-2, Memory-Efficient, or Math)."
-                )
-                print_once()
-            else:
-                print_once = _print_once(
-                    "Flash attention requested but Pytorch < 2.0 and SageAttention not found. Falling back to einsum."
-                )
-                print_once()
-        
-        self._init_messages = True
-
-    def forward(self, q, k, v):
-        q_len, k_len, device = q.shape[-2], k.shape[-2], q.device
-
-        # Печатаем сообщения инициализации при первом вызове
-        self._print_init_messages()
-
-        # Пробуем SageAttention если доступен
-        if self.use_sage and self.flash:
-            try:
-                # Исправленный вызов: убрали повторный try-except
-                out = sageattn(q, k, v, tensor_layout="HND", is_causal=False)
-                return out
-            except Exception as e:
-                print(f"SageAttention failed with error: {e}. Falling back.")
-                self.use_sage = False
-                if not self._sdpa_checked:
-                    if version.parse(torch.__version__) >= version.parse("2.0.0"):
-                        self.use_pytorch_sdpa = True
-                        print_once = _print_once(
-                            "Falling back to PyTorch SDPA."
-                        )
-                        print_once()
-                    else:
-                        print_once = _print_once("Falling back to einsum.")
-                        print_once()
-                    self._sdpa_checked = True
-
-        # Пробуем PyTorch SDPA если доступен
-        if self.use_pytorch_sdpa and self.flash:
-            try:
-                # Для PyTorch >= 2.0
-                if version.parse(torch.__version__) >= version.parse("2.0.0"):
-                    with torch.backends.cuda.sdp_kernel(
-                        enable_flash=True, enable_math=True, enable_mem_efficient=True
-                    ):
-                        out = F.scaled_dot_product_attention(
-                            q,
-                            k,
-                            v,
-                            attn_mask=None,
-                            dropout_p=self.dropout if self.training else 0.0,
-                            is_causal=False,
-                        )
-                    return out
-            except Exception as e:
-                print(f"PyTorch SDPA failed with error: {e}. Falling back to einsum.")
-                self.use_pytorch_sdpa = False
-
-        # Fallback на einsum (работает в PyTorch 1.13+)
-        scale = default(self.scale, q.shape[-1] ** -0.5)
-
-        sim = einsum(f"b h i d, b h j d -> b h i j", q, k) * scale
-
-        attn = sim.softmax(dim=-1)
-        attn = self.attn_dropout(attn)
-
-        out = einsum(f"b h i j, b h j d -> b h i d", attn, v)
-
+from functools import wraps
+from packaging import version
+from collections import namedtuple
+
+import os
+import torch
+from torch import nn, einsum
+import torch.nn.functional as F
+
+from einops import rearrange, reduce
+
+
+def _print_once(msg):
+    printed = False
+
+    def inner():
+        nonlocal printed
+        if not printed:
+            print(msg)
+            printed = True
+
+    return inner
+
+
+# Проверяем доступность SageAttention
+try:
+    from sageattention import sageattn
+    _has_sage_attention = True
+except ImportError:
+    _has_sage_attention = False
+    _print_sage_not_found = _print_once(
+        "SageAttention not found. Will fall back to PyTorch SDPA (if available) or manual einsum."
+    )
+    _print_sage_not_found()
+
+
+def exists(val):
+    return val is not None
+
+
+def default(v, d):
+    return v if exists(v) else d
+
+
+class Attend(nn.Module):
+    def __init__(self, dropout=0.0, flash=False, scale=None):
+        super().__init__()
+        self.scale = scale
+        self.dropout = dropout
+
+        self.use_sage = flash and _has_sage_attention
+        self.use_pytorch_sdpa = False
+        self._sdpa_checked = False
+        self.flash = flash
+        
+        # Инициализируем сообщения
+        self._init_messages = False
+        
+        if flash and not self.use_sage:
+            if not self._sdpa_checked:
+                if version.parse(torch.__version__) >= version.parse("2.0.0"):
+                    self.use_pytorch_sdpa = True
+                self._sdpa_checked = True
+
+        self.attn_dropout = nn.Dropout(dropout)
+
+    def _print_init_messages(self):
+        """Печатаем сообщения инициализации один раз"""
+        if self._init_messages:
+            return
+            
+        if self.flash:
+            if self.use_sage:
+                print_once = _print_once("Using SageAttention backend.")
+                print_once()
+            elif self.use_pytorch_sdpa:
+                print_once = _print_once(
+                    "Using PyTorch SDPA backend (FlashAttention-2, Memory-Efficient, or Math)."
+                )
+                print_once()
+            else:
+                print_once = _print_once(
+                    "Flash attention requested but Pytorch < 2.0 and SageAttention not found. Falling back to einsum."
+                )
+                print_once()
+        
+        self._init_messages = True
+
+    def forward(self, q, k, v):
+        q_len, k_len, device = q.shape[-2], k.shape[-2], q.device
+
+        # Печатаем сообщения инициализации при первом вызове
+        self._print_init_messages()
+
+        # Пробуем SageAttention если доступен
+        if self.use_sage and self.flash:
+            try:
+                # Исправленный вызов: убрали повторный try-except
+                out = sageattn(q, k, v, tensor_layout="HND", is_causal=False)
+                return out
+            except Exception as e:
+                print(f"SageAttention failed with error: {e}. Falling back.")
+                self.use_sage = False
+                if not self._sdpa_checked:
+                    if version.parse(torch.__version__) >= version.parse("2.0.0"):
+                        self.use_pytorch_sdpa = True
+                        print_once = _print_once(
+                            "Falling back to PyTorch SDPA."
+                        )
+                        print_once()
+                    else:
+                        print_once = _print_once("Falling back to einsum.")
+                        print_once()
+                    self._sdpa_checked = True
+
+        # Пробуем PyTorch SDPA если доступен
+        if self.use_pytorch_sdpa and self.flash:
+            try:
+                # Для PyTorch >= 2.0
+                if version.parse(torch.__version__) >= version.parse("2.0.0"):
+                    with torch.backends.cuda.sdp_kernel(
+                        enable_flash=True, enable_math=True, enable_mem_efficient=True
+                    ):
+                        out = F.scaled_dot_product_attention(
+                            q,
+                            k,
+                            v,
+                            attn_mask=None,
+                            dropout_p=self.dropout if self.training else 0.0,
+                            is_causal=False,
+                        )
+                    return out
+            except Exception as e:
+                print(f"PyTorch SDPA failed with error: {e}. Falling back to einsum.")
+                self.use_pytorch_sdpa = False
+
+        # Fallback на einsum (работает в PyTorch 1.13+)
+        scale = default(self.scale, q.shape[-1] ** -0.5)
+
+        sim = einsum(f"b h i d, b h j d -> b h i j", q, k) * scale
+
+        attn = sim.softmax(dim=-1)
+        attn = self.attn_dropout(attn)
+
+        out = einsum(f"b h i j, b h j d -> b h i d", attn, v)
+
         return out

mvsepless/models/bs_roformer/attend_sw.py

@@ -1,88 +1,88 @@
-import logging
-import os
-
-import torch
-import torch.nn.functional as F
-from packaging import version
-from torch import Tensor, einsum, nn
-from torch.nn.attention import SDPBackend, sdpa_kernel
-
-logger = logging.getLogger(__name__)
-
-
-class Attend(nn.Module):
-    def __init__(self, dropout: float = 0.0, flash: bool = False, scale=None):
-        super().__init__()
-        self.scale = scale
-        self.dropout = dropout
-        self.attn_dropout = nn.Dropout(dropout)
-
-        self.flash = flash
-        assert not (
-            flash and version.parse(torch.__version__) < version.parse("2.0.0")
-        ), "expected pytorch >= 2.0.0 to use flash attention"
-
-        self.cpu_backends = [
-            SDPBackend.FLASH_ATTENTION,
-            SDPBackend.EFFICIENT_ATTENTION,
-            SDPBackend.MATH,
-        ]
-        self.cuda_backends: list | None = None
-
-        if not torch.cuda.is_available() or not flash:
-            return
-
-        device_properties = torch.cuda.get_device_properties(torch.device("cuda"))
-        device_version = version.parse(
-            f"{device_properties.major}.{device_properties.minor}"
-        )
-
-        if device_version >= version.parse("8.0"):
-            if os.name == "nt":
-                cuda_backends = [SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH]
-                logger.info(f"windows detected, {cuda_backends=}")
-            else:
-                cuda_backends = [SDPBackend.FLASH_ATTENTION]
-                logger.info(f"gpu compute capability >= 8.0, {cuda_backends=}")
-        else:
-            cuda_backends = [SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH]
-            logger.info(f"gpu compute capability < 8.0, {cuda_backends=}")
-
-        self.cuda_backends = cuda_backends
-
-    def flash_attn(self, q: Tensor, k: Tensor, v: Tensor) -> Tensor:
-        _, _heads, _q_len, _, _k_len, is_cuda, _device = (
-            *q.shape,
-            k.shape[-2],
-            q.is_cuda,
-            q.device,
-        )  # type: ignore
-
-        if self.scale is not None:
-            default_scale = q.shape[-1] ** -0.5
-            q = q * (self.scale / default_scale)
-
-        backends = self.cuda_backends if is_cuda else self.cpu_backends
-        with sdpa_kernel(backends=backends):  # type: ignore
-            out = F.scaled_dot_product_attention(
-                q, k, v, dropout_p=self.dropout if self.training else 0.0
-            )
-
-        return out
-
-    def forward(self, q: Tensor, k: Tensor, v: Tensor) -> Tensor:
-        _q_len, _k_len, _device = q.shape[-2], k.shape[-2], q.device
-
-        scale = self.scale or q.shape[-1] ** -0.5
-
-        if self.flash:
-            return self.flash_attn(q, k, v)
-
-        sim = einsum("b h i d, b h j d -> b h i j", q, k) * scale
-
-        attn = sim.softmax(dim=-1)
-        attn = self.attn_dropout(attn)
-
-        out = einsum("b h i j, b h j d -> b h i d", attn, v)
-
-        return out
+import logging
+import os
+
+import torch
+import torch.nn.functional as F
+from packaging import version
+from torch import Tensor, einsum, nn
+from torch.nn.attention import SDPBackend, sdpa_kernel
+
+logger = logging.getLogger(__name__)
+
+
+class Attend(nn.Module):
+    def __init__(self, dropout: float = 0.0, flash: bool = False, scale=None):
+        super().__init__()
+        self.scale = scale
+        self.dropout = dropout
+        self.attn_dropout = nn.Dropout(dropout)
+
+        self.flash = flash
+        assert not (
+            flash and version.parse(torch.__version__) < version.parse("2.0.0")
+        ), "expected pytorch >= 2.0.0 to use flash attention"
+
+        self.cpu_backends = [
+            SDPBackend.FLASH_ATTENTION,
+            SDPBackend.EFFICIENT_ATTENTION,
+            SDPBackend.MATH,
+        ]
+        self.cuda_backends: list | None = None
+
+        if not torch.cuda.is_available() or not flash:
+            return
+
+        device_properties = torch.cuda.get_device_properties(torch.device("cuda"))
+        device_version = version.parse(
+            f"{device_properties.major}.{device_properties.minor}"
+        )
+
+        if device_version >= version.parse("8.0"):
+            if os.name == "nt":
+                cuda_backends = [SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH]
+                logger.info(f"windows detected, {cuda_backends=}")
+            else:
+                cuda_backends = [SDPBackend.FLASH_ATTENTION]
+                logger.info(f"gpu compute capability >= 8.0, {cuda_backends=}")
+        else:
+            cuda_backends = [SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH]
+            logger.info(f"gpu compute capability < 8.0, {cuda_backends=}")
+
+        self.cuda_backends = cuda_backends
+
+    def flash_attn(self, q: Tensor, k: Tensor, v: Tensor) -> Tensor:
+        _, _heads, _q_len, _, _k_len, is_cuda, _device = (
+            *q.shape,
+            k.shape[-2],
+            q.is_cuda,
+            q.device,
+        )  # type: ignore
+
+        if self.scale is not None:
+            default_scale = q.shape[-1] ** -0.5
+            q = q * (self.scale / default_scale)
+
+        backends = self.cuda_backends if is_cuda else self.cpu_backends
+        with sdpa_kernel(backends=backends):  # type: ignore
+            out = F.scaled_dot_product_attention(
+                q, k, v, dropout_p=self.dropout if self.training else 0.0
+            )
+
+        return out
+
+    def forward(self, q: Tensor, k: Tensor, v: Tensor) -> Tensor:
+        _q_len, _k_len, _device = q.shape[-2], k.shape[-2], q.device
+
+        scale = self.scale or q.shape[-1] ** -0.5
+
+        if self.flash:
+            return self.flash_attn(q, k, v)
+
+        sim = einsum("b h i d, b h j d -> b h i j", q, k) * scale
+
+        attn = sim.softmax(dim=-1)
+        attn = self.attn_dropout(attn)
+
+        out = einsum("b h i j, b h j d -> b h i d", attn, v)
+
+        return out

mvsepless/models/bs_roformer/bs_roformer.py

@@ -1,696 +1,696 @@
-from functools import partial
-
-import torch
-from torch import nn, einsum, tensor, Tensor
-from torch.nn import Module, ModuleList
-import torch.nn.functional as F
-
-from .attend import Attend
-
-try:
-    from .attend_sage import Attend as AttendSage
-except:
-    pass
-from torch.utils.checkpoint import checkpoint
-
-from beartype.typing import Tuple, Optional, List, Callable
-from beartype import beartype
-
-from rotary_embedding_torch import RotaryEmbedding
-
-from einops import rearrange, pack, unpack
-from einops.layers.torch import Rearrange
-
-
-def exists(val):
-    return val is not None
-
-
-def default(v, d):
-    return v if exists(v) else d
-
-
-def pack_one(t, pattern):
-    return pack([t], pattern)
-
-
-def unpack_one(t, ps, pattern):
-    return unpack(t, ps, pattern)[0]
-
-
-def l2norm(t):
-    return F.normalize(t, dim=-1, p=2)
-
-
-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
-
-
-class FeedForward(Module):
-    def __init__(self, dim, mult=4, dropout=0.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.0,
-        rotary_embed=None,
-        flash=True,
-        sage_attention=False,
-    ):
-        super().__init__()
-        self.heads = heads
-        self.scale = dim_head**-0.5
-        dim_inner = heads * dim_head
-
-        self.rotary_embed = rotary_embed
-
-        if sage_attention:
-            self.attend = AttendSage(flash=flash, dropout=dropout)
-        else:
-            self.attend = Attend(flash=flash, dropout=dropout)
-
-        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)
-
-        q, k, v = rearrange(
-            self.to_qkv(x), "b n (qkv h d) -> qkv b h n d", qkv=3, h=self.heads
-        )
-
-        if exists(self.rotary_embed):
-            q = self.rotary_embed.rotate_queries_or_keys(q)
-            k = self.rotary_embed.rotate_queries_or_keys(k)
-
-        out = self.attend(q, k, v)
-
-        gates = self.to_gates(x)
-        out = out * rearrange(gates, "b n h -> b h n 1").sigmoid()
-
-        out = rearrange(out, "b h n d -> b n (h d)")
-        return self.to_out(out)
-
-
-class LinearAttention(Module):
-
-    @beartype
-    def __init__(
-        self,
-        *,
-        dim,
-        dim_head=32,
-        heads=8,
-        scale=8,
-        flash=True,
-        dropout=0.0,
-        sage_attention=False,
-    ):
-        super().__init__()
-        dim_inner = dim_head * heads
-        self.norm = RMSNorm(dim)
-
-        self.to_qkv = nn.Sequential(
-            nn.Linear(dim, dim_inner * 3, bias=False),
-            Rearrange("b n (qkv h d) -> qkv b h d n", qkv=3, h=heads),
-        )
-
-        self.temperature = nn.Parameter(torch.ones(heads, 1, 1))
-
-        if sage_attention:
-            self.attend = AttendSage(scale=scale, dropout=dropout, flash=flash)
-        else:
-            self.attend = Attend(scale=scale, dropout=dropout, flash=flash)
-
-        self.to_out = nn.Sequential(
-            Rearrange("b h d n -> b n (h d)"), nn.Linear(dim_inner, dim, bias=False)
-        )
-
-    def forward(self, x):
-        x = self.norm(x)
-
-        q, k, v = self.to_qkv(x)
-
-        q, k = map(l2norm, (q, k))
-        q = q * self.temperature.exp()
-
-        out = self.attend(q, k, v)
-
-        return self.to_out(out)
-
-
-class Transformer(Module):
-    def __init__(
-        self,
-        *,
-        dim,
-        depth,
-        dim_head=64,
-        heads=8,
-        attn_dropout=0.0,
-        ff_dropout=0.0,
-        ff_mult=4,
-        norm_output=True,
-        rotary_embed=None,
-        flash_attn=True,
-        linear_attn=False,
-        sage_attention=False,
-    ):
-        super().__init__()
-        self.layers = ModuleList([])
-
-        for _ in range(depth):
-            if linear_attn:
-                attn = LinearAttention(
-                    dim=dim,
-                    dim_head=dim_head,
-                    heads=heads,
-                    dropout=attn_dropout,
-                    flash=flash_attn,
-                    sage_attention=sage_attention,
-                )
-            else:
-                attn = Attention(
-                    dim=dim,
-                    dim_head=dim_head,
-                    heads=heads,
-                    dropout=attn_dropout,
-                    rotary_embed=rotary_embed,
-                    flash=flash_attn,
-                    sage_attention=sage_attention,
-                )
-
-            self.layers.append(
-                ModuleList(
-                    [attn, 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)
-
-
-class BandSplit(Module):
-    @beartype
-    def __init__(self, dim, dim_inputs: Tuple[int, ...]):
-        super().__init__()
-        self.dim_inputs = dim_inputs
-        self.to_features = ModuleList([])
-
-        for dim_in in dim_inputs:
-            net = nn.Sequential(RMSNorm(dim_in), nn.Linear(dim_in, dim))
-
-            self.to_features.append(net)
-
-    def forward(self, x):
-        x = x.split(self.dim_inputs, dim=-1)
-
-        outs = []
-        for split_input, to_feature in zip(x, self.to_features):
-            split_output = to_feature(split_input)
-            outs.append(split_output)
-
-        return torch.stack(outs, dim=-2)
-
-
-def MLP(dim_in, dim_out, dim_hidden=None, depth=1, activation=nn.Tanh):
-    dim_hidden = default(dim_hidden, dim_in)
-
-    net = []
-    dims = (dim_in, *((dim_hidden,) * (depth - 1)), dim_out)
-
-    for ind, (layer_dim_in, layer_dim_out) in enumerate(zip(dims[:-1], dims[1:])):
-        is_last = ind == (len(dims) - 2)
-
-        net.append(nn.Linear(layer_dim_in, layer_dim_out))
-
-        if is_last:
-            continue
-
-        net.append(activation())
-
-    return nn.Sequential(*net)
-
-
-class MaskEstimator(Module):
-    @beartype
-    def __init__(self, dim, dim_inputs: Tuple[int, ...], depth, mlp_expansion_factor=4):
-        super().__init__()
-        self.dim_inputs = dim_inputs
-        self.to_freqs = ModuleList([])
-        dim_hidden = dim * mlp_expansion_factor
-
-        for dim_in in dim_inputs:
-            net = []
-
-            mlp = nn.Sequential(
-                MLP(dim, dim_in * 2, dim_hidden=dim_hidden, depth=depth), nn.GLU(dim=-1)
-            )
-
-            self.to_freqs.append(mlp)
-
-    def forward(self, x):
-        x = x.unbind(dim=-2)
-
-        outs = []
-
-        for band_features, mlp in zip(x, self.to_freqs):
-            freq_out = mlp(band_features)
-            outs.append(freq_out)
-
-        return torch.cat(outs, dim=-1)
-
-
-DEFAULT_FREQS_PER_BANDS = (
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    4,
-    4,
-    4,
-    4,
-    4,
-    4,
-    4,
-    4,
-    4,
-    4,
-    4,
-    4,
-    12,
-    12,
-    12,
-    12,
-    12,
-    12,
-    12,
-    12,
-    24,
-    24,
-    24,
-    24,
-    24,
-    24,
-    24,
-    24,
-    48,
-    48,
-    48,
-    48,
-    48,
-    48,
-    48,
-    48,
-    128,
-    129,
-)
-
-
-class BSRoformer(Module):
-
-    @beartype
-    def __init__(
-        self,
-        dim,
-        *,
-        depth,
-        stereo=False,
-        num_stems=1,
-        time_transformer_depth=2,
-        freq_transformer_depth=2,
-        linear_transformer_depth=0,
-        freqs_per_bands: Tuple[int, ...] = DEFAULT_FREQS_PER_BANDS,
-        dim_head=64,
-        heads=8,
-        attn_dropout=0.0,
-        ff_dropout=0.0,
-        flash_attn=True,
-        dim_freqs_in=1025,
-        stft_n_fft=2048,
-        stft_hop_length=512,
-        stft_win_length=2048,
-        stft_normalized=False,
-        stft_window_fn: Optional[Callable] = None,
-        zero_dc=True,
-        mask_estimator_depth=2,
-        multi_stft_resolution_loss_weight=1.0,
-        multi_stft_resolutions_window_sizes: Tuple[int, ...] = (
-            4096,
-            2048,
-            1024,
-            512,
-            256,
-        ),
-        multi_stft_hop_size=147,
-        multi_stft_normalized=False,
-        multi_stft_window_fn: Callable = torch.hann_window,
-        mlp_expansion_factor=4,
-        use_torch_checkpoint=False,
-        skip_connection=False,
-        sage_attention=False,
-    ):
-        super().__init__()
-
-        self.stereo = stereo
-        self.audio_channels = 2 if stereo else 1
-        self.num_stems = num_stems
-        self.use_torch_checkpoint = use_torch_checkpoint
-        self.skip_connection = skip_connection
-
-        self.layers = ModuleList([])
-
-        if sage_attention:
-            print("Use Sage Attention")
-
-        transformer_kwargs = dict(
-            dim=dim,
-            heads=heads,
-            dim_head=dim_head,
-            attn_dropout=attn_dropout,
-            ff_dropout=ff_dropout,
-            flash_attn=flash_attn,
-            norm_output=False,
-            sage_attention=sage_attention,
-        )
-
-        time_rotary_embed = RotaryEmbedding(dim=dim_head)
-        freq_rotary_embed = RotaryEmbedding(dim=dim_head)
-
-        for _ in range(depth):
-            tran_modules = []
-            if linear_transformer_depth > 0:
-                tran_modules.append(
-                    Transformer(
-                        depth=linear_transformer_depth,
-                        linear_attn=True,
-                        **transformer_kwargs,
-                    )
-                )
-            tran_modules.append(
-                Transformer(
-                    depth=time_transformer_depth,
-                    rotary_embed=time_rotary_embed,
-                    **transformer_kwargs,
-                )
-            )
-            tran_modules.append(
-                Transformer(
-                    depth=freq_transformer_depth,
-                    rotary_embed=freq_rotary_embed,
-                    **transformer_kwargs,
-                )
-            )
-            self.layers.append(nn.ModuleList(tran_modules))
-
-        self.final_norm = RMSNorm(dim)
-
-        self.stft_kwargs = dict(
-            n_fft=stft_n_fft,
-            hop_length=stft_hop_length,
-            win_length=stft_win_length,
-            normalized=stft_normalized,
-        )
-
-        self.stft_window_fn = partial(
-            default(stft_window_fn, torch.hann_window), stft_win_length
-        )
-
-        freqs = torch.stft(
-            torch.randn(1, 4096),
-            **self.stft_kwargs,
-            window=torch.ones(stft_win_length),
-            return_complex=True,
-        ).shape[1]
-
-        assert len(freqs_per_bands) > 1
-        assert (
-            sum(freqs_per_bands) == freqs
-        ), f"the number of freqs in the bands must equal {freqs} based on the STFT settings, but got {sum(freqs_per_bands)}"
-
-        freqs_per_bands_with_complex = tuple(
-            2 * f * self.audio_channels for f in freqs_per_bands
-        )
-
-        self.band_split = BandSplit(dim=dim, dim_inputs=freqs_per_bands_with_complex)
-
-        self.mask_estimators = nn.ModuleList([])
-
-        for _ in range(num_stems):
-            mask_estimator = MaskEstimator(
-                dim=dim,
-                dim_inputs=freqs_per_bands_with_complex,
-                depth=mask_estimator_depth,
-                mlp_expansion_factor=mlp_expansion_factor,
-            )
-
-            self.mask_estimators.append(mask_estimator)
-
-        self.zero_dc = zero_dc
-
-        self.multi_stft_resolution_loss_weight = multi_stft_resolution_loss_weight
-        self.multi_stft_resolutions_window_sizes = multi_stft_resolutions_window_sizes
-        self.multi_stft_n_fft = stft_n_fft
-        self.multi_stft_window_fn = multi_stft_window_fn
-
-        self.multi_stft_kwargs = dict(
-            hop_length=multi_stft_hop_size, normalized=multi_stft_normalized
-        )
-
-    def forward(self, raw_audio, target=None, return_loss_breakdown=False):
-
-        device = raw_audio.device
-
-        x_is_mps = True if device.type == "mps" else False
-
-        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 needs to be set to True if passing in audio signal that is stereo (channel dimension of 2). also need to be False if mono (channel dimension of 1)"
-
-        raw_audio, batch_audio_channel_packed_shape = pack_one(raw_audio, "* t")
-
-        stft_window = self.stft_window_fn(device=device)
-
-        try:
-            stft_repr = torch.stft(
-                raw_audio, **self.stft_kwargs, window=stft_window, return_complex=True
-            )
-        except:
-            stft_repr = torch.stft(
-                raw_audio.cpu() if x_is_mps else raw_audio,
-                **self.stft_kwargs,
-                window=stft_window.cpu() if x_is_mps else stft_window,
-                return_complex=True,
-            ).to(device)
-        stft_repr = torch.view_as_real(stft_repr)
-
-        stft_repr = unpack_one(stft_repr, batch_audio_channel_packed_shape, "* f t c")
-
-        stft_repr = rearrange(stft_repr, "b s f t c -> b (f s) t c")
-
-        x = rearrange(stft_repr, "b f t c -> b t (f c)")
-
-        if self.use_torch_checkpoint:
-            x = checkpoint(self.band_split, x, use_reentrant=False)
-        else:
-            x = self.band_split(x)
-
-        store = [None] * len(self.layers)
-        for i, transformer_block in enumerate(self.layers):
-
-            if len(transformer_block) == 3:
-                linear_transformer, time_transformer, freq_transformer = (
-                    transformer_block
-                )
-
-                x, ft_ps = pack([x], "b * d")
-                if self.use_torch_checkpoint:
-                    x = checkpoint(linear_transformer, x, use_reentrant=False)
-                else:
-                    x = linear_transformer(x)
-                (x,) = unpack(x, ft_ps, "b * d")
-            else:
-                time_transformer, freq_transformer = transformer_block
-
-            if self.skip_connection:
-                for j in range(i):
-                    x = x + store[j]
-
-            x = rearrange(x, "b t f d -> b f t d")
-            x, ps = pack([x], "* t d")
-
-            if self.use_torch_checkpoint:
-                x = checkpoint(time_transformer, x, use_reentrant=False)
-            else:
-                x = time_transformer(x)
-
-            (x,) = unpack(x, ps, "* t d")
-            x = rearrange(x, "b f t d -> b t f d")
-            x, ps = pack([x], "* f d")
-
-            if self.use_torch_checkpoint:
-                x = checkpoint(freq_transformer, x, use_reentrant=False)
-            else:
-                x = freq_transformer(x)
-
-            (x,) = unpack(x, ps, "* f d")
-
-            if self.skip_connection:
-                store[i] = x
-
-        x = self.final_norm(x)
-
-        num_stems = len(self.mask_estimators)
-
-        if self.use_torch_checkpoint:
-            mask = torch.stack(
-                [checkpoint(fn, x, use_reentrant=False) for fn in self.mask_estimators],
-                dim=1,
-            )
-        else:
-            mask = torch.stack([fn(x) for fn in self.mask_estimators], dim=1)
-        mask = rearrange(mask, "b n t (f c) -> b n f t c", c=2)
-
-        stft_repr = rearrange(stft_repr, "b f t c -> b 1 f t c")
-
-        stft_repr = torch.view_as_complex(stft_repr)
-        mask = torch.view_as_complex(mask)
-
-        stft_repr = stft_repr * mask
-
-        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 = stft_repr.index_fill(1, tensor(0, device=device), 0.0)
-
-        try:
-            recon_audio = torch.istft(
-                stft_repr,
-                **self.stft_kwargs,
-                window=stft_window,
-                return_complex=False,
-                length=raw_audio.shape[-1],
-            )
-        except:
-            recon_audio = torch.istft(
-                stft_repr.cpu() if x_is_mps else stft_repr,
-                **self.stft_kwargs,
-                window=stft_window.cpu() if x_is_mps else stft_window,
-                return_complex=False,
-                length=raw_audio.shape[-1],
-            ).to(device)
-
-        recon_audio = rearrange(
-            recon_audio, "(b n s) t -> b n s t", s=self.audio_channels, n=num_stems
-        )
-
-        if num_stems == 1:
-            recon_audio = rearrange(recon_audio, "b 1 s t -> b s t")
-
-        if not exists(target):
-            return recon_audio
-
-        if self.num_stems > 1:
-            assert target.ndim == 4 and target.shape[1] == self.num_stems
-
-        if target.ndim == 2:
-            target = rearrange(target, "... t -> ... 1 t")
-
-        target = target[..., : recon_audio.shape[-1]]
-
-        loss = F.l1_loss(recon_audio, target)
-
-        multi_stft_resolution_loss = 0.0
-
-        for window_size in self.multi_stft_resolutions_window_sizes:
-            res_stft_kwargs = dict(
-                n_fft=max(window_size, self.multi_stft_n_fft),
-                win_length=window_size,
-                return_complex=True,
-                window=self.multi_stft_window_fn(window_size, device=device),
-                **self.multi_stft_kwargs,
-            )
-
-            recon_Y = torch.stft(
-                rearrange(recon_audio, "... s t -> (... s) t"), **res_stft_kwargs
-            )
-            target_Y = torch.stft(
-                rearrange(target, "... s t -> (... s) t"), **res_stft_kwargs
-            )
-
-            multi_stft_resolution_loss = multi_stft_resolution_loss + F.l1_loss(
-                recon_Y, target_Y
-            )
-
-        weighted_multi_resolution_loss = (
-            multi_stft_resolution_loss * self.multi_stft_resolution_loss_weight
-        )
-
-        total_loss = loss + weighted_multi_resolution_loss
-
-        if not return_loss_breakdown:
-            return total_loss
-
-        return total_loss, (loss, multi_stft_resolution_loss)
+from functools import partial
+
+import torch
+from torch import nn, einsum, tensor, Tensor
+from torch.nn import Module, ModuleList
+import torch.nn.functional as F
+
+from .attend import Attend
+
+try:
+    from .attend_sage import Attend as AttendSage
+except:
+    pass
+from torch.utils.checkpoint import checkpoint
+
+from beartype.typing import Tuple, Optional, List, Callable
+from beartype import beartype
+
+from rotary_embedding_torch import RotaryEmbedding
+
+from einops import rearrange, pack, unpack
+from einops.layers.torch import Rearrange
+
+
+def exists(val):
+    return val is not None
+
+
+def default(v, d):
+    return v if exists(v) else d
+
+
+def pack_one(t, pattern):
+    return pack([t], pattern)
+
+
+def unpack_one(t, ps, pattern):
+    return unpack(t, ps, pattern)[0]
+
+
+def l2norm(t):
+    return F.normalize(t, dim=-1, p=2)
+
+
+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
+
+
+class FeedForward(Module):
+    def __init__(self, dim, mult=4, dropout=0.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.0,
+        rotary_embed=None,
+        flash=True,
+        sage_attention=False,
+    ):
+        super().__init__()
+        self.heads = heads
+        self.scale = dim_head**-0.5
+        dim_inner = heads * dim_head
+
+        self.rotary_embed = rotary_embed
+
+        if sage_attention:
+            self.attend = AttendSage(flash=flash, dropout=dropout)
+        else:
+            self.attend = Attend(flash=flash, dropout=dropout)
+
+        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)
+
+        q, k, v = rearrange(
+            self.to_qkv(x), "b n (qkv h d) -> qkv b h n d", qkv=3, h=self.heads
+        )
+
+        if exists(self.rotary_embed):
+            q = self.rotary_embed.rotate_queries_or_keys(q)
+            k = self.rotary_embed.rotate_queries_or_keys(k)
+
+        out = self.attend(q, k, v)
+
+        gates = self.to_gates(x)
+        out = out * rearrange(gates, "b n h -> b h n 1").sigmoid()
+
+        out = rearrange(out, "b h n d -> b n (h d)")
+        return self.to_out(out)
+
+
+class LinearAttention(Module):
+
+    @beartype
+    def __init__(
+        self,
+        *,
+        dim,
+        dim_head=32,
+        heads=8,
+        scale=8,
+        flash=True,
+        dropout=0.0,
+        sage_attention=False,
+    ):
+        super().__init__()
+        dim_inner = dim_head * heads
+        self.norm = RMSNorm(dim)
+
+        self.to_qkv = nn.Sequential(
+            nn.Linear(dim, dim_inner * 3, bias=False),
+            Rearrange("b n (qkv h d) -> qkv b h d n", qkv=3, h=heads),
+        )
+
+        self.temperature = nn.Parameter(torch.ones(heads, 1, 1))
+
+        if sage_attention:
+            self.attend = AttendSage(scale=scale, dropout=dropout, flash=flash)
+        else:
+            self.attend = Attend(scale=scale, dropout=dropout, flash=flash)
+
+        self.to_out = nn.Sequential(
+            Rearrange("b h d n -> b n (h d)"), nn.Linear(dim_inner, dim, bias=False)
+        )
+
+    def forward(self, x):
+        x = self.norm(x)
+
+        q, k, v = self.to_qkv(x)
+
+        q, k = map(l2norm, (q, k))
+        q = q * self.temperature.exp()
+
+        out = self.attend(q, k, v)
+
+        return self.to_out(out)
+
+
+class Transformer(Module):
+    def __init__(
+        self,
+        *,
+        dim,
+        depth,
+        dim_head=64,
+        heads=8,
+        attn_dropout=0.0,
+        ff_dropout=0.0,
+        ff_mult=4,
+        norm_output=True,
+        rotary_embed=None,
+        flash_attn=True,
+        linear_attn=False,
+        sage_attention=False,
+    ):
+        super().__init__()
+        self.layers = ModuleList([])
+
+        for _ in range(depth):
+            if linear_attn:
+                attn = LinearAttention(
+                    dim=dim,
+                    dim_head=dim_head,
+                    heads=heads,
+                    dropout=attn_dropout,
+                    flash=flash_attn,
+                    sage_attention=sage_attention,
+                )
+            else:
+                attn = Attention(
+                    dim=dim,
+                    dim_head=dim_head,
+                    heads=heads,
+                    dropout=attn_dropout,
+                    rotary_embed=rotary_embed,
+                    flash=flash_attn,
+                    sage_attention=sage_attention,
+                )
+
+            self.layers.append(
+                ModuleList(
+                    [attn, 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)
+
+
+class BandSplit(Module):
+    @beartype
+    def __init__(self, dim, dim_inputs: Tuple[int, ...]):
+        super().__init__()
+        self.dim_inputs = dim_inputs
+        self.to_features = ModuleList([])
+
+        for dim_in in dim_inputs:
+            net = nn.Sequential(RMSNorm(dim_in), nn.Linear(dim_in, dim))
+
+            self.to_features.append(net)
+
+    def forward(self, x):
+        x = x.split(self.dim_inputs, dim=-1)
+
+        outs = []
+        for split_input, to_feature in zip(x, self.to_features):
+            split_output = to_feature(split_input)
+            outs.append(split_output)
+
+        return torch.stack(outs, dim=-2)
+
+
+def MLP(dim_in, dim_out, dim_hidden=None, depth=1, activation=nn.Tanh):
+    dim_hidden = default(dim_hidden, dim_in)
+
+    net = []
+    dims = (dim_in, *((dim_hidden,) * (depth - 1)), dim_out)
+
+    for ind, (layer_dim_in, layer_dim_out) in enumerate(zip(dims[:-1], dims[1:])):
+        is_last = ind == (len(dims) - 2)
+
+        net.append(nn.Linear(layer_dim_in, layer_dim_out))
+
+        if is_last:
+            continue
+
+        net.append(activation())
+
+    return nn.Sequential(*net)
+
+
+class MaskEstimator(Module):
+    @beartype
+    def __init__(self, dim, dim_inputs: Tuple[int, ...], depth, mlp_expansion_factor=4):
+        super().__init__()
+        self.dim_inputs = dim_inputs
+        self.to_freqs = ModuleList([])
+        dim_hidden = dim * mlp_expansion_factor
+
+        for dim_in in dim_inputs:
+            net = []
+
+            mlp = nn.Sequential(
+                MLP(dim, dim_in * 2, dim_hidden=dim_hidden, depth=depth), nn.GLU(dim=-1)
+            )
+
+            self.to_freqs.append(mlp)
+
+    def forward(self, x):
+        x = x.unbind(dim=-2)
+
+        outs = []
+
+        for band_features, mlp in zip(x, self.to_freqs):
+            freq_out = mlp(band_features)
+            outs.append(freq_out)
+
+        return torch.cat(outs, dim=-1)
+
+
+DEFAULT_FREQS_PER_BANDS = (
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    4,
+    4,
+    4,
+    4,
+    4,
+    4,
+    4,
+    4,
+    4,
+    4,
+    4,
+    4,
+    12,
+    12,
+    12,
+    12,
+    12,
+    12,
+    12,
+    12,
+    24,
+    24,
+    24,
+    24,
+    24,
+    24,
+    24,
+    24,
+    48,
+    48,
+    48,
+    48,
+    48,
+    48,
+    48,
+    48,
+    128,
+    129,
+)
+
+
+class BSRoformer(Module):
+
+    @beartype
+    def __init__(
+        self,
+        dim,
+        *,
+        depth,
+        stereo=False,
+        num_stems=1,
+        time_transformer_depth=2,
+        freq_transformer_depth=2,
+        linear_transformer_depth=0,
+        freqs_per_bands: Tuple[int, ...] = DEFAULT_FREQS_PER_BANDS,
+        dim_head=64,
+        heads=8,
+        attn_dropout=0.0,
+        ff_dropout=0.0,
+        flash_attn=True,
+        dim_freqs_in=1025,
+        stft_n_fft=2048,
+        stft_hop_length=512,
+        stft_win_length=2048,
+        stft_normalized=False,
+        stft_window_fn: Optional[Callable] = None,
+        zero_dc=True,
+        mask_estimator_depth=2,
+        multi_stft_resolution_loss_weight=1.0,
+        multi_stft_resolutions_window_sizes: Tuple[int, ...] = (
+            4096,
+            2048,
+            1024,
+            512,
+            256,
+        ),
+        multi_stft_hop_size=147,
+        multi_stft_normalized=False,
+        multi_stft_window_fn: Callable = torch.hann_window,
+        mlp_expansion_factor=4,
+        use_torch_checkpoint=False,
+        skip_connection=False,
+        sage_attention=False,
+    ):
+        super().__init__()
+
+        self.stereo = stereo
+        self.audio_channels = 2 if stereo else 1
+        self.num_stems = num_stems
+        self.use_torch_checkpoint = use_torch_checkpoint
+        self.skip_connection = skip_connection
+
+        self.layers = ModuleList([])
+
+        if sage_attention:
+            print("Use Sage Attention")
+
+        transformer_kwargs = dict(
+            dim=dim,
+            heads=heads,
+            dim_head=dim_head,
+            attn_dropout=attn_dropout,
+            ff_dropout=ff_dropout,
+            flash_attn=flash_attn,
+            norm_output=False,
+            sage_attention=sage_attention,
+        )
+
+        time_rotary_embed = RotaryEmbedding(dim=dim_head)
+        freq_rotary_embed = RotaryEmbedding(dim=dim_head)
+
+        for _ in range(depth):
+            tran_modules = []
+            if linear_transformer_depth > 0:
+                tran_modules.append(
+                    Transformer(
+                        depth=linear_transformer_depth,
+                        linear_attn=True,
+                        **transformer_kwargs,
+                    )
+                )
+            tran_modules.append(
+                Transformer(
+                    depth=time_transformer_depth,
+                    rotary_embed=time_rotary_embed,
+                    **transformer_kwargs,
+                )
+            )
+            tran_modules.append(
+                Transformer(
+                    depth=freq_transformer_depth,
+                    rotary_embed=freq_rotary_embed,
+                    **transformer_kwargs,
+                )
+            )
+            self.layers.append(nn.ModuleList(tran_modules))
+
+        self.final_norm = RMSNorm(dim)
+
+        self.stft_kwargs = dict(
+            n_fft=stft_n_fft,
+            hop_length=stft_hop_length,
+            win_length=stft_win_length,
+            normalized=stft_normalized,
+        )
+
+        self.stft_window_fn = partial(
+            default(stft_window_fn, torch.hann_window), stft_win_length
+        )
+
+        freqs = torch.stft(
+            torch.randn(1, 4096),
+            **self.stft_kwargs,
+            window=torch.ones(stft_win_length),
+            return_complex=True,
+        ).shape[1]
+
+        assert len(freqs_per_bands) > 1
+        assert (
+            sum(freqs_per_bands) == freqs
+        ), f"the number of freqs in the bands must equal {freqs} based on the STFT settings, but got {sum(freqs_per_bands)}"
+
+        freqs_per_bands_with_complex = tuple(
+            2 * f * self.audio_channels for f in freqs_per_bands
+        )
+
+        self.band_split = BandSplit(dim=dim, dim_inputs=freqs_per_bands_with_complex)
+
+        self.mask_estimators = nn.ModuleList([])
+
+        for _ in range(num_stems):
+            mask_estimator = MaskEstimator(
+                dim=dim,
+                dim_inputs=freqs_per_bands_with_complex,
+                depth=mask_estimator_depth,
+                mlp_expansion_factor=mlp_expansion_factor,
+            )
+
+            self.mask_estimators.append(mask_estimator)
+
+        self.zero_dc = zero_dc
+
+        self.multi_stft_resolution_loss_weight = multi_stft_resolution_loss_weight
+        self.multi_stft_resolutions_window_sizes = multi_stft_resolutions_window_sizes
+        self.multi_stft_n_fft = stft_n_fft
+        self.multi_stft_window_fn = multi_stft_window_fn
+
+        self.multi_stft_kwargs = dict(
+            hop_length=multi_stft_hop_size, normalized=multi_stft_normalized
+        )
+
+    def forward(self, raw_audio, target=None, return_loss_breakdown=False):
+
+        device = raw_audio.device
+
+        x_is_mps = True if device.type == "mps" else False
+
+        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 needs to be set to True if passing in audio signal that is stereo (channel dimension of 2). also need to be False if mono (channel dimension of 1)"
+
+        raw_audio, batch_audio_channel_packed_shape = pack_one(raw_audio, "* t")
+
+        stft_window = self.stft_window_fn(device=device)
+
+        try:
+            stft_repr = torch.stft(
+                raw_audio, **self.stft_kwargs, window=stft_window, return_complex=True
+            )
+        except:
+            stft_repr = torch.stft(
+                raw_audio.cpu() if x_is_mps else raw_audio,
+                **self.stft_kwargs,
+                window=stft_window.cpu() if x_is_mps else stft_window,
+                return_complex=True,
+            ).to(device)
+        stft_repr = torch.view_as_real(stft_repr)
+
+        stft_repr = unpack_one(stft_repr, batch_audio_channel_packed_shape, "* f t c")
+
+        stft_repr = rearrange(stft_repr, "b s f t c -> b (f s) t c")
+
+        x = rearrange(stft_repr, "b f t c -> b t (f c)")
+
+        if self.use_torch_checkpoint:
+            x = checkpoint(self.band_split, x, use_reentrant=False)
+        else:
+            x = self.band_split(x)
+
+        store = [None] * len(self.layers)
+        for i, transformer_block in enumerate(self.layers):
+
+            if len(transformer_block) == 3:
+                linear_transformer, time_transformer, freq_transformer = (
+                    transformer_block
+                )
+
+                x, ft_ps = pack([x], "b * d")
+                if self.use_torch_checkpoint:
+                    x = checkpoint(linear_transformer, x, use_reentrant=False)
+                else:
+                    x = linear_transformer(x)
+                (x,) = unpack(x, ft_ps, "b * d")
+            else:
+                time_transformer, freq_transformer = transformer_block
+
+            if self.skip_connection:
+                for j in range(i):
+                    x = x + store[j]
+
+            x = rearrange(x, "b t f d -> b f t d")
+            x, ps = pack([x], "* t d")
+
+            if self.use_torch_checkpoint:
+                x = checkpoint(time_transformer, x, use_reentrant=False)
+            else:
+                x = time_transformer(x)
+
+            (x,) = unpack(x, ps, "* t d")
+            x = rearrange(x, "b f t d -> b t f d")
+            x, ps = pack([x], "* f d")
+
+            if self.use_torch_checkpoint:
+                x = checkpoint(freq_transformer, x, use_reentrant=False)
+            else:
+                x = freq_transformer(x)
+
+            (x,) = unpack(x, ps, "* f d")
+
+            if self.skip_connection:
+                store[i] = x
+
+        x = self.final_norm(x)
+
+        num_stems = len(self.mask_estimators)
+
+        if self.use_torch_checkpoint:
+            mask = torch.stack(
+                [checkpoint(fn, x, use_reentrant=False) for fn in self.mask_estimators],
+                dim=1,
+            )
+        else:
+            mask = torch.stack([fn(x) for fn in self.mask_estimators], dim=1)
+        mask = rearrange(mask, "b n t (f c) -> b n f t c", c=2)
+
+        stft_repr = rearrange(stft_repr, "b f t c -> b 1 f t c")
+
+        stft_repr = torch.view_as_complex(stft_repr)
+        mask = torch.view_as_complex(mask)
+
+        stft_repr = stft_repr * mask
+
+        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 = stft_repr.index_fill(1, tensor(0, device=device), 0.0)
+
+        try:
+            recon_audio = torch.istft(
+                stft_repr,
+                **self.stft_kwargs,
+                window=stft_window,
+                return_complex=False,
+                length=raw_audio.shape[-1],
+            )
+        except:
+            recon_audio = torch.istft(
+                stft_repr.cpu() if x_is_mps else stft_repr,
+                **self.stft_kwargs,
+                window=stft_window.cpu() if x_is_mps else stft_window,
+                return_complex=False,
+                length=raw_audio.shape[-1],
+            ).to(device)
+
+        recon_audio = rearrange(
+            recon_audio, "(b n s) t -> b n s t", s=self.audio_channels, n=num_stems
+        )
+
+        if num_stems == 1:
+            recon_audio = rearrange(recon_audio, "b 1 s t -> b s t")
+
+        if not exists(target):
+            return recon_audio
+
+        if self.num_stems > 1:
+            assert target.ndim == 4 and target.shape[1] == self.num_stems
+
+        if target.ndim == 2:
+            target = rearrange(target, "... t -> ... 1 t")
+
+        target = target[..., : recon_audio.shape[-1]]
+
+        loss = F.l1_loss(recon_audio, target)
+
+        multi_stft_resolution_loss = 0.0
+
+        for window_size in self.multi_stft_resolutions_window_sizes:
+            res_stft_kwargs = dict(
+                n_fft=max(window_size, self.multi_stft_n_fft),
+                win_length=window_size,
+                return_complex=True,
+                window=self.multi_stft_window_fn(window_size, device=device),
+                **self.multi_stft_kwargs,
+            )
+
+            recon_Y = torch.stft(
+                rearrange(recon_audio, "... s t -> (... s) t"), **res_stft_kwargs
+            )
+            target_Y = torch.stft(
+                rearrange(target, "... s t -> (... s) t"), **res_stft_kwargs
+            )
+
+            multi_stft_resolution_loss = multi_stft_resolution_loss + F.l1_loss(
+                recon_Y, target_Y
+            )
+
+        weighted_multi_resolution_loss = (
+            multi_stft_resolution_loss * self.multi_stft_resolution_loss_weight
+        )
+
+        total_loss = loss + weighted_multi_resolution_loss
+
+        if not return_loss_breakdown:
+            return total_loss
+
+        return total_loss, (loss, multi_stft_resolution_loss)

mvsepless/models/bs_roformer/bs_roformer_fno.py

@@ -1,704 +1,704 @@
-from functools import partial
-
-import torch
-from torch import nn, einsum, Tensor
-from torch.nn import Module, ModuleList
-import torch.nn.functional as F
-from neuralop.models import FNO1d
-
-from .attend import Attend
-
-try:
-    from .attend_sage import Attend as AttendSage
-except:
-    pass
-from torch.utils.checkpoint import checkpoint
-
-from beartype.typing import Tuple, Optional, List, Callable
-from beartype import beartype
-
-from rotary_embedding_torch import RotaryEmbedding
-
-from einops import rearrange, pack, unpack
-from einops.layers.torch import Rearrange
-
-
-def exists(val):
-    return val is not None
-
-
-def default(v, d):
-    return v if exists(v) else d
-
-
-def pack_one(t, pattern):
-    return pack([t], pattern)
-
-
-def unpack_one(t, ps, pattern):
-    return unpack(t, ps, pattern)[0]
-
-
-def l2norm(t):
-    return F.normalize(t, dim=-1, p=2)
-
-
-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
-
-
-class FeedForward(Module):
-    def __init__(self, dim, mult=4, dropout=0.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.0,
-        rotary_embed=None,
-        flash=True,
-        sage_attention=False,
-    ):
-        super().__init__()
-        self.heads = heads
-        self.scale = dim_head**-0.5
-        dim_inner = heads * dim_head
-
-        self.rotary_embed = rotary_embed
-
-        if sage_attention:
-            self.attend = AttendSage(flash=flash, dropout=dropout)
-        else:
-            self.attend = Attend(flash=flash, dropout=dropout)
-
-        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)
-
-        q, k, v = rearrange(
-            self.to_qkv(x), "b n (qkv h d) -> qkv b h n d", qkv=3, h=self.heads
-        )
-
-        if exists(self.rotary_embed):
-            q = self.rotary_embed.rotate_queries_or_keys(q)
-            k = self.rotary_embed.rotate_queries_or_keys(k)
-
-        out = self.attend(q, k, v)
-
-        gates = self.to_gates(x)
-        out = out * rearrange(gates, "b n h -> b h n 1").sigmoid()
-
-        out = rearrange(out, "b h n d -> b n (h d)")
-        return self.to_out(out)
-
-
-class LinearAttention(Module):
-
-    @beartype
-    def __init__(
-        self,
-        *,
-        dim,
-        dim_head=32,
-        heads=8,
-        scale=8,
-        flash=False,
-        dropout=0.0,
-        sage_attention=False,
-    ):
-        super().__init__()
-        dim_inner = dim_head * heads
-        self.norm = RMSNorm(dim)
-
-        self.to_qkv = nn.Sequential(
-            nn.Linear(dim, dim_inner * 3, bias=False),
-            Rearrange("b n (qkv h d) -> qkv b h d n", qkv=3, h=heads),
-        )
-
-        self.temperature = nn.Parameter(torch.ones(heads, 1, 1))
-
-        if sage_attention:
-            self.attend = AttendSage(scale=scale, dropout=dropout, flash=flash)
-        else:
-            self.attend = Attend(scale=scale, dropout=dropout, flash=flash)
-
-        self.to_out = nn.Sequential(
-            Rearrange("b h d n -> b n (h d)"), nn.Linear(dim_inner, dim, bias=False)
-        )
-
-    def forward(self, x):
-        x = self.norm(x)
-
-        q, k, v = self.to_qkv(x)
-
-        q, k = map(l2norm, (q, k))
-        q = q * self.temperature.exp()
-
-        out = self.attend(q, k, v)
-
-        return self.to_out(out)
-
-
-class Transformer(Module):
-    def __init__(
-        self,
-        *,
-        dim,
-        depth,
-        dim_head=64,
-        heads=8,
-        attn_dropout=0.0,
-        ff_dropout=0.0,
-        ff_mult=4,
-        norm_output=True,
-        rotary_embed=None,
-        flash_attn=True,
-        linear_attn=False,
-        sage_attention=False,
-    ):
-        super().__init__()
-        self.layers = ModuleList([])
-
-        for _ in range(depth):
-            if linear_attn:
-                attn = LinearAttention(
-                    dim=dim,
-                    dim_head=dim_head,
-                    heads=heads,
-                    dropout=attn_dropout,
-                    flash=flash_attn,
-                    sage_attention=sage_attention,
-                )
-            else:
-                attn = Attention(
-                    dim=dim,
-                    dim_head=dim_head,
-                    heads=heads,
-                    dropout=attn_dropout,
-                    rotary_embed=rotary_embed,
-                    flash=flash_attn,
-                    sage_attention=sage_attention,
-                )
-
-            self.layers.append(
-                ModuleList(
-                    [attn, 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)
-
-
-class BandSplit(Module):
-    @beartype
-    def __init__(self, dim, dim_inputs: Tuple[int, ...]):
-        super().__init__()
-        self.dim_inputs = dim_inputs
-        self.to_features = ModuleList([])
-
-        for dim_in in dim_inputs:
-            net = nn.Sequential(RMSNorm(dim_in), nn.Linear(dim_in, dim))
-
-            self.to_features.append(net)
-
-    def forward(self, x):
-        x = x.split(self.dim_inputs, dim=-1)
-
-        outs = []
-        for split_input, to_feature in zip(x, self.to_features):
-            split_output = to_feature(split_input)
-            outs.append(split_output)
-
-        return torch.stack(outs, dim=-2)
-
-
-def MLP(dim_in, dim_out, dim_hidden=None, depth=1, activation=nn.Tanh):
-    dim_hidden = default(dim_hidden, dim_in)
-
-    net = []
-    dims = (dim_in, *((dim_hidden,) * (depth - 1)), dim_out)
-
-    for ind, (layer_dim_in, layer_dim_out) in enumerate(zip(dims[:-1], dims[1:])):
-        is_last = ind == (len(dims) - 2)
-
-        net.append(nn.Linear(layer_dim_in, layer_dim_out))
-
-        if is_last:
-            continue
-
-        net.append(activation())
-
-    return nn.Sequential(*net)
-
-
-class MaskEstimator(Module):
-    @beartype
-    def __init__(self, dim, dim_inputs: Tuple[int, ...], depth, mlp_expansion_factor=4):
-        super().__init__()
-        self.dim_inputs = dim_inputs
-        self.to_freqs = ModuleList([])
-        dim_hidden = dim * mlp_expansion_factor
-
-        for dim_in in dim_inputs:
-            net = []
-
-            mlp = nn.Sequential(
-                FNO1d(
-                    n_modes_height=64,
-                    hidden_channels=dim,
-                    in_channels=dim,
-                    out_channels=dim_in * 2,
-                    lifting_channels=dim,
-                    projection_channels=dim,
-                    n_layers=3,
-                    separable=True,
-                ),
-                nn.GLU(dim=-2),
-            )
-
-            self.to_freqs.append(mlp)
-
-    def forward(self, x):
-        x = x.unbind(dim=-2)
-
-        outs = []
-
-        for band_features, mlp in zip(x, self.to_freqs):
-            band_features = rearrange(band_features, "b t c -> b c t")
-            with torch.autocast(device_type="cuda", enabled=False, dtype=torch.float32):
-                freq_out = mlp(band_features).float()
-            freq_out = rearrange(freq_out, "b c t -> b t c")
-            outs.append(freq_out)
-
-        return torch.cat(outs, dim=-1)
-
-
-DEFAULT_FREQS_PER_BANDS = (
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    4,
-    4,
-    4,
-    4,
-    4,
-    4,
-    4,
-    4,
-    4,
-    4,
-    4,
-    4,
-    12,
-    12,
-    12,
-    12,
-    12,
-    12,
-    12,
-    12,
-    24,
-    24,
-    24,
-    24,
-    24,
-    24,
-    24,
-    24,
-    48,
-    48,
-    48,
-    48,
-    48,
-    48,
-    48,
-    48,
-    128,
-    129,
-)
-
-
-class BSRoformer_FNO(Module):
-
-    @beartype
-    def __init__(
-        self,
-        dim,
-        *,
-        depth,
-        stereo=False,
-        num_stems=1,
-        time_transformer_depth=2,
-        freq_transformer_depth=2,
-        linear_transformer_depth=0,
-        freqs_per_bands: Tuple[int, ...] = DEFAULT_FREQS_PER_BANDS,
-        dim_head=64,
-        heads=8,
-        attn_dropout=0.0,
-        ff_dropout=0.0,
-        flash_attn=True,
-        dim_freqs_in=1025,
-        stft_n_fft=2048,
-        stft_hop_length=512,
-        stft_win_length=2048,
-        stft_normalized=False,
-        stft_window_fn: Optional[Callable] = None,
-        mask_estimator_depth=2,
-        multi_stft_resolution_loss_weight=1.0,
-        multi_stft_resolutions_window_sizes: Tuple[int, ...] = (
-            4096,
-            2048,
-            1024,
-            512,
-            256,
-        ),
-        multi_stft_hop_size=147,
-        multi_stft_normalized=False,
-        multi_stft_window_fn: Callable = torch.hann_window,
-        mlp_expansion_factor=4,
-        use_torch_checkpoint=False,
-        skip_connection=False,
-        sage_attention=False,
-    ):
-        super().__init__()
-
-        self.stereo = stereo
-        self.audio_channels = 2 if stereo else 1
-        self.num_stems = num_stems
-        self.use_torch_checkpoint = use_torch_checkpoint
-        self.skip_connection = skip_connection
-
-        self.layers = ModuleList([])
-
-        if sage_attention:
-            print("Use Sage Attention")
-
-        transformer_kwargs = dict(
-            dim=dim,
-            heads=heads,
-            dim_head=dim_head,
-            attn_dropout=attn_dropout,
-            ff_dropout=ff_dropout,
-            flash_attn=flash_attn,
-            norm_output=False,
-            sage_attention=sage_attention,
-        )
-
-        time_rotary_embed = RotaryEmbedding(dim=dim_head)
-        freq_rotary_embed = RotaryEmbedding(dim=dim_head)
-
-        for _ in range(depth):
-            tran_modules = []
-            if linear_transformer_depth > 0:
-                tran_modules.append(
-                    Transformer(
-                        depth=linear_transformer_depth,
-                        linear_attn=True,
-                        **transformer_kwargs,
-                    )
-                )
-            tran_modules.append(
-                Transformer(
-                    depth=time_transformer_depth,
-                    rotary_embed=time_rotary_embed,
-                    **transformer_kwargs,
-                )
-            )
-            tran_modules.append(
-                Transformer(
-                    depth=freq_transformer_depth,
-                    rotary_embed=freq_rotary_embed,
-                    **transformer_kwargs,
-                )
-            )
-            self.layers.append(nn.ModuleList(tran_modules))
-
-        self.final_norm = RMSNorm(dim)
-
-        self.stft_kwargs = dict(
-            n_fft=stft_n_fft,
-            hop_length=stft_hop_length,
-            win_length=stft_win_length,
-            normalized=stft_normalized,
-        )
-
-        self.stft_window_fn = partial(
-            default(stft_window_fn, torch.hann_window), stft_win_length
-        )
-
-        freqs = torch.stft(
-            torch.randn(1, 4096),
-            **self.stft_kwargs,
-            window=torch.ones(stft_win_length),
-            return_complex=True,
-        ).shape[1]
-
-        assert len(freqs_per_bands) > 1
-        assert (
-            sum(freqs_per_bands) == freqs
-        ), f"the number of freqs in the bands must equal {freqs} based on the STFT settings, but got {sum(freqs_per_bands)}"
-
-        freqs_per_bands_with_complex = tuple(
-            2 * f * self.audio_channels for f in freqs_per_bands
-        )
-
-        self.band_split = BandSplit(dim=dim, dim_inputs=freqs_per_bands_with_complex)
-
-        self.mask_estimators = nn.ModuleList([])
-
-        for _ in range(num_stems):
-            mask_estimator = MaskEstimator(
-                dim=dim,
-                dim_inputs=freqs_per_bands_with_complex,
-                depth=mask_estimator_depth,
-                mlp_expansion_factor=mlp_expansion_factor,
-            )
-
-            self.mask_estimators.append(mask_estimator)
-
-        self.multi_stft_resolution_loss_weight = multi_stft_resolution_loss_weight
-        self.multi_stft_resolutions_window_sizes = multi_stft_resolutions_window_sizes
-        self.multi_stft_n_fft = stft_n_fft
-        self.multi_stft_window_fn = multi_stft_window_fn
-
-        self.multi_stft_kwargs = dict(
-            hop_length=multi_stft_hop_size, normalized=multi_stft_normalized
-        )
-
-    def forward(self, raw_audio, target=None, return_loss_breakdown=False):
-
-        device = raw_audio.device
-
-        x_is_mps = True if device.type == "mps" else False
-
-        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 needs to be set to True if passing in audio signal that is stereo (channel dimension of 2). also need to be False if mono (channel dimension of 1)"
-
-        raw_audio, batch_audio_channel_packed_shape = pack_one(raw_audio, "* t")
-
-        stft_window = self.stft_window_fn(device=device)
-
-        try:
-            stft_repr = torch.stft(
-                raw_audio, **self.stft_kwargs, window=stft_window, return_complex=True
-            )
-        except:
-            stft_repr = torch.stft(
-                raw_audio.cpu() if x_is_mps else raw_audio,
-                **self.stft_kwargs,
-                window=stft_window.cpu() if x_is_mps else stft_window,
-                return_complex=True,
-            ).to(device)
-        stft_repr = torch.view_as_real(stft_repr)
-
-        stft_repr = unpack_one(stft_repr, batch_audio_channel_packed_shape, "* f t c")
-
-        stft_repr = rearrange(stft_repr, "b s f t c -> b (f s) t c")
-
-        x = rearrange(stft_repr, "b f t c -> b t (f c)")
-
-        if self.use_torch_checkpoint:
-            x = checkpoint(self.band_split, x, use_reentrant=False)
-        else:
-            x = self.band_split(x)
-
-        store = [None] * len(self.layers)
-        for i, transformer_block in enumerate(self.layers):
-
-            if len(transformer_block) == 3:
-                linear_transformer, time_transformer, freq_transformer = (
-                    transformer_block
-                )
-
-                x, ft_ps = pack([x], "b * d")
-                if self.use_torch_checkpoint:
-                    x = checkpoint(linear_transformer, x, use_reentrant=False)
-                else:
-                    x = linear_transformer(x)
-                (x,) = unpack(x, ft_ps, "b * d")
-            else:
-                time_transformer, freq_transformer = transformer_block
-
-            if self.skip_connection:
-                for j in range(i):
-                    x = x + store[j]
-
-            x = rearrange(x, "b t f d -> b f t d")
-            x, ps = pack([x], "* t d")
-
-            if self.use_torch_checkpoint:
-                x = checkpoint(time_transformer, x, use_reentrant=False)
-            else:
-                x = time_transformer(x)
-
-            (x,) = unpack(x, ps, "* t d")
-            x = rearrange(x, "b f t d -> b t f d")
-            x, ps = pack([x], "* f d")
-
-            if self.use_torch_checkpoint:
-                x = checkpoint(freq_transformer, x, use_reentrant=False)
-            else:
-                x = freq_transformer(x)
-
-            (x,) = unpack(x, ps, "* f d")
-
-            if self.skip_connection:
-                store[i] = x
-
-        x = self.final_norm(x)
-
-        num_stems = len(self.mask_estimators)
-
-        if self.use_torch_checkpoint:
-            mask = torch.stack(
-                [checkpoint(fn, x, use_reentrant=False) for fn in self.mask_estimators],
-                dim=1,
-            )
-        else:
-            mask = torch.stack([fn(x) for fn in self.mask_estimators], dim=1)
-        mask = rearrange(mask, "b n t (f c) -> b n f t c", c=2)
-
-        stft_repr = rearrange(stft_repr, "b f t c -> b 1 f t c")
-
-        stft_repr = torch.view_as_complex(stft_repr)
-        mask = torch.view_as_complex(mask)
-
-        stft_repr = stft_repr * mask
-
-        stft_repr = rearrange(
-            stft_repr, "b n (f s) t -> (b n s) f t", s=self.audio_channels
-        )
-
-        try:
-            recon_audio = torch.istft(
-                stft_repr,
-                **self.stft_kwargs,
-                window=stft_window,
-                return_complex=False,
-                length=raw_audio.shape[-1],
-            )
-        except:
-            recon_audio = torch.istft(
-                stft_repr.cpu() if x_is_mps else stft_repr,
-                **self.stft_kwargs,
-                window=stft_window.cpu() if x_is_mps else stft_window,
-                return_complex=False,
-                length=raw_audio.shape[-1],
-            ).to(device)
-
-        recon_audio = rearrange(
-            recon_audio, "(b n s) t -> b n s t", s=self.audio_channels, n=num_stems
-        )
-
-        if num_stems == 1:
-            recon_audio = rearrange(recon_audio, "b 1 s t -> b s t")
-
-        if not exists(target):
-            return recon_audio
-
-        if self.num_stems > 1:
-            assert target.ndim == 4 and target.shape[1] == self.num_stems
-
-        if target.ndim == 2:
-            target = rearrange(target, "... t -> ... 1 t")
-
-        target = target[..., : recon_audio.shape[-1]]
-
-        loss = F.l1_loss(recon_audio, target)
-
-        multi_stft_resolution_loss = 0.0
-
-        for window_size in self.multi_stft_resolutions_window_sizes:
-            res_stft_kwargs = dict(
-                n_fft=max(window_size, self.multi_stft_n_fft),
-                win_length=window_size,
-                return_complex=True,
-                window=self.multi_stft_window_fn(window_size, device=device),
-                **self.multi_stft_kwargs,
-            )
-
-            recon_Y = torch.stft(
-                rearrange(recon_audio, "... s t -> (... s) t"), **res_stft_kwargs
-            )
-            target_Y = torch.stft(
-                rearrange(target, "... s t -> (... s) t"), **res_stft_kwargs
-            )
-
-            multi_stft_resolution_loss = multi_stft_resolution_loss + F.l1_loss(
-                recon_Y, target_Y
-            )
-
-        weighted_multi_resolution_loss = (
-            multi_stft_resolution_loss * self.multi_stft_resolution_loss_weight
-        )
-
-        total_loss = loss + weighted_multi_resolution_loss
-
-        if not return_loss_breakdown:
-            return total_loss
-
-        return total_loss, (loss, multi_stft_resolution_loss)
+from functools import partial
+
+import torch
+from torch import nn, einsum, Tensor
+from torch.nn import Module, ModuleList
+import torch.nn.functional as F
+from neuralop.models import FNO1d
+
+from .attend import Attend
+
+try:
+    from .attend_sage import Attend as AttendSage
+except:
+    pass
+from torch.utils.checkpoint import checkpoint
+
+from beartype.typing import Tuple, Optional, List, Callable
+from beartype import beartype
+
+from rotary_embedding_torch import RotaryEmbedding
+
+from einops import rearrange, pack, unpack
+from einops.layers.torch import Rearrange
+
+
+def exists(val):
+    return val is not None
+
+
+def default(v, d):
+    return v if exists(v) else d
+
+
+def pack_one(t, pattern):
+    return pack([t], pattern)
+
+
+def unpack_one(t, ps, pattern):
+    return unpack(t, ps, pattern)[0]
+
+
+def l2norm(t):
+    return F.normalize(t, dim=-1, p=2)
+
+
+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
+
+
+class FeedForward(Module):
+    def __init__(self, dim, mult=4, dropout=0.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.0,
+        rotary_embed=None,
+        flash=True,
+        sage_attention=False,
+    ):
+        super().__init__()
+        self.heads = heads
+        self.scale = dim_head**-0.5
+        dim_inner = heads * dim_head
+
+        self.rotary_embed = rotary_embed
+
+        if sage_attention:
+            self.attend = AttendSage(flash=flash, dropout=dropout)
+        else:
+            self.attend = Attend(flash=flash, dropout=dropout)
+
+        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)
+
+        q, k, v = rearrange(
+            self.to_qkv(x), "b n (qkv h d) -> qkv b h n d", qkv=3, h=self.heads
+        )
+
+        if exists(self.rotary_embed):
+            q = self.rotary_embed.rotate_queries_or_keys(q)
+            k = self.rotary_embed.rotate_queries_or_keys(k)
+
+        out = self.attend(q, k, v)
+
+        gates = self.to_gates(x)
+        out = out * rearrange(gates, "b n h -> b h n 1").sigmoid()
+
+        out = rearrange(out, "b h n d -> b n (h d)")
+        return self.to_out(out)
+
+
+class LinearAttention(Module):
+
+    @beartype
+    def __init__(
+        self,
+        *,
+        dim,
+        dim_head=32,
+        heads=8,
+        scale=8,
+        flash=False,
+        dropout=0.0,
+        sage_attention=False,
+    ):
+        super().__init__()
+        dim_inner = dim_head * heads
+        self.norm = RMSNorm(dim)
+
+        self.to_qkv = nn.Sequential(
+            nn.Linear(dim, dim_inner * 3, bias=False),
+            Rearrange("b n (qkv h d) -> qkv b h d n", qkv=3, h=heads),
+        )
+
+        self.temperature = nn.Parameter(torch.ones(heads, 1, 1))
+
+        if sage_attention:
+            self.attend = AttendSage(scale=scale, dropout=dropout, flash=flash)
+        else:
+            self.attend = Attend(scale=scale, dropout=dropout, flash=flash)
+
+        self.to_out = nn.Sequential(
+            Rearrange("b h d n -> b n (h d)"), nn.Linear(dim_inner, dim, bias=False)
+        )
+
+    def forward(self, x):
+        x = self.norm(x)
+
+        q, k, v = self.to_qkv(x)
+
+        q, k = map(l2norm, (q, k))
+        q = q * self.temperature.exp()
+
+        out = self.attend(q, k, v)
+
+        return self.to_out(out)
+
+
+class Transformer(Module):
+    def __init__(
+        self,
+        *,
+        dim,
+        depth,
+        dim_head=64,
+        heads=8,
+        attn_dropout=0.0,
+        ff_dropout=0.0,
+        ff_mult=4,
+        norm_output=True,
+        rotary_embed=None,
+        flash_attn=True,
+        linear_attn=False,
+        sage_attention=False,
+    ):
+        super().__init__()
+        self.layers = ModuleList([])
+
+        for _ in range(depth):
+            if linear_attn:
+                attn = LinearAttention(
+                    dim=dim,
+                    dim_head=dim_head,
+                    heads=heads,
+                    dropout=attn_dropout,
+                    flash=flash_attn,
+                    sage_attention=sage_attention,
+                )
+            else:
+                attn = Attention(
+                    dim=dim,
+                    dim_head=dim_head,
+                    heads=heads,
+                    dropout=attn_dropout,
+                    rotary_embed=rotary_embed,
+                    flash=flash_attn,
+                    sage_attention=sage_attention,
+                )
+
+            self.layers.append(
+                ModuleList(
+                    [attn, 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)
+
+
+class BandSplit(Module):
+    @beartype
+    def __init__(self, dim, dim_inputs: Tuple[int, ...]):
+        super().__init__()
+        self.dim_inputs = dim_inputs
+        self.to_features = ModuleList([])
+
+        for dim_in in dim_inputs:
+            net = nn.Sequential(RMSNorm(dim_in), nn.Linear(dim_in, dim))
+
+            self.to_features.append(net)
+
+    def forward(self, x):
+        x = x.split(self.dim_inputs, dim=-1)
+
+        outs = []
+        for split_input, to_feature in zip(x, self.to_features):
+            split_output = to_feature(split_input)
+            outs.append(split_output)
+
+        return torch.stack(outs, dim=-2)
+
+
+def MLP(dim_in, dim_out, dim_hidden=None, depth=1, activation=nn.Tanh):
+    dim_hidden = default(dim_hidden, dim_in)
+
+    net = []
+    dims = (dim_in, *((dim_hidden,) * (depth - 1)), dim_out)
+
+    for ind, (layer_dim_in, layer_dim_out) in enumerate(zip(dims[:-1], dims[1:])):
+        is_last = ind == (len(dims) - 2)
+
+        net.append(nn.Linear(layer_dim_in, layer_dim_out))
+
+        if is_last:
+            continue
+
+        net.append(activation())
+
+    return nn.Sequential(*net)
+
+
+class MaskEstimator(Module):
+    @beartype
+    def __init__(self, dim, dim_inputs: Tuple[int, ...], depth, mlp_expansion_factor=4):
+        super().__init__()
+        self.dim_inputs = dim_inputs
+        self.to_freqs = ModuleList([])
+        dim_hidden = dim * mlp_expansion_factor
+
+        for dim_in in dim_inputs:
+            net = []
+
+            mlp = nn.Sequential(
+                FNO1d(
+                    n_modes_height=64,
+                    hidden_channels=dim,
+                    in_channels=dim,
+                    out_channels=dim_in * 2,
+                    lifting_channels=dim,
+                    projection_channels=dim,
+                    n_layers=3,
+                    separable=True,
+                ),
+                nn.GLU(dim=-2),
+            )
+
+            self.to_freqs.append(mlp)
+
+    def forward(self, x):
+        x = x.unbind(dim=-2)
+
+        outs = []
+
+        for band_features, mlp in zip(x, self.to_freqs):
+            band_features = rearrange(band_features, "b t c -> b c t")
+            with torch.autocast(device_type="cuda", enabled=False, dtype=torch.float32):
+                freq_out = mlp(band_features).float()
+            freq_out = rearrange(freq_out, "b c t -> b t c")
+            outs.append(freq_out)
+
+        return torch.cat(outs, dim=-1)
+
+
+DEFAULT_FREQS_PER_BANDS = (
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    4,
+    4,
+    4,
+    4,
+    4,
+    4,
+    4,
+    4,
+    4,
+    4,
+    4,
+    4,
+    12,
+    12,
+    12,
+    12,
+    12,
+    12,
+    12,
+    12,
+    24,
+    24,
+    24,
+    24,
+    24,
+    24,
+    24,
+    24,
+    48,
+    48,
+    48,
+    48,
+    48,
+    48,
+    48,
+    48,
+    128,
+    129,
+)
+
+
+class BSRoformer_FNO(Module):
+
+    @beartype
+    def __init__(
+        self,
+        dim,
+        *,
+        depth,
+        stereo=False,
+        num_stems=1,
+        time_transformer_depth=2,
+        freq_transformer_depth=2,
+        linear_transformer_depth=0,
+        freqs_per_bands: Tuple[int, ...] = DEFAULT_FREQS_PER_BANDS,
+        dim_head=64,
+        heads=8,
+        attn_dropout=0.0,
+        ff_dropout=0.0,
+        flash_attn=True,
+        dim_freqs_in=1025,
+        stft_n_fft=2048,
+        stft_hop_length=512,
+        stft_win_length=2048,
+        stft_normalized=False,
+        stft_window_fn: Optional[Callable] = None,
+        mask_estimator_depth=2,
+        multi_stft_resolution_loss_weight=1.0,
+        multi_stft_resolutions_window_sizes: Tuple[int, ...] = (
+            4096,
+            2048,
+            1024,
+            512,
+            256,
+        ),
+        multi_stft_hop_size=147,
+        multi_stft_normalized=False,
+        multi_stft_window_fn: Callable = torch.hann_window,
+        mlp_expansion_factor=4,
+        use_torch_checkpoint=False,
+        skip_connection=False,
+        sage_attention=False,
+    ):
+        super().__init__()
+
+        self.stereo = stereo
+        self.audio_channels = 2 if stereo else 1
+        self.num_stems = num_stems
+        self.use_torch_checkpoint = use_torch_checkpoint
+        self.skip_connection = skip_connection
+
+        self.layers = ModuleList([])
+
+        if sage_attention:
+            print("Use Sage Attention")
+
+        transformer_kwargs = dict(
+            dim=dim,
+            heads=heads,
+            dim_head=dim_head,
+            attn_dropout=attn_dropout,
+            ff_dropout=ff_dropout,
+            flash_attn=flash_attn,
+            norm_output=False,
+            sage_attention=sage_attention,
+        )
+
+        time_rotary_embed = RotaryEmbedding(dim=dim_head)
+        freq_rotary_embed = RotaryEmbedding(dim=dim_head)
+
+        for _ in range(depth):
+            tran_modules = []
+            if linear_transformer_depth > 0:
+                tran_modules.append(
+                    Transformer(
+                        depth=linear_transformer_depth,
+                        linear_attn=True,
+                        **transformer_kwargs,
+                    )
+                )
+            tran_modules.append(
+                Transformer(
+                    depth=time_transformer_depth,
+                    rotary_embed=time_rotary_embed,
+                    **transformer_kwargs,
+                )
+            )
+            tran_modules.append(
+                Transformer(
+                    depth=freq_transformer_depth,
+                    rotary_embed=freq_rotary_embed,
+                    **transformer_kwargs,
+                )
+            )
+            self.layers.append(nn.ModuleList(tran_modules))
+
+        self.final_norm = RMSNorm(dim)
+
+        self.stft_kwargs = dict(
+            n_fft=stft_n_fft,
+            hop_length=stft_hop_length,
+            win_length=stft_win_length,
+            normalized=stft_normalized,
+        )
+
+        self.stft_window_fn = partial(
+            default(stft_window_fn, torch.hann_window), stft_win_length
+        )
+
+        freqs = torch.stft(
+            torch.randn(1, 4096),
+            **self.stft_kwargs,
+            window=torch.ones(stft_win_length),
+            return_complex=True,
+        ).shape[1]
+
+        assert len(freqs_per_bands) > 1
+        assert (
+            sum(freqs_per_bands) == freqs
+        ), f"the number of freqs in the bands must equal {freqs} based on the STFT settings, but got {sum(freqs_per_bands)}"
+
+        freqs_per_bands_with_complex = tuple(
+            2 * f * self.audio_channels for f in freqs_per_bands
+        )
+
+        self.band_split = BandSplit(dim=dim, dim_inputs=freqs_per_bands_with_complex)
+
+        self.mask_estimators = nn.ModuleList([])
+
+        for _ in range(num_stems):
+            mask_estimator = MaskEstimator(
+                dim=dim,
+                dim_inputs=freqs_per_bands_with_complex,
+                depth=mask_estimator_depth,
+                mlp_expansion_factor=mlp_expansion_factor,
+            )
+
+            self.mask_estimators.append(mask_estimator)
+
+        self.multi_stft_resolution_loss_weight = multi_stft_resolution_loss_weight
+        self.multi_stft_resolutions_window_sizes = multi_stft_resolutions_window_sizes
+        self.multi_stft_n_fft = stft_n_fft
+        self.multi_stft_window_fn = multi_stft_window_fn
+
+        self.multi_stft_kwargs = dict(
+            hop_length=multi_stft_hop_size, normalized=multi_stft_normalized
+        )
+
+    def forward(self, raw_audio, target=None, return_loss_breakdown=False):
+
+        device = raw_audio.device
+
+        x_is_mps = True if device.type == "mps" else False
+
+        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 needs to be set to True if passing in audio signal that is stereo (channel dimension of 2). also need to be False if mono (channel dimension of 1)"
+
+        raw_audio, batch_audio_channel_packed_shape = pack_one(raw_audio, "* t")
+
+        stft_window = self.stft_window_fn(device=device)
+
+        try:
+            stft_repr = torch.stft(
+                raw_audio, **self.stft_kwargs, window=stft_window, return_complex=True
+            )
+        except:
+            stft_repr = torch.stft(
+                raw_audio.cpu() if x_is_mps else raw_audio,
+                **self.stft_kwargs,
+                window=stft_window.cpu() if x_is_mps else stft_window,
+                return_complex=True,
+            ).to(device)
+        stft_repr = torch.view_as_real(stft_repr)
+
+        stft_repr = unpack_one(stft_repr, batch_audio_channel_packed_shape, "* f t c")
+
+        stft_repr = rearrange(stft_repr, "b s f t c -> b (f s) t c")
+
+        x = rearrange(stft_repr, "b f t c -> b t (f c)")
+
+        if self.use_torch_checkpoint:
+            x = checkpoint(self.band_split, x, use_reentrant=False)
+        else:
+            x = self.band_split(x)
+
+        store = [None] * len(self.layers)
+        for i, transformer_block in enumerate(self.layers):
+
+            if len(transformer_block) == 3:
+                linear_transformer, time_transformer, freq_transformer = (
+                    transformer_block
+                )
+
+                x, ft_ps = pack([x], "b * d")
+                if self.use_torch_checkpoint:
+                    x = checkpoint(linear_transformer, x, use_reentrant=False)
+                else:
+                    x = linear_transformer(x)
+                (x,) = unpack(x, ft_ps, "b * d")
+            else:
+                time_transformer, freq_transformer = transformer_block
+
+            if self.skip_connection:
+                for j in range(i):
+                    x = x + store[j]
+
+            x = rearrange(x, "b t f d -> b f t d")
+            x, ps = pack([x], "* t d")
+
+            if self.use_torch_checkpoint:
+                x = checkpoint(time_transformer, x, use_reentrant=False)
+            else:
+                x = time_transformer(x)
+
+            (x,) = unpack(x, ps, "* t d")
+            x = rearrange(x, "b f t d -> b t f d")
+            x, ps = pack([x], "* f d")
+
+            if self.use_torch_checkpoint:
+                x = checkpoint(freq_transformer, x, use_reentrant=False)
+            else:
+                x = freq_transformer(x)
+
+            (x,) = unpack(x, ps, "* f d")
+
+            if self.skip_connection:
+                store[i] = x
+
+        x = self.final_norm(x)
+
+        num_stems = len(self.mask_estimators)
+
+        if self.use_torch_checkpoint:
+            mask = torch.stack(
+                [checkpoint(fn, x, use_reentrant=False) for fn in self.mask_estimators],
+                dim=1,
+            )
+        else:
+            mask = torch.stack([fn(x) for fn in self.mask_estimators], dim=1)
+        mask = rearrange(mask, "b n t (f c) -> b n f t c", c=2)
+
+        stft_repr = rearrange(stft_repr, "b f t c -> b 1 f t c")
+
+        stft_repr = torch.view_as_complex(stft_repr)
+        mask = torch.view_as_complex(mask)
+
+        stft_repr = stft_repr * mask
+
+        stft_repr = rearrange(
+            stft_repr, "b n (f s) t -> (b n s) f t", s=self.audio_channels
+        )
+
+        try:
+            recon_audio = torch.istft(
+                stft_repr,
+                **self.stft_kwargs,
+                window=stft_window,
+                return_complex=False,
+                length=raw_audio.shape[-1],
+            )
+        except:
+            recon_audio = torch.istft(
+                stft_repr.cpu() if x_is_mps else stft_repr,
+                **self.stft_kwargs,
+                window=stft_window.cpu() if x_is_mps else stft_window,
+                return_complex=False,
+                length=raw_audio.shape[-1],
+            ).to(device)
+
+        recon_audio = rearrange(
+            recon_audio, "(b n s) t -> b n s t", s=self.audio_channels, n=num_stems
+        )
+
+        if num_stems == 1:
+            recon_audio = rearrange(recon_audio, "b 1 s t -> b s t")
+
+        if not exists(target):
+            return recon_audio
+
+        if self.num_stems > 1:
+            assert target.ndim == 4 and target.shape[1] == self.num_stems
+
+        if target.ndim == 2:
+            target = rearrange(target, "... t -> ... 1 t")
+
+        target = target[..., : recon_audio.shape[-1]]
+
+        loss = F.l1_loss(recon_audio, target)
+
+        multi_stft_resolution_loss = 0.0
+
+        for window_size in self.multi_stft_resolutions_window_sizes:
+            res_stft_kwargs = dict(
+                n_fft=max(window_size, self.multi_stft_n_fft),
+                win_length=window_size,
+                return_complex=True,
+                window=self.multi_stft_window_fn(window_size, device=device),
+                **self.multi_stft_kwargs,
+            )
+
+            recon_Y = torch.stft(
+                rearrange(recon_audio, "... s t -> (... s) t"), **res_stft_kwargs
+            )
+            target_Y = torch.stft(
+                rearrange(target, "... s t -> (... s) t"), **res_stft_kwargs
+            )
+
+            multi_stft_resolution_loss = multi_stft_resolution_loss + F.l1_loss(
+                recon_Y, target_Y
+            )
+
+        weighted_multi_resolution_loss = (
+            multi_stft_resolution_loss * self.multi_stft_resolution_loss_weight
+        )
+
+        total_loss = loss + weighted_multi_resolution_loss
+
+        if not return_loss_breakdown:
+            return total_loss
+
+        return total_loss, (loss, multi_stft_resolution_loss)

mvsepless/models/bs_roformer/bs_roformer_hyperace.py

@@ -1,1122 +1,1122 @@
-from functools import partial
-
-import torch
-from torch import nn, einsum, Tensor
-from torch.nn import Module, ModuleList
-import torch.nn.functional as F
-
-from .attend import Attend
-
-try:
-    from .attend_sage import Attend as AttendSage
-except:
-    pass
-from torch.utils.checkpoint import checkpoint
-
-from beartype.typing import Tuple, Optional, List, Callable
-from beartype import beartype
-
-from rotary_embedding_torch import RotaryEmbedding
-
-from einops import rearrange, pack, unpack
-from einops.layers.torch import Rearrange
-import torchaudio
-
-
-def exists(val):
-    return val is not None
-
-
-def default(v, d):
-    return v if exists(v) else d
-
-
-def pack_one(t, pattern):
-    return pack([t], pattern)
-
-
-def unpack_one(t, ps, pattern):
-    return unpack(t, ps, pattern)[0]
-
-
-def l2norm(t):
-    return F.normalize(t, dim=-1, p=2)
-
-
-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
-
-
-class FeedForward(Module):
-    def __init__(self, dim, mult=4, dropout=0.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.0,
-        rotary_embed=None,
-        flash=True,
-        sage_attention=False,
-    ):
-        super().__init__()
-        self.heads = heads
-        self.scale = dim_head**-0.5
-        dim_inner = heads * dim_head
-
-        self.rotary_embed = rotary_embed
-
-        if sage_attention:
-            self.attend = AttendSage(flash=flash, dropout=dropout)
-        else:
-            self.attend = Attend(flash=flash, dropout=dropout)
-
-        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)
-
-        q, k, v = rearrange(
-            self.to_qkv(x), "b n (qkv h d) -> qkv b h n d", qkv=3, h=self.heads
-        )
-
-        if exists(self.rotary_embed):
-            q = self.rotary_embed.rotate_queries_or_keys(q)
-            k = self.rotary_embed.rotate_queries_or_keys(k)
-
-        out = self.attend(q, k, v)
-
-        gates = self.to_gates(x)
-        out = out * rearrange(gates, "b n h -> b h n 1").sigmoid()
-
-        out = rearrange(out, "b h n d -> b n (h d)")
-        return self.to_out(out)
-
-
-class LinearAttention(Module):
-
-    @beartype
-    def __init__(
-        self,
-        *,
-        dim,
-        dim_head=32,
-        heads=8,
-        scale=8,
-        flash=True,
-        dropout=0.0,
-        sage_attention=False,
-    ):
-        super().__init__()
-        dim_inner = dim_head * heads
-        self.norm = RMSNorm(dim)
-
-        self.to_qkv = nn.Sequential(
-            nn.Linear(dim, dim_inner * 3, bias=False),
-            Rearrange("b n (qkv h d) -> qkv b h d n", qkv=3, h=heads),
-        )
-
-        self.temperature = nn.Parameter(torch.ones(heads, 1, 1))
-
-        if sage_attention:
-            self.attend = AttendSage(scale=scale, dropout=dropout, flash=flash)
-        else:
-            self.attend = Attend(scale=scale, dropout=dropout, flash=flash)
-
-        self.to_out = nn.Sequential(
-            Rearrange("b h d n -> b n (h d)"), nn.Linear(dim_inner, dim, bias=False)
-        )
-
-    def forward(self, x):
-        x = self.norm(x)
-
-        q, k, v = self.to_qkv(x)
-
-        q, k = map(l2norm, (q, k))
-        q = q * self.temperature.exp()
-
-        out = self.attend(q, k, v)
-
-        return self.to_out(out)
-
-
-class Transformer(Module):
-    def __init__(
-        self,
-        *,
-        dim,
-        depth,
-        dim_head=64,
-        heads=8,
-        attn_dropout=0.0,
-        ff_dropout=0.0,
-        ff_mult=4,
-        norm_output=True,
-        rotary_embed=None,
-        flash_attn=True,
-        linear_attn=False,
-        sage_attention=False,
-    ):
-        super().__init__()
-        self.layers = ModuleList([])
-
-        for _ in range(depth):
-            if linear_attn:
-                attn = LinearAttention(
-                    dim=dim,
-                    dim_head=dim_head,
-                    heads=heads,
-                    dropout=attn_dropout,
-                    flash=flash_attn,
-                    sage_attention=sage_attention,
-                )
-            else:
-                attn = Attention(
-                    dim=dim,
-                    dim_head=dim_head,
-                    heads=heads,
-                    dropout=attn_dropout,
-                    rotary_embed=rotary_embed,
-                    flash=flash_attn,
-                    sage_attention=sage_attention,
-                )
-
-            self.layers.append(
-                ModuleList(
-                    [attn, 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)
-
-
-class BandSplit(Module):
-    @beartype
-    def __init__(self, dim, dim_inputs: Tuple[int, ...]):
-        super().__init__()
-        self.dim_inputs = dim_inputs
-        self.to_features = ModuleList([])
-
-        for dim_in in dim_inputs:
-            net = nn.Sequential(RMSNorm(dim_in), nn.Linear(dim_in, dim))
-
-            self.to_features.append(net)
-
-    def forward(self, x):
-
-        x = x.split(self.dim_inputs, dim=-1)
-
-        outs = []
-        for split_input, to_feature in zip(x, self.to_features):
-            split_output = to_feature(split_input)
-            outs.append(split_output)
-
-        x = torch.stack(outs, dim=-2)
-
-        return x
-
-
-class Conv(nn.Module):
-    def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True):
-        super().__init__()
-        self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False)
-        self.bn = nn.InstanceNorm2d(c2, affine=True, eps=1e-8)
-        self.act = nn.SiLU() if act else nn.Identity()
-
-    def forward(self, x):
-        return self.act(self.bn(self.conv(x)))
-
-
-def autopad(k, p=None):
-    if p is None:
-        p = k // 2 if isinstance(k, int) else [x // 2 for x in k]
-    return p
-
-
-class DSConv(nn.Module):
-    def __init__(self, c1, c2, k=3, s=1, p=None, act=True):
-        super().__init__()
-        self.dwconv = nn.Conv2d(c1, c1, k, s, autopad(k, p), groups=c1, bias=False)
-        self.pwconv = nn.Conv2d(c1, c2, 1, 1, 0, bias=False)
-        self.bn = nn.InstanceNorm2d(c2, affine=True, eps=1e-8)
-        self.act = nn.SiLU() if act else nn.Identity()
-
-    def forward(self, x):
-        return self.act(self.bn(self.pwconv(self.dwconv(x))))
-
-
-class DS_Bottleneck(nn.Module):
-    def __init__(self, c1, c2, k=3, shortcut=True):
-        super().__init__()
-        c_ = c1
-        self.dsconv1 = DSConv(c1, c_, k=3, s=1)
-        self.dsconv2 = DSConv(c_, c2, k=k, s=1)
-        self.shortcut = shortcut and c1 == c2
-
-    def forward(self, x):
-        return (
-            x + self.dsconv2(self.dsconv1(x))
-            if self.shortcut
-            else self.dsconv2(self.dsconv1(x))
-        )
-
-
-class DS_C3k(nn.Module):
-    def __init__(self, c1, c2, n=1, k=3, e=0.5):
-        super().__init__()
-        c_ = int(c2 * e)
-        self.cv1 = Conv(c1, c_, 1, 1)
-        self.cv2 = Conv(c1, c_, 1, 1)
-        self.cv3 = Conv(2 * c_, c2, 1, 1)
-        self.m = nn.Sequential(
-            *[DS_Bottleneck(c_, c_, k=k, shortcut=True) for _ in range(n)]
-        )
-
-    def forward(self, x):
-        return self.cv3(torch.cat((self.m(self.cv1(x)), self.cv2(x)), dim=1))
-
-
-class DS_C3k2(nn.Module):
-    def __init__(self, c1, c2, n=1, k=3, e=0.5):
-        super().__init__()
-        c_ = int(c2 * e)
-        self.cv1 = Conv(c1, c_, 1, 1)
-        self.m = DS_C3k(c_, c_, n=n, k=k, e=1.0)
-        self.cv2 = Conv(c_, c2, 1, 1)
-
-    def forward(self, x):
-        x_ = self.cv1(x)
-        x_ = self.m(x_)
-        return self.cv2(x_)
-
-
-class AdaptiveHyperedgeGeneration(nn.Module):
-    def __init__(self, in_channels, num_hyperedges, num_heads=8):
-        super().__init__()
-        self.num_hyperedges = num_hyperedges
-        self.num_heads = num_heads
-        self.head_dim = in_channels // num_heads
-
-        self.global_proto = nn.Parameter(torch.randn(num_hyperedges, in_channels))
-
-        self.context_mapper = nn.Linear(
-            2 * in_channels, num_hyperedges * in_channels, bias=False
-        )
-
-        self.query_proj = nn.Linear(in_channels, in_channels, bias=False)
-
-        self.scale = self.head_dim**-0.5
-
-    def forward(self, x):
-        B, N, C = x.shape
-
-        f_avg = F.adaptive_avg_pool1d(x.permute(0, 2, 1), 1).squeeze(-1)
-        f_max = F.adaptive_max_pool1d(x.permute(0, 2, 1), 1).squeeze(-1)
-        f_ctx = torch.cat((f_avg, f_max), dim=1)
-
-        delta_P = self.context_mapper(f_ctx).view(B, self.num_hyperedges, C)
-        P = self.global_proto.unsqueeze(0) + delta_P
-
-        z = self.query_proj(x)
-
-        z = z.view(B, N, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
-
-        P = P.view(B, self.num_hyperedges, self.num_heads, self.head_dim).permute(
-            0, 2, 3, 1
-        )
-
-        sim = (z @ P) * self.scale
-
-        s_bar = sim.mean(dim=1)
-
-        A = F.softmax(s_bar.permute(0, 2, 1), dim=-1)
-
-        return A
-
-
-class HypergraphConvolution(nn.Module):
-    def __init__(self, in_channels, out_channels):
-        super().__init__()
-        self.W_e = nn.Linear(in_channels, in_channels, bias=False)
-        self.W_v = nn.Linear(in_channels, out_channels, bias=False)
-        self.act = nn.SiLU()
-
-    def forward(self, x, A):
-        f_m = torch.bmm(A, x)
-        f_m = self.act(self.W_e(f_m))
-
-        x_out = torch.bmm(A.transpose(1, 2), f_m)
-        x_out = self.act(self.W_v(x_out))
-
-        return x + x_out
-
-
-class AdaptiveHypergraphComputation(nn.Module):
-    def __init__(self, in_channels, out_channels, num_hyperedges=8, num_heads=8):
-        super().__init__()
-        self.adaptive_hyperedge_gen = AdaptiveHyperedgeGeneration(
-            in_channels, num_hyperedges, num_heads
-        )
-        self.hypergraph_conv = HypergraphConvolution(in_channels, out_channels)
-
-    def forward(self, x):
-        B, C, H, W = x.shape
-        x_flat = x.flatten(2).permute(0, 2, 1)
-
-        A = self.adaptive_hyperedge_gen(x_flat)
-
-        x_out_flat = self.hypergraph_conv(x_flat, A)
-
-        x_out = x_out_flat.permute(0, 2, 1).view(B, -1, H, W)
-        return x_out
-
-
-class C3AH(nn.Module):
-    def __init__(self, c1, c2, num_hyperedges=8, num_heads=8, e=0.5):
-        super().__init__()
-        c_ = int(c1 * e)
-        self.cv1 = Conv(c1, c_, 1, 1)
-        self.cv2 = Conv(c1, c_, 1, 1)
-        self.ahc = AdaptiveHypergraphComputation(c_, c_, num_hyperedges, num_heads)
-        self.cv3 = Conv(2 * c_, c2, 1, 1)
-
-    def forward(self, x):
-        x_lateral = self.cv1(x)
-        x_ahc = self.ahc(self.cv2(x))
-        return self.cv3(torch.cat((x_ahc, x_lateral), dim=1))
-
-
-class HyperACE(nn.Module):
-    def __init__(
-        self,
-        in_channels: List[int],
-        out_channels: int,
-        num_hyperedges=8,
-        num_heads=8,
-        k=2,
-        l=1,
-        c_h=0.5,
-        c_l=0.25,
-    ):
-        super().__init__()
-
-        c2, c3, c4, c5 = in_channels
-        c_mid = c4
-
-        self.fuse_conv = Conv(c2 + c3 + c4 + c5, c_mid, 1, 1)
-
-        self.c_h = int(c_mid * c_h)
-        self.c_l = int(c_mid * c_l)
-        self.c_s = c_mid - self.c_h - self.c_l
-        assert self.c_s > 0, "Channel split error"
-
-        self.high_order_branch = nn.ModuleList(
-            [
-                C3AH(self.c_h, self.c_h, num_hyperedges, num_heads, e=1.0)
-                for _ in range(k)
-            ]
-        )
-        self.high_order_fuse = Conv(self.c_h * k, self.c_h, 1, 1)
-
-        self.low_order_branch = nn.Sequential(
-            *[DS_C3k(self.c_l, self.c_l, n=1, k=3, e=1.0) for _ in range(l)]
-        )
-
-        self.final_fuse = Conv(self.c_h + self.c_l + self.c_s, out_channels, 1, 1)
-
-    def forward(self, x: List[torch.Tensor]) -> torch.Tensor:
-        B2, B3, B4, B5 = x
-
-        B, _, H4, W4 = B4.shape
-
-        B2_resized = F.interpolate(
-            B2, size=(H4, W4), mode="bilinear", align_corners=False
-        )
-        B3_resized = F.interpolate(
-            B3, size=(H4, W4), mode="bilinear", align_corners=False
-        )
-        B5_resized = F.interpolate(
-            B5, size=(H4, W4), mode="bilinear", align_corners=False
-        )
-
-        x_b = self.fuse_conv(torch.cat((B2_resized, B3_resized, B4, B5_resized), dim=1))
-
-        x_h, x_l, x_s = torch.split(x_b, [self.c_h, self.c_l, self.c_s], dim=1)
-
-        x_h_outs = [m(x_h) for m in self.high_order_branch]
-        x_h_fused = self.high_order_fuse(torch.cat(x_h_outs, dim=1))
-
-        x_l_out = self.low_order_branch(x_l)
-
-        y = self.final_fuse(torch.cat((x_h_fused, x_l_out, x_s), dim=1))
-
-        return y
-
-
-class GatedFusion(nn.Module):
-    def __init__(self, in_channels):
-        super().__init__()
-        self.gamma = nn.Parameter(torch.zeros(1, in_channels, 1, 1))
-
-    def forward(self, f_in, h):
-        if f_in.shape[1] != h.shape[1]:
-            raise ValueError(f"Channel mismatch: f_in={f_in.shape}, h={h.shape}")
-        return f_in + self.gamma * h
-
-
-class Backbone(nn.Module):
-    def __init__(self, in_channels=256, base_channels=64, base_depth=3):
-        super().__init__()
-        c = base_channels
-        c2 = base_channels
-        c3 = 256
-        c4 = 384
-        c5 = 512
-        c6 = 768
-
-        self.stem = DSConv(in_channels, c2, k=3, s=(2, 1), p=1)
-
-        self.p2 = nn.Sequential(
-            DSConv(c2, c3, k=3, s=(2, 1), p=1), DS_C3k2(c3, c3, n=base_depth)
-        )
-
-        self.p3 = nn.Sequential(
-            DSConv(c3, c4, k=3, s=(2, 1), p=1), DS_C3k2(c4, c4, n=base_depth * 2)
-        )
-
-        self.p4 = nn.Sequential(
-            DSConv(c4, c5, k=3, s=(2, 1), p=1), DS_C3k2(c5, c5, n=base_depth * 2)
-        )
-
-        self.p5 = nn.Sequential(
-            DSConv(c5, c6, k=3, s=(2, 1), p=1), DS_C3k2(c6, c6, n=base_depth)
-        )
-
-        self.out_channels = [c3, c4, c5, c6]
-
-    def forward(self, x):
-        x = self.stem(x)
-        x2 = self.p2(x)
-        x3 = self.p3(x2)
-        x4 = self.p4(x3)
-        x5 = self.p5(x4)
-        return [x2, x3, x4, x5]
-
-
-class Decoder(nn.Module):
-    def __init__(
-        self,
-        encoder_channels: List[int],
-        hyperace_out_c: int,
-        decoder_channels: List[int],
-    ):
-        super().__init__()
-        c_p2, c_p3, c_p4, c_p5 = encoder_channels
-        c_d2, c_d3, c_d4, c_d5 = decoder_channels
-
-        self.h_to_d5 = Conv(hyperace_out_c, c_d5, 1, 1)
-        self.h_to_d4 = Conv(hyperace_out_c, c_d4, 1, 1)
-        self.h_to_d3 = Conv(hyperace_out_c, c_d3, 1, 1)
-        self.h_to_d2 = Conv(hyperace_out_c, c_d2, 1, 1)
-
-        self.fusion_d5 = GatedFusion(c_d5)
-        self.fusion_d4 = GatedFusion(c_d4)
-        self.fusion_d3 = GatedFusion(c_d3)
-        self.fusion_d2 = GatedFusion(c_d2)
-
-        self.skip_p5 = Conv(c_p5, c_d5, 1, 1)
-        self.skip_p4 = Conv(c_p4, c_d4, 1, 1)
-        self.skip_p3 = Conv(c_p3, c_d3, 1, 1)
-        self.skip_p2 = Conv(c_p2, c_d2, 1, 1)
-
-        self.up_d5 = DS_C3k2(c_d5, c_d4, n=1)
-        self.up_d4 = DS_C3k2(c_d4, c_d3, n=1)
-        self.up_d3 = DS_C3k2(c_d3, c_d2, n=1)
-
-        self.final_d2 = DS_C3k2(c_d2, c_d2, n=1)
-
-    def forward(self, enc_feats: List[torch.Tensor], h_ace: torch.Tensor):
-        p2, p3, p4, p5 = enc_feats
-
-        d5 = self.skip_p5(p5)
-        h_d5 = self.h_to_d5(F.interpolate(h_ace, size=d5.shape[2:], mode="bilinear"))
-        d5 = self.fusion_d5(d5, h_d5)
-
-        d5_up = F.interpolate(d5, size=p4.shape[2:], mode="bilinear")
-        d4_skip = self.skip_p4(p4)
-        d4 = self.up_d5(d5_up) + d4_skip
-
-        h_d4 = self.h_to_d4(F.interpolate(h_ace, size=d4.shape[2:], mode="bilinear"))
-        d4 = self.fusion_d4(d4, h_d4)
-
-        d4_up = F.interpolate(d4, size=p3.shape[2:], mode="bilinear")
-        d3_skip = self.skip_p3(p3)
-        d3 = self.up_d4(d4_up) + d3_skip
-
-        h_d3 = self.h_to_d3(F.interpolate(h_ace, size=d3.shape[2:], mode="bilinear"))
-        d3 = self.fusion_d3(d3, h_d3)
-
-        d3_up = F.interpolate(d3, size=p2.shape[2:], mode="bilinear")
-        d2_skip = self.skip_p2(p2)
-        d2 = self.up_d3(d3_up) + d2_skip
-
-        h_d2 = self.h_to_d2(F.interpolate(h_ace, size=d2.shape[2:], mode="bilinear"))
-        d2 = self.fusion_d2(d2, h_d2)
-
-        d2_final = self.final_d2(d2)
-
-        return d2_final
-
-
-class FreqPixelShuffle(nn.Module):
-    def __init__(self, in_channels, out_channels, scale=2):
-        super().__init__()
-        self.scale = scale
-        self.conv = DSConv(in_channels, out_channels * scale, k=3, s=1, p=1)
-        self.act = nn.SiLU()
-
-    def forward(self, x):
-        x = self.conv(x)
-        B, C_r, H, W = x.shape
-        out_c = C_r // self.scale
-
-        x = x.view(B, out_c, self.scale, H, W)
-
-        x = x.permute(0, 1, 3, 4, 2).contiguous()
-        x = x.view(B, out_c, H, W * self.scale)
-
-        return x
-
-
-class ProgressiveUpsampleHead(nn.Module):
-    def __init__(self, in_channels, out_channels, target_bins=1025):
-        super().__init__()
-        self.target_bins = target_bins
-
-        c = in_channels
-
-        self.block1 = FreqPixelShuffle(c, c, scale=2)
-        self.block2 = FreqPixelShuffle(c, c // 2, scale=2)
-        self.block3 = FreqPixelShuffle(c // 2, c // 2, scale=2)
-        self.block4 = FreqPixelShuffle(c // 2, c // 4, scale=2)
-
-        self.final_conv = nn.Conv2d(c // 4, out_channels, kernel_size=1, bias=False)
-
-    def forward(self, x):
-
-        x = self.block1(x)
-        x = self.block2(x)
-        x = self.block3(x)
-        x = self.block4(x)
-
-        if x.shape[-1] != self.target_bins:
-            x = F.interpolate(
-                x,
-                size=(x.shape[2], self.target_bins),
-                mode="bilinear",
-                align_corners=False,
-            )
-
-        x = self.final_conv(x)
-        return x
-
-
-class SegmModel(nn.Module):
-    def __init__(
-        self,
-        in_bands=62,
-        in_dim=256,
-        out_bins=1025,
-        out_channels=4,
-        base_channels=64,
-        base_depth=2,
-        num_hyperedges=16,
-        num_heads=8,
-    ):
-        super().__init__()
-
-        self.backbone = Backbone(
-            in_channels=in_dim, base_channels=base_channels, base_depth=base_depth
-        )
-        enc_channels = self.backbone.out_channels
-        c2, c3, c4, c5 = enc_channels
-
-        hyperace_in_channels = enc_channels
-        hyperace_out_channels = c4
-        self.hyperace = HyperACE(
-            hyperace_in_channels,
-            hyperace_out_channels,
-            num_hyperedges,
-            num_heads,
-            k=3,
-            l=2,
-        )
-
-        decoder_channels = [c2, c3, c4, c5]
-        self.decoder = Decoder(enc_channels, hyperace_out_channels, decoder_channels)
-
-        self.upsample_head = ProgressiveUpsampleHead(
-            in_channels=decoder_channels[0],
-            out_channels=out_channels,
-            target_bins=out_bins,
-        )
-
-    def forward(self, x):
-        H, W = x.shape[2:]
-
-        enc_feats = self.backbone(x)
-
-        h_ace_feats = self.hyperace(enc_feats)
-
-        dec_feat = self.decoder(enc_feats, h_ace_feats)
-
-        feat_time_restored = F.interpolate(
-            dec_feat, size=(H, dec_feat.shape[-1]), mode="bilinear", align_corners=False
-        )
-
-        out = self.upsample_head(feat_time_restored)
-
-        return out
-
-
-def MLP(dim_in, dim_out, dim_hidden=None, depth=1, activation=nn.Tanh):
-    dim_hidden = default(dim_hidden, dim_in)
-
-    net = []
-    dims = (dim_in, *((dim_hidden,) * (depth - 1)), dim_out)
-
-    for ind, (layer_dim_in, layer_dim_out) in enumerate(zip(dims[:-1], dims[1:])):
-        is_last = ind == (len(dims) - 2)
-
-        net.append(nn.Linear(layer_dim_in, layer_dim_out))
-
-        if is_last:
-            continue
-
-        net.append(activation())
-
-    return nn.Sequential(*net)
-
-
-class MaskEstimator(Module):
-    @beartype
-    def __init__(self, dim, dim_inputs: Tuple[int, ...], depth, mlp_expansion_factor=4):
-        super().__init__()
-        self.dim_inputs = dim_inputs
-        self.to_freqs = ModuleList([])
-        dim_hidden = dim * mlp_expansion_factor
-
-        for dim_in in dim_inputs:
-            net = []
-
-            mlp = nn.Sequential(
-                MLP(dim, dim_in * 2, dim_hidden=dim_hidden, depth=depth), nn.GLU(dim=-1)
-            )
-
-            self.to_freqs.append(mlp)
-
-        self.segm = SegmModel(
-            in_bands=len(dim_inputs), in_dim=dim, out_bins=sum(dim_inputs) // 4
-        )
-
-    def forward(self, x):
-        y = rearrange(x, "b t f c -> b c t f")
-        y = self.segm(y)
-        y = rearrange(y, "b c t f -> b t (f c)")
-
-        x = x.unbind(dim=-2)
-
-        outs = []
-
-        for band_features, mlp in zip(x, self.to_freqs):
-            freq_out = mlp(band_features)
-            outs.append(freq_out)
-
-        return torch.cat(outs, dim=-1) + y
-
-
-DEFAULT_FREQS_PER_BANDS = (
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    2,
-    4,
-    4,
-    4,
-    4,
-    4,
-    4,
-    4,
-    4,
-    4,
-    4,
-    4,
-    4,
-    12,
-    12,
-    12,
-    12,
-    12,
-    12,
-    12,
-    12,
-    24,
-    24,
-    24,
-    24,
-    24,
-    24,
-    24,
-    24,
-    48,
-    48,
-    48,
-    48,
-    48,
-    48,
-    48,
-    48,
-    128,
-    129,
-)
-
-
-class BSRoformerHyperACE(Module):
-
-    @beartype
-    def __init__(
-        self,
-        dim,
-        *,
-        depth,
-        stereo=False,
-        num_stems=1,
-        time_transformer_depth=2,
-        freq_transformer_depth=2,
-        linear_transformer_depth=0,
-        freqs_per_bands: Tuple[int, ...] = DEFAULT_FREQS_PER_BANDS,
-        dim_head=64,
-        heads=8,
-        attn_dropout=0.0,
-        ff_dropout=0.0,
-        flash_attn=True,
-        dim_freqs_in=1025,
-        stft_n_fft=2048,
-        stft_hop_length=512,
-        stft_win_length=2048,
-        stft_normalized=False,
-        stft_window_fn: Optional[Callable] = None,
-        mask_estimator_depth=2,
-        multi_stft_resolution_loss_weight=1.0,
-        multi_stft_resolutions_window_sizes: Tuple[int, ...] = (
-            4096,
-            2048,
-            1024,
-            512,
-            256,
-        ),
-        multi_stft_hop_size=147,
-        multi_stft_normalized=False,
-        multi_stft_window_fn: Callable = torch.hann_window,
-        mlp_expansion_factor=4,
-        use_torch_checkpoint=False,
-        skip_connection=False,
-        sage_attention=False,
-    ):
-        super().__init__()
-
-        self.stereo = stereo
-        self.audio_channels = 2 if stereo else 1
-        self.num_stems = num_stems
-        self.use_torch_checkpoint = use_torch_checkpoint
-        self.skip_connection = skip_connection
-
-        self.layers = ModuleList([])
-
-        if sage_attention:
-            print("Use Sage Attention")
-
-        transformer_kwargs = dict(
-            dim=dim,
-            heads=heads,
-            dim_head=dim_head,
-            attn_dropout=attn_dropout,
-            ff_dropout=ff_dropout,
-            flash_attn=flash_attn,
-            norm_output=False,
-            sage_attention=sage_attention,
-        )
-
-        time_rotary_embed = RotaryEmbedding(dim=dim_head)
-        freq_rotary_embed = RotaryEmbedding(dim=dim_head)
-
-        for _ in range(depth):
-            tran_modules = []
-            tran_modules.append(
-                Transformer(
-                    depth=time_transformer_depth,
-                    rotary_embed=time_rotary_embed,
-                    **transformer_kwargs,
-                )
-            )
-            tran_modules.append(
-                Transformer(
-                    depth=freq_transformer_depth,
-                    rotary_embed=freq_rotary_embed,
-                    **transformer_kwargs,
-                )
-            )
-            self.layers.append(nn.ModuleList(tran_modules))
-
-        self.final_norm = RMSNorm(dim)
-
-        self.stft_kwargs = dict(
-            n_fft=stft_n_fft,
-            hop_length=stft_hop_length,
-            win_length=stft_win_length,
-            normalized=stft_normalized,
-        )
-
-        self.stft_window_fn = partial(
-            default(stft_window_fn, torch.hann_window), stft_win_length
-        )
-
-        freqs = torch.stft(
-            torch.randn(1, 4096),
-            **self.stft_kwargs,
-            window=torch.ones(stft_win_length),
-            return_complex=True,
-        ).shape[1]
-
-        assert len(freqs_per_bands) > 1
-        assert (
-            sum(freqs_per_bands) == freqs
-        ), f"the number of freqs in the bands must equal {freqs} based on the STFT settings, but got {sum(freqs_per_bands)}"
-
-        freqs_per_bands_with_complex = tuple(
-            2 * f * self.audio_channels for f in freqs_per_bands
-        )
-
-        self.band_split = BandSplit(dim=dim, dim_inputs=freqs_per_bands_with_complex)
-
-        self.mask_estimators = nn.ModuleList([])
-
-        for _ in range(num_stems):
-            mask_estimator = MaskEstimator(
-                dim=dim,
-                dim_inputs=freqs_per_bands_with_complex,
-                depth=mask_estimator_depth,
-                mlp_expansion_factor=mlp_expansion_factor,
-            )
-
-            self.mask_estimators.append(mask_estimator)
-
-        self.multi_stft_resolution_loss_weight = multi_stft_resolution_loss_weight
-        self.multi_stft_resolutions_window_sizes = multi_stft_resolutions_window_sizes
-        self.multi_stft_n_fft = stft_n_fft
-        self.multi_stft_window_fn = multi_stft_window_fn
-
-        self.multi_stft_kwargs = dict(
-            hop_length=multi_stft_hop_size, normalized=multi_stft_normalized
-        )
-
-    def forward(self, raw_audio, target=None, return_loss_breakdown=False):
-
-        device = raw_audio.device
-
-        x_is_mps = True if device.type == "mps" else False
-
-        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 needs to be set to True if passing in audio signal that is stereo (channel dimension of 2). also need to be False if mono (channel dimension of 1)"
-
-        raw_audio, batch_audio_channel_packed_shape = pack_one(raw_audio, "* t")
-
-        stft_window = self.stft_window_fn(device=device)
-
-        try:
-            stft_repr = torch.stft(
-                raw_audio, **self.stft_kwargs, window=stft_window, return_complex=True
-            )
-        except:
-            stft_repr = torch.stft(
-                raw_audio.cpu() if x_is_mps else raw_audio,
-                **self.stft_kwargs,
-                window=stft_window.cpu() if x_is_mps else stft_window,
-                return_complex=True,
-            ).to(device)
-        stft_repr = torch.view_as_real(stft_repr)
-
-        stft_repr = unpack_one(stft_repr, batch_audio_channel_packed_shape, "* f t c")
-
-        stft_repr = rearrange(stft_repr, "b s f t c -> b (f s) t c")
-
-        x = rearrange(stft_repr, "b f t c -> b t (f c)")
-
-        x = self.band_split(x)
-
-        for i, transformer_block in enumerate(self.layers):
-
-            time_transformer, freq_transformer = transformer_block
-
-            x = rearrange(x, "b t f d -> b f t d")
-            x, ps = pack([x], "* t d")
-
-            x = time_transformer(x)
-
-            (x,) = unpack(x, ps, "* t d")
-            x = rearrange(x, "b f t d -> b t f d")
-            x, ps = pack([x], "* f d")
-
-            x = freq_transformer(x)
-
-            (x,) = unpack(x, ps, "* f d")
-
-        x = self.final_norm(x)
-
-        num_stems = len(self.mask_estimators)
-
-        mask = torch.stack([fn(x) for fn in self.mask_estimators], dim=1)
-        mask = rearrange(mask, "b n t (f c) -> b n f t c", c=2)
-
-        stft_repr = rearrange(stft_repr, "b f t c -> b 1 f t c")
-
-        stft_repr = torch.view_as_complex(stft_repr)
-        mask = torch.view_as_complex(mask)
-
-        stft_repr = stft_repr * mask
-
-        stft_repr = rearrange(
-            stft_repr, "b n (f s) t -> (b n s) f t", s=self.audio_channels
-        )
-
-        try:
-            recon_audio = torch.istft(
-                stft_repr,
-                **self.stft_kwargs,
-                window=stft_window,
-                return_complex=False,
-                length=raw_audio.shape[-1],
-            )
-        except:
-            recon_audio = torch.istft(
-                stft_repr.cpu() if x_is_mps else stft_repr,
-                **self.stft_kwargs,
-                window=stft_window.cpu() if x_is_mps else stft_window,
-                return_complex=False,
-                length=raw_audio.shape[-1],
-            ).to(device)
-
-        recon_audio = rearrange(
-            recon_audio, "(b n s) t -> b n s t", s=self.audio_channels, n=num_stems
-        )
-
-        if num_stems == 1:
-            recon_audio = rearrange(recon_audio, "b 1 s t -> b s t")
-
-        if not exists(target):
-            return recon_audio
-
-        if self.num_stems > 1:
-            assert target.ndim == 4 and target.shape[1] == self.num_stems
-
-        if target.ndim == 2:
-            target = rearrange(target, "... t -> ... 1 t")
-
-        target = target[..., : recon_audio.shape[-1]]
-
-        loss = F.l1_loss(recon_audio, target)
-
-        multi_stft_resolution_loss = 0.0
-
-        for window_size in self.multi_stft_resolutions_window_sizes:
-            res_stft_kwargs = dict(
-                n_fft=max(window_size, self.multi_stft_n_fft),
-                win_length=window_size,
-                return_complex=True,
-                window=self.multi_stft_window_fn(window_size, device=device),
-                **self.multi_stft_kwargs,
-            )
-
-            recon_Y = torch.stft(
-                rearrange(recon_audio, "... s t -> (... s) t"), **res_stft_kwargs
-            )
-            target_Y = torch.stft(
-                rearrange(target, "... s t -> (... s) t"), **res_stft_kwargs
-            )
-
-            multi_stft_resolution_loss = multi_stft_resolution_loss + F.l1_loss(
-                recon_Y, target_Y
-            )
-
-        weighted_multi_resolution_loss = (
-            multi_stft_resolution_loss * self.multi_stft_resolution_loss_weight
-        )
-
-        total_loss = loss + weighted_multi_resolution_loss
-
-        if not return_loss_breakdown:
-            return total_loss
-
-        return total_loss, (loss, multi_stft_resolution_loss)
+from functools import partial
+
+import torch
+from torch import nn, einsum, Tensor
+from torch.nn import Module, ModuleList
+import torch.nn.functional as F
+
+from .attend import Attend
+
+try:
+    from .attend_sage import Attend as AttendSage
+except:
+    pass
+from torch.utils.checkpoint import checkpoint
+
+from beartype.typing import Tuple, Optional, List, Callable
+from beartype import beartype
+
+from rotary_embedding_torch import RotaryEmbedding
+
+from einops import rearrange, pack, unpack
+from einops.layers.torch import Rearrange
+import torchaudio
+
+
+def exists(val):
+    return val is not None
+
+
+def default(v, d):
+    return v if exists(v) else d
+
+
+def pack_one(t, pattern):
+    return pack([t], pattern)
+
+
+def unpack_one(t, ps, pattern):
+    return unpack(t, ps, pattern)[0]
+
+
+def l2norm(t):
+    return F.normalize(t, dim=-1, p=2)
+
+
+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
+
+
+class FeedForward(Module):
+    def __init__(self, dim, mult=4, dropout=0.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.0,
+        rotary_embed=None,
+        flash=True,
+        sage_attention=False,
+    ):
+        super().__init__()
+        self.heads = heads
+        self.scale = dim_head**-0.5
+        dim_inner = heads * dim_head
+
+        self.rotary_embed = rotary_embed
+
+        if sage_attention:
+            self.attend = AttendSage(flash=flash, dropout=dropout)
+        else:
+            self.attend = Attend(flash=flash, dropout=dropout)
+
+        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)
+
+        q, k, v = rearrange(
+            self.to_qkv(x), "b n (qkv h d) -> qkv b h n d", qkv=3, h=self.heads
+        )
+
+        if exists(self.rotary_embed):
+            q = self.rotary_embed.rotate_queries_or_keys(q)
+            k = self.rotary_embed.rotate_queries_or_keys(k)
+
+        out = self.attend(q, k, v)
+
+        gates = self.to_gates(x)
+        out = out * rearrange(gates, "b n h -> b h n 1").sigmoid()
+
+        out = rearrange(out, "b h n d -> b n (h d)")
+        return self.to_out(out)
+
+
+class LinearAttention(Module):
+
+    @beartype
+    def __init__(
+        self,
+        *,
+        dim,
+        dim_head=32,
+        heads=8,
+        scale=8,
+        flash=True,
+        dropout=0.0,
+        sage_attention=False,
+    ):
+        super().__init__()
+        dim_inner = dim_head * heads
+        self.norm = RMSNorm(dim)
+
+        self.to_qkv = nn.Sequential(
+            nn.Linear(dim, dim_inner * 3, bias=False),
+            Rearrange("b n (qkv h d) -> qkv b h d n", qkv=3, h=heads),
+        )
+
+        self.temperature = nn.Parameter(torch.ones(heads, 1, 1))
+
+        if sage_attention:
+            self.attend = AttendSage(scale=scale, dropout=dropout, flash=flash)
+        else:
+            self.attend = Attend(scale=scale, dropout=dropout, flash=flash)
+
+        self.to_out = nn.Sequential(
+            Rearrange("b h d n -> b n (h d)"), nn.Linear(dim_inner, dim, bias=False)
+        )
+
+    def forward(self, x):
+        x = self.norm(x)
+
+        q, k, v = self.to_qkv(x)
+
+        q, k = map(l2norm, (q, k))
+        q = q * self.temperature.exp()
+
+        out = self.attend(q, k, v)
+
+        return self.to_out(out)
+
+
+class Transformer(Module):
+    def __init__(
+        self,
+        *,
+        dim,
+        depth,
+        dim_head=64,
+        heads=8,
+        attn_dropout=0.0,
+        ff_dropout=0.0,
+        ff_mult=4,
+        norm_output=True,
+        rotary_embed=None,
+        flash_attn=True,
+        linear_attn=False,
+        sage_attention=False,
+    ):
+        super().__init__()
+        self.layers = ModuleList([])
+
+        for _ in range(depth):
+            if linear_attn:
+                attn = LinearAttention(
+                    dim=dim,
+                    dim_head=dim_head,
+                    heads=heads,
+                    dropout=attn_dropout,
+                    flash=flash_attn,
+                    sage_attention=sage_attention,
+                )
+            else:
+                attn = Attention(
+                    dim=dim,
+                    dim_head=dim_head,
+                    heads=heads,
+                    dropout=attn_dropout,
+                    rotary_embed=rotary_embed,
+                    flash=flash_attn,
+                    sage_attention=sage_attention,
+                )
+
+            self.layers.append(
+                ModuleList(
+                    [attn, 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)
+
+
+class BandSplit(Module):
+    @beartype
+    def __init__(self, dim, dim_inputs: Tuple[int, ...]):
+        super().__init__()
+        self.dim_inputs = dim_inputs
+        self.to_features = ModuleList([])
+
+        for dim_in in dim_inputs:
+            net = nn.Sequential(RMSNorm(dim_in), nn.Linear(dim_in, dim))
+
+            self.to_features.append(net)
+
+    def forward(self, x):
+
+        x = x.split(self.dim_inputs, dim=-1)
+
+        outs = []
+        for split_input, to_feature in zip(x, self.to_features):
+            split_output = to_feature(split_input)
+            outs.append(split_output)
+
+        x = torch.stack(outs, dim=-2)
+
+        return x
+
+
+class Conv(nn.Module):
+    def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True):
+        super().__init__()
+        self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False)
+        self.bn = nn.InstanceNorm2d(c2, affine=True, eps=1e-8)
+        self.act = nn.SiLU() if act else nn.Identity()
+
+    def forward(self, x):
+        return self.act(self.bn(self.conv(x)))
+
+
+def autopad(k, p=None):
+    if p is None:
+        p = k // 2 if isinstance(k, int) else [x // 2 for x in k]
+    return p
+
+
+class DSConv(nn.Module):
+    def __init__(self, c1, c2, k=3, s=1, p=None, act=True):
+        super().__init__()
+        self.dwconv = nn.Conv2d(c1, c1, k, s, autopad(k, p), groups=c1, bias=False)
+        self.pwconv = nn.Conv2d(c1, c2, 1, 1, 0, bias=False)
+        self.bn = nn.InstanceNorm2d(c2, affine=True, eps=1e-8)
+        self.act = nn.SiLU() if act else nn.Identity()
+
+    def forward(self, x):
+        return self.act(self.bn(self.pwconv(self.dwconv(x))))
+
+
+class DS_Bottleneck(nn.Module):
+    def __init__(self, c1, c2, k=3, shortcut=True):
+        super().__init__()
+        c_ = c1
+        self.dsconv1 = DSConv(c1, c_, k=3, s=1)
+        self.dsconv2 = DSConv(c_, c2, k=k, s=1)
+        self.shortcut = shortcut and c1 == c2
+
+    def forward(self, x):
+        return (
+            x + self.dsconv2(self.dsconv1(x))
+            if self.shortcut
+            else self.dsconv2(self.dsconv1(x))
+        )
+
+
+class DS_C3k(nn.Module):
+    def __init__(self, c1, c2, n=1, k=3, e=0.5):
+        super().__init__()
+        c_ = int(c2 * e)
+        self.cv1 = Conv(c1, c_, 1, 1)
+        self.cv2 = Conv(c1, c_, 1, 1)
+        self.cv3 = Conv(2 * c_, c2, 1, 1)
+        self.m = nn.Sequential(
+            *[DS_Bottleneck(c_, c_, k=k, shortcut=True) for _ in range(n)]
+        )
+
+    def forward(self, x):
+        return self.cv3(torch.cat((self.m(self.cv1(x)), self.cv2(x)), dim=1))
+
+
+class DS_C3k2(nn.Module):
+    def __init__(self, c1, c2, n=1, k=3, e=0.5):
+        super().__init__()
+        c_ = int(c2 * e)
+        self.cv1 = Conv(c1, c_, 1, 1)
+        self.m = DS_C3k(c_, c_, n=n, k=k, e=1.0)
+        self.cv2 = Conv(c_, c2, 1, 1)
+
+    def forward(self, x):
+        x_ = self.cv1(x)
+        x_ = self.m(x_)
+        return self.cv2(x_)
+
+
+class AdaptiveHyperedgeGeneration(nn.Module):
+    def __init__(self, in_channels, num_hyperedges, num_heads=8):
+        super().__init__()
+        self.num_hyperedges = num_hyperedges
+        self.num_heads = num_heads
+        self.head_dim = in_channels // num_heads
+
+        self.global_proto = nn.Parameter(torch.randn(num_hyperedges, in_channels))
+
+        self.context_mapper = nn.Linear(
+            2 * in_channels, num_hyperedges * in_channels, bias=False
+        )
+
+        self.query_proj = nn.Linear(in_channels, in_channels, bias=False)
+
+        self.scale = self.head_dim**-0.5
+
+    def forward(self, x):
+        B, N, C = x.shape
+
+        f_avg = F.adaptive_avg_pool1d(x.permute(0, 2, 1), 1).squeeze(-1)
+        f_max = F.adaptive_max_pool1d(x.permute(0, 2, 1), 1).squeeze(-1)
+        f_ctx = torch.cat((f_avg, f_max), dim=1)
+
+        delta_P = self.context_mapper(f_ctx).view(B, self.num_hyperedges, C)
+        P = self.global_proto.unsqueeze(0) + delta_P
+
+        z = self.query_proj(x)
+
+        z = z.view(B, N, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
+
+        P = P.view(B, self.num_hyperedges, self.num_heads, self.head_dim).permute(
+            0, 2, 3, 1
+        )
+
+        sim = (z @ P) * self.scale
+
+        s_bar = sim.mean(dim=1)
+
+        A = F.softmax(s_bar.permute(0, 2, 1), dim=-1)
+
+        return A
+
+
+class HypergraphConvolution(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.W_e = nn.Linear(in_channels, in_channels, bias=False)
+        self.W_v = nn.Linear(in_channels, out_channels, bias=False)
+        self.act = nn.SiLU()
+
+    def forward(self, x, A):
+        f_m = torch.bmm(A, x)
+        f_m = self.act(self.W_e(f_m))
+
+        x_out = torch.bmm(A.transpose(1, 2), f_m)
+        x_out = self.act(self.W_v(x_out))
+
+        return x + x_out
+
+
+class AdaptiveHypergraphComputation(nn.Module):
+    def __init__(self, in_channels, out_channels, num_hyperedges=8, num_heads=8):
+        super().__init__()
+        self.adaptive_hyperedge_gen = AdaptiveHyperedgeGeneration(
+            in_channels, num_hyperedges, num_heads
+        )
+        self.hypergraph_conv = HypergraphConvolution(in_channels, out_channels)
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        x_flat = x.flatten(2).permute(0, 2, 1)
+
+        A = self.adaptive_hyperedge_gen(x_flat)
+
+        x_out_flat = self.hypergraph_conv(x_flat, A)
+
+        x_out = x_out_flat.permute(0, 2, 1).view(B, -1, H, W)
+        return x_out
+
+
+class C3AH(nn.Module):
+    def __init__(self, c1, c2, num_hyperedges=8, num_heads=8, e=0.5):
+        super().__init__()
+        c_ = int(c1 * e)
+        self.cv1 = Conv(c1, c_, 1, 1)
+        self.cv2 = Conv(c1, c_, 1, 1)
+        self.ahc = AdaptiveHypergraphComputation(c_, c_, num_hyperedges, num_heads)
+        self.cv3 = Conv(2 * c_, c2, 1, 1)
+
+    def forward(self, x):
+        x_lateral = self.cv1(x)
+        x_ahc = self.ahc(self.cv2(x))
+        return self.cv3(torch.cat((x_ahc, x_lateral), dim=1))
+
+
+class HyperACE(nn.Module):
+    def __init__(
+        self,
+        in_channels: List[int],
+        out_channels: int,
+        num_hyperedges=8,
+        num_heads=8,
+        k=2,
+        l=1,
+        c_h=0.5,
+        c_l=0.25,
+    ):
+        super().__init__()
+
+        c2, c3, c4, c5 = in_channels
+        c_mid = c4
+
+        self.fuse_conv = Conv(c2 + c3 + c4 + c5, c_mid, 1, 1)
+
+        self.c_h = int(c_mid * c_h)
+        self.c_l = int(c_mid * c_l)
+        self.c_s = c_mid - self.c_h - self.c_l
+        assert self.c_s > 0, "Channel split error"
+
+        self.high_order_branch = nn.ModuleList(
+            [
+                C3AH(self.c_h, self.c_h, num_hyperedges, num_heads, e=1.0)
+                for _ in range(k)
+            ]
+        )
+        self.high_order_fuse = Conv(self.c_h * k, self.c_h, 1, 1)
+
+        self.low_order_branch = nn.Sequential(
+            *[DS_C3k(self.c_l, self.c_l, n=1, k=3, e=1.0) for _ in range(l)]
+        )
+
+        self.final_fuse = Conv(self.c_h + self.c_l + self.c_s, out_channels, 1, 1)
+
+    def forward(self, x: List[torch.Tensor]) -> torch.Tensor:
+        B2, B3, B4, B5 = x
+
+        B, _, H4, W4 = B4.shape
+
+        B2_resized = F.interpolate(
+            B2, size=(H4, W4), mode="bilinear", align_corners=False
+        )
+        B3_resized = F.interpolate(
+            B3, size=(H4, W4), mode="bilinear", align_corners=False
+        )
+        B5_resized = F.interpolate(
+            B5, size=(H4, W4), mode="bilinear", align_corners=False
+        )
+
+        x_b = self.fuse_conv(torch.cat((B2_resized, B3_resized, B4, B5_resized), dim=1))
+
+        x_h, x_l, x_s = torch.split(x_b, [self.c_h, self.c_l, self.c_s], dim=1)
+
+        x_h_outs = [m(x_h) for m in self.high_order_branch]
+        x_h_fused = self.high_order_fuse(torch.cat(x_h_outs, dim=1))
+
+        x_l_out = self.low_order_branch(x_l)
+
+        y = self.final_fuse(torch.cat((x_h_fused, x_l_out, x_s), dim=1))
+
+        return y
+
+
+class GatedFusion(nn.Module):
+    def __init__(self, in_channels):
+        super().__init__()
+        self.gamma = nn.Parameter(torch.zeros(1, in_channels, 1, 1))
+
+    def forward(self, f_in, h):
+        if f_in.shape[1] != h.shape[1]:
+            raise ValueError(f"Channel mismatch: f_in={f_in.shape}, h={h.shape}")
+        return f_in + self.gamma * h
+
+
+class Backbone(nn.Module):
+    def __init__(self, in_channels=256, base_channels=64, base_depth=3):
+        super().__init__()
+        c = base_channels
+        c2 = base_channels
+        c3 = 256
+        c4 = 384
+        c5 = 512
+        c6 = 768
+
+        self.stem = DSConv(in_channels, c2, k=3, s=(2, 1), p=1)
+
+        self.p2 = nn.Sequential(
+            DSConv(c2, c3, k=3, s=(2, 1), p=1), DS_C3k2(c3, c3, n=base_depth)
+        )
+
+        self.p3 = nn.Sequential(
+            DSConv(c3, c4, k=3, s=(2, 1), p=1), DS_C3k2(c4, c4, n=base_depth * 2)
+        )
+
+        self.p4 = nn.Sequential(
+            DSConv(c4, c5, k=3, s=(2, 1), p=1), DS_C3k2(c5, c5, n=base_depth * 2)
+        )
+
+        self.p5 = nn.Sequential(
+            DSConv(c5, c6, k=3, s=(2, 1), p=1), DS_C3k2(c6, c6, n=base_depth)
+        )
+
+        self.out_channels = [c3, c4, c5, c6]
+
+    def forward(self, x):
+        x = self.stem(x)
+        x2 = self.p2(x)
+        x3 = self.p3(x2)
+        x4 = self.p4(x3)
+        x5 = self.p5(x4)
+        return [x2, x3, x4, x5]
+
+
+class Decoder(nn.Module):
+    def __init__(
+        self,
+        encoder_channels: List[int],
+        hyperace_out_c: int,
+        decoder_channels: List[int],
+    ):
+        super().__init__()
+        c_p2, c_p3, c_p4, c_p5 = encoder_channels
+        c_d2, c_d3, c_d4, c_d5 = decoder_channels
+
+        self.h_to_d5 = Conv(hyperace_out_c, c_d5, 1, 1)
+        self.h_to_d4 = Conv(hyperace_out_c, c_d4, 1, 1)
+        self.h_to_d3 = Conv(hyperace_out_c, c_d3, 1, 1)
+        self.h_to_d2 = Conv(hyperace_out_c, c_d2, 1, 1)
+
+        self.fusion_d5 = GatedFusion(c_d5)
+        self.fusion_d4 = GatedFusion(c_d4)
+        self.fusion_d3 = GatedFusion(c_d3)
+        self.fusion_d2 = GatedFusion(c_d2)
+
+        self.skip_p5 = Conv(c_p5, c_d5, 1, 1)
+        self.skip_p4 = Conv(c_p4, c_d4, 1, 1)
+        self.skip_p3 = Conv(c_p3, c_d3, 1, 1)
+        self.skip_p2 = Conv(c_p2, c_d2, 1, 1)
+
+        self.up_d5 = DS_C3k2(c_d5, c_d4, n=1)
+        self.up_d4 = DS_C3k2(c_d4, c_d3, n=1)
+        self.up_d3 = DS_C3k2(c_d3, c_d2, n=1)
+
+        self.final_d2 = DS_C3k2(c_d2, c_d2, n=1)
+
+    def forward(self, enc_feats: List[torch.Tensor], h_ace: torch.Tensor):
+        p2, p3, p4, p5 = enc_feats
+
+        d5 = self.skip_p5(p5)
+        h_d5 = self.h_to_d5(F.interpolate(h_ace, size=d5.shape[2:], mode="bilinear"))
+        d5 = self.fusion_d5(d5, h_d5)
+
+        d5_up = F.interpolate(d5, size=p4.shape[2:], mode="bilinear")
+        d4_skip = self.skip_p4(p4)
+        d4 = self.up_d5(d5_up) + d4_skip
+
+        h_d4 = self.h_to_d4(F.interpolate(h_ace, size=d4.shape[2:], mode="bilinear"))
+        d4 = self.fusion_d4(d4, h_d4)
+
+        d4_up = F.interpolate(d4, size=p3.shape[2:], mode="bilinear")
+        d3_skip = self.skip_p3(p3)
+        d3 = self.up_d4(d4_up) + d3_skip
+
+        h_d3 = self.h_to_d3(F.interpolate(h_ace, size=d3.shape[2:], mode="bilinear"))
+        d3 = self.fusion_d3(d3, h_d3)
+
+        d3_up = F.interpolate(d3, size=p2.shape[2:], mode="bilinear")
+        d2_skip = self.skip_p2(p2)
+        d2 = self.up_d3(d3_up) + d2_skip
+
+        h_d2 = self.h_to_d2(F.interpolate(h_ace, size=d2.shape[2:], mode="bilinear"))
+        d2 = self.fusion_d2(d2, h_d2)
+
+        d2_final = self.final_d2(d2)
+
+        return d2_final
+
+
+class FreqPixelShuffle(nn.Module):
+    def __init__(self, in_channels, out_channels, scale=2):
+        super().__init__()
+        self.scale = scale
+        self.conv = DSConv(in_channels, out_channels * scale, k=3, s=1, p=1)
+        self.act = nn.SiLU()
+
+    def forward(self, x):
+        x = self.conv(x)
+        B, C_r, H, W = x.shape
+        out_c = C_r // self.scale
+
+        x = x.view(B, out_c, self.scale, H, W)
+
+        x = x.permute(0, 1, 3, 4, 2).contiguous()
+        x = x.view(B, out_c, H, W * self.scale)
+
+        return x
+
+
+class ProgressiveUpsampleHead(nn.Module):
+    def __init__(self, in_channels, out_channels, target_bins=1025):
+        super().__init__()
+        self.target_bins = target_bins
+
+        c = in_channels
+
+        self.block1 = FreqPixelShuffle(c, c, scale=2)
+        self.block2 = FreqPixelShuffle(c, c // 2, scale=2)
+        self.block3 = FreqPixelShuffle(c // 2, c // 2, scale=2)
+        self.block4 = FreqPixelShuffle(c // 2, c // 4, scale=2)
+
+        self.final_conv = nn.Conv2d(c // 4, out_channels, kernel_size=1, bias=False)
+
+    def forward(self, x):
+
+        x = self.block1(x)
+        x = self.block2(x)
+        x = self.block3(x)
+        x = self.block4(x)
+
+        if x.shape[-1] != self.target_bins:
+            x = F.interpolate(
+                x,
+                size=(x.shape[2], self.target_bins),
+                mode="bilinear",
+                align_corners=False,
+            )
+
+        x = self.final_conv(x)
+        return x
+
+
+class SegmModel(nn.Module):
+    def __init__(
+        self,
+        in_bands=62,
+        in_dim=256,
+        out_bins=1025,
+        out_channels=4,
+        base_channels=64,
+        base_depth=2,
+        num_hyperedges=16,
+        num_heads=8,
+    ):
+        super().__init__()
+
+        self.backbone = Backbone(
+            in_channels=in_dim, base_channels=base_channels, base_depth=base_depth
+        )
+        enc_channels = self.backbone.out_channels
+        c2, c3, c4, c5 = enc_channels
+
+        hyperace_in_channels = enc_channels
+        hyperace_out_channels = c4
+        self.hyperace = HyperACE(
+            hyperace_in_channels,
+            hyperace_out_channels,
+            num_hyperedges,
+            num_heads,
+            k=3,
+            l=2,
+        )
+
+        decoder_channels = [c2, c3, c4, c5]
+        self.decoder = Decoder(enc_channels, hyperace_out_channels, decoder_channels)
+
+        self.upsample_head = ProgressiveUpsampleHead(
+            in_channels=decoder_channels[0],
+            out_channels=out_channels,
+            target_bins=out_bins,
+        )
+
+    def forward(self, x):
+        H, W = x.shape[2:]
+
+        enc_feats = self.backbone(x)
+
+        h_ace_feats = self.hyperace(enc_feats)
+
+        dec_feat = self.decoder(enc_feats, h_ace_feats)
+
+        feat_time_restored = F.interpolate(
+            dec_feat, size=(H, dec_feat.shape[-1]), mode="bilinear", align_corners=False
+        )
+
+        out = self.upsample_head(feat_time_restored)
+
+        return out
+
+
+def MLP(dim_in, dim_out, dim_hidden=None, depth=1, activation=nn.Tanh):
+    dim_hidden = default(dim_hidden, dim_in)
+
+    net = []
+    dims = (dim_in, *((dim_hidden,) * (depth - 1)), dim_out)
+
+    for ind, (layer_dim_in, layer_dim_out) in enumerate(zip(dims[:-1], dims[1:])):
+        is_last = ind == (len(dims) - 2)
+
+        net.append(nn.Linear(layer_dim_in, layer_dim_out))
+
+        if is_last:
+            continue
+
+        net.append(activation())
+
+    return nn.Sequential(*net)
+
+
+class MaskEstimator(Module):
+    @beartype
+    def __init__(self, dim, dim_inputs: Tuple[int, ...], depth, mlp_expansion_factor=4):
+        super().__init__()
+        self.dim_inputs = dim_inputs
+        self.to_freqs = ModuleList([])
+        dim_hidden = dim * mlp_expansion_factor
+
+        for dim_in in dim_inputs:
+            net = []
+
+            mlp = nn.Sequential(
+                MLP(dim, dim_in * 2, dim_hidden=dim_hidden, depth=depth), nn.GLU(dim=-1)
+            )
+
+            self.to_freqs.append(mlp)
+
+        self.segm = SegmModel(
+            in_bands=len(dim_inputs), in_dim=dim, out_bins=sum(dim_inputs) // 4
+        )
+
+    def forward(self, x):
+        y = rearrange(x, "b t f c -> b c t f")
+        y = self.segm(y)
+        y = rearrange(y, "b c t f -> b t (f c)")
+
+        x = x.unbind(dim=-2)
+
+        outs = []
+
+        for band_features, mlp in zip(x, self.to_freqs):
+            freq_out = mlp(band_features)
+            outs.append(freq_out)
+
+        return torch.cat(outs, dim=-1) + y
+
+
+DEFAULT_FREQS_PER_BANDS = (
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2,
+    4,
+    4,
+    4,
+    4,
+    4,
+    4,
+    4,
+    4,
+    4,
+    4,
+    4,
+    4,
+    12,
+    12,
+    12,
+    12,
+    12,
+    12,
+    12,
+    12,
+    24,
+    24,
+    24,
+    24,
+    24,
+    24,
+    24,
+    24,
+    48,
+    48,
+    48,
+    48,
+    48,
+    48,
+    48,
+    48,
+    128,
+    129,
+)
+
+
+class BSRoformerHyperACE(Module):
+
+    @beartype
+    def __init__(
+        self,
+        dim,
+        *,
+        depth,
+        stereo=False,
+        num_stems=1,
+        time_transformer_depth=2,
+        freq_transformer_depth=2,
+        linear_transformer_depth=0,
+        freqs_per_bands: Tuple[int, ...] = DEFAULT_FREQS_PER_BANDS,
+        dim_head=64,
+        heads=8,
+        attn_dropout=0.0,
+        ff_dropout=0.0,
+        flash_attn=True,
+        dim_freqs_in=1025,
+        stft_n_fft=2048,
+        stft_hop_length=512,
+        stft_win_length=2048,
+        stft_normalized=False,
+        stft_window_fn: Optional[Callable] = None,
+        mask_estimator_depth=2,
+        multi_stft_resolution_loss_weight=1.0,
+        multi_stft_resolutions_window_sizes: Tuple[int, ...] = (
+            4096,
+            2048,
+            1024,
+            512,
+            256,
+        ),
+        multi_stft_hop_size=147,
+        multi_stft_normalized=False,
+        multi_stft_window_fn: Callable = torch.hann_window,
+        mlp_expansion_factor=4,
+        use_torch_checkpoint=False,
+        skip_connection=False,
+        sage_attention=False,
+    ):
+        super().__init__()
+
+        self.stereo = stereo
+        self.audio_channels = 2 if stereo else 1
+        self.num_stems = num_stems
+        self.use_torch_checkpoint = use_torch_checkpoint
+        self.skip_connection = skip_connection
+
+        self.layers = ModuleList([])
+
+        if sage_attention:
+            print("Use Sage Attention")
+
+        transformer_kwargs = dict(
+            dim=dim,
+            heads=heads,
+            dim_head=dim_head,
+            attn_dropout=attn_dropout,
+            ff_dropout=ff_dropout,
+            flash_attn=flash_attn,
+            norm_output=False,
+            sage_attention=sage_attention,
+        )
+
+        time_rotary_embed = RotaryEmbedding(dim=dim_head)
+        freq_rotary_embed = RotaryEmbedding(dim=dim_head)
+
+        for _ in range(depth):
+            tran_modules = []
+            tran_modules.append(
+                Transformer(
+                    depth=time_transformer_depth,
+                    rotary_embed=time_rotary_embed,
+                    **transformer_kwargs,
+                )
+            )
+            tran_modules.append(
+                Transformer(
+                    depth=freq_transformer_depth,
+                    rotary_embed=freq_rotary_embed,
+                    **transformer_kwargs,
+                )
+            )
+            self.layers.append(nn.ModuleList(tran_modules))
+
+        self.final_norm = RMSNorm(dim)
+
+        self.stft_kwargs = dict(
+            n_fft=stft_n_fft,
+            hop_length=stft_hop_length,
+            win_length=stft_win_length,
+            normalized=stft_normalized,
+        )
+
+        self.stft_window_fn = partial(
+            default(stft_window_fn, torch.hann_window), stft_win_length
+        )
+
+        freqs = torch.stft(
+            torch.randn(1, 4096),
+            **self.stft_kwargs,
+            window=torch.ones(stft_win_length),
+            return_complex=True,
+        ).shape[1]
+
+        assert len(freqs_per_bands) > 1
+        assert (
+            sum(freqs_per_bands) == freqs
+        ), f"the number of freqs in the bands must equal {freqs} based on the STFT settings, but got {sum(freqs_per_bands)}"
+
+        freqs_per_bands_with_complex = tuple(
+            2 * f * self.audio_channels for f in freqs_per_bands
+        )
+
+        self.band_split = BandSplit(dim=dim, dim_inputs=freqs_per_bands_with_complex)
+
+        self.mask_estimators = nn.ModuleList([])
+
+        for _ in range(num_stems):
+            mask_estimator = MaskEstimator(
+                dim=dim,
+                dim_inputs=freqs_per_bands_with_complex,
+                depth=mask_estimator_depth,
+                mlp_expansion_factor=mlp_expansion_factor,
+            )
+
+            self.mask_estimators.append(mask_estimator)
+
+        self.multi_stft_resolution_loss_weight = multi_stft_resolution_loss_weight
+        self.multi_stft_resolutions_window_sizes = multi_stft_resolutions_window_sizes
+        self.multi_stft_n_fft = stft_n_fft
+        self.multi_stft_window_fn = multi_stft_window_fn
+
+        self.multi_stft_kwargs = dict(
+            hop_length=multi_stft_hop_size, normalized=multi_stft_normalized
+        )
+
+    def forward(self, raw_audio, target=None, return_loss_breakdown=False):
+
+        device = raw_audio.device
+
+        x_is_mps = True if device.type == "mps" else False
+
+        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 needs to be set to True if passing in audio signal that is stereo (channel dimension of 2). also need to be False if mono (channel dimension of 1)"
+
+        raw_audio, batch_audio_channel_packed_shape = pack_one(raw_audio, "* t")
+
+        stft_window = self.stft_window_fn(device=device)
+
+        try:
+            stft_repr = torch.stft(
+                raw_audio, **self.stft_kwargs, window=stft_window, return_complex=True
+            )
+        except:
+            stft_repr = torch.stft(
+                raw_audio.cpu() if x_is_mps else raw_audio,
+                **self.stft_kwargs,
+                window=stft_window.cpu() if x_is_mps else stft_window,
+                return_complex=True,
+            ).to(device)
+        stft_repr = torch.view_as_real(stft_repr)
+
+        stft_repr = unpack_one(stft_repr, batch_audio_channel_packed_shape, "* f t c")
+
+        stft_repr = rearrange(stft_repr, "b s f t c -> b (f s) t c")
+
+        x = rearrange(stft_repr, "b f t c -> b t (f c)")
+
+        x = self.band_split(x)
+
+        for i, transformer_block in enumerate(self.layers):
+
+            time_transformer, freq_transformer = transformer_block
+
+            x = rearrange(x, "b t f d -> b f t d")
+            x, ps = pack([x], "* t d")
+
+            x = time_transformer(x)
+
+            (x,) = unpack(x, ps, "* t d")
+            x = rearrange(x, "b f t d -> b t f d")
+            x, ps = pack([x], "* f d")
+
+            x = freq_transformer(x)
+
+            (x,) = unpack(x, ps, "* f d")
+
+        x = self.final_norm(x)
+
+        num_stems = len(self.mask_estimators)
+
+        mask = torch.stack([fn(x) for fn in self.mask_estimators], dim=1)
+        mask = rearrange(mask, "b n t (f c) -> b n f t c", c=2)
+
+        stft_repr = rearrange(stft_repr, "b f t c -> b 1 f t c")
+
+        stft_repr = torch.view_as_complex(stft_repr)
+        mask = torch.view_as_complex(mask)
+
+        stft_repr = stft_repr * mask
+
+        stft_repr = rearrange(
+            stft_repr, "b n (f s) t -> (b n s) f t", s=self.audio_channels
+        )
+
+        try:
+            recon_audio = torch.istft(
+                stft_repr,
+                **self.stft_kwargs,
+                window=stft_window,
+                return_complex=False,
+                length=raw_audio.shape[-1],
+            )
+        except:
+            recon_audio = torch.istft(
+                stft_repr.cpu() if x_is_mps else stft_repr,
+                **self.stft_kwargs,
+                window=stft_window.cpu() if x_is_mps else stft_window,
+                return_complex=False,
+                length=raw_audio.shape[-1],
+            ).to(device)
+
+        recon_audio = rearrange(
+            recon_audio, "(b n s) t -> b n s t", s=self.audio_channels, n=num_stems
+        )
+
+        if num_stems == 1:
+            recon_audio = rearrange(recon_audio, "b 1 s t -> b s t")
+
+        if not exists(target):
+            return recon_audio
+
+        if self.num_stems > 1:
+            assert target.ndim == 4 and target.shape[1] == self.num_stems
+
+        if target.ndim == 2:
+            target = rearrange(target, "... t -> ... 1 t")
+
+        target = target[..., : recon_audio.shape[-1]]
+
+        loss = F.l1_loss(recon_audio, target)
+
+        multi_stft_resolution_loss = 0.0
+
+        for window_size in self.multi_stft_resolutions_window_sizes:
+            res_stft_kwargs = dict(
+                n_fft=max(window_size, self.multi_stft_n_fft),
+                win_length=window_size,
+                return_complex=True,
+                window=self.multi_stft_window_fn(window_size, device=device),
+                **self.multi_stft_kwargs,
+            )
+
+            recon_Y = torch.stft(
+                rearrange(recon_audio, "... s t -> (... s) t"), **res_stft_kwargs
+            )
+            target_Y = torch.stft(
+                rearrange(target, "... s t -> (... s) t"), **res_stft_kwargs
+            )
+
+            multi_stft_resolution_loss = multi_stft_resolution_loss + F.l1_loss(
+                recon_Y, target_Y
+            )
+
+        weighted_multi_resolution_loss = (
+            multi_stft_resolution_loss * self.multi_stft_resolution_loss_weight
+        )
+
+        total_loss = loss + weighted_multi_resolution_loss
+
+        if not return_loss_breakdown:
+            return total_loss
+
+        return total_loss, (loss, multi_stft_resolution_loss)