initial deploy

.gitignore

@@ -0,0 +1,4 @@
+__pycache__/
+*.pyc
+.DS_Store
+banquet-repo/

README.md

@@ -1,15 +1,24 @@
 ---
 title: Banquet
-emoji: 📉
+emoji: 🍽️
 colorFrom: purple
 colorTo: green
 sdk: gradio
-sdk_version: 6.20.0
-python_version: '3.12'
+sdk_version: 5.28.0
+python_version: '3.11'
 app_file: app.py
 pinned: false
 license: mit
-short_description: todo
+short_description: Separate any instrument from a mix using an audio query
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Banquet
+
+Query-based music source separation: give it a mixture and a 10-second audio
+example of the instrument you want, and it extracts that instrument from the
+mix. No fixed vocals/drums/bass/other setup — any instrument you can supply
+an example of, including niche ones like reeds or organ.
+
+Paper: [A Stem-Agnostic Single-Decoder System for Music Source Separation
+Beyond Four Stems](https://arxiv.org/abs/2406.18747) (Watcharasupat & Lerch,
+ISMIR 2024).

app.py

@@ -0,0 +1,228 @@
+import sys
+
+sys.stdout.reconfigure(line_buffering=True)
+
+try:
+    import spaces
+except ImportError:
+    # keep @spaces.GPU usable as a no-op; ZeroGPU requires this exact name.
+    class spaces:
+        class GPU:
+            def __init__(self, func=None, duration=60):
+                self.func = func
+
+            def __call__(self, *args, **kwargs):
+                if self.func is not None:
+                    return self.func(*args, **kwargs)
+                func = args[0]
+                return func
+
+import os
+import threading
+from types import SimpleNamespace
+
+import torch
+import torchaudio
+import gradio as gr
+from audiotools import AudioSignal
+from pyharp import ModelCard, build_endpoint, load_audio, save_audio
+
+from core.models.ebase import EndToEndLightningSystem
+from core.models.e2e.bandit.bandit import PasstFiLMConditionedBandit
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+CKPT_PATH = os.path.join(os.path.dirname(__file__), "ev-pre-aug.ckpt")
+MODEL_FS = 44100
+QUERY_LENGTH_SECONDS = 10.0
+
+# inference chunking (default: chunk_size_seconds=6.0, hop_size_seconds=0.5,
+# batch_size=12, per repo config config/data/moisesdb-test.yml) — internal
+# windowing detail, not something a musician can meaningfully tune.
+CHUNK_SIZE_SECONDS = 6.0
+HOP_SIZE_SECONDS = 0.5
+INFERENCE_BATCH_SIZE = 12
+
+# architecture kwargs, per repo config config/models/bandit-query-pre.yml
+MODEL_KWARGS = dict(
+    in_channel=2,
+    band_type="musical",
+    n_bands=64,
+    additive_film=True,
+    multiplicative_film=True,
+    film_depth=2,
+    n_sqm_modules=8,
+    emb_dim=128,
+    rnn_dim=256,
+    bidirectional=True,
+    rnn_type="GRU",
+    mlp_dim=512,
+    hidden_activation="Tanh",
+    hidden_activation_kwargs=None,
+    complex_mask=True,
+    use_freq_weights=True,
+    n_fft=2048,
+    win_length=2048,
+    hop_length=512,
+    window_fn="hann_window",
+    wkwargs=None,
+    power=None,
+    center=True,
+    normalized=True,
+    pad_mode="reflect",
+    onesided=True,
+    fs=MODEL_FS,
+    # repo config points this at a stale training-cluster path used only to
+    # seed weights before training; our checkpoint below is loaded strict=True
+    # right after construction and overwrites all of these anyway.
+    pretrain_encoder=None,
+    freeze_encoder=False,
+)
+
+system = None
+model_ready = False  # has the model been moved onto the GPU yet?
+model_loading = True
+model_error = None
+
+
+def load_model():
+    """Build the model and load the checkpoint on CPU only. Do NOT call .to(DEVICE)
+    or otherwise touch CUDA here — ZeroGPU only intercepts CUDA calls made inside an
+    @spaces.GPU-decorated call, not from a background thread. This split is a no-op
+    on CPU-only hardware (DEVICE == "cpu"), so keep it even while testing on a
+    personal CPU-tier Space, before GPU code is added later."""
+    global system, model_loading, model_error
+    try:
+        model = PasstFiLMConditionedBandit(**MODEL_KWARGS)
+        system = EndToEndLightningSystem.load_from_checkpoint(
+            CKPT_PATH,
+            map_location="cpu",
+            strict=True,
+            model=model,
+            loss_handler=None,
+            metrics=None,
+            augmentation_handler=None,
+            inference_handler=SimpleNamespace(
+                fs=MODEL_FS,
+                chunk_size_seconds=CHUNK_SIZE_SECONDS,
+                hop_size_seconds=HOP_SIZE_SECONDS,
+                batch_size=INFERENCE_BATCH_SIZE,
+            ),
+            optimization_bundle=None,
+        )
+        system.eval()
+        print("Model loaded (CPU).")
+    except Exception as e:
+        model_error = str(e)
+        print(f"Load error: {e}")
+    finally:
+        model_loading = False
+
+
+threading.Thread(target=load_model, daemon=True).start()
+
+
+model_card = ModelCard(
+    name="Banquet",
+    description=(
+        "Extracts any instrument from a music mixture using a short audio "
+        "example as a query, instead of a fixed vocals/drums/bass/other setup."
+    ),
+    author="Karn N. Watcharasupat and Alexander Lerch",
+    tags=["source separation", "music"],
+)
+
+
+def _load_resampled(path: str) -> torch.Tensor:
+    """Loads an audio file and resamples it to the model's sample rate.
+    Returns a (channels, samples) float32 tensor."""
+    signal = load_audio(path)
+    audio = signal.audio_data.squeeze(0)
+    if signal.sample_rate != MODEL_FS:
+        audio = torchaudio.functional.resample(
+            audio, orig_freq=signal.sample_rate, new_freq=MODEL_FS
+        )
+    return audio
+
+
+def _ensure_stereo(audio: torch.Tensor) -> torch.Tensor:
+    """The model's architecture is built for a fixed 2-channel input; duplicate
+    mono uploads to stereo and drop any channels beyond the first two."""
+    if audio.shape[0] == 1:
+        audio = audio.repeat(2, 1)
+    elif audio.shape[0] > 2:
+        audio = audio[:2]
+    return audio
+
+
+def _fit_query_length(query: torch.Tensor) -> torch.Tensor:
+    """Truncates or tiles the query to exactly 10 seconds."""
+    target_len = int(QUERY_LENGTH_SECONDS * MODEL_FS)
+    if query.shape[-1] > target_len:
+        query = query[:, :target_len]
+    elif query.shape[-1] < target_len:
+        reps = target_len // query.shape[-1] + 1
+        query = query.repeat(1, reps)[:, :target_len]
+    return query
+
+
+@spaces.GPU
+@torch.inference_mode()
+def process_fn(mixture_path: str, query_path: str) -> str:
+    """Separates the instrument described by the query clip out of the mixture."""
+    global model_ready
+
+    if model_loading:
+        raise gr.Error("Model is still loading, please wait a moment and try again.")
+    if system is None:
+        raise gr.Error(f"Model failed to load: {model_error}")
+    if not model_ready:
+        system.to(DEVICE)  # only safe here, inside @spaces.GPU
+        model_ready = True
+
+    orig_fs = load_audio(mixture_path).sample_rate
+
+    mixture = _ensure_stereo(_load_resampled(mixture_path)).unsqueeze(0).to(DEVICE)
+    query = _fit_query_length(_load_resampled(query_path)).unsqueeze(0).to(DEVICE)
+
+    batch = {
+        "mixture": {"audio": mixture},
+        "query": {"audio": query},
+        "metadata": {"stem": ["target"]},
+        "estimates": {},
+    }
+
+    out = system.chunked_inference(batch)
+    estimate = out["estimates"]["target"]["audio"].squeeze(0).cpu()
+
+    if orig_fs != MODEL_FS:
+        estimate = torchaudio.functional.resample(
+            estimate, orig_freq=MODEL_FS, new_freq=orig_fs
+        )
+
+    output_signal = AudioSignal(estimate, sample_rate=orig_fs)
+    return save_audio(output_signal)
+
+
+with gr.Blocks() as demo:
+    input_components = [
+        gr.Audio(type="filepath", label="Mixture").harp_required(True),
+        gr.Audio(
+            type="filepath",
+            label="Query Example (~10s clip of the instrument you want extracted)",
+        ).harp_required(True),
+    ]
+    output_components = [
+        gr.Audio(type="filepath", label="Separated Audio").set_info(
+            "The instrument extracted from the mixture, matched to the query example."
+        ),
+    ]
+
+    build_endpoint(
+        model_card=model_card,
+        input_components=input_components,
+        output_components=output_components,
+        process_fn=process_fn,
+    )
+
+demo.queue().launch(pwa=True)

