test1

app.py

@@ -0,0 +1,131 @@
+import os
+import sys
+import time
+import shutil
+import argparse
+from datetime import datetime
+import gradio as gr
+os.system("pip install https://github.com/noblebarkrr/mvsepless/blob/bd611441e48e918650e6860738894673b3a1a5f1/fixed/audio_separator-0.32.0-py3-none-any.whl")
+from multi_inference import MVSEPLESS, OUTPUT_FORMATS
+from assets.translations import TRANSLATIONS, TRANSLATIONS_STEMS
+
+OUTPUT_DIR = os.path.join(os.getcwd(), "output")
+plugins_dir = os.path.join(os.getcwd(), "plugins")
+os.makedirs(plugins_dir, exist_ok=True)
+
+CURRENT_LANG = "ru"
+
+def t(key, **kwargs):
+    """Функция для получения перевода с подстановкой значений"""
+    lang = CURRENT_LANG
+    translation = TRANSLATIONS.get(lang, {}).get(key, key)
+    return translation.format(**kwargs) if kwargs else translation
+
+def t_stem(key, **kwargs):
+    """Функция для получения перевода с подстановкой значений"""
+    lang = CURRENT_LANG
+    translation = TRANSLATIONS_STEMS.get(lang, {}).get(key, key)
+    return translation.format(**kwargs) if kwargs else translation
+
+def gen_out_dir():
+    return os.path.join(OUTPUT_DIR, datetime.now().strftime("%Y%m%d_%H%M%S"))
+
+mvsepless = MVSEPLESS()
+
+def sep_wrapper(a, b, c, d, e, f, g, h):
+    results = mvsepless.separator(input_file=a, output_dir=gen_out_dir(), model_type=b, model_name=c, ext_inst=d, vr_aggr=e, output_format=f, output_bitrate=f'{g}k', call_method="cli", selected_stems=h)
+    stems = []
+    if results:
+        for i, (stem, output_file) in enumerate(results[:20]):
+            stems.append(gr.update(
+                visible=True,
+                label=t_stem(stem),
+                value=output_file
+            ))
+    
+    while len(stems) < 20:
+        stems.append(gr.update(visible=False, label=None, value=None))
+        
+    return tuple(stems)
+
+
+theme = gr.themes.Default(
+        primary_hue="violet",
+        secondary_hue="cyan",
+        neutral_hue="blue",
+        spacing_size="sm",
+        font=[gr.themes.GoogleFont("Tektur"), 'ui-sans-serif', 'system-ui', 'sans-serif'],
+    ).set(
+        body_text_color='*neutral_950',
+        body_text_color_subdued='*neutral_500',
+        background_fill_primary='*neutral_200',
+        background_fill_primary_dark='*neutral_800',
+        border_color_accent='*primary_950',
+        border_color_accent_dark='*neutral_700',
+        border_color_accent_subdued='*primary_500',
+        border_color_primary='*primary_800',
+        border_color_primary_dark='*neutral_400',
+        color_accent_soft='*primary_100',
+        color_accent_soft_dark='*neutral_800',
+        link_text_color='*secondary_700',
+        link_text_color_active='*secondary_700',
+        link_text_color_hover='*secondary_800',
+        link_text_color_visited='*secondary_600',
+        link_text_color_visited_dark='*secondary_700',
+        block_background_fill='*background_fill_secondary',
+        block_background_fill_dark='*neutral_950',
+        block_label_background_fill='*secondary_400',
+        block_label_text_color='*neutral_800',
+        panel_background_fill='*background_fill_primary',
+        checkbox_background_color='*background_fill_secondary',
+        checkbox_label_background_fill_dark='*neutral_900',
+        input_background_fill_dark='*neutral_900',
+        input_background_fill_focus='*neutral_100',
+        input_background_fill_focus_dark='*neutral_950',
+        button_small_radius='*radius_sm',
+        button_secondary_background_fill='*neutral_400',
+        button_secondary_background_fill_dark='*neutral_500',
+        button_secondary_background_fill_hover_dark='*neutral_950'
+    )
+
+
+def create_app():
+    with gr.Row():
+        with gr.Column():
+            input_audio = gr.Audio(label=t("select_file"), interactive=True, type="filepath")
+            input_audio_path = gr.Textbox(label=t("audio_path"), info=t("audio_path_info"), interactive=True)
+        with gr.Column():
+            with gr.Row():
+                model_type = gr.Dropdown(label=t("model_type"), choices=mvsepless.get_mt(), value=mvsepless.get_mt()[0], interactive=True, filterable=False)
+                model_name = gr.Dropdown(label=t("model_name"), choices=mvsepless.get_mn(mvsepless.get_mt()[0]), value=mvsepless.get_mn(mvsepless.get_mt()[0])[0], interactive=True, filterable=False)
+            target_instrument = gr.Textbox(label=t("target_instrument"), value=mvsepless.get_tgt_inst(mvsepless.get_mt()[0], mvsepless.get_mn(mvsepless.get_mt()[0])[0]), interactive=False)
+            vr_aggr = gr.Slider(0, 100, step=1, label=t("vr_aggressiveness"), visible=False, value=5, interactive=True)
+            extract_instrumental = gr.Checkbox(label=t("extract_instrumental"), value=True, interactive=True)
+            stems_list = gr.CheckboxGroup(label=t("stems_list"), info=t("stems_info", target_instrument="vocals"), choices=mvsepless.get_stems(mvsepless.get_mt()[0], mvsepless.get_mn(mvsepless.get_mt()[0])[0]), value=None, interactive=False)
+            with gr.Row():
+                output_format, output_bitrate = gr.Dropdown(label=t("output_format"), choices=OUTPUT_FORMATS, value="mp3", interactive=True, filterable=False), gr.Slider(32, 320, step=1, label=t("bitrate"), value=320, interactive=True)
+            separate_btn = gr.Button(t("separate_btn"), variant="primary", interactive=True)
+    download_via_zip_btn = gr.DownloadButton(label="Download via zip", visible=False, interactive=True)
+    output_stems = []
+    for _ in range(10):
+        with gr.Row():
+            audio1 = gr.Audio(visible=False, interactive=False, type="filepath", show_download_button=True)
+            audio2 = gr.Audio(visible=False, interactive=False, type="filepath", show_download_button=True)
+            output_stems.extend([audio1, audio2])
+
+    input_audio.upload(fn=(lambda x: gr.update(value=x)), inputs=input_audio, outputs=input_audio_path)
+    model_type.change(fn=(lambda x: gr.update(choices=mvsepless.get_mn(x), value=mvsepless.get_mn(x)[0])), inputs=model_type, outputs=model_name).then(fn=(lambda x: (gr.update(visible=False if x in ["vr", "mdx"] else True), gr.update(visible=True if x == "vr" else False))), inputs=model_type, outputs=[extract_instrumental, vr_aggr])
+    model_name.change(fn=(lambda x, y: gr.update(choices=mvsepless.get_stems(x, y), value=None)), inputs=[model_type, model_name], outputs=stems_list).then(fn=(lambda x, y: (gr.update(interactive=True if mvsepless.get_tgt_inst(x, y) == None else None, info=t("stems_info", target_instrument=mvsepless.get_tgt_inst(x, y)) if mvsepless.get_tgt_inst(x, y) is not None else t("stems_info2")), gr.update(value=mvsepless.get_tgt_inst(x, y)), gr.update(value=True if mvsepless.get_tgt_inst(x, y) is not None else False))), inputs=[model_type, model_name], outputs=[stems_list, target_instrument, extract_instrumental])
+    separate_btn.click(fn=sep_wrapper, inputs=[input_audio_path, model_type, model_name, extract_instrumental, vr_aggr, output_format, output_bitrate, stems_list], outputs=output_stems, show_progress_on=input_audio)
+
+
+CURRENT_LANG = ru
+css = """
+.fixed-height { height: 160px !important; min-height: 160px !important; }
+.fixed-height2 { height: 250px !important; min-height: 250px !important; }
+"""
+
+with gr.Blocks(theme=theme, css=css) as app:
+    create_app()
+
+app.launch(allowed_paths=["/"], server_port=7860, share=False)

