Убраны комментарии и отформатирован код

MVSepLess_Epsilon_Colab.ipynb

@@ -51,7 +51,7 @@
         "prodigyopt\n",
         "torch_log_wmse\n",
         "rotary_embedding_torch\n",
-        "gradio\n",
+        "gradio<=6.0\n",
         "omegaconf\n",
         "beartype\n",
         "spafe\n",
@@ -395,6 +395,18 @@
         "!python mvsepless/model_manager.py vbach list"
       ]
     },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "#@title Удаление голосовой модели\n",
+        "%cd $mvsepless_dir\n",
+        "voicemodel_name = \"\" # @param {\"type\":\"string\",\"placeholder\":\"Имя модели\"}\n",
+        "!python mvsepless/model_manager.py vbach remove --model_name \"$voicemodel_name\""
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -461,9 +473,7 @@
     },
     {
       "cell_type": "markdown",
-      "metadata": {
-        "id": "4GSGHM4rYSVp"
-      },
+      "metadata": {},
       "source": [
         "## Инференс"
       ]
@@ -485,7 +495,7 @@
         "# @markdown ---\n",
         "# @markdown  ### Hubert\n",
         "# @markdown * Стэк\n",
-        "stack = \"transformers\" # @param [\"fairseq\",\"transformers\"]\n",
+        "stack = \"fairseq\" # @param [\"fairseq\",\"transformers\"]\n",
         "# @markdown * Имя модели для fairseq\n",
         "fairseq_embedder = \"hubert_base\" # @param [\"hubert_base\",\"contentvec_base\",\"korean_hubert_base\",\"chinese_hubert_base\",\"portuguese_hubert_base\",\"japanese_hubert_base\"]\n",
         "# @markdown * Имя модели для transformers\n",
@@ -497,7 +507,7 @@
         "# @markdown * Стерео режим\n",
         "stereo_mode = \"mono\" # @param [\"mono\",\"left/right\",\"sim/dif\"]\n",
         "# @markdown * Метод определения тона\n",
-        "method_pitch = \"rmvpe+\" # @param [\"rmvpe+\",\"mangio-crepe\",\"fcpe\"]\n",
+        "method_pitch = \"rmvpe+\" # @param [\"rmvpe+\",\"mangio-crepe\",\"mangio-crepe-tiny\",\"fcpe\",'harvest\",\"pm\",\"pyin\"]\n",
         "# @markdown * Изменение высоты тона (полутона)\n",
         "pitch = 0 # @param {\"type\":\"slider\",\"min\":-48,\"max\":48,\"step\":1}\n",
         "# @markdown * Длина шага (для mangio-crepe)\n",

mvsepless/__main__.py

@@ -11,18 +11,59 @@ if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="MVSepless")
     subparsers = parser.add_subparsers(dest="command")
     app_parser = subparsers.add_parser("app", help="Приложение MVSepless")
-    app_parser.add_argument("--port", type=int, default=None, help="Порт для запуска сервера Gradio.")
-    app_parser.add_argument("--share", action="store_true", help="Создать публичную ссылку для приложения Gradio.")
+    app_parser.add_argument(
+        "--port", type=int, default=None, help="Порт для запуска сервера Gradio."
+    )
+    app_parser.add_argument(
+        "--share",
+        action="store_true",
+        help="Создать публичную ссылку для приложения Gradio.",
+    )
     cli_parser = subparsers.add_parser("cli", help="CLI MVSepless Lite")
-    cli_parser.add_argument("--input", type=str, required=True, help="Входной аудиофайл или каталог.")
-    cli_parser.add_argument("--output_dir", type=str, default=None, help="Каталог для выходных файлов.")
-    cli_parser.add_argument("--model_type", type=str, default="mel_band_roformer", help="Тип модели разделения.")
-    cli_parser.add_argument("--model_name", type=str, default="Mel-Band-Roformer_Vocals_kimberley_jensen", help="Имя модели разделения.")
-    cli_parser.add_argument("--ext_inst", action="store_true", help="Извлечь инструментал.")
-    cli_parser.add_argument("--output_format", type=str, default="mp3", choices=Separator.audio.output_formats, help="Формат выходного файла.")
-    cli_parser.add_argument("--output_bitrate", type=str, default="320k", help="Битрейт выходного файла.")
-    cli_parser.add_argument("--template", type=str, default="NAME (STEM) MODEL", help="Шаблон именования выходных файлов.")
-    cli_parser.add_argument("--selected_stems", type=str, nargs='*', default=None, help="Выбранные стемы для разделения.")
+    cli_parser.add_argument(
+        "--input", type=str, required=True, help="Входной аудиофайл или каталог."
+    )
+    cli_parser.add_argument(
+        "--output_dir", type=str, default=None, help="Каталог для выходных файлов."
+    )
+    cli_parser.add_argument(
+        "--model_type",
+        type=str,
+        default="mel_band_roformer",
+        help="Тип модели разделения.",
+    )
+    cli_parser.add_argument(
+        "--model_name",
+        type=str,
+        default="Mel-Band-Roformer_Vocals_kimberley_jensen",
+        help="Имя модели разделения.",
+    )
+    cli_parser.add_argument(
+        "--ext_inst", action="store_true", help="Извлечь инструментал."
+    )
+    cli_parser.add_argument(
+        "--output_format",
+        type=str,
+        default="mp3",
+        choices=Separator.audio.output_formats,
+        help="Формат выходного файла.",
+    )
+    cli_parser.add_argument(
+        "--output_bitrate", type=str, default="320k", help="Битрейт выходного файла."
+    )
+    cli_parser.add_argument(
+        "--template",
+        type=str,
+        default="NAME (STEM) MODEL",
+        help="Шаблон именования выходных файлов.",
+    )
+    cli_parser.add_argument(
+        "--selected_stems",
+        type=str,
+        nargs="*",
+        default=None,
+        help="Выбранные стемы для разделения.",
+    )
     args = parser.parse_args()
 
     if args.command == "app":
@@ -39,7 +80,13 @@ if __name__ == "__main__":
             ],
         )
         mvsepless_lite_app = mvsepless_app(theme=theme)
-        mvsepless_lite_app.launch(server_name="0.0.0.0", server_port=args.port, share=args.share, allowed_paths=["/"], debug=True)
+        mvsepless_lite_app.launch(
+            server_name="0.0.0.0",
+            server_port=args.port,
+            share=args.share,
+            allowed_paths=["/"],
+            debug=True,
+        )
     elif args.command == "cli":
         input_data = args.input
         if os.path.isdir(input_data):
@@ -62,6 +109,6 @@ if __name__ == "__main__":
             output_format=args.output_format,
             output_bitrate=args.output_bitrate,
             template=args.template,
-            selected_stems=args.selected_stems
+            selected_stems=args.selected_stems,
         )
-        print("Разделение завершено.")
+        print("Разделение завершено.")

mvsepless/audio.py

