Upload 101 files

mvsepless/__init__.py

@@ -0,0 +1,1579 @@
+import os
+import sys
+import shutil
+import logging
+import zipfile
+import importlib.util
+from pathlib import Path
+logging.basicConfig(level=logging.WARNING)
+
+script_dir = os.path.dirname(os.path.abspath(__file__))
+os.chdir(script_dir)
+
+if not __package__:
+    from model_manager import MvseplessModelManager
+    from audio import Audio, Inverter
+    from namer import Namer
+    from vbach_infer import vbach_inference, model_manager as VbachModel
+    from downloader import dw_file, dw_yt_dlp
+    from ensemble import ensemble_audio_files
+else:
+    from .model_manager import MvseplessModelManager
+    from .audio import Audio, Inverter
+    from .namer import Namer
+    from .vbach_infer import vbach_inference, model_manager as VbachModel
+    from .downloader import dw_file, dw_yt_dlp
+    from .ensemble import ensemble_audio_files
+from typing import Literal
+import gradio as gr
+import pandas as pd
+import subprocess
+import json
+import threading
+import queue
+import time
+import argparse
+from datetime import datetime
+import tempfile
+import ast
+
+class MVSEPLESS:
+    audio = Audio()
+    inverter = Inverter()
+    namer = Namer()
+    model_manager = MvseplessModelManager()
+    vbach_model_manager = VbachModel
+
+class Separator(MVSEPLESS):
+
+    class OutputReader:
+        def __init__(self, debug=False):
+            self.debug = debug
+
+        def parse_json_line(self, line):
+            try:
+                return json.loads(line)
+            except json.JSONDecodeError:
+                return None
+
+        def reaction_line(self, line, progress, add_text):
+            _add_text = ""
+            if add_text != "" or add_text is not None:
+                _add_text = f"| {add_text}"
+
+            data = self.parse_json_line(line)
+            if data is None:
+                return None
+            elif "reading" in data:
+                progress(0.05, desc=f"Чтение файла {_add_text}")
+                print("Чтение файла")
+                return None
+            elif "processing" in data:
+                progress_a = data["processing"]
+                processed = progress_a.get("processed", 0)
+                total = progress_a.get("total", 1)
+                # Исправлено: убрано деление на ноль
+                if total > 0:
+                    progress_ratio = min(0.89, 0.05 + (processed / total * 0.85))  # Оставляем место для этапа записи
+                    percent = int((processed / total) * 100)
+                    progress(progress_ratio, desc=f"Обработано: {percent}% {_add_text}")
+                    print(f"\rОбработано: {percent}%", end="")
+                return None
+            elif "writing" in data:
+                progress(0.9, desc="Запись результатов")
+                print(f"\rЗапись в файл {data['writing']}", end="")
+                return None
+            elif "done" in data:
+                progress(1.0, desc=f"Завершено {_add_text}")
+                print("\rЗавершено", end="\n")
+                return data["done"]
+            elif "error" in data:
+                raise Exception(data["error"])
+
+        def read_stream_to_queue(self, stream, queue_obj, stream_name):
+            """Чтение потока вывода подпроцесса и запись в очередь"""
+            try:
+                for line in iter(stream.readline, ''):
+                    line = line.strip()
+                    if line:
+                        if self.debug:
+                            print(f"[{stream_name}] {line}")  # Отладочный вывод
+                        queue_obj.put(line)
+                stream.close()
+            except Exception as e:
+                print(f"Error reading {stream_name}: {e}")
+
+    output_reader = OutputReader()
+
+    def separator_model_loader(self, model_type: str, model_name: str, mdx_denoise: bool, vr_aggr: bool, progress) -> tuple[int, str, str]:
+
+        if model_type in [
+            "mel_band_roformer",
+            "bs_roformer",
+            "mdx23c",
+            "mdxnet",
+            "vr",
+            "scnet",
+            "htdemucs",
+            "bandit",
+            "bandit_v2",
+        ]:
+            info = self.model_manager.models_info[model_type].get(model_name, None)
+            if not info:
+                raise ValueError(f"Модель {model_name} не найдена для типа {model_type}")
+
+            id = self.model_manager.get_id(model_type, model_name)
+            conf, ckpt = self.model_manager.download_model(
+                self.model_manager.models_cache_dir,
+                model_name,
+                model_type,
+                info["checkpoint_url"],
+                info["config_url"],
+            )
+            if model_type != "htdemucs":
+                self.model_manager.conf_editor(conf, mdx_denoise, vr_aggr, model_type)
+
+            return id, conf, ckpt
+
+        else:
+            raise ValueError("Неподдерживаемый тип модели")
+
+    def separator_base(
+        self,
+        input_file: str,
+        output_dir: str,
+        model_type: Literal[
+            "mel_band_roformer",
+            "bs_roformer",
+            "mdx23c",
+            "mdxnet",
+            "scnet",
+            "htdemucs",
+            "bandit",
+            "bandit_v2",
+        ] = "mel_band_roformer",
+        model_name: str = "Mel-Band-Roformer_Vocals_kimberley_jensen",
+        ext_inst: bool = True,
+        output_format: Literal["mp3", "wav", "flac", "ogg", "opus", "m4a", "aac", "ac3", "aiff"] = "mp3",
+        output_bitrate: str = "320k",
+        template: str = "NAME_(STEM)_MODEL",
+        selected_stems: list = None,
+        ckpt: str = None,
+        conf: str = None,
+        id: int = None,
+        progress: any = gr.Progress(track_tqdm=True),
+        add_text_progress: str = ""
+    ) -> list[tuple[str, str]]:
+
+        if model_type in [
+            "mel_band_roformer",
+            "bs_roformer",
+            "mdx23c",
+            "mdxnet",
+            "vr",
+            "scnet",
+            "htdemucs",
+            "bandit",
+            "bandit_v2",
+        ]:
+
+            cmd = [
+                os.sys.executable,
+                "-m",
+                "infer",
+                "--input", input_file,
+                "--store_dir", output_dir,
+                "--model_type", model_type,
+                "--model_name", model_name,
+                "--model_id", str(id),
+                "--config_path", conf,
+                "--start_check_point", ckpt,
+                "--output_format", output_format,
+                "--output_bitrate", str(output_bitrate),
+                "--template", template
+            ]
+            if ext_inst:
+                cmd.append("--extract_instrumental")
+            if selected_stems:
+                cmd.append("--selected_instruments")
+                cmd.extend(selected_stems)
+
+            try:
+                process = subprocess.Popen(
+                    cmd,
+                    stdout=subprocess.PIPE,
+                    stderr=subprocess.PIPE,
+                    text=True,
+                    bufsize=1,
+                    universal_newlines=True,
+                    encoding='utf-8',
+                    errors='replace'
+                )
+                
+                # Создаем очереди для stdout и stderr
+                stdout_queue = queue.Queue()
+                stderr_queue = queue.Queue()
+                
+                # Запускаем потоки для чтения stdout и stderr
+                stdout_thread = threading.Thread(
+                    target=self.output_reader.read_stream_to_queue, 
+                    args=(process.stdout, stdout_queue, "stdout")
+                )
+                stderr_thread = threading.Thread(
+                    target=self.output_reader.read_stream_to_queue, 
+                    args=(process.stderr, stderr_queue, "stderr")
+                )
+                
+                stdout_thread.daemon = True
+                stderr_thread.daemon = True
+                
+                stdout_thread.start()
+                stderr_thread.start()
+                
+                results = {'output': None, 'error': None}
+                process_completed = False
+                
+                # Основной цикл обработки сообщений
+                while not process_completed:
+                    # Проверяем завершение процесса
+                    if process.poll() is not None:
+                        process_completed = True
+                    
+                    # Обрабатываем сообщения из stdout
+                    try:
+                        stdout_line = stdout_queue.get_nowait()
+                        result = self.output_reader.reaction_line(stdout_line, progress, add_text_progress)
+                        if result is not None:
+                            results['output'] = result
+                            break
+                    except queue.Empty:
+                        pass
+                    
+                    # Обрабатываем сообщения из stderr
+                    try:
+                        stderr_line = stderr_queue.get_nowait()
+                        result = self.output_reader.reaction_line(stderr_line, progress, add_text_progress)
+                        if result is not None:
+                            results['output'] = result
+                            break
+                    except queue.Empty:
+                        pass
+                    
+                    # Если процесс еще работает, ждем немного перед следующей проверкой
+                    if not process_completed:
+                        time.sleep(0.1)
+                
+                # Дополнительная обработка оставшихся сообщений после завершения процесса
+                for _ in range(10):  # Проверяем несколько раз на случай задержек
+                    try:
+                        stdout_line = stdout_queue.get_nowait()
+                        result = self.output_reader.reaction_line(stdout_line, progress, add_text_progress)
+                        if result is not None:
+                            results['output'] = result
+                            break
+                    except queue.Empty:
+                        pass
+                        
+                    try:
+                        stderr_line = stderr_queue.get_nowait()
+                        result = self.output_reader.reaction_line(stderr_line, progress, add_text_progress)
+                        if result is not None:
+                            results['output'] = result
+                            break
+                    except queue.Empty:
+                        pass
+                    
+                    time.sleep(0.1)
+                
+                # Проверяем результаты
+                if results.get('error'):
+                    raise Exception(results['error'])
+                    
+                if results.get('output'):
+                    return results['output']
+                    
+                # Если процесс завершился с ошибкой
+                if process.returncode != 0:
+                    # Пытаемся получить последние сообщения об ошибках
+                    error_messages = []
+                    try:
+                        while True:
+                            error_msg = stderr_queue.get_nowait()
+                            error_messages.append(error_msg)
+                    except queue.Empty:
+                        pass
+                    
+                    error_text = "\n".join(error_messages[-5:])  # Последние 5 сообщений
+                    raise Exception(f"Процесс завершился с ошибкой. Код возврата: {process.returncode}. Сообщения об ошибках:\n{error_text}")
+                    
+            except Exception as e:
+                raise e
+            finally:
+                # Гарантируем завершение процесса
+                try:
+                    if process.poll() is None:
+                        process.terminate()
+                        process.wait(timeout=5)
+                except:
+                    try:
+                        process.kill()
+                    except:
+                        pass
+        else:
+            raise ValueError("Неподдерживаемый тип модели")
+
+    def separate(
+        self,
+        input: str | list = None,
+        output_dir: str = None,
+        model_type: Literal[
+            "mel_band_roformer",
+            "bs_roformer",
+            "mdx23c",
+            "mdxnet",
+            "vr",
+            "scnet",
+            "htdemucs",
+            "bandit",
+            "bandit_v2",
+        ] = "mel_band_roformer",
+        model_name: str = "Mel-Band-Roformer_Vocals_kimberley_jensen",
+        ext_inst: bool = True,
+        output_format: Literal["mp3", "wav", "flac", "ogg", "opus", "m4a", "aac", "ac3", "aiff"] = "mp3",
+        output_bitrate: str = "320k",
+        template: str = "NAME_(STEM)_MODEL",
+        selected_stems: list = None,
+        add_settings: dict = {"mdx_denoise": False, "vr_aggr": 5, "add_single_sep_text_progress": None},
+        progress: any = gr.Progress(track_tqdm=True)
+    ) -> list[tuple[str, str]] | list[str, list[tuple[str, str]]]:
+
+        progress(0, desc="Начало обработки")
+
+        # Валидация параметров
+        if output_format not in self.audio.output_formats:
+            output_format = "flac"
+
+        if output_dir is None:
+            output_dir = os.getcwd()
+
+        if output_dir:
+            output_dir = os.path.abspath(output_dir)
+
+        if selected_stems is None:
+            selected_stems = []
+
+        if not input:
+            raise ValueError("Входной файл не указан")
+        
+        if "STEM" not in template and template is not None:
+            template = template + "_STEM_"
+        if not template:
+            template = "mvsepless_NAME_(STEM)"
+
+        os.makedirs(output_dir, exist_ok=True)
+
+        mdx_denoise = add_settings.get("mdx_denoise", False)
+
+        vr_aggr = add_settings.get("vr_aggr", 5)
+
+        add_progress_text_custom = add_settings.get("add_single_sep_text_progress", "")
+
+        id, conf, ckpt = self.separator_model_loader(model_type, model_name, mdx_denoise, vr_aggr, progress)
+
+        if isinstance(input, str):
+            if not os.path.exists(input):
+                raise ValueError(f"Входной файл не найден: {input}")
+            
+            if not self.audio.check(input):
+                raise ValueError("Входной файл не содержит аудио")
+            
+            basename = os.path.splitext(os.path.basename(input))[0]
+            seped = self.separator_base(input_file=input,
+                        output_dir=output_dir,
+                        model_type=model_type,
+                        model_name=model_name,
+                        ext_inst=ext_inst,
+                        output_format=output_format,
+                        output_bitrate=output_bitrate,
+                        template=template,
+                        selected_stems=selected_stems,
+                        ckpt=ckpt,
+                        conf=conf,
+                        id=id,
+                        progress=progress,
+                        add_text_progress=add_progress_text_custom)
+            return seped
+        
+        elif isinstance(input, list):
+            results = []
+            for i, f in enumerate(input, 1):
+                print(f"Файл {i} из {len(input)}: {f}")
+                if os.path.exists(f):
+                    if self.audio.check(f):
+                        basename = os.path.splitext(os.path.basename(f))[0]
+                        seped = self.separator_base(input_file=f,
+                            output_dir=output_dir,
+                            model_type=model_type,
+                            model_name=model_name,
+                            ext_inst=ext_inst,
+                            output_format=output_format,
+                            output_bitrate=output_bitrate,
+                            template=template,
+                            selected_stems=selected_stems,
+                            ckpt=ckpt,
+                            conf=conf,
+                            id=id,
+                            progress=progress,
+                            add_text_progress=f"({i} из {len(input)}) ")
+                        results.append([basename, seped])
+            return results
+        
+
+    def UI(self, output_base_dir, output_temp_dir_check):
+        default_mts = self.model_manager.get_mt()
+        default_mt = self.model_manager.get_mt()[0]
+        default_mns = self.model_manager.get_mn(default_mt)
+        default_mn = default_mns[0]
+        default_stems = self.model_manager.get_stems(default_mt, default_mn)
+        default_tgt_inst = self.model_manager.get_tgt_inst(default_mt, default_mn)
+        with gr.Blocks():
+            with gr.Row():
+                with gr.Column():
+                    with gr.Group(visible=False) as add_inputs:
+                        input_path = gr.Textbox(label="Путь к входному файлу", interactive=True)
+                        add_inputs_btn = gr.Button("Добавить файл", variant="primary")
+                    with gr.Group(visible=False) as add_inputs_from_url:
+                        input_url = gr.Textbox(label="URL входного файла", interactive=True)
+                        with gr.Row():
+                            inputs_url_format = gr.Dropdown(label="Формат входного файла", interactive=True,
+                                                        choices=self.audio.output_formats,
+                                                        value="mp3", filterable=False)
+                            inputs_url_bitrate = gr.Slider(label="Битрейт входного файла", minimum=64, maximum=512, step=32, value=320, interactive=True)
+                        with gr.Row():
+                            inputs_url_cookie = gr.UploadButton(label="Файл cookie (необязательно)", interactive=True, type="filepath", file_count="single", file_types=[".txt", ".cookies"], variant="secondary")
+                            add_inputs_url_btn = gr.Button("Добавить файл", variant="primary")
+                    with gr.Row(visible=True) as add_buttons_row:
+                        add_path_btn = gr.Button("Добавить файл по пути", variant="secondary")
+                        add_url_btn = gr.Button("Добавить файл по URL", variant="secondary")
+                    with gr.Group():
+                        input_audio = gr.File(label="Входные аудио", interactive=True, type="filepath", file_count="multiple", file_types=[f".{of}" for of in self.audio.input_formats])
+                        sep_state = gr.Textbox(label="Состояние разделения", interactive=False, value="", visible=False)
+                        status = gr.Textbox(container=False, lines=3, interactive=False, max_lines=3)
+                        input_preview_check = gr.Checkbox(label="Показать плееры для входных аудио", interactive=True, value=False)
+                        @gr.render(inputs=[input_preview_check, input_audio])
+                        def show_input_players(preview, audios):
+                            if preview:
+                                if audios:
+                                    with gr.Group():
+                                        for file in audios:
+                                            gr.Audio(label=os.path.splitext(os.path.basename(file))[0], value=file, interactive=False, show_download_button=False, type="filepath")
+                with gr.Column():
+                    with gr.Group():
+                        with gr.Row():
+                            model_type = gr.Dropdown(label="Тип модели", interactive=True, filterable=False,
+                                                choices=default_mts,
+                                                value=default_mt)
+                            model_name = gr.Dropdown(label="Имя модели", interactive=True, filterable=False,
+                                                choices=default_mns, value=default_mn)
+                        with gr.Group():
+                            extract_instrumental = gr.Checkbox(label="Извлечь инструментал", interactive=True, value=True) 
+                            selected_stems = gr.CheckboxGroup(label="Выбранные стемы для разделения", interactive=False,
+                                                            choices=default_stems, value=[])
+                            
+                        with gr.Accordion(label="Дополнительные настройки", open=False):
+                            vr_aggr_slider = gr.Slider(label="Сила подавления для VR моделей", minimum=-100, maximum=100, value=5, step=1)
+                            mdx_denoise_check = gr.Checkbox(label="Включить шумоподавление для MDX-NET моделей (это повышает потребление памяти в два раза)", value=False)
+                        with gr.Row():
+                            output_format = gr.Dropdown(label="Формат выходного файла", interactive=True,
+                                                    choices=self.audio.output_formats,
+                                                    value="mp3", filterable=False)
+                            output_bitrate = gr.Slider(label="Битрейт выходного файла", minimum=64, maximum=512, step=32, value=320, interactive=True)
+
+                        template = gr.Textbox(label="Шаблон именования выходных файлов", interactive=True, value="NAME (STEM) MODEL", info="Используйте ключи: \nNAME - имя входного файла без расширения, \nSTEM - имя стема, \nMODEL - имя модели разделения")
+                        separate_btn = gr.Button("Разделить", variant="primary")
+
+            @gr.render(inputs=[sep_state], triggers=[sep_state.change])
+            def players(state):
+                def create_archive_advanced(file_list, archive_name="archive.zip"):
+                    """
+                    Создает архив с расширенной обработкой ошибок
+                    """
+                    try:
+                        print("Генерация ZIP-архива с результатами разделения...")
+                        with zipfile.ZipFile(archive_name, 'w', zipfile.ZIP_DEFLATED) as zipf:
+                            successful_files = 0
+                            
+                            for basename, stems in file_list:
+                                for stem_name, stem_path in stems:
+                                    try:
+                                        if os.path.exists(stem_path) and os.path.isfile(stem_path):
+                                            basename_ = os.path.basename(stem_path)
+                                            zipf.write(stem_path, basename_)
+                                            successful_files += 1
+                                            print(f"✓ Добавлен: {stem_path} -> {basename}")
+                                        else:
+                                            print(f"✗ Файл не найден или не является файлом: {stem_path}")
+                                            
+                                    except Exception as e:
+                                        print(f"✗ Ошибка при добавлении {stem_path}: {e}")
+                            
+                            print(f"\nАрхив создан: {archive_name}")
+                            print(f"Успешно добавлено файлов: {successful_files}")
+                            return os.path.abspath(archive_name)
+                            
+                    except Exception as e:
+                        print(f"Ошибка при создании архива: {e}")
+                if state != "":
+                    state_loaded = ast.literal_eval(state)
+                    archive_stems = create_archive_advanced(state_loaded, os.path.join(tempfile.tempdir, f"mvsepless_output_{datetime.now().strftime('%Y%m%d_%H%M%S')}.zip"))
+                    for basename, stems in state_loaded:
+                        with gr.Group():
+                            gr.Markdown(f"<h4><center>{basename}</center></h4>")
+                            for stem_name, stem_path in stems:
+                                with gr.Row(equal_height=True):
+                                    output_stem = gr.Audio(value=stem_path, label=stem_name, type="filepath", interactive=False, show_download_button=True, scale=15)
+                                    reuse_btn = gr.Button("Использовать снова", variant="secondary")
+                                    @reuse_btn.click(
+                                        inputs=[output_stem, input_audio],
+                                        outputs=input_audio
+                                    )
+                                    def reuse_fn(stem_audio, input_a):
+                                        if input_a is None:
+                                            input_a = []
+                                        if isinstance(input_a, str):
+                                            input_a = [input_a]
+                                        if os.path.exists(stem_audio):
+                                            if self.audio.check(stem_audio):
+                                                input_a.append(stem_audio)
+                                        return input_a
+                                    
+                    gr.DownloadButton(label="Скачать как ZIP", value=archive_stems, interactive=True)
+                                    
+            @add_inputs_btn.click(
+                inputs=[input_path, input_audio],
+                outputs=[add_inputs, input_audio, add_buttons_row])
+            def add_inputs_fn(input_p, input_a):
+                if input_p and os.path.exists(input_p):
+                    if input_a is None:
+                        input_a = []
+                    if isinstance(input_a, str):
+                        input_a = [input_a]
+                    if self.audio.check(input_p):
+                        input_a.append(input_p)
+                    return gr.update(visible=False), gr.update(value=input_a), gr.update(visible=True)
+                return gr.update(visible=False), gr.update(value=input_a), gr.update(visible=True)
+
+            @add_inputs_url_btn.click(
+                inputs=[input_url, input_audio, inputs_url_format, inputs_url_bitrate, inputs_url_cookie],
+                outputs=[add_inputs_from_url, input_audio, add_buttons_row])
+            def add_inputs_from_url_fn(input_u, input_a, fmt, br, cookie):
+                if input_u:
+                    if input_a is None:
+                        input_a = []
+                    if isinstance(input_a, str):
+                        input_a = [input_a]
+                    downloaded_file = dw_yt_dlp(
+                        url=input_u,
+                        output_format=fmt,
+                        output_bitrate=str(int(br)),
+                        cookie=cookie
+                    )
+                    if downloaded_file and os.path.exists(downloaded_file):
+                        if self.audio.check(downloaded_file):
+                            input_a.append(downloaded_file)
+                        return gr.update(visible=False), gr.update(value=input_a), gr.update(visible=True)
+                return gr.update(visible=False), gr.update(value=input_a), gr.update(visible=True)
+            
+            add_path_btn.click(
+                lambda: (gr.update(visible=True), gr.update(visible=False)),
+                outputs=[add_inputs, add_buttons_row])
+
+            add_url_btn.click(
+                lambda: (gr.update(visible=True), gr.update(visible=False)),
+                outputs=[add_inputs_from_url, add_buttons_row])
+
+            model_type.change(lambda x: gr.update(choices=self.model_manager.get_mn(x), value=self.model_manager.get_mn(x)[0]),
+                            inputs=model_type, outputs=model_name)
+            model_name.change(lambda mt, mn: (gr.update(choices=self.model_manager.get_stems(mt, mn), value=[], interactive=False if self.model_manager.get_tgt_inst(mt, mn) else True), gr.update(value=True if self.model_manager.get_tgt_inst(mt, mn) else False)), inputs=[model_type, model_name], outputs=[selected_stems, extract_instrumental])
+            output_format.change(lambda x: gr.update(visible=False if x in ["wav", "flac", "aiff"] else True), inputs=output_format, outputs=output_bitrate)
+            inputs_url_format.change(lambda x: gr.update(visible=False if x in ["wav", "flac", "aiff"] else True), inputs=inputs_url_format, outputs=inputs_url_bitrate)
+            @separate_btn.click(
+                inputs=[
+                    input_audio, model_type, model_name,
+                    extract_instrumental, output_format, output_bitrate,
+                    template, selected_stems, output_base_dir, output_temp_dir_check, mdx_denoise_check, vr_aggr_slider
+                ],
+                outputs=[sep_state, status],
+                show_progress="full"
+            )
+            def wrap(i, mt, mn, ei, of, ob, t, stems, o_dir, temp_save, mdx_denoise, vr_aggr, progress=gr.Progress(track_tqdm=True)):
+                if o_dir.strip() != "" and not temp_save:
+                    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+                    o = os.path.join(o_dir, f"mvsepless_outputs_{timestamp}")
+                    os.makedirs(o, exist_ok=True)
+                else:
+                    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+                    o = tempfile.mkdtemp(prefix=f"mvsepless_outputs_{timestamp}_")
+                    os.makedirs(o, exist_ok=True)
+                results = self.separate(i, o, mt, mn, ei, of, ob, t, stems, add_settings={"mdx_denoise": mdx_denoise, "vr_aggr": int(vr_aggr)}, progress=progress)
+                return str(results), ""
+        
+class AutoEnsembless(Separator):
+
+    class ModelManager(MVSEPLESS):
+        def __init__(self):
+            self.data = []
+            self.ensemble_methods = ("min_fft", "max_fft", "avg_fft", "median_fft")
+            self.ensemble_invert_methods_map = {"min_fft": "max_fft", "max_fft": "min_fft", "avg_fft": "avg_fft", "median_fft": "median_fft"}
+            self.dir_presets = os.path.join(tempfile.tempdir, "presets")
+            os.makedirs(self.dir_presets, exist_ok=True)
+
+        def save(self, name):
+            if not name:
+                name = "ensembless_preset"
+            filepath = os.path.join(self.dir_presets, f"{self.namer.short(self.namer.sanitize(name), length=50)}.json")
+            with open(filepath, "w") as f:
+                json.dump(self.data, f, indent=4, ensure_ascii=False)
+            return filepath
+
+        def load(self, filepath):
+            with open(filepath, "r") as f:
+                ensemble_data_temp = json.load(f)
+            self.data = []
+            for (mt, mn, s_stem, i_stem, weight) in ensemble_data_temp:
+                if {mt, mn} not in [{model[0], model[1]} for model in self.data]:
+                    self.data.append((mt, mn, s_stem, i_stem, weight))
+
+        def add(self, mt, mn, s_stem, i_stem, weight):
+            if {mt, mn} not in [{model[0], model[1]} for model in self.data]:
+                if s_stem and i_stem:
+                    self.data.append((mt, mn, s_stem, i_stem, weight))
+
+        def replace(self, mt, mn, s_stem, i_stem, weight, index=1):
+            if self.data:
+                len_data = len(self.data)
+                if index >= 1:
+                    if index <= len_data:
+                        self.data[index - 1] = (mt, mn, s_stem, i_stem, weight)
+                elif index == 0:
+                    self.data[0] = (mt, mn, s_stem, i_stem, weight)
+
+        def remove(self, index=1):
+            if self.data:
+                len_data = len(self.data)
+                if index >= 1:
+                    if index <= len_data:
+                        del self.data[index - 1]
+                elif index == 0:
+                    del self.data[0]
+
+        def clear(self):
+            self.data = []
+
+        def get_df(self):
+            if not self.data:
+                columns = ["#", "Имя модели", "Основной стем", "Инверсия", "Вес"]
+                return pd.DataFrame(columns=columns)
+
+            data = []
+            for i, model in enumerate(self.data):
+                data.append(
+                    [
+                        f"{i+1}",
+                        model[1],
+                        model[2],
+                        model[3],
+                        model[4],
+                    ]
+                )
+            columns = ["#", "Имя модели", "Основной стем", "Инверсия", "Вес"]
+            return pd.DataFrame(data, columns=columns)
+    
+    def UI(self, output_base_dir, output_temp_dir_check):
+        ensemble_model_manager = self.ModelManager()
+        def get_stems(mt, mn):
+            stems = []
+            for stem in self.model_manager.get_stems(mt, mn):
+                stems.append(stem)
+            
+            if not self.model_manager.get_tgt_inst(mt, mn):
+                if set(stems) == {"bass", "drums", "other", "vocals"} or set(stems) == {"bass", "drums", "other", "vocals", "piano", "guitar"}:
+                    stems.append("instrumental +")
+                    stems.append("instrumental -")
+
+            return stems
+
+        def get_invert_stems(mt, mn, s_stem):
+            orig_stems = []
+            stems = []
+            for stem in self.model_manager.get_stems(mt, mn):
+                orig_stems.append(stem)
+
+            for stem in orig_stems:
+                if stem != s_stem:
+                    stems.append(stem)
+            
+            if not self.model_manager.get_tgt_inst(mt, mn):
+                if len(orig_stems) > 2:
+                    if s_stem not in ["instrumental +", "instrumental -"]:
+                        stems.append("inverted +")
+                        stems.append("inverted -")
+
+            return stems
+
+        default_model = {
+            "mt": self.model_manager.get_mt(),
+            "mn": self.model_manager.get_mn(self.model_manager.get_mt()[0]),
+            "stem": get_stems(
+                self.model_manager.get_mt()[0],
+                self.model_manager.get_mn(self.model_manager.get_mt()[0])[0],
+            ),
+            "invert_stem": get_invert_stems(
+                self.model_manager.get_mt()[0],
+                self.model_manager.get_mn(self.model_manager.get_mt()[0])[0],
+                "vocals",
+            ),
+            "weight": 1,
+        }
+
+        gr.Markdown("<h3>Пресет</h3>")
+        with gr.Group():
+            with gr.Row(equal_height=True):
+                export_preset_name = gr.Textbox(
+                    label="Имя пресета",
+                    interactive=True,
+                    value="ensembless_preset", scale=9
+                )
+                export_btn = gr.DownloadButton("Экспорт", variant="secondary", scale=3, interactive=True)
+                import_btn = gr.UploadButton(
+                    "Импорт", file_types=[".json"], file_count="single", scale=3, interactive=True
+                )
+        gr.Markdown("<h3>Ансамбль</h3>")
+        with gr.Row():
+            with gr.Column(scale=3): # логика добавлеия моделей
+                model_type = gr.Dropdown(label="Тип модели", choices=default_model["mt"], value=default_model["mt"][0], interactive=True, filterable=False)
+                model_name = gr.Dropdown(label="Имя модели", choices=default_model["mn"], value=default_model["mn"][0], interactive=True, filterable=False)
+                primary_stem = gr.Dropdown(label="Основной стем", choices=default_model["stem"], value=default_model["stem"][0], interactive=True, filterable=False)
+                secondary_stem = gr.Dropdown(label="Инверсия", choices=default_model["invert_stem"], value=default_model["invert_stem"][0], interactive=True, filterable=False)
+                weight = gr.Slider(label="Вес", minimum=0, maximum=10, step=0.01, value=1, interactive=True)
+                @model_type.change(
+                    inputs=[model_type],
+                    outputs=[model_name]
+                )
+                def update_model_names(mt):
+                    model_names = self.model_manager.get_mn(mt)
+                    new_mn = model_names[0] if model_names else ""
+                    
+                    return gr.update(choices=model_names, value=new_mn)
+                @model_name.change(
+                    inputs=[model_type, model_name],
+                    outputs=[primary_stem, secondary_stem]
+                )
+                def update_stems_after_model_change(mt, mn):
+                    stems = get_stems(mt, mn)
+                    invert_stems = get_invert_stems(mt, mn, stems[0]) if stems else []
+                    
+                    new_s_stem = stems[0] if stems else ""
+                    new_i_stem = invert_stems[0] if invert_stems else ""
+                    
+                    return (
+                        gr.update(choices=stems, value=new_s_stem),
+                        gr.update(choices=invert_stems, value=new_i_stem)
+                    )
+                @primary_stem.change(
+                    inputs=[model_type, model_name, primary_stem],
+                    outputs=[secondary_stem]
+                )
+                def update_invert_stems(mt, mn, s_stem):
+                    stems = get_invert_stems(mt, mn, s_stem)
+                    new_i_stem = stems[0] if stems else ""
+                    return gr.update(choices=stems, value=new_i_stem)
+                
+                model_add_button = gr.Button("Добавить", interactive=True)
+            with gr.Column(scale=10):
+                df = gr.DataFrame(
+                    value=ensemble_model_manager.get_df(),
+                    headers=["#", "Имя модели", "Основной стем", "Инверсия", "Вес"],
+                    datatype=["number", "str", "str", "str", "number"],
+                    interactive=False
+                )
+
+                with gr.Group():
+                    with gr.Row(equal_height=True):
+                        with gr.Column():
+                            model_index = gr.Number(label="Индекс модели", value=1, interactive=True)
+                            model_clear_btn = gr.Button("Очистить", variant="stop", interactive=True)
+                        with gr.Column():
+                            model_replace_btn = gr.Button("Заменить", variant="primary", interactive=True)
+                            model_delete_btn = gr.Button("Удалить", variant="stop", interactive=True)
+
+                @model_add_button.click(
+                    inputs=[model_type, model_name, primary_stem, secondary_stem, weight],
+                    outputs=df
+                )
+                def add_model_to_auto_ensemble(mt, mn, s_stem, i_stem, weight):
+                    ensemble_model_manager.add(mt, mn, s_stem, i_stem, weight)
+                    return ensemble_model_manager.get_df()
+
+                @model_replace_btn.click(
+                    inputs=[model_type, model_name, primary_stem, secondary_stem, weight, model_index],
+                    outputs=df
+                )
+                def replace_model_to_auto_ensemble(mt, mn, s_stem, i_stem, weight, index):
+                    ensemble_model_manager.replace(mt, mn, s_stem, i_stem, weight, index)
+                    return ensemble_model_manager.get_df()
+
+                @model_delete_btn.click(
+                    inputs=[model_index],
+                    outputs=df
+                )
+                def delete_model_to_auto_ensemble(index):
+                    ensemble_model_manager.remove(index)
+                    return ensemble_model_manager.get_df()
+
+                @model_clear_btn.click(
+                    outputs=df
+                )
+                def clear_model_to_auto_ensemble():
+                    ensemble_model_manager.clear()
+                    return ensemble_model_manager.get_df()
+
+                gr.on(fn=ensemble_model_manager.get_df, outputs=df)
+
+                df.change(
+                    fn=ensemble_model_manager.save,
+                    inputs=export_preset_name,
+                    outputs=export_btn
+                )
+
+                export_preset_name.change(
+                    fn=ensemble_model_manager.save,
+                    inputs=export_preset_name,
+                    outputs=export_btn
+                )
+
+                @import_btn.upload(
+                    inputs=import_btn,
+                    outputs=df
+                )
+                def load_ensemble_preset(filepath):
+                    ensemble_model_manager.load(filepath)
+                    return ensemble_model_manager.get_df()
+        with gr.Row():
+            with gr.Column():
+                gr.Markdown("<h3>Входное аудио</h3>")
+                with gr.Group():
+                    with gr.Group(visible=False) as add_inputs:
+                        input_path = gr.Textbox(label="Путь к входному файлу", interactive=True)
+                        add_inputs_btn = gr.Button("Загрузить файл", variant="primary")
+                    with gr.Group(visible=False) as add_inputs_from_url:
+                        input_url = gr.Textbox(label="URL входного файла", interactive=True)
+                        with gr.Row():
+                            inputs_url_format = gr.Dropdown(label="Формат входного файла", interactive=True,
+                                                        choices=self.audio.output_formats,
+                                                        value="mp3", filterable=False)
+                            inputs_url_bitrate = gr.Slider(label="Битрейт входного файла", minimum=64, maximum=512, step=32, value=320, interactive=True)
+                        with gr.Row():
+                            inputs_url_cookie = gr.UploadButton(label="Файл cookie (необязательно)", interactive=True, type="filepath", file_count="single", file_types=[".txt", ".cookies"], variant="secondary")
+                            add_inputs_url_btn = gr.Button("Загрузить файл", variant="primary")
+                    with gr.Row(visible=True) as add_buttons_row:
+                        add_path_btn = gr.Button("Загрузить файл по пути", variant="secondary")
+                        add_url_btn = gr.Button("Загрузить файл по URL", variant="secondary")
+                    with gr.Group():
+                        input_audio = gr.File(label="Входное аудио", interactive=True, type="filepath", file_count="single", file_types=[f".{of}" for of in self.audio.input_formats])
+            with gr.Column():
+                gr.Markdown("<h3>Настройки</h3>")
+                with gr.Group():
+                    method = gr.Dropdown(label="Алгоритм склеивания", choices=["min_fft", "max_fft", "avg_fft", "median_fft"], value="avg_fft", filterable=False)
+                    invert_ensemble = gr.Checkbox(label="Инверсия ансамбля", interactive=True, value=False)
+                    output_format = gr.Dropdown(label="Формат выходного файла", interactive=True,
+                                            choices=self.audio.output_formats,
+                                            value="mp3", filterable=False)
+                    run_btn = gr.Button("Создать ансамбль", variant="primary", interactive=True)
+
+        with gr.Row():
+            with gr.Column():
+                gr.Markdown("<h3>Результаты</h3>")
+                output_audio = gr.Audio(label="Результат", type="filepath", interactive=False, show_download_button=True)
+                output_audio_wav = gr.Textbox(label="Результат в WAV", interactive=False, visible=False)
+                with gr.Group():
+                    invert_method = gr.Radio(
+                        choices=["waveform", "spectrogram"],
+                        label="Метод создания инверсии",
+                        value="waveform",
+                    )
+                    invert_btn = gr.Button("Инвертировать")
+                output_inverted_audio = gr.Audio(label="Инверсия", type="filepath", interactive=False, show_download_button=True)
+                @invert_btn.click(inputs=[input_audio, output_audio_wav, invert_method, output_format], outputs=[output_inverted_audio])
+                def invert_result_ensemble(input_file, output_file, method, out_format):
+                    if input_file and output_file:
+                        o_dir = os.path.dirname(output_file)
+                        basename = os.path.splitext(os.path.basename(input_file))[0]
+                        output_path = os.path.join(o_dir, f"ensembless_{self.namer.short(basename, length=50)}_{method}_invert.{out_format}")
+                        inverted = self.inverter.process_audio(audio1_path=input_file, audio2_path=output_file, out_format=out_format, method=method, output_path=output_path)
+                        return inverted
+                    else:
+                        return None
+
+            with gr.Column():
+                gr.Markdown("<h3>Исходники ансамбля (WAV)</h3>")
+                output_source_files = gr.Files(type="filepath", interactive=False, show_label=False)
+                output_source_preview_check = gr.Checkbox(label="Показать плееры для исходников ансамбля", interactive=True, value=False)
+                @gr.render(inputs=[output_source_preview_check, output_source_files])
+                def show_output_auto_ensemble_players(preview, audios):
+                    if preview:
+                        if audios:
+                            with gr.Group():
+                                for file in audios:
+                                    gr.Audio(label=os.path.splitext(os.path.basename(file))[0], value=file, interactive=False, show_download_button=False, type="filepath")
+
+        @run_btn.click(
+            inputs=[input_audio, method, output_format, invert_ensemble, output_base_dir, output_temp_dir_check],
+            outputs=[output_audio, output_audio_wav, output_inverted_audio, output_source_files]
+        )
+        def auto_ensemble_run(input_file, method, out_format, invert_ensemble, o_dir, temp_save, progress=gr.Progress(track_tqdm=True)):
+            ensemble_state = ensemble_model_manager.data
+            invert_methods_map = ensemble_model_manager.ensemble_invert_methods_map
+            if not input_file:
+                return None, None, None, None, []
+            if not os.path.exists(input_file):
+                return None, None, None, None, []
+            if not self.audio.check(input_file):
+                return None, None, None, None, []
+            if o_dir.strip() != "" and not temp_save:
+                timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+                o = os.path.join(o_dir, f"ensembless_outputs_{timestamp}")
+                os.makedirs(o, exist_ok=True)
+            else:
+                timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+                o = tempfile.mkdtemp(prefix=f"ensembless_outputs_{timestamp}_")
+                os.makedirs(o, exist_ok=True)
+
+            basename = os.path.splitext(os.path.basename(input_file))[0]
+            def invert_weights(weights):
+                total_weight = sum(weights)
+                return [total_weight - w for w in weights]
+            # print(json.dumps(ensemble_state, indent=4, ensure_ascii=False))
+            success_separations = []  # list[tuple[str, str, float]] [(s_stem, i_stem, weight)]
+            ensemble_sources_list = [] # list[str, str, str, ...]
+            if ensemble_state:
+                total_ensemble_models = len(ensemble_state)
+                for i, model in enumerate(ensemble_state, start=1):
+
+                    ens_mt = model[0]
+                    ens_mn = model[1]
+                    ens_s_stem = model[2]
+                    ens_i_stem = model[3]
+                    weight = model[4]
+
+                    s_stem = None #path to primary stem
+                    i_stem = None #path to invert stem
+
+                    try:
+                        result_seped_auto_ensemble = self.separate(input=input_file, output_dir=os.path.join(o, ens_mn), model_type=ens_mt, model_name=ens_mn, ext_inst=True, template="NAME - MODEL - STEM", output_format="wav", add_settings={"add_single_sep_text_progress": f"{i} из {total_ensemble_models}"}, progress=progress)
+                        if result_seped_auto_ensemble:
+                            for stem, path in result_seped_auto_ensemble:
+                                ensemble_sources_list.append(path)
+                                if stem == ens_s_stem:
+                                    s_stem = path
+                                elif stem == ens_i_stem:
+                                    i_stem = path
+
+                        if invert_ensemble:
+                            if not i_stem:
+                                result_seped_auto_ensemble_invert = self.separate(input=input_file, output_dir=os.path.join(o, f"{ens_mn}_invert"), model_type=ens_mt, model_name=ens_mn, ext_inst=True, template="NAME - MODEL - STEM", output_format="wav", selected_stems=[ens_s_stem], add_settings={"add_single_sep_text_progress": f"{i} из {total_ensemble_models} (инверт.)"}, progress=progress)
+                                if result_seped_auto_ensemble_invert:
+                                    for stem, path in result_seped_auto_ensemble:
+                                        if stem == ens_i_stem:
+                                            i_stem = path
+                                            ensemble_sources_list.append(path)
+
+                    except Exception as e:
+                        print(f"\nПроизошла ошибка при разделении: {e}")
+                        progress(0, desc="Произошла ошибка при разделении, модель пропускается...")
+                        continue
+                    finally:
+                        if s_stem:
+                            success_separations.append((ens_mn, s_stem, i_stem, weight))
+
+            ensemble_sources_stems = []
+            ensemble_sources_invert_stems = []
+            weights = []
+
+            for out_mn, out_s_stem, out_i_stem, out_weight in success_separations:
+                ensemble_sources_stems.append(out_s_stem)
+                ensemble_sources_invert_stems.append(out_i_stem)
+                weights.append(out_weight)
+
+
+            auto_ensemble_invout_file = None
+            auto_ensemble_invout_file_wav = None
+
+            auto_ensemble_output_name = f"ensembless_{self.namer.short(basename, length=50)}_{len(ensemble_sources_stems)}_{method}"
+            auto_ensemble_inverted_output_name = f"ensembless_{self.namer.short(basename, length=50)}_{len(ensemble_sources_stems)}_{invert_methods_map[method]}_invert"
+            auto_ensemble_out_file, auto_ensemble_out_file_wav = ensemble_audio_files(files=ensemble_sources_stems, weights=weights, output=os.path.join(o, auto_ensemble_output_name), ensemble_type=method, out_format=out_format, add_wav=True)
+
+            if invert_ensemble:
+                auto_ensemble_invout_file, auto_ensemble_invout_file_wav = ensemble_audio_files(files=ensemble_sources_invert_stems, weights=invert_weights(weights), output=os.path.join(o, auto_ensemble_inverted_output_name), ensemble_type=invert_methods_map[method], out_format=out_format, add_wav=True)
+            
+            return auto_ensemble_out_file, auto_ensemble_out_file_wav, auto_ensemble_invout_file, ensemble_sources_list
+
+        @add_inputs_btn.click(
+            inputs=[input_path, input_audio],
+            outputs=[add_inputs, input_audio, add_buttons_row])
+        def add_inputs_fn(input_p, input_a):
+            if input_p and os.path.exists(input_p):
+                if input_a is None:
+                    input_a = None
+                if self.audio.check(input_p):
+                    input_a = input_p
+                return gr.update(visible=False), gr.update(value=input_a), gr.update(visible=True)
+            return gr.update(visible=False), gr.update(value=input_a), gr.update(visible=True)
+
+        @add_inputs_url_btn.click(
+            inputs=[input_url, input_audio, inputs_url_format, inputs_url_bitrate, inputs_url_cookie],
+            outputs=[add_inputs_from_url, input_audio, add_buttons_row])
+        def add_inputs_from_url_fn(input_u, input_a, fmt, br, cookie):
+            if input_u:
+                if input_a is None:
+                    input_a = None
+                downloaded_file = dw_yt_dlp(
+                    url=input_u,
+                    output_format=fmt,
+                    output_bitrate=str(int(br)),
+                    cookie=cookie
+                )
+                if downloaded_file and os.path.exists(downloaded_file):
+                    if self.audio.check(downloaded_file):
+                        input_a = downloaded_file
+                    return gr.update(visible=False), gr.update(value=input_a), gr.update(visible=True)
+            return gr.update(visible=False), gr.update(value=input_a), gr.update(visible=True)
+        
+        add_path_btn.click(
+            lambda: (gr.update(visible=True), gr.update(visible=False)),
+            outputs=[add_inputs, add_buttons_row])
+
+        add_url_btn.click(
+            lambda: (gr.update(visible=True), gr.update(visible=False)),
+            outputs=[add_inputs_from_url, add_buttons_row])
+
+        inputs_url_format.change(lambda x: gr.update(visible=False if x in ["wav", "flac", "aiff"] else True), inputs=inputs_url_format, outputs=inputs_url_bitrate)
+
+class ManualEnsembless(MVSEPLESS):
+    def UI(self, output_base_dir, output_temp_dir_check):
+        with gr.Row():
+            with gr.Column():
+                with gr.Group(visible=False) as add_ensemble_inputs:
+                    input_ensemble_path = gr.Textbox(label="Путь к входному файлу", interactive=True)
+                    add_ensemble_inputs_btn = gr.Button("Добавить файл", variant="primary")
+                add_ensemble_path_btn = gr.Button("Добавить файл по пути", variant="secondary")
+                input_ensemble_files = gr.File(label="Входное аудио", interactive=True, type="filepath", file_count="multiple", file_types=[f".{of}" for of in self.audio.input_formats])
+                input_ensemble_preview_check = gr.Checkbox(label="Показать плееры для входных аудио", interactive=True, value=False)
+                @gr.render(inputs=[input_ensemble_preview_check, input_ensemble_files])
+                def show_input_ensemble_players(preview, audios):
+                    if preview:
+                        if audios:
+                            with gr.Group():
+                                for file in audios:
+                                    gr.Audio(label=os.path.splitext(os.path.basename(file))[0], value=file, interactive=False, show_download_button=False, type="filepath")
+
+            with gr.Column():
+                @gr.render(inputs=[input_ensemble_files])
+                def input_ensemble_files_fn(input_files):
+                    check_ensemble_files_status = f"""Анализ входных файлов
+---"""
+                    hz_ = []
+                    err_list = []
+                    if input_files:
+                        for file in input_files:
+                            basename = os.path.splitext(os.path.basename(file))[0]
+                            if os.path.exists(file):
+                                if self.audio.check(file):
+                                    info = self.audio.get_info(file)
+                                    hz = info[0].get("sample_rate")
+                                    check_ensemble_files_status += f"\n{basename} - {hz} hz"
+                                    hz_.append(hz)
+                                else:
+                                    check_ensemble_files_status += f"\n{basename} - Нет аудио"
+                                    err_list.append(file)
+                            else:
+                                check_ensemble_files_status += f"\n{basename} - Файл не найден"
+                                err_list.append(file)
+
+                    check_ensemble_files_result = f"Действительных файлов: {len(hz_)}"
+
+                    all_same = True
+
+                    common_rate = None
+
+                    for hz_hz in hz_:
+                        if common_rate is None:
+                            common_rate = hz_hz
+                        elif common_rate != hz_hz:
+                            all_same = False
+
+                    if hz_ and len(hz_) > 1:
+                        check_ensemble_files_result += "\nВсе действительные файлы имеют одинаковую частоту дискретизации" if all_same else "\nОшибка! Все действительные файлы имеют РАЗНУЮ частоту дискретизации"
+                    else:
+                        check_ensemble_files_result += "\nДля создания ансамбля нужно загрузить, как минимум - 2 файла, содержащие аудио"
+
+
+                    check_ensemble_files_status += f"\n \n{check_ensemble_files_result}"
+
+                    gr.Textbox(container=False, lines=len(check_ensemble_files_status.split("\n")), interactive=False, value=check_ensemble_files_status)
+
+        weights = gr.Textbox(label="Веса", value="1.0,1.0")
+
+        method = gr.Dropdown(label="Алгоритм склеивания", choices=["min_fft", "max_fft", "avg_fft", "median_fft"], value="avg_fft", filterable=False)
+
+        output_format = gr.Dropdown(label="Формат выходного файла", interactive=True,
+                                choices=self.audio.output_formats,
+                                value="mp3", filterable=False)
+
+        output_manual_ensemble_filename = gr.Textbox(label="Имя выходного файла", value="ensemble", interactive=True)
+
+        make_manual_ensemble_btn = gr.Button(value="Создать ансамбль", variant="primary")
+
+        manual_ensemble_output_audio = gr.Audio(label="Результат", type="filepath", interactive=False, show_download_button=True)
+
+        @make_manual_ensemble_btn.click(
+            inputs=[input_ensemble_files, method, output_format, output_base_dir, output_temp_dir_check, output_manual_ensemble_filename, weights], outputs=manual_ensemble_output_audio
+        )
+        def make_manual_ensemble_fn(input_files_list, method, out_format, o_dir, temp_save, o_filename, weights: str):
+            if o_dir.strip() != "" and not temp_save:
+                timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+                o = os.path.join(o_dir, f"ensembless_outputs_{timestamp}")
+                os.makedirs(o, exist_ok=True)
+            else:
+                timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+                o = tempfile.mkdtemp(prefix=f"ensembless_outputs_{timestamp}_")
+                os.makedirs(o, exist_ok=True)
+
+            o_filename = self.namer.sanitize(o_filename)
+            o_filename = self.namer.short(o_filename)
+            
+            output_file = ensemble_audio_files(files=input_files_list, output=os.path.join(o, o_filename), weights=[float(x) for x in weights.split(",")], ensemble_type=method, out_format=out_format)
+            return output_file
+                    
+        add_ensemble_path_btn.click(
+            lambda: (gr.update(visible=True), gr.update(visible=False)),
+            outputs=[add_ensemble_inputs, add_ensemble_path_btn])
+
+        @add_ensemble_inputs_btn.click(
+            inputs=[input_ensemble_path, input_ensemble_files],
+            outputs=[add_ensemble_inputs, input_ensemble_files, add_ensemble_path_btn])
+        def add_ensemble_inputs_fn(input_p, input_a):
+            if input_p and os.path.exists(input_p):
+                if input_a is None:
+                    input_a = []
+                if isinstance(input_a, str):
+                    input_a = [input_a]
+                if self.audio.check(input_p):
+                    input_a.append(input_p)
+                return gr.update(visible=False), gr.update(value=input_a), gr.update(visible=True)
+            return gr.update(visible=False), gr.update(value=input_a), gr.update(visible=True)
+
+class Inverter_UI(MVSEPLESS):
+    def UI(self):
+        with gr.Group():
+            with gr.Row():
+                original_audio = gr.File(label="Оригинал", interactive=True, type="filepath", file_count="single", file_types=[f".{of}" for of in self.audio.input_formats])
+                stem_audio = gr.File(label="Cтем, который будет вычтен из оригинала", interactive=True, type="filepath", file_count="single", file_types=[f".{of}" for of in self.audio.input_formats])
+            with gr.Group():
+                output_format = gr.Dropdown(label="Формат выходного файла", interactive=True,
+                                                    choices=self.audio.output_formats,
+                                                    value="mp3", filterable=False)
+                method = gr.Radio(
+                    choices=["waveform", "spectrogram"],
+                    label="Метод вычитания",
+                    value="waveform",
+                )
+                btn = gr.Button("Вычесть")
+        output_audio = gr.Audio(label="Инверсия", type="filepath", interactive=False, show_download_button=True)
+        @btn.click(inputs=[original_audio, stem_audio, method, output_format], outputs=[output_audio])
+        def invert_result_ensemble(input_file, output_file, method, out_format):
+            if input_file and output_file:
+                o_dir = tempfile.mkdtemp(suffix="_inverter")
+                basename = os.path.splitext(os.path.basename(input_file))[0]
+                output_path = os.path.join(o_dir, f"inverter_{self.namer.short(basename, length=50)}_{method}.{out_format}")
+                inverted = self.inverter.process_audio(audio1_path=input_file, audio2_path=output_file, out_format=out_format, method=method, output_path=output_path)
+                return inverted
+            else:
+                return None
+
+class Vbach(MVSEPLESS):
+    pitch_methods = ("rmvpe+", "fcpe", "mangio-crepe")
+    hop_length_values = (8, 512)
+    index_rates_values = (0, 1)
+    filter_radius_values = (0, 7)
+    protect_values = (0, 0.5)
+    rms_values = (0, 1)
+    f0_min_values = (50, 3000)
+    f0_max_values = (300, 6000)
+
+    def UI(self):
+        with gr.Tab("Инференс"):
+            with gr.Row():
+                with gr.Column():
+                    with gr.Group():
+                        input_audio = gr.File(label="Входные аудио", interactive=True, type="filepath", file_count="multiple", file_types=[f".{of}" for of in self.audio.input_formats])
+                        converted_state = gr.Textbox(label="Состояние разделения", interactive=False, value="", visible=False)
+                        status = gr.Textbox(container=False, lines=3, interactive=False, max_lines=3)
+                        input_preview_check = gr.Checkbox(label="Показать плееры для входных аудио", interactive=True, value=False)
+                        @gr.render(inputs=[input_preview_check, input_audio])
+                        def show_input_players(preview, audios):
+                            if preview:
+                                if audios:
+                                    with gr.Group():
+                                        for file in audios:
+                                            gr.Audio(label=os.path.splitext(os.path.basename(file))[0], value=file, interactive=False, show_download_button=False, type="filepath")
+
+                with gr.Column():
+                    with gr.Group():
+                        model_name = gr.Dropdown(label="Имя модели", interactive=True)
+                        model_list_refresh_btn = gr.Button("Обновить", variant="secondary", interactive=True)
+                        @model_list_refresh_btn.click(
+                            outputs=[model_name]
+                        )
+                        def refresh_list_voice_models():
+                            models = []
+                            models = self.vbach_model_manager.parse_voice_models()
+                            first_model = None
+                            if len(models) > 0:
+                                first_model = models[0]
+                            return gr.update(choices=models, value=first_model)
+                    with gr.Group():
+                        pitch_method = gr.Radio(label="Метод извлечения высоты тона", choices=self.pitch_methods, value=self.pitch_methods[0], interactive=True)
+                        pitch = gr.Slider(label="Высота тона", minimum=-48, maximum=48, step=0.5, value=0, interactive=True)
+                        hop_length = gr.Slider(label="Длина шага", info="Длина шага влияет на точность передачи высоты тона\nЧем меньше длина шага - тем точнее будет передана высота тона", minimum=self.hop_length_values[0], maximum=self.hop_length_values[1], step=8, value=128, interactive=True, visible=False)
+                        @pitch_method.change(
+                            inputs=[pitch_method],
+                            outputs=[hop_length]
+                        )
+                        def show_mangio_crepe_hop_length(pitch_method):
+                            return gr.update(visible=True if pitch_method in ["mangio-crepe"] else False)
+                        stereo_mode = gr.Radio(
+                            choices=["mono", "left/right", "sim/dif"],
+                            label="Стерео режим",
+                            info="mono - монофоническая обработка аудио, \nleft/right - обработка левого и правого каналов отдельно, \nsim/dif - обработка фантомного центра и стерео-базы, разделенную на левый и правый каналы",
+                            value="mono",
+                            interactive=True
+                        )
+                    with gr.Accordion(label="Дополнительные настройки",open=False):
+                        with gr.Group():
+                            with gr.Row():
+                                index_rate = gr.Slider(label="Влияние индекса", info="Чем ниже значение, тем больше голос похож на исходный; чем выше, тем ближе к модели", minimum=self.index_rates_values[0], maximum=self.index_rates_values[1], step=0.05, value=0, interactive=True)
+                                filter_radius = gr.Slider(label="Радиус фильтра", info="Сглаживает результаты извлечения тона\nМожет снизить дыхание и шумы на выходе", minimum=self.filter_radius_values[0], maximum=self.filter_radius_values[1], step=1, value=3, interactive=True)
+                            with gr.Row():      
+                                rms = gr.Slider(label="Соотношение огибающих громкости", info="Значение 0 - огибающая громкости как у входного аудио, 1 - как у выходного сигнала", minimum=self.rms_values[0], maximum=self.rms_values[1], step=0.05, value=0.25, interactive=True)
+                                protect = gr.Slider(label="Защита согласных", info="Предовращает роботизацию дыхания и согласных (Может влиять на четкость речи)\nЗначение 0.5 - выключает защиту, 0 - максимальная защита", minimum=self.protect_values[0], maximum=self.protect_values[1], step=0.05, value=0.35, interactive=True)
+                        with gr.Group():
+                            with gr.Row():
+                                f0_min = gr.Slider(label="Нижний предел диапазона определения высоты тона", minimum=self.f0_min_values[0], maximum=self.f0_min_values[1], step=10, value=50, interactive=True)
+                                f0_max = gr.Slider(label="Верхний предел диапазона определения высоты тона", minimum=self.f0_max_values[0], maximum=self.f0_max_values[1], step=10, value=1100, interactive=True)
+
+                    with gr.Group():
+                        output_name = gr.Textbox(label="Имя выходного файла", interactive=True, value="NAME - MODEL - F0METHOD - PITCH")
+                        format_output_name_check = gr.Checkbox(label="Форматировать имя", info="Используйте ключи: \nNAME - имя входного файла без расширения, \nPITCH - высота тона, \nF0METHOD - метод извлечения высота тона, \nMODEL - имя голосовой модели", value=True, interactive=True)
+                        output_format = gr.Dropdown(label="Формат выходного файла", interactive=True, choices=self.audio.output_formats, value=self.audio.output_formats[0], filterable=False)
+                        convert_btn = gr.Button("Преобразовать", variant="primary", interactive=True)
+
+            
+            @convert_btn.click(
+                inputs=[
+                    input_audio, 
+                    model_name, 
+                    pitch_method, 
+                    pitch, 
+                    hop_length, 
+                    index_rate, 
+                    filter_radius,
+                    rms, 
+                    protect, 
+                    f0_min, 
+                    f0_max, 
+                    output_name, 
+                    format_output_name_check, 
+                    output_format,
+                    stereo_mode
+                ], outputs=[converted_state, status]
+            )
+            def vbach_convert_batch(ifl, mn, pm, p, hl, ir, fr, rms, pr, f0min, f0max, on, fn, of, sm):
+                output_converted_files = []
+                progress = gr.Progress()
+                if ifl:
+                    for i, file in enumerate(ifl, start=1):
+                        try:
+                            print(f"Файл {i} из {len(ifl)}: {file}")
+                            progress(progress=(i / len(ifl)), desc=f"Файл {i} из {len(ifl)}")
+                            out_conv = vbach_inference(input_file=file, model_name=mn, output_dir=tempfile.mkdtemp(), output_name=on, format_name=True if len(ifl) > 1 else fn, output_format=of, pitch=p, method_pitch=pm, output_bitrate=320, add_params={ "index_rate": ir,"filter_radius": fr,"protect": pr,"rms": rms,"mangio_crepe_hop_length": hl,"f0_min": f0min,"f0_max": f0max,"stereo_mode": sm })
+                            output_converted_files.append(out_conv)
+                        except Exception as e:
+                            print(e)
+                return str(output_converted_files), None
+
+            @gr.render(inputs=[converted_state])
+            def show_players_converted(state):
+                if state != "":
+                    output_converted_files = ast.literal_eval(state)
+                    if output_converted_files:
+                        with gr.Group():
+                            for conv_file in output_converted_files:
+                                basename = os.path.splitext(os.path.basename(conv_file))[0]
+                                gr.Audio(
+                                    label=basename,
+                                    value=conv_file,
+                                    type="filepath",
+                                    interactive=False,
+                                    show_download_button=True,
+                                )
+
+        with gr.TabItem("Менеджер"):
+            with gr.TabItem("Загрузить по ссылке"):
+                with gr.TabItem("Через zip файл"):
+                    with gr.Row():
+                        with gr.Column(variant="panel"):
+                            url_zip = gr.Text(label="Ссылка на zip файл")
+                            with gr.Group():
+                                url_zip_model_name = gr.Text(
+                                    label="Имя модели",
+                                )
+                                url_zip_download_btn = gr.Button("Загрузить", variant="primary")
+
+                            url_zip_output = gr.Text(label="Статус", interactive=False, lines=5)
+                            url_zip_download_btn.click(
+                                (lambda x, y: self.vbach_model_manager.install_model_zip(x, self.namer.short(self.namer.sanitize(y), length=40), "url")),
+                                inputs=[url_zip, url_zip_model_name],
+                                outputs=url_zip_output,
+                            )
+
+                with gr.TabItem("Через отдельные файлы"):
+                    with gr.Row():
+                        with gr.Column(variant="panel"):
+                            url_pth = gr.Text(label="Ссылка на *.pth файл")
+                            url_index = gr.Text(label="Ссылка на *.index файл (необязательно)")
+                            with gr.Group():
+                                url_file_model_name = gr.Text(
+                                    label="Имя модели",
+                                )
+                                url_file_download_btn = gr.Button("Загрузить", variant="primary")
+
+                            url_file_output = gr.Text(label="Статус", interactive=False, lines=5)
+                            url_file_download_btn.click(
+                                (lambda x, y, z: self.vbach_model_manager.install_model_files(x, y, self.namer.short(self.namer.sanitize(z), length=40), "url")),
+                                inputs=[url_index, url_pth, url_file_model_name],
+                                outputs=url_file_output,
+                            )
+
+            with gr.Tab("Загрузить с устройства"):
+                with gr.Tab("Через zip файл"):
+                    with gr.Row():
+                        with gr.Column():
+                            local_zip = gr.File(
+                                label="zip файл", file_types=[".zip"], file_count="single"
+                            )
+                        with gr.Column(variant="panel"):
+                            with gr.Group():
+                                local_zip_model_name = gr.Text(
+                                    label="Имя модели",
+                                )
+                                local_zip_upload_btn = gr.Button("Загрузить", variant="primary")
+
+                            local_zip_output = gr.Text(label="Статус", interactive=False, lines=5)
+                            local_zip_upload_btn.click(
+                                (lambda x, y: self.vbach_model_manager.install_model_zip(x, self.namer.short(self.namer.sanitize(y), length=40), "local")),
+                                inputs=[local_zip, local_zip_model_name],
+                                outputs=local_zip_output,
+                            )
+
+                with gr.TabItem("Через отдельные файлы"):
+                    with gr.Group():
+                        with gr.Row():
+                            local_pth = gr.File(
+                                label="*.pth файл", file_types=[".pth"], file_count="single"
+                            )
+                            local_index = gr.File(
+                                label="*.index файл (необязательно)", file_types=[".index"], file_count="single"
+                            )
+                        with gr.Column(variant="panel"):
+                            with gr.Group():
+                                local_file_model_name = gr.Text(
+                                    label="Имя модели",
+                                )
+                                local_file_upload_btn = gr.Button("Загрузить", variant="primary")
+
+                            local_file_output = gr.Text(
+                                label="Статус", interactive=False
+                            )
+                            local_file_upload_btn.click(
+                                (lambda x, y, z: self.vbach_model_manager.install_model_files(x, y, self.namer.short(self.namer.sanitize(z), length=40), "local")),
+                                inputs=[local_index, local_pth, local_file_model_name],
+                                outputs=local_file_output,
+                            )
+
+            with gr.TabItem("Удалить модель"):
+                with gr.Column(variant="panel"):
+                    with gr.Group():
+                        delete_model_name = gr.Dropdown(
+                        label="Имя модели",
+                        choices=self.vbach_model_manager.parse_voice_models(),
+                        interactive=True,
+                        filterable=False
+                        )
+                        delete_refresh_btn = gr.Button("Обновить")
+                        @delete_refresh_btn.click(inputs=None, outputs=delete_model_name)
+                        def refresh_list_voice_models():
+                            models = []
+                            models = self.vbach_model_manager.parse_voice_models()
+                            first_model = None
+                            if len(models) > 0:
+                                first_model = models[0]
+                            return gr.update(choices=models, value=first_model)
+
+                        delete_output = gr.Text(
+                            label="Статус", interactive=False, lines=5
+                        )
+                        delete_btn = gr.Button("Удалить")
+                        delete_btn.click(
+                            fn=self.vbach_model_manager.del_voice_model,
+                            inputs=delete_model_name,
+                            outputs=delete_output
+                        )
+
+            @gr.on(fn="decorator", inputs=None, outputs=[delete_model_name, model_name])
+            def refresh_list_voice_models():
+                models = []
+                models = self.vbach_model_manager.parse_voice_models()
+                first_model = None
+                if len(models) > 0:
+                    first_model = models[0]
+                return gr.update(choices=models, value=first_model), gr.update(choices=models, value=first_model)
+
+
+class PluginManager(Separator):
+    plugins_dir = os.path.join(script_dir, "plugins")
+    os.makedirs(plugins_dir, exist_ok=True)
+
+    def restart_after_install_plugin(self):
+        subprocess.Popen([os.sys.executable] + sys.argv)
+        os._exit(0)
+
+    def parse_plugins(self):
+        for plugin_file in os.listdir(self.plugins_dir):
+            # Пропускаем не-Python файлы и __init__.py
+            if not plugin_file.endswith('.py') or plugin_file == '__init__.py':
+                continue
+                
+            # Получаем имя модуля без расширения
+            plugin_module_name = os.path.splitext(plugin_file)[0]
+            
+            try:
+                # Определяем путь импорта в зависимости от структуры проекта
+                if __package__:
+                    # Если мы в пакете, используем абсолютный импорт
+                    plugin_module = importlib.import_module(f".plugins.{plugin_module_name}", package=__package__)
+                else:
+                    # Если не в пакете, пробуем разные варианты
+                    try:
+                        # Попробуем абсолютный импорт
+                        plugin_module = importlib.import_module(f"plugins.{plugin_module_name}")
+                    except ImportError:
+                        # Если не сработало, загружаем из файла
+                        plugin_path = os.path.join(self.plugins_dir, plugin_file)
+                        spec = importlib.util.spec_from_file_location(plugin_module_name, plugin_path)
+                        plugin_module = importlib.util.module_from_spec(spec)
+                        spec.loader.exec_module(plugin_module)
+                
+                # Получаем класс Plugin из модуля
+                plugin_class = getattr(plugin_module, 'Plugin')
+                
+                # Создаем экземпляр плагина
+                plugin_instance = plugin_class()
+                
+                # Создаем UI плагина
+                with gr.Tab(plugin_instance.name):
+                    plugin_instance.UI()
+                    
+            except Exception as e:
+                print(f"Ошибка загрузки плагина {plugin_module_name}: {e}")
+                continue
+
+    def UI(self):
+        with gr.Tab("Установка"):
+            upload_plugins_files = gr.File(label="Загрузить плагины", file_types=[".py"], file_count="multiple", interactive=True)
+            install_plugins_btn = gr.Button("Установить", interactive=True)
+            
+            @install_plugins_btn.click(
+                inputs=[upload_plugins_files]
+            )
+            def upload_plugin_list(files):
+                if not files:
+                    return
+                for file in files:
+                    try:
+                        # Копируем только .py файлы
+                        if file.name.endswith('.py'):
+                            shutil.copy(
+                                file, os.path.join(self.plugins_dir, os.path.basename(file).replace(" ", "_"))
+                            )
+                    except Exception as e:
+                        print(f"Ошибка копирования файла {file}: {e}")
+                time.sleep(2)
+                self.restart_after_install_plugin()
+
+        self.parse_plugins()
+
+def mvsepless_app(theme):
+    css = None
+    with gr.Blocks(theme=theme, css=css, title="Разделение музыки и вокала") as app:
+
+        output_base_dir_state = gr.State(value=os.path.join(os.getcwd(), "outputs"))
+        output_temp_dir_state = gr.State(value=False)
+        
+        with gr.Tab("Инференс"):
+            Separator().UI(output_base_dir_state, output_temp_dir_state)
+
+        with gr.Tab("Ансамбль"):
+            with gr.Tab("Авто-ансамбль"):
+                AutoEnsembless().UI(output_base_dir_state, output_temp_dir_state)
+
+            with gr.Tab("Ручной ансамбль"):
+                ManualEnsembless().UI(output_base_dir_state, output_temp_dir_state)
+
+        with gr.Tab("Вычитание"):
+            Inverter_UI().UI()
+
+        with gr.Tab("Преобразование"):
+            Vbach().UI()
+
+        with gr.Tab("Плагины"):
+            PluginManager().UI()
+
+        with gr.Tab("Настройки"):
+            with gr.Column():
+                output_base_dir_ui = gr.Textbox(
+                    label="Базовый каталог для выходных файлов", 
+                    interactive=True, 
+                    value=os.path.join(os.getcwd(), "outputs"),
+                    lines=5
+                )
+                output_temp_dir_check_ui = gr.Checkbox(
+                    label="Использовать временный каталог для выходных файлов", 
+                    interactive=True, 
+                    value=False
+                )
+                
+                # Связываем UI элементы с состоянием
+                output_base_dir_ui.change(
+                    lambda x: x, 
+                    inputs=[output_base_dir_ui], 
+                    outputs=[output_base_dir_state]
+                )
+                output_temp_dir_check_ui.change(
+                    lambda x: x, 
+                    inputs=[output_temp_dir_check_ui], 
+                    outputs=[output_temp_dir_state]
+                )
+
+    return app
+