assets/translations.py

@@ -0,0 +1,171 @@
+TRANSLATIONS = {
+    "ru": {
+        "app_title": "MVSEPLESS",
+        "separation": "Разделение",
+        "plugins": "Плагины",
+        "select_file": "Выберите файл",
+        "audio_path": "Путь к файлу",
+        "audio_path_info": "Здесь можно ввести путь к файлу, либо загрузить его выше и получить путь к загруженному файлу",
+        "model_type": "Тип модели",
+        "model_name": "Имя модели",
+        "vr_aggressiveness": "Агрессивность для VR моделей",
+        "extract_instrumental": "Извлечь инструментал",
+        "stems_list": "Список стемов",
+        "output_format": "Формат вывода",
+        "separate_btn": "Разделить",
+        "upload": "Загрузка плагинов (.py)",
+        "upload_btn": "Загрузить",
+        "loading_plugin": "Загружается плагин: {name}",
+        "error_loading_plugin": "Произошла ошибка при загрузке плагина: {e}",
+        "target_instrument": "Целевой инструмент",
+        "stems_info": "Выбор стемов недоступен\nДля извлечения второго стема включите \"Извлечь инструментал\"",
+        "stems_info2": "Для получения остатка (при выбранных стемах), включите \"Извлечь инструментал\"",
+        "bitrate": "Битрейт (Кбит/сек)"
+    },
+    "en": {
+        "app_title": "MVSEPLESS",
+        "separation": "Separation",
+        "plugins": "Plugins",
+        "select_file": "Select File",
+        "audio_path": "Audio path",
+        "audio_path_info": "You can enter the file path here, or upload it above and get the path to the uploaded file.",
+        "model_type": "Model Type",
+        "model_name": "Model Name",
+        "vr_aggressiveness": "Aggressiveness for VR Models",
+        "extract_instrumental": "Extract Instrumental",
+        "stems_list": "Stems List",
+        "output_format": "Output Format",
+        "separate_btn": "Separate",
+        "upload": "Upload plugins (.py)",
+        "upload_btn": "Upload",
+        "loading_plugin": "Loading plugin: {name}",
+        "error_loading_plugin": "As error occured loading plugin: {e}",
+        "target_instrument": "Target instrument",
+        "stems_info": "Stem selection unavailable\nEnable \"Extract Instrumental\" to extract the second stem",
+        "stems_info2": "To extract the residual (with selected_stems), enable \"Extract Instrumental\"",
+        "bitrate": "Bitrate (Kbit/sec)"
+    }
+}
+
+TRANSLATIONS_STEMS = {
+    "ru": {
+        "vocals": "Вокал",
+        "Vocals": "Вокал",
+        "other": "Другое",
+        "Other": "Другое",
+        "Instrumental": "Инструментал",
+        "instrumnetal": "Инструментал",
+        "instrumental +": "Инструментал +",
+        "instrumental -": "Инструментал -",
+        "Bleed": "Фон",
+        "Guitar": "Гитара",
+        "drums": "Барабаны",
+        "bass": "Бас",
+        "karaoke": "Караоке",
+        "reverb": "Реверберация",
+        "noreverb": "Без реверберации",
+        "aspiration": "Придыхание",
+        "dry": "Сухой звук",
+        "crowd": "Толпа",
+        "percussions": "Перкуссия",
+        "piano": "Пианино",
+        "guitar": "Гитара",
+        "male": "Мужской",
+        "female": "Женский",
+        "kick": "Кик",
+        "snare": "Малый барабан",
+        "toms": "Том-томы",
+        "hh": "Хай-хэт",
+        "ride": "Райд",
+        "crash": "Крэш",
+        "similarity": "Сходство",
+        "difference": "Различие",
+        "inst": "Инструмент",
+        "orch": "Оркестр",
+        "No Woodwinds": "Без деревянных духовых",
+        "Woodwinds": "Деревянные духовые",
+        "No Echo": "Без эха",
+        "Echo": "Эхо",
+        "No Reverb": "Без реверберации",
+        "Reverb": "Реверберация",
+        "Noise": "Шум",
+        "No Noise": "Без шума",
+        "Dry": "Сухой звук",
+        "No Dry": "Не сухой звук",
+        "Breath": "Дыхание",
+        "No Breath": "Без дыхания",
+        "No Crowd": "Без толпы",
+        "Crowd": "Толпа",
+        "No Other": "Без другого",
+        "Bass": "Бас",
+        "No Bass": "Без баса",
+        "Drums": "Барабаны",
+        "No Drums": "Без барабанов",
+        "speech": "Речь",
+        "music": "Музыка",
+        "effects": "Эффекты",
+        "sfx": "Звуковые эффекты",
+        "inverted +": "Инверсия +",
+        "inverted -": "Инверсия -"
+    },
+    "en": {
+        "vocals": "Vocals",
+        "Vocals": "Vocals",
+        "other": "Other",
+        "Other": "Other",
+        "Instrumental": "Instrumental",
+        "instrumnetal": "Instrumental",
+        "instrumental +": "Instrumental +",
+        "instrumental -": "Instrumental -",
+        "Bleed": "Bleed",
+        "Guitar": "Guitar",
+        "drums": "Drums",
+        "bass": "Bass",
+        "karaoke": "Karaoke",
+        "reverb": "Reverb",
+        "noreverb": "No reverb",
+        "aspiration": "Aspiration",
+        "dry": "Dry",
+        "crowd": "Crowd",
+        "percussions": "Percussions",
+        "piano": "Piano",
+        "guitar": "Guitar",
+        "male": "Male",
+        "female": "Female",
+        "kick": "Kick",
+        "snare": "Snare",
+        "toms": "Toms",
+        "hh": "Hi-hat",
+        "ride": "Ride",
+        "crash": "Crash",
+        "similarity": "Similarity",
+        "difference": "Difference",
+        "inst": "Instrument",
+        "orch": "Orchestra",
+        "No Woodwinds": "No Woodwinds",
+        "Woodwinds": "Woodwinds",
+        "No Echo": "No Echo",
+        "Echo": "Echo",
+        "No Reverb": "No Reverb",
+        "Reverb": "Reverb",
+        "Noise": "Noise",
+        "No Noise": "No Noise",
+        "Dry": "Dry",
+        "No Dry": "No Dry",
+        "Breath": "Breath",
+        "No Breath": "No Breath",
+        "No Crowd": "No Crowd",
+        "Crowd": "Crowd",
+        "No Other": "No Other",
+        "Bass": "Bass",
+        "No Bass": "No Bass",
+        "Drums": "Drums",
+        "No Drums": "No Drums",
+        "speech": "Speech",
+        "music": "Music",
+        "effects": "Effects",
+        "sfx": "SFX",
+        "inverted +": "Inverted +",
+        "inverted -": "Inverted -"
+    }
+}