@@ -1,781 +1,789 @@
-import os
-from pathlib import Path
-import sys
-import json
-import subprocess
-import numpy as np
-from typing import Literal
-from collections.abc import Callable
-from pathlib import Path
-from numpy.typing import DTypeLike
-import tempfile
-import librosa
-if not __package__:
-    from namer import Namer
-else:
-    from .namer import Namer
-class NotInputFileSpecified(Exception): pass
-class NotOutputFileSpecified(Exception): pass
-class NotSupportedDataType(Exception): pass
-class ErrorDecode(Exception): pass
-class ErrorEncode(Exception): pass
-class NotSupportedFormat(Exception): pass
-class SampleRateError(Exception): pass
-class FileIsNotAudio(Exception): pass
-
-class Audio(Namer):
-    def __init__(self):
-        """
-Чтение и запись аудио файла через ffmpeg
-
-Поддерживаемые типы данных: - int16, int32, float32, float64
-        """
-        super().__init__()
-        self.ffmpeg_path = os.environ.get("MVSEPLESS_FFMPEG", "ffmpeg")
-        self.ffprobe_path = os.environ.get("MVSEPLESS_FFPROBE", "ffprobe")
-        self.output_formats = ("mp3", "wav", "flac", "ogg", "opus", "m4a", "aac", "ac3", "aiff")
-        self.input_formats = ("mp3", "wav", "flac", "ogg", "opus", "m4a", "aac", "ac3", "aiff", "mp4", "mkv", "webm", "avi", "mov", "ts")
-        self.supported_dtypes = ("int16", "int32", "float32", "float64")
-        self.dtypes_dict = {
-            "int16": "s16le",
-            "int32": "s32le",
-            "float32": "f32le",
-            "float64": "f64le",
-            np.int16: "s16le",
-            np.int32: "s32le",
-            np.float32: "f32le",
-            np.float64: "f64le",
-        }
-        self.bitrate_limit = {
-            "mp3": {"min": 8, "max": 320},
-            "aac": {"min": 8, "max": 512},
-            "m4a": {"min": 8, "max": 512},
-            "ac3": {"min": 32, "max": 640},
-            "ogg": {"min": 64, "max": 500},
-            "opus": {"min": 6, "max": 512},
-        }
-        self.sample_rates = {
-            "mp3": {
-                "supported": (44100, 48000, 32000, 22050, 24000, 16000, 11025, 12000, 8000)
-            },
-            "opus": {"supported": (48000, 24000, 16000, 12000, 8000)},
-            "m4a": {
-                "supported": (
-                    96000,
-                    88200,
-                    64000,
-                    48000,
-                    44100,
-                    32000,
-                    24000,
-                    22050,
-                    16000,
-                    12000,
-                    11025,
-                    8000,
-                    7350,
-                )
-            },
-            "aac": {
-                "supported": (
-                    96000,
-                    88200,
-                    64000,
-                    48000,
-                    44100,
-                    32000,
-                    24000,
-                    22050,
-                    16000,
-                    12000,
-                    11025,
-                    8000,
-                    7350,
-                )
-            },
-            "ac3": {
-                "supported": (
-                    48000,
-                    44100,
-                    32000,
-                )
-            },
-            "ogg": {"min": 6, "max": 192000},
-            "wav": {"min": 0, "max": float("inf")},
-            "aiff": {"min": 0, "max": float("inf")},
-            "flac": {"min": 0, "max": 192000},
-        }
-        self.check_ffmpeg()
-        self.check_ffprobe()
-
-    def check_ffmpeg(self):
-        """
-        Проверяет, установлен ли ffmpeg?
-        """
-        try:
-            ffmpeg_version_output = subprocess.check_output(
-                [self.ffmpeg_path, "-version"], text=True
-            )
-        except FileNotFoundError:
-            if "PYTEST_CURRENT_TEST" not in os.environ:
-                raise FileNotFoundError("FFMPEG не установлен. Укажите путь к установленному FFMPEG через переменную окружения MVSEPLESS_FFMPEG")
-
-    def check_ffprobe(self):
-        """
-        Проверяет, установлен ли ffprobe?
-        """
-        try:
-            ffmpeg_version_output = subprocess.check_output(
-                [self.ffprobe_path, "-version"], text=True
-            )
-        except FileNotFoundError:
-            if "PYTEST_CURRENT_TEST" not in os.environ:
-                raise FileNotFoundError("FFPROBE не установлен. Укажите путь к установленному FFPROBE через переменную окружения MVSEPLESS_FFPROBE")
-
-
-    def fit_sr(
-        self, 
-        f: str | Literal["mp3", "wav", "flac", "ogg", "opus", "m4a", "aac", "ac3", "aiff"] = "mp3",
-        sr: int = 44100
-    ) -> int:
-        """
-        Исправляет значение частоты дисректизации выходного файла
-
-        Параметры:
-            f: Формат вывода
-            sr: Частота дискретизации (целое число)
-        Возвращает:
-            sr: Исправленная частота дискретизации
-        """
-        format_info = self.sample_rates.get(f.lower())
-
-        if not format_info:
-            return None  # Формат не найден
-
-        if "supported" in format_info:
-            # Для форматов с конкретным списком
-            supported_rates = format_info["supported"]
-            if sr in supported_rates:
-                return sr
-
-            # Находим ближайшую поддерживаемую частоту
-            return min(supported_rates, key=lambda x: abs(x - sr))
-
-        elif "min" in format_info and "max" in format_info:
-            # Для форматов с диапазоном - обрезаем до границ
-            min_rate = format_info["min"]
-            max_rate = format_info["max"]
-
-            if sr < min_rate:
-                return min_rate
-            elif sr > max_rate:
-                return max_rate
-            else:
-                return sr
-
-        return None
-    
-    def fit_br(
-        self, 
-        f: str | Literal["mp3", "wav", "flac", "ogg", "opus", "m4a", "aac", "ac3", "aiff"] = "mp3", 
-        br: int = 320
-        ) -> int:
-        """
-        Исправляет значение битрейта выходного файла
-
-        Параметры:
-            f: Формат вывода
-            br: Битрейт (целое число)
-        Возвращает:
-            br: Исправленный битрейт
-        """
-        if f not in self.bitrate_limit:
-            raise NotSupportedFormat(f"Формат {f} не поддерживается")
-
-        limits = self.bitrate_limit[f]
-
-        if br < limits["min"]:
-            return limits["min"]
-        elif br > limits["max"]:
-            return limits["max"]
-        else:
-            return br
-        
-    def get_info(
-        self,
-        i: str | os.PathLike | Callable | None = None,
-    ) -> dict[int, dict[int, float]]:
-        """
-        Получает информацию о аудио потоках из файла напрямую через FFMPEG
-
-        Параметры:
-            i: Путь к выходному файлу
-        Возвращает:
-            audio_info: Словарь с информацией о аудиопотоках вида:
-
-                {Номер потока: 
-                    {
-                        "sample_rate": Частота дисректизации (является целым числом), 
-                        "duration": Длительность аудиопотока (является числом с плавающей точкой)
-                    }
-                }
-        """
-        audio_info = {}
-        if i:
-            if isinstance(i, Path):
-                i = str(i)
-            if os.path.exists(i):
-                cmd = [self.ffprobe_path, "-i", i, "-v", "quiet", "-hide_banner", 
-                    "-show_entries", "stream=index,sample_rate,duration", "-select_streams", "a", "-of", "json"]
-                
-                process = subprocess.Popen(
-                    cmd,
-                    stdin=subprocess.PIPE,
-                    stdout=subprocess.PIPE,
-                    stderr=subprocess.PIPE,
-                )
-
-                stdout, stderr = process.communicate()
-
-                if process.returncode != 0:
-                    print(f"STDERR: {stderr.decode('utf-8')}")
-                    print(f"STDOUT: {stdout.decode('utf-8')}")
-
-                json_output = json.loads(stdout)
-                streams = json_output["streams"]
-                if not streams:
-                    pass
-                    
-                else:
-                    for a, stream in enumerate(streams):
-                        audio_info[a] = {
-                            "sample_rate": int(stream.get("sample_rate", 0)),
-                            "duration": float(stream.get("duration", 0))
-                        }
-
-                return audio_info
-                
-            else:
-                raise FileExistsError("Указанного файла не существует")
-            
-        else:
-            raise NotInputFileSpecified("Не указан путь к файлу")
-
-    def check(
-        self,
-        i: str | os.PathLike | Callable | None = None
-    ) -> bool:
-        """
-        Проверяет, является ли файл аудио или видео файлом, поддерживаемым ffmpeg
-
-        Параметры:
-            i: Путь к выходному файлу
-        Возвращает:
-            is_audio_video: Булево значение, является ли файл аудио или видео файлом
-        """
-        if i:
-            if isinstance(i, Path):
-                i = str(i)
-            if os.path.exists(i):
-                info = self.get_info(i=i)
-                if info:
-                    list_streams = list(info.keys())
-                    if len(list_streams) > 0:
-                        if info[0].get("sample_rate") > 0:
-                            return True
-                        else:
-                            return False
-                    else:
-                        return False
-                else:
-                    return False
-            else:
-                raise FileExistsError("Указанного файла не существует")
-        else:
-            raise NotInputFileSpecified("Не указан путь к файлу")
-
-    def read(
-        self,
-        i: str | os.PathLike | Callable | None = None, 
-        sr: int | None = None, 
-        mono: bool = False,
-        dtype: DTypeLike = np.float32,
-        s: int = 0
-    ) -> tuple[np.ndarray, int, float]:
-        """
-        Читает аудио-файл, преобразовывая его в массив с аудио данными напрямую через FFMPEG
-        Является заменой soundfile.read() и librosa.load()
-
-        Параметры:
-            i: Путь к выходному файлу
-            sr: Целевая частота дискретизации (Если не указана, то используется частота дискретизации входного файла)
-            mono: Конвертация в моно (по умолчанию отключена)
-            dtype: Тип данных (поддерживаются типы: int16, int32, float32, float64; по умолчанию - float32)
-            s: Номер аудиопотока (по умолчанию 0)
-        Возвращает:
-            audio_array: Массив с аудио данными
-            sr: Частота дискретизации массива
-            duration: Длительность аудио (количество сэмплов / частота дискретизации)
-        """
-        output_format = self.dtypes_dict.get(dtype, None)
-        if not output_format:
-            raise NotSupportedDataType(f"Этот тип данных не поддерживается {dtype}")
-        if i:
-            if isinstance(i, Path):
-                i = str(i)
-            if os.path.exists(i):
-                audio_info = self.get_info(i=i)
-                list_streams = list(audio_info.keys())
-                if audio_info.get(s, False):
-                    stream = s
-                else:
-                    if len(list_streams) > 0:
-                        stream = 0
-                    else:
-                        raise FileIsNotAudio("В входном файле нет аудио потоков")
-                
-                sample_rate_input = audio_info[stream]["sample_rate"]
-                if sample_rate_input == 0:
-                    raise FileIsNotAudio("В входном файле нет аудио потоков")
-
-                cmd = [
-                    self.ffmpeg_path,
-                    "-i", i,
-                    "-map", f"0:a:{stream}", "-vn",
-                    "-f", output_format,
-                    "-ac", "1" if mono else "2",
-                ]
-                
-                if sr:
-                    cmd.extend(["-ar", str(sr)])
-                else:
-                    sr = sample_rate_input
-                
-                cmd.append("pipe:1")
-                
-                process = subprocess.Popen(
-                    cmd,
-                    stdout=subprocess.PIPE,
-                    stderr=subprocess.PIPE,
-                    bufsize=10**8
-                )
-                
-                try:
-
-                    raw_audio, stderr = process.communicate(timeout=300)
-                    
-                    if process.returncode != 0:
-                        raise ErrorDecode(f"FFmpeg error: {stderr.decode()}")
-                        
-                except subprocess.TimeoutExpired:
-                    process.kill()
-                    raise ErrorDecode("FFmpeg timeout при чтении файла")
-                
-                audio_array = np.frombuffer(raw_audio, dtype=dtype)
-                
-                channels = 1 if mono else 2
-                audio_array = audio_array.reshape((-1, channels)).T
-                if audio_array.ndim > 1 and channels == 1:
-                    audio_array = np.mean(audio_array, axis=tuple(range(audio_array.ndim - 1)))
-
-                len_samples = float(audio_array.shape[-1])
-
-                duration = len_samples / sr
-
-                print(f"Частота дискретизации: {sr}")
-                
-                return audio_array.copy(), sr, duration
-            else:
-                raise FileExistsError("Указанного файла не существует")
-                
-        else:
-            raise NotInputFileSpecified("Не указан путь к файлу")
-    
-    def write(
-        self, 
-        o: str | os.PathLike | Callable | None = None, 
-        array: np.ndarray = np.array([], dtype=np.float32),
-        sr: int = 44100, 
-        of: str | Literal["mp3", "wav", "flac", "ogg", "opus", "m4a", "aac", "ac3", "aiff"] | None = None, 
-        br: str | int | None = None
-    ) -> str:
-        """
-        Записывает numpy-массив с аудио данными в файл напрямую через ffmpeg.
-        Является заменой soundfile.write()
-
-        Параметры:
-            o: Путь к выходному файлу
-            array: Массив с аудио данными (поддерживаются типы: int16, int32, float32, float64)
-            sr: Частота дискретизации массива
-            of: Формат вывода (по умолчанию mp3)
-            br: Битрейт для кодеков, сжимающих аудио с потерями
-        Возвращает:
-            o: Путь к выходному файлу
-        """
-        if isinstance(array, np.ndarray):
-
-            if len(array.shape) == 1:
-                array = array.reshape(-1, 1)
-            elif len(array.shape) == 2:
-                if array.shape[0] == 2:
-                    array = array.T
-            else:
-                raise ValueError("numpy-массив должен быть либо одномерным, либо двухмерным")
-
-            if array.dtype == np.int16:
-                input_format = "s16le"
-            elif array.dtype == np.int32:
-                input_format = "s32le"
-            elif array.dtype == np.float32:
-                input_format = "f32le"
-            elif array.dtype == np.float64:
-                input_format = "f64le"
-            else:
-                raise NotSupportedDataType(f"Этот тип данных не поддерживается {array.dtype}")
-
-            if array.shape[1] == 1:
-                audio_bytes = array.tobytes()
-
-                channels = 1
-                
-            elif array.shape[1] == 2:
-                audio_bytes = array.tobytes()
-
-                channels = 2
-            else:
-                raise ValueError("numpy-массив должен содержать 1 или 2 канала")
-
-        else:
-            raise ValueError("Вход должен быть numpy-массивом")
-
-        if o:
-            if isinstance(o, Path):
-                o = str(o)
-            output_dir = os.path.dirname(o)
-            output_base = os.path.basename(o)
-            output_name, output_ext = os.path.splitext(output_base)
-            if output_dir != "":
-                os.makedirs(output_dir, exist_ok=True)
-            if output_ext == "":
-                if of:
-                    o += f".{of}"
-                else:
-                    o += f".mp3"
-            elif output_ext == ".":
-                if of:
-                    o += f"{of}"
-                else:
-                    o += f"mp3"
-        else:
-            raise NotOutputFileSpecified("Не указан путь к выходному файлу")
-        
-        if of:
-            if of in self.output_formats:
-                output_name, output_ext = os.path.splitext(o)
-                if output_ext == f".{of}":
-                    pass
-                else:
-                    o = f"{os.path.join(output_dir, output_name)}.{of}"
-            else:
-                raise NotSupportedFormat(f"Неподдерживаемый формат: {of}")
-        else:
-            of = os.path.splitext(o)[1].strip(".")
-            if of in self.output_formats:
-                pass
-            else:
-                raise NotSupportedFormat(f"Неподдерживаемый формат: {of}")
-
-        if sr:
-            if isinstance(sr, int):
-                sample_rate_fixed = self.fit_sr(f=of, sr=sr)
-            elif isinstance(sr, float):
-                sr = int(sr)
-                sample_rate_fixed = self.fit_sr(f=of, sr=sr)
-            else:
-                raise SampleRateError(f"Частота дискретизации должна быть числом\n\nЗначение: {sr}\nТип: {type(sr)}")
-        else:
-            raise SampleRateError("Не указана частота дискретизации")
-
-        bitrate_fixed = "320k"
-
-        if of not in ["wav", "flac", "aiff"]:
-            if br:
-                if isinstance(br, int):
-                    bitrate_fixed = self.fit_br(f=of, br=br)
-                elif isinstance(br, float):
-                    bitrate_fixed = self.fit_br(f=of, br=int(br))
-                elif isinstance(br, str):
-                    bitrate_fixed = self.fit_br(f=of, br=int(br.strip("k").strip("K")))
-                else:
-                    bitrate_fixed = self.fit_br(f=of, br=320)
-            else:
-                bitrate_fixed = self.fit_br(of, 320)
-
-        format_settings = {
-            "wav": [
-                "-c:a",
-                "pcm_f32le",
-                "-sample_fmt",
-                "flt",
-            ],
-            "aiff": [
-                "-c:a",
-                "pcm_f32be",
-                "-sample_fmt",
-                "flt",
-            ],
-            "flac": [
-                "-c:a",
-                "flac",
-                "-compression_level",
-                "12",
-                "-sample_fmt",
-                "s32",
-            ],
-            "mp3": [
-                "-c:a",
-                "libmp3lame",
-                "-b:a",
-                f"{bitrate_fixed}k",
-            ],
-            "ogg": [
-                "-c:a",
-                "libvorbis",
-                "-b:a",
-                f"{bitrate_fixed}k",
-            ],
-            "opus": [
-                "-c:a",
-                "libopus",
-                "-b:a",
-                f"{bitrate_fixed}k",
-            ],
-            "m4a": [
-                "-c:a",
-                "aac",
-                "-b:a",
-                f"{bitrate_fixed}k",
-            ],
-            "aac": [
-                "-c:a",
-                "aac",
-                "-b:a",
-                f"{bitrate_fixed}k",
-            ],
-            "ac3": [
-                "-c:a",
-                "ac3",
-                "-b:a",
-                f"{bitrate_fixed}k",
-            ],
-        }
-
-        cmd = [
-            self.ffmpeg_path,
-            "-y",
-            "-f",
-            input_format,
-            "-ar",
-            str(sr),
-            "-ac",
-            str(channels),
-            "-i",
-            "pipe:0",
-            "-ac",
-            str(channels),
-        ]
-
-        cmd.extend(["-ar", str(sample_rate_fixed)])
-        cmd.extend(format_settings[of])
-        o_dir, o_base = os.path.split(o)
-        o_base_n, o_base_ext = os.path.splitext(o_base)
-        o_base_n = self.sanitize(o_base_n)
-        o_base_n = self.short(o_base_n)
-        o = os.path.join(o_dir, f"{o_base_n}{o_base_ext}")
-        o = self.iter(o)
-        cmd.append(o)
-
-        process = subprocess.Popen(
-            cmd,
-            stdin=subprocess.PIPE,
-            stdout=subprocess.PIPE,
-            stderr=subprocess.PIPE,
-        )
-
-        try:
-            stdout, stderr = process.communicate(input=audio_bytes, timeout=300)
-        except subprocess.TimeoutExpired:
-            process.kill()
-            raise ErrorEncode("FFmpeg timeout: операция заняла слишком много времени")
-
-        if process.returncode != 0:
-            raise ErrorEncode(f"FFmpeg завершился с ошибкой (код: {process.returncode})")
-
-        return os.path.abspath(o)
-
-class Inverter(Audio):
-    def __init__(self):
-        super().__init__()
-        self.test = "test"
-        self.w_types = [
-            "boxcar",  # Прямоугольное окно
-            "triang",  # Треугольное окно
-            "blackman",  # Окно Блэкмана
-            "hamming",  # Окно Хэмминга
-            "hann",  # Окно Ханна
-            "bartlett",  # Окно Бартлетта
-            "flattop",  # Окно с плоской вершиной
-            "parzen",  # Окно Парзена
-            "bohman",  # Окно Бохмана
-            "blackmanharris",  # Окно Блэкмана-Харриса
-            "nuttall",  # Окно Нуттала
-            "barthann",  # Окно Бартлетта-Ханна
-            "cosine",  # Косинусное окно
-            "exponential",  # Экспоненциальное окно
-            "tukey",  # Окно Туки
-            "taylor",  # Окно Тейлора
-            "lanczos",  # Окно Ланцоша
-        ]
-
-    def load_audio(self, filepath):
-        """Загрузка аудиофайла с помощью librosa"""
-        try:
-            y, sr, _ = self.read(i=filepath, sr=None, mono=False)
-            return y, sr
-        except Exception as e:
-            print(f"Ошибка загрузки аудио: {e}")
-            return None, None
-
-    def process_channel(self, y1_ch, y2_ch, sr, method, w_size=2048, overlap=2, w_type="hann"):
-        """Обработка одного аудиоканала"""
-        HOP_LENGTH = w_size // overlap
-        if method == "waveform":
-            return y1_ch - y2_ch
-
-        elif method == "spectrogram":
-            # Вычисляем спектрограммы
-            S1 = librosa.stft(
-                y1_ch, n_fft=w_size, hop_length=HOP_LENGTH, win_length=w_size
-            )
-            S2 = librosa.stft(
-                y2_ch, n_fft=w_size, hop_length=HOP_LENGTH, win_length=w_size
-            )
-
-            # Амплитудные спектрограммы
-            mag1 = np.abs(S1)
-            mag2 = np.abs(S2)
-
-            # Спектральное вычитание
-            mag_result = np.maximum(mag1 - mag2, 0)
-
-            # Сохраняем фазовую информацию исходного сигнала
-            phase = np.angle(S1)
-
-            # Комбинируем амплитуду результата с фазой
-            S_result = mag_result * np.exp(1j * phase)
-
-            # Обратное преобразование
-            return librosa.istft(
-                S_result,
-                n_fft=w_size,
-                hop_length=HOP_LENGTH,
-                win_length=w_size,
-                length=len(y1_ch),
-            )
-
-    def process_audio(self, audio1_path, audio2_path, out_format, method, output_path="./inverted.mp3", w_size=2048, overlap=2, w_type="hann"):
-        # Загрузка аудиофайлов
-        y1, sr1 = self.load_audio(audio1_path)
-        y2, sr2 = self.load_audio(audio2_path)
-
-        if sr1 is None or sr2 is None:
-            raise Exception("Произошла ошибка при чтении файлов")
-
-        # Определяем количество каналов
-        channels1 = 1 if y1.ndim == 1 else y1.shape[0]
-        channels2 = 1 if y2.ndim == 1 else y2.shape[0]
-
-        # Преобразование в форму (samples, channels)
-        if channels1 > 1:
-            y1 = y1.T  # (channels, samples) -> (samples, channels)
-        else:
-            y1 = y1.reshape(-1, 1)
-
-        if channels2 > 1:
-            y2 = y2.T  # (channels, samples) -> (samples, channels)
-        else:
-            y2 = y2.reshape(-1, 1)
-
-        if sr1 != sr2:
-            if channels2 > 1:
-                # Ресемплинг для каждого канала отдельно
-                y2_resampled_list = []
-                for c in range(channels2):
-                    channel_resampled = librosa.resample(
-                        y2[:, c], orig_sr=sr2, target_sr=sr1
-                    )
-                    y2_resampled_list.append(channel_resampled)
-                
-                # Находим минимальную длину среди всех каналов
-                min_channel_length = min(len(ch) for ch in y2_resampled_list)
-                
-                # Обрезаем все каналы до одинаковой длины и собираем в массив
-                y2_resampled = np.zeros((min_channel_length, channels2), dtype=np.float32)
-                for c, channel in enumerate(y2_resampled_list):
-                    y2_resampled[:, c] = channel[:min_channel_length]
-                
-                y2 = y2_resampled
-            else:
-                y2 = librosa.resample(y2[:, 0], orig_sr=sr2, target_sr=sr1)
-                y2 = y2.reshape(-1, 1)
-            sr2 = sr1
-
-        # Приводим к одинаковой длине
-        min_len = min(len(y1), len(y2))
-        y1 = y1[:min_len]
-        y2 = y2[:min_len]
-
-        # Обрабатываем каждый канал отдельно
-        result_channels = []
-
-        # Если основной сигнал моно, а удаляемый стерео - преобразуем удаляемый в моно
-        if channels1 == 1 and channels2 > 1:
-            y2 = y2.mean(axis=1, keepdims=True)
-            channels2 = 1
-
-        for c in range(channels1):
-            # Выбираем канал для основного сигнала
-            y1_ch = y1[:, c]
-
-            # Выбираем канал для удаляемого сигнала
-            if channels2 == 1:
-                y2_ch = y2[:, 0]
-            else:
-                # Если каналов удаляемого сигнала больше, используем соответствующий канал
-                y2_ch = y2[:, min(c, channels2 - 1)]
-
-            # Обрабатываем канал
-            result_ch = self.process_channel(y1_ch, y2_ch, sr1, method, w_size=w_size, overlap=overlap, w_type=w_type)
-            result_channels.append(result_ch)
-
-        # Собираем каналы в один массив
-        if len(result_channels) > 1:
-            result = np.column_stack(result_channels)
-        else:
-            result = np.array(result_channels[0])
-
-        # Нормализация (предотвращение клиппинга)
-        if result.ndim > 1:
-            # Для многоканального аудио нормализуем каждый канал отдельно
-            for c in range(result.shape[1]):
-                channel = result[:, c]
-                max_val = np.max(np.abs(channel))
-                if max_val > 0:
-                    result[:, c] = channel * 0.9 / max_val
-        else:
-            max_val = np.max(np.abs(result))
-            if max_val > 0:
-                result = result * 0.9 / max_val
-
-        inverted = self.write(o=output_path, array=result.T, sr=sr1, of=out_format, br="320k")
-        return inverted
+import os
+from pathlib import Path
+import sys
+import json
+import subprocess
+import numpy as np
+from typing import Literal
+from collections.abc import Callable
+from pathlib import Path
+from numpy.typing import DTypeLike
+import tempfile
+import librosa
+
+if not __package__:
+    from namer import Namer
+else:
+    from .namer import Namer
+
+
+class NotInputFileSpecified(Exception):
+    pass
+
+
+class NotOutputFileSpecified(Exception):
+    pass
+
+
+class NotSupportedDataType(Exception):
+    pass
+
+
+class ErrorDecode(Exception):
+    pass
+
+
+class ErrorEncode(Exception):
+    pass
+
+
+class NotSupportedFormat(Exception):
+    pass
+
+
+class SampleRateError(Exception):
+    pass
+
+
+class FileIsNotAudio(Exception):
+    pass
+
+
+class Audio(Namer):
+    def __init__(self):
+        super().__init__()
+        self.ffmpeg_path = os.environ.get("MVSEPLESS_FFMPEG", "ffmpeg")
+        self.ffprobe_path = os.environ.get("MVSEPLESS_FFPROBE", "ffprobe")
+        self.output_formats = (
+            "mp3",
+            "wav",
+            "flac",
+            "ogg",
+            "opus",
+            "m4a",
+            "aac",
+            "ac3",
+            "aiff",
+        )
+        self.input_formats = (
+            "mp3",
+            "wav",
+            "flac",
+            "ogg",
+            "opus",
+            "m4a",
+            "aac",
+            "ac3",
+            "aiff",
+            "mp4",
+            "mkv",
+            "webm",
+            "avi",
+            "mov",
+            "ts",
+        )
+        self.supported_dtypes = ("int16", "int32", "float32", "float64")
+        self.dtypes_dict = {
+            "int16": "s16le",
+            "int32": "s32le",
+            "float32": "f32le",
+            "float64": "f64le",
+            np.int16: "s16le",
+            np.int32: "s32le",
+            np.float32: "f32le",
+            np.float64: "f64le",
+        }
+        self.bitrate_limit = {
+            "mp3": {"min": 8, "max": 320},
+            "aac": {"min": 8, "max": 512},
+            "m4a": {"min": 8, "max": 512},
+            "ac3": {"min": 32, "max": 640},
+            "ogg": {"min": 64, "max": 500},
+            "opus": {"min": 6, "max": 512},
+        }
+        self.sample_rates = {
+            "mp3": {
+                "supported": (
+                    44100,
+                    48000,
+                    32000,
+                    22050,
+                    24000,
+                    16000,
+                    11025,
+                    12000,
+                    8000,
+                )
+            },
+            "opus": {"supported": (48000, 24000, 16000, 12000, 8000)},
+            "m4a": {
+                "supported": (
+                    96000,
+                    88200,
+                    64000,
+                    48000,
+                    44100,
+                    32000,
+                    24000,
+                    22050,
+                    16000,
+                    12000,
+                    11025,
+                    8000,
+                    7350,
+                )
+            },
+            "aac": {
+                "supported": (
+                    96000,
+                    88200,
+                    64000,
+                    48000,
+                    44100,
+                    32000,
+                    24000,
+                    22050,
+                    16000,
+                    12000,
+                    11025,
+                    8000,
+                    7350,
+                )
+            },
+            "ac3": {
+                "supported": (
+                    48000,
+                    44100,
+                    32000,
+                )
+            },
+            "ogg": {"min": 6, "max": 192000},
+            "wav": {"min": 0, "max": float("inf")},
+            "aiff": {"min": 0, "max": float("inf")},
+            "flac": {"min": 0, "max": 192000},
+        }
+        self.check_ffmpeg()
+        self.check_ffprobe()
+
+    def check_ffmpeg(self):
+        try:
+            ffmpeg_version_output = subprocess.check_output(
+                [self.ffmpeg_path, "-version"], text=True
+            )
+        except FileNotFoundError:
+            if "PYTEST_CURRENT_TEST" not in os.environ:
+                raise FileNotFoundError(
+                    "FFMPEG не установлен. Укажите путь к установленному FFMPEG через переменную окружения MVSEPLESS_FFMPEG"
+                )
+
+    def check_ffprobe(self):
+        try:
+            ffmpeg_version_output = subprocess.check_output(
+                [self.ffprobe_path, "-version"], text=True
+            )
+        except FileNotFoundError:
+            if "PYTEST_CURRENT_TEST" not in os.environ:
+                raise FileNotFoundError(
+                    "FFPROBE не установлен. Укажите путь к установленному FFPROBE через переменную окружения MVSEPLESS_FFPROBE"
+                )
+
+    def fit_sr(
+        self,
+        f: (
+            str
+            | Literal["mp3", "wav", "flac", "ogg", "opus", "m4a", "aac", "ac3", "aiff"]
+        ) = "mp3",
+        sr: int = 44100,
+    ) -> int:
+        format_info = self.sample_rates.get(f.lower())
+
+        if not format_info:
+            return None
+
+        if "supported" in format_info:
+            supported_rates = format_info["supported"]
+            if sr in supported_rates:
+                return sr
+
+            return min(supported_rates, key=lambda x: abs(x - sr))
+
+        elif "min" in format_info and "max" in format_info:
+            min_rate = format_info["min"]
+            max_rate = format_info["max"]
+
+            if sr < min_rate:
+                return min_rate
+            elif sr > max_rate:
+                return max_rate
+            else:
+                return sr
+
+        return None
+
+    def fit_br(
+        self,
+        f: (
+            str
+            | Literal["mp3", "wav", "flac", "ogg", "opus", "m4a", "aac", "ac3", "aiff"]
+        ) = "mp3",
+        br: int = 320,
+    ) -> int:
+        if f not in self.bitrate_limit:
+            raise NotSupportedFormat(f"Формат {f} не поддерживается")
+
+        limits = self.bitrate_limit[f]
+
+        if br < limits["min"]:
+            return limits["min"]
+        elif br > limits["max"]:
+            return limits["max"]
+        else:
+            return br
+
+    def get_info(
+        self,
+        i: str | os.PathLike | Callable | None = None,
+    ) -> dict[int, dict[int, float]]:
+        audio_info = {}
+        if i:
+            if isinstance(i, Path):
+                i = str(i)
+            if os.path.exists(i):
+                cmd = [
+                    self.ffprobe_path,
+                    "-i",
+                    i,
+                    "-v",
+                    "quiet",
+                    "-hide_banner",
+                    "-show_entries",
+                    "stream=index,sample_rate,duration",
+                    "-select_streams",
+                    "a",
+                    "-of",
+                    "json",
+                ]
+
+                process = subprocess.Popen(
+                    cmd,
+                    stdin=subprocess.PIPE,
+                    stdout=subprocess.PIPE,
+                    stderr=subprocess.PIPE,
+                )
+
+                stdout, stderr = process.communicate()
+
+                if process.returncode != 0:
+                    print(f"STDERR: {stderr.decode('utf-8')}")
+                    print(f"STDOUT: {stdout.decode('utf-8')}")
+
+                json_output = json.loads(stdout)
+                streams = json_output["streams"]
+                if not streams:
+                    pass
+
+                else:
+                    for a, stream in enumerate(streams):
+                        audio_info[a] = {
+                            "sample_rate": int(stream.get("sample_rate", 0)),
+                            "duration": float(stream.get("duration", 0)),
+                        }
+
+                return audio_info
+
+            else:
+                raise FileExistsError("Указанного файла не существует")
+
+        else:
+            raise NotInputFileSpecified("Не указан путь к файлу")
+
+    def check(self, i: str | os.PathLike | Callable | None = None) -> bool:
+        if i:
+            if isinstance(i, Path):
+                i = str(i)
+            if os.path.exists(i):
+                info = self.get_info(i=i)
+                if info:
+                    list_streams = list(info.keys())
+                    if len(list_streams) > 0:
+                        if info[0].get("sample_rate") > 0:
+                            return True
+                        else:
+                            return False
+                    else:
+                        return False
+                else:
+                    return False
+            else:
+                raise FileExistsError("Указанного файла не существует")
+        else:
+            raise NotInputFileSpecified("Не указан путь к файлу")
+
+    def read(
+        self,
+        i: str | os.PathLike | Callable | None = None,
+        sr: int | None = None,
+        mono: bool = False,
+        dtype: DTypeLike = np.float32,
+        s: int = 0,
+    ) -> tuple[np.ndarray, int, float]:
+        output_format = self.dtypes_dict.get(dtype, None)
+        if not output_format:
+            raise NotSupportedDataType(f"Этот тип данных не поддерживается {dtype}")
+        if i:
+            if isinstance(i, Path):
+                i = str(i)
+            if os.path.exists(i):
+                audio_info = self.get_info(i=i)
+                list_streams = list(audio_info.keys())
+                if audio_info.get(s, False):
+                    stream = s
+                else:
+                    if len(list_streams) > 0:
+                        stream = 0
+                    else:
+                        raise FileIsNotAudio("В входном файле нет аудио потоков")
+
+                sample_rate_input = audio_info[stream]["sample_rate"]
+                if sample_rate_input == 0:
+                    raise FileIsNotAudio("В входном файле нет аудио потоков")
+
+                cmd = [
+                    self.ffmpeg_path,
+                    "-i",
+                    i,
+                    "-map",
+                    f"0:a:{stream}",
+                    "-vn",
+                    "-f",
+                    output_format,
+                    "-ac",
+                    "1" if mono else "2",
+                ]
+
+                if sr:
+                    cmd.extend(["-ar", str(sr)])
+                else:
+                    sr = sample_rate_input
+
+                cmd.append("pipe:1")
+
+                process = subprocess.Popen(
+                    cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, bufsize=10**8
+                )
+
+                try:
+
+                    raw_audio, stderr = process.communicate(timeout=300)
+
+                    if process.returncode != 0:
+                        raise ErrorDecode(f"FFmpeg error: {stderr.decode()}")
+
+                except subprocess.TimeoutExpired:
+                    process.kill()
+                    raise ErrorDecode("FFmpeg timeout при чтении файла")
+
+                audio_array = np.frombuffer(raw_audio, dtype=dtype)
+
+                channels = 1 if mono else 2
+                audio_array = audio_array.reshape((-1, channels)).T
+                if audio_array.ndim > 1 and channels == 1:
+                    audio_array = np.mean(
+                        audio_array, axis=tuple(range(audio_array.ndim - 1))
+                    )
+
+                len_samples = float(audio_array.shape[-1])
+
+                duration = len_samples / sr
+
+                print(f"Частота дискретизации: {sr}")
+
+                return audio_array.copy(), sr, duration
+            else:
+                raise FileExistsError("Указанного файла не существует")
+
+        else:
+            raise NotInputFileSpecified("Не указан путь к файлу")
+
+    def write(
+        self,
+        o: str | os.PathLike | Callable | None = None,
+        array: np.ndarray = np.array([], dtype=np.float32),
+        sr: int = 44100,
+        of: (
+            str
+            | Literal["mp3", "wav", "flac", "ogg", "opus", "m4a", "aac", "ac3", "aiff"]
+            | None
+        ) = None,
+        br: str | int | None = None,
+    ) -> str:
+        if isinstance(array, np.ndarray):
+
+            if len(array.shape) == 1:
+                array = array.reshape(-1, 1)
+            elif len(array.shape) == 2:
+                if array.shape[0] == 2:
+                    array = array.T
+            else:
+                raise ValueError(
+                    "numpy-массив должен быть либо одномерным, либо двухмерным"
+                )
+
+            if array.dtype == np.int16:
+                input_format = "s16le"
+            elif array.dtype == np.int32:
+                input_format = "s32le"
+            elif array.dtype == np.float32:
+                input_format = "f32le"
+            elif array.dtype == np.float64:
+                input_format = "f64le"
+            else:
+                raise NotSupportedDataType(
+                    f"Этот тип данных не поддерживается {array.dtype}"
+                )
+
+            if array.shape[1] == 1:
+                audio_bytes = array.tobytes()
+
+                channels = 1
+
+            elif array.shape[1] == 2:
+                audio_bytes = array.tobytes()
+
+                channels = 2
+            else:
+                raise ValueError("numpy-массив должен содержать 1 или 2 канала")
+
+        else:
+            raise ValueError("Вход должен быть numpy-массивом")
+
+        if o:
+            if isinstance(o, Path):
+                o = str(o)
+            output_dir = os.path.dirname(o)
+            output_base = os.path.basename(o)
+            output_name, output_ext = os.path.splitext(output_base)
+            if output_dir != "":
+                os.makedirs(output_dir, exist_ok=True)
+            if output_ext == "":
+                if of:
+                    o += f".{of}"
+                else:
+                    o += f".mp3"
+            elif output_ext == ".":
+                if of:
+                    o += f"{of}"
+                else:
+                    o += f"mp3"
+        else:
+            raise NotOutputFileSpecified("Не указан путь к выходному файлу")
+
+        if of:
+            if of in self.output_formats:
+                output_name, output_ext = os.path.splitext(o)
+                if output_ext == f".{of}":
+                    pass
+                else:
+                    o = f"{os.path.join(output_dir, output_name)}.{of}"
+            else:
+                raise NotSupportedFormat(f"Неподдерживаемый формат: {of}")
+        else:
+            of = os.path.splitext(o)[1].strip(".")
+            if of in self.output_formats:
+                pass
+            else:
+                raise NotSupportedFormat(f"Неподдерживаемый формат: {of}")
+
+        if sr:
+            if isinstance(sr, int):
+                sample_rate_fixed = self.fit_sr(f=of, sr=sr)
+            elif isinstance(sr, float):
+                sr = int(sr)
+                sample_rate_fixed = self.fit_sr(f=of, sr=sr)
+            else:
+                raise SampleRateError(
+                    f"Частота дискретизации должна быть числом\n\nЗначение: {sr}\nТип: {type(sr)}"
+                )
+        else:
+            raise SampleRateError("Не указана частота дискретизации")
+
+        bitrate_fixed = "320k"
+
+        if of not in ["wav", "flac", "aiff"]:
+            if br:
+                if isinstance(br, int):
+                    bitrate_fixed = self.fit_br(f=of, br=br)
+                elif isinstance(br, float):
+                    bitrate_fixed = self.fit_br(f=of, br=int(br))
+                elif isinstance(br, str):
+                    bitrate_fixed = self.fit_br(f=of, br=int(br.strip("k").strip("K")))
+                else:
+                    bitrate_fixed = self.fit_br(f=of, br=320)
+            else:
+                bitrate_fixed = self.fit_br(of, 320)
+
+        format_settings = {
+            "wav": [
+                "-c:a",
+                "pcm_f32le",
+                "-sample_fmt",
+                "flt",
+            ],
+            "aiff": [
+                "-c:a",
+                "pcm_f32be",
+                "-sample_fmt",
+                "flt",
+            ],
+            "flac": [
+                "-c:a",
+                "flac",
+                "-compression_level",
+                "12",
+                "-sample_fmt",
+                "s32",
+            ],
+            "mp3": [
+                "-c:a",
+                "libmp3lame",
+                "-b:a",
+                f"{bitrate_fixed}k",
+            ],
+            "ogg": [
+                "-c:a",
+                "libvorbis",
+                "-b:a",
+                f"{bitrate_fixed}k",
+            ],
+            "opus": [
+                "-c:a",
+                "libopus",
+                "-b:a",
+                f"{bitrate_fixed}k",
+            ],
+            "m4a": [
+                "-c:a",
+                "aac",
+                "-b:a",
+                f"{bitrate_fixed}k",
+            ],
+            "aac": [
+                "-c:a",
+                "aac",
+                "-b:a",
+                f"{bitrate_fixed}k",
+            ],
+            "ac3": [
+                "-c:a",
+                "ac3",
+                "-b:a",
+                f"{bitrate_fixed}k",
+            ],
+        }
+
+        cmd = [
+            self.ffmpeg_path,
+            "-y",
+            "-f",
+            input_format,
+            "-ar",
+            str(sr),
+            "-ac",
+            str(channels),
+            "-i",
+            "pipe:0",
+            "-ac",
+            str(channels),
+        ]
+
+        cmd.extend(["-ar", str(sample_rate_fixed)])
+        cmd.extend(format_settings[of])
+        o_dir, o_base = os.path.split(o)
+        o_base_n, o_base_ext = os.path.splitext(o_base)
+        o_base_n = self.sanitize(o_base_n)
+        o_base_n = self.short(o_base_n)
+        o = os.path.join(o_dir, f"{o_base_n}{o_base_ext}")
+        o = self.iter(o)
+        cmd.append(o)
+
+        process = subprocess.Popen(
+            cmd,
+            stdin=subprocess.PIPE,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+        )
+
+        try:
+            stdout, stderr = process.communicate(input=audio_bytes, timeout=300)
+        except subprocess.TimeoutExpired:
+            process.kill()
+            raise ErrorEncode("FFmpeg timeout: операция заняла слишком много времени")
+
+        if process.returncode != 0:
+            raise ErrorEncode(
+                f"FFmpeg завершился с ошибкой (код: {process.returncode})"
+            )
+
+        return os.path.abspath(o)
+
+
+class Inverter(Audio):
+    def __init__(self):
+        super().__init__()
+        self.test = "test"
+        self.w_types = [
+            "boxcar",
+            "triang",
+            "blackman",
+            "hamming",
+            "hann",
+            "bartlett",
+            "flattop",
+            "parzen",
+            "bohman",
+            "blackmanharris",
+            "nuttall",
+            "barthann",
+            "cosine",
+            "exponential",
+            "tukey",
+            "taylor",
+            "lanczos",
+        ]
+
+    def load_audio(self, filepath):
+        try:
+            y, sr, _ = self.read(i=filepath, sr=None, mono=False)
+            return y, sr
+        except Exception as e:
+            print(f"Ошибка загрузки аудио: {e}")
+            return None, None
+
+    def process_channel(
+        self, y1_ch, y2_ch, sr, method, w_size=2048, overlap=2, w_type="hann"
+    ):
+        HOP_LENGTH = w_size // overlap
+        if method == "waveform":
+            return y1_ch - y2_ch
+
+        elif method == "spectrogram":
+            S1 = librosa.stft(
+                y1_ch, n_fft=w_size, hop_length=HOP_LENGTH, win_length=w_size
+            )
+            S2 = librosa.stft(
+                y2_ch, n_fft=w_size, hop_length=HOP_LENGTH, win_length=w_size
+            )
+
+            mag1 = np.abs(S1)
+            mag2 = np.abs(S2)
+
+            mag_result = np.maximum(mag1 - mag2, 0)
+
+            phase = np.angle(S1)
+
+            S_result = mag_result * np.exp(1j * phase)
+
+            return librosa.istft(
+                S_result,
+                n_fft=w_size,
+                hop_length=HOP_LENGTH,
+                win_length=w_size,
+                length=len(y1_ch),
+            )
+
+    def process_audio(
+        self,
+        audio1_path,
+        audio2_path,
+        out_format,
+        method,
+        output_path="./inverted.mp3",
+        w_size=2048,
+        overlap=2,
+        w_type="hann",
+    ):
+        y1, sr1 = self.load_audio(audio1_path)
+        y2, sr2 = self.load_audio(audio2_path)
+
+        if sr1 is None or sr2 is None:
+            raise Exception("Произошла ошибка при чтении файлов")
+
+        channels1 = 1 if y1.ndim == 1 else y1.shape[0]
+        channels2 = 1 if y2.ndim == 1 else y2.shape[0]
+
+        if channels1 > 1:
+            y1 = y1.T
+        else:
+            y1 = y1.reshape(-1, 1)
+
+        if channels2 > 1:
+            y2 = y2.T
+        else:
+            y2 = y2.reshape(-1, 1)
+
+        if sr1 != sr2:
+            if channels2 > 1:
+                y2_resampled_list = []
+                for c in range(channels2):
+                    channel_resampled = librosa.resample(
+                        y2[:, c], orig_sr=sr2, target_sr=sr1
+                    )
+                    y2_resampled_list.append(channel_resampled)
+
+                min_channel_length = min(len(ch) for ch in y2_resampled_list)
+
+                y2_resampled = np.zeros(
+                    (min_channel_length, channels2), dtype=np.float32
+                )
+                for c, channel in enumerate(y2_resampled_list):
+                    y2_resampled[:, c] = channel[:min_channel_length]
+
+                y2 = y2_resampled
+            else:
+                y2 = librosa.resample(y2[:, 0], orig_sr=sr2, target_sr=sr1)
+                y2 = y2.reshape(-1, 1)
+            sr2 = sr1
+
+        min_len = min(len(y1), len(y2))
+        y1 = y1[:min_len]
+        y2 = y2[:min_len]
+
+        result_channels = []
+
+        if channels1 == 1 and channels2 > 1:
+            y2 = y2.mean(axis=1, keepdims=True)
+            channels2 = 1
+
+        for c in range(channels1):
+            y1_ch = y1[:, c]
+
+            if channels2 == 1:
+                y2_ch = y2[:, 0]
+            else:
+                y2_ch = y2[:, min(c, channels2 - 1)]
+
+            result_ch = self.process_channel(
+                y1_ch, y2_ch, sr1, method, w_size=w_size, overlap=overlap, w_type=w_type
+            )
+            result_channels.append(result_ch)
+
+        if len(result_channels) > 1:
+            result = np.column_stack(result_channels)
+        else:
+            result = np.array(result_channels[0])
+
+        if result.ndim > 1:
+            for c in range(result.shape[1]):
+                channel = result[:, c]
+                max_val = np.max(np.abs(channel))
+                if max_val > 0:
+                    result[:, c] = channel * 0.9 / max_val
+        else:
+            max_val = np.max(np.abs(result))
+            if max_val > 0:
+                result = result * 0.9 / max_val
+
+        inverted = self.write(
+            o=output_path, array=result.T, sr=sr1, of=out_format, br="320k"
+        )
+        return inverted