core/models/e2e/bandit/bandit.py

@@ -0,0 +1,619 @@
+from typing import Dict, List, Optional, Tuple
+from core.models.e2e.bandit.bandsplit import BandSplitModule
+from core.models.e2e.bandit.maskestim import OverlappingMaskEstimationModule
+from core.models.e2e.bandit.tfmodel import SeqBandModellingModule
+from core.models.e2e.bandit.utils import MusicalBandsplitSpecification
+from core.models.e2e.querier.passt import Passt, PasstWrapper
+from core.types import InputType, OperationMode, SimpleishNamespace
+from torch import Tensor, nn
+import torch
+
+from core.models.e2e.base import BaseEndToEndModule
+from core.models.e2e.conditioners.film import FiLM
+
+import torchaudio as ta
+
+
+class BaseBandit(BaseEndToEndModule):
+
+    def __init__(
+        self,
+        in_channel: 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,
+        fs: int = 44100,
+    ):
+        super().__init__()
+
+        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_channel=in_channel,
+            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_channel: 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_channel=in_channel,
+            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",
+    ):
+        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: InputType, mode: OperationMode = OperationMode.TRAIN):
+
+        with torch.no_grad():
+            x = self.stft(batch.mixture.audio)
+            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)
+
+        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_channel: 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,
+    ):
+        super().__init__(
+            in_channel=in_channel,
+            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.instantiate_mask_estim(
+            in_channel=in_channel,
+            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_channel: 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 = True,
+    ):
+        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_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
+            }
+        )
+
+    def separate(self, batch):
+
+        x, q, length = self.encode(batch)
+
+        for stem, mem in self.mask_estim.items():
+            m = mem(q)
+            s = self.mask(x, m)
+            s = torch.reshape(s, x.shape)
+            batch.estimates[stem] = SimpleishNamespace(
+                audio=self.istft(s, length), spectrogram=s
+            )
+
+        return batch
+
+
+class BaseConditionedBandit(BaseBandit):
+    query_encoder: nn.Module
+
+    def __init__(
+        self,
+        in_channel: 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",
+        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,
+    ):
+        super().__init__(
+            in_channel=in_channel,
+            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.instantiate_mask_estim(
+            in_channel=in_channel,
+            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_channel: int,
+        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 = True,
+    ):
+        if hidden_activation_kwargs is None:
+            hidden_activation_kwargs = {}
+
+        assert n_freq is not None
+
+        self.mask_estim = 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_channel=in_channel,
+            hidden_activation=hidden_activation,
+            hidden_activation_kwargs=hidden_activation_kwargs,
+            complex_mask=complex_mask,
+            use_freq_weights=use_freq_weights,
+        )
+
+    def separate(self, batch):
+
+        x, q, length = self.encode(batch)
+
+        q = self.adapt_query(q, batch)
+
+        m = self.mask_estim(q)
+        s = self.mask(x, m)
+        s = torch.reshape(s, x.shape)
+        batch.estimates["target"] = SimpleishNamespace(
+            audio=self.istft(s, length), spectrogram=s
+        )
+
+        return batch
+
+    def adapt_query(self, q, batch):
+        raise NotImplementedError
+
+
+class PasstFiLMConditionedBandit(BaseConditionedBandit):
+
+    def __init__(
+        self,
+        in_channel: int,
+        band_type: str = "musical",
+        n_bands: int = 64,
+        additive_film: bool = True,
+        multiplicative_film: bool = True,
+        film_depth: int = 2,
+        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,
+        pretrain_encoder = None,
+        freeze_encoder = False
+    ):
+        super().__init__(
+            in_channel=in_channel,
+            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,
+            mlp_dim=mlp_dim,
+            hidden_activation=hidden_activation,
+            hidden_activation_kwargs=hidden_activation_kwargs,
+            complex_mask=complex_mask,
+            use_freq_weights=use_freq_weights,
+            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.query_encoder = Passt(
+            original_fs=fs,
+            passt_fs=32000,
+        )
+        
+        self.film = FiLM(
+            self.query_encoder.PASST_EMB_DIM,
+            emb_dim,
+            additive=additive_film,
+            multiplicative=multiplicative_film,
+            depth=film_depth,
+        )
+        
+        if pretrain_encoder is not None:
+            self.load_pretrained_encoder(pretrain_encoder)
+            
+            for p in self.band_split.parameters():
+                p.requires_grad = not freeze_encoder
+                
+            for p in self.tf_model.parameters():
+                p.requires_grad = not freeze_encoder
+            
+        
+        
+    def load_pretrained_encoder(self, path):
+        
+        state_dict = torch.load(path, map_location="cpu")["state_dict"]
+        
+        state_dict_ = {k.replace("model.", "") if k.startswith("model.") else k: v for k, v in state_dict.items()}
+        
+        state_dict = {}
+
+        for k, v in state_dict_.items():
+            if "mask_estim" in k:
+                continue
+            
+            if "tf_seqband" in k:
+                k = k.replace("tf_seqband", "tf_model.seqband")
+            
+            state_dict[k] = v
+            
+        
+        res = self.load_state_dict(state_dict, strict=False)
+        
+        for k in res.unexpected_keys:
+            if "mask_estim" in k:
+                continue
+            print(f"Unexpected key: {k}")
+        
+        for k in res.missing_keys:
+            print(f"Missing key: {k}")
+            for kw in ["band_split", "tf_model"]:
+                if kw in k:
+                    raise ValueError(f"Missing key: {k}")
+                
+            for kw in ["mask_estim", "query_encoder"]:
+                if kw in k:
+                    continue
+            
+
+
+    def adapt_query(self, q, batch):
+        
+        w = self.query_encoder(batch.query.audio)
+        q = torch.permute(q, (0, 3, 1, 2)) # (batch, n_band, n_time, emb_dim) -> (batch, emb_dim, n_band, n_time)
+        q = self.film(q, w)
+        q = torch.permute(q, (0, 2, 3, 1)) # -> (batch, n_band, n_time, emb_dim)
+        
+        return q
+
+
+    def optimized_forward(self, batch: InputType, mode: OperationMode = OperationMode.TRAIN):
+
+        with torch.no_grad():
+            x = self.stft(batch.mixture.audio)
+            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.optimized_separate(batch)
+
+        return batch
+
+
+    def optimized_separate(self, batch):
+
+        x, q, length = self.encode(batch)
+
+        for stem, query in batch.query.items():
+
+            batch_ = SimpleishNamespace(**batch.__dict__)
+            batch_.query = query
+
+            q = self.adapt_query(q, batch_)
+
+            m = self.mask_estim(q)
+            s = self.mask(x, m)
+            s = torch.reshape(s, x.shape)
+            batch.estimates[stem] = SimpleishNamespace(
+                audio=self.istft(s, length), spectrogram=s
+            )
+
+        return batch

core/models/e2e/bandit/bandsplit.py

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

core/models/e2e/bandit/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 core.models.e2e.bandit.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

core/models/e2e/bandit/tfmodel.py

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

core/models/e2e/bandit/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)