multi_inference.py

@@ -0,0 +1,303 @@
+import os
+import time
+import shutil
+import sys
+import gc
+import argparse
+import json
+import subprocess
+from datetime import datetime
+from tabulate import tabulate
+
+SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(SCRIPT_DIR)
+os.chdir(SCRIPT_DIR)
+
+from model_list import models_data
+from utils.preedit_config import conf_editor
+from utils.download_models import download_model
+
+MODELS_CACHE_DIR = os.path.join(SCRIPT_DIR, "separator", "models_cache")
+MODEL_TYPES = ["mel_band_roformer", "bs_roformer", "mdx23c", "scnet", "htdemucs", "bandit", "bandit_v2", "vr", "mdx"]
+OUTPUT_FORMATS = ["mp3", "wav", "flac", "ogg", "opus", "m4a", "aac", "aiff"]
+
+class MVSEPLESS:
+    def __init__(self):
+        self.models_cache_dir = os.path.join(SCRIPT_DIR, "separator", "models_cache")
+        self.model_types = MODEL_TYPES
+        self.output_formats = OUTPUT_FORMATS
+
+    def get_mt(self):
+        return list(models_data.keys())
+       
+    def get_mn(self, model_type):
+        return list(models_data[model_type].keys())
+       
+    def get_stems(self, model_type, model_name):
+        stems = models_data[model_type][model_name]["stems"]
+        return stems
+    
+    def get_tgt_inst(self, model_type, model_name):
+        target_instrument = models_data[model_type][model_name]["target_instrument"]
+        return target_instrument
+
+    def display_models_info(self, filter: str = None):
+        print("\nAvailable Models Information:")
+        print("=" * 50)
+        
+        for model_type in models_data:
+            print(f"\nModel Type: {model_type.upper()}")
+            print("-" * 50)
+            
+            table_data = []
+            headers = ["Model Name", "Stems", "Target Instrument", "Primary Stem"]
+            
+            for model_name in models_data[model_type]:
+                model_info = models_data[model_type][model_name]
+                
+                if filter and filter not in model_info.get('stems', []):
+                    continue
+                    
+                stems = "\n".join(model_info.get('stems', [])) if 'stems' in model_info else "N/A"
+                target = model_info.get('target_instrument', "N/A")
+                primary = model_info.get('primary_stem', "N/A")
+                
+                table_data.append([model_name, stems, target, primary])
+            
+            print(tabulate(table_data, headers=headers, tablefmt="grid"))
+            print()
+            
+    def separator(
+        self,
+        input_file: str = None,
+        output_dir: str = None,
+        model_type: str = "mel_band_roformer",
+        model_name: str = "Mel-Band-Roformer_Vocals_kimberley_jensen",
+        ext_inst: bool = False,
+        vr_aggr: int = 5,
+        output_format: str = "wav",
+        output_bitrate: str = "320k",
+        template: str = "NAME_(STEM)_MODEL",
+        call_method: str = "cli",
+        selected_stems: list = None
+    ):
+        if selected_stems is None:
+            selected_stems = []
+            
+        if not input_file:
+            print("Please, input path to input file")
+            return [("None", "/none/none.mp3")]
+    
+        if not os.path.exists(input_file):
+            print("Input file not exist")
+            return [("None", "/none/none.mp3")]
+    
+        if "STEM" not in template:
+            template = template + "_STEM"
+    
+        print(f"Starting inference: {model_type}/{model_name}, bitrate={output_bitrate}, method={call_method}, stems={selected_stems}")
+        os.makedirs(output_dir, exist_ok=True)
+    
+        if model_type in ["mel_band_roformer", "bs_roformer", "mdx23c", "scnet", "htdemucs", "bandit", "bandit_v2"]:
+            try:
+                info = models_data[model_type][model_name]
+            except KeyError:
+                print("Model not exist")
+                return [("None", "/none/none.mp3")]
+                
+            conf, ckpt = download_model(self.models_cache_dir, model_name, model_type, 
+                                      info["checkpoint_url"], info["config_url"])
+            if model_type != "htdemucs":
+                conf_editor(conf)
+    
+            if call_method == "cli":
+                cmd = ["python", "-m", "separator.msst_separator", f'--input "{input_file}"', 
+                      f'--store_dir "{output_dir}"', f'--model_type "{model_type}"', 
+                      f'--model_name "{model_name}"', f'--config_path "{conf}"', 
+                      f'--start_check_point "{ckpt}"', f'--output_format "{output_format}"', 
+                      f'--output_bitrate "{output_bitrate}"', f'--template "{template}"', 
+                      "--save_results_info"]
+                if ext_inst:
+                    cmd.append("--extract_instrumental")
+                if selected_stems:
+                    instruments = " ".join(f'"{s}"' for s in selected_stems)
+                    cmd.append(f'--selected_instruments {instruments}')
+                subprocess.run(" ".join(cmd), shell=True, check=True)
+    
+                results_path = os.path.join(output_dir, "results.json")
+                if os.path.exists(results_path):
+                    with open(results_path, encoding="utf-8") as f:
+                        return json.load(f)
+                return [("None", "/none/none.mp3")]
+    
+            elif call_method == "direct":
+                from separator.msst_separator import mvsep_offline
+                try:
+                    return mvsep_offline(
+                        input_path=input_file, store_dir=output_dir, model_type=model_type, 
+                        config_path=conf, start_check_point=ckpt, extract_instrumental=ext_inst, 
+                        output_format=output_format, output_bitrate=output_bitrate, 
+                        model_name=model_name, template=template, selected_instruments=selected_stems
+                    )
+                except Exception as e:
+                    print(e)
+                    return [("None", "/none/none.mp3")]
+    
+        elif model_type in ["vr", "mdx"]:
+            try:
+                info = models_data[model_type][model_name]
+            except KeyError:
+                print("Model not exist")
+                return [("None", "/none/none.mp3")]
+                
+            if model_type == "vr" and info.get("custom_vr", False):
+                conf, ckpt = download_model(self.models_cache_dir, model_name, model_type, 
+                                          info["checkpoint_url"], info["config_url"])
+                primary_stem = info["primary_stem"]
+    
+                if call_method == "cli":
+                    cmd = ["python", "-m", "separator.uvr_sep", "custom_vr", 
+                          f'--input_file "{input_file}"', f'--ckpt_path "{ckpt}"', 
+                          f'--config_path "{conf}"', f'--bitrate "{output_bitrate}"', 
+                          f'--model_name "{model_name}"', f'--template "{template}"', 
+                          f'--output_format "{output_format}"', f'--primary_stem "{primary_stem}"', 
+                          f'--aggression {vr_aggr}', f'--output_dir "{output_dir}"']
+                    if selected_stems:
+                        instruments = " ".join(f'"{s}"' for s in selected_stems)
+                        cmd.append(f'--selected_instruments {instruments}')
+                    subprocess.run(" ".join(cmd), shell=True, check=True)
+    
+                    results_path = os.path.join(output_dir, "results.json")
+                    if os.path.exists(results_path):
+                        with open(results_path, encoding="utf-8") as f:
+                            return json.load(f)
+                    return [("None", "/none/none.mp3")]
+    
+                elif call_method == "direct":
+                    from separator.uvr_sep import custom_vr_separate
+                    try:
+                        return custom_vr_separate(
+                            input_file=input_file, ckpt_path=ckpt, config_path=conf, 
+                            bitrate=output_bitrate, model_name=model_name, template=template, 
+                            output_format=output_format, primary_stem=primary_stem, 
+                            aggression=vr_aggr, output_dir=output_dir, 
+                            selected_instruments=selected_stems
+                        )
+                    except Exception as e:
+                        print(e)
+                        return [("None", "/none/none.mp3")]
+            else:
+                if call_method == "cli":
+                    cmd = ["python", "-m", "separator.uvr_sep", "uvr", 
+                          f'--input_file "{input_file}"', f'--output_dir "{output_dir}"', 
+                          f'--template "{template}"', f'--bitrate "{output_bitrate}"', 
+                          f'--model_dir "{self.models_cache_dir}"', f'--model_type "{model_type}"', 
+                          f'--model_name "{model_name}"', f'--output_format "{output_format}"', 
+                          f'--aggression {vr_aggr}']
+                    if selected_stems:
+                        instruments = " ".join(f'"{s}"' for s in selected_stems)
+                        cmd.append(f'--selected_instruments {instruments}')
+                    subprocess.run(" ".join(cmd), shell=True, check=True)
+    
+                    results_path = os.path.join(output_dir, "results.json")
+                    if os.path.exists(results_path):
+                        with open(results_path, encoding="utf-8") as f:
+                            return json.load(f)
+                    return [("None", "/none/none.mp3")]
+    
+                elif call_method == "direct":
+                    from separator.uvr_sep import non_custom_uvr_inference
+                    try:
+                        return non_custom_uvr_inference(
+                            input_file=input_file, output_dir=output_dir, template=template, 
+                            bitrate=output_bitrate, model_dir=self.models_cache_dir, 
+                            model_type=model_type, model_name=model_name, 
+                            output_format=output_format, aggression=vr_aggr, 
+                            selected_instruments=selected_stems
+                        )
+                    except Exception as e:
+                        print(e)
+                        return [("None", "/none/none.mp3")]
+    
+        print("Unsupported model type")
+        return [("None", "/none/none.mp3")]
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Multi-inference for separation audio in Google Colab")
+    subparsers = parser.add_subparsers(dest='command', required=True, help='Sub-command help')
+    
+    list_models = subparsers.add_parser('list', help='List of exist models')
+    list_models.add_argument("-l_filter", "--list_filter", type=str, default=None, help="Show models in list only with specified stem")
+
+    separate = subparsers.add_parser('separate', help='Separate I/O params')
+    separate.add_argument("-i", "--input", type=str, required=True, help="Input file or directory")
+    separate.add_argument("-o", "--output", type=str, required=True, help="Output directory")
+    separate.add_argument("-mt", "--model_type", type=str, required=True, choices=MODEL_TYPES, help="Model type")
+    separate.add_argument("-mn", "--model_name", type=str, required=True, help="Model name")
+    separate.add_argument("-inst", "--instrumental", action='store_true', help="Extract instrumental")
+    separate.add_argument("-stems", "--stems", nargs="+", help="Select output stems")
+    separate.add_argument("-bitrate", "--bitrate", type=str, default="320k", help="Output bitrate")
+    separate.add_argument("-of", "--format", type=str, default="mp3", help="Output format")
+    separate.add_argument("-vr_aggr", "--vr_arch_aggressive", type=int, default=5, help="Aggression for VR ARCH models")
+    separate.add_argument('--template', type=str, default='NAME_STEM', help='Template naming of output files')
+    separate.add_argument("-l_out", "--list_output", action='store_true', help="Show list output files")
+    
+    return parser.parse_args()
+
+if __name__ == "__main__":
+    args = parse_args()
+    mvsepless = MVSEPLESS()
+    
+    if args.command == 'list':
+        mvsepless.display_models_info(args.list_filter)
+        
+    elif args.command == 'separate':
+        if os.path.isfile(args.input):
+            results = mvsepless.separator(
+                input_file=args.input,
+                output_dir=args.output,
+                model_type=args.model_type,
+                model_name=args.model_name,
+                ext_inst=args.instrumental,
+                vr_aggr=args.vr_arch_aggressive,
+                output_format=args.format,
+                output_bitrate=args.bitrate,
+                template=args.template,
+                call_method="cli",
+                selected_stems=args.stems
+            )
+            if args.list_output:
+                print("Results\n")
+                for stem, path in results:
+                    print(f"Stem - {stem}\nPath - {path}\n")
+                    
+        elif os.path.isdir(args.input):
+            batch_results = []
+            for file in os.listdir(args.input):
+                abs_path_file = os.path.join(args.input, file)
+                if os.path.isfile(abs_path_file):
+                    base_name = os.path.splitext(os.path.basename(abs_path_file))[0]
+                    output_subdir = os.path.join(args.output, base_name)
+                    
+                    results = mvsepless.separator(
+                        input_file=abs_path_file,
+                        output_dir=output_subdir,
+                        model_type=args.model_type,
+                        model_name=args.model_name,
+                        ext_inst=args.instrumental,
+                        vr_aggr=args.vr_arch_aggressive,
+                        output_format=args.format,
+                        output_bitrate=args.bitrate,
+                        template=args.template,
+                        call_method="cli",
+                        selected_stems=args.stems
+                    )
+                    batch_results.append((base_name, results))
+                    
+            if args.list_output:
+                print("Results\n")
+                for name, stems in batch_results:
+                    print(f"Name - {name}")
+                    for stem, path in stems:
+                        print(f"  Stem - {stem}\n  Path - {path}\n")
+