mvsepless/__main__.py

@@ -0,0 +1,67 @@
+import argparse
+import gradio as gr
+import os
+
+if not __package__:
+    from __init__ import mvsepless_app, Separator
+else:
+    from .__init__ import mvsepless_app, Separator
+
+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.")
+    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="Выбранные стемы для разделения.")
+    args = parser.parse_args()
+
+    if args.command == "app":
+        theme = gr.themes.Citrus(
+            primary_hue="teal",
+            secondary_hue="blue",
+            neutral_hue="blue",
+            spacing_size="sm",
+            font=[
+                gr.themes.GoogleFont("Montserrat"),
+                "ui-sans-serif",
+                "system-ui",
+                "sans-serif",
+            ],
+        )
+        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)
+    elif args.command == "cli":
+        input_data = args.input
+        if os.path.isdir(input_data):
+            list_valid_files = []
+            for file in os.listdir(args.input):
+                if os.path.isfile(os.path.join(args.input, file)):
+                    if Separator.audio.check(os.path.join(args.input, file)):
+                        list_valid_files.append(os.path.join(args.input, file))
+
+            input_files = list_valid_files
+        else:
+            input_files = input_data
+
+        results = Separator().separate(
+            input=input_files,
+            output_dir=args.output_dir,
+            model_type=args.model_type,
+            model_name=args.model_name,
+            ext_inst=args.ext_inst,
+            output_format=args.output_format,
+            output_bitrate=args.output_bitrate,
+            template=args.template,
+            selected_stems=args.selected_stems
+        )
+        print("Разделение завершено.")