mvsepless/downloader.py

@@ -1,92 +1,90 @@
-import os
-import yt_dlp
-import time
-from tqdm import tqdm
-import urllib.request
-
-DOWNLOAD_DIR = os.environ.get(
-    "MVSEPLESS_DOWNLOAD_DIR", os.path.join(os.getcwd(), "downloaded")
-)
-
-def dw_file(url_model: str, local_path: str, retries: int = 180):
-    dir_name = os.path.dirname(local_path)
-    if dir_name != "":
-        os.makedirs(dir_name, exist_ok=True)
-    
-    class TqdmUpTo(tqdm):
-        def update_to(self, b=1, bsize=1, tsize=None):
-            if tsize is not None:
-                self.total = tsize
-            self.update(b * bsize - self.n)
-    
-    for attempt in range(retries):
-        try:
-            with TqdmUpTo(
-                unit="B",
-                unit_scale=True,
-                unit_divisor=1024,
-                miniters=1,
-                desc=os.path.basename(local_path),
-            ) as t:
-                urllib.request.urlretrieve(
-                    url_model, local_path, reporthook=t.update_to
-                )
-            # Если дошли сюда - загрузка успешна, прерываем цикл
-            break
-        except Exception as e:
-            print(f"Попытка {attempt + 1}/{retries} не удалась. Ошибка: {e}")
-            if attempt < retries - 1:  # Если это не последняя попытка
-                print("Повторная попытка...")
-                time.sleep(2)  # Небольшая задержка перед повторной попыткой
-            else:
-                print("Все попытки загрузки завершились неудачно")
-                raise  # Пробрасываем исключение дальше
-
-def dw_yt_dlp(
-    url,
-    output_dir=None,
-    cookie=None,
-    output_format="mp3",
-    output_bitrate="320",
-    title=None,
-):
-    # Подготовка шаблона имени файла
-    outtmpl = "%(title)s.%(ext)s" if title is None else f"{title}.%(ext)s"
-
-    ydl_opts = {
-        "format": "bestaudio/best",
-        "outtmpl": os.path.join(DOWNLOAD_DIR if not output_dir else output_dir, outtmpl),
-        "postprocessors": [
-            {
-                "key": "FFmpegExtractAudio",
-                "preferredcodec": output_format,
-                "preferredquality": output_bitrate,
-            }
-        ],
-        "noplaylist": True,  # Скачивать только одно видео, не плейлист
-        "quiet": True,  # Отключить вывод в консоль
-        "no_warnings": True,  # Скрыть предупреждения
-    }
-
-    # Добавляем cookies если указаны
-    if cookie and os.path.exists(cookie):
-        ydl_opts["cookiefile"] = cookie
-
-    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
-        try:
-            info = ydl.extract_info(url, download=True)
-            if "_type" in info and info["_type"] == "playlist":
-                # Для плейлистов берем первое видео
-                entry = info["entries"][0]
-                filename = ydl.prepare_filename(entry)
-            else:
-                # Для одиночного видео
-                filename = ydl.prepare_filename(info)
-
-            # Заменяем оригинальное расширение на выбранный формат
-            base, _ = os.path.splitext(filename)
-            audio_file = base + f".{output_format}"
-
-            return os.path.join(DOWNLOAD_DIR, audio_file)
-        except Exception as e:
-            return None
+import os
+import yt_dlp
+import time
+from tqdm import tqdm
+import urllib.request
+
+DOWNLOAD_DIR = os.environ.get(
+    "MVSEPLESS_DOWNLOAD_DIR", os.path.join(os.getcwd(), "downloaded")
+)
+
+
+def dw_file(url_model: str, local_path: str, retries: int = 180):
+    dir_name = os.path.dirname(local_path)
+    if dir_name != "":
+        os.makedirs(dir_name, exist_ok=True)
+
+    class TqdmUpTo(tqdm):
+        def update_to(self, b=1, bsize=1, tsize=None):
+            if tsize is not None:
+                self.total = tsize
+            self.update(b * bsize - self.n)
+
+    for attempt in range(retries):
+        try:
+            with TqdmUpTo(
+                unit="B",
+                unit_scale=True,
+                unit_divisor=1024,
+                miniters=1,
+                desc=os.path.basename(local_path),
+            ) as t:
+                urllib.request.urlretrieve(
+                    url_model, local_path, reporthook=t.update_to
+                )
+            break
+        except Exception as e:
+            print(f"Попытка {attempt + 1}/{retries} не удалась. Ошибка: {e}")
+            if attempt < retries - 1:
+                print("Повторная попытка...")
+                time.sleep(2)
+            else:
+                print("Все попытки загрузки завершились неудачно")
+                raise
+
+
+def dw_yt_dlp(
+    url,
+    output_dir=None,
+    cookie=None,
+    output_format="mp3",
+    output_bitrate="320",
+    title=None,
+):
+    outtmpl = "%(title)s.%(ext)s" if title is None else f"{title}.%(ext)s"
+
+    ydl_opts = {
+        "format": "bestaudio/best",
+        "outtmpl": os.path.join(
+            DOWNLOAD_DIR if not output_dir else output_dir, outtmpl
+        ),
+        "postprocessors": [
+            {
+                "key": "FFmpegExtractAudio",
+                "preferredcodec": output_format,
+                "preferredquality": output_bitrate,
+            }
+        ],
+        "noplaylist": True,
+        "quiet": True,
+        "no_warnings": True,
+    }
+
+    if cookie and os.path.exists(cookie):
+        ydl_opts["cookiefile"] = cookie
+
+    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
+        try:
+            info = ydl.extract_info(url, download=True)
+            if "_type" in info and info["_type"] == "playlist":
+                entry = info["entries"][0]
+                filename = ydl.prepare_filename(entry)
+            else:
+                filename = ydl.prepare_filename(info)
+
+            base, _ = os.path.splitext(filename)
+            audio_file = base + f".{output_format}"
+
+            return os.path.join(DOWNLOAD_DIR, audio_file)
+        except Exception as e:
+            return None

mvsepless/ensemble.py

@@ -1,224 +1,206 @@
-# coding: utf-8
-__author__ = "Roman Solovyev (ZFTurbo): https://github.com/ZFTurbo/"
-
-import os
-import sys
-import librosa
-import tempfile
-import numpy as np
-import argparse
-if not __package__:
-    from audio import Audio
-    from namer import Namer
-else:
-    from .audio import Audio
-    from .namer import Namer
-
-audio = Audio()
-
-def stft(wave, nfft, hl):
-    wave_left = np.asfortranarray(wave[0])
-    wave_right = np.asfortranarray(wave[1])
-    spec_left = librosa.stft(wave_left, n_fft=nfft, hop_length=hl)
-    spec_right = librosa.stft(wave_right, n_fft=nfft, hop_length=hl)
-    spec = np.asfortranarray([spec_left, spec_right])
-    return spec
-
-
-def istft(spec, hl, length):
-    spec_left = np.asfortranarray(spec[0])
-    spec_right = np.asfortranarray(spec[1])
-    wave_left = librosa.istft(spec_left, hop_length=hl, length=length)
-    wave_right = librosa.istft(spec_right, hop_length=hl, length=length)
-    wave = np.asfortranarray([wave_left, wave_right])
-    return wave
-
-
-def absmax(a, *, axis):
-    dims = list(a.shape)
-    dims.pop(axis)
-    indices = np.ogrid[tuple(slice(0, d) for d in dims)]
-    argmax = np.abs(a).argmax(axis=axis)
-    # Convert indices to list before insertion
-    indices = list(indices)
-    indices.insert(axis % len(a.shape), argmax)
-    return a[tuple(indices)]
-
-
-def absmin(a, *, axis):
-    dims = list(a.shape)
-    dims.pop(axis)
-    indices = np.ogrid[tuple(slice(0, d) for d in dims)]
-    argmax = np.abs(a).argmin(axis=axis)
-    indices.insert((len(a.shape) + axis) % len(a.shape), argmax)
-    return a[tuple(indices)]
-
-
-def lambda_max(arr, axis=None, key=None, keepdims=False):
-    idxs = np.argmax(key(arr), axis)
-    if axis is not None:
-        idxs = np.expand_dims(idxs, axis)
-        result = np.take_along_axis(arr, idxs, axis)
-        if not keepdims:
-            result = np.squeeze(result, axis=axis)
-        return result
-    else:
-        return arr.flatten()[idxs]
-
-
-def lambda_min(arr, axis=None, key=None, keepdims=False):
-    idxs = np.argmin(key(arr), axis)
-    if axis is not None:
-        idxs = np.expand_dims(idxs, axis)
-        result = np.take_along_axis(arr, idxs, axis)
-        if not keepdims:
-            result = np.squeeze(result, axis=axis)
-        return result
-    else:
-        return arr.flatten()[idxs]
-
-
-def average_waveforms(pred_track, weights, algorithm):
-    """
-    :param pred_track: shape = (num, channels, length)
-    :param weights: shape = (num, )
-    :param algorithm: One of avg_wave, median_wave, min_wave, max_wave, avg_fft, median_fft, min_fft, max_fft
-    :return: averaged waveform in shape (channels, length)
-    """
-
-    pred_track = np.array(pred_track)
-    final_length = pred_track.shape[-1]
-
-    mod_track = []
-    for i in range(pred_track.shape[0]):
-        if algorithm == "avg_wave":
-            mod_track.append(pred_track[i] * weights[i])
-        elif algorithm in ["median_wave", "min_wave", "max_wave"]:
-            mod_track.append(pred_track[i])
-        elif algorithm in ["avg_fft", "min_fft", "max_fft", "median_fft"]:
-            spec = stft(pred_track[i], nfft=2048, hl=1024)
-            if algorithm in ["avg_fft"]:
-                mod_track.append(spec * weights[i])
-            else:
-                mod_track.append(spec)
-    pred_track = np.array(mod_track)
-
-    if algorithm in ["avg_wave"]:
-        pred_track = pred_track.sum(axis=0)
-        pred_track /= np.array(weights).sum().T
-    elif algorithm in ["median_wave"]:
-        pred_track = np.median(pred_track, axis=0)
-    elif algorithm in ["min_wave"]:
-        pred_track = np.array(pred_track)
-        pred_track = lambda_min(pred_track, axis=0, key=np.abs)
-    elif algorithm in ["max_wave"]:
-        pred_track = np.array(pred_track)
-        pred_track = lambda_max(pred_track, axis=0, key=np.abs)
-    elif algorithm in ["avg_fft"]:
-        pred_track = pred_track.sum(axis=0)
-        pred_track /= np.array(weights).sum()
-        pred_track = istft(pred_track, 1024, final_length)
-    elif algorithm in ["min_fft"]:
-        pred_track = np.array(pred_track)
-        pred_track = lambda_min(pred_track, axis=0, key=np.abs)
-        pred_track = istft(pred_track, 1024, final_length)
-    elif algorithm in ["max_fft"]:
-        pred_track = np.array(pred_track)
-        pred_track = absmax(pred_track, axis=0)
-        pred_track = istft(pred_track, 1024, final_length)
-    elif algorithm in ["median_fft"]:
-        pred_track = np.array(pred_track)
-        pred_track = np.median(pred_track, axis=0)
-        pred_track = istft(pred_track, 1024, final_length)
-    return pred_track
-
-
-def ensemble_audio_files(
-    files, output="res.wav", ensemble_type="avg_wave", weights=None, out_format="wav", add_wav=False
-) -> str | tuple[str, str]:
-    """
-    Основная функция для объединения аудиофайлов
-
-    :param files: список путей к аудиофайлам
-    :param output: путь для сохранения результ��та
-    :param ensemble_type: метод объединения (avg_wave, median_wave, min_wave, max_wave, avg_fft, median_fft, min_fft, max_fft)
-    :param weights: список весов для каждого файла (None для равных весов)
-    :return: None
-    """
-    print("Алгоритм склеивания: {}".format(ensemble_type))
-    print("Количество входных файлов: {}".format(len(files)))
-    if weights is not None:
-        weights = np.array(weights)
-    else:
-        weights = np.ones(len(files))
-    print("Весы: {}".format(weights))
-    print("Имя выходного файла: {}".format(output))
-
-    data = []
-    sr = None
-    max_length = 0
-    max_channels = 0
-
-    # Первый проход: определяем максимальную длину и количество каналов
-    for f in files:
-        if not os.path.isfile(f):
-            print("Не удается найти файл: {}. Check paths.".format(f))
-            exit()
-        print("Читается файл: {}".format(f))
-        wav, current_sr, _ = audio.read(i=f, sr=None, mono=False)
-        if sr is None:
-            sr = current_sr
-        elif sr != current_sr:
-            print("Частота дискретизации на всех файлах должна быть одинаковой")
-            exit()
-
-        # Определяем количество каналов
-        if wav.ndim == 1:
-            channels = 1
-            length = len(wav)
-        else:
-            channels = wav.shape[0]
-            length = wav.shape[1]
-
-        max_length = max(max_length, length)
-        max_channels = max(max_channels, channels)
-        print("Форма сигнала: {} частота дискретизации: {}".format(wav.shape, sr))
-
-    # Второй проход: обработка и выравнивание файлов
-    for f in files:
-        wav, current_sr, _ = audio.read(i=f, sr=None, mono=False)
-
-        # Обработка каналов
-        if wav.ndim == 1:
-            # Моно -> стерео
-            wav = np.vstack([wav, wav])
-        elif wav.shape[0] == 1:
-            # Один канал -> стерео
-            wav = np.vstack([wav[0], wav[0]])
-        elif wav.shape[0] > 2:
-            # Более 2 каналов -> берем первые два
-            wav = wav[:2, :]
-
-        # Выравнивание длины
-        if wav.shape[1] < max_length:
-            pad_width = ((0, 0), (0, max_length - wav.shape[1]))
-            wav = np.pad(wav, pad_width, mode="constant")
-        elif wav.shape[1] > max_length:
-            wav = wav[:, :max_length]
-
-        data.append(wav)
-
-    data = np.array(data)
-    res = average_waveforms(data, weights, ensemble_type)
-    print("Форма результата: {}".format(res.shape))
-
-    output_wav = f"{output}_orig.wav"
-    output = f"{output}.{out_format}"
-
-    output = audio.write(o=output, array=res.T, sr=sr, of=out_format, br="320k")
-    if add_wav:
-        output_wav = audio.write(o=output_wav, array=res.T, sr=sr, of="wav")
-        return output, output_wav
-
-    else:
-        return output
+__author__ = "Roman Solovyev (ZFTurbo): https://github.com/ZFTurbo/"
+
+import os
+import sys
+import librosa
+import tempfile
+import numpy as np
+import argparse
+
+if not __package__:
+    from audio import Audio
+    from namer import Namer
+else:
+    from .audio import Audio
+    from .namer import Namer
+
+audio = Audio()
+
+
+def stft(wave, nfft, hl):
+    wave_left = np.asfortranarray(wave[0])
+    wave_right = np.asfortranarray(wave[1])
+    spec_left = librosa.stft(wave_left, n_fft=nfft, hop_length=hl)
+    spec_right = librosa.stft(wave_right, n_fft=nfft, hop_length=hl)
+    spec = np.asfortranarray([spec_left, spec_right])
+    return spec
+
+
+def istft(spec, hl, length):
+    spec_left = np.asfortranarray(spec[0])
+    spec_right = np.asfortranarray(spec[1])
+    wave_left = librosa.istft(spec_left, hop_length=hl, length=length)
+    wave_right = librosa.istft(spec_right, hop_length=hl, length=length)
+    wave = np.asfortranarray([wave_left, wave_right])
+    return wave
+
+
+def absmax(a, *, axis):
+    dims = list(a.shape)
+    dims.pop(axis)
+    indices = np.ogrid[tuple(slice(0, d) for d in dims)]
+    argmax = np.abs(a).argmax(axis=axis)
+    indices = list(indices)
+    indices.insert(axis % len(a.shape), argmax)
+    return a[tuple(indices)]
+
+
+def absmin(a, *, axis):
+    dims = list(a.shape)
+    dims.pop(axis)
+    indices = np.ogrid[tuple(slice(0, d) for d in dims)]
+    argmax = np.abs(a).argmin(axis=axis)
+    indices.insert((len(a.shape) + axis) % len(a.shape), argmax)
+    return a[tuple(indices)]
+
+
+def lambda_max(arr, axis=None, key=None, keepdims=False):
+    idxs = np.argmax(key(arr), axis)
+    if axis is not None:
+        idxs = np.expand_dims(idxs, axis)
+        result = np.take_along_axis(arr, idxs, axis)
+        if not keepdims:
+            result = np.squeeze(result, axis=axis)
+        return result
+    else:
+        return arr.flatten()[idxs]
+
+
+def lambda_min(arr, axis=None, key=None, keepdims=False):
+    idxs = np.argmin(key(arr), axis)
+    if axis is not None:
+        idxs = np.expand_dims(idxs, axis)
+        result = np.take_along_axis(arr, idxs, axis)
+        if not keepdims:
+            result = np.squeeze(result, axis=axis)
+        return result
+    else:
+        return arr.flatten()[idxs]
+
+
+def average_waveforms(pred_track, weights, algorithm):
+
+    pred_track = np.array(pred_track)
+    final_length = pred_track.shape[-1]
+
+    mod_track = []
+    for i in range(pred_track.shape[0]):
+        if algorithm == "avg_wave":
+            mod_track.append(pred_track[i] * weights[i])
+        elif algorithm in ["median_wave", "min_wave", "max_wave"]:
+            mod_track.append(pred_track[i])
+        elif algorithm in ["avg_fft", "min_fft", "max_fft", "median_fft"]:
+            spec = stft(pred_track[i], nfft=2048, hl=1024)
+            if algorithm in ["avg_fft"]:
+                mod_track.append(spec * weights[i])
+            else:
+                mod_track.append(spec)
+    pred_track = np.array(mod_track)
+
+    if algorithm in ["avg_wave"]:
+        pred_track = pred_track.sum(axis=0)
+        pred_track /= np.array(weights).sum().T
+    elif algorithm in ["median_wave"]:
+        pred_track = np.median(pred_track, axis=0)
+    elif algorithm in ["min_wave"]:
+        pred_track = np.array(pred_track)
+        pred_track = lambda_min(pred_track, axis=0, key=np.abs)
+    elif algorithm in ["max_wave"]:
+        pred_track = np.array(pred_track)
+        pred_track = lambda_max(pred_track, axis=0, key=np.abs)
+    elif algorithm in ["avg_fft"]:
+        pred_track = pred_track.sum(axis=0)
+        pred_track /= np.array(weights).sum()
+        pred_track = istft(pred_track, 1024, final_length)
+    elif algorithm in ["min_fft"]:
+        pred_track = np.array(pred_track)
+        pred_track = lambda_min(pred_track, axis=0, key=np.abs)
+        pred_track = istft(pred_track, 1024, final_length)
+    elif algorithm in ["max_fft"]:
+        pred_track = np.array(pred_track)
+        pred_track = absmax(pred_track, axis=0)
+        pred_track = istft(pred_track, 1024, final_length)
+    elif algorithm in ["median_fft"]:
+        pred_track = np.array(pred_track)
+        pred_track = np.median(pred_track, axis=0)
+        pred_track = istft(pred_track, 1024, final_length)
+    return pred_track
+
+
+def ensemble_audio_files(
+    files,
+    output="res.wav",
+    ensemble_type="avg_wave",
+    weights=None,
+    out_format="wav",
+    add_wav=False,
+) -> str | tuple[str, str]:
+    print("Алгоритм склеивания: {}".format(ensemble_type))
+    print("Количество входных файлов: {}".format(len(files)))
+    if weights is not None:
+        weights = np.array(weights)
+    else:
+        weights = np.ones(len(files))
+    print("Весы: {}".format(weights))
+    print("Имя выходного файла: {}".format(output))
+
+    data = []
+    sr = None
+    max_length = 0
+    max_channels = 0
+
+    for f in files:
+        if not os.path.isfile(f):
+            print("Не удается найти файл: {}. Check paths.".format(f))
+            exit()
+        print("Читается файл: {}".format(f))
+        wav, current_sr, _ = audio.read(i=f, sr=None, mono=False)
+        if sr is None:
+            sr = current_sr
+        elif sr != current_sr:
+            print("Частота дискретизации на всех файлах должна быть одинаковой")
+            exit()
+
+        if wav.ndim == 1:
+            channels = 1
+            length = len(wav)
+        else:
+            channels = wav.shape[0]
+            length = wav.shape[1]
+
+        max_length = max(max_length, length)
+        max_channels = max(max_channels, channels)
+        print("Форма сигнала: {} частота дискретизации: {}".format(wav.shape, sr))
+
+    for f in files:
+        wav, current_sr, _ = audio.read(i=f, sr=None, mono=False)
+
+        if wav.ndim == 1:
+            wav = np.vstack([wav, wav])
+        elif wav.shape[0] == 1:
+            wav = np.vstack([wav[0], wav[0]])
+        elif wav.shape[0] > 2:
+            wav = wav[:2, :]
+
+        if wav.shape[1] < max_length:
+            pad_width = ((0, 0), (0, max_length - wav.shape[1]))
+            wav = np.pad(wav, pad_width, mode="constant")
+        elif wav.shape[1] > max_length:
+            wav = wav[:, :max_length]
+
+        data.append(wav)
+
+    data = np.array(data)
+    res = average_waveforms(data, weights, ensemble_type)
+    print("Форма результата: {}".format(res.shape))
+
+    output_wav = f"{output}_orig.wav"
+    output = f"{output}.{out_format}"
+
+    output = audio.write(o=output, array=res.T, sr=sr, of=out_format, br="320k")
+    if add_wav:
+        output_wav = audio.write(o=output_wav, array=res.T, sr=sr, of="wav")
+        return output, output_wav
+
+    else:
+        return output