requirements.txt

@@ -0,0 +1,50 @@
+torch==2.6.0
+torchvision==0.21.0
+torchaudio==2.6.0
+torchcrepe==0.0.23
+numpy==2.0.2
+pandas==2.2.2
+scipy==1.15.3
+librosa==0.9.1
+matplotlib==3.9.0
+tqdm==4.67.1
+einops==0.8.1
+protobuf==5.29.4
+soundfile==0.13.1
+pydub==0.25.1
+pyloudnorm==0.1.1
+praat-parselmouth==0.4.5
+webrtcvad==2.0.10
+edge-tts==7.0.2
+audiomentations==0.24.0
+pedalboard==0.8.1
+ffmpeg-python==0.2.0
+faiss-cpu==1.11
+ml_collections==1.1.0
+timm==1.0.15
+wandb==0.19.11
+accelerate==1.7.0
+bitsandbytes==0.46.0
+tokenizers==0.19
+huggingface-hub==0.28.1
+transformers==4.41
+https://github.com/noblebarkrr/mvsepless/blob/bd611441e48e918650e6860738894673b3a1a5f1/fixed/fairseq_fixed-0.13.0-cp311-cp311-linux_x86_64.whl
+torchseg==0.0.1a4
+demucs==4.0.0
+asteroid==0.7.0
+prodigyopt==1.1.2
+torch_log_wmse==0.3.0
+rotary_embedding_torch==0.6.5
+local-attention==1.11.1
+tenacity==9.1.2
+gradio==5.38.2
+omegaconf==2.3.0
+beartype==0.18.5
+spafe==0.3.2
+torch_audiomentations==0.12.0
+auraloss==0.4.0
+onnxruntime-gpu>=1.17
+yt_dlp
+python-magic
+pyngrok
+