mvsepless/audio.py

@@ -0,0 +1,781 @@
+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, 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_f32le",
+                "-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

mvsepless/downloader.py

@@ -0,0 +1,92 @@
+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 = 30):
+    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

mvsepless/ensemble.py

@@ -0,0 +1,224 @@
+# 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

mvsepless/infer.py

@@ -0,0 +1,623 @@
+import os
+import sys
+import json
+import argparse
+import time
+from datetime import datetime
+import gc
+import glob
+import yaml
+import torch
+import numpy as np
+import soundfile as sf
+import torch.nn as nn
+
+from typing import Literal
+
+from audio import Audio
+from namer import Namer
+
+namer = Namer()
+audio = Audio()
+
+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  # избегаем деления на ноль
+    scale_factor = peak / current_peak
+    return audio * scale_factor
+
+
+gc.enable()
+
+
+def cleanup_model(model):
+    try:
+        if isinstance(model, torch.nn.DataParallel):
+            model = model.module
+
+        model.to("cpu")
+
+        for name, param in list(model.named_parameters()):
+            del param
+        for name, buf in list(model.named_buffers()):
+            del buf
+
+        del model
+
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            torch.cuda.ipc_collect()
+
+        gc.collect()
+        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.flush()
+
+
+def once_inference(
+    path: str = None,
+    model: any = None,
+    config: any = None,
+    device: any = None,
+    model_type: str = None,
+    extract_instrumental: bool = False,
+    detailed_pbar: bool = False,
+    output_format: Literal[
+        "mp3", "wav", "flac", "ogg", "opus", "m4a", "aac", "aiff"
+    ] = "mp3",
+    output_bitrate: str = "320k",
+    use_tta: bool = False,
+    verbose: bool = False,
+    model_name: str = None,
+    sample_rate: int = 44100,
+    instruments: list = [],
+    store_dir: str = None,
+    template: str = None,
+    selected_instruments: list = [],
+    model_id: int = 0,
+):
+    results = []
+    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.flush()
+    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.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.flush()
+        return results
+
+    mix_orig = mix.copy()
+
+    mean = std = None
+    if config.inference.get("normalize", False):
+        mono = mix.mean(0)
+        mean = mono.mean()
+        std = mono.std()
+        mix = (mix - mean) / std
+
+    if use_tta:
+        track_proc_list = [mix.copy(), mix[::-1].copy(), -1.0 * mix.copy()]
+    else:
+        track_proc_list = [mix.copy()]
+    full_result = []
+    for m in track_proc_list:
+        try:
+            waveforms = demix(
+                config, model, m, device, pbar=detailed_pbar, model_type=model_type
+            )
+
+            full_result.append(waveforms)
+        except Exception as e:
+            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.flush()
+        return results
+
+    waveforms = full_result[0]
+    for i in range(1, len(full_result)):
+        d = full_result[i]
+        for el in d:
+            if i == 2:
+                waveforms[el] += -1.0 * d[el]
+            elif i == 1:
+                waveforms[el] += d[el][::-1].copy()
+            else:
+                waveforms[el] += d[el]
+    for el in waveforms:
+        waveforms[el] /= len(full_result)
+
+    if (
+        extract_instrumental and config.training.target_instrument is not None
+    ):  # Если включен "Extract Instrumental / Извлечь инструментал" и найден целевой инструмент
+        second_stem = [
+            s
+            for s in config.training.instruments
+            if s != config.training.target_instrument
+        ]
+        if second_stem:
+            second_stem_key = second_stem[0]
+            if second_stem_key not in instruments:
+                instruments.append(second_stem_key)
+            waveforms[second_stem_key] = mix_orig - waveforms[instruments[0]]
+
+    elif (
+        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:
+
+            waveforms["inverted -"] = mix_orig.copy()
+            for instr in instruments:
+                if instr in waveforms:
+                    waveforms["inverted -"] -= waveforms[
+                        instr
+                    ]  # стем "inverted -": вычитание выбранного стема из оригинального сигнала (не всегда хорошо)
+
+            if "inverted -" not in instruments:
+                instruments.append("inverted -")
+
+            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 +": сложение не выбранных инструментов в один стем
+                if "inverted +" not in instruments:
+                    instruments.append("inverted +")
+
+            peak = np.max(np.abs(waveforms["inverted -"]))
+            waveforms["inverted +"] = normalize_peak(waveforms["inverted +"], peak)
+
+    elif (
+        extract_instrumental
+        and not selected_instruments
+        and config.training.target_instrument is None
+        and (
+            all(
+                instr in config.training.instruments
+                for instr in ["bass", "drums", "other", "vocals"]
+            )
+            or all(
+                instr in config.training.instruments
+                for instr in ["bass", "drums", "other", "vocals", "piano", "guitar"]
+            )
+        )
+    ):
+
+        waveforms["instrumental -"] = mix_orig.copy()
+        waveforms["instrumental -"] -= waveforms[
+            "vocals"
+        ]  # стем "inverted -": вычитание выбранного стема из оригинального сигнала (не всегда хорошо)
+
+        if "instrumental -" not in instruments:
+            instruments.append("instrumental -")
+
+        all_instruments = config.training.instruments
+        non_vocal_stems = [s for s in all_instruments if s not in ["vocals"]]
+        if non_vocal_stems:
+            waveforms["instrumental +"] = np.zeros_like(mix_orig)
+            for stem in non_vocal_stems:
+                if stem in waveforms:
+                    waveforms["instrumental +"] += waveforms[
+                        stem
+                    ]  # стем "inverted +": сложение не выбранных инструментов в один стем
+            if "instrumental +" not in instruments:
+                instruments.append("instrumental +")
+
+        peak = np.max(np.abs(waveforms["instrumental -"]))
+        waveforms["instrumental +"] = normalize_peak(waveforms["instrumental +"], peak)
+
+    template = namer.sanitize(template)
+    template = namer.dedup_template(template, keys=["NAME", "MODEL", "STEM", "ID"])
+    template = namer.short(template, length=40)
+
+    for instr in instruments:
+        try:
+            estimates = waveforms[instr].T
+            if mean is not None and std is not None:
+                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)
+            custom_name = namer.template(
+                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.flush()
+
+            output_path = audio.write(
+                o=output_path, array=estimates, sr=sr, of=output_format, br=output_bitrate
+            )  # запись стема в аудио файл с помощью универсальной функции
+
+            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.flush()
+        gc.collect()
+
+    del mix, mix_orig, waveforms, full_result
+    gc.collect()
+
+    return results
+
+
+def run_inference(
+    model: any = None,
+    config: any = None,
+    input_path: str = None,
+    store_dir: str = None,
+    device: any = None,
+    model_type: str = None,
+    extract_instrumental: bool = False,
+    disable_detailed_pbar: bool = False,
+    output_format: Literal[
+        "mp3", "wav", "flac", "ogg", "opus", "m4a", "aac", "aiff"
+    ] = "mp3",
+    output_bitrate: str = "320k",
+    use_tta: bool = False,
+    verbose: bool = False,
+    model_name: str = None,
+    template: str = "NAME_STEM",
+    selected_instruments: list = [],
+    model_id: int = 0,
+):
+    start_time = time.time()
+    if model_type != "vr":
+        model.eval()
+    sample_rate = 44100
+    if "sample_rate" in config.audio:
+        sample_rate = config.audio["sample_rate"]
+
+    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.flush()
+    else:
+        if selected_instruments is not None and selected_instruments != []:
+            instruments = [
+                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.flush()
+
+    os.makedirs(store_dir, exist_ok=True)
+
+    detailed_pbar = not disable_detailed_pbar
+
+    results = once_inference(
+        path=input_path,
+        model=model,
+        config=config,
+        device=device,
+        model_type=model_type,
+        extract_instrumental=extract_instrumental,
+        detailed_pbar=detailed_pbar,
+        output_format=output_format,
+        output_bitrate=output_bitrate,
+        use_tta=use_tta,
+        verbose=verbose,
+        model_name=model_name,
+        sample_rate=sample_rate,
+        instruments=instruments,
+        store_dir=store_dir,
+        template=template,
+        selected_instruments=selected_instruments,
+        model_id=model_id,
+    )
+
+    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.flush()
+    sys.stdout.write(json.dumps({"done": results}, ensure_ascii=False) + '\n')
+    sys.stdout.flush()
+    return results
+
+
+def load_model(model_type, config_path, start_check_point, device_ids, force_cpu=False):
+    device = "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.flush()
+        device = "cuda"
+
+        if device_ids is None:
+            device = "cuda:0"
+        elif isinstance(device_ids, (list, tuple)):
+            device = f"cuda:{device_ids[0]}" if device_ids else "cuda:0"
+        elif isinstance(device_ids, bool):
+            device = "cuda:0"
+        else:
+            device = f"cuda:{int(device_ids)}"
+    elif torch.backends.mps.is_available():
+        device = "mps"
+
+    sys.stdout.write(json.dumps({"device": device}, ensure_ascii=False) + '\n')
+    sys.stdout.flush()
+
+    model_load_start_time = time.time()
+    torch.backends.cudnn.benchmark = True
+
+    model, config = get_model_from_config(model_type, config_path)
+
+    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, 
+            primary_stem=config.training.instruments[0],
+            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.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.flush()
+            
+            if model_type in ["htdemucs", "apollo"]:
+                state_dict = torch.load(
+                    start_check_point, map_location=device, weights_only=False
+                )
+                if "state" in state_dict:
+                    state_dict = state_dict["state"]
+                if "state_dict" in state_dict:
+                    state_dict = state_dict["state_dict"]
+            else:
+                try:
+                    state_dict = torch.load(
+                        start_check_point, map_location=device, weights_only=True
+                    )
+                except torch.serialization.pickle.UnpicklingError:
+                    with torch.serialization.safe_globals([torch._C._nn.gelu]):
+                        state_dict = torch.load(
+                            start_check_point, map_location=device, weights_only=True
+                        )
+            try:
+                model.load_state_dict(state_dict)
+            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.flush()
+
+        if (
+            isinstance(device_ids, (list, tuple))
+            and len(device_ids) > 1
+            and not force_cpu
+            and torch.cuda.is_available()
+        ):
+            model = nn.DataParallel(model, device_ids=[int(d) for d in device_ids])
+
+        model = model.to(device)
+
+        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.flush()
+
+        return model, config, device
+
+
+def mvsep_offline(
+    input_path,
+    store_dir,
+    model_type,
+    config_path,
+    start_check_point,
+    extract_instrumental,
+    output_format,
+    output_bitrate,
+    model_name,
+    template,
+    device_ids=None,
+    disable_detailed_pbar=False,
+    use_tta=False,
+    force_cpu=False,
+    verbose=False,
+    selected_instruments=None,
+    model_id=0,
+):
+    model, config, device = load_model(
+        model_type, config_path, start_check_point, device_ids, force_cpu
+    )
+
+    results = run_inference(
+        model=model,
+        config=config,
+        input_path=input_path,
+        store_dir=store_dir,
+        device=device,
+        model_type=model_type,
+        extract_instrumental=extract_instrumental,
+        disable_detailed_pbar=disable_detailed_pbar,
+        output_format=output_format,
+        output_bitrate=output_bitrate,
+        use_tta=use_tta,
+        verbose=verbose,
+        model_name=model_name,
+        template=template,
+        selected_instruments=selected_instruments,
+        model_id=model_id,
+    )
+
+    if model_type != "vr":
+        cleanup_model(model)
+    del config
+    gc.collect()
+    return results
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        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,
+        default="htdemucs",
+        choices=[
+            "mel_band_roformer",
+            "bs_roformer",
+            "mdx23c",
+            "scnet",
+            "htdemucs",
+            "bandit",
+            "bandit_v2",
+            "mdxnet",
+            "vr"
+        ],
+        help="Тип модели (по умолчанию: htdemucs)",
+    )
+    parser.add_argument(
+        "--config_path",
+        type=str,
+        required=True,
+        help="Путь к конфигурационному файлу модели",
+    )
+    parser.add_argument(
+        "--start_check_point", type=str, required=True, help="Путь к чекпоинту модели"
+    )
+
+    # Параметры вывода
+    parser.add_argument(
+        "--output_format",
+        type=str,
+        default="wav",
+        choices=audio.output_formats,
+        help="Формат выходных файлов",
+    )
+    parser.add_argument(
+        "--output_bitrate", type=str, required=True, help="Битрейт выходного файла"
+    )
+
+    parser.add_argument(
+        "--selected_instruments",
+        nargs="+",
+        help="Список стемов для сохранения (например: vocals drums)",
+    )
+    parser.add_argument(
+        "--extract_instrumental",
+        action="store_true",
+        help="Извлечь инструментальную версию",
+    )
+    parser.add_argument(
+        "--template",
+        type=str,
+        default="NAME_STEM",
+        help="Шаблон для имен выходных файлов",
+    )
+    parser.add_argument(
+        "--model_name",
+        type=str,
+        default="model",
+        help="Имя модели для шаблона имен файлов",
+    )
+    parser.add_argument("-m_id", "--model_id", type=int, required=True, help="Model ID")
+    parser.add_argument(
+        "--device_ids", nargs="+", help="ID GPU устройств для использования"
+    )
+    parser.add_argument(
+        "--force_cpu", action="store_true", help="Принудительно использовать CPU"
+    )
+    parser.add_argument(
+        "--use_tta", action="store_true", help="Использовать тестовую аугментацию"
+    )
+    parser.add_argument(
+        "--disable_detailed_pbar",
+        action="store_true",
+        help="Отключить детальный прогресс-бар",
+    )
+    parser.add_argument("--verbose", action="store_true", help="Подробный вывод")
+
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+
+    device_ids = None
+    if args.device_ids:
+        device_ids = [int(x) for x in args.device_ids]
+
+    results = mvsep_offline(
+        input_path=args.input,
+        store_dir=args.store_dir,
+        model_type=args.model_type,
+        config_path=args.config_path,
+        start_check_point=args.start_check_point,
+        extract_instrumental=args.extract_instrumental,
+        output_format=args.output_format,
+        output_bitrate=args.output_bitrate,
+        model_name=args.model_name,
+        template=args.template,
+        device_ids=device_ids,
+        disable_detailed_pbar=args.disable_detailed_pbar,
+        use_tta=args.use_tta,
+        force_cpu=args.force_cpu,
+        verbose=args.verbose,
+        selected_instruments=args.selected_instruments,
+        model_id=args.model_id,
+    )
+
+
+if __name__ == "__main__":
+    main()

mvsepless/infer_utils.py

@@ -0,0 +1,382 @@
+# coding: utf-8
+__author__ = "Roman Solovyev (ZFTurbo): https://github.com/ZFTurbo/"
+
+import sys
+import json
+import numpy as np
+import torch
+import torch.nn as nn
+import yaml
+import librosa
+import torch.nn.functional as F
+from ml_collections import ConfigDict
+from omegaconf import OmegaConf
+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":
+                config = OmegaConf.load(config_path)
+            else:
+                config = ConfigDict(yaml.load(f, Loader=yaml.FullLoader))
+            return config
+    except FileNotFoundError:
+        raise FileNotFoundError(f"Configuration file not found at {config_path}")
+    except Exception as e:
+        raise ValueError(f"Error loading configuration: {e}")
+
+
+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":
+        from models.mdx23c_tfc_tdf_v3 import TFC_TDF_net
+
+        model = TFC_TDF_net(config)
+
+    elif model_type == "mdxnet":
+        from models.mdx_net import MDXNet
+
+        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":
+        from models.demucs4ht import get_model
+
+        model = get_model(config)
+
+    elif model_type == "mel_band_roformer":
+        if hasattr(config, "windowed"): # Это не нарушает совместимость со обычными моделями на Mel-Band Roformer 
+            from models.windowed_roformer.model import MelBandRoformerWSA
+
+            model = MelBandRoformerWSA(**dict(config.model))
+
+        else:
+            from models.bs_roformer import MelBandRoformer
+
+            model = MelBandRoformer(**dict(config.model))
+
+    elif model_type == "bs_roformer":
+        if hasattr(config.model, "use_shared_bias"):
+            from models.bs_roformer import BSRoformer_SW
+
+            model = BSRoformer_SW(**dict(config.model))
+        elif hasattr(config.model, "fno"):
+            from models.bs_roformer import BSRoformer_FNO
+
+            model = BSRoformer_FNO(**dict(config.model))
+        else:
+            from models.bs_roformer import BSRoformer
+
+            model = BSRoformer(**dict(config.model))
+
+    elif model_type == "bandit":
+        from models.bandit.core.model import MultiMaskMultiSourceBandSplitRNNSimple
+
+        model = MultiMaskMultiSourceBandSplitRNNSimple(**config.model)
+
+    elif model_type == "bandit_v2":
+        from models.bandit_v2.bandit import Bandit
+
+        model = Bandit(**config.kwargs)
+    elif model_type == "scnet_unofficial":
+        from models.scnet_unofficial import SCNet
+
+        model = SCNet(**config.model)
+    elif model_type == "scnet":
+        from models.scnet import SCNet
+
+        model = SCNet(**config.model)
+    else:
+        raise ValueError(f"Unknown model type: {model_type}")
+
+    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)
+
+    window = torch.ones(window_size)
+    window[-fade_size:] = fadeout
+    window[:fade_size] = fadein
+    return window
+
+def demix_mdxnet(
+    config: Any,
+    model: Any,
+    mix: np.ndarray,
+    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)
+        result = processed_wav.cpu().numpy()
+        inv_result = inv_processed_wav.cpu().numpy()
+        result_separation = (result + -inv_result) * 0.5
+    else:
+        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)
+
+    return {stem_name: result_separation}  # Перемещаем на CPU для возврата
+
+def demix_vr(
+    config: Any,
+    model: Any,
+    mix: np.ndarray,
+    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)
+
+def demix_demucs(config, model, mix, device, pbar=False):
+    mix = torch.tensor(mix, dtype=torch.float32)
+    chunk_size = config.training.samplerate * config.training.segment
+    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)  # Создаём окно
+
+    batch_size = config.inference.batch_size
+    use_amp = getattr(config.training, "use_amp", True)
+
+    with torch.cuda.amp.autocast(enabled=use_amp):
+        with torch.inference_mode():
+            req_shape = (num_instruments,) + mix.shape
+            result = torch.zeros(req_shape, dtype=torch.float32)
+            counter = torch.zeros(req_shape, dtype=torch.float32)
+
+            i = 0
+            batch_data = []
+            batch_locations = []
+
+            while i < mix.shape[1]:
+                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
+                )
+
+                batch_data.append(part)
+                batch_locations.append((i, chunk_len))
+                i += step
+
+                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:  # Первый чанк, без fade-in
+                        window[:fade_size] = 1
+                    elif i >= mix.shape[1]:  # Последний чанк, без fade-out
+                        window[-fade_size:] = 1
+
+                    for j, (start, seg_len) in enumerate(batch_locations):
+                        result[..., start : start + seg_len] += (
+                            x[j, ..., :seg_len].cpu() * window[..., :seg_len]
+                        )
+                        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.flush()
+
+                    batch_data.clear()
+                    batch_locations.clear()
+
+            estimated_sources = result / counter
+            estimated_sources = estimated_sources.cpu().numpy()
+            np.nan_to_num(estimated_sources, copy=False, nan=0.0)
+
+    if num_instruments <= 1:
+        return estimated_sources
+    else:
+        instruments = config.training.instruments
+        return {k: v for k, v in zip(instruments, estimated_sources)}
+
+def demix_generic(
+    config: ConfigDict,
+    model: torch.nn.Module,
+    mix: torch.Tensor,
+    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)
+    num_instruments = len(instruments)
+    num_overlap = config.inference.num_overlap
+
+    fade_size = chunk_size // 10
+    step = chunk_size // num_overlap
+    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")
+
+    batch_size = config.inference.batch_size
+    use_amp = getattr(config.training, "use_amp", True)
+
+    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)
+
+            i = 0
+            batch_data = []
+            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
+                )
+
+                batch_data.append(part)
+                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
+                        window[:fade_size] = 1
+                    elif i >= mix.shape[1]:  # Last audio chunk, no fadeout
+                        window[-fade_size:] = 1
+
+                    for j, (start, seg_len) in enumerate(batch_locations):
+                        result[..., start : start + seg_len] += (
+                            x[j, ..., :seg_len].cpu() * window[..., :seg_len]
+                        )
+                        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.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,
+    mix: np.ndarray,
+    device: torch.device,
+    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:
+        return config.training.instruments
+
+
+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