mvsepless/infer.py

@@ -20,17 +20,13 @@ from namer import Namer
 namer = Namer()
 audio = Audio()
 
-from infer_utils import (
-    prefer_target_instrument,
-    demix,
-    get_model_from_config
-)
+from infer_utils import prefer_target_instrument, demix, get_model_from_config
 
 
 def normalize_peak(audio, peak):
     current_peak = np.max(np.abs(audio))
     if current_peak == 0:
-        return audio  # избегаем деления на ноль
+        return audio
     scale_factor = peak / current_peak
     return audio * scale_factor
 
@@ -57,10 +53,18 @@ def cleanup_model(model):
             torch.cuda.ipc_collect()
 
         gc.collect()
-        sys.stdout.write(json.dumps({"cleanup": "Модель выгружена из памяти"}, ensure_ascii=False) + '\n')
+        sys.stdout.write(
+            json.dumps({"cleanup": "Модель выгружена из памяти"}, ensure_ascii=False)
+            + "\n"
+        )
         sys.stdout.flush()
     except Exception as e:
-        sys.stdout.write(json.dumps({"error": f"Ошибка при выгрузке модели: {str(e)}"}, ensure_ascii=False) + '\n')
+        sys.stdout.write(
+            json.dumps(
+                {"error": f"Ошибка при выгрузке модели: {str(e)}"}, ensure_ascii=False
+            )
+            + "\n"
+        )
         sys.stdout.flush()
 
 
@@ -87,22 +91,32 @@ def once_inference(
     model_id: int = 0,
 ):
     results = []
-    sys.stdout.write(json.dumps({"reading": path}, ensure_ascii=False) + '\n')
+    sys.stdout.write(json.dumps({"reading": path}, ensure_ascii=False) + "\n")
     sys.stdout.flush()
-    sys.stdout.write(json.dumps({"selected_stems": selected_instruments}, ensure_ascii=False) + '\n')
+    sys.stdout.write(
+        json.dumps({"selected_stems": selected_instruments}, ensure_ascii=False) + "\n"
+    )
     sys.stdout.flush()
-    sys.stdout.write(json.dumps({"stems": list(instruments)}, ensure_ascii=False) + '\n')
+    sys.stdout.write(
+        json.dumps({"stems": list(instruments)}, ensure_ascii=False) + "\n"
+    )
     sys.stdout.flush()
 
     if config.training.target_instrument is not None:
-        sys.stdout.write(json.dumps({"target_instrument": config.training.target_instrument}, ensure_ascii=False) + '\n')
+        sys.stdout.write(
+            json.dumps(
+                {"target_instrument": config.training.target_instrument},
+                ensure_ascii=False,
+            )
+            + "\n"
+        )
         sys.stdout.flush()
 
     try:
         mix, sr, _ = audio.read(i=path, sr=sample_rate, mono=False)
     except Exception as e:
         error_msg = f"Не удалось прочитать аудио: {path}\nОшибка: {e}"
-        sys.stdout.write(json.dumps({"error": error_msg}, ensure_ascii=False) + '\n')
+        sys.stdout.write(json.dumps({"error": error_msg}, ensure_ascii=False) + "\n")
         sys.stdout.flush()
         return results
 