separator/audio_writer.py

@@ -0,0 +1,85 @@
+from pydub import AudioSegment
+import numpy as np
+
+def write_audio_file(output_file_path, numpy_array, sample_rate, output_format, bitrate):
+    """
+    Записывает аудиофайл из numpy массива в указанном формате с помощью pydub.
+    
+    Параметры:
+        output_file_path (str): Путь для сохранения файла (без расширения)
+        numpy_array (numpy.ndarray): Аудиоданные в виде numpy массива
+        sample_rate (int): Частота дискретизации (в Гц)
+        output_format (str): Формат выходного файла ('mp3', 'flac', 'wav', 'aiff', 'm4a', 'aac', 'ogg', 'opus')
+        encoder_settings (dict, optional): Cловарь с настройками кодировки аудио
+    """
+    try:
+        # Проверка и нормализация входных данных
+        if not isinstance(numpy_array, np.ndarray):
+            raise ValueError("Input must be a numpy array")
+        
+        # Преобразование в правильную форму (samples, channels)
+        if len(numpy_array.shape) == 1:
+            numpy_array = numpy_array.reshape(-1, 1)  # Моно
+        elif len(numpy_array.shape) == 2:
+            if numpy_array.shape[0] == 2:  # Если (channels, samples)
+                numpy_array = numpy_array.T  # Транспонируем в (samples, channels)
+        else:
+            raise ValueError("Input array must be 1D or 2D")
+        
+        # Нормализация до диапазона [-1.0, 1.0] если нужно
+        if np.issubdtype(numpy_array.dtype, np.floating):
+            numpy_array = np.clip(numpy_array, -1.0, 1.0)
+            numpy_array = (numpy_array * 32767).astype(np.int16)
+        elif numpy_array.dtype != np.int16:
+            numpy_array = numpy_array.astype(np.int16)
+        
+        # Создание AudioSegment
+        if numpy_array.shape[1] == 1:  # Моно
+            audio_segment = AudioSegment(
+                numpy_array.tobytes(),
+                frame_rate=sample_rate,
+                sample_width=2,  # 16-bit = 2 bytes
+                channels=1
+            )
+        else:  # Стерео
+            # Для стерео нужно чередовать байты левого и правого каналов
+            interleaved = np.empty((numpy_array.shape[0] * 2,), dtype=np.int16)
+            interleaved[0::2] = numpy_array[:, 0]  # Левый канал
+            interleaved[1::2] = numpy_array[:, 1]  # Правый канал
+            audio_segment = AudioSegment(
+                interleaved.tobytes(),
+                frame_rate=sample_rate,
+                sample_width=2,
+                channels=2
+            )
+        
+        # Формирование параметров экспорта
+        
+        parameters = {}            
+        if bitrate:
+            parameters['bitrate'] = bitrate
+        
+        # Поддержка различных форматов
+        format_mapping = {
+            'mp3': 'mp3',
+            'flac': 'flac',
+            'wav': 'wav',
+            'aiff': 'aiff',
+            'm4a': 'ipod',  # для m4a в pydub используется кодек ipod
+            'aac': 'adts',  # для aac в pydub используется adts
+            'ogg': 'ogg',
+            'opus': 'opus'
+        }
+        
+        if output_format not in format_mapping:
+            raise ValueError(f"Unsupported format: {output_format}. Supported formats are: {list(format_mapping.keys())}")
+        
+        # Добавление расширения файла, если его нет
+        if not output_file_path.lower().endswith(f'.{output_format}'):
+            output_file_path = f"{output_file_path}.{output_format}"
+        
+        # Экспорт в нужный формат
+        audio_segment.export(output_file_path, format=format_mapping[output_format], **parameters)
+    
+    except Exception as e:
+        raise RuntimeError(f"Error writing audio file: {str(e)}")