mvsepless/model_manager.py

@@ -0,0 +1,540 @@
+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.hubert_path = os.path.join(script_dir, "embedders", "hubert_base.pt")
+        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], ["https://huggingface.co/Politrees/RVC_resources/resolve/main/embedders/hubert_base.pt", self.hubert_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_and_load()
+        pass
+    
+    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_file(self, url_model, local_path):
+        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)
+
+        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
+            )
+
+    def download_requirements(self):
+        for url, file in self.requirements:
+            if not os.path.exists(file):
+                self.download_file(url_model=url, local_path=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:
+                self.download_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("Нет установленных моделей")
+
+
+
+
+
+

mvsepless/models/bandit/core/__init__.py

@@ -0,0 +1,691 @@
+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,
+        )

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

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

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

@@ -0,0 +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

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

@@ -0,0 +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

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

@@ -0,0 +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

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

@@ -0,0 +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")

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

@@ -0,0 +1,68 @@
+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

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

@@ -0,0 +1,366 @@
+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

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

@@ -0,0 +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()

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

@@ -0,0 +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

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

@@ -0,0 +1,273 @@
+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_)

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

@@ -0,0 +1,226 @@
+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()

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

@@ -0,0 +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'
+  - 'Traffic Experiment - Sirens'

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

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

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