core/models/e2e/base.py

@@ -0,0 +1,34 @@
+import warnings
+from typing import Any, Dict, Optional, Tuple, Union
+import pytorch_lightning as pl
+
+#from audiocraft.models import encodec
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from ...types import (
+    BatchedInputOutput,
+    InputType,
+    OperationMode,
+    OutputType,
+    SimpleishNamespace,
+    TensorCollection
+)
+
+import torchaudio as ta
+
+
+class BaseEndToEndModule(pl.LightningModule):
+
+    def __init__(
+        self,
+    ) -> None:
+        super().__init__()
+
+
+if __name__ == "__main__":
+    model = BaseEndToEndModule()
+    print(model)
+    print(model.__class__.__name__)
+    print(model.__module__)

core/models/e2e/conditioners/base.py

@@ -0,0 +1,27 @@
+from torch import nn
+
+
+class Conditioning(nn.Module):
+    def __init__(
+        self, cond_embedding_dim: int, channels: int, channels_per_group: int = 16
+    ):
+        super().__init__()
+
+        self.channels = channels
+        self.cond_embedding_dim = cond_embedding_dim
+        self.channels_per_group = channels_per_group
+
+        self.gn = nn.GroupNorm(self.channels // self.channels_per_group, self.channels)
+
+    def forward(self, x, w):
+        raise NotImplementedError
+
+
+class PassThroughConditioning(Conditioning):
+    def __init__(
+        self, cond_embedding_dim: int, channels: int, channels_per_group: int = 16
+    ):
+        super().__init__(cond_embedding_dim, channels, channels_per_group)
+
+    def forward(self, x, w):
+        return self.gn(x)

core/models/e2e/conditioners/film.py

@@ -0,0 +1,197 @@
+import math
+import torch
+from core.models.e2e.conditioners.base import Conditioning
+
+
+from torch import nn
+from torch.nn.modules import activation as activation_
+
+
+class FiLM(Conditioning):
+    def __init__(
+        self,
+        cond_embedding_dim: int,
+        channels: int,
+        additive: bool = True,
+        multiplicative: bool = False,
+        depth: int = 1,
+        activation: str = "ELU",
+        channels_per_group: int = 16,
+    ):
+        super().__init__(
+            channels=channels,
+            channels_per_group=channels_per_group,
+            cond_embedding_dim=cond_embedding_dim,
+        )
+
+        self.additive = additive
+        self.multiplicative = multiplicative
+        self.depth = depth
+        self.activation = activation
+
+        Activation = activation_.__dict__[activation]
+
+        if self.multiplicative:
+
+            if depth == 1:
+                self.gamma = nn.Linear(self.cond_embedding_dim, self.channels)
+            else:
+                layers = [nn.Linear(self.cond_embedding_dim, self.channels)]
+                for _ in range(depth - 1):
+                    layers += [Activation(), nn.Linear(self.channels, self.channels)]
+                self.gamma = nn.Sequential(*layers)
+        else:
+            self.gamma = None
+
+        if self.additive:
+            if depth == 1:
+                self.beta = nn.Linear(self.cond_embedding_dim, self.channels)
+            else:
+                layers = [nn.Linear(self.cond_embedding_dim, self.channels)]
+                for _ in range(depth - 1):
+                    layers += [Activation(), nn.Linear(self.channels, self.channels)]
+                self.beta = nn.Sequential(*layers)
+        else:
+            self.beta = None
+
+    def forward(self, x, w):
+
+        x = self.gn(x)
+
+        if self.multiplicative:
+            gamma = self.gamma(w)
+
+            if len(x.shape) == 4:
+                gamma = gamma[:, :, None, None]
+            elif len(x.shape) == 3:
+                gamma = gamma[:, :, None]
+            elif len(x.shape) == 2:
+                pass
+            else:
+                raise ValueError(f"Invalid shape for input tensor: {x.shape}")
+
+            x = gamma * x
+
+        if self.additive:
+            beta = self.beta(w)
+            if len(x.shape) == 4:
+                beta = beta[:, :, None, None]
+            elif len(x.shape) == 3:
+                beta = beta[:, :, None]
+            elif len(x.shape) == 2:
+                pass
+            else:
+                raise ValueError(f"Invalid shape for input tensor: {x.shape}")
+
+            x = x + beta
+
+        return x
+        
+class CosineSimiliarity(Conditioning):
+    def __init__(self, cond_embedding_dim: int, channels: int, channels_per_group: int = 16):
+        super().__init__(cond_embedding_dim, channels, channels_per_group)
+        
+        self.csim = nn.CosineSimilarity(dim=1)
+        self.proj = nn.Linear(self.cond_embedding_dim, self.channels * self.channels)
+        
+    def forward(self, x, w):
+        
+        
+        x = self.gn(x)
+
+        gamma = self.gamma(w)
+
+        if len(x.shape) == 4:
+            gamma = gamma[:, :, None, None]
+        elif len(x.shape) == 3:
+            gamma = gamma[:, :, None]
+        elif len(x.shape) == 2:
+            pass
+        else:
+            raise ValueError(f"Invalid shape for input tensor: {x.shape}")
+        
+        c = self.csim(gamma, x)
+        
+        x = c[:, None, ...] * x
+
+        
+        
+
+
+class GeneralizedBilinear(nn.Bilinear):
+    def __init__(self, in1_features: int, in2_features: int, out_features: int, bias: bool = True, device=None, dtype=None) -> None:
+        super().__init__(in1_features, in2_features, out_features, bias, device, dtype)
+        
+    def forward(self, x1, x2):
+        
+        out = torch.einsum(
+            "bc...,acd,bd->ba...", x1, self.weight, x2
+        )
+        
+        if self.bias is not None:
+            ndim = out.ndim
+            bias = torch.reshape(self.bias, (1, -1) + (1,) * (ndim - 2))
+            
+            out = out + bias
+            
+        return out
+           
+
+class BilinearFiLM(Conditioning):
+    def __init__(
+        self,
+        cond_embedding_dim: int,
+        channels: int,
+        additive: bool = True,
+        multiplicative: bool = False,
+        depth: int = 2,
+        activation: str = "ELU",
+        channels_per_group: int = 16,
+    ):
+        super().__init__(
+            channels=channels,
+            channels_per_group=channels_per_group,
+            cond_embedding_dim=cond_embedding_dim,
+        )
+
+        self.additive = additive
+        self.multiplicative = multiplicative
+        self.depth = depth
+        assert depth == 2, "Only depth 2 is supported for BilinearFiLM"
+        self.activation = activation
+
+        Activation = activation_.__dict__[activation]
+
+        if self.multiplicative:
+            self.gamma_proj = nn.Sequential(
+                nn.Linear(self.cond_embedding_dim, self.channels),
+                Activation(),
+            )
+            self.gamma_bilinear = GeneralizedBilinear(self.channels, self.channels, self.channels)
+        else:
+            self.gamma = None
+
+        if self.additive:
+            self.beta_proj = nn.Sequential(
+                nn.Linear(self.cond_embedding_dim, self.channels),
+                Activation(),
+            )
+            self.beta_bilinear = GeneralizedBilinear(self.channels, self.channels, self.channels)
+        else:
+            self.beta = None
+
+    def forward(self, x, w):
+
+        x = self.gn(x)
+
+        if self.multiplicative:
+            gamma = self.gamma_proj(w)
+            gamma = self.gamma_bilinear(x, gamma)
+            x = gamma * x
+
+        if self.additive:
+            beta = self.beta_proj(w)
+            beta = self.beta_bilinear(x, beta)
+            x = x + beta
+
+        return x

core/models/e2e/querier/passt.py

@@ -0,0 +1,93 @@
+import torch
+import torchaudio as ta
+from hear21passt.base import get_basic_model
+from torch import nn
+
+class Passt(nn.Module):
+
+    PASST_EMB_DIM: int = 768
+    PASST_FS: int = 32000
+
+    def __init__(
+        self,
+        original_fs: int=44100,
+        passt_fs: int=PASST_FS,
+    ):
+        super().__init__()
+
+        self.passt = get_basic_model(mode="embed_only", arch="openmic").eval()
+        self.resample = ta.transforms.Resample(
+            orig_freq=original_fs, new_freq=passt_fs
+        ).eval()
+
+        for p in self.passt.parameters():
+            p.requires_grad = False
+
+    def forward(self, x):
+        """
+        Forward pass of the PasstWrapper model.
+
+        Args:
+            qspec (torch.Tensor): Query spectrogram.
+            qaudio (torch.Tensor): Query audio.
+
+        Returns:
+            torch.Tensor: Embedding output.
+        """
+        with torch.no_grad():
+            x = torch.mean(x, dim=1)
+            x = self.resample(x)
+
+            specs = self.passt.mel(x)[..., :998]
+            specs = specs[:, None, ...]
+            _, z = self.passt.net(specs)
+
+        return z
+
+
+class PasstWrapper(nn.Module):
+
+    PASST_EMB_DIM: int = 768
+    PASST_FS: int = 32000
+
+    def __init__(
+        self,
+        cond_emb_dim: int = 384,
+        original_cond_emb_dim=PASST_EMB_DIM,
+        original_fs: int=44100,
+        passt_fs: int=PASST_FS,
+    ):
+        super().__init__()
+        self.cond_emb_dim = cond_emb_dim
+
+        self.passt = get_basic_model(mode="embed_only", arch="openmic").eval()
+        self.proj = nn.Linear(original_cond_emb_dim, cond_emb_dim) if cond_emb_dim is not None else nn.Identity()
+        self.resample = ta.transforms.Resample(
+            orig_freq=original_fs, new_freq=passt_fs
+        ).eval()
+
+        for p in self.passt.parameters():
+            p.requires_grad = False
+
+    def forward(self, qspec, qaudio):
+        """
+        Forward pass of the PasstWrapper model.
+
+        Args:
+            qspec (torch.Tensor): Query spectrogram.
+            qaudio (torch.Tensor): Query audio.
+
+        Returns:
+            torch.Tensor: Embedding output.
+        """
+        with torch.no_grad():
+            x = torch.mean(qaudio, dim=1)
+            x = self.resample(x)
+
+            specs = self.passt.mel(x)[..., :998]
+            specs = specs[:, None, ...]
+            _, z = self.passt.net(specs)
+
+        z = self.proj(z)
+
+        return z

core/models/ebase.py

@@ -0,0 +1,506 @@
+import math
+import os.path
+from collections import defaultdict
+from itertools import chain, combinations
+from pprint import pprint
+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 torch import nn, optim
+from torch.optim import lr_scheduler
+from torch.optim.lr_scheduler import LRScheduler
+from tqdm import tqdm
+
+from torch.nn import functional as F
+
+from core.types import BatchedInputOutput, OperationMode, RawInputType, SimpleishNamespace
+from core.types import (
+    InputType,
+    OutputType,
+    LossOutputType,
+    MetricOutputType,
+    ModelType,
+    OptimizerType,
+    SchedulerType,
+    MetricType,
+    LossType,
+    OptimizationBundle,
+    LossHandler,
+    MetricHandler,
+    AugmentationHandler,
+    InferenceHandler,
+)
+
+
+class EndToEndLightningSystem(pl.LightningModule):
+    def __init__(
+        self,
+        model: ModelType,
+        loss_handler: LossHandler,
+        metrics: MetricHandler,
+        augmentation_handler: AugmentationHandler,
+        inference_handler: InferenceHandler,
+        optimization_bundle: OptimizationBundle,
+        fast_run: bool = False,
+        commitment_weight: float = 1.0,
+        batch_size: Optional[int] = None,
+        effective_batch_size: Optional[int] = None,
+    ) -> None:
+        super().__init__()
+
+        self.model = model
+
+        self.loss = loss_handler
+
+        self.metrics = metrics
+        self.optimization = optimization_bundle
+        self.augmentation = augmentation_handler
+        self.inference = inference_handler
+
+        self.fast_run = fast_run
+
+        self.model.fast_run = fast_run
+
+        self.commitment_weight = commitment_weight
+
+        self.batch_size = batch_size
+        self.effective_batch_size = effective_batch_size if effective_batch_size is not None else batch_size
+        self.accum_ratio = self.effective_batch_size // self.batch_size if self.effective_batch_size is not None else 1
+
+        self.output_dir = None
+        self.split_size = None
+
+    def configure_optimizers(self) -> Any:
+        optimizer = self.optimization.optimizer.cls(
+            self.model.parameters(),
+            **self.optimization.optimizer.kwargs
+        )
+
+        ret = {
+            "optimizer": optimizer,
+        }
+
+        if self.optimization.scheduler is not None:
+            scheduler = self.optimization.scheduler.cls(
+                optimizer,
+                **self.optimization.scheduler.kwargs
+            )
+            ret["lr_scheduler"] = scheduler
+
+        return ret
+
+    def compute_loss(
+        self,
+        batch: BatchedInputOutput,
+        mode=OperationMode.TRAIN
+    ) -> LossOutputType:
+        loss_dict = self.loss(batch)
+        return loss_dict
+
+    # TODO: move to a metric handler
+    def update_metrics(
+        self,
+        batch: BatchedInputOutput,
+        mode: OperationMode = OperationMode.TRAIN,
+    ) -> None:
+        metrics: MetricType = self.metrics.get_mode(mode)
+
+        for stem, metric in metrics.items():
+            if stem not in batch.estimates.keys():
+                continue
+            metric.update(batch)
+
+    # TODO: move to a metric handler
+    def compute_metrics(self, mode: OperationMode) -> MetricOutputType:
+        metrics: MetricType = self.metrics.get_mode(mode)
+
+        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
+
+    # TODO: move to a metric handler
+    def reset_metrics(self, mode: OperationMode) -> None:
+        metrics: MetricType = self.metrics.get_mode(mode)
+
+        for _, metric in metrics.items():
+            metric.reset()
+
+    def forward(self, batch: RawInputType) -> Tuple[InputType, OutputType]:
+        batch = self.model(batch)
+        return batch
+
+    def common_step(
+        self, batch: RawInputType, mode: OperationMode, batch_idx: int = -1
+    ) -> Tuple[OutputType, LossOutputType]:
+        batch = BatchedInputOutput.from_dict(batch)
+        batch = self.forward(batch)
+
+        loss_dict = self.compute_loss(batch, mode=mode)
+
+        if not self.fast_run:
+            with torch.no_grad():
+                self.update_metrics(batch, mode=mode)
+
+        return loss_dict
+
+    def training_step(self, batch: RawInputType, batch_idx: int) -> LossOutputType:
+        # augmented_batch = self.augmentation(batch, mode=OperationMode.TRAIN)
+
+        self.model.train()
+
+        loss_dict = self.common_step(batch, mode=OperationMode.TRAIN, batch_idx=batch_idx)
+
+        self.log_dict_with_prefix(loss_dict, prefix=OperationMode.TRAIN, prog_bar=True)
+
+        return loss_dict
+
+    def on_train_batch_end(
+        self, outputs: OutputType, batch: RawInputType, batch_idx: int
+    ) -> None:
+
+        if self.fast_run:
+            return
+
+        if (batch_idx + 1) % self.accum_ratio == 0:
+            metric_dict = self.compute_metrics(mode=OperationMode.TRAIN)
+            self.log_dict_with_prefix(metric_dict, prefix=OperationMode.TRAIN)
+            self.reset_metrics(mode=OperationMode.TRAIN)
+
+    @torch.inference_mode()
+    def validation_step(
+        self, batch: RawInputType, batch_idx: int, dataloader_idx: int = 0
+    ) -> Dict[str, Any]:
+
+        self.model.eval()
+
+        with torch.inference_mode():
+            loss_dict = self.common_step(batch, mode=OperationMode.VAL)
+
+        self.log_dict_with_prefix(loss_dict, prefix=OperationMode.VAL)
+
+        return loss_dict
+
+    def on_validation_epoch_start(self) -> None:
+        self.reset_metrics(mode=OperationMode.VAL)
+
+    def on_validation_epoch_end(self) -> None:
+        if self.fast_run:
+            return
+
+        metric_dict = self.compute_metrics(mode=OperationMode.VAL)
+        self.log_dict_with_prefix(
+            metric_dict, OperationMode.VAL, prog_bar=True, add_dataloader_idx=False
+        )
+        self.reset_metrics(mode=OperationMode.VAL)
+
+    
+    def save_to_audio(self, batch: BatchedInputOutput, batch_idx: int) -> None:
+        
+        batch_size = batch["mixture"]["audio"].shape[0]
+        
+        assert batch_size == 1, "Batch size must be 1 for inference"
+        
+        metadata = batch.metadata
+        
+        song_id = metadata["mix"][0]
+        stem = metadata["stem"][0]
+        
+        log_dir = os.path.join(self.logger.log_dir, "audio")
+        
+        os.makedirs(os.path.join(log_dir, song_id), exist_ok=True)
+        
+        audio = batch.estimates[stem]["audio"]
+        
+        audio = audio.squeeze(0).cpu().numpy()
+        
+        audio_path = os.path.join(log_dir, song_id, f"{stem}.wav")
+        
+        ta.save(audio_path, torch.tensor(audio), self.inference.fs)
+
+    def save_vdbo_to_audio(self, batch: BatchedInputOutput, batch_idx: int) -> None:
+        
+        batch_size = batch["mixture"]["audio"].shape[0]
+        
+        assert batch_size == 1, "Batch size must be 1 for inference"
+        
+        metadata = batch.metadata
+        
+        song_id = metadata["song_id"][0]
+        
+        log_dir = os.path.join(self.logger.log_dir, "audio")
+        
+        os.makedirs(os.path.join(log_dir, song_id), exist_ok=True)
+        
+        for stem, audio in batch.estimates.items():
+            audio = audio["audio"]
+            audio = audio.squeeze(0).cpu().numpy()
+            
+            audio_path = os.path.join(log_dir, song_id, f"{stem}.wav")
+            
+            ta.save(audio_path, torch.tensor(audio), self.inference.fs)
+
+    @torch.inference_mode()
+    def chunked_inference(
+        self, batch: RawInputType, batch_idx: int = -1, dataloader_idx: int = 0
+    ) -> BatchedInputOutput:
+        batch = BatchedInputOutput.from_dict(batch)
+        
+        audio = batch["mixture"]["audio"]
+        
+        b, c, n_samples = audio.shape
+        
+        assert b == 1
+
+        fs = self.inference.fs
+
+        chunk_size = int(self.inference.chunk_size_seconds * fs)
+        hop_size = int(self.inference.hop_size_seconds * fs)
+        
+        batch_size = self.inference.batch_size
+        
+        overlap = chunk_size - hop_size
+        
+        scaler = chunk_size / (2 * hop_size)
+
+        n_chunks = int(math.ceil(
+            (n_samples + 4 * overlap - chunk_size) / hop_size
+        )) + 1
+        
+        pad = (n_chunks - 1) * hop_size + chunk_size - n_samples
+
+        # print(audio.shape)
+        audio = F.pad(
+            audio,
+            pad=(2 * overlap, 2 * overlap + pad),
+            mode="reflect"
+        )
+        padded_length = audio.shape[-1]
+        audio = audio.reshape(c, 1, -1, 1)
+        
+        chunked_audio = F.unfold(
+            audio,
+            kernel_size=(chunk_size, 1), 
+            stride=(hop_size, 1)
+        ) # (c, chunk_size, n_chunk)
+
+        # print(chunked_audio.shape)
+
+        chunked_audio = chunked_audio.permute(2, 0, 1).reshape(-1, c, chunk_size)
+        
+        n_chunks = chunked_audio.shape[0]
+        
+        n_batch = math.ceil(n_chunks / batch_size)
+
+        outputs = []
+        
+        for i in tqdm(range(n_batch)):
+            start = i * batch_size
+            end = min((i + 1) * batch_size, n_chunks)
+            
+            chunked_batch = SimpleishNamespace(
+                mixture={
+                    "audio": chunked_audio[start:end]
+                },
+                query=batch["query"],
+                estimates=batch["estimates"]
+            )
+            
+            output = self.forward(chunked_batch)
+            outputs.append(output.estimates["target"]["audio"])
+
+        output = torch.cat(outputs, dim=0) # (n_chunks, c, chunk_size)
+        window = torch.hann_window(chunk_size, device=self.device).reshape(1, 1, chunk_size)
+        output = output * window / scaler
+
+        output = torch.permute(output, (1, 2, 0))
+
+        output = F.fold(
+            output,
+            output_size=(padded_length, 1),
+            kernel_size=(chunk_size, 1),
+            stride=(hop_size, 1)
+        ) # (c, 1, t, 1)
+
+        output = output[None, :, 0, 2*overlap: n_samples + 2*overlap, 0]
+
+        stem = batch.metadata["stem"][0]
+
+        batch["estimates"][stem] = {
+            "audio": output
+        }
+
+        return batch
+
+    def chunked_vdbo_inference(
+        self, batch: RawInputType, batch_idx: int = -1, dataloader_idx: int = 0
+    ) -> BatchedInputOutput:
+        batch = BatchedInputOutput.from_dict(batch)
+        
+        audio = batch["mixture"]["audio"]
+        
+        b, c, n_samples = audio.shape
+        
+        assert b == 1
+
+        fs = self.inference.fs
+
+        chunk_size = int(self.inference.chunk_size_seconds * fs)
+        hop_size = int(self.inference.hop_size_seconds * fs)
+        
+        batch_size = self.inference.batch_size
+        
+        overlap = chunk_size - hop_size
+        
+        scaler = chunk_size / (2 * hop_size)
+
+        n_chunks = int(math.ceil(
+            (n_samples + 4 * overlap - chunk_size) / hop_size
+        )) + 1
+        
+        pad = (n_chunks - 1) * hop_size + chunk_size - n_samples
+
+        # print(audio.shape)
+        audio = F.pad(
+            audio,
+            pad=(2 * overlap, 2 * overlap + pad),
+            mode="reflect"
+        )
+        padded_length = audio.shape[-1]
+        audio = audio.reshape(c, 1, -1, 1)
+        
+        chunked_audio = F.unfold(
+            audio,
+            kernel_size=(chunk_size, 1), 
+            stride=(hop_size, 1)
+        ) # (c, chunk_size, n_chunk)
+
+        # print(chunked_audio.shape)
+
+        chunked_audio = chunked_audio.permute(2, 0, 1).reshape(-1, c, chunk_size)
+        
+        n_chunks = chunked_audio.shape[0]
+        
+        n_batch = math.ceil(n_chunks / batch_size)
+
+        outputs = defaultdict(list)
+        
+        for i in tqdm(range(n_batch)):
+            start = i * batch_size
+            end = min((i + 1) * batch_size, n_chunks)
+            
+            chunked_batch = SimpleishNamespace(
+                mixture={
+                    "audio": chunked_audio[start:end]
+                },
+                estimates=batch["estimates"]
+            )
+            
+            output = self.forward(chunked_batch)
+            
+            for stem, estimate in output.estimates.items():
+                outputs[stem].append(estimate["audio"])
+
+        for stem, outputs_ in outputs.items():                
+
+            output = torch.cat(outputs_, dim=0) # (n_chunks, c, chunk_size)
+            window = torch.hann_window(chunk_size, device=self.device).reshape(1, 1, chunk_size)
+            output = output * window / scaler
+
+            output = torch.permute(output, (1, 2, 0))
+
+            output = F.fold(
+                output,
+                output_size=(padded_length, 1),
+                kernel_size=(chunk_size, 1),
+                stride=(hop_size, 1)
+            ) # (c, 1, t, 1)
+
+            output = output[None, :, 0, 2*overlap: n_samples + 2*overlap, 0]
+
+            batch["estimates"][stem] = {
+                "audio": output
+            }
+
+        return batch
+
+
+    def on_test_epoch_start(self) -> None:
+        self.reset_metrics(mode=OperationMode.TEST)
+
+    def test_step(
+        self, batch: RawInputType, batch_idx: int, dataloader_idx: int = 0
+    ) -> Any:
+
+        self.model.eval()
+        
+        if "query" in batch.keys():
+            batch = self.chunked_inference(batch, batch_idx, dataloader_idx)
+        else:
+            batch = self.chunked_vdbo_inference(batch, batch_idx, dataloader_idx)
+        
+        self.reset_metrics(mode=OperationMode.TEST)
+        self.update_metrics(batch, mode=OperationMode.TEST)
+        metrics = self.compute_metrics(mode=OperationMode.TEST)
+        # metrics["song_id"] = batch.metadata["mix"][0]
+        self.log_dict_with_prefix(metrics, OperationMode.TEST, 
+                                  on_step=True, on_epoch=False, prog_bar=True)
+        self.reset_metrics(mode=OperationMode.TEST)
+
+        # pprint(metrics)
+
+        return batch
+
+    def on_test_epoch_end(self) -> None:
+        self.reset_metrics(mode=OperationMode.TEST)
+
+    def set_output_path(self, output_dir: str) -> None:
+        self.output_dir = output_dir
+
+    def predict_step(
+        self, batch: RawInputType, batch_idx: int, dataloader_idx: int = 0
+    ) -> Any:
+
+        self.model.eval()
+        
+        if "query" in batch.keys():    
+            batch = self.chunked_inference(batch, batch_idx, dataloader_idx)
+            
+            self.save_to_audio(batch, batch_idx)
+        else:
+            batch = self.chunked_vdbo_inference(batch, batch_idx, dataloader_idx)
+            self.save_vdbo_to_audio(batch, batch_idx)
+
+    def load_state_dict(
+        self, state_dict: Mapping[str, Any], strict: bool = False
+    ) -> Any:
+        return super().load_state_dict(state_dict, strict=False)
+
+    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,
+            # on_step=True,
+            # on_epoch=False,
+        )
+
+        self.logger.save()