separator/ensemble.py

@@ -0,0 +1,192 @@
+# coding: utf-8
+__author__ = 'Roman Solovyev (ZFTurbo): https://github.com/ZFTurbo/'
+
+import os
+import sys
+import librosa
+import tempfile
+import soundfile as sf
+import numpy as np
+import argparse
+from separator.audio_writer import write_audio_file
+
+
+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"):
+    """
+    Основная функция для объединения аудиофайлов
+    
+    :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('Ensemble type: {}'.format(ensemble_type))
+    print('Number of input files: {}'.format(len(files)))
+    if weights is not None:
+        weights = np.array(weights)
+    else:
+        weights = np.ones(len(files))
+    print('Weights: {}'.format(weights))
+    print('Output file: {}'.format(output))
+    
+    data = []
+    sr = None
+    for f in files:
+        if not os.path.isfile(f):
+            print('Error. Can\'t find file: {}. Check paths.'.format(f))
+            exit()
+        print('Reading file: {}'.format(f))
+        wav, current_sr = librosa.load(f, sr=None, mono=False)
+        if sr is None:
+            sr = current_sr
+        elif sr != current_sr:
+            print('Error: Sample rates must be equal for all files')
+            exit()
+        print("Waveform shape: {} sample rate: {}".format(wav.shape, sr))
+        data.append(wav)
+    
+    data = np.array(data)
+    res = average_waveforms(data, weights, ensemble_type)
+    print('Result shape: {}'.format(res.shape))
+
+    output_wav = f"{output}_orig.wav"
+    output = f"{output}.{out_format}"
+    
+    if out_format in ["wav", "flac"]:
+    
+        sf.write(output, res.T, sr, subtype='PCM_16')
+        sf.write(output_wav, res.T, sr, subtype='PCM_16')
+
+    elif out_format in ["mp3", "m4a", "aac", "ogg", "opus", "aiff"]:
+
+        write_audio_file(output, res.T, sr, out_format, "320k")
+        sf.write(output_wav, res.T, sr, subtype='PCM_16')
+
+    return output, output_wav
+
+
+    
+# input_settings = [("demucs / v4", 1.0, "vocals"), ("mel_band_roformer / mel_4_stems", 0.5, "vocals")]
+
+# out, wav = ensembless(input_audio, input_settings, "max_fft", format)
+    
+    
+    

separator/models/bandit/core/__init__.py

@@ -0,0 +1,744 @@
+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 models.bandit.core import loss, metrics as metrics_, model
+from models.bandit.core.data._types import BatchedDataDict
+from models.bandit.core.data.augmentation import BaseAugmentor, StemAugmentor
+from models.bandit.core.utils import audio as audio_
+from models.bandit.core.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,
+        )

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

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

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

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

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

@@ -0,0 +1,107 @@
+from abc import ABC
+from typing import Any, Dict, Union
+
+import torch
+import torch_audiomentations as tam
+from torch import nn
+
+from models.bandit.core.data._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

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

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

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

@@ -0,0 +1,69 @@
+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 models.bandit.core.data._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"
+        )

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

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

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

@@ -0,0 +1,392 @@
+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 models.bandit.core.data._types import AudioDict, DataDict
+from models.bandit.core.data.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

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

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

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

@@ -0,0 +1,77 @@
+import os.path
+from typing import Mapping, Optional
+
+import pytorch_lightning as pl
+
+from models.bandit.core.data.musdb.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

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

@@ -0,0 +1,280 @@
+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 models.bandit.core.data._types import AudioDict, DataDict
+from models.bandit.core.data.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_)

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

@@ -0,0 +1,238 @@
+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 core.data._types import DataDict
+from core.data.musdb.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()

separator/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'

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

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

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

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

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

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

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

@@ -0,0 +1,113 @@
+from typing import Any, Dict, Optional
+
+import torch
+from torch import nn
+from torch.nn.modules.loss import _Loss
+
+from models.bandit.core.loss._multistem import MultiStemWrapper
+from models.bandit.core.loss._complex import ReImL1Loss, ReImL2Loss, ReImLossWrapper
+from models.bandit.core.loss.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
+        )

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

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

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

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

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

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

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

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

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

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

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

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

separator/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 models.bandit.core.model.bsrnn.bandsplit import BandSplitModule
+from models.bandit.core.model.bsrnn.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

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

@@ -0,0 +1,139 @@
+from typing import List, Tuple
+
+import torch
+from torch import nn
+
+from models.bandit.core.model.bsrnn.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

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

@@ -0,0 +1,661 @@
+from typing import Dict, List, Optional, Tuple
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from models.bandit.core.model.bsrnn import BandsplitCoreBase
+from models.bandit.core.model.bsrnn.bandsplit import BandSplitModule
+from models.bandit.core.model.bsrnn.maskestim import (
+    MaskEstimationModule,
+    OverlappingMaskEstimationModule
+)
+from models.bandit.core.model.bsrnn.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

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

@@ -0,0 +1,347 @@
+import warnings
+from typing import Dict, List, Optional, Tuple, Type
+
+import torch
+from torch import nn
+from torch.nn.modules import activation
+
+from models.bandit.core.model.bsrnn.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

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

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

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

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

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

@@ -0,0 +1,882 @@
+from pprint import pprint
+from typing import Dict, List, Optional, Tuple, Union
+
+import torch
+from torch import nn
+
+from models.bandit.core.model._spectral import _SpectralComponent
+from models.bandit.core.model.bsrnn.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,
+        )

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

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

separator/models/bandit/model_from_config.py

@@ -0,0 +1,31 @@
+import sys
+import os.path
+import torch
+
+code_path = os.path.dirname(os.path.abspath(__file__)) + '/'
+sys.path.append(code_path)
+
+import yaml
+from ml_collections import ConfigDict
+
+torch.set_float32_matmul_precision("medium")
+
+
+def get_model(
+    config_path,
+    weights_path,
+    device,
+):
+    from models.bandit.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

separator/models/bandit_v2/bandit.py

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

separator/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

separator/models/bandit_v2/film.py

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

separator/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

separator/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)

separator/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)

separator/models/bs_roformer/__init__.py

@@ -0,0 +1,2 @@
+from models.bs_roformer.bs_roformer import BSRoformer
+from models.bs_roformer.mel_band_roformer import MelBandRoformer