@@ -0,0 +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))

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

@@ -0,0 +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)

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

@@ -0,0 +1,95 @@
+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)

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

@@ -0,0 +1,139 @@
+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

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

@@ -0,0 +1,9 @@
+from .snr import (
+    ChunkMedianScaleInvariantSignalDistortionRatio,
+    ChunkMedianScaleInvariantSignalNoiseRatio,
+    ChunkMedianSignalDistortionRatio,
+    ChunkMedianSignalNoiseRatio,
+    SafeSignalDistortionRatio,
+)
+
+# from .mushra import EstimatedMushraScore

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

@@ -0,0 +1,443 @@
+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.
+    """

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

@@ -0,0 +1,127 @@
+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,
+        )

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

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

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

@@ -0,0 +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,
+        )

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

@@ -0,0 +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

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

@@ -0,0 +1,135 @@
+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

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

@@ -0,0 +1,651 @@
+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

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

@@ -0,0 +1,351 @@
+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

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

@@ -0,0 +1,320 @@
+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

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

@@ -0,0 +1,525 @@
+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)

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

@@ -0,0 +1,829 @@
+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,
+        )

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

@@ -0,0 +1,412 @@
+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_))

mvsepless/models/bandit/model_from_config.py

@@ -0,0 +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

mvsepless/models/bandit_v2/bandit.py

@@ -0,0 +1,363 @@
+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

mvsepless/models/bandit_v2/bandsplit.py

@@ -0,0 +1,130 @@
+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

mvsepless/models/bandit_v2/film.py

@@ -0,0 +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)

mvsepless/models/bandit_v2/maskestim.py

@@ -0,0 +1,281 @@
+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

mvsepless/models/bandit_v2/tfmodel.py

@@ -0,0 +1,145 @@
+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)

mvsepless/models/bandit_v2/utils.py

@@ -0,0 +1,523 @@
+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)