core/types/__init__.py

@@ -0,0 +1,163 @@
+from types import SimpleNamespace
+from typing import Any, Dict, Optional, TypedDict
+
+import torch
+from torch import nn, optim
+import torchmetrics as tm
+
+
+class OperationMode:
+    TRAIN = "train"
+    VAL = "val"
+    TEST = "test"
+    PREDICT = "predict"
+
+
+RawInputType = Dict
+
+
+def nested_dict_to_nested_namespace(d: dict) -> SimpleNamespace:
+    d_ = d.copy()
+
+    for k, v in d.items():
+        if isinstance(v, dict):
+            v = nested_dict_to_nested_namespace(v)
+
+        d_[k] = v
+
+    return SimpleNamespace(**d_)
+
+
+RawInputType = TypedDict(
+    "RawInputType",
+    {
+        "mixture": torch.Tensor,
+        "sources": Dict[str, torch.Tensor],
+        "estimates": Optional[Dict[str, torch.Tensor]],
+        "metadata": Dict[str, Any],
+    },
+    total=False,
+)
+
+
+def input_dict(
+    mixture: torch.Tensor = None,
+    sources: Dict[str, torch.Tensor] = None,
+    query: torch.Tensor = None,
+    metadata: Dict[str, Any] = None,
+    modality: str = "audio",
+) -> RawInputType:
+
+    out = {
+        "estimates": {
+            k: {
+                modality: torch.empty(
+                    0,
+                )
+            }
+            for k, v in sources.items()
+        }
+    }
+
+    if mixture is not None:
+        out["mixture"] = {modality: torch.from_numpy(mixture).to(torch.float32)}
+
+    if sources is not None:
+        out["sources"] = {k: {modality: torch.from_numpy(v).to(torch.float32)} for k, v in sources.items()}
+
+    if query is not None:
+        out["query"] = {modality: torch.from_numpy(query).to(torch.float32)}
+
+    if metadata is not None:
+        out["metadata"] = metadata
+
+    return out
+
+
+class SimpleishNamespace(SimpleNamespace):
+    def __init__(self, **kwargs: Any) -> None:
+        kwargs_ = kwargs.copy()
+
+        for k, v in kwargs.items():
+            if isinstance(v, dict):
+                v = SimpleishNamespace(**v)
+
+            kwargs_[k] = v
+
+        super().__init__(**kwargs_)
+
+    def copy(self) -> "SimpleishNamespace":
+        return SimpleishNamespace(**{k: v for k, v in self.__dict__.items()})
+
+    def add_subnamespace(self, name: str, **kwargs: Any) -> None:
+        if hasattr(self, name):
+            raise ValueError(f"Namespace already has attribute {name}")
+
+        setattr(self, name, SimpleishNamespace(**kwargs))
+
+    def keys(self):
+        return self.__dict__.keys()
+
+    def __getitem__(self, key: str) -> Any:
+        return self.__dict__[key]
+
+    def __setitem__(self, key: str, value: Any) -> None:
+        self.__dict__[key] = value
+
+    def items(self):
+        return self.__dict__.items()
+
+
+class BatchedInputOutput(SimpleishNamespace):
+    mixture: torch.Tensor
+    sources: Dict[str, torch.Tensor]
+    estimates: Optional[Dict[str, torch.Tensor]]
+    metadata: Dict[str, Any]
+
+    def __init__(self, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+
+    @classmethod
+    def from_dict(cls, d: dict) -> "BatchedInputOutput":
+        return cls(**d)
+
+    def to_dict(self) -> dict:
+        return self.__dict__
+
+
+class TensorCollection(SimpleishNamespace):
+    def __init__(self, **kwargs: torch.Tensor) -> None:
+        super().__init__(**kwargs)
+
+    def apply(self, func: Any, *args: Any, **kwargs: Any) -> "TensorCollection":
+        return TensorCollection(
+            **{k: func(v, *args, **kwargs) for k, v in self.__dict__.items()}
+        )
+
+    def as_stacked_tensor(self, dim: int = 0) -> torch.Tensor:
+        return torch.stack(list(self.__dict__.values()), dim=dim)
+
+    def as_concatenated_tensor(self, dim: int = 0) -> torch.Tensor:
+        return torch.cat(list(self.__dict__.values()), dim=dim)
+
+    def __getitem__(self, key: str) -> torch.Tensor:
+        return self.__dict__[key]
+
+
+InputType = BatchedInputOutput
+OutputType = BatchedInputOutput
+LossOutputType = Any
+MetricOutputType = Any
+
+ModelType = nn.Module
+OptimizerType = optim.Optimizer
+SchedulerType = optim.lr_scheduler._LRScheduler
+MetricType = tm.Metric
+LossType = nn.Module
+
+OptimizationBundle = Any
+
+LossHandler = Any
+MetricHandler = Any
+AugmentationHandler = Any
+InferenceHandler = Any

ev-pre-aug.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:657295888781e62ef50593002720d2edb3858b9e5bbfabf0c54f715a0da4b9e2
+size 645470187

requirements.txt

@@ -0,0 +1,7 @@
+git+https://github.com/TEAMuP-dev/pyharp.git@v0.3.0
+# model-specific deps below:
+pytorch-lightning
+torchmetrics
+librosa
+hear21passt
+timm