@@ -128,13 +142,19 @@ def once_inference(
 
             full_result.append(waveforms)
         except Exception as e:
-            sys.stdout.write(json.dumps({"error": f"Ошибка при демиксе: {e}"}, ensure_ascii=False) + '\n')
+            sys.stdout.write(
+                json.dumps({"error": f"Ошибка при демиксе: {e}"}, ensure_ascii=False)
+                + "\n"
+            )
             sys.stdout.flush()
         del m
         gc.collect()
 
     if not full_result:
-        sys.stdout.write(json.dumps({"error": "Пустой результат демикса."}, ensure_ascii=False) + '\n')
+        sys.stdout.write(
+            json.dumps({"error": "Пустой результат демикса."}, ensure_ascii=False)
+            + "\n"
+        )
         sys.stdout.flush()
         return results
 
@@ -151,9 +171,7 @@ def once_inference(
     for el in waveforms:
         waveforms[el] /= len(full_result)
 
-    if (
-        extract_instrumental and config.training.target_instrument is not None
-    ):  # Если включен "Extract Instrumental / Извлечь инструментал" и найден целевой инструмент
+    if extract_instrumental and config.training.target_instrument is not None:
         second_stem = [
             s
             for s in config.training.instruments
@@ -169,8 +187,7 @@ def once_inference(
         extract_instrumental
         and selected_instruments
         and config.training.target_instrument is None
-    ):  # Если включен "Extract Instrumental / Извлечь инструментал" и выбраны инструменты, то создаются стемы "inverted -" и "inverted +" (если не найден целевого инструмент)
-
+    ):
 
         all_instruments = config.training.instruments
         if len(all_instruments) > 2:
@@ -178,21 +195,19 @@ def once_inference(
             waveforms["inverted -"] = mix_orig.copy()
             for instr in instruments:
                 if instr in waveforms:
-                    waveforms["inverted -"] -= waveforms[
-                        instr
-                    ]  # стем "inverted -": вычитание выбранного стема из оригинального сигнала (не всегда хорошо)
+                    waveforms["inverted -"] -= waveforms[instr]
 
             if "inverted -" not in instruments:
                 instruments.append("inverted -")
 
-            unselected_stems = [s for s in all_instruments if s not in selected_instruments]
+            unselected_stems = [
+                s for s in all_instruments if s not in selected_instruments
+            ]
             if unselected_stems:
                 waveforms["inverted +"] = np.zeros_like(mix_orig)
                 for stem in unselected_stems:
                     if stem in waveforms:
-                        waveforms["inverted +"] += waveforms[
-                            stem
-                        ]  # стем "inverted +": сложение не выбранных инструментов в один стем
+                        waveforms["inverted +"] += waveforms[stem]
                 if "inverted +" not in instruments:
                     instruments.append("inverted +")
 
@@ -216,9 +231,7 @@ def once_inference(
     ):
 
         waveforms["instrumental -"] = mix_orig.copy()
-        waveforms["instrumental -"] -= waveforms[
-            "vocals"
-        ]  # стем "inverted -": вычитание выбранного стема из оригинального сигнала (не всегда хорошо)
+        waveforms["instrumental -"] -= waveforms["vocals"]
 
         if "instrumental -" not in instruments:
             instruments.append("instrumental -")
@@ -229,9 +242,7 @@ def once_inference(
             waveforms["instrumental +"] = np.zeros_like(mix_orig)
             for stem in non_vocal_stems:
                 if stem in waveforms:
-                    waveforms["instrumental +"] += waveforms[
-                        stem
-                    ]  # стем "inverted +": сложение не выбранных инструментов в один стем
+                    waveforms["instrumental +"] += waveforms[stem]
             if "instrumental +" not in instruments:
                 instruments.append("instrumental +")
 
@@ -249,25 +260,40 @@ def once_inference(
                 estimates = estimates * std + mean
 
             file_name = os.path.splitext(os.path.basename(path))[0]
-            file_name_shorted = namer.short_input_name_template(template, STEM=instr, MODEL=model_name, ID=model_id, NAME=file_name)
+            file_name_shorted = namer.short_input_name_template(
+                template, STEM=instr, MODEL=model_name, ID=model_id, NAME=file_name
+            )
             custom_name = namer.template(
-                template, STEM=instr, MODEL=model_name, ID=model_id, NAME=file_name_shorted
+                template,
+                STEM=instr,
+                MODEL=model_name,
+                ID=model_id,
+                NAME=file_name_shorted,
             )
             output_path = os.path.join(store_dir, f"{custom_name}.{output_format}")
 
-            sys.stdout.write(json.dumps({"writing": output_path}, ensure_ascii=False) + '\n')
+            sys.stdout.write(
+                json.dumps({"writing": output_path}, ensure_ascii=False) + "\n"
+            )
             sys.stdout.flush()
 
             output_path = audio.write(
-                o=output_path, array=estimates, sr=sr, of=output_format, br=output_bitrate
-            )  # запись стема в аудио файл с помощью универсальной функции
+                o=output_path,
+                array=estimates,
+                sr=sr,
+                of=output_format,
+                br=output_bitrate,
+            )
 
-            results.append(
-                (instr, output_path)
-            )  # запись информации о разделении: (название стема, путь к файлу)
+            results.append((instr, output_path))
             del estimates
         except Exception as e:
-            sys.stdout.write(json.dumps({"error": f"Ошибка при обработке {instr}: {e}"}, ensure_ascii=False) + '\n')
+            sys.stdout.write(
+                json.dumps(
+                    {"error": f"Ошибка при обработке {instr}: {e}"}, ensure_ascii=False
+                )
+                + "\n"
+            )
             sys.stdout.flush()
         gc.collect()
 
@@ -307,7 +333,15 @@ def run_inference(
     instruments = prefer_target_instrument(config)
 
     if config.training.target_instrument is not None:
-        sys.stdout.write(json.dumps({"info": "Целевой инструмент найден в конфигурации модели. Выбранные стемы будут проигнорированы."}, ensure_ascii=False) + '\n')
+        sys.stdout.write(
+            json.dumps(
+                {
+                    "info": "Целевой инструмент найден в конфигурации модели. Выбранные стемы будут проигнорированы."
+                },
+                ensure_ascii=False,
+            )
+            + "\n"
+        )
         sys.stdout.flush()
     else:
         if selected_instruments is not None and selected_instruments != []:
@@ -315,7 +349,10 @@ def run_inference(
                 instr for instr in instruments if instr in selected_instruments
             ]
             if verbose:
-                sys.stdout.write(json.dumps({"selected_stems": instruments}, ensure_ascii=False) + '\n')
+                sys.stdout.write(
+                    json.dumps({"selected_stems": instruments}, ensure_ascii=False)
+                    + "\n"
+                )
                 sys.stdout.flush()
 
     os.makedirs(store_dir, exist_ok=True)
@@ -345,9 +382,11 @@ def run_inference(
 
     time.sleep(1)
     time_taken = time.time() - start_time
-    sys.stdout.write(json.dumps({"time": f"{time_taken:.2f} сек."}, ensure_ascii=False) + '\n')
+    sys.stdout.write(
+        json.dumps({"time": f"{time_taken:.2f} сек."}, ensure_ascii=False) + "\n"
+    )
     sys.stdout.flush()
-    sys.stdout.write(json.dumps({"done": results}, ensure_ascii=False) + '\n')
+    sys.stdout.write(json.dumps({"done": results}, ensure_ascii=False) + "\n")
     sys.stdout.flush()
     return results
 
@@ -357,7 +396,15 @@ def load_model(model_type, config_path, start_check_point, device_ids, force_cpu
     if force_cpu:
         device = "cpu"
     elif torch.cuda.is_available():
-        sys.stdout.write(json.dumps({"info": "Разделение выполняется на ядрах CUDA. Для выполнения на процессоре установите force_cpu=True."}, ensure_ascii=False) + '\n')
+        sys.stdout.write(
+            json.dumps(
+                {
+                    "info": "Разделение выполняется на ядрах CUDA. Для выполнения на процессоре установите force_cpu=True."
+                },
+                ensure_ascii=False,
+            )
+            + "\n"
+        )
         sys.stdout.flush()
         device = "cuda"
 
@@ -372,7 +419,7 @@ def load_model(model_type, config_path, start_check_point, device_ids, force_cpu
     elif torch.backends.mps.is_available():
         device = "mps"
 
-    sys.stdout.write(json.dumps({"device": device}, ensure_ascii=False) + '\n')
+    sys.stdout.write(json.dumps({"device": device}, ensure_ascii=False) + "\n")
     sys.stdout.flush()
 
     model_load_start_time = time.time()
@@ -382,29 +429,30 @@ def load_model(model_type, config_path, start_check_point, device_ids, force_cpu
 
     if model_type == "vr":
         model.load_checkpoint(start_check_point, device)
-        model.settings(enable_post_process=False, 
-            post_process_threshold=config.inference.post_process_threshold, 
-            batch_size=config.inference.batch_size, 
-            window_size=config.inference.window_size, 
-            high_end_process=config.inference.high_end_process, 
+        model.settings(
+            enable_post_process=False,
+            post_process_threshold=config.inference.post_process_threshold,
+            batch_size=config.inference.batch_size,
+            window_size=config.inference.window_size,
+            high_end_process=config.inference.high_end_process,
             primary_stem=config.training.instruments[0],
-            secondary_stem=config.training.instruments[1]
+            secondary_stem=config.training.instruments[1],
         )
         return model, config, device
 
     elif model_type == "mdxnet":
         if start_check_point != "":
-            sys.stdout.write(json.dumps({"checkpoint": start_check_point}) + '\n')
+            sys.stdout.write(json.dumps({"checkpoint": start_check_point}) + "\n")
             sys.stdout.flush()
             model.init_onnx_session(start_check_point, device)
 
         return model, config, device
-    
+
     else:
         if start_check_point != "":
-            sys.stdout.write(json.dumps({"checkpoint": start_check_point}) + '\n')
+            sys.stdout.write(json.dumps({"checkpoint": start_check_point}) + "\n")
             sys.stdout.flush()
-            
+
             if model_type in ["htdemucs", "apollo"]:
                 state_dict = torch.load(
                     start_check_point, map_location=device, weights_only=False
@@ -428,7 +476,10 @@ def load_model(model_type, config_path, start_check_point, device_ids, force_cpu
             except RuntimeError:
                 model.load_state_dict(state_dict, strict=False)
 
-        sys.stdout.write(json.dumps({"stems": list(config.training.instruments)}, ensure_ascii=False) + '\n')
+        sys.stdout.write(
+            json.dumps({"stems": list(config.training.instruments)}, ensure_ascii=False)
+            + "\n"
+        )
         sys.stdout.flush()
 
         if (
@@ -443,7 +494,10 @@ def load_model(model_type, config_path, start_check_point, device_ids, force_cpu
 
         load_time = time.time() - model_load_start_time
 
-        sys.stdout.write(json.dumps({"model_load_time": f"{load_time:.2f} сек."}, ensure_ascii=False) + '\n')
+        sys.stdout.write(
+            json.dumps({"model_load_time": f"{load_time:.2f} сек."}, ensure_ascii=False)
+            + "\n"
+        )
         sys.stdout.flush()
 
         return model, config, device
@@ -503,13 +557,11 @@ def parse_args():
         description="Модифицированный Music-Source-Separation-Training для разделения аудио на источники"
     )
 
-    # Обязательные аргументы
     parser.add_argument("--input", type=str, help="Путь к входному файлу или папке")
     parser.add_argument(
         "--store_dir", type=str, required=True, help="Путь для сохранения результатов"
     )
 
-    # Основные параметры модели
     parser.add_argument(
         "--model_type",
         type=str,
@@ -523,7 +575,7 @@ def parse_args():
             "bandit",
             "bandit_v2",
             "mdxnet",
-            "vr"
+            "vr",
         ],
         help="Тип модели (по умолчанию: htdemucs)",
     )
@@ -537,7 +589,6 @@ def parse_args():
         "--start_check_point", type=str, required=True, help="Путь к чекпоинту модели"
     )
 
-    # Параметры вывода
     parser.add_argument(
         "--output_format",
         type=str,
@@ -620,4 +671,4 @@ def main():
 
 
 if __name__ == "__main__":
-    main()
+    main()

mvsepless/infer_utils.py

@@ -1,4 +1,3 @@
-# coding: utf-8
 __author__ = "Roman Solovyev (ZFTurbo): https://github.com/ZFTurbo/"
 
 import sys
@@ -15,9 +14,6 @@ from typing import Dict, List, Tuple, Any, List, Optional
 
 
 def load_config(model_type: str, config_path: str) -> Any:
-    """
-    Load the configuration from the specified path based on the model type.
-    """
     try:
         with open(config_path, "r") as f:
             if model_type == "htdemucs":
@@ -32,9 +28,6 @@ def load_config(model_type: str, config_path: str) -> Any:
 
 
 def get_model_from_config(model_type: str, config_path: str) -> Tuple:
-    """
-    Load the model specified by the model type and configuration file.
-    """
     config = load_config(model_type, config_path)
 
     if model_type == "mdx23c":
@@ -47,11 +40,9 @@ def get_model_from_config(model_type: str, config_path: str) -> Tuple:
 
         model = MDXNet(**dict(config.model))
 
-    # В функции get_model_from_config добавьте:
-
     elif model_type == "vr":
         from models.vr_arch import VRNet
-        # Передаем instruments из config.training в модель
+
         model = VRNet(**dict(config.model))
 
     elif model_type == "htdemucs":
@@ -60,7 +51,7 @@ def get_model_from_config(model_type: str, config_path: str) -> Tuple:
         model = get_model(config)
 
     elif model_type == "mel_band_roformer":
-        if hasattr(config, "windowed"): # Это не нарушает совместимость со обычными моделями на Mel-Band Roformer 
+        if hasattr(config, "windowed"):
             from models.windowed_roformer.model import MelBandRoformerWSA
 
             model = MelBandRoformerWSA(**dict(config.model))
@@ -114,10 +105,8 @@ def get_model_from_config(model_type: str, config_path: str) -> Tuple:
 
     return model, config
 
+
 def _getWindowingArray(window_size: int, fade_size: int) -> torch.Tensor:
-    """
-    Generate a windowing array with a linear fade-in at the beginning and a fade-out at the end.
-    """
     fadein = torch.linspace(0, 1, fade_size)
     fadeout = torch.linspace(1, 0, fade_size)
 
@@ -126,6 +115,7 @@ def _getWindowingArray(window_size: int, fade_size: int) -> torch.Tensor:
     window[:fade_size] = fadein
     return window
 
+
 def demix_mdxnet(
     config: Any,
     model: Any,
@@ -133,18 +123,17 @@ def demix_mdxnet(
     device: torch.device,
     pbar: bool = False,
 ) -> Dict[str, np.ndarray]:
-    """
-    MDX-Net specific demixing function с поддержкой overlap
-    """
     mix_tensor = torch.tensor(mix, dtype=torch.float32)
     inv_mix_tensor = torch.tensor(-mix, dtype=torch.float32)
-    
+
     num_overlap = config.inference.num_overlap
     denoise = config.inference.denoise
     stem_name = model.primary_stem
     if denoise:
         processed_wav = model.process_wave(mix_tensor, device, num_overlap, pbar=pbar)
-        inv_processed_wav = model.process_wave(inv_mix_tensor, device, num_overlap, pbar=pbar)
+        inv_processed_wav = model.process_wave(
+            inv_mix_tensor, device, num_overlap, pbar=pbar
+        )
         result = processed_wav.cpu().numpy()
         inv_result = inv_processed_wav.cpu().numpy()
         result_separation = (result + -inv_result) * 0.5
@@ -152,9 +141,12 @@ def demix_mdxnet(
         processed_wav = model.process_wave(mix_tensor, device, num_overlap, pbar=pbar)
         result_separation = processed_wav.cpu().numpy()
 
-    result_separation = np.nan_to_num(result_separation, nan=0.0, posinf=0.0, neginf=0.0)
+    result_separation = np.nan_to_num(
+        result_separation, nan=0.0, posinf=0.0, neginf=0.0
+    )
+
+    return {stem_name: result_separation}
 
-    return {stem_name: result_separation}  # Перемещаем на CPU для возврата
 
 def demix_vr(
     config: Any,
@@ -163,12 +155,10 @@ def demix_vr(
     device: torch.device,
     pbar: bool = False,
 ) -> Dict[str, np.ndarray]:
-    """
-    VR-specific demixing function that processes the entire audio at once
-    since VR architecture doesn't support chunk-based processing
-    """
-    # Convert to tensor and add batch dimension
-    return model.demix(mix, config.audio.sample_rate, device, config.inference.aggression)
+    return model.demix(
+        mix, config.audio.sample_rate, device, config.inference.aggression
+    )
+
 
 def demix_demucs(config, model, mix, device, pbar=False):
     mix = torch.tensor(mix, dtype=torch.float32)
@@ -176,8 +166,8 @@ def demix_demucs(config, model, mix, device, pbar=False):
     num_instruments = len(config.training.instruments)
     num_overlap = config.inference.num_overlap
     step = chunk_size // num_overlap
-    fade_size = chunk_size // 10  # Добавляем fade_size для оконной функции
-    windowing_array = _getWindowingArray(chunk_size, fade_size)  # Создаём окно
+    fade_size = chunk_size // 10
+    windowing_array = _getWindowingArray(chunk_size, fade_size)
 
     batch_size = config.inference.batch_size
     use_amp = getattr(config.training, "use_amp", True)
@@ -209,9 +199,9 @@ def demix_demucs(config, model, mix, device, pbar=False):
                     x = model(arr)
 
                     window = windowing_array.clone()
-                    if i - step == 0:  # Первый чанк, без fade-in
+                    if i - step == 0:
                         window[:fade_size] = 1
-                    elif i >= mix.shape[1]:  # Последний чанк, без fade-out
+                    elif i >= mix.shape[1]:
                         window[-fade_size:] = 1
 
                     for j, (start, seg_len) in enumerate(batch_locations):
@@ -220,10 +210,14 @@ def demix_demucs(config, model, mix, device, pbar=False):
                         )
                         counter[..., start : start + seg_len] += window[..., :seg_len]
 
-                    # Output progress
                     processed = min(i, mix.shape[1])
                     total = mix.shape[1]
-                    sys.stdout.write(json.dumps({"processing": {"processed": processed, "total": total}}) + '\n')
+                    sys.stdout.write(
+                        json.dumps(
+                            {"processing": {"processed": processed, "total": total}}
+                        )
+                        + "\n"
+                    )
                     sys.stdout.flush()
 
                     batch_data.clear()
@@ -239,6 +233,7 @@ def demix_demucs(config, model, mix, device, pbar=False):
         instruments = config.training.instruments
         return {k: v for k, v in zip(instruments, estimated_sources)}
 
+
 def demix_generic(
     config: ConfigDict,
     model: torch.nn.Module,
@@ -246,9 +241,6 @@ def demix_generic(
     device: torch.device,
     pbar: bool = False,
 ) -> Dict[str, np.ndarray]:
-    """
-    Generic demixing function for models that support chunk-based processing
-    """
     mix = torch.tensor(mix, dtype=torch.float32)
     chunk_size = config.audio.chunk_size
     instruments = prefer_target_instrument(config)
@@ -260,8 +252,7 @@ def demix_generic(
     border = chunk_size - step
     length_init = mix.shape[-1]
     windowing_array = _getWindowingArray(chunk_size, fade_size)
-    
-    # Add padding to handle edge artifacts
+
     if length_init > 2 * border and border > 0:
         mix = nn.functional.pad(mix, (border, border), mode="reflect")
 
@@ -270,7 +261,6 @@ def demix_generic(
 
     with torch.cuda.amp.autocast(enabled=use_amp):
         with torch.inference_mode():
-            # Initialize result and counter tensors
             req_shape = (num_instruments,) + mix.shape
             result = torch.zeros(req_shape, dtype=torch.float32)
             counter = torch.zeros(req_shape, dtype=torch.float32)
@@ -280,10 +270,9 @@ def demix_generic(
             batch_locations = []
 
             while i < mix.shape[1]:
-                # Extract chunk and apply padding if necessary
                 part = mix[:, i : i + chunk_size].to(device)
                 chunk_len = part.shape[-1]
-                
+
                 pad_mode = "reflect" if chunk_len > chunk_size // 2 else "constant"
                 part = nn.functional.pad(
                     part, (0, chunk_size - chunk_len), mode=pad_mode, value=0
@@ -293,15 +282,14 @@ def demix_generic(
                 batch_locations.append((i, chunk_len))
                 i += step
 
-                # Process batch if it's full or the end is reached
                 if len(batch_data) >= batch_size or i >= mix.shape[1]:
                     arr = torch.stack(batch_data, dim=0)
                     x = model(arr)
 
                     window = windowing_array.clone()
-                    if i - step == 0:  # First audio chunk, no fadein
+                    if i - step == 0:
                         window[:fade_size] = 1
-                    elif i >= mix.shape[1]:  # Last audio chunk, no fadeout
+                    elif i >= mix.shape[1]:
                         window[-fade_size:] = 1
 
                     for j, (start, seg_len) in enumerate(batch_locations):
@@ -310,27 +298,30 @@ def demix_generic(
                         )
                         counter[..., start : start + seg_len] += window[..., :seg_len]
 
-                    # Output progress
                     processed = min(i, mix.shape[1])
                     total = mix.shape[1]
-                    sys.stdout.write(json.dumps({"processing": {"processed": processed, "total": total}}, ensure_ascii=False) + '\n')
+                    sys.stdout.write(
+                        json.dumps(
+                            {"processing": {"processed": processed, "total": total}},
+                            ensure_ascii=False,
+                        )
+                        + "\n"
+                    )
                     sys.stdout.flush()
 
                     batch_data.clear()
                     batch_locations.clear()
 
-            # Compute final estimated sources
             estimated_sources = result / counter
             estimated_sources = estimated_sources.cpu().numpy()
             np.nan_to_num(estimated_sources, copy=False, nan=0.0)
 
-            # Remove padding
             if length_init > 2 * border and border > 0:
                 estimated_sources = estimated_sources[..., border:-border]
 
-    # Return the result as a dictionary
     return {k: v for k, v in zip(instruments, estimated_sources)}
 
+
 def demix(
     config: ConfigDict,
     model: torch.nn.Module,
@@ -339,28 +330,17 @@ def demix(
     model_type: str,
     pbar: bool = False,
 ) -> Dict[str, np.ndarray]:
-    """
-    Unified function for audio source separation with support for multiple processing modes.
-    """
-    # Handle different model types
     if model_type == "vr":
         return demix_vr(config, model, mix, device, pbar)
     elif model_type == "mdxnet":
         return demix_mdxnet(config, model, mix, device, pbar)
     elif model_type == "htdemucs":
-        # HTDemucs uses its own processing
         return demix_demucs(config, model, mix, device, pbar)
     else:
-        # Generic processing for other models
         return demix_generic(config, model, mix, device, pbar)
 
 
 def prefer_target_instrument(config: ConfigDict) -> List[str]:
-    """
-    Return the list of target instruments based on the configuration.
-    If a specific target instrument is specified in the configuration,
-    it returns a list with that instrument. Otherwise, it returns the list of instruments.
-    """
     if config.training.get("target_instrument"):
         return [config.training.target_instrument]
     else:
@@ -370,21 +350,13 @@ def prefer_target_instrument(config: ConfigDict) -> List[str]:
 def prefer_target_instrument_test(
     config: ConfigDict, selected_instruments: Optional[List[str]] = None
 ) -> List[str]:
-    """
-    Return the list of target instruments based on the configuration and selected instruments.
-    If selected_instruments is specified, returns the intersection with available instruments.
-    Otherwise, if a target instrument is specified, returns it, else returns all instruments.
-    """
     available_instruments = config.training.instruments
 
     if selected_instruments is not None:
-        # Return only selected instruments that are available
         return [
             instr for instr in selected_instruments if instr in available_instruments
         ]
     elif config.training.get("target_instrument"):
-        # Default behavior if no selection - return target instrument
         return [config.training.target_instrument]
     else:
-        # If no target and no selection, return all instruments
-        return available_instruments
+        return available_instruments

mvsepless/model_manager.py

@@ -1,609 +1,682 @@
-import os
-import sys
-import json
-import yaml
-from tabulate import tabulate
-import shutil
-from tqdm import tqdm
-import urllib.request
-import gdown
-import requests
-import zipfile
-import tempfile
-import secrets
-import string
-import argparse
-from typing import Dict, Any
-script_dir = os.path.dirname(os.path.abspath(__file__))
-if not __package__:
-    from downloader import dw_file
-else:
-    from .downloader import dw_file
-
-def generate_secure_random(length=10):
-    """Генерирует криптографически безопасную случайную строку"""
-    characters = string.ascii_letters + string.digits
-    return ''.join(secrets.choice(characters) for _ in range(length))
-
-class MvseplessModelManager:
-    def __init__(
-        self,
-        models_info_path=os.path.join(script_dir, "models.json"),
-        cache_dir=os.path.join(script_dir, "mvsepless_models_cache"),
-    ):
-        self.models_cache_dir = cache_dir
-        self.models_info_path = models_info_path
-        with open(self.models_info_path, "r", encoding="utf-8") as f:
-            models_info = json.load(f)
-        self.models_info = models_info
-
-    def get_mt(self):
-        return list(self.models_info.keys())
-
-    def get_mn(self, model_type):
-        try:
-            mt = self.models_info.get(model_type, None)
-            if mt:
-                return list(self.models_info[model_type].keys())
-            return []
-        except (KeyError, TypeError):
-            return []
-
-    def get_stems(self, model_type, model_name):
-        try:
-            mt = self.models_info.get(model_type, None)
-            if mt:
-                mn = self.models_info[model_type].get(model_name, None)
-                if mn and "stems" in self.models_info[model_type][model_name]:
-                    return self.models_info[model_type][model_name]["stems"]
-            return []
-        except (KeyError, TypeError):
-            return []
-
-    def get_id(self, model_type, model_name):
-        try:
-            mt = self.models_info.get(model_type, None)
-            if mt:
-                mn = self.models_info[model_type].get(model_name, None)
-                if mn and "id" in self.models_info[model_type][model_name]:
-                    return self.models_info[model_type][model_name]["id"]
-            return 0
-        except (KeyError, TypeError):
-            return 0
-
-    def get_tgt_inst(self, model_type, model_name):
-        try:
-            mt = self.models_info.get(model_type, None)
-            if mt:
-                mn = self.models_info[model_type].get(model_name, None)
-                if mn and "target_instrument" in self.models_info[model_type][model_name]:
-                    return self.models_info[model_type][model_name]["target_instrument"]
-            return None
-        except (KeyError, TypeError):
-            return None
-
-    def display_models_info(self, filter: str = None):
-        # Собираем данные для таблицы
-        table_data = []
-        headers = [
-            "Тип модели",
-            "ID",
-            "Имя модели",
-            "Стемы",
-            "Целевой инструмент",
-        ]
-
-        for model_type, models in self.models_info.items():
-            for model_name, model_info in models.items():
-                try:
-                    stems_list = model_info.get("stems", [])
-                    id = model_info.get("id", "н/д")
-                    # Применяем фильтр (регистронезависимо)
-                    if filter:
-                        filter_lower = filter.lower()
-                        if not any(filter_lower == s.lower() for s in stems_list):
-                            continue
-
-                    # Подготавливаем данные для строки таблицы
-                    row = [
-                        model_type,
-                        id,
-                        model_name,
-                        ", ".join(stems_list) or "н/д",
-                        model_info.get("target_instrument", "н/д"),
-                    ]
-                    table_data.append(row)
-                except (KeyError, TypeError, AttributeError) as e:
-                    print(f"Ошибка при обработке модели {model_type}/{model_name}: {e}")
-                    continue
-
-        # Выводим результат
-        if table_data:
-            print(tabulate(table_data, headers=headers, tablefmt="grid"))
-        else:
-            print("Нет моделей, которые содержат указанный стем")
-
-    def download_model(
-        self, model_paths, model_name, model_type, ckpt_url, conf_url
-    ):
-        model_dir = os.path.join(model_paths, model_type)
-        os.makedirs(model_dir, exist_ok=True)
-
-        config_path = None
-        checkpoint_path = None
-
-        if model_type == "mel_band_roformer":
-            config_path = os.path.join(model_dir, f"{model_name}_config.yaml")
-            checkpoint_path = os.path.join(model_dir, f"{model_name}.ckpt")
-
-        elif model_type == "vr":
-            config_path = os.path.join(model_dir, f"{model_name}.yaml")
-            checkpoint_path = os.path.join(model_dir, f"{model_name}.pth")
-
-        elif model_type == "mdxnet":
-            config_path = os.path.join(model_dir, f"{model_name}.yaml")
-            checkpoint_path = os.path.join(model_dir, f"{model_name}.onnx")
-
-        elif model_type == "bs_roformer":
-            config_path = os.path.join(model_dir, f"{model_name}_config.yaml")
-            checkpoint_path = os.path.join(model_dir, f"{model_name}.ckpt")
-
-        elif model_type == "mdx23c":
-            config_path = os.path.join(model_dir, f"{model_name}_config.yaml")
-            checkpoint_path = os.path.join(model_dir, f"{model_name}.ckpt")
-
-        elif model_type == "scnet":
-            config_path = os.path.join(model_dir, f"{model_name}_config.yaml")
-            checkpoint_path = os.path.join(model_dir, f"{model_name}.ckpt")
-
-        elif model_type == "bandit":
-            config_path = os.path.join(model_dir, f"{model_name}_config.yaml")
-            checkpoint_path = os.path.join(model_dir, f"{model_name}.chpt")
-
-        elif model_type == "bandit_v2":
-            config_path = os.path.join(model_dir, f"{model_name}_config.yaml")
-            checkpoint_path = os.path.join(model_dir, f"{model_name}.ckpt")
-
-        elif model_type == "htdemucs":
-            config_path = os.path.join(model_dir, f"{model_name}_config.yaml")
-            checkpoint_path = os.path.join(model_dir, f"{model_name}.th")
-
-        else:
-            raise ValueError(
-                f"{self.I18N_helper.t('error_unsupported_model_type')}: {model_type}"
-            )
-
-        # Проверяем, что пути заданы (на всякий случай)
-        if config_path is None or checkpoint_path is None:
-            raise RuntimeError()
-
-        # Если файлы уже есть — пропускаем загрузку
-        if os.path.exists(checkpoint_path) and os.path.exists(config_path):
-            if os.path.getsize(checkpoint_path) == 0 or os.path.getsize(checkpoint_path) == 0:
-                for local_path, url_model in [
-                    (checkpoint_path, ckpt_url),
-                    (config_path, conf_url),
-                ]:
-                    if not os.path.exists(local_path):
-
-                        dw_file(url_model, local_path)
-            else:
-                pass
-        else:
-            for local_path, url_model in [
-                (checkpoint_path, ckpt_url),
-                (config_path, conf_url),
-            ]:
-                if not os.path.exists(local_path):
-
-                    dw_file(url_model, local_path)
-
-        return config_path, checkpoint_path
-
-    def conf_editor(self, config_path, mdx_denoise, vr_aggr, model_type):
-
-        class IndentDumper(yaml.Dumper):
-            def increase_indent(self, flow=False, indentless=False):
-                return super(IndentDumper, self).increase_indent(flow, False)
-
-        def tuple_constructor(loader, node):
-            # Load the sequence of values from the YAML node
-            values = loader.construct_sequence(node)
-            # Return a tuple constructed from the sequence
-            return tuple(values)
-
-        # Register the constructor with PyYAML
-        yaml.SafeLoader.add_constructor(
-            "tag:yaml.org,2002:python/tuple", tuple_constructor
-        )
-
-        def conf_edit(config_path, mdx_denoise, vr_aggr, model_type):
-            with open(config_path, "r") as f:
-                data = yaml.load(f, Loader=yaml.SafeLoader)
-
-            # handle cases where 'use_amp' is missing from config:
-            if "use_amp" not in data.keys():
-                data["training"]["use_amp"] = True
-
-            if model_type != "vr":
-                if data["inference"]["num_overlap"] != 2:
-                    data["inference"]["num_overlap"] = 2
-
-            if data["inference"]["batch_size"] != 1:
-                data["inference"]["batch_size"] = 1
-
-            if model_type == "mdxnet":
-                data["inference"]["denoise"] = mdx_denoise
-
-            elif model_type == "vr":
-                data["inference"]["aggression"] = vr_aggr
-
-            with open(config_path, "w") as f:
-                yaml.dump(
-                    data,
-                    f,
-                    default_flow_style=False,
-                    sort_keys=False,
-                    Dumper=IndentDumper,
-                    allow_unicode=True,
-                )
-
-        conf_edit(config_path, mdx_denoise, vr_aggr, model_type)
-
-class VbachModelManager:
-    def __init__(self):
-        self.rmvpe_path = os.path.join(script_dir, "predictors", "rmvpe.pt")
-        self.fcpe_path = os.path.join(script_dir, "predictors", "fcpe.pt")
-        self.custom_fairseq_huberts_dir = os.path.join(script_dir, "custom_fairseq_embedders")
-        self.custom_transformers_huberts_dir = os.path.join(script_dir, "custom_transformers_embedders")
-        self.huberts_fairseq_dict = {
-            "hubert_base": {
-                "url": "https://huggingface.co/Politrees/RVC_resources/resolve/main/embedders/hubert_base.pt",
-                "local_path": os.path.join(self.custom_fairseq_huberts_dir, "hubert_base.pt")
-            },
-            "contentvec_base": {
-                "url": "https://huggingface.co/Politrees/RVC_resources/resolve/main/embedders/contentvec_base.pt",
-                "local_path": os.path.join(self.custom_fairseq_huberts_dir, "contentvec_base.pt")
-            },
-            "korean_hubert_base": {
-                "url": "https://huggingface.co/Politrees/RVC_resources/resolve/main/embedders/korean_hubert_base.pt",
-                "local_path": os.path.join(self.custom_fairseq_huberts_dir, "korean_hubert_base.pt")
-            },
-            "chinese_hubert_base": {
-                "url": "https://huggingface.co/Politrees/RVC_resources/resolve/main/embedders/chinese_hubert_base.pt",
-                "local_path": os.path.join(self.custom_fairseq_huberts_dir, "chinese_hubert_base.pt")
-            },
-            "portuguese_hubert_base": {
-                "url": "https://huggingface.co/Politrees/RVC_resources/resolve/main/embedders/portuguese_hubert_base.pt",
-                "local_path": os.path.join(self.custom_fairseq_huberts_dir, "portuguese_hubert_base.pt")
-            },
-            "japanese_hubert_base": {
-                "url": "https://huggingface.co/Politrees/RVC_resources/resolve/main/embedders/japanese_hubert_base.pt",
-                "local_path": os.path.join(self.custom_fairseq_huberts_dir, "japanese_hubert_base.pt")
-            }
-        }
-        self.huberts_transformers_dict = {
-            "contentvec": {
-                "base_dir": os.path.join(self.custom_transformers_huberts_dir, "contentvec"),
-                "url_bin": "https://huggingface.co/IAHispano/Applio/resolve/main/Resources/embedders/contentvec/pytorch_model.bin",
-                "url_json": "https://huggingface.co/IAHispano/Applio/resolve/main/Resources/embedders/contentvec/config.json",
-                "local_bin": os.path.join(self.custom_transformers_huberts_dir, "contentvec", "pytorch_model.bin"),
-                "local_json": os.path.join(self.custom_transformers_huberts_dir, "contentvec", "config.json")
-            },
-            "spin": {
-                "base_dir": os.path.join(self.custom_transformers_huberts_dir, "spin"),
-                "url_bin": "https://huggingface.co/IAHispano/Applio/resolve/main/Resources/embedders/spin/pytorch_model.bin",
-                "url_json": "https://huggingface.co/IAHispano/Applio/resolve/main/Resources/embedders/spin/config.json",
-                "local_bin": os.path.join(self.custom_transformers_huberts_dir, "spin", "pytorch_model.bin"),
-                "local_json": os.path.join(self.custom_transformers_huberts_dir, "spin", "config.json")
-            },
-            "spin-v2": {
-                "base_dir": os.path.join(self.custom_transformers_huberts_dir, "spinv2"),
-                "url_bin": "https://huggingface.co/IAHispano/Applio/resolve/main/Resources/embedders/spin-v2/pytorch_model.bin",
-                "url_json": "https://huggingface.co/IAHispano/Applio/resolve/main/Resources/embedders/spin-v2/config.json",
-                "local_bin": os.path.join(self.custom_transformers_huberts_dir, "spinv2", "pytorch_model.bin"),
-                "local_json": os.path.join(self.custom_transformers_huberts_dir, "spinv2", "config.json")
-            },
-            "chinese-hubert-base": {
-                "base_dir": os.path.join(self.custom_transformers_huberts_dir, "chinese_hubert_base"),
-                "url_bin": "https://huggingface.co/IAHispano/Applio/resolve/main/Resources/embedders/chinese_hubert_base/pytorch_model.bin",
-                "url_json": "https://huggingface.co/IAHispano/Applio/resolve/main/Resources/embedders/chinese_hubert_base/config.json",
-                "local_bin": os.path.join(self.custom_transformers_huberts_dir, "chinese_hubert_base", "pytorch_model.bin"),
-                "local_json": os.path.join(self.custom_transformers_huberts_dir, "chinese_hubert_base", "config.json")
-            },
-            "japanese-hubert-base": {
-                "base_dir": os.path.join(self.custom_transformers_huberts_dir, "japanese_hubert_base"),
-                "url_bin": "https://huggingface.co/IAHispano/Applio/resolve/main/Resources/embedders/japanese_hubert_base/pytorch_model.bin",
-                "url_json": "https://huggingface.co/IAHispano/Applio/resolve/main/Resources/embedders/japanese_hubert_base/config.json",
-                "local_bin": os.path.join(self.custom_transformers_huberts_dir, "japanese_hubert_base", "pytorch_model.bin"),
-                "local_json": os.path.join(self.custom_transformers_huberts_dir, "japanese_hubert_base", "config.json")
-            },
-            "korean-hubert-base": {
-                "base_dir": os.path.join(self.custom_transformers_huberts_dir, "korean_hubert_base"),
-                "url_bin": "https://huggingface.co/IAHispano/Applio/resolve/main/Resources/embedders/korean_hubert_base/pytorch_model.bin",
-                "url_json": "https://huggingface.co/IAHispano/Applio/resolve/main/Resources/embedders/korean_hubert_base/config.json",
-                "local_bin": os.path.join(self.custom_transformers_huberts_dir, "korean_hubert_base", "pytorch_model.bin"),
-                "local_json": os.path.join(self.custom_transformers_huberts_dir, "korean_hubert_base", "config.json")
-            }
-        }
-        self.requirements = [["https://huggingface.co/Politrees/RVC_resources/resolve/main/predictors/rmvpe.pt", self.rmvpe_path], ["https://huggingface.co/Politrees/RVC_resources/resolve/main/predictors/fcpe.pt", self.fcpe_path]]
-        self.voicemodels_dir = os.path.join(script_dir, "vbach_models_cache")
-        os.makedirs(self.voicemodels_dir, exist_ok=True)
-        self.voicemodels_info = os.path.join(self.voicemodels_dir, "vbach_models.json")
-        self.voicemodels: Dict[str, Dict[str, str]] = {}
-        self.download_requirements()
-        self.check_hubert("hubert_base")
-        self.check_and_load()
-        pass
-
-    def check_hubert(self, embedder_name):
-        if embedder_name in self.huberts_fairseq_dict:
-            if not os.path.exists(self.huberts_fairseq_dict[embedder_name]["local_path"]):
-                dw_file(self.huberts_fairseq_dict[embedder_name]["url"], self.huberts_fairseq_dict[embedder_name]["local_path"])
-            return self.huberts_fairseq_dict[embedder_name]["local_path"]
-        else:
-            return None
-        
-    def check_hubert_transformers(self, embedder_name):
-        if embedder_name in self.huberts_transformers_dict:
-            os.makedirs(self.huberts_transformers_dict[embedder_name]["base_dir"], exist_ok=True)
-            if not os.path.exists(self.huberts_transformers_dict[embedder_name]["local_bin"]) and not os.path.exists(self.huberts_transformers_dict[embedder_name]["local_json"]):
-                dw_file(self.huberts_transformers_dict[embedder_name]["url_bin"], self.huberts_transformers_dict[embedder_name]["local_bin"])
-                dw_file(self.huberts_transformers_dict[embedder_name]["url_json"], self.huberts_transformers_dict[embedder_name]["local_json"])
-            return self.huberts_transformers_dict[embedder_name]["base_dir"]
-        else:
-            return None
-    
-    def write_voicemodels_info(self):
-        with open(self.voicemodels_info, "w") as f:
-            json.dump(self.voicemodels, f, indent=4)
-
-    def load_voicemodels_info(self):
-        with open(self.voicemodels_info, "r") as f:
-            return json.load(f)
-
-    def add_voice_model(
-        self,
-        name,
-        pth_path,
-        index_path,
-    ):
-        self.voicemodels[name] = {"pth": pth_path, "index": index_path}
-        self.write_voicemodels_info()
-
-    def del_voice_model(
-        self, name
-    ):
-        if name in self.parse_voice_models():
-            pth = self.voicemodels[name].get("pth", None)
-            index = self.voicemodels[name].get("index", None)
-            if index:
-                os.remove(index)
-            if pth:
-                os.remove(pth)
-            del self.voicemodels[name]
-            self.write_voicemodels_info()
-            return f"Модель {name} удалена"
-        else:
-            return f"Модель не была удалена, как так её не существует"
-
-    def parse_voice_models(self):
-        list_models = list(self.voicemodels.keys())
-        return list_models
-    
-    def parse_pth_and_index(self, name):
-        pth = self.voicemodels[name].get("pth", None)
-        index = self.voicemodels[name].get("index", None)
-        return pth, index
-    
-    def check_and_load(self):
-        if os.path.exists(self.voicemodels_info):
-            self.voicemodels = self.load_voicemodels_info()
-        else:
-            self.write_voicemodels_info()
-
-    def clear_voicemodels_info(self):
-        self.voicemodels: Dict[str, Dict[str, str]] = {}
-        self.write_voicemodels_info()
-    
-    def download_requirements(self):
-        for url, file in self.requirements:
-            if not os.path.exists(file):
-                dw_file(url, file)
-
-    def download_voice_model_file(self, url, zip_name):
-        try:
-            if "drive.google.com" in url:
-                self.download_from_google_drive(url, zip_name)
-            elif "pixeldrain.com" in url:
-                self.download_from_pixeldrain(url, zip_name)
-            elif "disk.yandex.ru" in url or "yadi.sk" in url:
-                self.download_from_yandex(url, zip_name)
-            else:
-                dw_file(url, zip_name)
-        except Exception as e:
-            print(e)
-
-    def download_from_google_drive(self, url, zip_name):
-        file_id = (
-            url.split("file/d/")[1].split("/")[0]
-            if "file/d/" in url
-            else url.split("id=")[1].split("&")[0]
-        )
-        gdown.download(id=file_id, output=str(zip_name), quiet=False)
-
-    def download_from_pixeldrain(self, url, zip_name):
-        file_id = url.split("pixeldrain.com/u/")[1]
-        response = requests.get(f"https://pixeldrain.com/api/file/{file_id}")
-        with open(zip_name, "wb") as f:
-            f.write(response.content)
-
-    def download_from_yandex(self, url, zip_name):
-        yandex_public_key = f"download?public_key={url}"
-        yandex_api_url = f"https://cloud-api.yandex.net/v1/disk/public/resources/{yandex_public_key}"
-        response = requests.get(yandex_api_url)
-        if response.status_code == 200:
-            download_link = response.json().get("href")
-            urllib.request.urlretrieve(download_link, zip_name)
-        else:
-            print(response.status_code)
-
-    def extract_zip(self, zip_name, model_name):
-        model_dir = os.path.join(self.voicemodels_dir, f"{model_name}_{generate_secure_random(17)}")
-        os.makedirs(model_dir, exist_ok=True)
-        try:
-            with zipfile.ZipFile(zip_name, "r") as zip_ref:
-                zip_ref.extractall(model_dir)
-            os.remove(zip_name)
-
-            added_voice_models = []
-
-            index_filepath, model_filepaths = None, []
-            for root, _, files in os.walk(model_dir):
-                for name in files:
-                    file_path = os.path.join(root, name)
-                    if name.endswith(".index") and os.stat(file_path).st_size > 1024 * 100:
-                        index_filepath = file_path
-                    if name.endswith(".pth") and os.stat(file_path).st_size > 1024 * 1024 * 20:
-                        model_filepaths.append(file_path)
-
-            if len(model_filepaths) == 1:
-                self.add_voice_model(model_name, model_filepaths[0], index_filepath)
-                added_voice_models.append(model_name)
-            else:
-                for i, pth in enumerate(model_filepaths):
-                    self.add_voice_model(f"{model_name}_{i + 1}", pth, index_filepath)
-                    added_voice_models.append(f"{model_name}_{i + 1}")
-            list_models_str = '\n'.join(added_voice_models)
-            return f"Добавленные модели:\n{list_models_str}"
-        except Exception as e:
-            return f"Произошла ошибка при загрузке модели: {e}"
-
-    def install_model_zip(self, zip, model_name, mode="url"):
-        if model_name in self.parse_voice_models():
-            print("Эта модель уже есть в списке установленных моделей. Она будут перезаписана")
-        if mode == "url":
-            with tempfile.TemporaryDirectory(prefix="vbach_temp_model", ignore_cleanup_errors=True) as tmp:
-                zip_path = os.path.join(tmp, "model.zip")
-                self.download_voice_model_file(zip, zip_path)
-                status = self.extract_zip(zip_path, model_name)
-        if mode == "local":
-            status = self.extract_zip(zip, model_name)
-        return status
-
-    def install_model_files(self, index, pth, model_name, mode="url"):
-        if model_name in self.parse_voice_models():
-            print("Эта модель уже есть в списке установленных моделей. Она будут перезаписана")
-        model_dir = os.path.join(self.voicemodels_dir, f"{model_name}_{generate_secure_random(17)}")
-        os.makedirs(model_dir, exist_ok=True)
-        local_index_path = None
-        local_pth_path = None
-        try:
-            if mode == "url":
-                if index:
-                    local_index_path = os.path.join(model_dir, "model.index")
-                    self.download_voice_model_file(index, local_index_path)
-                if pth:
-                    local_pth_path = os.path.join(model_dir, "model.pth")
-                    self.download_voice_model_file(pth, local_pth_path)
-
-            if mode == "local":
-                if index:
-                    if os.path.exists(index):
-                        local_index_path = os.path.join(model_dir, os.path.basename(index))
-                        shutil.copy(index, local_index_path)
-                if pth:
-                    if os.path.exists(pth):
-                        local_pth_path = os.path.join(model_dir, os.path.basename(pth))
-                        shutil.copy(pth, local_pth_path)
-
-            self.add_voice_model(model_name, local_pth_path, local_index_path)
-            return f"Модель {model_name} добавлена"
-        except Exception as e:
-            return f"Произошла ошибка при загрузке модели: {e}"
-        
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Менеджер моделей")
-    subparsers = parser.add_subparsers(title="subcommands", dest="command", required=True)
-    
-    # Mvsepless subcommand
-    mvsepless_parser = subparsers.add_parser("mvsepless", help="Скачивание моделей в MVSepLess")
-    mvsepless_parser.add_argument("--model_type", required=True, help="Тип модели")
-    mvsepless_parser.add_argument("--model_name", required=True, help="Имя модели")
-    
-    # Vbach subcommand
-    vbach_parser = subparsers.add_parser("vbach", help="Установка голосовых моделей в Vbach")
-    vbach_subparsers = vbach_parser.add_subparsers(title="vbach_commands", dest="vbach_command", required=True)
-    
-    # Vbach install_local
-    install_local_parser = vbach_subparsers.add_parser("install_local", help="Установка голосовой модели по локальным файлам")
-    install_local_parser.add_argument("--model_name", required=True, help="Имя голосовой модели")
-    install_local_parser.add_argument("--pth", required=True, help="Путь к *.pth файлу")
-    install_local_parser.add_argument("--index", required=False, help="Путь к *.index файлу")
-    
-    # Vbach install_url_zip
-    install_url_zip_parser = vbach_subparsers.add_parser("install_url_zip", help="Установка голосовой модели по URL (архив с файлами)")
-    install_url_zip_parser.add_argument("--model_name", required=True, help="Имя голосовой модели")
-    install_url_zip_parser.add_argument("--url", required=True, help="URL *.zip файла")
-    
-    # Vbach install_url_files
-    install_url_files_parser = vbach_subparsers.add_parser("install_url_files", help="Установка голосовой модели по URL (отдельные файлы)")
-    install_url_files_parser.add_argument("--model_name", required=True, help="Имя голосовой модели")
-    install_url_files_parser.add_argument("--pth_url", required=True, help="URL *.pth файла")
-    install_url_files_parser.add_argument("--index_url", required=False, help="URL *.index файла")
-    
-    # Vbach list
-    list_parser = vbach_subparsers.add_parser("list", help="List installed voice models")
-    
-    args = parser.parse_args()
-
-    if args.command == "mvsepless":
-
-        _model_manager = MvseplessModelManager()
-        info = _model_manager.models_info[args.model_type].get(args.model_name, None)
-        if not info:
-            raise ValueError(f"Модель {args.model_name} не найдена для типа {args.model_type}")
-        conf, ckpt = _model_manager.download_model(
-            _model_manager.models_cache_dir,
-            args.model_name,
-            args.model_type,
-            info["checkpoint_url"],
-            info["config_url"],
-        )
-        
-    elif args.command == "vbach":
-        model_manager = VbachModelManager()
-        
-        if args.vbach_command == "install_local":
-            status = model_manager.install_model_files(
-                args.index, args.pth, args.model_name, mode="local"
-            )
-            print(status)
-            
-        elif args.vbach_command == "install_url_zip":
-            status = model_manager.install_model_zip(
-                args.url, args.model_name, mode="url"
-            )
-            print(status)
-            
-        elif args.vbach_command == "install_url_files":
-            status = model_manager.install_model_files(
-                args.index_url, args.pth_url, args.model_name, mode="url"
-            )
-            print(status)
-            
-        elif args.vbach_command == "list":
-            models = model_manager.parse_voice_models()
-            if models:
-                print("Установленные модели:")
-                for model in models:
-                    print(f"  - {model}")
-            else:
-                print("Нет установленных моделей")
-
-
-
-
-
-
+import os
+import sys
+import json
+import yaml
+from tabulate import tabulate
+import shutil
+from tqdm import tqdm
+import urllib.request
+import gdown
+import requests
+import zipfile
+import tempfile
+import secrets
+import string
+import argparse
+from typing import Dict, Any
+
+script_dir = os.path.dirname(os.path.abspath(__file__))
+if not __package__:
+    from downloader import dw_file
+else:
+    from .downloader import dw_file
+
+
+def generate_secure_random(length=10):
+    characters = string.ascii_letters + string.digits
+    return "".join(secrets.choice(characters) for _ in range(length))
+
+
+class MvseplessModelManager:
+    def __init__(
+        self,
+        models_info_path=os.path.join(script_dir, "models.json"),
+        cache_dir=os.path.join(script_dir, "mvsepless_models_cache"),
+    ):
+        self.models_cache_dir = cache_dir
+        self.models_info_path = models_info_path
+        with open(self.models_info_path, "r", encoding="utf-8") as f:
+            models_info = json.load(f)
+        self.models_info = models_info
+
+    def get_mt(self):
+        return [mt for mt in self.models_info]
+
+    def get_mn(self, model_type):
+        return [mn for mn in self.models_info.get(model_type, [])]
+
+    def get_stems(self, model_type, model_name):
+        return [
+            stem
+            for stem in self.models_info.get(model_type)
+            .get(model_name)
+            .get("stems", [])
+        ]
+
+    def get_id(self, model_type, model_name):
+        return self.models_info.get(model_type).get(model_name).get("id", 0)
+
+    def get_tgt_inst(self, model_type, model_name):
+        return (
+            self.models_info.get(model_type)
+            .get(model_name)
+            .get("target_instrument", None)
+        )
+
+    def get_category(self, model_type, model_name):
+        return self.models_info.get(model_type).get(model_name).get("category", "")
+
+    def display_models_info(self, filter: str = None):
+        table_data = []
+        headers = [
+            "Тип модели",
+            "ID",
+            "Имя модели",
+            "Стемы",
+            "Целевой инструмент",
+        ]
+
+        for model_type, models in self.models_info.items():
+            for model_name, model_info in models.items():
+                try:
+                    stems_list = model_info.get("stems", [])
+                    id = model_info.get("id", "н/д")
+                    if filter:
+                        filter_lower = filter.lower()
+                        if not any(filter_lower == s.lower() for s in stems_list):
+                            continue
+
+                    row = [
+                        model_type,
+                        id,
+                        model_name,
+                        ", ".join(stems_list) or "н/д",
+                        model_info.get("target_instrument", "н/д"),
+                    ]
+                    table_data.append(row)
+                except (KeyError, TypeError, AttributeError) as e:
+                    print(f"Ошибка при обработке модели {model_type}/{model_name}: {e}")
+                    continue
+
+        if table_data:
+            print(tabulate(table_data, headers=headers, tablefmt="grid"))
+        else:
+            print("Нет моделей, которые содержат указанный стем")
+
+    def download_model(self, model_paths, model_name, model_type, ckpt_url, conf_url):
+        model_dir = os.path.join(model_paths, model_type)
+        os.makedirs(model_dir, exist_ok=True)
+
+        config_path = os.path.join(model_dir, f"{model_name}_config.yaml")
+        checkpoint_path = os.path.join(
+            model_dir,
+            f"{model_name}.onnx" if model_type == "mdxnet" else f"{model_name}.ckpt",
+        )
+
+        if config_path is None or checkpoint_path is None:
+            raise RuntimeError()
+
+        if os.path.exists(checkpoint_path) and os.path.exists(config_path):
+            if (
+                os.path.getsize(checkpoint_path) == 0
+                or os.path.getsize(checkpoint_path) == 0
+            ):
+                for local_path, url_model in [
+                    (checkpoint_path, ckpt_url),
+                    (config_path, conf_url),
+                ]:
+                    if not os.path.exists(local_path):
+
+                        dw_file(url_model, local_path)
+            else:
+                pass
+        else:
+            for local_path, url_model in [
+                (checkpoint_path, ckpt_url),
+                (config_path, conf_url),
+            ]:
+                if not os.path.exists(local_path):
+
+                    dw_file(url_model, local_path)
+
+        return config_path, checkpoint_path
+
+    def conf_editor(self, config_path, mdx_denoise, vr_aggr, model_type):
+
+        class IndentDumper(yaml.Dumper):
+            def increase_indent(self, flow=False, indentless=False):
+                return super(IndentDumper, self).increase_indent(flow, False)
+
+        def tuple_constructor(loader, node):
+            values = loader.construct_sequence(node)
+            return tuple(values)
+
+        yaml.SafeLoader.add_constructor(
+            "tag:yaml.org,2002:python/tuple", tuple_constructor
+        )
+
+        def conf_edit(config_path, mdx_denoise, vr_aggr, model_type):
+            with open(config_path, "r") as f:
+                data = yaml.load(f, Loader=yaml.SafeLoader)
+
+            if "use_amp" not in data.keys():
+                data["training"]["use_amp"] = True
+
+            if model_type != "vr":
+                if data["inference"]["num_overlap"] != 2:
+                    data["inference"]["num_overlap"] = 2
+
+            if data["inference"]["batch_size"] != 1:
+                data["inference"]["batch_size"] = 1
+
+            if model_type == "mdxnet":
+                data["inference"]["denoise"] = mdx_denoise
+
+            elif model_type == "vr":
+                data["inference"]["aggression"] = vr_aggr
+
+            with open(config_path, "w") as f:
+                yaml.dump(
+                    data,
+                    f,
+                    default_flow_style=False,
+                    sort_keys=False,
+                    Dumper=IndentDumper,
+                    allow_unicode=True,
+                )
+
+        conf_edit(config_path, mdx_denoise, vr_aggr, model_type)
+
+
+class VbachModelManager:
+    def __init__(self):
+        self.rmvpe_path = os.path.join(script_dir, "predictors", "rmvpe.pt")
+        self.fcpe_path = os.path.join(script_dir, "predictors", "fcpe.pt")
+        self.custom_fairseq_huberts_dir = os.path.join(
+            script_dir, "custom_fairseq_embedders"
+        )
+        self.custom_transformers_huberts_dir = os.path.join(
+            script_dir, "custom_transformers_embedders"
+        )
+        self.huberts_fairseq_dict = {
+            "hubert_base": {
+                "url": "https://huggingface.co/Politrees/RVC_resources/resolve/main/embedders/hubert_base.pt",
+                "local_path": os.path.join(
+                    self.custom_fairseq_huberts_dir, "hubert_base.pt"
+                ),
+            },
+            "contentvec_base": {
+                "url": "https://huggingface.co/Politrees/RVC_resources/resolve/main/embedders/contentvec_base.pt",
+                "local_path": os.path.join(
+                    self.custom_fairseq_huberts_dir, "contentvec_base.pt"
+                ),
+            },
+            "korean_hubert_base": {
+                "url": "https://huggingface.co/Politrees/RVC_resources/resolve/main/embedders/korean_hubert_base.pt",
+                "local_path": os.path.join(
+                    self.custom_fairseq_huberts_dir, "korean_hubert_base.pt"
+                ),
+            },
+            "chinese_hubert_base": {
+                "url": "https://huggingface.co/Politrees/RVC_resources/resolve/main/embedders/chinese_hubert_base.pt",
+                "local_path": os.path.join(
+                    self.custom_fairseq_huberts_dir, "chinese_hubert_base.pt"
+                ),
+            },
+            "portuguese_hubert_base": {
+                "url": "https://huggingface.co/Politrees/RVC_resources/resolve/main/embedders/portuguese_hubert_base.pt",
+                "local_path": os.path.join(
+                    self.custom_fairseq_huberts_dir, "portuguese_hubert_base.pt"
+                ),
+            },
+            "japanese_hubert_base": {
+                "url": "https://huggingface.co/Politrees/RVC_resources/resolve/main/embedders/japanese_hubert_base.pt",
+                "local_path": os.path.join(
+                    self.custom_fairseq_huberts_dir, "japanese_hubert_base.pt"
+                ),
+            },
+        }
+        self.huberts_transformers_dict = {
+            "contentvec": {
+                "base_dir": os.path.join(
+                    self.custom_transformers_huberts_dir, "contentvec"
+                ),
+                "url_bin": "https://huggingface.co/IAHispano/Applio/resolve/main/Resources/embedders/contentvec/pytorch_model.bin",
+                "url_json": "https://huggingface.co/IAHispano/Applio/resolve/main/Resources/embedders/contentvec/config.json",
+                "local_bin": os.path.join(
+                    self.custom_transformers_huberts_dir,
+                    "contentvec",
+                    "pytorch_model.bin",
+                ),
+                "local_json": os.path.join(
+                    self.custom_transformers_huberts_dir, "contentvec", "config.json"
+                ),
+            },
+            "spin": {
+                "base_dir": os.path.join(self.custom_transformers_huberts_dir, "spin"),
+                "url_bin": "https://huggingface.co/IAHispano/Applio/resolve/main/Resources/embedders/spin/pytorch_model.bin",
+                "url_json": "https://huggingface.co/IAHispano/Applio/resolve/main/Resources/embedders/spin/config.json",
+                "local_bin": os.path.join(
+                    self.custom_transformers_huberts_dir, "spin", "pytorch_model.bin"
+                ),
+                "local_json": os.path.join(
+                    self.custom_transformers_huberts_dir, "spin", "config.json"
+                ),
+            },
+            "spin-v2": {
+                "base_dir": os.path.join(
+                    self.custom_transformers_huberts_dir, "spinv2"
+                ),
+                "url_bin": "https://huggingface.co/IAHispano/Applio/resolve/main/Resources/embedders/spin-v2/pytorch_model.bin",
+                "url_json": "https://huggingface.co/IAHispano/Applio/resolve/main/Resources/embedders/spin-v2/config.json",
+                "local_bin": os.path.join(
+                    self.custom_transformers_huberts_dir, "spinv2", "pytorch_model.bin"
+                ),
+                "local_json": os.path.join(
+                    self.custom_transformers_huberts_dir, "spinv2", "config.json"
+                ),
+            },
+            "chinese-hubert-base": {
+                "base_dir": os.path.join(
+                    self.custom_transformers_huberts_dir, "chinese_hubert_base"
+                ),
+                "url_bin": "https://huggingface.co/IAHispano/Applio/resolve/main/Resources/embedders/chinese_hubert_base/pytorch_model.bin",
+                "url_json": "https://huggingface.co/IAHispano/Applio/resolve/main/Resources/embedders/chinese_hubert_base/config.json",
+                "local_bin": os.path.join(
+                    self.custom_transformers_huberts_dir,
+                    "chinese_hubert_base",
+                    "pytorch_model.bin",
+                ),
+                "local_json": os.path.join(
+                    self.custom_transformers_huberts_dir,
+                    "chinese_hubert_base",
+                    "config.json",
+                ),
+            },
+            "japanese-hubert-base": {
+                "base_dir": os.path.join(
+                    self.custom_transformers_huberts_dir, "japanese_hubert_base"
+                ),
+                "url_bin": "https://huggingface.co/IAHispano/Applio/resolve/main/Resources/embedders/japanese_hubert_base/pytorch_model.bin",
+                "url_json": "https://huggingface.co/IAHispano/Applio/resolve/main/Resources/embedders/japanese_hubert_base/config.json",
+                "local_bin": os.path.join(
+                    self.custom_transformers_huberts_dir,
+                    "japanese_hubert_base",
+                    "pytorch_model.bin",
+                ),
+                "local_json": os.path.join(
+                    self.custom_transformers_huberts_dir,
+                    "japanese_hubert_base",
+                    "config.json",
+                ),
+            },
+            "korean-hubert-base": {
+                "base_dir": os.path.join(
+                    self.custom_transformers_huberts_dir, "korean_hubert_base"
+                ),
+                "url_bin": "https://huggingface.co/IAHispano/Applio/resolve/main/Resources/embedders/korean_hubert_base/pytorch_model.bin",
+                "url_json": "https://huggingface.co/IAHispano/Applio/resolve/main/Resources/embedders/korean_hubert_base/config.json",
+                "local_bin": os.path.join(
+                    self.custom_transformers_huberts_dir,
+                    "korean_hubert_base",
+                    "pytorch_model.bin",
+                ),
+                "local_json": os.path.join(
+                    self.custom_transformers_huberts_dir,
+                    "korean_hubert_base",
+                    "config.json",
+                ),
+            },
+        }
+        self.requirements = [
+            [
+                "https://huggingface.co/Politrees/RVC_resources/resolve/main/predictors/rmvpe.pt",
+                self.rmvpe_path,
+            ],
+            [
+                "https://huggingface.co/Politrees/RVC_resources/resolve/main/predictors/fcpe.pt",
+                self.fcpe_path,
+            ],
+        ]
+        self.voicemodels_dir = os.path.join(script_dir, "vbach_models_cache")
+        os.makedirs(self.voicemodels_dir, exist_ok=True)
+        self.voicemodels_info = os.path.join(self.voicemodels_dir, "vbach_models.json")
+        self.voicemodels: Dict[str, Dict[str, str]] = {}
+        self.download_requirements()
+        self.check_hubert("hubert_base")
+        self.check_and_load()
+        pass
+
+    def check_hubert(self, embedder_name):
+        if embedder_name in self.huberts_fairseq_dict:
+            if not os.path.exists(
+                self.huberts_fairseq_dict[embedder_name]["local_path"]
+            ):
+                dw_file(
+                    self.huberts_fairseq_dict[embedder_name]["url"],
+                    self.huberts_fairseq_dict[embedder_name]["local_path"],
+                )
+            return self.huberts_fairseq_dict[embedder_name]["local_path"]
+        else:
+            return None
+
+    def check_hubert_transformers(self, embedder_name):
+        if embedder_name in self.huberts_transformers_dict:
+            os.makedirs(
+                self.huberts_transformers_dict[embedder_name]["base_dir"], exist_ok=True
+            )
+            if not os.path.exists(
+                self.huberts_transformers_dict[embedder_name]["local_bin"]
+            ) and not os.path.exists(
+                self.huberts_transformers_dict[embedder_name]["local_json"]
+            ):
+                dw_file(
+                    self.huberts_transformers_dict[embedder_name]["url_bin"],
+                    self.huberts_transformers_dict[embedder_name]["local_bin"],
+                )
+                dw_file(
+                    self.huberts_transformers_dict[embedder_name]["url_json"],
+                    self.huberts_transformers_dict[embedder_name]["local_json"],
+                )
+            return self.huberts_transformers_dict[embedder_name]["base_dir"]
+        else:
+            return None
+
+    def write_voicemodels_info(self):
+        with open(self.voicemodels_info, "w") as f:
+            json.dump(self.voicemodels, f, indent=4)
+
+    def load_voicemodels_info(self):
+        with open(self.voicemodels_info, "r") as f:
+            return json.load(f)
+
+    def add_voice_model(
+        self,
+        name,
+        pth_path,
+        index_path,
+    ):
+        self.voicemodels[name] = {"pth": pth_path, "index": index_path}
+        self.write_voicemodels_info()
+
+    def del_voice_model(self, name):
+        if name in self.parse_voice_models():
+            pth = self.voicemodels[name].get("pth", None)
+            index = self.voicemodels[name].get("index", None)
+            if index:
+                os.remove(index)
+            if pth:
+                os.remove(pth)
+            del self.voicemodels[name]
+            self.write_voicemodels_info()
+            return f"Модель {name} удалена"
+        else:
+            return f"Модель не была удалена, как так её не существует"
+
+    def parse_voice_models(self):
+        list_models = list(self.voicemodels.keys())
+        return list_models
+
+    def parse_pth_and_index(self, name):
+        pth = self.voicemodels[name].get("pth", None)
+        index = self.voicemodels[name].get("index", None)
+        return pth, index
+
+    def check_and_load(self):
+        if os.path.exists(self.voicemodels_info):
+            self.voicemodels = self.load_voicemodels_info()
+        else:
+            self.write_voicemodels_info()
+
+    def clear_voicemodels_info(self):
+        self.voicemodels: Dict[str, Dict[str, str]] = {}
+        self.write_voicemodels_info()
+
+    def download_requirements(self):
+        for url, file in self.requirements:
+            if not os.path.exists(file):
+                dw_file(url, file)
+
+    def download_voice_model_file(self, url, zip_name):
+        try:
+            if "drive.google.com" in url:
+                self.download_from_google_drive(url, zip_name)
+            elif "pixeldrain.com" in url:
+                self.download_from_pixeldrain(url, zip_name)
+            elif "disk.yandex.ru" in url or "yadi.sk" in url:
+                self.download_from_yandex(url, zip_name)
+            else:
+                dw_file(url, zip_name)
+        except Exception as e:
+            print(e)
+
+    def download_from_google_drive(self, url, zip_name):
+        file_id = (
+            url.split("file/d/")[1].split("/")[0]
+            if "file/d/" in url
+            else url.split("id=")[1].split("&")[0]
+        )
+        gdown.download(id=file_id, output=str(zip_name), quiet=False)
+
+    def download_from_pixeldrain(self, url, zip_name):
+        file_id = url.split("pixeldrain.com/u/")[1]
+        response = requests.get(f"https://pixeldrain.com/api/file/{file_id}")
+        with open(zip_name, "wb") as f:
+            f.write(response.content)
+
+    def download_from_yandex(self, url, zip_name):
+        yandex_public_key = f"download?public_key={url}"
+        yandex_api_url = (
+            f"https://cloud-api.yandex.net/v1/disk/public/resources/{yandex_public_key}"
+        )
+        response = requests.get(yandex_api_url)
+        if response.status_code == 200:
+            download_link = response.json().get("href")
+            urllib.request.urlretrieve(download_link, zip_name)
+        else:
+            print(response.status_code)
+
+    def extract_zip(self, zip_name, model_name):
+        model_dir = os.path.join(
+            self.voicemodels_dir, f"{model_name}_{generate_secure_random(17)}"
+        )
+        os.makedirs(model_dir, exist_ok=True)
+        try:
+            with zipfile.ZipFile(zip_name, "r") as zip_ref:
+                zip_ref.extractall(model_dir)
+            os.remove(zip_name)
+
+            added_voice_models = []
+
+            index_filepath, model_filepaths = None, []
+            for root, _, files in os.walk(model_dir):
+                for name in files:
+                    file_path = os.path.join(root, name)
+                    if (
+                        name.endswith(".index")
+                        and os.stat(file_path).st_size > 1024 * 100
+                    ):
+                        index_filepath = file_path
+                    if (
+                        name.endswith(".pth")
+                        and os.stat(file_path).st_size > 1024 * 1024 * 20
+                    ):
+                        model_filepaths.append(file_path)
+
+            if len(model_filepaths) == 1:
+                self.add_voice_model(model_name, model_filepaths[0], index_filepath)
+                added_voice_models.append(model_name)
+            else:
+                for i, pth in enumerate(model_filepaths):
+                    self.add_voice_model(f"{model_name}_{i + 1}", pth, index_filepath)
+                    added_voice_models.append(f"{model_name}_{i + 1}")
+            list_models_str = "\n".join(added_voice_models)
+            return f"Добавленные модели:\n{list_models_str}"
+        except Exception as e:
+            return f"Произошла ошибка при загрузке модели: {e}"
+
+    def install_model_zip(self, zip, model_name, mode="url"):
+        if model_name in self.parse_voice_models():
+            print(
+                "Эта модель уже есть в списке установленных моделей. Она будут перезаписана"
+            )
+        if mode == "url":
+            with tempfile.TemporaryDirectory(
+                prefix="vbach_temp_model", ignore_cleanup_errors=True
+            ) as tmp:
+                zip_path = os.path.join(tmp, "model.zip")
+                self.download_voice_model_file(zip, zip_path)
+                status = self.extract_zip(zip_path, model_name)
+        if mode == "local":
+            status = self.extract_zip(zip, model_name)
+        return status
+
+    def install_model_files(self, index, pth, model_name, mode="url"):
+        if model_name in self.parse_voice_models():
+            print(
+                "Эта модель уже есть в списке установленных моделей. Она будут перезаписана"
+            )
+        model_dir = os.path.join(
+            self.voicemodels_dir, f"{model_name}_{generate_secure_random(17)}"
+        )
+        os.makedirs(model_dir, exist_ok=True)
+        local_index_path = None
+        local_pth_path = None
+        try:
+            if mode == "url":
+                if index:
+                    local_index_path = os.path.join(model_dir, "model.index")
+                    self.download_voice_model_file(index, local_index_path)
+                if pth:
+                    local_pth_path = os.path.join(model_dir, "model.pth")
+                    self.download_voice_model_file(pth, local_pth_path)
+
+            if mode == "local":
+                if index:
+                    if os.path.exists(index):
+                        local_index_path = os.path.join(
+                            model_dir, os.path.basename(index)
+                        )
+                        shutil.copy(index, local_index_path)
+                if pth:
+                    if os.path.exists(pth):
+                        local_pth_path = os.path.join(model_dir, os.path.basename(pth))
+                        shutil.copy(pth, local_pth_path)
+
+            self.add_voice_model(model_name, local_pth_path, local_index_path)
+            return f"Модель {model_name} добавлена"
+        except Exception as e:
+            return f"Произошла ошибка при загрузке модели: {e}"
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Менеджер моделей")
+    subparsers = parser.add_subparsers(
+        title="subcommands", dest="command", required=True
+    )
+
+    mvsepless_parser = subparsers.add_parser(
+        "mvsepless", help="Скачивание моделей в MVSepLess"
+    )
+    mvsepless_parser.add_argument("--model_type", required=True, help="Тип модели")
+    mvsepless_parser.add_argument("--model_name", required=True, help="Имя модели")
+
+    vbach_parser = subparsers.add_parser(
+        "vbach", help="Установка голосовых моделей в Vbach"
+    )
+    vbach_subparsers = vbach_parser.add_subparsers(
+        title="vbach_commands", dest="vbach_command", required=True
+    )
+
+    install_local_parser = vbach_subparsers.add_parser(
+        "install_local", help="Установка голосовой модели по локальным файлам"
+    )
+    install_local_parser.add_argument(
+        "--model_name", required=True, help="Имя голосовой модели"
+    )
+    install_local_parser.add_argument("--pth", required=True, help="Путь к *.pth файлу")
+    install_local_parser.add_argument(
+        "--index", required=False, help="Путь к *.index файлу"
+    )
+
+    install_url_zip_parser = vbach_subparsers.add_parser(
+        "install_url_zip", help="Установка голосовой модели по URL (архив с файлами)"
+    )
+    install_url_zip_parser.add_argument(
+        "--model_name", required=True, help="Имя голосовой модели"
+    )
+    install_url_zip_parser.add_argument("--url", required=True, help="URL *.zip файла")
+
+    install_url_files_parser = vbach_subparsers.add_parser(
+        "install_url_files", help="Установка голосовой модели по URL (отдельные файлы)"
+    )
+    install_url_files_parser.add_argument(
+        "--model_name", required=True, help="Имя голосовой модели"
+    )
+    install_url_files_parser.add_argument(
+        "--pth_url", required=True, help="URL *.pth файла"
+    )
+    install_url_files_parser.add_argument(
+        "--index_url", required=False, help="URL *.index файла"
+    )
+
+    list_parser = vbach_subparsers.add_parser(
+        "list", help="Список установленных моделей"
+    )
+
+    remove_voice_model = vbach_subparsers.add_parser("remove", help="Удаление модели")
+    remove_voice_model.add_argument(
+        "--model_name", required=True, help="Имя голосовой модели"
+    )
+
+    args = parser.parse_args()
+
+    if args.command == "mvsepless":
+
+        _model_manager = MvseplessModelManager()
+        info = _model_manager.models_info.get(args.model_type).get(args.model_name, None)
+        if not info:
+            raise ValueError(
+                f"Модель {args.model_name} не найдена для типа {args.model_type}"
+            )
+        conf, ckpt = _model_manager.download_model(
+            _model_manager.models_cache_dir,
+            args.model_name,
+            args.model_type,
+            info["checkpoint_url"],
+            info["config_url"],
+        )
+
+    elif args.command == "vbach":
+        model_manager = VbachModelManager()
+
+        if args.vbach_command == "install_local":
+            status = model_manager.install_model_files(
+                args.index, args.pth, args.model_name, mode="local"
+            )
+            print(status)
+
+        elif args.vbach_command == "install_url_zip":
+            status = model_manager.install_model_zip(
+                args.url, args.model_name, mode="url"
+            )
+            print(status)
+
+        elif args.vbach_command == "install_url_files":
+            status = model_manager.install_model_files(
+                args.index_url, args.pth_url, args.model_name, mode="url"
+            )
+            print(status)
+
+        elif args.vbach_command == "list":
+            models = model_manager.parse_voice_models()
+            if models:
+                print("Установленные модели:")
+                for model in models:
+                    print(f"  - {model}")
+            else:
+                print("Нет установленных моделей")
+
+        elif args.vbach_command == "remove":
+            status = model_manager.del_voice_model(args.model_name)
+            print(status)

mvsepless/models/bandit/core/__init__.py

@@ -1,691 +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 deepspeed.ops.adam import DeepSpeedCPUAdam
-# from geoopt import optim as gooptim
-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
-
-# from pandas.io.json._normalize import nested_to_record
-
-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__:
-            # print(config["kwargs"])
-            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"
-
-            # print(f"matching for {stem}")
-            if mode == "train":
-                metric.update(
-                    output["audio"][stem],  # .cpu(),
-                    batch["audio"][stem],  # .cpu()
-                )
-            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
-
-                # print(batch["audio"].keys())
-
-                if matched:
-                    # print(f"matched {stem}!")
-                    if stem == "mne" and "mne" not in output["audio"]:
-                        output["audio"]["mne"] = (
-                            output["audio"]["music"] + output["audio"]["effects"]
-                        )
-
-                    metric.update(
-                        output["audio"][stem],  # .cpu(),
-                        batch["audio"][stem],  # .cpu(),
-                    )
-
-                    # print(metric.compute())
-
-    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)
-        # print(batch)
-        # print(output)
-        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:
-                # print(f"Switching to validation dataloader {dataloader_idx}")
-                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.logger.save()
-        # print(self.val_prefix, "Validation metrics:", metric_dict)
-        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,
-            }
-        )  # TODO: support track properly
-
-        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}"
-
-        # print(batch["audio"].keys())
-
-        if curr_test_prefix != self.test_prefix:
-            # print(f"Switching to test dataloader {dataloader_idx}")
-            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)
-            # print(output)
-            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.logger.save()
-        # print(self.test_prefix, "Test metrics:", metric_dict)
-        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
-            # 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
-
-        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,68 +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,
-    # augmentor=None
-) -> 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)
-
-    # if augmentor is not None:
-    #     train_dataset = AugmentedDataset(train_dataset, augmentor)
-
-    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,366 +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:
-            # noinspection PyUnresolvedReferences
-            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))
-        ]
-        # pprint(list(enumerate(files)))
-        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)
-        # self.index_lock = index
-        audio = self.get_audio(identifier)
-        # self.index_lock = None
-
-        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))
-        ]
-        # pprint(list(enumerate(files)))
-        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
-
-        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,273 +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)
-        # noinspection PyUnresolvedReferences
-
-        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(".")]
-        # pprint(list(enumerate(files)))
-        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
-
-        # if np.random.rand() > self.apply_probability:
-        #     return super().__getitem__(index)
-
-        audio = {}
-        identifier = self.get_identifier(index)
-
-        # assert self.target_stem in self.stems_no_mixture
-        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 stem == self.target_stem:
-
-            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:
-                # db_scale = "-inf"
-                linear_scale = 0.0
-            else:
-                db_scale = np.random.uniform(*self.random_scale_range_db)
-                linear_scale = np.power(10, db_scale / 20)
-                # db_scale = f"{db_scale:+2.1f}"
-            # print(linear_scale)
-            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}"}
-
-
-# if __name__ == "__main__":
-#
-#     from pprint import pprint
-#     from tqdm.auto import tqdm
-#
-#     for split_ in ["train", "val", "test"]:
-#         ds = MUSDB18SadOnTheFlyAugmentedDataset(
-#             data_root="$DATA_ROOT/MUSDB18/HQ/saded",
-#             split=split_,
-#             target_stem="vocals"
-#         )
-#
-#         print(split_, len(ds))
-#
-#         for track_ in tqdm(ds):
-#             track_["audio"] = {
-#                 k: v.shape for k, v in track_["audio"].items()
-#             }
-#         pprint(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,226 +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
-
-                # ta.save(os.path.join(outpath, f"{stem}.wav"), segment, self.fs)
-
-                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,
-    # output_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,95 +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
-
-    # TODO: add better type hints
-    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,139 +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)
-        # 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"
-            )
-        # Step 1. Zero-mean norm
-        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
-        # Step 2. Pair-wise SI-SDR.
-        if self.sdr_type in ["sisdr", "sdsdr"]:
-            # [batch, 1]
-            dot = torch.sum(est_target * target, dim=[1, 2], keepdim=True)
-            # [batch, 1]
-            s_target_energy = torch.sum(target**2, dim=[1, 2], keepdim=True) + self.EPS
-            # [batch, time]
-            scaled_target = dot * target / s_target_energy
-        else:
-            # [batch, time]
-            scaled_target = target
-        if self.sdr_type in ["sdsdr", "snr"]:
-            e_noise = est_target - target
-        else:
-            e_noise = est_target - scaled_target
-        # [batch]
-
-        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,9 +1,7 @@
-from .snr import (
-    ChunkMedianScaleInvariantSignalDistortionRatio,
-    ChunkMedianScaleInvariantSignalNoiseRatio,
-    ChunkMedianSignalDistortionRatio,
-    ChunkMedianSignalNoiseRatio,
-    SafeSignalDistortionRatio,
-)
-
-# from .mushra import EstimatedMushraScore
+from .snr import (
+    ChunkMedianScaleInvariantSignalDistortionRatio,
+    ChunkMedianScaleInvariantSignalNoiseRatio,
+    ChunkMedianSignalDistortionRatio,
+    ChunkMedianSignalNoiseRatio,
+    SafeSignalDistortionRatio,
+)

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

@@ -1,443 +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:
-    """The metric defined by ITU-T P.862 is often called 'PESQ score', which is defined
-    for narrow-band signals and has a value range of [-0.5, 4.5] exactly. Here, we use the metric
-    defined by ITU-T P.862.2, commonly known as 'wide-band PESQ' and will be referred to as "PESQ score".
-
-    Args:
-        x (float): Narrow-band PESQ score.
-
-    Returns:
-        (float): Wide-band PESQ score.
-    """
-    return 0.999 + (4.999 - 0.999) / (1 + math.exp(-1.3669 * x + 3.8224))
-
-
-PESQRange: Tuple[float, float] = (
-    1.0,  # P.862.2 uses a different input filter than P.862, and the lower bound of
-    # the raw score is not -0.5 anymore. It's hard to figure out the true lower bound.
-    # We are using 1.0 as a reasonable approximation.
-    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):
-    """Encoder module that transform 1D waveform to 2D representations.
-
-    Args:
-        feat_dim (int, optional): The feature dimension after Encoder module. (Default: 512)
-        win_len (int, optional): kernel size in the Conv1D layer. (Default: 32)
-    """
-
-    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:
-        """Apply waveforms to convolutional layer and ReLU layer.
-
-        Args:
-            x (torch.Tensor): Input waveforms. Tensor with dimensions `(batch, time)`.
-
-        Returns:
-            (torch,Tensor): Feature Tensor with dimensions `(batch, channel, frame)`.
-        """
-        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:
-        # input shape: batch, seq, dim
-        out, _ = self.rnn(x)
-        out = self.proj(out)
-        return out
-
-
-class DPRNN(nn.Module):
-    """*Dual-path recurrent neural networks (DPRNN)* :cite:`luo2020dual`.
-
-    Args:
-        feat_dim (int, optional): The feature dimension after Encoder module. (Default: 64)
-        hidden_dim (int, optional): Hidden dimension in the RNN layer of DPRNN. (Default: 128)
-        num_blocks (int, optional): Number of DPRNN layers. (Default: 6)
-        rnn_type (str, optional): Type of RNN in DPRNN. Valid options are ["RNN", "LSTM", "GRU"]. (Default: "LSTM")
-        d_model (int, optional): The number of expected features in the input. (Default: 256)
-        chunk_size (int, optional): Chunk size of input for DPRNN. (Default: 100)
-        chunk_stride (int, optional): Stride of chunk input for DPRNN. (Default: 50)
-    """
-
-    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]:
-        # input shape: (B, N, T)
-        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):
-    """Speech Quality and Intelligibility Measures (SQUIM) model that predicts **objective** metric scores
-    for speech enhancement (e.g., STOI, PESQ, and SI-SDR).
-
-    Args:
-        encoder (torch.nn.Module): Encoder module to transform 1D waveform to 2D feature representation.
-        dprnn (torch.nn.Module): DPRNN module to model sequential feature.
-        branches (torch.nn.ModuleList): Transformer branches in which each branch estimate one objective metirc score.
-    """
-
-    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]:
-        """
-        Args:
-            x (torch.Tensor): Input waveforms. Tensor with dimensions `(batch, time)`.
-
-        Returns:
-            List(torch.Tensor): List of score Tenosrs. Each Tensor is with dimension `(batch,)`.
-        """
-        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:
-    """Create branch module after DPRNN model for predicting metric score.
-
-    Args:
-        d_model (int): The number of expected features in the input.
-        nhead (int): Number of heads in the multi-head attention model.
-        metric (str): The metric name to predict.
-
-    Returns:
-        (nn.Module): Returned module to predict corresponding metric score.
-    """
-    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:
-    """Build a custome :class:`torchaudio.prototype.models.SquimObjective` model.
-
-    Args:
-        feat_dim (int, optional): The feature dimension after Encoder module.
-        win_len (int): Kernel size in the Encoder module.
-        d_model (int): The number of expected features in the input.
-        nhead (int): Number of heads in the multi-head attention model.
-        hidden_dim (int): Hidden dimension in the RNN layer of DPRNN.
-        num_blocks (int): Number of DPRNN layers.
-        rnn_type (str): Type of RNN in DPRNN. Valid options are ["RNN", "LSTM", "GRU"].
-        chunk_size (int): Chunk size of input for DPRNN.
-        chunk_stride (int or None, optional): Stride of chunk input for DPRNN.
-    """
-    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:
-    """Build :class:`torchaudio.prototype.models.SquimObjective` model with default arguments."""
-    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:
-        """Construct the SquimObjective model, and load the pretrained weight.
-
-        The weight file is downloaded from the internet and cached with
-        :func:`torch.hub.load_state_dict_from_url`
-
-        Args:
-            dl_kwargs (dictionary of keyword arguments): Passed to :func:`torch.hub.load_state_dict_from_url`.
-
-        Returns:
-            Variation of :py:class:`~torchaudio.models.SquimObjective`.
-        """
-        model = squim_objective_base()
-        model.load_state_dict(self._get_state_dict(dl_kwargs))
-        model.eval()
-        return model
-
-    @property
-    def sample_rate(self):
-        """Sample rate of the audio that the model is trained on.
-
-        :type: float
-        """
-        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,127 +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.zero_mean = zero_mean
-        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])
-
-                # print(preds.shape, chunk_snr.shape)
-
-                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,135 +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):
-        # xb = (batch, n_time, in_chan, reim * band_width)
-
-        batch, n_time, in_chan, ribw = xb.shape
-        xb = self.norm(xb.reshape(batch, n_time, in_chan * ribw))
-        # (batch, n_time, in_chan * reim * band_width)
-
-        if not self.treat_channel_as_feature:
-            xb = xb.reshape(batch, n_time, in_chan, ribw)
-            # (batch, n_time, in_chan, reim * band_width)
-
-        zb = self.fc(xb)
-        # (batch, n_time, emb_dim)
-        # OR
-        # (batch, n_time, in_chan, emb_dim_per_chan)
-
-        if not self.treat_channel_as_feature:
-            batch, n_time, in_chan, emb_dim_per_chan = zb.shape
-            # (batch, n_time, in_chan, emb_dim_per_chan)
-            zb = zb.reshape((batch, n_time, in_chan * emb_dim_per_chan))
-
-        return zb  # (batch, n_time, emb_dim)
-
-
-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
-        # list of [fstart, fend) in index.
-        # Note that fend is exclusive.
-        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):
-        # x = complex spectrogram (batch, in_chan, n_freq, n_time)
-
-        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)  # batch, in_chan, n_freq, n_time, 2
-        xr = torch.permute(xr, (0, 3, 1, 4, 2))  # batch, n_time, in_chan, 2, n_freq
-        batch, n_time, in_chan, reim, band_width = xr.shape
-        for i, nfm in enumerate(self.norm_fc_modules):
-            # print(f"bandsplit/band{i:02d}")
-            fstart, fend = self.band_specs[i]
-            xb = xr[..., fstart:fend]
-            # (batch, n_time, in_chan, reim, band_width)
-            xb = torch.reshape(xb, (batch, n_time, in_chan, -1))
-            # (batch, n_time, in_chan, reim * band_width)
-            # z.append(nfm(xb))  # (batch, n_time, emb_dim)
-            z[:, i, :, :] = nfm(xb.contiguous())
-
-        # z = torch.stack(z, dim=1)
-
-        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,651 +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):
-        # x = complex spectrogram (batch, in_chan, n_freq, n_time)
-        # print(x.shape)
-        batch, in_chan, n_freq, n_time = x.shape
-        x = torch.reshape(x, (-1, 1, n_freq, n_time))
-
-        z = self.band_split(x)  # (batch, emb_dim, n_band, n_time)
-
-        # if torch.any(torch.isnan(z)):
-        #     raise ValueError("z nan")
-
-        # print(z)
-        q = self.tf_model(z)  # (batch, emb_dim, n_band, n_time)
-        # print(q)
-
-        # if torch.any(torch.isnan(q)):
-        #     raise ValueError("q nan")
-
-        out = {}
-
-        for stem, mem in self.mask_estim.items():
-            m = mem(q, cond=cond)
-
-            # if torch.any(torch.isnan(m)):
-            #     raise ValueError("m nan", stem)
-
-            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):
-        # x = complex spectrogram (batch, in_chan, n_freq, n_time)
-        z = self.band_split(x)  # (batch, emb_dim, n_band, n_time)
-        q = self.tf_model(z)  # (batch, emb_dim, n_band, n_time)
-        m = self.mask_estim(q)  # (batch, in_chan, n_freq, n_time)
-
-        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]
-
-        # print(mm.shape, am.shape, x.shape, m.shape)
-
-        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):
-        # x.shape = (batch, n_channel, n_freq, n_time)
-        # m.shape = (kernel_freq, kernel_time, batch, n_channel, n_freq, n_time)
-
-        _, 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):
-        # x.shape = (batch, n_channel, n_freq, n_time)
-        # m.shape = (kernel_freq, kernel_time, batch, n_channel, n_freq, n_time)
-
-        s = torch.zeros_like(x)
-
-        _, n_channel, n_freq, n_time = x.shape
-        kernel_freq, kernel_time, _, _, _, _ = m.shape
-
-        # print(x.shape, 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))
-
-                # if `df` > 0:
-                #     x[:, :, :df, :] = 0
-                # elif `df` < 0:
-                #     x[:, :, df:, :] = 0
-
-                # if `dt` > 0:
-                #     x[:, :, :, :dt] = 0
-                # elif `dt` < 0:
-                #     x[:, :, :, dt:] = 0
-
-                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,351 +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):
-        # print(mb.shape)
-        batch, n_time, _ = mb.shape
-        if self.complex_mask:
-            mb = mb.reshape(
-                batch, n_time, self.in_channel, self.bandwidth, self.reim
-            ).contiguous()
-            # print(mb.shape)
-            mb = torch.view_as_complex(mb)  # (batch, n_time, in_channel, bandwidth)
-        else:
-            mb = mb.reshape(batch, n_time, self.in_channel, self.bandwidth)
-
-        mb = torch.permute(mb, (0, 2, 3, 1))  # (batch, in_channel, bandwidth, n_time)
-
-        return mb
-
-    def forward(self, qb):
-        # qb = (batch, n_time, emb_dim)
-
-        # if torch.any(torch.isnan(qb)):
-        #     raise ValueError("qb0")
-
-        qb = self.norm(qb)  # (batch, n_time, emb_dim)
-
-        # if torch.any(torch.isnan(qb)):
-        #     raise ValueError("qb1")
-
-        qb = self.hidden(qb)  # (batch, n_time, mlp_dim)
-        # if torch.any(torch.isnan(qb)):
-        #     raise ValueError("qb2")
-        mb = self.output(qb)  # (batch, n_time, bandwidth * in_channel * reim)
-        # if torch.any(torch.isnan(qb)):
-        #     raise ValueError("mb")
-        mb = self.reshape_output(mb)  # (batch, in_channel, bandwidth, n_time)
-
-        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)  # (batch, n_time, emb_dim)
-        qb = self.hidden(qb)  # (batch, n_time, mlp_dim)
-        mmb = self.output(qb)  # (batch, n_time, bandwidth * in_channel * reim)
-        mmb = self.reshape_output(mmb)  # (batch, in_channel, bandwidth, n_time)
-        amb = self.output2(qb)  # (batch, n_time, bandwidth * in_channel * reim)
-        amb = self.reshape_output(amb)  # (batch, in_channel, bandwidth, n_time)
-
-        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):
-            # print(f"maskestim/{b:02d}")
-            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)
-
-        # if cond_dim > 0:
-        #     raise NotImplementedError
-
-        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):
-        # q = (batch, n_bands, n_time, emb_dim)
-
-        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
-        )  # [n_bands  * (batch, in_channel, bandwidth, n_time)]
-
-        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):
-        # q = (batch, n_bands, n_time, emb_dim)
-
-        masks = self.compute_masks(
-            q
-        )  # [n_bands  * (batch, in_channel, bandwidth, n_time)]
-
-        # TODO: currently this requires band specs to have no gap and no overlap
-        masks = torch.concat(masks, dim=2)  # (batch, in_channel, n_freq, n_time)
-
-        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,320 +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:
-        # n_group is the size of the 2nd dim
-        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):
-        # z = (batch, n_uncrossed, n_across, emb_dim)
-
-        z0 = torch.clone(z)
-
-        # print(z.device)
-
-        if self.use_layer_norm:
-            z = self.norm(z)  # (batch, n_uncrossed, n_across, emb_dim)
-        else:
-            z = torch.permute(
-                z, (0, 3, 1, 2)
-            )  # (batch, emb_dim, n_uncrossed, n_across)
-
-            z = self.norm(z)  # (batch, emb_dim, n_uncrossed, n_across)
-
-            z = torch.permute(
-                z, (0, 2, 3, 1)
-            )  # (batch, n_uncrossed, n_across, emb_dim)
-
-        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
-            ]  # (batch * n_uncrossed, n_across, dir_rnn_dim)
-
-            z = torch.reshape(z, (batch, n_uncrossed, n_across, -1))
-            # (batch, n_uncrossed, n_across, dir_rnn_dim)
-        else:
-            # Note: this is EXTREMELY SLOW
-            zlist = []
-            for i in range(n_uncrossed):
-                zi = self.rnn(z[:, i, :, :])[0]  # (batch, n_across, emb_dim)
-                zlist.append(zi)
-
-            z = torch.stack(zlist, dim=1)  # (batch, n_uncrossed, n_across, dir_rnn_dim)
-
-        z = self.fc(z)  # (batch, n_uncrossed, n_across, emb_dim)
-
-        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):
-        # z = (batch, n_bands, n_time, emb_dim)
-
-        if self.parallel_mode:
-            for sbm_pair in self.seqband:
-                # z: (batch, n_bands, n_time, emb_dim)
-                sbm_t, sbm_f = sbm_pair[0], sbm_pair[1]
-                zt = sbm_t(z)  # (batch, n_bands, n_time, emb_dim)
-                zf = sbm_f(z.transpose(1, 2))  # (batch, n_time, n_bands, emb_dim)
-                z = zt + zf.transpose(1, 2)
-        else:
-            for sbm in self.seqband:
-                z = sbm(z)
-                z = z.transpose(1, 2)
-
-                # (batch, n_bands, n_time, emb_dim)
-                #   --> (batch, n_time, n_bands, emb_dim)
-                # OR
-                # (batch, n_time, n_bands, emb_dim)
-                #   --> (batch, n_bands, n_time, emb_dim)
-
-        q = z
-        return q  # (batch, n_bands, n_time, emb_dim)
-
-
-class ResidualTransformer(nn.Module):
-    def __init__(
-        self,
-        emb_dim: int = 128,
-        rnn_dim: int = 256,
-        bidirectional: bool = True,
-        dropout: float = 0.0,
-    ) -> None:
-        # n_group is the size of the 2nd dim
-        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
-        )  # (batch, n_uncrossed, n_across, emb_dim)
-        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 = (batch, n_bands, n_time, emb_dim)
-        z = self.norm(z)  # (batch, n_bands, n_time, emb_dim)
-
-        for sbm in self.seqband:
-            z = sbm(z)
-            z = z.transpose(1, 2)
-
-            # (batch, n_bands, n_time, emb_dim)
-            #   --> (batch, n_time, n_bands, emb_dim)
-            # OR
-            # (batch, n_time, n_bands, emb_dim)
-            #   --> (batch, n_bands, n_time, emb_dim)
-
-        q = z
-        return q  # (batch, n_bands, n_time, emb_dim)
-
-
-class ResidualConvolution(nn.Module):
-    def __init__(
-        self,
-        emb_dim: int = 128,
-        rnn_dim: int = 256,
-        bidirectional: bool = True,
-        dropout: float = 0.0,
-    ) -> None:
-        # n_group is the size of the 2nd dim
-        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):
-        # z = (batch, n_uncrossed, n_across, emb_dim)
-
-        z0 = torch.clone(z)
-
-        z = self.norm(z)  # (batch, n_uncrossed, n_across, emb_dim)
-        z = self.conv(z)  # (batch, n_uncrossed, n_across, emb_dim)
-        z = self.fc(z)  # (batch, n_uncrossed, n_across, emb_dim)
-        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 = (batch, n_bands, n_time, emb_dim)
-
-        z = torch.permute(z, (0, 3, 1, 2))  # (batch, emb_dim, n_bands, n_time)
-
-        z = self.seqband(z)  # (batch, emb_dim, n_bands, n_time)
-
-        z = torch.permute(z, (0, 2, 3, 1))  # (batch, n_bands, n_time, emb_dim)
-
-        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,525 +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):
-    # pprint(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)  # (1, n_freqs)
-        normalized_mel_fb = self.filterbank / weight_per_bin  # (n_mels, n_freqs)
-
-        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
-    # init freqs
-    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)
-
-    # calculate mel freq bins
-    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)
-
-    # create filterbank
-    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:
-    # freq bins
-    A = (1000 * np.log(10)) / (24.7 * 4.37)
-    all_freqs = torch.linspace(0, fs // 2, n_freqs)
-
-    # calculate mel freq bins
-    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
-
-    # create filterbank
-    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,829 +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):
-        # with torch.no_grad():
-        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,412 +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
-# @torch.compile
-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
-
-    # @torch.jit.script
-
-
-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
-
-    # @torch.jit.script
-    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(tqdm(chunks_in)):
-        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}
-
-    #
-    # def old_forward(
-    #         self,
-    #         audio: torch.Tensor,
-    #         model_fn: Callable[[torch.Tensor], Dict[str, torch.Tensor]],
-    # ):
-    #
-    #     n_samples = audio.shape[-1]
-    #     original_batch_size = audio.shape[0]
-    #
-    #     padded_audio, n_chunks = self.prepare(audio)
-    #
-    #     ndim = padded_audio.ndim
-    #     broadcaster = [1 for _ in range(ndim - 1)] + [self.chunk_size]
-    #
-    #     outputs = defaultdict(
-    #             lambda: torch.zeros_like(
-    #                     padded_audio, device=audio.device, dtype=torch.float64
-    #             )
-    #     )
-    #
-    #     all_chunks_out = []
-    #     len_chunks_in = []
-    #
-    #     batch_size_ = int(self.batch_size // original_batch_size)
-    #     for b in range(int(np.ceil(n_chunks / batch_size_))):
-    #         chunks_in = []
-    #         for j in range(batch_size_):
-    #             i = b * batch_size_ + j
-    #             if i == n_chunks:
-    #                 break
-    #
-    #             start = i * hop_size
-    #             end = start + self.chunk_size
-    #             chunk_in = padded_audio[..., start:end]
-    #             chunks_in.append(chunk_in)
-    #
-    #         chunks_in = torch.concat(chunks_in, dim=0)
-    #         chunks_out = model_fn(chunks_in)
-    #         all_chunks_out.append(chunks_out)
-    #         len_chunks_in.append(len(chunks_in))
-    #
-    #     for b, (chunks_out, lci) in enumerate(
-    #             zip(all_chunks_out, len_chunks_in)
-    #     ):
-    #         for stem in chunks_out:
-    #             for j in range(lci // original_batch_size):
-    #                 i = b * batch_size_ + j
-    #
-    #                 if self.fade_edge_frames:
-    #                     window = self.standard_window
-    #                 else:
-    #                     if i == 0:
-    #                         window = self.first_window
-    #                     elif i == n_chunks - 1:
-    #                         window = self.last_window
-    #                     else:
-    #                         window = self.standard_window
-    #
-    #                 start = i * hop_size
-    #                 end = start + self.chunk_size
-    #
-    #                 chunk_out = chunks_out[stem][j * original_batch_size: (j + 1) * original_batch_size,
-    #                             ...]
-    #                 contrib = window.view(*broadcaster) * chunk_out
-    #                 outputs[stem][..., start:end] = (
-    #                         outputs[stem][..., start:end] + contrib
-    #                 )
-    #
-    #     if self.fade_edge_frames:
-    #         pad_front, pad_back = self.edge_frame_pad_sizes
-    #         outputs = {k: v[..., pad_front:-pad_back] for k, v in
-    #                    outputs.items()}
-    #
-    #     outputs = {k: v[..., :n_samples].to(audio.dtype) for k, v in
-    #                outputs.items()}
-    #
-    #     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)
-
-        # using nn.Parameters allows lightning to take care of devices for us
-        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)
-        # print(f"hop multiplier: {self.hop_multiplier}")
-
-        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,  # dtype=torch.float64
-            )
-            / 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,363 +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"):
-        # Model takes mono as input we give stereo, so we do process of each channel independently
-        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:
-                # We need to obtain stereo again
-                r = batch["estimates"][s]["audio"].view(
-                    -1, init_shape[1], init_shape[2]
-                )
-                b.append(r)
-            # And we need to return back tensor and not independent stems
-            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)  # (batch, emb_dim, n_band, n_time)
-        q = self.tf_model(z)  # (batch, emb_dim, n_band, n_time)
-
-        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,130 +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
-        # list of [fstart, fend) in index.
-        # Note that fend is exclusive.
-        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):
-        # x = complex spectrogram (batch, in_chan, n_freq, n_time)
-
-        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,281 +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):
-        # print(mb.shape)
-        batch, n_time, _ = mb.shape
-        if self.complex_mask:
-            mb = mb.reshape(
-                batch, n_time, self.in_channels, self.bandwidth, self.reim
-            ).contiguous()
-            # print(mb.shape)
-            mb = torch.view_as_complex(mb)  # (batch, n_time, in_channels, bandwidth)
-        else:
-            mb = mb.reshape(batch, n_time, self.in_channels, self.bandwidth)
-
-        mb = torch.permute(mb, (0, 2, 3, 1))  # (batch, in_channels, bandwidth, n_time)
-
-        return mb
-
-    def forward(self, qb):
-        # qb = (batch, n_time, emb_dim)
-        # qb = self.norm(qb)  # (batch, n_time, emb_dim)
-        # qb = self.hidden(qb)  # (batch, n_time, mlp_dim)
-        # mb = self.output(qb)  # (batch, n_time, bandwidth * in_channels * reim)
-
-        mb = checkpoint_sequential(self.combined, 2, qb, use_reentrant=False)
-        mb = self.reshape_output(mb)  # (batch, in_channels, bandwidth, n_time)
-
-        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):
-            # print(f"maskestim/{b:02d}")
-            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):
-        # q = (batch, n_bands, n_time, emb_dim)
-
-        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):
-        # q = (batch, n_bands, n_time, emb_dim)
-
-        masks = self.compute_masks(
-            q
-        )  # [n_bands  * (batch, in_channels, bandwidth, n_time)]
-
-        # TODO: currently this requires band specs to have no gap and no overlap
-        masks = torch.concat(masks, dim=2)  # (batch, in_channels, n_freq, n_time)
-
-        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,145 +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:
-        # n_group is the size of the 2nd dim
-        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):
-        # z = (batch, n_uncrossed, n_across, emb_dim)
-
-        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)  # (batch, n_uncrossed, n_across, emb_dim)
-
-        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):
-        # z = (batch, n_bands, n_time, emb_dim)
-
-        if self.parallel_mode:
-            for sbm_pair in self.seqband:
-                # z: (batch, n_bands, n_time, emb_dim)
-                sbm_t, sbm_f = sbm_pair[0], sbm_pair[1]
-                zt = sbm_t(z)  # (batch, n_bands, n_time, emb_dim)
-                zf = sbm_f(z.transpose(1, 2))  # (batch, n_time, n_bands, emb_dim)
-                z = zt + zf.transpose(1, 2)
-        else:
-            z = checkpoint_sequential(
-                self.seqband, self.n_modules, z, use_reentrant=False
-            )
-
-        q = z
-        return q  # (batch, n_bands, n_time, emb_dim)
+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,523 +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
-
-# from spafe.fbanks import bark_fbanks
-# from spafe.utils.converters import erb2hz, hz2bark, hz2erb
-
-
-def band_widths_from_specs(band_specs):
-    return [e - i for i, e in band_specs]
-
-
-def check_nonzero_bandwidth(band_specs):
-    # pprint(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)  # (1, n_freqs)
-        normalized_mel_fb = self.filterbank / weight_per_bin  # (n_mels, n_freqs)
-
-        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
-    # init freqs
-    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)
-
-#     # calculate mel freq bins
-#     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)
-
-#     # create filterbank
-#     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:
-#     # freq bins
-#     A = (1000 * np.log(10)) / (24.7 * 4.37)
-#     all_freqs = torch.linspace(0, fs // 2, n_freqs)
-
-#     # calculate mel freq bins
-#     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
-
-#     # create filterbank
-#     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 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,6 +1,6 @@
-from .bs_roformer import BSRoformer
-from .bs_roformer_sw import BSRoformer_SW
-from .bs_roformer_fno import BSRoformer_FNO
-from .bs_roformer_hyperace import BSRoformerHyperACE
-from .bs_roformer_hyperace2 import BSRoformerHyperACE_2
-from .mel_band_roformer import MelBandRoformer
+from .bs_roformer import BSRoformer
+from .bs_roformer_sw import BSRoformer_SW
+from .bs_roformer_fno import BSRoformer_FNO
+from .bs_roformer_hyperace import BSRoformerHyperACE
+from .bs_roformer_hyperace2 import BSRoformerHyperACE_2
+from .mel_band_roformer import MelBandRoformer

mvsepless/models/bs_roformer/attend.py

@@ -1,126 +1,120 @@
-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
-
-# constants
-
-FlashAttentionConfig = namedtuple('FlashAttentionConfig', ['enable_flash', 'enable_math', 'enable_mem_efficient'])
-
-# helpers
-
-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)
-
-# main class
-
-class Attend(nn.Module):
-    def __init__(
-        self,
-        dropout = 0.,
-        flash = 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')), 'in order to use flash attention, you must be using pytorch 2.0 or above'
-
-        # determine efficient attention configs for cuda and cpu
-
-        self.cpu_config = FlashAttentionConfig(True, True, True)
-        self.cuda_config = 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':
-                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)
-
-    def flash_attn(self, q, k, v):
-        _, heads, q_len, _, k_len, is_cuda, device = *q.shape, k.shape[-2], q.is_cuda, q.device
-
-        if exists(self.scale):
-            default_scale = q.shape[-1] ** -0.5
-            q = q * (self.scale / default_scale)
-
-        # Check if there is a compatible device for flash attention
-
-        config = self.cuda_config if is_cuda else self.cpu_config
-
-        # pytorch 2.0 flash attn: q, k, v, mask, dropout, softmax_scale
-
-        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.
-            )
-
-        return out
-
-    def forward(self, q, k, v):
-        """
-        einstein notation
-        b - batch
-        h - heads
-        n, i, j - sequence length (base sequence length, source, target)
-        d - feature dimension
-        """
-
-        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:
-            return self.flash_attn(q, k, v)
-
-        # similarity
-
-        sim = einsum(f"b h i d, b h j d -> b h i j", q, k) * scale
-
-        # attention
-
-        attn = sim.softmax(dim=-1)
-        attn = self.attn_dropout(attn)
-
-        # aggregate values
-
-        out = einsum(f"b h i j, b h j d -> b h i d", attn, v)
-
-        return out
+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
+        assert not (
+            flash and version.parse(torch.__version__) < version.parse("2.0.0")
+        ), "in order to use flash attention, you must be using pytorch 2.0 or above"
+
+        self.cpu_config = FlashAttentionConfig(True, True, True)
+        self.cuda_config = 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":
+                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)
+
+    def flash_attn(self, q, k, v):
+        _, heads, q_len, _, k_len, is_cuda, device = (
+            *q.shape,
+            k.shape[-2],
+            q.is_cuda,
+            q.device,
+        )
+
+        if exists(self.scale):
+            default_scale = q.shape[-1] ** -0.5
+            q = q * (self.scale / default_scale)
+
+        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:
+            return self.flash_attn(q, k, v)
+
+        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