initial deploy

.gitignore

@@ -0,0 +1,12 @@
+__pycache__/
+*.pyc
+.DS_Store
+*.pyo
+*.egg-info/
+.eggs/
+dist/
+build/
+*.log
+outputs/
+.hydra/
+wandb/

README.md

@@ -1,15 +1,28 @@
 ---
 title: MEGAMI
-emoji: 🦀
+emoji: 🎚
 colorFrom: yellow
 colorTo: gray
 sdk: gradio
-sdk_version: 6.19.0
-python_version: '3.13'
+sdk_version: 5.28.0
+python_version: "3.10"
 app_file: app.py
 pinned: false
 license: cc-by-nc-sa-4.0
 short_description: Automatic Mixing via Generative Model of Effect Embeddings
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# MEGAMI — Automatic Music Mixing
+
+Upload 2–6 **dry (unprocessed) mono/stereo instrument stems** (at least ~12 s each) and MEGAMI will generate a professionally mixed stereo output.
+
+**Requirements:**
+- Each track must be at least 11.9 s long (525 312 samples @ 44 100 Hz)
+- Supported formats: WAV, FLAC, MP3 (auto-resampled to 44 100 Hz)
+- At least 2 tracks required
+
+**How it works:**
+1. MEGAMI’s diffusion model (FxGenerator) generates audio-effect embeddings from the stems
+2. The effect processor (FxProcessor) applies those embeddings to produce a polished stereo mix
+
+Paper: [arxiv.org/abs/2511.08040](https://arxiv.org/abs/2511.08040) | License: CC BY-NC-SA 4.0

app.py

@@ -0,0 +1,209 @@
+import os
+import shutil
+import tempfile
+import threading
+
+import numpy as np
+import torch
+import torchaudio
+import soundfile as sf
+import gradio as gr
+from pyharp import ModelCard, build_endpoint
+
+# ---- Model card ----
+
+model_card = ModelCard(
+    name="MEGAMI",
+    description=(
+        "Automatic music mixing via a generative model of audio effect embeddings."
+        "Upload 2–6 dry (unprocessed) instrument stems and MEGAMI will generate "
+        "a professionally mixed stereo output."
+    ),
+    author="Eloi Moliner, Marco A. Martínez-Ramírez, Junghyun Koo, Wei-Hsiang Liao, Kin Wai Cheuk, Joan Serrà, Vesa Välimäki, Yuki Mitsufuji",
+    tags=["mixing", "music-production", "generative", "diffusion"],
+)
+
+# ---- Background model loading ----
+
+_inference = None
+_load_error = None
+_model_ready = threading.Event()
+
+MIN_AUDIO_LEN = 525312   # samples @ 44100 Hz ≈ 11.9 s
+SAMPLE_RATE   = 44100
+
+
+def _download_checkpoints():
+    import urllib.request
+    from huggingface_hub import hf_hub_download
+
+    os.makedirs("checkpoints", exist_ok=True)
+
+    github_base = "https://github.com/SonyResearch/MEGAMI/releases/download/v0/"
+    for fname in ["FxGenerator_public.pt", "FxProcessor_public.pt", "CLAP_DA_public.pt"]:
+        dest = os.path.join("checkpoints", fname)
+        if not os.path.exists(dest):
+            print(f"[download] {fname} …")
+            urllib.request.urlretrieve(github_base + fname, dest)
+            print(f"[download] {fname} done.")
+
+    # The MEGAMI README cites the original LAION-CLAP file as the source for this checkpoint;
+    # no patching script is documented — we download the original and save under the expected filename
+    clap_dest = "checkpoints/music_audioset_epoch_15_esc_90.14.patched.pt"
+    if not os.path.exists(clap_dest):
+        print("[download] music_audioset_epoch_15_esc_90.14.patched.pt …")
+        cached = hf_hub_download(
+            repo_id="lukewys/laion_clap",
+            filename="music_audioset_epoch_15_esc_90.14.pt",
+        )
+        shutil.copy(cached, clap_dest)
+        print("[download] music_audioset_epoch_15_esc_90.14.patched.pt done.")
+
+    fxenc_path = "checkpoints/fxenc_plusplus_default.pt"
+    if not os.path.exists(fxenc_path):
+        print("[download] fxenc_plusplus_default.pt …")
+        cached = hf_hub_download(
+            repo_id="yytung/fxencoder-plusplus",
+            filename="fxenc_plusplus_default.pt",
+        )
+        shutil.copy(cached, fxenc_path)
+        print("[download] fxenc_plusplus_default.pt done.")
+
+    print("[download] All checkpoints ready.")
+
+
+def _load_model():
+    global _inference, _load_error
+    try:
+        _download_checkpoints()
+
+        import omegaconf
+        from inference.inference import Inference
+
+        method_args = omegaconf.OmegaConf.create(
+            {
+                "FxGenerator_code": "public",
+                "FxProcessor_code": "public",
+                "T": 30,
+                "cfg_scale": 1.0,
+                "Schurn": 0,
+            }
+        )
+        _inference = Inference(method_args=method_args)
+        print("[model] Inference object ready.")
+    except Exception as exc:
+        import traceback
+        _load_error = str(exc)
+        print(f"[model] Load FAILED: {exc}")
+        traceback.print_exc()
+    finally:
+        _model_ready.set()
+
+
+threading.Thread(target=_load_model, daemon=True).start()
+
+# ---- Process function ----
+
+def process_fn(track1, track2, track3, track4, track5, track6, steps):
+    _model_ready.wait()
+
+    if _load_error:
+        raise gr.Error(f"Model failed to load: {_load_error}")
+
+    track_paths = [t for t in [track1, track2, track3, track4, track5, track6] if t is not None]
+    if len(track_paths) < 2:
+        raise gr.Error("Please upload at least 2 tracks.")
+
+    from utils.feature_extractors.dsp_features import compute_log_rms_gated_max
+
+    dry_tracks   = []
+    dry_segments = []
+
+    for path in track_paths:
+        audio, file_sr = sf.read(path, dtype="float32")
+        if audio.ndim == 1:
+            audio = np.stack([audio, audio], axis=-1)
+        elif audio.shape[1] == 1:
+            audio = np.concatenate([audio, audio], axis=-1)
+
+        audio_tensor = torch.from_numpy(audio).T  # (channels, samples)
+        if file_sr != SAMPLE_RATE:
+            audio_tensor = torchaudio.functional.resample(audio_tensor, file_sr, SAMPLE_RATE)
+
+        x_mono = audio_tensor.mean(dim=0, keepdim=True)  # (1, samples)
+
+        if x_mono.shape[-1] < MIN_AUDIO_LEN:
+            needed_seconds = MIN_AUDIO_LEN / SAMPLE_RATE
+            got_seconds    = x_mono.shape[-1] / SAMPLE_RATE
+            raise gr.Error(
+                f"Each track must be at least {needed_seconds:.1f} s long "
+                f"(got {got_seconds:.1f} s)."
+            )
+
+        x_mono = x_mono.to(_inference.device)
+        segment = _inference.select_high_energy_segment(x_mono)
+
+        dry_tracks.append(x_mono)
+        dry_segments.append(segment)
+
+    x_dry          = torch.stack(dry_tracks)    # (N, 1, L)
+    x_dry_segments = torch.stack(dry_segments)  # (N, 1, MIN_AUDIO_LEN)
+
+    rms    = compute_log_rms_gated_max(x_dry_segments)
+    silent = (rms < -60).squeeze()
+    if silent.ndim == 0:
+        silent = silent.unsqueeze(0)
+    if silent.all():
+        raise gr.Error("All uploaded tracks appear to be silent.")
+    if silent.any():
+        x_dry          = x_dry[~silent]
+        x_dry_segments = x_dry_segments[~silent]
+
+    _inference.method_args.T = int(steps)  # slider returns float, but T is used as a loop bound so casting to int
+
+    embedding_preds = _inference.generate_Fx(x_dry_segments, num_samples=1)
+    fx_embeddings   = _inference.embedding_post_processing(embedding_preds)
+    fx_embedding    = fx_embeddings[0]
+
+    y_final = _inference.apply_effects(x_dry.clone(), fx_embedding)  # (N, 2, L)
+    mix     = y_final.sum(dim=0)                             # (2, L)
+
+    peak = torch.max(torch.abs(mix)).clamp(min=1e-8)  # avoiding div by 0 for silent output
+    mix  = mix / peak
+
+    output_file = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
+    sf.write(
+        output_file.name,
+        mix.cpu().clamp(-1, 1).numpy().T,
+        SAMPLE_RATE,
+        subtype="PCM_16",
+    )
+    return output_file.name
+
+
+# ---- Gradio interface ----
+
+inputs = [
+    gr.Audio(type="filepath", label="Track 1 (required)").harp_required(True),
+    gr.Audio(type="filepath", label="Track 2 (required)").harp_required(True),
+    gr.Audio(type="filepath", label="Track 3"),
+    gr.Audio(type="filepath", label="Track 4"),
+    gr.Audio(type="filepath", label="Track 5"),
+    gr.Audio(type="filepath", label="Track 6"),
+    gr.Slider(
+        minimum=10,
+        maximum=100,
+        step=10,
+        value=30,
+        label="Diffusion Steps (higher = better quality, slower)",
+    ),
+]
+
+outputs = [
+    gr.Audio(type="filepath", label="MEGAMI Mix"),
+]
+
+demo = build_endpoint(inputs, outputs, process_fn, model_card)
+
+if __name__ == "__main__":
+    demo.queue().launch(pwa=True)

conf/conf_FxGenerator_Public.yaml

@@ -0,0 +1,71 @@
+model_dir: "checkpoints"
+
+diff_params:
+  _target_: "diff_params.edm_multitrack_embs.EDM_Multitrack_Embeddings"
+  type: "ve_karras"
+  content_encoder_type: "CLAP"
+  style_encoder_type: "FxEncoder++_DynamicFeatures"
+  sample_rate: 44100
+  cfg_dropout_prob: 0.2
+  sde_hp:
+    sigma_data: 0.025
+    sigma_min: 5e-4
+    sigma_max: 5
+    max_sigma: 5
+    rho: 10
+  CLAP_args:
+    ckpt_path: "checkpoints/music_audioset_epoch_15_esc_90.14.patched.pt"
+    distance_type: "cosine"
+    normalize: True
+    use_adaptor: True
+    adaptor_checkpoint: "checkpoints/CLAP_DA_public.pt"
+    adaptor_type: "MLP_CLAP_regressor"
+    add_noise: True
+    noise_sigma: 0.011
+  fx_encoder_plusplus_args:
+    distance_type: "cosine"
+    ckpt_path: "checkpoints/fxenc_plusplus_default.pt"
+  context_signal: "wet"
+  apply_fxnormaug: False
+  fxnormaug_train: null
+  fxnormaug_inference: null
+  default_shape: [1, 3, 64, 33]
+
+network:
+  _target_: "networks.dit_multitrack.DiffusionTransformer"
+  io_channels: 64
+  embed_dim: 512
+  depth: 16
+  num_heads: 16
+  cond_token_dim: 64
+  global_cond_dim: 0
+  input_concat_dim: 0
+  project_cond_tokens: false
+  transformer_type: "continuous_transformer"
+  pos_emb_strategy: "concatenation"
+  pos_emb_dim: 64
+  pos_emb_type: "one-hot"
+  pos_emb_crossattn_strategy: "concatenation"
+  pos_emb_crossattn_dim: 64
+  pos_emb_crossattn_type: "one-hot"
+  use_taxonomy_in_pos_emb: False
+
+tester:
+  checkpoint: null
+  sampler:
+    _target_: "inference.sampler_euler_heun_multitrack.SamplerEulerHeun"
+  sampling_params:
+    same_as_training: False
+    sde_hp:
+      sigma_data: 0.025
+      sigma_min: 5e-4
+      sigma_max: 5
+      rho: 10
+    Schurn: 0
+    Snoise: 1
+    Stmin: 0
+    Stmax: 10
+    order: 1
+    T: 51
+    schedule: "edm"
+  cfg_scale: 1.0

conf/conf_FxProcessor_Public.yaml

@@ -0,0 +1,40 @@
+model_dir: "checkpoints"
+
+exp:
+  apply_fxnorm: True
+  fxnorm:
+    _target_: "utils.fx_normalization.fxnorm_v2_public.FxNormAug"
+    sample_rate: 44100
+    device: "cpu"
+    features_path: "utils/fx_normalization/features.npy"
+  use_gated_RMSnorm: True
+  RMS_norm: -16
+
+network:
+  _target_: "networks.blackbox_TCN.TCNModel"
+  ninputs: 1
+  noutputs: 2
+  cond_dim: 2112
+  nblocks: 14
+  kernel_size: 15
+  stride: 1
+  dilation_growth: 2
+  channel_growth: 1
+  channel_width: 128
+  stack_size: 15
+  grouped: False
+  causal: False
+  skip_connections: False
+  use_CLAP: True
+  CLAP_args:
+    ckpt_path: "checkpoints/music_audioset_epoch_15_esc_90.14.patched.pt"
+    distance_type: "cosine"
+    normalize: True
+    use_adaptor: False
+    adaptor_checkpoint: null
+    adaptor_type: null
+    add_noise: False
+    noise_sigma: null
+
+tester:
+  checkpoint: null

diff_params/edm_multitrack_embs.py

@@ -0,0 +1,432 @@
+# Copyright (c) 2025 Sony Research
+# Licensed under CC BY-NC-SA 4.0
+# See LICENSE file for details
+
+import torch
+import torch.nn.functional as F
+import sys
+import math
+import time
+import os
+import einops
+import numpy as np
+
+from utils.multitrack_utils import multitrack_batched_processing
+from utils.data_utils import apply_RMS_normalization
+
+class EDM_Multitrack_Embeddings:
+    """
+        This implements the time-frequency domain diffusion
+        Definition of the diffusion parameterization, following ( Karras et al., "Elucidating...", 2022). 
+        This includes only the utilities needed for training, not for sampling.
+    """
+
+    def __init__(self,
+        type,
+        sde_hp,
+        default_shape,
+        cfg_dropout_prob=0.2,
+        sample_rate=44100,
+        context_signal="dry",
+        content_encoder_type="CLAP",
+        style_encoder_type="fx_encoder2048+AFv6",
+        *args,
+        **kwargs
+        ):
+
+        self.type = type
+        self.sde_hp = sde_hp
+
+        self.sigma_data = self.sde_hp.sigma_data  # depends on the training data!! precalculated variance of the dataset
+        self.sigma_min = self.sde_hp.sigma_min
+        self.sigma_max = self.sde_hp.sigma_max
+        self.rho = self.sde_hp.rho
+
+        self.default_shape = torch.Size(default_shape)
+
+        try:
+            self.max_t = self.sde_hp.max_sigma
+        except Exception as e:
+            print(e)
+            print("max_sigma not defined, please add it. It should be the highest sigma value seen during training")
+
+        device=torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.device=device
+
+        self.cfg_dropout_prob = cfg_dropout_prob
+
+        self.context_signal=context_signal
+
+        self.sample_rate=sample_rate
+
+        self.prepare_content_encoder(content_encoder_type, sample_rate, *args, **kwargs)
+        self.prepare_style_encoder(style_encoder_type, *args, **kwargs)
+
+    def prepare_content_encoder(self, type, sample_rate, *args, **kwargs):
+
+        if type=="CLAP":
+            CLAP_args= kwargs.get("CLAP_args", None)
+            assert CLAP_args is not None, "CLAP_args must be provided for CLAP AE"
+
+            # Save original path
+            from utils.feature_extractors.load_features import load_CLAP
+            CLAP_encoder= load_CLAP(CLAP_args, device=self.device)
+
+            def encode_fn(x, *args, **kwargs):
+                x=x.to(self.device)
+
+                type= kwargs.get("type", None)
+                z=CLAP_encoder(x, type) #shape (B, C)
+
+                z=z.view(z.shape[0], 64, -1) #shape (B, 64, N)
+
+                z=z.permute(0, 2, 1) #shape (B, N, 64)
+
+                return z
+            
+            self.content_encode_fn=encode_fn
+
+        else:
+            raise NotImplementedError(f"AE type {AE_type} not implemented")
+
+
+    def prepare_style_encoder(self, type, *args, **kwargs):
+
+        if type=="FxEncoder++_DynamicFeatures":
+
+            Fxencoder_plusplus_args=kwargs.get("fx_encoder_plusplus_args", None)
+
+            from utils.feature_extractors.load_features import load_fx_encoder_plusplus_2048
+            feat_extractor = load_fx_encoder_plusplus_2048(Fxencoder_plusplus_args, self.device)
+
+            from utils.feature_extractors.AF_features_embedding import AF_fourier_embedding
+            AFembedding= AF_fourier_embedding(device=self.device)
+
+            def fxencode_fn(x):
+                """
+                x: tensor of shape [B, C, L] where B is the batch size, C is the number of channels and L is the length of the audio
+                """
+                z=feat_extractor(x)
+                z=torch.nn.functional.normalize(z, dim=-1, p=2)
+                z=z*math.sqrt(z.shape[-1])  # rescale to keep the same scale
+
+                z_af,_=AFembedding.encode(x)
+                z_af=z_af* math.sqrt(z_af.shape[-1])  # rescale to keep the same scale
+
+
+                z_all= torch.cat([z, z_af], dim=-1)
+
+                #now L2 normalize
+
+                norm_z= z_all/ math.sqrt(z_all.shape[-1])  # normalize by dividing by sqrt(dim) to keep the same scale
+
+                norm_z=norm_z.view(norm_z.shape[0], 64, -1)  # Reshape to [B, 64, L//64] where L//64 is the number of frames
+
+                return norm_z
+
+            def reshape_fn(embed):
+                """
+                embed: tensor of shape [B, 64, L//64] where B is the batch size
+                """
+                embed=embed.view(embed.shape[0], -1)
+
+                return embed
+
+
+            self.style_encode_fn=fxencode_fn
+            self.style_reshape=reshape_fn       
+
+        else:
+            raise NotImplementedError(f"FX encoder type {type} not implemented")
+
+    def sample_time_training(self, N):
+        """
+        For training, getting t according to a similar criteria as sampling. Simpler and safer to what Karras et al. did
+        Args:
+            N (int): batch size
+        """
+        a = torch.rand(N)
+        t = (self.sigma_max**(1/self.rho) +a *(self.sigma_min**(1/self.rho) - self.sigma_max**(1/self.rho)))**self.rho
+
+        return t
+
+    def sample_prior(self, shape=None, t=None, dtype=None):
+        """
+        Just sample some gaussian noise, nothing more
+        Args:
+            shape (tuple): shape of the noise to sample, something like (B,T)
+        """
+        assert shape is not None
+        if t is not None:
+            n = torch.randn(shape).to(t.device) * t
+        else:
+            n = torch.randn(shape)
+        return n
+
+    def cskip(self, sigma):
+        """
+        Just one of the preconditioning parameters
+        Args:
+            sigma (float): noise level (equal to timestep is sigma=t, which is our default)
+        
+        """
+        return self.sigma_data**2 *(sigma**2+self.sigma_data**2)**-1
+
+    def cout(self, sigma):
+        """
+        Just one of the preconditioning parameters
+        Args:
+            sigma (float): noise level (equal to timestep is sigma=t, which is our default)
+        """
+        return sigma*self.sigma_data* (self.sigma_data**2+sigma**2)**(-0.5)
+
+    def cin(self, sigma):
+        """
+        Just one of the preconditioning parameters
+        Args:
+            sigma (float): noise level (equal to timestep is sigma=t, which is our default)
+        """
+        return (self.sigma_data**2+sigma**2)**(-0.5)
+
+    def cnoise(self, sigma):
+        """
+        preconditioning of the noise embedding
+        Args:
+            sigma (float): noise level (equal to timestep is sigma=t, which is our default)
+        """
+        return (1/4)*torch.log(sigma)
+
+    def lambda_w(self, sigma):
+        """
+        Score matching loss weighting
+        """
+        return (sigma*self.sigma_data)**(-2) * (self.sigma_data**2+sigma**2)
+
+    def prepare_train_preconditioning(self, x, t, n=None, *args, **kwargs):
+
+        mu, sigma = self._mean(x, t), self._std(t).unsqueeze(-1)
+        sigma = sigma.view(*sigma.size(), *(1,)*(x.ndim - sigma.ndim))
+        if n is None:
+            n=self.sample_prior(shape=x.shape).to(x.device)
+        x_perturbed = mu + sigma *n
+        #self.sample_prior(x.shape).to(x.device)
+
+        cskip = self.cskip(sigma)
+        cout = self.cout(sigma)
+        cin = self.cin(sigma)
+        cnoise = self.cnoise(sigma.squeeze())
+
+        #check if cnoise is a scalar, if so, repeat it
+        if len(cnoise.shape) == 0:
+            cnoise = cnoise.repeat(x.shape[0],)
+        else:
+            cnoise = cnoise.view(x.shape[0],)
+
+        target = 1/cout * (x - cskip * x_perturbed)
+
+        return cin * x_perturbed, target, cnoise
+
+    def loss_fn(self, net,  sample=None,  sample_aug=None, context=None, clusters=None, taxonomy=None, masks=None, *args, **kwargs):
+        """
+        Loss function, which is the mean squared error between the denoised latent and the clean latent
+        Args:
+            net (nn.Module): Model of the denoiser
+            x (Tensor): shape: (B,T) Intermediate noisy latent to denoise
+            sigma (float): noise level (equal to timestep is sigma=t, which is our default)
+        """
+
+        start=time.time()
+        y=sample
+
+        t = self.sample_time_training(y.shape[0]).to(y.device)
+
+        if self.context_signal == "wet":
+            if sample_aug is not None:
+                context = sample_aug
+            else:
+                context = y.clone()  # use the wet signal as context
+
+        else:
+            assert context is not None, "Context must be provided if context_signal is not 'wet'"
+        
+        a=time.time
+
+        with torch.no_grad():
+            
+            y_style=self.style_encode(y,  masks=masks, taxonomy=taxonomy)
+
+            if context is not None:
+                z, x=self.transform_forward(context,  is_condition=True, clusters=clusters, masks=masks, taxonomy=taxonomy, is_wet=(self.context_signal == "wet"))
+                if self.cfg_dropout_prob > 0.0:
+                    null_embed = torch.zeros_like(z, device=z.device)
+                    #dropout context with probability cfg_dropout_prob
+                    mask = torch.rand(z.shape[0], device=z.device) < self.cfg_dropout_prob
+                    z = torch.where(mask.view(-1,1,1,1), null_embed, z)
+
+
+            input, target, cnoise = self.prepare_train_preconditioning(y_style, t )
+
+
+            if len(cnoise.shape)==1:
+                cnoise=cnoise.unsqueeze(-1)
+            if input.ndim==2:
+                input=input.unsqueeze(1)
+
+        estimate = net(input, cnoise, cross_attn_cond=z, taxonomy=taxonomy, mask=masks, cross_attn_cond_mask=masks) 
+        
+        if target.ndim==2 and estimate.ndim==3:
+            estimate=estimate.squeeze(1)
+
+        error=torch.square(torch.abs(estimate-target))
+
+        # do not propagate the error of the padded tracks
+        error= error*masks.view(masks.shape[0], masks.shape[1], 1, 1) 
+
+        compensating_scalar= torch.numel(masks)/ torch.sum(masks, dim=(0,1), keepdim=False).clamp(min=1.0)
+        error= error * compensating_scalar.view(-1, 1, 1, 1)
+
+        return error, self._std(t), x, y
+
+
+    def get_null_embed(self, context):
+        null_embed = torch.zeros_like(context, device=context.device)
+        return null_embed
+
+
+    def denoiser(self, xn , net, t, cond=None,cfg_scale=1.0, masks=None, taxonomy=None, **kwargs):
+        """
+        This method does the whole denoising step, which implies applying the model and the preconditioning
+        Args:
+            x (Tensor): shape: (CQT shape?) Intermediate noisy latent to denoise
+            model (nn.Module): Model of the denoiser
+            sigma (float): noise level (equal to timestep is sigma=t, which is our default)
+        """
+        sigma = self._std(t).unsqueeze(-1)
+        sigma = sigma.view(*sigma.size(), *(1,)*(xn.ndim - sigma.ndim))
+
+        cskip = self.cskip(sigma)
+        cout = self.cout(sigma)
+        cin = self.cin(sigma)
+        cnoise = self.cnoise(sigma.squeeze())
+
+        #check if cnoise is a scalar, if so, repeat it
+        if len(cnoise.shape) == 0:
+            cnoise = cnoise.repeat(xn.shape[0],).unsqueeze(-1)
+        else:
+            cnoise = cnoise.view(xn.shape[0],).unsqueeze(-1)
+
+
+        x_in=cin*xn
+
+        if cfg_scale == 1.0:
+            net_out=net(x_in, cnoise.to(torch.float32), cross_attn_cond=cond, mask=masks, taxonomy=taxonomy,  cross_attn_cond_mask=masks)  #this will crash because of broadcasting problems, debug later!
+        else:
+            null_embed = self.get_null_embed(cond)
+
+            inputs_cond= torch.cat([cond, null_embed], dim=0)
+
+            x_in_cat= torch.cat([x_in, x_in], dim=0)
+
+            cnoise= torch.cat([cnoise, cnoise], dim=0)
+
+            masks_in= torch.cat([masks, masks], dim=0) if masks is not None else None
+
+
+            net_out_batch=net(x_in_cat, cnoise.to(torch.float32), cross_attn_cond=inputs_cond , mask=masks_in,  cross_attn_cond_mask=masks_in)  #this will crash because of broadcasting problems, debug later!0
+
+            cond_output, uncond_output = torch.chunk(net_out_batch, 2, dim=0)
+
+            net_out = uncond_output + (cond_output - uncond_output) * cfg_scale
+
+
+        x_hat=cskip*xn + cout*net_out   
+        x_hat=x_hat* masks.view(masks.shape[0], masks.shape[1], 1, 1) 
+
+        return x_hat
+
+    def style_encode(self, x,  masks=None, taxonomy=None, use_adaptor=False):
+        """
+        Encode the input audio using the style encoder
+        Args:
+            x (Tensor): shape: (B,N, C, T) Audio to encode
+            masks (Tensor): shape: (B, N) Mask indicating which tracks are present in the batch
+        """
+        def apply_fxenc(x_masked, taxonommy=None):
+            x_emb=self.style_encode_fn(x_masked)
+
+            return x_emb
+
+        assert masks is not None, "masks must be provided for style encoding"
+
+        output_emb=multitrack_batched_processing( x, taxonomy=taxonomy ,function=apply_fxenc, class_dependent=False, masks=masks)
+
+        return output_emb
+
+    def _mean(self, x, t):
+        return x
+    
+    def _std(self, t):
+        return t
+    
+    def _ode_integrand(self, x, t, score):
+        return -t * score
+    
+    def transform_inverse(self, z):
+        #shape is (B, N, C, T)
+        B, N, C, T = z.shape
+        # Reshape z to (B*N, C, T)
+        z_reshaped = einops.rearrange(z, "b n c t -> (b n) c t")
+        z_reshaped= self.style_reshape(z_reshaped)
+
+        # Reshape back to (B, N, C)
+
+        z_out = einops.rearrange(z_reshaped, "(b n) c -> b n c", b=B, n=N)
+
+        return z_out
+        
+    def preprocessor(self, x, is_test=False, taxonomy=None):
+            """
+                x: tensor of shape (BxN, C, T) where B is the batch size, N is the number of tracks, C is the number of channels and T is the number of time steps
+                taxonomy: list of lists of strings, where each string is the taxonomy of the track with length BxN. It may be useful if we want to apply different augmentations depending on the taxonomy of the track.
+            """
+            if x.shape[1] == 2:
+                        x = torch.mean(x, dim=1, keepdim=True).expand(-1, 2, -1)  # convert to stereo if it is mono
+            elif x.shape[1] == 1:  # if context is mono, we expand it to stereo
+                        x = x.expand(-1, 2, -1)
+                    
+            if not is_test:
+                #random flip
+                if np.random.rand() > 0.5:
+                    x = -x
+
+            #rms normalize context to -25 dB
+            x= apply_RMS_normalization(x, -25, device=self.device)
+
+            return x
+
+    def Tweedie2score(self, tweedie, xt, t, *args, **kwargs):
+        return (tweedie - self._mean(xt, t)) / self._std(t)**2
+
+    def score2Tweedie(self, score, xt, t, *args, **kwargs):
+        return self._std(t)**2 * score + self._mean(xt, t)
+
+    def transform_forward(self, x, y=None,  is_condition=False, is_test=False, clusters=None, masks=None, taxonomy=None, is_wet=False):
+
+        assert masks is not None
+
+        #convert y to mono and rms normalize itdd
+
+        def prerprocess_and_encode(x_masked, taxonomy=None):
+            if is_condition:
+
+                x_masked=self.preprocessor(x_masked, is_test=is_test, taxonomy=taxonomy)
+    
+            with torch.no_grad():
+                x_emb=self.content_encode_fn(x_masked, type="wet" if is_wet else "dry")
+
+            return x_emb, x_masked
+
+        z, x_out=multitrack_batched_processing(
+            x, taxonomy=taxonomy, function=prerprocess_and_encode, class_dependent=False, masks=masks, number_outputs=2
+        )
+        return z, x_out

inference/inference.py

@@ -0,0 +1,517 @@
+# Copyright (c) 2025 Sony Research
+# Licensed under CC BY-NC-SA 4.0
+# See LICENSE file for details
+
+import sys
+import os
+from utils.feature_extractors.dsp_features import compute_log_rms_gated_max
+
+import pyloudnorm as pyln
+
+import hydra
+import torch
+import torchaudio
+from hydra import initialize, compose
+import utils.training_utils as tr_utils
+
+import torch
+import omegaconf
+import math
+
+import soundfile as sf
+from utils.data_utils import apply_RMS_normalization
+
+import numpy as np
+import glob
+
+from utils.data_utils import read_wav_segment
+
+
+def load_audio( file, start=None, end=None, stereo=True):
+
+            x, fs=read_wav_segment(file, start, end)
+            if stereo:
+                if len(x.shape)==1:
+                    #print( "dry not stereo , doubling channels", x_dry.shape)
+                    x=x[:,np.newaxis]
+                    x= np.concatenate((x, x), axis=-1)
+                elif len(x.shape)==2 and x.shape[-1]==1:
+                    #print( "dry not stereo , doubling channels", x_dry.shape)
+                    x = np.concatenate((x, x), axis=-1)
+
+            x=torch.from_numpy(x).permute(1,0)
+
+            return x, fs
+
+
+class Inference:
+    def __init__(
+        self,
+        method_args=None,
+        path_benchmark="/add/the/path/to/the/benchmark/data/here",
+        load_segment_length=525312, #segment length used for loading and extract CLAP embeddings
+        processor_segment_length=525312, #segment length used for loading and extract CLAP embeddings
+        processor_overlap=8192,
+    ):
+
+        self.method_args = method_args
+        self.path_benchmark = path_benchmark
+        self.load_segment_length=load_segment_length
+        self.processor_segment_length=processor_segment_length
+        self.processor_overlap=processor_overlap
+
+        self.FxGenerator_code = method_args.FxGenerator_code
+        self.FxProcessor_code = method_args.FxProcessor_code
+
+        self.config_file_rel = "../conf"
+        # self.config_path="/home/eloi/projects/project_mfm_eloi/src/conf"
+
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        self.load_FxGenerator()  # Load the S1 model
+        self.load_FxProcessor()  # Load the S2 model
+        self.prepare_feature_extractors()  # Prepare the feature extractors
+
+    def load_FxGenerator(self):
+        if self.FxGenerator_code == "public":
+            config_name = "conf_FxGenerator_Public.yaml"
+            model_dir = "checkpoints"
+            ckpt = "FxGenerator_public.pt"
+        else:
+            raise ValueError(f"Unknown FxGenerator_code: {self.FxGenerator_code}")
+
+        overrides = [
+            f"model_dir={model_dir}",
+            f"tester.checkpoint={ckpt}",
+        ]
+
+        with initialize(version_base=None, config_path=self.config_file_rel):
+            args = compose(config_name=config_name, overrides=overrides)
+
+        if not os.path.exists(args.model_dir):
+            raise Exception(f"Model directory {args.model_dir} does not exist")
+
+        diff_params = hydra.utils.instantiate(args.diff_params)
+
+        network = hydra.utils.instantiate(args.network)
+        network = network.to(self.device)
+        state_dict = torch.load(
+            os.path.join(args.model_dir, args.tester.checkpoint),
+            map_location=self.device,
+            weights_only=False,
+        )
+
+        tr_utils.load_state_dict(state_dict, ema=network)
+
+        self.sampler = hydra.utils.instantiate(
+            args.tester.sampler,
+            network,
+            diff_params,
+            args,
+        )
+
+    def load_FxProcessor(self):
+
+        ### Loading effects model ###
+
+        if self.FxProcessor_code == "public":
+            config_name = "conf_FxProcessor_Public.yaml"
+            model_dir = "checkpoints"
+            ckpt = "FxProcessor_public.pt"
+        else:
+            raise ValueError(f"Unknown FxProcessor_code: {self.FxProcessor_code}")
+
+        overrides = [
+            f"model_dir={model_dir}",
+            f"tester.checkpoint={ckpt}",
+        ]
+
+        with initialize(version_base=None, config_path=self.config_file_rel):
+            args = compose(config_name=config_name, overrides=overrides)
+
+        if not os.path.exists(args.model_dir):
+            raise Exception(f"Model directory {args.model_dir} does not exist")
+
+        fx_model = hydra.utils.instantiate(args.network)
+        self.fx_model = fx_model.to(self.device)
+        state_dict = torch.load(
+            os.path.join(args.model_dir, args.tester.checkpoint),
+            map_location=self.device,
+            weights_only=False,
+        )
+
+        tr_utils.load_state_dict(state_dict, network=fx_model)
+
+        if args.exp.apply_fxnorm:
+            print("Applying fx_normalizer")
+            if "public" in self.FxProcessor_code:
+                fx_normalizer = hydra.utils.instantiate(args.exp.fxnorm, device=str(self.device))
+                self.fx_normalizer = lambda x: fx_normalizer(
+                    x, use_gate=args.exp.use_gated_RMSnorm, RMS=args.exp.RMS_norm
+                )
+            else:
+                self.fx_normalizer = hydra.utils.instantiate(args.exp.fxnorm)
+
+        else:
+            print("No fx_normalizer specified, using identity function")
+            self.fx_normalizer = (
+                lambda x: x
+            )  # identity function if no fx_normalizer is specified
+
+
+    def prepare_feature_extractors(self):
+
+        ### preparing feature extractor ###
+
+        Fxencoder_kwargs = omegaconf.OmegaConf.create(
+            {
+                "ckpt_path": "checkpoints/fxenc_plusplus_default.pt"
+            }
+        )
+
+        from utils.feature_extractors.load_features import load_fx_encoder_plusplus_2048
+
+        feat_extractor = load_fx_encoder_plusplus_2048(Fxencoder_kwargs, self.device)
+
+        from utils.feature_extractors.AF_features_embedding import AF_fourier_embedding
+
+        AFembedding = AF_fourier_embedding(device=self.device)
+
+        def FxEnc(x):
+            """
+            x: tensor of shape [B, C, L] where B is the batch size, C is the number of channels and L is the length of the audio
+            """
+            z = feat_extractor(x)
+            z = torch.nn.functional.normalize(
+                z, dim=-1, p=2
+            )  # normalize to unit variance
+            z = z * math.sqrt(z.shape[-1])  # rescale to keep the same scale
+
+            z_af, _ = AFembedding.encode(x)
+            z_af = z_af * math.sqrt(z_af.shape[-1])  # rescale to keep the same scale
+
+            z_all = torch.cat([z, z_af], dim=-1)
+
+            # now L2 normalize
+
+            norm_z = z_all / math.sqrt(
+                z_all.shape[-1]
+            )  # normalize by dividing by sqrt(dim) to keep the same scale
+
+            return norm_z
+
+        self.FxEnc = FxEnc
+
+        def embedding_post_processing(z):
+            """
+            L2 normalize each of the features in z
+            """
+            z_fxenc = z[
+                ..., :2048
+            ]  # assuming the FxEncoder features are the first 2048 dimensions
+            z_af = z[
+                ..., 2048:
+            ]  # assuming the AF features are the last 2048 dimensions
+
+            z_fxenc = torch.nn.functional.normalize(
+                z_fxenc, dim=-1, p=2
+            )  # normalize to unit variance
+            z_af = torch.nn.functional.normalize(z_af, dim=-1, p=2)
+
+            z_fxenc = z_fxenc * math.sqrt(
+                z_fxenc.shape[-1]
+            )  # rescale to keep the same scale
+            z_af = z_af * math.sqrt(z_af.shape[-1])  # rescale to
+
+            z_all = torch.cat([z_fxenc, z_af], dim=-1)
+
+            return z_all / math.sqrt(
+                z_all.shape[-1]
+            )  # normalize by dividing by sqrt(dim) to keep the same scale
+
+        self.embedding_post_processing = embedding_post_processing
+
+        def get_log_rms_from_z(z):
+
+            z = z * math.sqrt(z.shape[-1])  # rescale to keep the same scale
+            AF = z[..., 2048:]  # assuming the AF features are the last 2048 dimensions
+            AF = AF / math.sqrt(AF.shape[-1])  # normalize to unit variance
+
+            features = AFembedding.decode(AF)
+            log_rms = features[0]
+
+            return log_rms
+
+        def generate_Fx(
+            x, input_type="dry", num_samples=1, T=30, cfg_scale=1.0, Schurn=10
+        ):
+            N, C, L = (
+                x.shape
+            )  # B is the batch size, N is the number of tracks, C is the number of channels and L is the length of the audio
+            B = 1
+
+            shape = self.sampler.diff_params.default_shape
+            shape = [
+                num_samples,
+                N,
+                *shape[2:],
+            ]  # B is the batch size, we want to sample B samples
+
+            masks_fwd = torch.ones(
+                (B, N), dtype=torch.bool, device=self.device
+            )  # Create masks for all tracks, assuming all tracks are present
+            masks_diff = torch.ones(
+                (num_samples, N), dtype=torch.bool, device=self.device
+            )  # Create masks for all tracks, assuming all tracks are present
+
+            self.sampler.T = T
+            self.sampler.Schurn = Schurn  # Set the Schurn parameter for the sampler
+
+            with torch.no_grad():
+                is_wet = "wet" in input_type
+                cond, x_preprocessed = self.sampler.diff_params.transform_forward(
+                    x.unsqueeze(0),
+                    is_condition=True,
+                    is_test=True,
+                    masks=masks_fwd,
+                    is_wet=is_wet,
+                )
+                cond = cond.expand(
+                    shape[0], -1, -1, -1
+                )  # Expand the condition to match the batch size
+                preds, noise_init = self.sampler.predict_conditional(
+                    shape,
+                    cond=cond.contiguous(),
+                    cfg_scale=cfg_scale,
+                    device=self.device,
+                    masks=masks_diff,
+                )
+
+            return preds
+
+        self.generate_Fx = lambda x, num_samples: generate_Fx(
+            x,
+            input_type="wet",
+            num_samples=num_samples,
+            T=self.method_args.T,
+            cfg_scale=self.method_args.cfg_scale,
+            Schurn=self.method_args.Schurn,
+        )
+
+
+        def apply_rms(y_hat, z_pred):
+            """
+            Apply RMS normalization to the generated audio y_hat based on the predicted features z_pred.
+            """
+            pred_logrms = get_log_rms_from_z(
+                z_pred
+            )  # get the log RMS from the generated features
+            pred_rms = 10 ** (pred_logrms / 20)  # convert log RMS to linear scale
+
+            log_rms_y_hat = compute_log_rms_gated_max(
+                y_hat, sample_rate=44100, threshold=-60
+            )
+
+            rms_y_hat = 10 ** (log_rms_y_hat / 20)  # convert log RMS to linear scale
+
+            gain = pred_rms / (
+                rms_y_hat + 1e-6
+            )  # Compute the gain to apply to the generated audio
+
+            y_final = y_hat * gain.unsqueeze(-1)
+
+            return y_final
+
+        self.apply_rms = apply_rms
+
+        def apply_effects(x, z_pred):
+            segment_length = self.processor_segment_length
+            overlap = self.processor_overlap
+            total_length = x.shape[-1]
+            batch_size = x.shape[0]
+        
+            # Normalize input and conditioning outside block loop
+            x_norm = x.mean(dim=1, keepdim=True)
+        
+            if total_length > segment_length:
+                y_final = torch.zeros((batch_size, 2, total_length), device=x.device, dtype=x.dtype)
+        
+                hann = torch.hann_window(overlap * 2, device=x.device, dtype=x.dtype)
+                hann_left = hann[:overlap].view(1, 1, -1)
+                hann_right = hann[overlap:].view(1, 1, -1)
+
+                step = segment_length - overlap
+                positions = list(range(0, total_length - overlap, step))
+                for i, start in enumerate(positions):
+                    end = min(start + segment_length, total_length)
+
+                    seg_x_norm = x_norm[..., start:end]
+
+                    #check activity in seg_x_norm
+
+                    rms_dry_segment=compute_log_rms_gated_max(seg_x_norm, sample_rate=44100)  # Compute the log RMS of the dry audio
+                    indices_non_silent = torch.where(rms_dry_segment > -45)[0]  # Identify silent tracks
+
+
+                    seg_x_norm_non_silent = seg_x_norm[indices_non_silent]
+                    z_pred_non_silent = z_pred[indices_non_silent]
+
+                    if "public" in self.FxProcessor_code:
+                        seg_x_norm_non_silent = self.fx_normalizer(seg_x_norm_non_silent)
+                    else:
+                        seg_x_norm_non_silent = apply_RMS_normalization(seg_x_norm_non_silent, -25.0, device=self.device, use_gate=True)
+                        seg_x_norm_non_silent = self.fx_normalizer(seg_x_norm_non_silent, use_gate=True)
+        
+                    with torch.no_grad():
+                        seg_y_hat_non_silent=torch.zeros((seg_x_norm_non_silent.shape[0],2, seg_x_norm_non_silent.shape[2]), device=x.device, dtype=x.dtype)
+                        #I thought it may be better (but less efficient) to run it like this instead of in parallel. To avoid OOM issues.
+                        for i in range(seg_x_norm_non_silent.shape[0]):
+                            seg_y_hat_non_silent[i] = self.fx_model(seg_x_norm_non_silent[i].unsqueeze(0), z_pred_non_silent[i].unsqueeze(0)).squeeze(0)
+
+                    seg_y_hat_non_silent = apply_rms(seg_y_hat_non_silent, z_pred_non_silent)
+                    
+                    #fill with zeros the silent segments
+                    seg_y_hat= torch.zeros((seg_x_norm.shape[0], seg_y_hat_non_silent.shape[1], seg_y_hat_non_silent.shape[2]), device=x.device, dtype=x.dtype)
+                    seg_y_hat[indices_non_silent]=seg_y_hat_non_silent
+
+        
+                    seg_len = end - start
+        
+                    if i == 0:
+                        # First segment
+                        y_final[..., start:end-overlap] += seg_y_hat[..., :seg_len-overlap]
+                        y_final[..., end-overlap:end] += seg_y_hat[..., seg_len-overlap:] * hann_right
+                    elif end == total_length:
+                        # Last segment
+                        y_final[..., start:start+overlap] += seg_y_hat[..., :overlap] * hann_left
+                        y_final[..., start+overlap:end] += seg_y_hat[..., overlap:]
+                    else:
+                        # Middle segments
+                        y_final[..., start:start+overlap] += seg_y_hat[..., :overlap] * hann_left
+                        y_final[..., start+overlap:end-overlap] += seg_y_hat[..., overlap:seg_len-overlap]
+                        y_final[..., end-overlap:end] += seg_y_hat[..., seg_len-overlap:] * hann_right
+        
+                return y_final
+        
+            else:
+                with torch.no_grad():
+                    y_hat=torch.zeros((x_norm.shape[0], 2, x_norm.shape[2]), device=x.device, dtype=x.dtype)
+                    for i in range(x_norm.shape[0]):
+                        y_hat[i] = self.fx_model(x_norm[i].unsqueeze(0), z_pred[i].unsqueeze(0)).squeeze(0)
+                y_final = apply_rms(y_hat, z_pred)
+                return y_final        
+
+        self.apply_effects = apply_effects
+
+    def select_high_energy_segment(self, x_dry, seq_length=525312):
+            C, L = x_dry.shape
+            
+            track = x_dry
+                
+            # Calculate energy for windows of size seq_length
+            num_windows = L - seq_length + 1
+            max_energy = 0
+            max_energy_start = 0
+                
+            for i in range(0, num_windows, 1000):  # Step by 1000 for efficiency
+                    segment = track[..., i:i+seq_length]
+                    energy = (segment ** 2).sum()
+                    
+                    if energy > max_energy:
+                        max_energy = energy
+                        max_energy_start = i
+                
+            # Fine-tune search around the best region
+            fine_start = max(0, max_energy_start - 1000)
+            fine_end = min(L - seq_length + 1, max_energy_start + 1000)
+                
+            for i in range(fine_start, fine_end):
+                    segment = track[..., i:i+seq_length]
+                    energy = (segment ** 2).sum()
+                    
+                    if energy > max_energy:
+                        max_energy = energy
+                        max_energy_start = i
+                
+            return x_dry[:, max_energy_start:max_energy_start+seq_length]
+
+
+    def run_inference_single_song(self, exp_name="test_Sep28", directory=None, num_samples=1):
+        """
+        Run the inference on a single example
+        """
+
+        dry_files=glob.glob(os.path.join(directory, "*.wav"))
+
+        assert len(dry_files) > 0, f"No .wav files found in {directory}"
+        print(f"Found {len(dry_files)} dry files in {directory}")
+        print(dry_files)
+
+        dry_tracks=[]
+        dry_tracks_segments=[]
+
+        for f in dry_files:
+            x_dry_i, fs=load_audio(str(f), stereo=True)
+            x_dry_i=x_dry_i.to(self.device)
+
+            if fs!=self.sampler.diff_params.sample_rate:
+                x_dry_i=torchaudio.functional.resample(x_dry_i, orig_sr=fs, target_sr=self.sampler.diff_params.sample_rate)
+            #stereo to mono
+            x_dry_i = x_dry_i.mean(dim=0, keepdim=True)
+            dry_tracks.append(x_dry_i )
+
+            if x_dry_i.shape[-1] >= self.load_segment_length:
+                # search for each track  the segment of seq_length size that has highest energy
+                # x_dry_i shape is (N, C, L) where N is the number of tracks, C is the number of channels and L is the length of the audio
+                x_dry_i_segment = self.select_high_energy_segment(x_dry_i, seq_length=self.load_segment_length)
+            else:
+                raise ValueError(f"Input audio {f} is too short, needs to be at least {self.load_segment_length/self.sampler.diff_params.sample_rate:.2f} seconds.")
+
+            dry_tracks_segments.append(x_dry_i_segment)
+        
+        x_dry=torch.stack(dry_tracks, dim=0)  # shape (N, C, L) where N is the number of tracks, C is the number of channels and L is the length of the audio
+        x_dry_segments= torch.stack(dry_tracks_segments, dim=0)  # shape (N, C, L) where N is the number of tracks, C is the number of channels and L is the length of the audio
+
+        #first check if all tracks in x_dry have activity (RMS > -60 dBFS)
+        rms_dry=compute_log_rms_gated_max(x_dry_segments, sample_rate=44100)  # Compute the log RMS of the dry audio
+        silent_tracks = rms_dry < -60  # Identify silent tracks
+        silent_tracks = silent_tracks.squeeze()  # Remove singleton dimensions
+
+        if silent_tracks.any():
+            print(f"Removing {silent_tracks.sum()} silent tracks from the input audio.")
+            #shape before removing silent tracks is (N, C, L)
+            x_dry = x_dry[~silent_tracks]  # Remove silent tracks
+            x_dry_segments = x_dry_segments[~silent_tracks]  # Remove silent tracks
+
+        preds = self.generate_Fx(x_dry_segments, num_samples)
+
+        z_pred = self.embedding_post_processing(
+            preds
+        )  # post-process the generated features
+
+        del self.sampler
+
+        for i in range(num_samples):
+            z_i = z_pred[
+                i
+            ]  # Randomly sample 100 features from the generated features
+
+            y_final = self.apply_effects(
+                    x_dry.clone(), z_i
+            )  # Apply the effects to the input audio
+            y_hat_mixture = y_final.sum(dim=0, keepdim=False)
+    
+            # peak normalization of y_hat_mixture
+            peak = torch.max(torch.abs(y_hat_mixture))
+            y_hat_mixture /= peak  # Normalize the audio to [-1, 1]
+    
+            filename="MEGAMI_inference"+f"_sample{i}.wav"
+            os.makedirs(f"{directory}/{exp_name}", exist_ok=True)
+            sf.write(
+                f"{directory}/{exp_name}/{filename}",
+                y_hat_mixture.cpu().clamp(-1, 1).numpy().T,
+                44100,
+                subtype="PCM_16",
+            )
+
+

inference/sampler.py

@@ -0,0 +1,29 @@
+
+import omegaconf
+import os
+
+
+class Sampler():
+
+    def __init__(self, model, diff_params, args):
+
+        self.model = model.eval() #is it ok to do this here?
+        self.diff_params = diff_params #same as training, useful if we need to apply a wrapper or something
+        self.args=args
+        if self.args.tester.sampling_params.same_as_training:
+            self.sde_hp = diff_params.sde_hp
+        else:
+            self.sde_hp = self.args.tester.sampling_params.sde_hp
+
+        self.T = self.args.tester.sampling_params.T
+        self.step_counter = 0
+
+
+    #def setup_wandb(self):
+    #     config=omegaconf.OmegaConf.to_container(
+    #         self.args, resolve=True, throw_on_missing=True
+    #     )
+    #     self.wandb_run=wandb.init(project=self.args.logging.wandb.project, entity=self.args.logging.wandb.entity, config=config)
+    #     self.wandb_run.name=self.args.tester.wandb.run_name +os.path.basename(self.args.model_dir)+"_"+self.args.exp.exp_name+"_"+self.wandb_run.id
+
+

inference/sampler_euler_heun_multitrack.py

@@ -0,0 +1,302 @@
+import torch
+from inference.sampler import Sampler
+
+
+class SamplerEulerHeun(Sampler):
+
+    def __init__(self, model, diff_params, args):
+        super().__init__(model, diff_params, args)
+
+        # stochasticity parameters
+        self.Schurn = self.args.tester.sampling_params.Schurn
+        self.Snoise = self.args.tester.sampling_params.Snoise
+        self.Stmin = self.args.tester.sampling_params.Stmin
+        self.Stmax = self.args.tester.sampling_params.Stmax
+
+        # order of the sampler
+        self.order = self.args.tester.sampling_params.order
+        self.cond=None
+        self.cfg_scale = 1.0
+
+    def predict_DPS(
+            self,
+            shape,  # observations (lowpssed signal) Tensor with shape ??
+            cond=None,
+            cfg_scale=1.0,
+            device=None,  # device
+            apply_inverse_transform=True,  # whether to apply inverse transform
+            taxonomy=None,  # taxonomy for the conditional input
+            masks=None,  # masks for the conditional input
+            fwd_operator=None,
+            zeta=None,
+            dtype=torch.float32,  # data type
+    ):
+        self.cond = cond
+        assert self.cond is not None, "Conditional input is None"
+
+        self.taxonomy = taxonomy
+        self.masks = masks
+
+        self.cfg_scale = cfg_scale
+
+        # get the noise schedule
+        t = self.create_schedule().to(device).to(torch.float32)
+
+        # sample prior
+        x = self.diff_params.sample_prior(t=t[0], shape=shape, dtype=dtype)
+
+        # parameter for langevin stochasticity, if Schurn is 0, gamma will be 0 to, so the sampler will be deterministic
+        gamma = self.get_gamma(t).to(device)
+
+        for i in range(0, self.T-1, 1):
+            self.step_counter = i
+            x, x_den = self.step_DPS(x, t[i], t[i + 1], gamma[i], fwd_operator, zeta)
+
+
+        if apply_inverse_transform:
+            with torch.no_grad():
+                x_den_wave=self.diff_params.transform_inverse(x_den.detach())
+
+            return x_den_wave.detach(), None
+        else:
+            return x_den.detach(), None
+
+
+    def predict_conditional(
+            self,
+            shape,  # observations (lowpssed signal) Tensor with shape ??
+            cond=None,
+            cfg_scale=1.0,
+            device=None,  # device
+            apply_inverse_transform=True,  # whether to apply inverse transform
+            taxonomy=None,  # taxonomy for the conditional input
+            masks=None,  # masks for the conditional input
+            dtype=torch.float32,  # data type
+    ):
+        self.cond = cond
+        assert self.cond is not None, "Conditional input is None"
+
+        self.taxonomy = taxonomy
+        self.masks = masks
+
+        self.cfg_scale = cfg_scale
+
+        # get the noise schedule
+        t = self.create_schedule().to(device).to(torch.float32)
+
+        # sample prior
+        x = self.diff_params.sample_prior(t=t[0], shape=shape, dtype=dtype)
+
+        # parameter for langevin stochasticity, if Schurn is 0, gamma will be 0 to, so the sampler will be deterministic
+        gamma = self.get_gamma(t).to(device)
+
+        for i in range(0, self.T-1, 1):
+            self.step_counter = i
+            x, x_den = self.step(x, t[i], t[i + 1], gamma[i])
+
+
+        if apply_inverse_transform:
+            with torch.no_grad():
+                x_den_wave=self.diff_params.transform_inverse(x_den.detach())
+
+            return x_den_wave.detach(), None
+        else:
+            return x_den.detach(), None
+
+
+    def predict_unconditional(
+            self,
+            shape,  # observations (lowpssed signal) Tensor with shape ??
+            device
+    ):
+        self.y = None
+        self.degradation = None
+
+        return self.predict(shape, device)
+
+    def get_gamma(self, t):
+        """
+        Get the parameter gamma that defines the stochasticity of the sampler
+        Args
+            t (Tensor): shape: (N_steps, ) Tensor of timesteps, from which we will compute gamma
+        """
+        N = t.shape[0]
+        gamma = torch.zeros(t.shape).to(t.device)
+
+        # If desired, only apply stochasticity between a certain range of noises Stmin is 0 by default and Stmax is a huge number by default. (Unless these parameters are specified, this does nothing)
+        indexes = torch.logical_and(t > self.Stmin, t < self.Stmax)
+
+        # We use Schurn=5 as the default in our experiments
+        gamma[indexes] = gamma[indexes] + torch.min(torch.Tensor([self.Schurn / N, 2 ** (1 / 2) - 1]))
+
+        return gamma
+
+    def get_Tweedie_estimate(self, x, t_i):
+
+        if x.ndim==2:
+            x_=x.unsqueeze(1)
+        elif x.ndim==3:
+            pass
+
+        if self.cond is not None:
+            x_hat = self.diff_params.denoiser(x, self.model, t_i, cond=self.cond, cfg_scale=self.cfg_scale, taxonomy=self.taxonomy, masks=self.masks)
+        else:
+            x_hat = self.diff_params.denoiser(x, self.model, t_i, taxonomy=self.taxonomy, masks=self.masks)
+
+        return x_hat
+
+    def Tweedie2score(self, tweedie, xt, t):
+        return self.diff_params.Tweedie2score(tweedie, xt, t)
+
+    def score2Tweedie(self, score, xt, t):
+        return self.diff_params.score2Tweedie(score, xt, t)
+
+    def stochastic_timestep(self, x, t, gamma, Snoise=1):
+        t_hat = t + gamma * t  # if gamma_sig[i]==0 this is a deterministic step, make sure it doed not crash
+        t_hat=torch.clamp(t_hat, 0, self.diff_params.max_t)
+        epsilon = torch.randn(x.shape).to(x.device) * Snoise  # sample Gaussiannoise, Snoise is 1 by default
+        if t_hat <= t:
+            x_hat = x
+            #print(f"t_hat<=t, gamma {gamma}")
+        else:
+            #print(t_hat, t)
+            x_hat = x + ((t_hat ** 2 - t ** 2) ** (1 / 2)) * epsilon  # Perturb data
+        return x_hat, t_hat
+
+    def step_DPS(self, x_i, t_i, t_iplus1, gamma_i , fwd_operator=None, zeta=None):
+
+            #with torch.no_grad():
+
+            x_hat, t_hat=self.stochastic_timestep(x_i, t_i, gamma_i)
+
+            x_hat.requires_grad=True
+
+            x_den = self.get_Tweedie_estimate(x_hat, t_hat)
+
+            #optionally L2 normalize here...
+
+            #compute likelihood score
+            loss=fwd_operator(x_den)
+
+            loss.backward(retain_graph=False)
+
+            grads= x_hat.grad
+
+            #lets normalize the grads
+            norm_factor= torch.sqrt(torch.tensor(x_hat.view(-1).shape[0])).to(x_hat.device)
+            normguide=torch.norm(grads)/ norm_factor
+            zeta= zeta/(normguide+1e-8)
+            lh_score=-zeta*grads/t_hat
+
+            x_hat.detach_()  # detach x_hat to avoid accumulating gradients
+            x_den.detach_()  # detach x_den to avoid accumulating gradients
+
+            #compute normal score
+            score = self.Tweedie2score(x_den, x_hat, t_hat)
+
+            ode_integrand = self.diff_params._ode_integrand(x_hat, t_hat, score+ lh_score)
+
+            dt = t_iplus1 - t_hat
+
+            x_iplus1 = x_hat + dt * ode_integrand
+
+            return x_iplus1, x_den
+
+    def step(self, x_i, t_i, t_iplus1, gamma_i ):
+
+        with torch.no_grad():
+            x_hat, t_hat = self.stochastic_timestep(x_i, t_i, gamma_i)
+            x_den = self.get_Tweedie_estimate(x_hat, t_hat)
+            score = self.Tweedie2score(x_den, x_hat, t_hat)
+            ode_integrand = self.diff_params._ode_integrand(x_hat, t_hat, score)
+            dt = t_iplus1 - t_hat
+
+            if t_iplus1 != 0 and self.order == 2:  # second order correction
+                t_prime = t_iplus1
+                x_prime = x_hat + dt * ode_integrand
+                x_den = self.get_Tweedie_estimate(x_prime, t_prime)
+                score = self.Tweedie2score(x_den, x_prime, t_prime)
+                ode_integrand_next = self.diff_params._ode_integrand(x_prime, t_prime, score)
+                ode_integrand_midpoint = .5 * (ode_integrand + ode_integrand_next)
+                x_iplus1 = x_hat + dt * ode_integrand_midpoint
+
+            else:
+                x_iplus1 = x_hat + dt * ode_integrand
+
+            return x_iplus1, x_den
+    
+    def get_domain_shape(self, shape, device):
+
+        x=torch.zeros(shape, dtype=torch.float32).to(device)
+        X=self.diff_params.transform_forward(x)
+
+        return X.shape, X.dtype
+
+    def predict(
+            self,
+            shape,  # observations (lowpssed signal) Tensor with shape ??
+            device,  # lambda function
+            dtype=torch.float32,  # data type
+            apply_inverse_transform=True  # whether to apply inverse transform
+    ):
+
+        # get the noise schedule
+        t = self.create_schedule().to(device).to(torch.float32)
+
+        # sample prior
+        x = self.diff_params.sample_prior(t=t[0], shape=shape, dtype=dtype)
+
+        # parameter for langevin stochasticity, if Schurn is 0, gamma will be 0 to, so the sampler will be deterministic
+        gamma = self.get_gamma(t).to(device)
+
+        for i in range(0, self.T-1, 1):
+            self.step_counter = i
+            x, x_den = self.step(x, t[i], t[i + 1], gamma[i])
+
+
+        if apply_inverse_transform:
+            with torch.no_grad():
+                x_den_wave=self.diff_params.transform_inverse(x_den.detach())
+
+            return x_den_wave.detach(), None
+        else:
+            return x_den.detach(), None
+
+    def create_schedule(self, sigma_min=None, sigma_max=None, rho=None, T=None):
+        """
+        EDM schedule by default
+        """
+        if T is None:
+            T=self.T
+
+        if self.args.tester.sampling_params.schedule == "edm":
+            if sigma_min is None:
+                sigma_min = self.sde_hp.sigma_min
+            if sigma_max is None:
+                sigma_max = self.sde_hp.sigma_max
+            if rho is None:
+                rho = self.sde_hp.rho
+            a = torch.arange(0, T)
+            t = (sigma_max**(1/rho) + a/(T-1) *(sigma_min**(1/rho) - sigma_max**(1/rho)))**rho
+            t[-1] = 0
+            return t
+        
+        elif self.args.tester.sampling_params.schedule == "song":
+            if sigma_min is None:
+                sigma_min = self.sde_hp.sigma_min
+            if sigma_max is None:
+                sigma_max = self.sde_hp.sigma_max
+            if rho is None:
+                rho = self.sde_hp.rho
+            eps = 0. if not "t_eps" in self.args.tester.diff_params.keys() else self.args.tester.diff_params.t_eps
+            a = torch.arange(eps, T+1)
+            t = sigma_min**2 * (sigma_max / sigma_min)**(2*a)
+            t[-1] = 0
+            return t    
+        elif self.args.tester.sampling_params.schedule == "FM":
+            t = torch.linspace(1, 0, T+1)
+            return t
+
+        else:
+            raise NotImplementedError(f"schedule {self.args.tester.sampling_params.schedule} not implemented")
+

inference/validator_FxProcessor.py

@@ -0,0 +1,436 @@
+# Copyright (c) 2025 Sony Research
+# Licensed under CC BY-NC-SA 4.0
+# See LICENSE file for details
+
+from datetime import date
+import math
+import io
+import matplotlib.pyplot as plt
+from functools import partial
+import re
+import torch
+import os
+import wandb
+import copy
+from glob import glob
+from tqdm import tqdm
+import omegaconf
+import hydra
+import utils.log as utils_logging
+import utils.training_utils as tr_utils
+
+import soundfile as sf
+import numpy as np
+import torchaudio
+
+from utils.data_utils import apply_RMS_normalization
+
+
+class ValidatorFxProcessor:
+    def __init__(
+        self,
+        args,
+        network,
+        test_set_dict=None,
+        device=None,
+        in_training=False,
+    ):
+        self.args = args
+        self.network = network
+        self.device = device
+        self.test_set_dict = test_set_dict
+
+        self.use_wandb = False  # hardcoded for now
+        self.in_training = in_training
+
+        if in_training:
+            self.use_wandb = True
+            # Will inherit wandb_run from Trainer
+        else:  # If we use the tester in training, we will log in WandB in the Trainer() class, no need to create all these paths
+            torch.backends.cudnn.benchmark = True
+            if self.device is None:
+                self.device = torch.device(
+                    "cuda" if torch.cuda.is_available() else "cpu"
+                )
+
+            self.setup_wandb()
+
+        if self.args.tester.compute_metrics:
+            self.metrics_dict = self.prepare_metrics(self.args.tester.metrics)
+        else:
+            self.metrics_dict = {}
+
+
+        self.RMS_norm = (
+            self.args.exp.RMS_norm
+        )  # Use fixed RMS for evaluation, hardcoded for now
+
+        if self.args.exp.style_encoder_type == "FxEncoder++_DynamicFeatures":
+
+            Fxencoder_kwargs = self.args.exp.fx_encoder_plusplus_args
+
+            from utils.feature_extractors.load_features import (
+                load_fx_encoder_plusplus_2048,
+            )
+
+            feat_extractor = load_fx_encoder_plusplus_2048(
+                Fxencoder_kwargs, self.device
+            )
+
+            from utils.feature_extractors.AF_features_embedding import (
+                AF_fourier_embedding,
+            )
+
+            AFembedding = AF_fourier_embedding(device=self.device)
+
+            def fxencode_fn(x):
+                """
+                x: tensor of shape [B, C, L] where B is the batch size, C is the number of channels and L is the length of the audio
+                """
+                z = feat_extractor(x)
+                z = torch.nn.functional.normalize(
+                    z, dim=-1, p=2
+                )  # normalize to unit variance
+                z = z * math.sqrt(z.shape[-1])  # rescale to keep the same scale
+
+                z_af, _ = AFembedding.encode(x)
+                # embedding is l2 normalized, normalize to unit variance
+                z_af = z_af * math.sqrt(
+                    z_af.shape[-1]
+                )  # rescale to keep the same scale
+
+                # concatenate z and z_af (rescaling with sqrt(dim) to keep the same scale)
+                z_all = torch.cat([z, z_af], dim=-1)
+
+                norm_z = z_all / math.sqrt(
+                    z_all.shape[-1]
+                )  # normalize by dividing by sqrt(dim) to keep the same scale
+
+                return norm_z
+
+            self.style_encode = fxencode_fn
+        else:
+            raise NotImplementedError(
+                "Only FxEncoder++_DynamicFeatures is implemented for now"
+            )
+
+        if self.args.exp.apply_fxnorm:
+            self.fx_normalizer = hydra.utils.instantiate(self.args.exp.fxnorm)
+
+    def setup_wandb(self):
+        """
+        Configure wandb, open a new run and log the configuration.
+        """
+        config = omegaconf.OmegaConf.to_container(
+            self.args, resolve=True, throw_on_missing=True
+        )
+        self.wandb_run = wandb.init(
+            project="testing" + self.args.tester.wandb.project,
+            entity=self.args.tester.wandb.entity,
+            config=config,
+            tags=self.args.tester.wandb.tags,
+        )
+        # wandb.watch(self.network,
+        #            log_freq=self.args.logging.heavy_log_interval)
+
+        self.wandb_run.name = self.args.tester.wandb.run_name
+        self.use_wandb = True
+
+    def setup_wandb_run(self, run):
+        # get the wandb run object from outside (in trainer.py or somewhere else)
+        self.wandb_run = run
+        self.use_wandb = True
+
+    def load_latest_checkpoint(self):
+        # load the latest checkpoint from self.args.model_dir
+        try:
+            # find latest checkpoint_id
+            save_basename = f"{self.args.exp.exp_name}-*.pt"
+            save_name = f"{self.args.model_dir}/{save_basename}"
+            list_weights = glob(save_name)
+            id_regex = re.compile(f"{self.args.exp.exp_name}-(\d*)\.pt")
+            list_ids = [
+                int(id_regex.search(weight_path).groups()[0])
+                for weight_path in list_weights
+            ]
+            checkpoint_id = max(list_ids)
+
+            state_dict = torch.load(
+                f"{self.args.model_dir}/{self.args.exp.exp_name}-{checkpoint_id}.pt",
+                map_location=self.device,
+            )
+            try:
+                self.network.load_state_dict(state_dict["network"])
+            except Exception as e:
+                print(e)
+                print("Failed to load in strict mode, trying again without strict mode")
+                self.network.load_state_dict(state_dict["model"], strict=False)
+
+            print(f"Loaded checkpoint {checkpoint_id}")
+            return True
+        except (FileNotFoundError, ValueError):
+            raise ValueError("No checkpoint found")
+
+    def load_checkpoint(self, path):
+        state_dict = torch.load(path, map_location=self.device, weights_only=False)
+        print("state_dict keys:", state_dict.keys())
+        try:
+            self.it = state_dict["it"]
+        except:
+            self.it = 0
+
+        print(f"loading checkpoint {self.it}")
+        return tr_utils.load_state_dict(state_dict, network=self.network)
+
+    def log_figure(self, fig, name: str, step=None):
+        # Save the figure to a buffer
+
+        self.wandb_run.log(
+            {name: wandb.Image(fig)}, step=step if step is not None else self.it
+        )
+
+    def log_metric(self, value, name: str, step=None):
+        # print("logging metric it:", self.it, "name:", name)
+        self.wandb_run.log({name: value}, step=step if step is not None else self.it)
+
+    def log_audio(self, pred, name: str, it=None):
+        if it is None:
+            it = self.it
+        if self.use_wandb:
+            pred = pred.permute(1, 0)
+            self.wandb_run.log(
+                {
+                    name: wandb.Audio(
+                        pred.detach().cpu().numpy(),
+                        sample_rate=self.args.exp.sample_rate,
+                    )
+                },
+                step=it,
+            )
+
+    # ------------- UNCONDITIONAL SAMPLING ---------------#
+
+    ##############################
+    ### UNCONDITIONAL SAMPLING ###
+    ##############################
+
+    def prepare_metrics(self, metrics):
+        metrics_dict = {}
+        for metric in metrics:
+                print(f"Preparing metric {metric}")
+                if "pairwise" in metric:
+                    from utils.evaluation.pairwise_metrics import metric_factory
+
+                    metrics_dict[metric] = metric_factory(
+                        metric, self.args.exp.sample_rate, **self.args.tester
+                    )
+
+        return metrics_dict
+
+    def test_paired(self, mode, exp_description=""):
+
+        # self.it = 0
+        for k, test_set in self.test_set_dict.items():
+
+            print(f"Testing on {k} set", k)
+
+            assert len(test_set) != 0, "No samples found in test set"
+
+            dict_y = {}
+            dict_x = {}
+            dict_y_hat = {}
+
+            i = 0
+
+            for x, y in tqdm(test_set):
+
+                B, C, T = y.shape
+
+                x = x.to(self.device).float()
+                if x.shape[-1] > self.args.exp.audio_len:
+                    x = x[:, :, : self.args.exp.audio_len]
+                elif x.shape[-1] < self.args.exp.audio_len:
+                    raise ValueError(
+                        f"Sample length {x.shape[-1]} is less than expected {self.args.exp.audio_len}"
+                    )
+
+                if mode == "paired":
+                    y = y.to(self.device).float()
+
+                    if y.shape[-1] > self.args.exp.audio_len:
+                        y = y[:, :, : self.args.exp.audio_len]
+                    elif y.shape[-1] < self.args.exp.audio_len:
+                        raise ValueError(
+                            f"Sample length {y.shape[-1]} is less than expected {self.args.exp.audio_len}"
+                        )
+
+                if x.shape[1] == 2:
+                    x = x.mean(
+                        dim=1, keepdim=True
+                    )  # expand to [B*N, 1, L] to keep the shape consistent
+
+                # RMS normalization of x and y
+
+                if self.args.exp.apply_fxnorm:
+                    x = self.fx_normalizer(
+                        x,
+                        RMS=self.args.exp.RMS_norm,
+                        use_gate=self.args.exp.use_gated_RMSnorm,
+                    )
+                else:
+                    x = apply_RMS_normalization(
+                        x,
+                        self.args.exp.RMS_norm,
+                        use_gate=self.args.exp.use_gated_RMSnorm,
+                    )
+
+                if "baseline" in mode:
+                    if mode == "baseline_dry":
+                        preds = x  # Just return the dry vocals as baseline
+                    elif mode == "baseline_autoencoder":
+                        preds = self.autoencoder_reconstruction(
+                            y
+                        )  # Just return the dry vocals as baseline
+                    elif mode == "baseline_random":
+                        raise NotImplementedError(
+                            "Baseline random sampling not implemented yet"
+                        )
+                        pass
+                else:
+                    with torch.no_grad():
+                        z = self.style_encode(y)
+                        try:
+                            preds = self.network(
+                                x, z
+                            )  # Get the predictions from the network
+                        except Exception as e:
+                            print(f"Error during inference: {e}")
+                            continue
+                        print(
+                            "y_pred",
+                            preds.shape,
+                            preds.std(),
+                            preds.mean(),
+                            preds.min(),
+                            preds.max(),
+                        )
+
+                    is_nan = torch.isnan(preds).any()
+                    if is_nan:
+                        num_nan = torch.sum(torch.isnan(x)).item()
+                        print(
+                            f"Number of NaN values in sample_x: {num_nan} of {x.numel()}"
+                        )
+
+                y = apply_RMS_normalization(
+                    y, self.args.exp.RMS_norm, use_gate=self.args.exp.use_gated_RMSnorm
+                )
+
+                for b in range(B):
+                    if self.use_wandb:
+                        if (
+                            i < self.args.tester.wandb.num_examples_to_log
+                        ):  # Log only first 10 samples
+                            self.log_audio(
+                                preds[b], f"pred_wet_{k}_{mode}_{i}", it=self.it
+                            )  # Just log first sample
+                            self.log_audio(
+                                y[b], f"original_wet_{k}_{mode}_{i}", it=self.it
+                            )  # Just log first sample
+                            self.log_audio(
+                                x[b], f"original_dry_{k}_{mode}_{i}", it=self.it
+                            )  # Just log first sample
+
+                    dict_y[i] = y[b].detach().cpu().numpy()
+                    dict_x[i] = x[b].detach().cpu().numpy()
+                    dict_y_hat[i] = preds[b].detach().cpu().numpy()
+
+                    i += 1
+
+            if self.args.tester.compute_metrics:
+                for metric in self.metrics_dict.keys():
+                    try:
+                        print(f"Computing metric {metric}")
+                        result, result_dict = self.metrics_dict[metric].compute(
+                            dict_y, dict_y_hat, dict_x
+                        )
+
+                        if self.use_wandb:
+                            if result is not None:
+                                self.log_metric(
+                                    result, metric + "_" + k + "_" + mode, step=self.it
+                                )
+
+                            for key, value in result_dict.items():
+                                if "figure" in key:
+                                    # log figure as an image
+                                    self.log_figure(
+                                        value, key + "_" + k + "_" + mode, step=self.it
+                                    )
+                                else:
+                                    self.log_metric(
+                                        value, key + "_" + k + "_" + mode, step=self.it
+                                    )
+
+                    except Exception as e:
+                        print(f"Error computing metric {metric}: {e}")
+                        continue
+
+    def prepare_directories(self, mode, unconditional=False, string=None):
+
+        today = date.today()
+        self.paths = {}
+        if (
+            "overriden_name" in self.args.tester.keys()
+            and self.args.tester.overriden_name is not None
+        ):
+            self.path_sampling = os.path.join(
+                self.args.model_dir, self.args.tester.overriden_name
+            )
+        else:
+            self.path_sampling = os.path.join(
+                self.args.model_dir, "test" + today.strftime("%d_%m_%Y")
+            )
+        if not os.path.exists(self.path_sampling):
+            os.makedirs(self.path_sampling)
+
+        self.paths[mode] = os.path.join(
+            self.path_sampling, mode, self.args.exp.exp_name
+        )
+
+        if not os.path.exists(self.paths[mode]):
+            os.makedirs(self.paths[mode])
+        if string is None:
+            string = ""
+
+        self.paths[mode + "wet_original"] = os.path.join(
+            self.paths[mode], string + "wet_original"
+        )
+        if not os.path.exists(self.paths[mode + "wet_original"]):
+            os.makedirs(self.paths[mode + "wet_original"])
+        self.paths[mode + "dry"] = os.path.join(self.paths[mode], string + "dry")
+        if not os.path.exists(self.paths[mode + "dry"]):
+            os.makedirs(self.paths[mode + "dry"])
+        self.paths[mode + "wet_estimate"] = os.path.join(
+            self.paths[mode], string + "wet_estimate"
+        )
+        if not os.path.exists(self.paths[mode + "wet_estimate"]):
+            os.makedirs(self.paths[mode + "wet_estimate"])
+
+    def save_experiment_args(self, mode):
+        with open(
+            os.path.join(self.paths[mode], ".argv"), "w"
+        ) as f:  # Keep track of the arguments we used for this experiment
+            omegaconf.OmegaConf.save(config=self.args, f=f.name)
+
+    def do_test(self, it=0):
+
+        self.it = it
+        for m in self.args.tester.modes:
+            if m == "paired":
+                if not self.in_training:
+                    self.prepare_directories(m, unconditional=True)
+                    self.save_experiment_args(m)
+                self.test_paired(m)
+            else:
+                print("Warning: unknown mode: ", m)

networks/MLP_CLAP_regressor.py

@@ -0,0 +1,23 @@
+import torch
+import torch.nn as nn
+
+class MLP_CLAP_regressor(nn.Module):
+    """
+    A simple MLP regressor that uses CLAP features as input.
+    """
+
+    def __init__(self, dim=512, hidden_dim=512):
+        super(MLP_CLAP_regressor, self).__init__()
+
+        self.model = nn.Sequential(
+            nn.Linear(dim, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, dim)
+        )
+
+    def forward(self, x):
+
+        emb= self.model(x)
+        #l2 normalization
+        return nn.functional.normalize(emb, p=2, dim=-1)
+    

networks/blackbox_TCN.py

@@ -0,0 +1,346 @@
+""" 
+    Adapted from: https://github.com/SonyResearch/ITO-Master
+"""
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.nn.init as init
+
+import os
+import sys
+
+import omegaconf
+import torchaudio
+from utils.feature_extractors.load_features import load_CLAP
+
+# 1-dimensional convolutional layer
+# in the order of conv -> norm -> activation
+class Conv1d_layer(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, \
+                                    stride=1, \
+                                    padding="SAME", dilation=1, bias=True, \
+                                    norm="batch", activation="relu", \
+                                    mode="conv"):
+        super(Conv1d_layer, self).__init__()
+        
+        self.conv1d = nn.Sequential()
+
+        ''' padding '''
+        if mode=="deconv":
+            padding = int(dilation * (kernel_size-1) / 2)
+            out_padding = 0 if stride==1 else 1
+        elif mode=="conv" or "alias_free" in mode:
+            if padding == "SAME":
+                pad = int((kernel_size-1) * dilation)
+                l_pad = int(pad//2)
+                r_pad = pad - l_pad
+                padding_area = (l_pad, r_pad)
+            elif padding == "VALID":
+                padding_area = (0, 0)
+            else:
+                pass
+
+        ''' convolutional layer '''
+        if mode=="deconv":
+            self.conv1d.add_module("deconv1d", nn.ConvTranspose1d(in_channels, out_channels, kernel_size, \
+                                                            stride=stride, padding=padding, output_padding=out_padding, \
+                                                            dilation=dilation, \
+                                                            bias=bias))
+        elif mode=="conv":
+            self.conv1d.add_module(f"{mode}1d_pad", nn.ReflectionPad1d(padding_area))
+            self.conv1d.add_module(f"{mode}1d", nn.Conv1d(in_channels, out_channels, kernel_size, \
+                                                            stride=stride, padding=0, \
+                                                            dilation=dilation, \
+                                                            bias=bias))
+        elif "alias_free" in mode:
+            if "up" in mode:
+                up_factor = stride * 2
+                down_factor = 2
+            elif "down" in mode:
+                up_factor = 2
+                down_factor = stride * 2
+            else:
+                raise ValueError("choose alias-free method : 'up' or 'down'")
+            # procedure : conv -> upsample -> lrelu -> low-pass filter -> downsample
+            # the torchaudio.transforms.Resample's default resampling_method is 'sinc_interpolation' which performs low-pass filter during the process
+            # details at https://pytorch.org/audio/stable/transforms.html
+            self.conv1d.add_module(f"{mode}1d_pad", nn.ReflectionPad1d(padding_area))
+            self.conv1d.add_module(f"{mode}1d", nn.Conv1d(in_channels, out_channels, kernel_size, \
+                                                            stride=1, padding=0, \
+                                                            dilation=dilation, \
+                                                            bias=bias))
+            self.conv1d.add_module(f"{mode}upsample", torchaudio.transforms.Resample(orig_freq=1, new_freq=up_factor))
+            self.conv1d.add_module(f"{mode}lrelu", nn.LeakyReLU())
+            self.conv1d.add_module(f"{mode}downsample", torchaudio.transforms.Resample(orig_freq=down_factor, new_freq=1))
+
+        ''' normalization '''
+        if norm=="batch":
+            self.conv1d.add_module("batch_norm", nn.BatchNorm1d(out_channels))
+            # self.conv1d.add_module("batch_norm", nn.SyncBatchNorm(out_channels))
+
+        ''' activation '''
+        if 'alias_free' not in mode:
+            if activation=="relu":
+                self.conv1d.add_module("relu", nn.ReLU())
+            elif activation=="lrelu":
+                self.conv1d.add_module("lrelu", nn.LeakyReLU())
+
+
+    def forward(self, input):
+        # input shape should be : batch x channel x height x width
+        output = self.conv1d(input)
+        return output
+
+
+# compute receptive field
+def compute_receptive_field(kernels, strides, dilations):
+    rf = 0
+    for i in range(len(kernels)):
+        rf += rf * strides[i] + (kernels[i]-strides[i]) * dilations[i]
+    return rf
+
+# Feature-wise Linear Modulation
+class FiLM(nn.Module):
+    def __init__(self, condition_len=2048, feature_len=1024):
+        super(FiLM, self).__init__()
+        self.film_fc = nn.Linear(condition_len, feature_len*2)
+        self.feat_len = feature_len
+
+    def forward(self, feature, condition):
+        film_factor = self.film_fc(condition).unsqueeze(-1)
+        r, b = torch.split(film_factor, self.feat_len, dim=1)
+        return r*feature + b
+
+
+class ConvBlock(nn.Module):
+    def __init__(self, dimension, layer_num, \
+                        in_channels, out_channels, \
+                        kernel_size, \
+                        stride=1, padding="SAME", \
+                        dilation=1, \
+                        bias=True, \
+                        norm="batch", \
+                        activation="relu", last_activation="relu", \
+                        mode="conv"):
+        super(ConvBlock, self).__init__()
+
+        conv_block = []
+        if dimension==1:
+            for i in range(layer_num-1):
+                conv_block.append(Conv1d_layer(in_channels, in_channels, kernel_size, padding=padding, dilation=dilation, bias=bias, norm=norm, activation=activation))
+            conv_block.append(Conv1d_layer(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, bias=bias, norm=norm, activation=last_activation, mode=mode))
+        elif dimension==2:
+            for i in range(layer_num-1):
+                conv_block.append(Conv2d_layer(in_channels, in_channels, kernel_size, padding=padding, dilation=dilation, bias=bias, norm=norm, activation=activation))
+            conv_block.append(Conv2d_layer(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, bias=bias, norm=norm, activation=last_activation, mode=mode))
+        self.conv_block = nn.Sequential(*conv_block)
+
+    def forward(self, input):
+        return self.conv_block(input)
+
+
+class TCNBlock(torch.nn.Module):
+    def __init__(self, 
+                in_ch, 
+                out_ch, 
+                kernel_size=3, 
+                stride=1, 
+                dilation=1, 
+                cond_dim=2048, 
+                grouped=False, 
+                causal=False,
+                conditional=False, 
+                **kwargs):
+        super(TCNBlock, self).__init__()
+
+        self.in_ch = in_ch
+        self.out_ch = out_ch
+        self.kernel_size = kernel_size
+        self.dilation = dilation
+        self.grouped = grouped
+        self.causal = causal
+        self.conditional = conditional
+
+        groups = out_ch if grouped and (in_ch % out_ch == 0) else 1
+
+        self.pad_length = ((kernel_size-1)*dilation) if self.causal else ((kernel_size-1)*dilation)//2
+        self.conv1 = torch.nn.Conv1d(in_ch, 
+                                     out_ch, 
+                                     kernel_size=kernel_size, 
+                                     stride=stride,
+                                     padding=self.pad_length,
+                                     dilation=dilation,
+                                     groups=groups,
+                                     bias=False)
+        if grouped:
+            self.conv1b = torch.nn.Conv1d(out_ch, out_ch, kernel_size=1)
+
+        if conditional:
+            self.film = FiLM(cond_dim, out_ch)
+        self.bn = torch.nn.BatchNorm1d(out_ch)
+
+        self.relu = torch.nn.LeakyReLU()
+
+        if out_ch % in_ch == 0:
+            self.res = torch.nn.Conv1d(in_ch, 
+                                   out_ch, 
+                                   kernel_size=1,
+                                   stride=stride,
+                                   groups=in_ch,
+                                   bias=False)
+        else:
+            self.res = torch.nn.Conv1d(in_ch, 
+                                   out_ch, 
+                                   kernel_size=1,
+                                   stride=stride,
+                                   groups=1,
+                                   bias=False)
+
+    def forward(self, x, p):
+        x_in = x
+
+        x = self.relu(self.bn(self.conv1(x)))
+        #print("p", p.shape)
+        x = self.film(x, p)
+
+        x_res = self.res(x_in)
+
+        if self.causal:
+            x = x[..., :-self.pad_length]
+        x += x_res
+
+        return x
+
+class FourierFeatures(nn.Module):
+    def __init__(self, in_features, out_features, std=1.):
+        super().__init__()
+        assert out_features % 2 == 0
+        self.weight = nn.Parameter(torch.randn(
+            [out_features // 2, in_features]) * std)
+
+    def forward(self, input):
+        f = 2 * math.pi * input @ self.weight.T
+        return torch.cat([f.cos(), f.sin()], dim=-1)
+
+class TCNModel(nn.Module):
+    """ Temporal convolutional network with conditioning module.
+        Args:
+            nparams (int): Number of conditioning parameters.
+            ninputs (int): Number of input channels (mono = 1, stereo 2). Default: 1
+            noutputs (int): Number of output channels (mono = 1, stereo 2). Default: 1
+            nblocks (int): Number of total TCN blocks. Default: 10
+            kernel_size (int): Width of the convolutional kernels. Default: 3
+            dialation_growth (int): Compute the dilation factor at each block as dilation_growth ** (n % stack_size). Default: 1
+            channel_growth (int): Compute the output channels at each black as in_ch * channel_growth. Default: 2
+            channel_width (int): When channel_growth = 1 all blocks use convolutions with this many channels. Default: 64
+            stack_size (int): Number of blocks that constitute a single stack of blocks. Default: 10
+            grouped (bool): Use grouped convolutions to reduce the total number of parameters. Default: False
+            causal (bool): Causal TCN configuration does not consider future input values. Default: False
+            skip_connections (bool): Skip connections from each block to the output. Default: False
+        """
+    def __init__(self, 
+                 ninputs=1,
+                 noutputs=2,
+                 nblocks=14, 
+                 kernel_size=15, 
+                 stride=1,
+                 dilation_growth=2, 
+                 channel_growth=1, 
+                 channel_width=128, 
+                 stack_size=15,
+                 cond_dim=2048,
+                 grouped=False,
+                 causal=False,
+                 skip_connections=False,
+                 use_CLAP=False,
+                 CLAP_args=None,
+                 ):
+        super(TCNModel, self).__init__()
+
+        self.use_CLAP = use_CLAP
+        if self.use_CLAP:
+            assert CLAP_args is not None, "CLAP_args must be provided for CLAP AE"
+            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            CLAP_encoder= load_CLAP(CLAP_args, device=device)
+            cond_dim+= 512
+            def merge_CLAP_embeddings(x, emb):
+
+                clap_embedding = CLAP_encoder(x, type="dry")
+                #l2 normalize the clap embedding
+                clap_embedding = F.normalize(clap_embedding, p=2, dim=-1)
+
+                return torch.cat((emb, clap_embedding), dim=-1)
+            
+            self.merge_CLAP_embeddings = merge_CLAP_embeddings
+
+
+        self.hparams = {
+            "ninputs": ninputs,
+            "noutputs": noutputs,
+            "nblocks": nblocks,
+            "kernel_size": kernel_size,
+            "stride": stride,
+            "dilation_growth": dilation_growth,
+            "channel_growth": channel_growth,
+            "channel_width": channel_width,
+            "stack_size": stack_size,
+            "cond_dim": cond_dim,
+            "grouped": grouped,
+            "causal": causal,
+            "skip_connections": skip_connections,
+        }
+        self.hparams= omegaconf.OmegaConf.create(self.hparams)
+
+        self.blocks = torch.nn.ModuleList()
+        for n in range(nblocks):
+            in_ch = out_ch if n > 0 else ninputs
+            
+            if self.hparams.channel_growth > 1:
+                out_ch = in_ch * self.hparams.channel_growth 
+            else:
+                out_ch = self.hparams.channel_width
+
+            dilation = self.hparams.dilation_growth ** (n % self.hparams.stack_size)
+            cur_stride = stride[n] if isinstance(stride, list) else stride
+            self.blocks.append(TCNBlock(in_ch, 
+                                        out_ch, 
+                                        kernel_size=self.hparams.kernel_size, 
+                                        stride=cur_stride, 
+                                        dilation=dilation,
+                                        padding="same" if self.hparams.causal else "valid",
+                                        causal=self.hparams.causal,
+                                        cond_dim=cond_dim, 
+                                        grouped=self.hparams.grouped,
+                                        conditional=True ))
+
+        self.output = torch.nn.Conv1d(out_ch, noutputs, kernel_size=1)
+
+    def forward(self, x, cond):
+        # iterate over blocks passing conditioning
+        if self.use_CLAP:
+            with torch.no_grad():
+                cond = self.merge_CLAP_embeddings(x, cond)
+
+        for idx, block in enumerate(self.blocks):
+            # for SeFa
+            if isinstance(cond, list):
+                x = block(x, cond[idx])
+            else:
+                x = block(x, cond)
+            skips = 0
+
+        # out = torch.tanh(self.output(x + skips))
+        out = torch.clamp(self.output(x + skips), min=-1, max=1)
+
+        return out
+
+    def compute_receptive_field(self):
+        """ Compute the receptive field in samples."""
+        rf = self.hparams.kernel_size
+        for n in range(1,self.hparams.nblocks):
+            dilation = self.hparams.dilation_growth ** (n % self.hparams.stack_size)
+            rf = rf + ((self.hparams.kernel_size-1) * dilation)
+        return rf
+
+

networks/dit_multitrack.py

@@ -0,0 +1,488 @@
+
+# adapted from  https://github.com/Stability-AI/stable-audio-tools/blob/main/stable_audio_tools/models/dit.py
+
+import typing as tp
+import math
+import torch
+
+from einops import rearrange
+from torch import nn
+from torch.nn import functional as F
+
+class FourierFeatures(nn.Module):
+    def __init__(self, in_features, out_features, std=1.):
+        super().__init__()
+        assert out_features % 2 == 0
+        self.weight = nn.Parameter(torch.randn(
+            [out_features // 2, in_features]) * std)
+
+    def forward(self, input):
+        f = 2 * math.pi * input @ self.weight.T
+        return torch.cat([f.cos(), f.sin()], dim=-1)
+
+class OneHotPositionalEmbedding(nn.Module):
+    def __init__(self, max_seq_len):
+        super().__init__()
+        self.max_seq_len = max_seq_len
+
+    def forward(self, x, pos=None, seq_start_pos=None):
+        seq_len, device = x.shape[1], x.device
+        assert seq_len <= self.max_seq_len, f'you are passing in a sequence length of {seq_len} but your one-hot positional embedding has a max sequence length of {self.max_seq_len}'
+
+        if pos is None:
+            pos = torch.arange(seq_len, device=device)
+
+        if seq_start_pos is not None:
+            pos = (pos - seq_start_pos[..., None]).clamp(min=0)
+
+        pos_emb = F.one_hot(pos, num_classes=self.max_seq_len).to(x.dtype)
+        return pos_emb
+
+
+
+class DiffusionTransformer(nn.Module):
+    def __init__(self, 
+        io_channels=32, 
+        patch_size=1,
+        embed_dim=768,
+        cond_token_dim=0,
+        cond_token_proj_dim=64,
+        project_cond_tokens=False,
+        global_cond_dim=0,
+        project_global_cond=True,
+        input_concat_dim=0,
+        prepend_cond_dim=0,
+        depth=12,
+        num_heads=8,
+        transformer_type: tp.Literal["x-transformers", "continuous_transformer"] = "x-transformers",
+        global_cond_type: tp.Literal["prepend", "adaLN"] = "prepend",
+        timestep_cond_type: tp.Literal["global", "input_concat"] = "global",
+        timestep_embed_dim=None,
+        pos_emb_strategy="concatenation",
+        pos_emb_dim=None,
+        pos_emb_type="one-hot",
+        pos_emb_crossattn_strategy="concatenation",
+        pos_emb_crossattn_dim=None,
+        pos_emb_crossattn_type="one-hot",
+        use_taxonomy_in_pos_emb=True,
+        max_num_tracks=14,#used for one-hot positional embeddings
+        **kwargs):
+
+        super().__init__()
+        
+        self.cond_token_dim = cond_token_dim
+
+        # Timestep embeddings
+        self.timestep_cond_type = timestep_cond_type
+
+        timestep_features_dim = 256
+
+        self.timestep_features = FourierFeatures(1, timestep_features_dim)
+
+        if timestep_cond_type == "global":
+            timestep_embed_dim = embed_dim
+        elif timestep_cond_type == "input_concat":
+            assert timestep_embed_dim is not None, "timestep_embed_dim must be specified if timestep_cond_type is input_concat"
+            input_concat_dim += timestep_embed_dim
+
+        self.to_timestep_embed = nn.Sequential(
+            nn.Linear(timestep_features_dim, timestep_embed_dim, bias=True),
+            nn.SiLU(),
+            nn.Linear(timestep_embed_dim, timestep_embed_dim, bias=True),
+        )
+
+
+        self.project_cond_tokens = project_cond_tokens
+        if cond_token_dim > 0:
+            # Conditioning tokens
+            
+            if self.project_cond_tokens:
+                self.to_cond_embed = nn.Sequential(
+                    nn.Linear(cond_token_dim, cond_token_proj_dim, bias=False),
+                    nn.SiLU(),
+                    nn.Linear(cond_token_proj_dim, cond_token_proj_dim, bias=False)
+                )
+                cond_embed_dim = cond_token_proj_dim
+            else:
+                cond_embed_dim = cond_token_dim 
+        else:
+            cond_embed_dim = 0
+
+        if global_cond_dim > 0:
+            # Global conditioning
+            global_embed_dim = global_cond_dim if not project_global_cond else embed_dim
+            self.to_global_embed = nn.Sequential(
+                nn.Linear(global_cond_dim, global_embed_dim, bias=False),
+                nn.SiLU(),
+                nn.Linear(global_embed_dim, global_embed_dim, bias=False)
+            )
+
+        if prepend_cond_dim > 0:
+            # Prepend conditioning
+            self.to_prepend_embed = nn.Sequential(
+                nn.Linear(prepend_cond_dim, embed_dim, bias=False),
+                nn.SiLU(),
+                nn.Linear(embed_dim, embed_dim, bias=False)
+            )
+
+        self.input_concat_dim = input_concat_dim
+
+        dim_in = io_channels + self.input_concat_dim
+
+        self.patch_size = patch_size
+
+        # Transformer
+
+        self.transformer_type = transformer_type
+
+        self.global_cond_type = global_cond_type
+
+
+        if pos_emb_strategy == "concatenation":
+
+            assert pos_emb_dim is not None, "pos_emb_dim must be specified if pos_emb_strategy is concatenation"
+            if pos_emb_type == "one-hot":
+                # One-hot positional embeddings
+                self.pos_emb = OneHotPositionalEmbedding(pos_emb_dim-max_num_tracks)
+                self.extra_dim = pos_emb_dim
+
+                def concat_pos_emb(x, pos=None, seq_start_pos=None, taxonomy=None):
+                    B, N, T, C = x.shape
+                    pos_emb = self.pos_emb(x.view(-1,x.shape[-2], x.shape[-1]), pos=pos, seq_start_pos=seq_start_pos)
+                    assert pos_emb.shape[-1] == pos_emb_dim-max_num_tracks, f"pos_emb shape mismatch: {pos_emb.shape[-1]} != {pos_emb_dim}"
+                    assert pos_emb.shape[-2] == T, f"pos_emb sequence length mismatch: {pos_emb.shape[-2]} != {x.shape[-2]}"
+
+                    pos_emb=pos_emb.unsqueeze(0).unsqueeze(0).expand(B,N,T,-1)
+
+
+                    assert pos_emb.ndim == 4, f"pos_emb must be 2D or 3D, got {pos_emb.ndim}"
+                    assert pos_emb.shape[0] == x.shape[0], f"pos_emb batch size mismatch: {pos_emb.shape[0]} != {x.shape[0]}"
+                    
+                    pos_emb_track=torch.zeros((B, N, T, max_num_tracks), device=x.device, dtype=x.dtype)
+                    for i in range(B):
+                        for j in range(N):
+                            if use_taxonomy_in_pos_emb:
+                                raise NotImplementedError("use_taxonomy_in_pos_emb is not implemented for pos_emb_type 'one-hot'")
+                                assert taxonomy is not None, "taxonomy must be provided if use_taxonomy_in_pos_emb is True"
+                                if taxonomy[i][j]=="92":
+                                    pos_emb_track[i, j, :, :]= F.one_hot(torch.tensor([0], device=x.device), num_classes=3).to(x.dtype).expand(T, -1)
+                                elif taxonomy[i][j]=="2":
+                                    pos_emb_track[i, j, :, :]= F.one_hot(torch.tensor([1], device=x.device), num_classes=3).to(x.dtype).expand(T, -1)
+                                elif taxonomy[i][j]=="11":
+                                    pos_emb_track[i, j, :, :]= F.one_hot(torch.tensor([2], device=x.device), num_classes=3).to(x.dtype).expand(T, -1)
+                            else:
+                                if j >= max_num_tracks:
+                                    j= j% max_num_tracks
+                                pos_emb_track[i, j, :, :]= F.one_hot(torch.tensor([j], device=x.device), num_classes=max_num_tracks).to(x.dtype).expand(T, -1)
+
+                    return torch.cat((x, pos_emb, pos_emb_track), dim=-1)
+                
+                self.pos_emb_fn = concat_pos_emb
+                self.remove_pos_emb = lambda x: x[..., :-pos_emb_dim] if pos_emb_dim > 0 else x
+
+            else:
+                raise ValueError(f"Unknown pos_emb_type: {pos_emb_type}")
+        else:
+            raise ValueError(f"Unknown pos_emb_strategy: {pos_emb_strategy}")
+        
+
+        if pos_emb_crossattn_strategy == "concatenation":
+            assert pos_emb_crossattn_dim is not None, "pos_emb_crossattn_dim must be specified if pos_emb_crossattn_strategy is concatenation"
+            if pos_emb_crossattn_type == "one-hot":
+                # One-hot positional embeddings for cross-attention
+                self.pos_emb_crossattn = OneHotPositionalEmbedding(pos_emb_crossattn_dim-max_num_tracks)
+                self.crossattn_extra_dim = pos_emb_crossattn_dim
+
+                if not self.project_cond_tokens:
+                    def concat_pos_emb_crossattn(x, pos=None, seq_start_pos=None, taxonomy=None):
+                        B, N, T, C = x.shape
+    
+                        pos_emb = self.pos_emb_crossattn(x.view(-1, T,C), pos=pos, seq_start_pos=seq_start_pos)
+                        assert pos_emb.shape[-1] == pos_emb_crossattn_dim-max_num_tracks, f"pos_emb shape mismatch: {pos_emb.shape[-1]} != {pos_emb_crossattn_dim}"
+                        assert pos_emb.shape[-2] == T, f"pos_emb sequence length mismatch: {pos_emb.shape[-2]} != {x.shape[-2]}"
+    
+                        pos_emb=pos_emb.unsqueeze(0).unsqueeze(0).expand(B,N,T,-1)
+    
+                        #if pos_emb.ndim == 3:
+                        assert pos_emb.ndim == 4, f"pos_emb must be 2D or 3D, got {pos_emb.ndim}"
+                        assert pos_emb.shape[0] == x.shape[0], f"pos_emb batch size mismatch: {pos_emb.shape[0]} != {x.shape[0]}"
+    
+                        pos_emb_track=torch.zeros((B, N, T, max_num_tracks), device=x.device, dtype=x.dtype)
+                        for i in range(B):
+                            for j in range(N):
+                                if use_taxonomy_in_pos_emb:
+                                    raise NotImplementedError("use_taxonomy_in_pos_emb is not implemented for pos_emb_crossattn_type 'one-hot'")
+                                    assert taxonomy is not None, "taxonomy must be provided if use_taxonomy_in_pos_emb is True"
+                                    if taxonomy[i][j]=="92":
+                                        pos_emb_track[i, j, :, :]= F.one_hot(torch.tensor([0], device=x.device), num_classes=3).to(x.dtype).expand(T, -1)
+                                    elif taxonomy[i][j]=="2":
+                                        pos_emb_track[i, j, :, :]= F.one_hot(torch.tensor([1], device=x.device), num_classes=3).to(x.dtype).expand(T, -1)
+                                    elif taxonomy[i][j]=="11":
+                                        pos_emb_track[i, j, :, :]= F.one_hot(torch.tensor([2], device=x.device), num_classes=3).to(x.dtype).expand(T, -1)
+                                else:
+                                    if j >= max_num_tracks:
+                                        j= j% max_num_tracks
+                                    pos_emb_track[i, j, :, :]= F.one_hot(torch.tensor([j], device=x.device), num_classes=max_num_tracks).to(x.dtype).expand(T, -1)
+    
+                        return torch.cat((x, pos_emb, pos_emb_track), dim=-1)
+                else:
+                    def concat_pos_emb_crossattn(x, pos=None, seq_start_pos=None, taxonomy=None):
+                        #print("x shape",x.shape)
+
+                        B, N, T, C = x.shape
+
+                        assert T*C== cond_token_dim, f"cond_token_proj_dim must match T*C, got {cond_token_dim} != {T*C}"
+
+                        #rehape to B, N, 1, T*C
+                        x= rearrange(x, "b n t c -> b n 1 (t c)")
+    
+                        #pos_emb = self.pos_emb_crossattn(x.view(-1, T,C), pos=pos, seq_start_pos=seq_start_pos)
+                        #assert pos_emb.shape[-1] == pos_emb_crossattn_dim-3, f"pos_emb shape mismatch: {pos_emb.shape[-1]} != {pos_emb_crossattn_dim}"
+                        #assert pos_emb.shape[-2] == T, f"pos_emb sequence length mismatch: {pos_emb.shape[-2]} != {x.shape[-2]}"
+    
+                        #pos_emb=pos_emb.unsqueeze(0).unsqueeze(0).expand(B,N,T,-1)
+    
+                        #if pos_emb.ndim == 3:
+                        #assert pos_emb.ndim == 4, f"pos_emb must be 2D or 3D, got {pos_emb.ndim}"
+                        #assert pos_emb.shape[0] == x.shape[0], f"pos_emb batch size mismatch: {pos_emb.shape[0]} != {x.shape[0]}"
+
+                        x=rearrange(x, "b n 1 c -> (b n) 1 c")
+                        x=self.to_cond_embed(x)
+                        x=rearrange(x, "(b n) 1 c -> b n 1 c", b=B, n=N)
+    
+                        pos_emb_track=torch.zeros((B, N, 1, pos_emb_crossattn_dim), device=x.device, dtype=x.dtype)
+                        for i in range(B):
+                            for j in range(N):
+                                if use_taxonomy_in_pos_emb:
+                                    assert taxonomy is not None, "taxonomy must be provided if use_taxonomy_in_pos_emb is True"
+                                    if taxonomy[i][j]=="92":
+                                        pos_emb_track[i, j, :, :]= F.one_hot(torch.tensor([0], device=x.device), num_classes=pos_emb_crossattn_dim).to(x.dtype)
+                                    elif taxonomy[i][j]=="2":
+                                        pos_emb_track[i, j, :, :]= F.one_hot(torch.tensor([1], device=x.device), num_classes=pos_emb_crossattn_dim).to(x.dtype)
+                                    elif taxonomy[i][j]=="11":
+                                        pos_emb_track[i, j, :, :]= F.one_hot(torch.tensor([2], device=x.device), num_classes=pos_emb_crossattn_dim).to(x.dtype)
+                                else:
+                                    pos_emb_track[i, j, :, :]= F.one_hot(torch.tensor([j], device=x.device), num_classes=pos_emb_crossattn_dim).to(x.dtype)
+    
+                        return torch.cat((x,  pos_emb_track), dim=-1)
+
+
+
+                self.pos_emb_crossattn_fn = concat_pos_emb_crossattn
+
+            else:
+                raise ValueError(f"Unknown pos_emb_type: {pos_emb_crossattn_type}")
+
+
+
+        global_dim = None
+
+        if self.global_cond_type == "adaLN":
+            # The global conditioning is projected to the embed_dim already at this point
+            global_dim = embed_dim
+
+        from networks.transformer import ContinuousTransformer
+
+        self.transformer = ContinuousTransformer(
+            dim=embed_dim,
+            depth=depth,
+            num_heads= num_heads,
+            dim_in=dim_in + pos_emb_dim,
+            dim_out=io_channels ,
+            cross_attend = cond_token_dim > 0,
+            cond_token_dim = cond_embed_dim + pos_emb_crossattn_dim,
+            global_cond_dim=global_dim,
+            **kwargs
+        )
+             
+
+        self.preprocess_conv = nn.Conv1d(dim_in, dim_in, 1, bias=False)
+        nn.init.zeros_(self.preprocess_conv.weight)
+        self.postprocess_conv = nn.Conv1d(io_channels, io_channels, 1, bias=False)
+        nn.init.zeros_(self.postprocess_conv.weight)
+
+
+
+
+    def _forward(
+        self, 
+        x, 
+        t, 
+        mask=None,
+        cross_attn_cond=None,
+        cross_attn_cond_mask=None,
+        input_concat_cond=None,
+        global_embed=None,
+        prepend_cond=None,
+        prepend_cond_mask=None,
+        return_info=False,
+        **kwargs):
+
+        t=t.squeeze(1)
+
+        #if cross_attn_cond is not None:
+        #    cross_attn_cond = self.to_cond_embed(cross_attn_cond)
+
+        if global_embed is not None:
+            # Project the global conditioning to the embedding dimension
+            global_embed = self.to_global_embed(global_embed)
+
+        prepend_inputs = None 
+        prepend_mask = None
+        prepend_length = 0
+        if prepend_cond is not None:
+            # Project the prepend conditioning to the embedding dimension
+            prepend_cond = self.to_prepend_embed(prepend_cond)
+            
+            prepend_inputs = prepend_cond
+            if prepend_cond_mask is not None:
+                prepend_mask = prepend_cond_mask
+
+
+        # Get the batch of timestep embeddings
+        timestep_embed = self.to_timestep_embed(self.timestep_features(t[:, None])) # (b, embed_dim)
+
+        # Timestep embedding is considered a global embedding. Add to the global conditioning if it exists
+
+        if self.timestep_cond_type == "global":
+            if global_embed is not None:
+                global_embed = global_embed + timestep_embed
+            else:
+                global_embed = timestep_embed
+        elif self.timestep_cond_type == "input_concat":
+            x = torch.cat([x, timestep_embed.unsqueeze(1).expand(-1, -1, x.shape[2])], dim=1)
+
+        # Add the global_embed to the prepend inputs if there is no global conditioning support in the transformer
+        if self.global_cond_type == "prepend" and global_embed is not None:
+            if prepend_inputs is None:
+                # Prepend inputs are just the global embed, and the mask is all ones
+                prepend_inputs = global_embed.unsqueeze(1)
+                prepend_mask = torch.ones((x.shape[0], 1), device=x.device, dtype=torch.bool)
+            else:
+                # Prepend inputs are the prepend conditioning + the global embed
+                prepend_inputs = torch.cat([prepend_inputs, global_embed.unsqueeze(1)], dim=1)
+                prepend_mask = torch.cat([prepend_mask, torch.ones((x.shape[0], 1), device=x.device, dtype=torch.bool)], dim=1)
+
+            prepend_length = prepend_inputs.shape[1]
+
+        extra_args = {}
+
+        if self.global_cond_type == "adaLN":
+            extra_args["global_cond"] = global_embed
+
+
+        output = self.transformer(x, prepend_embeds=prepend_inputs, context=cross_attn_cond, context_mask=cross_attn_cond_mask, mask=mask, prepend_mask=prepend_mask, return_info=return_info, **extra_args, **kwargs)
+
+        if return_info:
+            output, info = output
+       
+        output=output[:,prepend_length:,:]
+
+
+        if return_info:
+            return output, info
+
+        return output
+
+    def forward(
+        self, 
+        x, 
+        t, 
+        cross_attn_cond=None,
+        cross_attn_cond_mask=None,
+        input_concat_cond=None,
+        global_embed=None,
+        taxonomy=None,
+        mask=None,
+        return_info=False,
+        **kwargs):
+
+
+        model_dtype = next(self.parameters()).dtype
+        
+        x = x.to(model_dtype)
+
+        t = t.to(model_dtype)
+
+        if cross_attn_cond is not None:
+            cross_attn_cond = cross_attn_cond.to(model_dtype)
+
+        if input_concat_cond is not None:
+            input_concat_cond = input_concat_cond.to(model_dtype)
+            # Interpolate input_concat_cond to the same length as x
+            assert input_concat_cond.ndim == 4, f"input_concat_cond must be 4D, got {input_concat_cond.ndim}"
+            assert input_concat_cond.shape[0] == x.shape[0]
+            assert input_concat_cond.shape[1] == x.shape[1]
+            assert input_concat_cond.shape[-1] == x.shape[-1]
+            assert input_concat_cond.shape[-2] == self.input_concat_dim, f"input_concat_cond shape mismatch: {input_concat_cond.shape[-2]} != {self.input_concat_dim}"
+
+            x = torch.cat([x, input_concat_cond], dim=-2)
+
+        if global_embed is not None:
+            global_embed = global_embed.to(model_dtype)
+
+        if cross_attn_cond_mask is not None:
+            cross_attn_cond_mask = cross_attn_cond_mask.bool()
+
+            #cross_attn_cond_mask = None # Temporarily disabling conditioning masks due to kernel issue for flash attention
+
+
+        orig_shape = x.shape
+        x= rearrange(x, "b n c t -> (b n) c t")
+
+
+        x = self.preprocess_conv(x) + x
+
+        x=x.view(orig_shape)
+
+        x= rearrange(x, "b n c t -> b n t c")
+        #shape of contecxt is already [B, N, T, C] so no need to rearrange 
+
+        orig_shape = x.shape
+
+        x= self.pos_emb_fn(x, taxonomy=taxonomy)
+        cross_attn_cond= self.pos_emb_crossattn_fn(cross_attn_cond, taxonomy=taxonomy) 
+
+
+        x=rearrange(x, "b n t c -> b (n t) c")
+
+        cross_attn_cond_orig_shape = cross_attn_cond.shape
+        cross_attn_cond = rearrange(cross_attn_cond, "b n t c -> b (n t) c")
+
+        # rehape to [B, N \times T, C] for the transformer
+
+
+        # mask has shape [B, N ], I need to expand it to [B, N, T] for the convolution
+        mask= mask.unsqueeze(-1).expand(orig_shape[0], orig_shape[1], orig_shape[2]) 
+        mask= rearrange(mask, "b n t -> b (n t)")
+
+        cross_attn_cond_mask = cross_attn_cond_mask.unsqueeze(-1).expand(cross_attn_cond_orig_shape[0], cross_attn_cond_orig_shape[1], cross_attn_cond_orig_shape[2])
+        cross_attn_cond_mask = rearrange(cross_attn_cond_mask, "b n t -> b (n t)")
+
+
+        out= self._forward(
+            x,
+            t,
+            cross_attn_cond=cross_attn_cond, 
+            cross_attn_cond_mask=cross_attn_cond_mask, 
+            input_concat_cond=input_concat_cond, 
+            global_embed=global_embed, 
+            mask=mask,
+            return_info=return_info,
+            **kwargs
+        )
+
+        #print("out shape", out.shape)
+
+        out = rearrange(out, "b t c -> b c t")
+        out= self.postprocess_conv(out) + out
+        out = rearrange(out, "b c t -> b t c")
+
+        #print("out shape after postprocess", out.shape)
+
+        #now we reshape... 
+        out = rearrange(out, "b (n t) c -> b n t c", n=orig_shape[1], t=orig_shape[2])
+
+        out=rearrange(out, "b n t c -> b n c t")
+
+        return out

networks/transformer.py

@@ -0,0 +1,932 @@
+#adapted from https://github.com/Stability-AI/stable-audio-tools/blob/main/stable_audio_tools/models/transformer.py
+
+from functools import reduce, partial
+from packaging import version
+
+from einops import rearrange, repeat
+from einops.layers.torch import Rearrange
+import torch
+import torch.nn.functional as F
+from torch import nn, einsum
+from torch.amp import autocast
+from torch.nn.utils.parametrizations import weight_norm
+from typing import Callable, Literal
+
+#from .utils import compile
+
+import os
+enable_torch_compile = os.environ.get("ENABLE_TORCH_COMPILE", "0") == "1"
+
+def compile(function, *args, **kwargs):
+    
+    if enable_torch_compile:
+        try:
+            return torch.compile(function, *args, **kwargs)
+        except RuntimeError:
+            return function
+
+    return function
+
+try:
+    from flash_attn import flash_attn_func, flash_attn_kvpacked_func
+except ImportError as e:
+    print(e)
+    print('flash_attn not installed, disabling Flash Attention')
+    flash_attn_kvpacked_func = None
+    flash_attn_func = None
+
+try:
+    import natten
+except ImportError:
+    natten = None
+
+def checkpoint(function, *args, **kwargs):
+    kwargs.setdefault("use_reentrant", False)
+    return torch.utils.checkpoint.checkpoint(function, *args, **kwargs)
+
+
+# Copied and modified from https://github.com/lucidrains/x-transformers/blob/main/x_transformers/attend.py under MIT License
+# License can be found in LICENSES/LICENSE_XTRANSFORMERS.txt
+
+def create_causal_mask(i, j, device):
+    return torch.ones((i, j), device = device, dtype = torch.bool).triu(j - i + 1)
+
+def or_reduce(masks):
+    head, *body = masks
+    for rest in body:
+        head = head | rest
+    return head
+
+class FourierFeatures(nn.Module):
+    def __init__(self, in_features, out_features, std=1.):
+        super().__init__()
+        assert out_features % 2 == 0
+        self.weight = nn.Parameter(torch.randn(
+            [out_features // 2, in_features]) * std)
+
+    def forward(self, input):
+        f = 2 * math.pi * input @ self.weight.T
+        return torch.cat([f.cos(), f.sin()], dim=-1)
+
+
+class AbsolutePositionalEmbedding(nn.Module):
+    def __init__(self, dim, max_seq_len, fixed = False, seed = 42):
+        super().__init__()
+        self.scale = dim ** -0.5
+        self.max_seq_len = max_seq_len
+        self.fixed=fixed
+        if fixed:
+            torch.manual_seed(seed)
+            emb = torch.randn(max_seq_len, dim)
+            self.register_buffer('emb', emb)  # replaces nn.Embedding
+        else:
+            self.emb = nn.Embedding(max_seq_len, dim)
+
+    def forward(self, x, pos = None, seq_start_pos = None):
+        seq_len, device = x.shape[1], x.device
+        assert seq_len <= self.max_seq_len, f'you are passing in a sequence length of {seq_len} but your absolute positional embedding has a max sequence length of {self.max_seq_len}'
+
+        if pos is None:
+            pos = torch.arange(seq_len, device = device)
+
+        if seq_start_pos is not None:
+            pos = (pos - seq_start_pos[..., None]).clamp(min = 0)
+
+        if not self.fixed:
+            pos_emb = self.emb(pos)
+        else:
+            pos_emb = self.emb[pos]
+        pos_emb = pos_emb * self.scale
+        return pos_emb
+
+class ScaledSinusoidalEmbedding(nn.Module):
+    def __init__(self, dim, theta = 10000):
+        super().__init__()
+        assert (dim % 2) == 0, 'dimension must be divisible by 2'
+        self.scale = nn.Parameter(torch.ones(1) * dim ** -0.5)
+
+        half_dim = dim // 2
+        freq_seq = torch.arange(half_dim).float() / half_dim
+        inv_freq = theta ** -freq_seq
+        self.register_buffer('inv_freq', inv_freq, persistent = False)
+
+    def forward(self, x, pos = None, seq_start_pos = None):
+        seq_len, device = x.shape[1], x.device
+
+        if pos is None:
+            pos = torch.arange(seq_len, device = device)
+
+        if seq_start_pos is not None:
+            pos = pos - seq_start_pos[..., None]
+
+        emb = einsum('i, j -> i j', pos, self.inv_freq)
+        emb = torch.cat((emb.sin(), emb.cos()), dim = -1)
+        return emb * self.scale
+    
+class RotaryEmbedding(nn.Module):
+    def __init__(
+        self,
+        dim,
+        use_xpos = False,
+        scale_base = 512,
+        interpolation_factor = 1.,
+        base = 10000,
+        base_rescale_factor = 1.
+    ):
+        super().__init__()
+        # proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning
+        # has some connection to NTK literature
+        # https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_have/
+        base *= base_rescale_factor ** (dim / (dim - 2))
+
+        inv_freq = 1. / (base ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer('inv_freq', inv_freq)
+
+        assert interpolation_factor >= 1.
+        self.interpolation_factor = interpolation_factor
+
+        if not use_xpos:
+            self.register_buffer('scale', None)
+            return
+
+        scale = (torch.arange(0, dim, 2) + 0.4 * dim) / (1.4 * dim)
+
+        self.scale_base = scale_base
+        self.register_buffer('scale', scale)
+
+    def forward_from_seq_len(self, seq_len):
+        device = self.inv_freq.device
+
+        t = torch.arange(seq_len, device = device)
+        return self.forward(t)
+
+    @autocast("cuda", enabled = False)
+    def forward(self, t):
+        device = self.inv_freq.device
+
+        t = t.to(torch.float32)
+
+        t = t / self.interpolation_factor
+
+        freqs = torch.einsum('i , j -> i j', t, self.inv_freq)
+        freqs = torch.cat((freqs, freqs), dim = -1)
+
+        if self.scale is None:
+            return freqs, 1.
+
+        power = (torch.arange(seq_len, device = device) - (seq_len // 2)) / self.scale_base
+        scale = self.scale ** rearrange(power, 'n -> n 1')
+        scale = torch.cat((scale, scale), dim = -1)
+
+        return freqs, scale
+
+def rotate_half(x):
+    x = rearrange(x, '... (j d) -> ... j d', j = 2)
+    x1, x2 = x.unbind(dim = -2)
+    return torch.cat((-x2, x1), dim = -1)
+
+@autocast("cuda", enabled = False)
+def apply_rotary_pos_emb(t, freqs, scale = 1):
+    out_dtype = t.dtype
+
+    # cast to float32 if necessary for numerical stability
+    dtype = reduce(torch.promote_types, (t.dtype, freqs.dtype, torch.float32))
+    rot_dim, seq_len = freqs.shape[-1], t.shape[-2]
+    freqs, t = freqs.to(dtype), t.to(dtype)
+    freqs = freqs[-seq_len:, :]
+
+    if t.ndim == 4 and freqs.ndim == 3:
+        freqs = rearrange(freqs, 'b n d -> b 1 n d')
+
+    # partial rotary embeddings, Wang et al. GPT-J
+    t, t_unrotated = t[..., :rot_dim], t[..., rot_dim:]
+    t = (t * freqs.cos() * scale) + (rotate_half(t) * freqs.sin() * scale)
+
+    t, t_unrotated = t.to(out_dtype), t_unrotated.to(out_dtype)
+
+    return torch.cat((t, t_unrotated), dim = -1)
+
+# norms
+class LayerNorm(nn.Module):
+    def __init__(self, dim, bias=False, fix_scale=False, force_fp32=False, eps=1e-5):
+        """
+        bias-less layernorm has been shown to be more stable. most newer models have moved towards rmsnorm, also bias-less
+        """
+        super().__init__()
+
+        if fix_scale:
+            self.register_buffer("gamma", torch.ones(dim))
+        else:
+            self.gamma = nn.Parameter(torch.ones(dim))
+
+        if bias:
+            self.beta = nn.Parameter(torch.zeros(dim))
+        else:
+            self.register_buffer("beta", torch.zeros(dim))
+
+        self.eps = eps
+
+        self.force_fp32 = force_fp32
+
+    def forward(self, x):
+        if not self.force_fp32:
+            return F.layer_norm(x, x.shape[-1:], weight=self.gamma, bias=self.beta, eps=self.eps)
+        else:
+            output = F.layer_norm(x.float(), x.shape[-1:], weight=self.gamma.float(), bias=self.beta.float(), eps=self.eps)
+            return output.to(x.dtype)
+
+class LayerScale(nn.Module):
+    def __init__(self, dim, init_val = 1e-2):
+        super().__init__()
+        self.scale = nn.Parameter(torch.full([dim], init_val))
+    def forward(self, x):
+        return x * self.scale
+
+
+# feedforward
+
+class GLU(nn.Module):
+    def __init__(
+        self,
+        dim_in,
+        dim_out,
+        activation: Callable,
+        use_conv = False,
+        conv_kernel_size = 3,
+    ):
+        super().__init__()
+        self.act = activation
+        self.proj = nn.Linear(dim_in, dim_out * 2) if not use_conv else nn.Conv1d(dim_in, dim_out * 2, conv_kernel_size, padding = (conv_kernel_size // 2))
+        self.use_conv = use_conv
+
+    def forward(self, x):
+        if self.use_conv:
+            x = rearrange(x, 'b n d -> b d n')
+            x = self.proj(x)
+            x = rearrange(x, 'b d n -> b n d')
+        else:
+            x = self.proj(x)
+
+        x, gate = x.chunk(2, dim = -1)
+        return x * self.act(gate)
+
+class FeedForward(nn.Module):
+    def __init__(
+        self,
+        dim,
+        dim_out = None,
+        mult = 4,
+        no_bias = False,
+        glu = True,
+        use_conv = False,
+        conv_kernel_size = 3,
+        zero_init_output = True,
+    ):
+        super().__init__()
+        inner_dim = int(dim * mult)
+
+        # Default to SwiGLU
+
+        activation = nn.SiLU()
+
+        dim_out = dim if dim_out is None else dim_out
+
+        if glu:
+            linear_in = GLU(dim, inner_dim, activation)
+        else:
+            linear_in = nn.Sequential(
+                Rearrange('b n d -> b d n') if use_conv else nn.Identity(),
+                nn.Linear(dim, inner_dim, bias = not no_bias) if not use_conv else nn.Conv1d(dim, inner_dim, conv_kernel_size, padding = (conv_kernel_size // 2), bias = not no_bias),
+                Rearrange('b n d -> b d n') if use_conv else nn.Identity(),
+                activation
+            )
+
+        linear_out = nn.Linear(inner_dim, dim_out, bias = not no_bias) if not use_conv else nn.Conv1d(inner_dim, dim_out, conv_kernel_size, padding = (conv_kernel_size // 2), bias = not no_bias)
+
+        # init last linear layer to 0
+        if zero_init_output:
+            nn.init.zeros_(linear_out.weight)
+            if not no_bias:
+                nn.init.zeros_(linear_out.bias)
+
+
+        self.ff = nn.Sequential(
+            linear_in,
+            Rearrange('b d n -> b n d') if use_conv else nn.Identity(),
+            linear_out,
+            Rearrange('b n d -> b d n') if use_conv else nn.Identity(),
+        )
+
+    #@compile
+    def forward(self, x):
+        return self.ff(x)
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        dim,
+        dim_heads = 64,
+        dim_context = None,
+        causal = False,
+        zero_init_output=True,
+        qk_norm: Literal['l2', 'ln', 'none'] = 'none',
+        natten_kernel_size = None,
+        sliding_window = [-1, -1],
+        feat_scale = False
+    ):
+        super().__init__()
+        self.dim = dim
+        self.dim_heads = dim_heads
+        self.causal = causal
+
+        dim_kv = dim_context if dim_context is not None else dim
+        
+        self.num_heads = dim // dim_heads
+        self.kv_heads = dim_kv // dim_heads
+        #print("num_heads", self.num_heads, "kv_heads", self.kv_heads)
+        #print("dim", dim, "dim_heads", dim_heads, "dim_kv", dim_kv, dim/dim_heads, dim_kv/dim_heads)
+        assert  dim % dim_heads == 0, f'dim {dim} must be divisible by dim_heads {dim_heads}'
+        assert dim_kv % dim_heads == 0, f'dim_kv {dim_kv} must be divisible by dim_heads {dim_heads}'
+
+
+        if dim_context is not None:
+            self.to_q = nn.Linear(dim, dim, bias=False)
+            self.to_kv = nn.Linear(dim_kv, dim_kv * 2, bias=False)
+        else:
+            self.to_qkv = nn.Linear(dim, dim * 3, bias=False)
+
+        self.to_out = nn.Linear(dim, dim, bias=False)
+
+        if zero_init_output:
+            nn.init.zeros_(self.to_out.weight)
+
+        if qk_norm not in ['l2', 'ln', 'none']:
+            raise ValueError(f'qk_norm must be one of ["l2", "ln", "none"], got {qk_norm}')
+            
+        self.qk_norm = qk_norm
+
+        if self.qk_norm == "ln":
+            self.q_norm = nn.LayerNorm(dim_heads, elementwise_affine=True, eps=1.0e-6)
+            self.k_norm = nn.LayerNorm(dim_heads, elementwise_affine=True, eps=1.0e-6)
+
+        # Using 1d neighborhood attention
+        self.natten_kernel_size = natten_kernel_size
+        if natten_kernel_size is not None:
+            return
+
+        self.use_pt_flash = torch.cuda.is_available() and version.parse(torch.__version__) >= version.parse('2.0.0')
+
+        self.use_fa_flash = torch.cuda.is_available() and flash_attn_func is not None
+
+        self.sdp_kwargs = dict(
+            enable_flash = True,
+            enable_math = True,
+            enable_mem_efficient = True
+        )
+
+        self.sliding_window = sliding_window
+        if not (sliding_window[0] == -1 and sliding_window[1] == -1)  and not self.use_fa_flash:
+            print('Sliding window is being used, but Flash Attention is not. Please install Flash Attention to get correct results')
+
+        self.feat_scale = feat_scale
+
+        if self.feat_scale:
+            self.lambda_dc = nn.Parameter(torch.zeros(dim))
+            self.lambda_hf = nn.Parameter(torch.zeros(dim))
+
+    def flash_attn(
+            self,
+            q, 
+            k, 
+            v,
+            mask = None,
+            causal = None
+    ):
+        batch, heads, q_len, _, k_len, device = *q.shape, k.shape[-2], q.device
+        kv_heads = k.shape[1]
+        # Recommended for multi-query single-key-value attention by Tri Dao
+        # kv shape torch.Size([1, 512, 64]) -> torch.Size([1, 8, 512, 64])
+
+        if heads != kv_heads:
+            # Repeat interleave kv_heads to match q_heads
+            heads_per_kv_head = heads // kv_heads
+            assert heads % kv_heads == 0, f'heads {heads} must be divisible by kv_heads {kv_heads}'
+            k, v = map(lambda t: t.repeat_interleave(heads_per_kv_head, dim = 1), (k, v))
+
+        if k.ndim == 3:
+            k = rearrange(k, 'b ... -> b 1 ...').expand_as(q)
+
+        if v.ndim == 3:
+            v = rearrange(v, 'b ... -> b 1 ...').expand_as(q)
+
+        causal = self.causal if causal is None else causal
+
+        if q_len == 1 and causal:
+            causal = False
+        
+        if mask is not None:
+            assert mask.ndim == 4
+            mask = mask.expand(batch, heads, q_len, k_len)
+
+        # handle kv cache - this should be bypassable in updated flash attention 2
+
+        if k_len > q_len and causal:
+            causal_mask = self.create_causal_mask(q_len, k_len, device = device)
+            if mask is None:
+                mask = ~causal_mask
+            else:
+                mask = mask & ~causal_mask
+            causal = False
+
+        # manually handle causal mask, if another mask was given
+
+        row_is_entirely_masked = None
+
+        if mask is not None and causal:
+            causal_mask = self.create_causal_mask(q_len, k_len, device = device)
+            mask = mask & ~causal_mask
+
+            # protect against an entire row being masked out
+
+            row_is_entirely_masked = ~mask.any(dim = -1)
+            mask[..., 0] = mask[..., 0] | row_is_entirely_masked
+
+            causal = False
+        
+        #with torch.backends.cuda.sdp_kernel(**self.sdp_kwargs):
+        #print("q.shape", q.shape, "k.shape", k.shape, "v.shape", v.shape, "mask.shape", mask.shape if mask is not None else None, "causal", causal)
+        out = F.scaled_dot_product_attention(
+            q, k, v,
+            attn_mask = mask,
+            is_causal = causal
+        )
+
+        # for a row that is entirely masked out, should zero out the output of that row token
+
+        if row_is_entirely_masked is not None:
+            out = out.masked_fill(row_is_entirely_masked[..., None], 0.)
+
+        return out
+
+    @compile
+    def apply_qk_layernorm(self, q, k):
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+        return q, k
+
+    #@compile
+    def forward(
+        self,
+        x,
+        context = None,
+        mask = None,
+        context_mask = None,
+        rotary_pos_emb = None,
+        causal = None
+    ):
+        h, kv_h, has_context = self.num_heads, self.kv_heads, context is not None
+
+        #print("h", h, "kv_h", kv_h, "has_context", has_context)
+        #print("mask", mask, "context_mask", context_mask)
+
+        kv_input = context if has_context else x
+
+        if hasattr(self, 'to_q'):
+            # Use separate linear projections for q and k/v
+            q = self.to_q(x)
+            q = rearrange(q, 'b n (h d) -> b h n d', h = h)
+
+            k, v = self.to_kv(kv_input).chunk(2, dim=-1)
+
+            k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = kv_h), (k, v))
+        else:
+            # Use fused linear projection
+            q, k, v = self.to_qkv(x).chunk(3, dim=-1)
+            q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = h), (q, k, v))
+        
+        # Normalize q and k for cosine sim attention
+        if self.qk_norm == "l2":
+            q = F.normalize(q, dim=-1)
+            k = F.normalize(k, dim=-1)
+        elif self.qk_norm == "ln":
+            q, k = self.apply_qk_layernorm(q, k)
+
+        if rotary_pos_emb is not None and not has_context:
+            freqs, _ = rotary_pos_emb
+
+            q_dtype = q.dtype
+            k_dtype = k.dtype
+
+            q = q.to(torch.float32)
+            k = k.to(torch.float32)
+            freqs = freqs.to(torch.float32)
+
+            q = apply_rotary_pos_emb(q, freqs)
+            k = apply_rotary_pos_emb(k, freqs)
+
+            q = q.to(q_dtype)
+            k = k.to(k_dtype)
+        
+        input_mask = context_mask 
+
+        if input_mask is None and not has_context:
+            input_mask = mask
+
+        # determine masking
+        masks = []
+        final_attn_mask = None # The mask that will be applied to the attention matrix, taking all masks into account
+
+        if input_mask is not None:
+            input_mask = rearrange(input_mask, 'b j -> b 1 1 j')
+            masks.append(~input_mask)
+
+        # Other masks will be added here later
+
+        if len(masks) > 0:
+            final_attn_mask = ~or_reduce(masks)
+
+        n, device = q.shape[-2], q.device
+
+        causal = self.causal if causal is None else causal
+
+        if n == 1 and causal:
+            causal = False
+
+        if self.natten_kernel_size is not None:
+            if natten is None:
+                raise ImportError('natten not installed, please install natten to use neighborhood attention')
+            
+            dtype_in = q.dtype
+            q, k, v = map(lambda t: t.to(torch.float32), (q, k, v))
+
+            attn = natten.functional.na1d_qk(q, k, kernel_size = self.natten_kernel_size, dilation=1)
+
+            if final_attn_mask is not None:
+                attn = attn.masked_fill(final_attn_mask, -torch.finfo(attn.dtype).max)
+
+            attn = F.softmax(attn, dim=-1, dtype=torch.float32)
+
+            out = natten.functional.na1d_av(attn, v, kernel_size = self.natten_kernel_size, dilation=1).to(dtype_in)
+
+        # Prioritize Flash Attention 2
+        elif self.use_fa_flash:
+            assert final_attn_mask is None, 'masking not yet supported for Flash Attention 2'
+            # Flash Attention 2 requires FP16 inputs
+            fa_dtype_in = q.dtype
+
+            q, k, v = map(lambda t: rearrange(t, 'b h n d -> b n h d'), (q, k, v))
+
+            if fa_dtype_in != torch.float16 and fa_dtype_in != torch.bfloat16:
+                q, k, v = map(lambda t: t.to(torch.float16), (q, k, v))
+            
+            out = flash_attn_func(q, k, v, causal = causal, window_size=self.sliding_window)
+            
+            out = rearrange(out.to(fa_dtype_in), 'b n h d -> b h n d')
+
+        # Fall back to PyTorch implementation
+        elif self.use_pt_flash:
+            out = self.flash_attn(q, k, v, causal = causal, mask = final_attn_mask)
+
+        else:
+            # Fall back to custom implementation
+
+            if h != kv_h:
+                # Repeat interleave kv_heads to match q_heads
+                heads_per_kv_head = h // kv_h
+                k, v = map(lambda t: t.repeat_interleave(heads_per_kv_head, dim = 1), (k, v))
+
+            scale = 1. / (q.shape[-1] ** 0.5)
+
+            kv_einsum_eq = 'b j d' if k.ndim == 3 else 'b h j d'
+
+            dots = einsum(f'b h i d, {kv_einsum_eq} -> b h i j', q, k) * scale
+            
+            i, j, dtype = *dots.shape[-2:], dots.dtype
+
+            mask_value = -torch.finfo(dots.dtype).max
+
+            if final_attn_mask is not None:
+                dots = dots.masked_fill(~final_attn_mask, mask_value)
+
+            if causal:
+                causal_mask = self.create_causal_mask(i, j, device = device)
+                dots = dots.masked_fill(causal_mask, mask_value)
+
+            attn = F.softmax(dots, dim=-1, dtype=torch.float32)
+            attn = attn.type(dtype)
+
+            out = einsum(f'b h i j, {kv_einsum_eq} -> b h i d', attn, v)
+
+        # merge heads
+        out = rearrange(out, ' b h n d -> b n (h d)')
+
+        # Communicate between heads
+        
+        # with autocast(enabled = False):
+        #     out_dtype = out.dtype
+        #     out = out.to(torch.float32)
+        #     out = self.to_out(out).to(out_dtype)
+        out = self.to_out(out)
+
+        if self.feat_scale:
+            out_dc = out.mean(dim=-2, keepdim=True)
+            out_hf = out - out_dc
+
+            # Selectively modulate DC and high frequency components
+            out = out + self.lambda_dc * out_dc + self.lambda_hf * out_hf
+
+        if mask is not None:
+            mask = rearrange(mask, 'b n -> b n 1')
+            out = out.masked_fill(~mask, 0.)
+
+        return out
+
+class ConformerModule(nn.Module):
+    def __init__(
+        self,
+        dim,
+        norm_kwargs = {},
+    ):     
+
+        super().__init__()
+
+        self.dim = dim
+        
+        self.in_norm = LayerNorm(dim, **norm_kwargs)
+        self.pointwise_conv = nn.Conv1d(dim, dim, kernel_size=1, bias=False)
+        self.glu = GLU(dim, dim, nn.SiLU())
+        self.depthwise_conv = nn.Conv1d(dim, dim, kernel_size=17, groups=dim, padding=8, bias=False)
+        self.mid_norm = LayerNorm(dim, **norm_kwargs) # This is a batch norm in the original but I don't like batch norm
+        self.swish = nn.SiLU()
+        self.pointwise_conv_2 = nn.Conv1d(dim, dim, kernel_size=1, bias=False)
+
+    #@compile
+    def forward(self, x):
+        x = self.in_norm(x)
+        x = rearrange(x, 'b n d -> b d n')
+        x = self.pointwise_conv(x)
+        x = rearrange(x, 'b d n -> b n d')
+        x = self.glu(x)
+        x = rearrange(x, 'b n d -> b d n')
+        x = self.depthwise_conv(x)
+        x = rearrange(x, 'b d n -> b n d')
+        x = self.mid_norm(x)
+        x = self.swish(x)
+        x = rearrange(x, 'b n d -> b d n')
+        x = self.pointwise_conv_2(x)
+        x = rearrange(x, 'b d n -> b n d')
+
+        return x
+
+class TransformerBlock(nn.Module):
+    def __init__(
+            self,
+            dim,
+            dim_heads = 64,
+            cross_attend = False,
+            dim_context = None,
+            global_cond_dim = None,
+            causal = False,
+            zero_init_branch_outputs = True,
+            conformer = False,
+            layer_ix = -1,
+            remove_norms = False,
+            add_rope = False,
+            layer_scale = False,
+            attn_kwargs = {},
+            ff_kwargs = {},
+            norm_kwargs = {}
+    ):
+        
+        super().__init__()
+        self.dim = dim
+        self.dim_heads = dim_heads
+        self.cross_attend = cross_attend
+        self.dim_context = dim_context
+        self.causal = causal
+       
+        if layer_scale and zero_init_branch_outputs:
+            print('zero_init_branch_outputs is redundant with layer_scale, setting zero_init_branch_outputs to False')
+            zero_init_branch_outputs = False
+            
+        self.pre_norm = LayerNorm(dim,**norm_kwargs) if not remove_norms else nn.Identity()
+
+        self.add_rope = add_rope
+
+        self.self_attn = Attention(
+            dim,
+            dim_heads = dim_heads,
+            causal = causal,
+            zero_init_output=zero_init_branch_outputs,
+            **attn_kwargs
+        )
+
+        self.self_attn_scale = LayerScale(dim) if layer_scale else nn.Identity()
+
+        self.cross_attend = cross_attend
+        if cross_attend:
+            self.cross_attend_norm = LayerNorm(dim, **norm_kwargs) if not remove_norms else nn.Identity()
+            self.cross_attn = Attention(
+                dim,
+                dim_heads = dim_heads,
+                dim_context=dim_context,
+                causal = causal,
+                zero_init_output=zero_init_branch_outputs,
+                **attn_kwargs
+            )
+            self.cross_attn_scale = LayerScale(dim) if layer_scale else nn.Identity()
+        
+        self.ff_norm = LayerNorm(dim, **norm_kwargs) if not remove_norms else nn.Identity()
+        self.ff = FeedForward(dim, zero_init_output=zero_init_branch_outputs, **ff_kwargs)
+        self.ff_scale = LayerScale(dim) if layer_scale else nn.Identity()
+
+        self.layer_ix = layer_ix
+
+        self.conformer = None
+        if conformer:
+            self.conformer = ConformerModule(dim, norm_kwargs=norm_kwargs)
+            self.conformer_scale = LayerScale(dim) if layer_scale else nn.Identity()
+
+        self.global_cond_dim = global_cond_dim
+
+        if global_cond_dim is not None:
+            self.to_scale_shift_gate = nn.Parameter(torch.randn(6*dim)/dim**0.5)
+
+        self.rope = RotaryEmbedding(max(dim_heads // 2, 32)) if add_rope else None
+
+    @compile
+    def forward(
+        self,
+        x,
+        context = None,
+        global_cond=None,
+        mask = None,
+        context_mask = None,
+        rotary_pos_emb = None
+    ):
+        if rotary_pos_emb is None and self.add_rope:
+            rotary_pos_emb = self.rope.forward_from_seq_len(x.shape[-2])
+
+        if self.global_cond_dim is not None and self.global_cond_dim > 0 and global_cond is not None:
+            
+            scale_self, shift_self, gate_self, scale_ff, shift_ff, gate_ff = (self.to_scale_shift_gate + global_cond).unsqueeze(1).chunk(6, dim=-1)
+
+            # self-attention with adaLN
+            residual = x
+            x = self.pre_norm(x)
+            x = x * (1 + scale_self) + shift_self
+            x = self.self_attn(x, mask = mask, rotary_pos_emb = rotary_pos_emb)
+            x = x * torch.sigmoid(1 - gate_self)
+            x = self.self_attn_scale(x)
+            x = x + residual
+
+            #print("self.cross_attend", self.cross_attend)
+            #print("context", context.shape)
+            if context is not None and self.cross_attend:
+                #print("mask", mask)
+                #print("context_mask", context_mask)
+
+                x = x + self.cross_attn_scale(self.cross_attn(self.cross_attend_norm(x), context = context, context_mask = context_mask))
+            
+            if self.conformer is not None:
+                x = x + self.conformer_scale(self.conformer(x))
+
+            # feedforward with adaLN
+            residual = x
+            x = self.ff_norm(x)
+            x = x * (1 + scale_ff) + shift_ff
+            x = self.ff(x)
+            x = x * torch.sigmoid(1 - gate_ff)
+            x = self.ff_scale(x)
+            x = x + residual
+
+        else:
+            x = x + self.self_attn_scale(self.self_attn(self.pre_norm(x), mask = mask, rotary_pos_emb = rotary_pos_emb))
+
+            if context is not None and self.cross_attend:
+                #print(x.shape, context.shape, context_mask.shape if context_mask is not None else None)
+                #print(self.dim, self.dim_context, self.dim_heads)
+                x = x + self.cross_attn_scale(self.cross_attn(self.cross_attend_norm(x), context = context, context_mask = context_mask))
+                    
+            if self.conformer is not None:
+                x = x + self.conformer_scale(self.conformer(x))
+
+            x = x + self.ff_scale(self.ff(self.ff_norm(x)))
+        return x
+        
+class ContinuousTransformer(nn.Module):
+    def __init__(
+        self,
+        dim,
+        depth,
+        *,
+        dim_in = None,
+        dim_out = None,
+        num_heads = 64,
+        cross_attend=False,
+        cond_token_dim=None,
+        final_cross_attn_ix=-1,
+        global_cond_dim=None,
+        causal=False,
+        zero_init_branch_outputs=True,
+        conformer=False,
+        **kwargs
+        ):
+
+        super().__init__()
+
+        self.dim = dim
+        self.depth = depth
+        self.causal = causal
+        self.layers = nn.ModuleList([])
+
+        self.project_in = nn.Linear(dim_in, dim, bias=False) if dim_in is not None else nn.Identity()
+        self.project_out = nn.Linear(dim, dim_out, bias=False) if dim_out is not None else nn.Identity()
+
+
+        
+
+        self.global_cond_embedder = None
+        if global_cond_dim is not None:
+            self.global_cond_embedder = nn.Sequential(
+                nn.Linear(global_cond_dim, dim),
+                nn.SiLU(),
+                nn.Linear(dim, dim * 6)
+            )
+
+        self.final_cross_attn_ix = final_cross_attn_ix
+
+        for i in range(depth):
+            should_cross_attend = cross_attend and (self.final_cross_attn_ix == -1 or i <= (self.final_cross_attn_ix))
+            self.layers.append(
+                TransformerBlock(
+                    dim,
+                    dim_heads = (dim) // num_heads,
+                    cross_attend = should_cross_attend,
+                    dim_context = cond_token_dim,
+                    global_cond_dim = global_cond_dim,
+                    causal = causal,
+                    zero_init_branch_outputs = zero_init_branch_outputs,
+                    conformer=conformer,
+                    layer_ix=i,
+                    **kwargs
+                )
+            )
+
+
+    def forward(
+        self,
+        x,
+        mask = None,
+        prepend_embeds = None,
+        prepend_mask = None,
+        global_cond = None,
+        return_info = False,
+        context = None,
+        context_mask = None,
+        **kwargs
+    ):
+        batch, seq, device = *x.shape[:2], x.device
+
+        model_dtype = next(self.parameters()).dtype
+        x = x.to(model_dtype)
+
+        info = {
+            "hidden_states": [],
+        }
+
+
+        x = self.project_in(x)
+        
+        if prepend_embeds is not None:
+            prepend_length, prepend_dim = prepend_embeds.shape[1:]
+
+            assert prepend_dim == x.shape[-1], 'prepend dimension must match sequence dimension'
+
+            x = torch.cat((prepend_embeds, x), dim = -2)
+
+            if prepend_mask is not None or mask is not None:
+                mask = mask if mask is not None else torch.ones((batch, seq), device = device, dtype = torch.bool)
+                prepend_mask = prepend_mask if prepend_mask is not None else torch.ones((batch, prepend_length), device = device, dtype = torch.bool)
+
+                mask = torch.cat((prepend_mask, mask), dim = -1)
+
+
+        if global_cond is not None and self.global_cond_embedder is not None:
+            global_cond = self.global_cond_embedder(global_cond)
+
+        # Iterate over the transformer layers
+        for layer in self.layers:
+            #x = layer(x, rotary_pos_emb = rotary_pos_emb, global_cond=global_cond, **kwargs)
+            x = checkpoint(layer, x, rotary_pos_emb = None, global_cond=global_cond, context=context, mask=mask, context_mask=context_mask, **kwargs)
+
+            if return_info:
+                info["hidden_states"].append(x)
+
+        x = self.project_out(x)
+
+
+        if return_info:
+            return x, info
+        
+        return x

requirements.txt

@@ -0,0 +1,15 @@
+git+https://github.com/TEAMuP-dev/pyharp.git@v0.3.0
+torch
+torchaudio
+omegaconf>=2.3.0
+hydra-core>=1.3.0
+numpy
+scipy
+soundfile
+pyloudnorm
+einops
+packaging
+librosa
+transformers
+wget
+huggingface-hub

utils/MSS_loss.py

@@ -0,0 +1,260 @@
+"""
+    Implementation of objective functions used in the task 'ITO-Master'
+    https://github.com/SonyResearch/ITO-Master/blob/master/ito_master/modules/loss.py
+"""
+
+import numpy as np
+import math
+import torch
+import torch.nn.functional as F
+import torch.nn as nn
+import auraloss
+import torchaudio
+import warnings
+
+
+import os
+import sys
+
+currentdir = os.path.dirname(os.path.realpath(__file__))
+sys.path.append(os.path.dirname(currentdir))
+
+class FrontEnd(nn.Module):
+    def __init__(self, channel='stereo', \
+                        n_fft=2048, \
+                        n_mels=128, \
+                        sample_rate=44100, \
+                        hop_length=None, \
+                        win_length=None, \
+                        window="hann", \
+                        eps=1e-7, \
+                        device=torch.device("cpu")):
+        super(FrontEnd, self).__init__()
+        self.channel = channel
+        self.n_fft = n_fft
+        self.n_mels = n_mels
+        self.sample_rate = sample_rate
+        self.hop_length = n_fft//4 if hop_length==None else hop_length
+        self.win_length = n_fft if win_length==None else win_length
+        self.eps = eps
+        if window=="hann":
+            self.window = torch.hann_window(window_length=self.win_length, periodic=True).to(device)
+        elif window=="hamming":
+            self.window = torch.hamming_window(window_length=self.win_length, periodic=True).to(device)
+        self.melscale_transform = torchaudio.transforms.MelScale(n_mels=self.n_mels, \
+                                                                    sample_rate=self.sample_rate, \
+                                                                    n_stft=self.n_fft//2+1).to(device)
+
+
+    def forward(self, input, mode):
+        # front-end function which channel-wise combines all demanded features
+        # input shape : batch x channel x raw waveform
+        # output shape : batch x channel x frequency x time
+        phase_output = None
+
+        front_output_list = []
+        for cur_mode in mode:
+            # Real & Imaginary
+            if cur_mode=="cplx":
+                if self.channel=="mono":
+                    output = torch.stft(input, n_fft=self.n_fft, hop_length=self.hop_length, win_length=self.win_length, window=self.window)
+                elif self.channel=="stereo":
+                    output_l = torch.stft(input[:,0], n_fft=self.n_fft, hop_length=self.hop_length, win_length=self.win_length, window=self.window)
+                    output_r = torch.stft(input[:,1], n_fft=self.n_fft, hop_length=self.hop_length, win_length=self.win_length, window=self.window)
+                    output = torch.cat((output_l, output_r), axis=-1)
+                if input.shape[-1] % round(self.n_fft/4) == 0:
+                    output = output[:, :, :-1]
+                if self.n_fft % 2 == 0:
+                    output = output[:, :-1]
+                front_output_list.append(output.permute(0, 3, 1, 2))
+            # Magnitude & Phase or Mel
+            elif "mag" in cur_mode or "mel" in cur_mode:
+                if self.channel=="mono":
+                    cur_cplx = torch.stft(input, n_fft=self.n_fft, hop_length=self.hop_length, win_length=self.win_length, window=self.window, return_complex=True)
+                    output = self.mag(cur_cplx).unsqueeze(-1)[..., 0:1]
+                    if "mag_phase" in cur_mode:
+                        phase = self.phase(cur_cplx)
+                    if "mel" in cur_mode:
+                        output = self.melscale_transform(output.squeeze(-1)).unsqueeze(-1)
+                elif self.channel=="stereo":
+                    cplx_l = torch.stft(input[:,0], n_fft=self.n_fft, hop_length=self.hop_length, win_length=self.win_length, window=self.window, return_complex=True)
+                    cplx_r = torch.stft(input[:,1], n_fft=self.n_fft, hop_length=self.hop_length, win_length=self.win_length, window=self.window, return_complex=True)
+                    mag_l = self.mag(cplx_l).unsqueeze(-1)
+                    mag_r = self.mag(cplx_r).unsqueeze(-1)
+                    output = torch.cat((mag_l, mag_r), axis=-1)
+                    if "mag_phase" in cur_mode:
+                        phase_l = self.phase(cplx_l).unsqueeze(-1)
+                        phase_r = self.phase(cplx_r).unsqueeze(-1)
+                        output = torch.cat((mag_l, phase_l, mag_r, phase_r), axis=-1)
+                    if "mel" in cur_mode:
+                        output = torch.cat((self.melscale_transform(mag_l.squeeze(-1)).unsqueeze(-1), self.melscale_transform(mag_r.squeeze(-1)).unsqueeze(-1)), axis=-1)
+
+                if "log" in cur_mode:
+                    output = torch.log(output+self.eps)
+
+                if input.shape[-1] % round(self.n_fft/4) == 0:
+                    output = output[:, :, :-1]
+                if cur_mode!="mel" and self.n_fft % 2 == 0: # discard highest frequency
+                    output = output[:, 1:]
+                front_output_list.append(output.permute(0, 3, 1, 2))
+
+        # combine all demanded features
+        if not front_output_list:
+            raise NameError("NameError at FrontEnd: check using features for front-end")
+        elif len(mode)!=1:
+            for i, cur_output in enumerate(front_output_list):
+                if i==0:
+                    front_output = cur_output
+                else:
+                    front_output = torch.cat((front_output, cur_output), axis=1)
+        else:
+            front_output = front_output_list[0]
+            
+        return front_output
+
+
+    def mag(self, cplx_input, eps=1e-07):
+        # mag_summed = cplx_input.pow(2.).sum(-1) + eps
+        mag_summed = cplx_input.real.pow(2.) + cplx_input.imag.pow(2.) + eps
+        return mag_summed.pow(0.5)
+
+
+    def phase(self, cplx_input, ):
+        return torch.atan2(cplx_input.imag, cplx_input.real)
+        # return torch.angle(cplx_input)
+
+
+
+# Multi-Scale Spectral Loss proposed at the paper "DDSP: DIFFERENTIABLE DIGITAL SIGNAL PROCESSING" (https://arxiv.org/abs/2001.04643)
+#   we extend this loss by applying it to mid/side channels
+class MultiScale_Spectral_Loss_MidSide_DDSP(nn.Module):
+    def __init__(self, mode='midside', \
+                        reduce=True, \
+                        n_filters=None, \
+                        windows_size=None, \
+                        hops_size=None, \
+                        window="hann", \
+                        eps=1e-7, \
+                        device=torch.device("cpu")):
+        super(MultiScale_Spectral_Loss_MidSide_DDSP, self).__init__()
+        self.mode = mode
+        self.eps = eps
+        self.mid_weight = 0.5   # value in the range of 0.0 ~ 1.0
+        self.logmag_weight = 0.1
+
+        if n_filters is None:
+            n_filters = [4096, 2048, 1024, 512]
+        if windows_size is None:
+            windows_size = [4096, 2048, 1024, 512]
+        if hops_size is None:
+            hops_size = [1024, 512, 256, 128]
+
+        self.multiscales = []
+        for i in range(len(windows_size)):
+            cur_scale = {'window_size' : float(windows_size[i])}
+            if self.mode=='midside':
+                cur_scale['front_end'] = FrontEnd(channel='mono', \
+                                                    n_fft=n_filters[i], \
+                                                    hop_length=hops_size[i], \
+                                                    win_length=windows_size[i], \
+                                                    window=window, \
+                                                    device=device)
+            elif self.mode=='ori':
+                cur_scale['front_end'] = FrontEnd(channel='stereo', \
+                                                    n_fft=n_filters[i], \
+                                                    hop_length=hops_size[i], \
+                                                    win_length=windows_size[i], \
+                                                    window=window, \
+                                                    device=device)
+            self.multiscales.append(cur_scale)
+
+        self.reduce=reduce
+        self.objective_l1 = nn.L1Loss(reduce=reduce)
+        self.objective_l2 = nn.MSELoss(reduce=reduce)
+
+
+    def forward(self, est_targets, targets):
+        if self.mode=='midside':
+            return self.forward_midside(est_targets, targets)
+        elif self.mode=='ori':
+            return self.forward_ori(est_targets, targets)
+
+
+    def forward_ori(self, est_targets, targets):
+        if self.reduce:
+            total_mag_loss = 0.0
+            total_logmag_loss = 0.0
+        else:
+            total_mag_loss=torch.zeros(est_targets.shape[0], 1).to(est_targets.device)
+            total_logmag_loss=torch.zeros(est_targets.shape[0], 1).to(est_targets.device)
+        for cur_scale in self.multiscales:
+            est_mag = cur_scale['front_end'](est_targets, mode=["mag"])
+            tgt_mag = cur_scale['front_end'](targets, mode=["mag"])
+
+            mag_loss = self.magnitude_loss(est_mag, tgt_mag)
+            logmag_loss = self.log_magnitude_loss(est_mag, tgt_mag)
+            if self.reduce:
+                total_mag_loss += mag_loss
+                total_logmag_loss += logmag_loss
+            else:
+                total_logmag_loss += logmag_loss.mean((1, 2, 3)).unsqueeze(-1)
+                total_mag_loss += mag_loss.mean((1, 2, 3)).unsqueeze(-1)
+        # return total_loss
+        return (1-self.logmag_weight)*total_mag_loss + \
+                (self.logmag_weight)*total_logmag_loss
+
+
+    def forward_midside(self, est_targets, targets):
+        est_mid, est_side = self.to_mid_side(est_targets)
+        tgt_mid, tgt_side = self.to_mid_side(targets)
+        if self.reduce:
+            total_mag_loss = 0.0
+            total_logmag_loss = 0.0
+        else:
+            total_logmag_loss=torch.zeros(est_targets.shape[0], 1).to(est_targets.device)
+            total_mag_loss=torch.zeros(est_targets.shape[0], 1).to(est_targets.device)
+
+        for cur_scale in self.multiscales:
+            est_mid_mag = cur_scale['front_end'](est_mid, mode=["mag"])
+            est_side_mag = cur_scale['front_end'](est_side, mode=["mag"])
+            tgt_mid_mag = cur_scale['front_end'](tgt_mid, mode=["mag"])
+            tgt_side_mag = cur_scale['front_end'](tgt_side, mode=["mag"])
+
+            mag_loss = self.mid_weight*self.magnitude_loss(est_mid_mag, tgt_mid_mag) + \
+                        (1-self.mid_weight)*self.magnitude_loss(est_side_mag, tgt_side_mag)
+            logmag_loss = self.mid_weight*self.log_magnitude_loss(est_mid_mag, tgt_mid_mag) + \
+                        (1-self.mid_weight)*self.log_magnitude_loss(est_side_mag, tgt_side_mag)
+
+                #take mean over all dimensions except batch
+            if self.reduce:
+                total_mag_loss += mag_loss
+                total_logmag_loss += logmag_loss
+            else:
+                total_mag_loss += mag_loss.mean((1, 2, 3)).unsqueeze(-1)
+                #mean over dims 1, 2, 3
+                total_logmag_loss += logmag_loss.mean((1, 2, 3)).unsqueeze(-1)
+        # return total_loss
+        return (1-self.logmag_weight)*total_mag_loss + \
+                (self.logmag_weight)*total_logmag_loss
+
+
+    def to_mid_side(self, stereo_in):
+        mid = stereo_in[:,0] + stereo_in[:,1]
+        side = stereo_in[:,0] - stereo_in[:,1]
+        return mid, side
+
+
+    def magnitude_loss(self, est_mag_spec, tgt_mag_spec):
+        if self.reduce:
+            return torch.norm(self.objective_l1(est_mag_spec, tgt_mag_spec))
+        else:
+            return self.objective_l1(est_mag_spec, tgt_mag_spec)
+
+
+    def log_magnitude_loss(self, est_mag_spec, tgt_mag_spec):
+        est_log_mag_spec = torch.log10(est_mag_spec+self.eps)
+        tgt_log_mag_spec = torch.log10(tgt_mag_spec+self.eps)
+        return self.objective_l2(est_log_mag_spec, tgt_log_mag_spec)
+
+
+

utils/collators.py

@@ -0,0 +1,262 @@
+
+
+import torch
+import torch.nn.functional as F
+import torchaudio
+
+def collate_multitrack_sim(batch, max_tracks=None):
+
+        x= [ data_i[0] for data_i in batch ]  # x is a list of tensors, each tensor is a track
+        clusters=[ torch.tensor(data_i[1]) for data_i in batch ]  # cluster is a list of tensors, each tensor is a cluster
+        taxonomies=[ data_i[2] for data_i in batch ]  # taxonomy is a
+
+        if max_tracks is None:
+            max_tracks = max(track.shape[0] for track in x)  # Find the maximum number of tracks in the batch
+        else:
+            if max_tracks > max(track.shape[0] for track in x):
+                print(f"Warning: max_tracks is set to {max_tracks}, but the maximum number of tracks in the batch is {max(track.shape[0] for track in x)}. I dont know what will happen, consider increasing max_tracks," )
+
+        # Pad each examples with zeros to the maximum length (dimension 0)
+
+        padded_x = []
+        padded_taxonomies = []
+        masks = torch.zeros((len(x), max_tracks), dtype=torch.bool)  # Create a mask tensor
+
+        for i in range(len(x)):
+            current_x = x[i]
+            current_taxonomies = taxonomies[i]
+            
+            # Get current number of tracks
+            current_tracks = current_x.shape[0]
+
+            masks[i, :current_tracks] = 1  # Set mask for current tracks
+            
+            if current_tracks < max_tracks:
+                # Pad dimension N (first dimension)
+                # For a tensor of shape [N, C, L], we need to pad the first dimension
+                # F.pad expects padding as (pad_left, pad_right, pad_top, pad_bottom, pad_front, pad_back)
+                # To pad the first dimension, we need to specify padding for the third-from-last dimension
+                pad_size = (0, 0,  # last dim (L): no padding
+                            0, 0,  # second-to-last dim (C): no padding
+                            0, max_tracks - current_tracks)  # third-to-last dim (N): pad at the end
+
+                padded_x.append(F.pad(current_x, pad_size))
+            
+                # Pad taxonomies
+                padded_taxonomies.append(current_taxonomies + [None] * (max_tracks - len(current_taxonomies)))
+            elif current_tracks > max_tracks :
+                raise ValueError(f"Number of tracks {current_tracks} exceeds maximum allowed {max_tracks}. that is impossible")
+            else:
+                padded_x.append(current_x)
+                padded_taxonomies.append(current_taxonomies[:max_tracks])
+            
+
+
+        x_stacked = torch.stack(padded_x, dim=0)  # Shape: [B, max_tracks, C, L]
+
+        clusters_stacked = torch.stack(clusters, dim=0)  # Shape: [B, max_tracks]
+
+        return {
+            'x': x_stacked,
+            'clusters': clusters_stacked,
+            'taxonomies': padded_taxonomies,
+            "masks": masks
+        }
+
+def collate_multitrack(batch, max_tracks=None, sample_rate=None, segment_length=None, device=None):
+
+        x= [ data_i[0] for data_i in batch ]  # x is a list of tensors, each tensor is a track
+
+        paths=[ data_i[1] for data_i in batch ]  # paths is a list of paths to the audio files, each path is a string
+        fs= [ data_i[2] for data_i in batch ]  # fs is a list of sample rates, each sample rate is an integer
+
+        if max_tracks is None:
+            max_tracks = max(track.shape[0] for track in x)  # Find the maximum number of tracks in the batch
+        else:
+            if max_tracks < max(track.shape[0] for track in x):
+                print(f"Warning: max_tracks is set to {max_tracks}, but the maximum number of tracks in the batch is {max(track.shape[0] for track in x)}. I will crop the last tracks.  I hope the order is random..." )
+
+                for i in range(len(x)):
+                    if x[i].shape[0] > max_tracks:
+                        print("cropping x[i] to max_tracks")
+                        x[i] = x[i][:max_tracks]
+
+        is_x_none= any(x_i is None for x_i in x)
+
+        assert not (is_x_none ), "Either x or y should be None, but not both. This is a bug in the collator"
+
+        #now resample audio to 44100 Hz and cut to segment length
+        for i in range(len(x)):
+            x[i] = x[i].to(device)  # Move to device
+
+            if fs[i] != sample_rate:
+                if fs[i] == 48000 and sample_rate == 44100:
+                    #print(f"Resampling audio from {fs[i]} Hz to {sample_rate} Hz")
+                    x[i]=torchaudio.functional.resample(x[i], orig_freq=160, new_freq=147)
+                else:
+                    #print(f"Resampling audio from {fs[i]} Hz to {sample_rate} Hz")
+                    x[i]=torchaudio.functional.resample(x[i], fs[i], sample_rate)
+
+            assert segment_length is not None, "segment_length should be set to the length of the audio segment in samples"
+            if x[i].shape[-1] > segment_length:
+                x[i] = x[i][..., :segment_length]
+            elif x[i].shape[-1] < segment_length:
+                raise ValueError(f"Audio length {x[i].shape[-1]} is less than segment length {segment_length}. Please check your data.")
+
+        # Pad each examples with zeros to the maximum length (dimension 0)
+        padded_x = []
+        masks = torch.zeros((len(x), max_tracks), dtype=torch.bool)  # Create a mask tensor
+
+
+        for i in range(len(x)):
+            if not is_x_none:
+                current_x = x[i]
+            
+            # Get current number of tracks
+            current_tracks = current_x.shape[0]
+
+            masks[i, :current_tracks] = 1  # Set mask for current tracks
+            
+            if current_tracks < max_tracks:
+                # Pad dimension N (first dimension)
+                # For a tensor of shape [N, C, L], we need to pad the first dimension
+                # F.pad expects padding as (pad_left, pad_right, pad_top, pad_bottom, pad_front, pad_back)
+                # To pad the first dimension, we need to specify padding for the third-from-last dimension
+                pad_size = (0, 0,  # last dim (L): no padding
+                            0, 0,  # second-to-last dim (C): no padding
+                            0, max_tracks - current_tracks)  # third-to-last dim (N): pad at the end
+
+
+                if not is_x_none:
+                    padded_x.append(F.pad(current_x, pad_size))
+                
+            elif current_tracks > max_tracks :
+                raise ValueError(f"Number of tracks {current_tracks} exceeds maximum allowed {max_tracks}. that is impossible")
+            else:
+                if not is_x_none:
+                    padded_x.append(current_x)
+
+
+        if not is_x_none:
+            x_stacked = torch.stack(padded_x, dim=0)  # Shape: [B, max_tracks, C, L]
+
+        return {
+            'y': x_stacked.to(device),  # Shape: [B, max_tracks, C, L]
+            "masks": masks.to(device),  # Shape: [B, max_tracks]
+            "paths": paths,
+            "fs": fs
+        }
+
+def collate_multitrack_paired(batch, max_tracks=None, sample_rate=None, segment_length=None, device=None):
+
+        x= [ data_i[0] for data_i in batch ]  # x is a list of tensors, each tensor is a track
+        y= [ data_i[1] for data_i in batch ]  # x is a list of tensors, each tensor is a track
+
+        taxonomies=[ data_i[2] for data_i in batch ]  # taxonomy is a
+
+        paths=[ data_i[3] for data_i in batch ]  # paths is a list of paths to the audio files, each path is a string
+
+        if max_tracks is None:
+            max_tracks = max(track.shape[0] for track in x)  # Find the maximum number of tracks in the batch
+        else:
+            if max_tracks < max(track.shape[0] for track in x):
+                print(f"Warning: max_tracks is set to {max_tracks}, but the maximum number of tracks in the batch is {max(track.shape[0] for track in x)}. I will crop the last tracks.  I hope the order is random..." )
+
+                for i in range(len(x)):
+                    if x[i].shape[0] > max_tracks:
+                        print("cropping x[i] to max_tracks")
+                        x[i] = x[i][:max_tracks]
+                    if y[i].shape[0] > max_tracks:
+                        print("cropping y[i] to max_tracks")
+                        y[i] = y[i][:max_tracks]
+                    if len(taxonomies[i]) > max_tracks:
+                        print("cropping taxonomies[i] to max_tracks")
+                        taxonomies[i] = taxonomies[i][:max_tracks]
+
+
+
+        is_x_none= any(x_i is None for x_i in x)
+        is_y_none= any(y_i is None for y_i in y)
+
+        assert not (is_x_none and is_y_none), "Either x or y should be None, but not both. This is a bug in the collator"
+
+        # Pad each examples with zeros to the maximum length (dimension 0)
+        padded_x = []
+        padded_y = []
+        padded_taxonomies = []
+        masks = torch.zeros((len(x), max_tracks), dtype=torch.bool)  # Create a mask tensor
+
+        for i in range(len(x)):
+            
+            if not is_x_none:
+                current_x = x[i]
+            if not is_y_none:
+                current_y = y[i]    
+            current_taxonomies = taxonomies[i]
+            
+            # Get current number of tracks
+            if is_x_none:
+                current_tracks = current_y.shape[0]
+            else:
+                current_tracks = current_x.shape[0]
+                if not is_y_none:
+                    assert current_tracks == current_y.shape[0], f"Number of tracks in x ({current_tracks}) does not match number of tracks in y ({current_y.shape[0]})"
+
+            masks[i, :current_tracks] = 1  # Set mask for current tracks
+            
+            if current_tracks < max_tracks:
+                # Pad dimension N (first dimension)
+                # For a tensor of shape [N, C, L], we need to pad the first dimension
+                # F.pad expects padding as (pad_left, pad_right, pad_top, pad_bottom, pad_front, pad_back)
+                # To pad the first dimension, we need to specify padding for the third-from-last dimension
+                pad_size = (0, 0,  # last dim (L): no padding
+                            0, 0,  # second-to-last dim (C): no padding
+                            0, max_tracks - current_tracks)  # third-to-last dim (N): pad at the end
+
+
+                if not is_x_none:
+                    padded_x.append(F.pad(current_x, pad_size))
+                if not is_y_none:
+                    padded_y.append(F.pad(current_y, pad_size))
+            
+                # Pad taxonomies
+                padded_taxonomies.append(current_taxonomies + [None] * (max_tracks - len(current_taxonomies)))
+            elif current_tracks > max_tracks :
+                raise ValueError(f"Number of tracks {current_tracks} exceeds maximum allowed {max_tracks}. that is impossible")
+            else:
+                if not is_x_none:
+                    padded_x.append(current_x)
+                if not is_y_none:
+                    padded_y.append(current_y)
+
+                padded_taxonomies.append(current_taxonomies[:max_tracks])
+            
+
+        if not is_x_none:
+            x_stacked = torch.stack(padded_x, dim=0)  # Shape: [B, max_tracks, C, L]
+
+        if not is_y_none:
+            y_stacked = torch.stack(padded_y, dim=0)  # Shape: [B, max_tracks, C, L]
+
+        if is_x_none:
+            return {
+                'y': y_stacked,
+                'taxonomies': padded_taxonomies,
+                "masks": masks,
+                "paths": paths
+            }
+        if is_y_none:
+            return {
+                'x': x_stacked,
+                'taxonomies': padded_taxonomies,
+                "masks": masks,
+                "paths": paths
+            }
+
+        return {
+            'x': x_stacked,
+            'y': y_stacked,
+            'taxonomies': padded_taxonomies,
+            "masks": masks,
+            "paths": paths
+        }

utils/common_audioeffects.py

@@ -0,0 +1,1729 @@
+"""
+Audio effects for data augmentation.
+
+Several audio effects can be combined into an augmentation chain.
+
+Important note: We assume that the parallelization during training is done using
+                multi-processing and not multi-threading. Hence, we do not need the
+                `@sox.sox_context()` decorators as discussed in this
+                [thread](https://github.com/pseeth/soxbindings/issues/4).
+
+Section 2, TL21
+AI Speech and Sound Group, SL1
+"""
+
+from itertools import permutations
+import os
+import io
+import functools
+import lameenc
+import logging
+import numpy as np
+import pymixconsole as pymc
+from pymixconsole.parameter import Parameter
+from pymixconsole.parameter_list import ParameterList
+from pymixconsole.processor import Processor
+from random import shuffle
+from scipy.signal import oaconvolve
+import soundfile as sf
+#import soxbindings as sox
+from typing import List, Optional, Tuple, Union, Dict
+from numba import jit
+
+#from common_dataprocessing import sample_data
+
+# prevent pysox from logging warnings regarding non-opimal timestretch factors
+logging.getLogger('sox').setLevel(logging.ERROR)
+
+# set maximum peak value if we pass a signal through SOX
+MAX_SOX_PROCESSING_PEAK = 0.707
+
+
+def convert_audio2data(x):
+    """
+    Convert audio data from the format it was stored in to float32.
+
+    Args:
+        x (Numpy array): input with `x.dtype` either `np.int16`, `np.int32`, `np.float32` or `np.float64`.
+
+    Returns:
+        Numpy array: output with values in [-1., 1.) where `dtype` is `np.float32`.
+    """
+    if x.dtype in [np.float32, np.float64]:
+        return x.astype(dtype=np.float32)
+    else:
+        return (x.astype(dtype=np.float64) / (1. + np.iinfo(x.dtype).max)).astype(np.float32)
+
+
+
+def sample_data(data: Tuple[int, Union[np.ndarray, functools.partial]],
+                start: int = 0, length: Optional[int] = None) -> np.ndarray:
+    """
+    Load one stem specified by `data`.
+
+    Returns the audio beginning from `start` either up to the end (if `length` is None)
+    or until the provided `length`. For the case that `start + length > n_samples`, we do a wrap-around and
+    load the remaining samples from the beginning of `data`.
+
+    Args:
+        data: Data with shape (n_samples, data).
+        start: Start index.
+        length: Length of sample. If `length` is not None, `length` samples are returned (possibly with a wrap-around).
+            Otherwise, everything until the end of `data` is returned.
+
+    Returns:
+        samples: data with shape `n_samples x n_channels`
+    """
+    n_samples, audio = data
+
+    # determine whether we have to load the audio or whether it was already loaded
+    is_loaded = True if type(audio) is np.ndarray else False
+
+    # if `length` is not None, then only select subset
+    do_wrap_around = False
+    if length is not None:
+        if start + length > n_samples:
+            # we need to wrap around and concatenate `start:` and `:stop`
+            do_wrap_around = True
+            stop = length - (n_samples - start)
+        else:
+            # no wrap around as it is inside the array/file boundaries
+            stop = start + length
+    else:
+        stop = None
+
+    if is_loaded:
+        if not do_wrap_around:
+            samples = convert_audio2data(audio[start:stop])
+        else:
+            samples = np.vstack((convert_audio2data(audio[start:]),
+                                 convert_audio2data(audio[:stop])))
+    else:
+        if not do_wrap_around:
+            samples = convert_audio2data(audio(start=start, stop=stop)[0])
+        else:
+            samples = np.vstack((convert_audio2data(audio(start=start)[0]),
+                                 convert_audio2data(audio(stop=stop)[0])))
+
+    return samples
+
+
+# Monkey-Patch `Processor` for convenience
+# (a) Allow `None` as blocksize if processor can work on variable-length audio
+def new_init(self, name, parameters, block_size, sample_rate=None, normalize=None, dtype='float32'):
+    """
+    Initialize processor.
+
+    Args:
+        self: Reference to object
+        name (str): Name of processor.
+        parameters (parameter_list): Parameters for this processor.
+        block_size (int): Size of blocks for blockwise processing.
+            Can also be `None` if full audio can be processed at once.
+        sample_rate (int): Sample rate of input audio. Use `None` if effect is independent of this value.
+        normalize (str): Defines, whether the processed signal is normalized.
+            Possible values are `'rms'`, `'max'` and `None`.
+        dtype (str): data type of samples
+
+    Raises:
+        ValueError: If `normalize` is not equal to `'rms'`, `'max'` or `False`.
+    """
+    self.name = name
+    self.parameters = parameters
+    self.block_size = block_size
+    self.sample_rate = sample_rate
+    self.dtype = dtype
+    if normalize not in [None, 'rms', 'max']:
+        raise ValueError(f'Unknown value {normalize} for `normalize`. Must be either `rms`, `max` or `None`')
+    self.normalize = normalize
+
+
+# (b) make code simpler
+def new_update(self, parameter_name):
+    """
+    Update processor after randomization of parameters.
+
+    Args:
+        self: Reference to object.
+        parameter_name (str): Parameter whose value has changed.
+    """
+    pass
+
+
+# (c) representation for nice print
+def new_repr(self):
+    """
+    Create human-readable representation.
+
+    Args:
+        self: Reference to object.
+
+    Returns:
+        string representation of object.
+    """
+    return f'Processor(name={self.name!r}, parameters={self.parameters!r}'
+
+
+Processor.__init__ = new_init
+Processor.__repr__ = new_repr
+Processor.update = new_update
+
+
+class AugmentationChain:
+    """Basic audio Fx chain which is used for data augmentation."""
+
+    def __init__(self,
+                 fxs: Optional[List[Tuple[Union[Processor, 'AugmentationChain'], float]]] = [],
+                 shuffle: Optional[bool] = False,
+                 apply_to: Optional[str] = 'both'):
+        """
+        Create augmentation chain from the dictionary `fxs`.
+
+        Args:
+            fxs (list of tuples): Each tuple has three elements:
+                First tuple element is an instance of `pymc.processor` or `AugmentationChain` that
+                we want to use for data augmentation.
+                Second element gives probability that effect should be applied.
+            shuffle (bool): If `True` then order of Fx are changed whenever chain is applied.
+            apply_to (str): Apply the chain to both input and target or one of them only.
+                Possible values are `'both'`, `'input'` and `target`.
+
+        Raises:
+            ValueError: If `apply_to` is not equal to `both`, `input` or `target`.
+        """
+        self.fxs = fxs
+        self.shuffle = shuffle
+        if apply_to not in ['both', 'input', 'target']:
+            raise ValueError(f'Unknown value {apply_to} for `apply_to`. Must be `both`, `input` or `target`')
+        else:
+            self.apply_to = apply_to
+
+    def apply_processor(self, x, processor: Processor):
+        """
+        Pass audio in `x` through `processor` and output the respective processed audio.
+
+        Args:
+            x (Numpy array): Input audio of shape `n_samples` x `n_channels`.
+            processor (Processor): Audio effect that we want to apply.
+
+        Returns:
+            Numpy array: Processed audio of shape `n_samples` x `n_channels` (same size as `x')
+        """
+        n_samples_input = x.shape[0]
+
+        if processor.block_size is None:
+            y = processor.process(x)
+        else:
+            # make sure that n_samples is a multiple of `processor.block_size`
+            if x.shape[0] % processor.block_size != 0:
+                n_pad = processor.block_size - x.shape[0] % processor.block_size
+                x = np.pad(x, ((0, n_pad), (0, 0)), mode='reflective')
+
+            y = np.zeros_like(x)
+            for idx in range(0, x.shape[0], processor.block_size):
+                y[idx:idx+processor.block_size, :] = processor.process(x[idx:idx+processor.block_size, :])
+
+        if processor.normalize is not None:
+            if processor.normalize == 'rms':  # normalize output energy such that it is the same as the input energy
+                scale = np.sqrt(np.mean(np.square(x)) / np.maximum(1e-7, np.mean(np.square(y))))
+            elif processor.normalize == 'max':  # normalize output signal by its max. amplitude
+                scale = (1 + 1e-7)/(np.max(np.abs(y)) + 1e-7)
+            y *= scale
+
+        # return audio of same length as x
+        return y[:n_samples_input, :]
+
+    def __call__(self, input_x, target_x):
+        """
+        Apply augmentation chain to audio in `input_x` and `target_x`.
+
+        Args:
+            input_x (Numpy array): Audio samples of shape `n_samples` x `n_channels`.
+            target_x (Numpy array): Audio samples of shape `n_samples` x `n_channels`.
+
+        Returns:
+            input_y (Numpy array): Processed audio of same shape as `input_x` where effects have been applied.
+            target_y (Numpy array): Processed audio of same shape as `target_x` where effects have been applied.
+        """
+        # randomly shuffle effect order if `self.shuffle` is True
+        if self.shuffle:
+            shuffle(self.fxs)
+
+        input_y = input_x
+        target_y = target_x
+
+        # check whether we only need to process once later
+        if self.apply_to == 'both' and np.allclose(input_y, target_y):
+            is_input_equal_target = True
+        else:
+            is_input_equal_target = False
+
+        # apply effects with probabilities given in `self.fxs`
+        for fx, p in self.fxs:
+            if np.random.rand() < p:
+                if isinstance(fx, Processor):
+                    # randomize all effect parameters (also calls `update()` for each processor)
+                    fx.randomize()
+                    # apply processor dependent on `apply_to`
+                    if self.apply_to == 'both':
+                        input_y = self.apply_processor(input_y, fx)
+                        if not is_input_equal_target:
+                            target_y = self.apply_processor(target_y, fx)
+                    elif self.apply_to == 'input':
+                        input_y = self.apply_processor(input_y, fx)
+                    elif self.apply_to == 'target':
+                        target_y = self.apply_processor(target_y, fx)
+                else:
+                    # apply effect chain
+                    if is_input_equal_target:
+                        target_y = input_y
+                    input_y, target_y = fx(input_y, target_y)
+                    # check whether input and target are still the same
+                    if not fx.apply_to == 'both':
+                        is_input_equal_target = False
+
+        if is_input_equal_target:
+            target_y = input_y
+
+        return input_y, target_y
+
+    def __repr__(self):
+        """
+        Human-readable representation.
+
+        Returns:
+            string representation of object.
+        """
+        return f'AugmentationChain(fxs={self.fxs!r}, shuffle={self.shuffle!r})'
+
+
+# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% DISTORTION %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
+def hard_clip(x, threshold_dB, drive):
+    """
+    Hard clip distortion.
+
+    Args:
+        x: input audio
+        threshold_dB: threshold
+        drive: drive
+
+    Returns:
+        (Numpy array): distorted audio
+    """
+    drive_linear = np.power(10., drive / 20.).astype(np.float32)
+    threshold_linear = 10. ** (threshold_dB / 20.)
+    return np.clip(x * drive_linear, -threshold_linear, threshold_linear)
+
+
+def overdrive(x, drive, colour, sample_rate):
+    """
+    Overdrive distortion.
+
+    Args:
+        x: input audio
+        drive: Controls the amount of distortion (dB).
+        colour: Controls the amount of even harmonic content in the output(dB)
+        sample_rate: sampling rate
+
+    Returns:
+        (Numpy array): distorted audio
+    """
+    scale = np.max(np.abs(x))
+    if scale > MAX_SOX_PROCESSING_PEAK:
+        clips = True
+        x = x * (MAX_SOX_PROCESSING_PEAK / scale)
+    else:
+        clips = False
+
+    tfm = sox.Transformer()
+    tfm.overdrive(gain_db=drive, colour=colour)
+    y = tfm.build_array(input_array=x, sample_rate_in=sample_rate).astype(np.float32)
+
+    if clips:
+        y *= scale / MAX_SOX_PROCESSING_PEAK  # rescale output to original scale
+    return y
+
+
+def hyperbolic_tangent(x, drive):
+    """
+    Hyperbolic Tanh distortion.
+
+    Args:
+        x: input audio
+        drive: drive
+
+    Returns:
+        (Numpy array): distorted audio
+    """
+    drive_linear = np.power(10., drive / 20.).astype(np.float32)
+    return np.tanh(2. * x * drive_linear)
+
+
+def soft_sine(x, drive):
+    """
+    Soft sine distortion.
+
+    Args:
+        x: input audio
+        drive: drive
+
+    Returns:
+        (Numpy array): distorted audio
+    """
+    drive_linear = np.power(10., drive / 20.).astype(np.float32)
+    y = np.clip(x * drive_linear, -np.pi/4.0, np.pi/4.0)
+    return np.sin(2. * y)
+
+
+def bit_crusher(x, bits):
+    """
+    Bit crusher distortion.
+
+    Args:
+        x: input audio
+        bits: bits
+
+    Returns:
+        (Numpy array): distorted audio
+    """
+    return np.rint(x * (2 ** bits)) / (2 ** bits)
+
+
+class Distortion(Processor):
+    """
+    Distortion processor.
+
+    Processor parameters:
+        sample_rate (int): Current (assumed) sample rate of the audio.
+        mode (str): Currently supports the following five modes: hard_clip, waveshaper, soft_sine, tanh, bit_crusher.
+            Each mode has different parameters such as threshold, factor, or bits.
+        threshold (float): threshold
+        drive (float): drive
+        factor (float): factor
+        limit_range (float): limit range
+        bits (int): bits
+    """
+
+    def __init__(self, sample_rates, name='Distortion', parameters=None, **kwargs):
+        """
+        Initialize processor.
+
+        Args:
+            sample_rates (list of ints): sample rates of audio.
+            name (str): Name of processor.
+            parameters (parameter_list): Parameters for this processor.
+            **kwargs (dict): Keyword arguments that are passed to the Processor class.
+        """
+        super().__init__(name, None, block_size=None, **kwargs)
+        if not parameters:
+            self.parameters = ParameterList()
+            self.parameters.add(Parameter('sample_rate', 44100, 'string', options=sample_rates))
+            self.parameters.add(Parameter('mode', 'hard_clip', 'string',
+                                          options=['hard_clip',
+                                                   'overdrive',
+                                                   'soft_sine',
+                                                   'tanh',
+                                                   'bit_crusher']))
+            self.parameters.add(Parameter('threshold', 0.0, 'float',
+                                          units='dB', maximum=0.0, minimum=-20.0))
+            self.parameters.add(Parameter('drive', 0.0, 'float',
+                                          units='dB', maximum=20.0, minimum=0.0))
+            self.parameters.add(Parameter('colour', 20.0, 'float',
+                                          maximum=100.0, minimum=0.0))
+            self.parameters.add(Parameter('bits', 12, 'int',
+                                          maximum=12, minimum=8))
+        else:
+            self.parameters = parameters
+
+    def process(self, x):
+        """
+        Process audio.
+
+        Args:
+            x (Numpy array): input audio of size `n_samples x n_channels`.
+
+        Returns:
+            (Numpy array): distorted audio of size `n_samples x n_channels`.
+        """
+        if self.parameters.mode.value == 'hard_clip':
+            y = hard_clip(x, self.parameters.threshold.value, self.parameters.drive.value)
+        elif self.parameters.mode.value == 'overdrive':
+            y = overdrive(x, self.parameters.drive.value,
+                          self.parameters.colour.value, self.parameters.sample_rate.value)
+        elif self.parameters.mode.value == 'soft_sine':
+            y = soft_sine(x, self.parameters.drive.value)
+        elif self.parameters.mode.value == 'tanh':
+            y = hyperbolic_tangent(x, self.parameters.drive.value)
+        elif self.parameters.mode.value == 'bit_crusher':
+            y = bit_crusher(x, self.parameters.bits.value)
+
+        return y
+
+
+# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% EQUALISER %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
+class Equaliser(Processor):
+    """
+    Five band parametric equaliser (two shelves and three central bands).
+
+    All gains are set in dB values and range from `MIN_GAIN` dB to `MAX_GAIN` dB.
+    This processor is implemented as cascade of five biquad IIR filters
+    that are implemented using the infamous cookbook formulae from RBJ.
+
+    Processor parameters:
+        sample_rate (int): Current (assumed) sample rate of the audio.
+        low_shelf_gain (float), low_shelf_freq (float)
+        first_band_gain (float), first_band_freq (float), first_band_q (float)
+        second_band_gain (float), second_band_freq (float), second_band_q (float)
+        third_band_gain (float), third_band_freq (float), third_band_q (float)
+
+    original from https://github.com/csteinmetz1/pymixconsole/blob/master/pymixconsole/processors/equaliser.py
+    """
+
+    def __init__(self, n_channels, sample_rates, gain_range=(-15.0, 15.0), q_range=(0.1, 2.0), hard_clip=False,
+                 name='Equaliser', parameters=None, **kwargs):
+        """
+        Initialize processor.
+
+        Args:
+            n_channels (int): Number of audio channels.
+            sample_rates (list of ints): Sample rates of audio.
+            gain_range (tuple of floats): minimum and maximum gain that can be used.
+            q_range (tuple of floats): minimum and maximum q value.
+            hard_clip (bool): Whether we clip to [-1, 1.] after processing.
+            name (str): Name of processor.
+            parameters (parameter_list): Parameters for this processor.
+            **kwargs (dict): Keyword arguments that are passed to the Processor class.
+        """
+        super().__init__(name, parameters=parameters, block_size=None, **kwargs)
+
+        self.n_channels = n_channels
+
+        MIN_GAIN, MAX_GAIN = gain_range
+        MIN_Q, MAX_Q = q_range
+
+        if not parameters:
+            self.parameters = ParameterList()
+            self.parameters.add(Parameter('sample_rate', 44100, 'string', options=sample_rates))
+            # low shelf parameters -------
+            self.parameters.add(Parameter('low_shelf_gain', 0.0, 'float', minimum=MIN_GAIN, maximum=MAX_GAIN))
+            self.parameters.add(Parameter('low_shelf_freq', 80.0, 'float', minimum=30.0, maximum=200.0))
+            # first band parameters ------
+            self.parameters.add(Parameter('first_band_gain', 0.0, 'float', minimum=MIN_GAIN, maximum=MAX_GAIN))
+            self.parameters.add(Parameter('first_band_freq', 400.0, 'float', minimum=200.0, maximum=1000.0))
+            self.parameters.add(Parameter('first_band_q', 0.7, 'float', minimum=MIN_Q, maximum=MAX_Q))
+            # second band parameters -----
+            self.parameters.add(Parameter('second_band_gain', 0.0, 'float', minimum=MIN_GAIN, maximum=MAX_GAIN))
+            self.parameters.add(Parameter('second_band_freq', 2000.0, 'float', minimum=1000.0, maximum=3000.0))
+            self.parameters.add(Parameter('second_band_q', 0.7, 'float', minimum=MIN_Q, maximum=MAX_Q))
+            # third band parameters ------
+            self.parameters.add(Parameter('third_band_gain', 0.0, 'float', minimum=MIN_GAIN, maximum=MAX_GAIN))
+            self.parameters.add(Parameter('third_band_freq', 4000.0, 'float', minimum=3000.0, maximum=8000.0))
+            self.parameters.add(Parameter('third_band_q', 0.7, 'float', minimum=MIN_Q, maximum=MAX_Q))
+            # high shelf parameters ------
+            self.parameters.add(Parameter('high_shelf_gain', 0.0, 'float', minimum=MIN_GAIN, maximum=MAX_GAIN))
+            self.parameters.add(Parameter('high_shelf_freq', 8000.0, 'float', minimum=5000.0, maximum=10000.0))
+        else:
+            self.parameters = parameters
+
+        self.bands = ['low_shelf', 'first_band', 'second_band', 'third_band', 'high_shelf']
+        self.filters = self.setup_filters()
+        self.hard_clip = hard_clip
+
+    def setup_filters(self):
+        """
+        Create IIR filters.
+
+        Returns:
+            IIR filters
+        """
+        filters = {}
+
+        for band in self.bands:
+
+            G = getattr(self.parameters, band + '_gain').value
+            fc = getattr(self.parameters, band + '_freq').value
+            rate = self.parameters.sample_rate.value
+
+            if band in ['low_shelf', 'high_shelf']:
+                Q = 0.707
+                filter_type = band
+            else:
+                Q = getattr(self.parameters, band + '_q').value
+                filter_type = 'peaking'
+
+            filters[band] = pymc.components.iirfilter.IIRfilter(G, Q, fc, rate, filter_type, n_channels=self.n_channels)
+
+        return filters
+
+    def update_filter(self, band):
+        """
+        Update filters.
+
+        Args:
+            band (str): Band that should be updated.
+        """
+        self.filters[band].G = getattr(self.parameters, band + '_gain').value
+        self.filters[band].fc = getattr(self.parameters, band + '_freq').value
+        self.filters[band].rate = self.parameters.sample_rate.value
+
+        if band in ['first_band', 'second_band', 'third_band']:
+            self.filters[band].Q = getattr(self.parameters, band + '_q').value
+
+    def update(self, parameter_name=None):
+        """
+        Update processor after randomization of parameters.
+
+        Args:
+            parameter_name (str): Parameter whose value has changed.
+        """
+        if parameter_name is not None:
+            bands = ['_'.join(parameter_name.split('_')[:2])]
+        else:
+            bands = self.bands
+
+        for band in bands:
+            self.update_filter(band)
+
+    def process(self, x):
+        """
+        Process audio.
+
+        Args:
+            x (Numpy array): input audio of size `n_samples x n_channels`.
+
+        Returns:
+            (Numpy array): equalized audio of size `n_samples x n_channels`.
+        """
+        for _band, iirfilter in self.filters.items():
+            iirfilter.reset_state()
+            x = iirfilter.apply_filter(x)
+
+        if self.hard_clip:
+            x = np.clip(x, -1.0, 1.0)
+
+        # make sure that we have float32 as IIR filtering returns float64
+        x = x.astype(np.float32)
+
+        # make sure that we have two dimensions (if `n_channels == 1`)
+        if x.ndim == 1:
+            x = x[:, np.newaxis]
+
+        return x
+
+
+# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% COMPRESSOR %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
+@jit(nopython=True)
+def compressor_process(x, threshold, attack_time, release_time, ratio, makeup_gain, sample_rate):
+    """
+    Apply compressor.
+
+    Args:
+        x (Numpy array): audio data.
+        threshold: threshold in dB.
+        attack_time: attack_time in ms.
+        release_time: release_time in ms.
+        ratio: ratio.
+        makeup_gain: makeup_gain.
+        sample_rate: sample rate.
+
+    Returns:
+        compressed audio.
+    """
+    M = x.shape[0]
+    x_g = np.zeros(M)
+    x_l = np.zeros(M)
+    y_g = np.zeros(M)
+    y_l = np.zeros(M)
+    c = np.zeros(M)
+    yL_prev = 0.
+
+    alpha_attack = np.exp(-1/(0.001 * sample_rate * attack_time))
+    alpha_release = np.exp(-1/(0.001 * sample_rate * release_time))
+
+    for i in np.arange(M):
+        if np.abs(x[i]) < 0.000001:
+            x_g[i] = -120.0
+        else:
+            x_g[i] = 20 * np.log10(np.abs(x[i]))
+
+        if x_g[i] >= threshold:
+            y_g[i] = threshold + (x_g[i] - threshold) / ratio
+        else:
+            y_g[i] = x_g[i]
+
+        x_l[i] = x_g[i] - y_g[i]
+
+        if x_l[i] > yL_prev:
+            y_l[i] = alpha_attack * yL_prev + (1 - alpha_attack) * x_l[i]
+        else:
+            y_l[i] = alpha_release * yL_prev + (1 - alpha_release) * x_l[i]
+
+        c[i] = np.power(10.0, (makeup_gain - y_l[i]) / 20.0)
+        yL_prev = y_l[i]
+
+    y = x * c
+
+    return y
+
+
+class Compressor(Processor):
+    """
+    Single band stereo dynamic range compressor.
+
+    Processor parameters:
+        sample_rate (int): Current (assumed) sample rate of the audio.
+        threshold (float)
+        attack_time (float)
+        release_time (float)
+        ratio (float)
+        makeup_gain (float)
+    """
+
+    def __init__(self, sample_rates, name='Compressor', parameters=None, **kwargs):
+        """
+        Initialize processor.
+
+        Args:
+            sample_rates (list of ints): Sample rates of input audio.
+            name (str): Name of processor.
+            parameters (parameter_list): Parameters for this processor.
+            **kwargs (dict): Keyword arguments that are passed to the Processor class.
+        """
+        super().__init__(name=name, parameters=parameters, block_size=None, **kwargs)
+
+        if not parameters:
+            self.parameters = ParameterList()
+            self.parameters.add(Parameter('sample_rate', 44100, 'string', options=sample_rates))
+            self.parameters.add(Parameter('threshold', 0.0, 'float', units='dB', minimum=-40.0, maximum=0.0))
+            self.parameters.add(Parameter('attack_time', 2.0, 'float', units='ms', minimum=0.03, maximum=30.0))
+            self.parameters.add(Parameter('release_time', 50.0, 'float', units='ms', minimum=50.0, maximum=100.0))
+            self.parameters.add(Parameter('ratio', 2.0, 'float', minimum=2.0, maximum=10.0))
+            self.parameters.add(Parameter('makeup_gain', 0.0, 'float', units='dB', minimum=-3.0, maximum=6.0))
+        else:
+            self.parameters = parameters
+
+    def process(self, x):
+        """
+        Process audio.
+
+        Args:
+            x (Numpy array): input audio of size `n_samples x n_channels`.
+
+        Returns:
+            (Numpy array): compressed audio of size `n_samples x n_channels`.
+        """
+        if not self.parameters.threshold.value == 0.0:
+            y = np.zeros_like(x)
+
+            for ch in range(x.shape[1]):
+                y[:, ch] = compressor_process(x[:, ch],
+                                              self.parameters.threshold.value,
+                                              self.parameters.attack_time.value,
+                                              self.parameters.release_time.value,
+                                              self.parameters.ratio.value,
+                                              self.parameters.makeup_gain.value,
+                                              self.parameters.sample_rate.value)
+        else:
+            y = x
+
+        return y
+
+
+# %%%%%%%%%%%%%%%%%%%%%%%%%% CONVOLUTIONAL REVERB %%%%%%%%%%%%%%%%%%%%%%%%%%%%%
+class ConvolutionalReverb(Processor):
+    """
+    Convolutional Reverb.
+
+    Important: Due to convolving the audio sequence with some impulse response, we should ignore the
+               first/last samples of the augmented audio sequence using `config['AUGMENTER_PADDING']`.
+
+    Processor parameters:
+        sample_rate (int): Current (assumed) sample rate of the audio.
+        wet_dry (float): Wet/dry ratio.
+        decay (float): Applies a fade out to the impulse response.
+        pre_delay (float): Value in ms. Shifts the IR in time and allows.
+            A positive value produces a traditional delay between the dry signal and the wet.
+            A negative delay is, in reality, zero delay, but effectively trims off the start of IR,
+            so the reverb response begins at a point further in.
+    """
+
+    def __init__(self, impulse_responses, sample_rates, name='ConvolutionalReverb', parameters=None, **kwargs):
+        """
+        Initialize processor.
+
+        Args:
+            impulse_responses (list): List with impulse responses created by `common_dataprocessing.create_dataset`
+            sample_rates (list of ints): Sample rates that we should assume (used for fade-out computation)
+            name (str): Name of processor.
+            parameters (parameter_list): Parameters for this processor.
+            **kwargs (dict): Keyword arguments that are passed to the Processor class.
+
+        Raises:
+            ValueError: if no impulse responses are provided.
+        """
+        super().__init__(name=name, parameters=parameters, block_size=None, **kwargs)
+
+        if impulse_responses is None:
+            raise ValueError('List of impulse responses must be provided for ConvolutionalReverb processor.')
+        self.impulse_responses = impulse_responses
+
+        if not parameters:
+            self.parameters = ParameterList()
+            self.parameters.add(Parameter('sample_rate', 44100, 'string', options=sample_rates))
+            self.parameters.add(Parameter('index', 0, 'int', minimum=0, maximum=len(impulse_responses)))
+            self.parameters.add(Parameter('wet_dry', 1.0, 'float', minimum=0.1, maximum=1.0))
+            self.parameters.add(Parameter('decay', 1.0, 'float', minimum=0.1, maximum=1.0))
+            self.parameters.add(Parameter('pre_delay', 0, 'int', units='ms', minimum=-100, maximum=100))
+        else:
+            self.parameters = parameters
+
+    def update(self, parameter_name=None):
+        """
+        Update processor after randomization of parameters.
+
+        Args:
+            parameter_name (str): Parameter whose value has changed.
+        """
+        # copy IR from current index (to avoid modifying it in-place)
+        self.h = np.copy(sample_data(self.impulse_responses.loc[self.parameters.index.value]['impulse_response']))
+
+        # fade out the impulse based on the decay setting (starting from peak value) - constant 20ms fade-out
+        if self.parameters.decay.value < 1.:
+            idx_peak = np.argmax(np.max(np.abs(self.h), axis=1), axis=0)
+            fstart = np.minimum(self.h.shape[0],
+                                idx_peak + int(self.parameters.decay.value * (self.h.shape[0] - idx_peak)))
+            fstop = np.minimum(self.h.shape[0], fstart + int(0.020*self.parameters.sample_rate.value))
+            flen = fstop - fstart
+
+            fade = np.arange(1, flen+1, dtype=self.dtype)/flen
+            fade = np.power(0.1, fade * 5)
+            self.h[fstart:fstop, :] *= fade[:, np.newaxis]
+            self.h = self.h[:fstop]
+
+    def process(self, x):
+        """
+        Process audio.
+
+        Args:
+            x (Numpy array): input audio of size `n_samples x n_channels`.
+
+        Returns:
+            (Numpy array): reverbed audio of size `n_samples x n_channels`.
+        """
+        # reshape IR to the correct size
+        n_channels = x.shape[1]
+        if self.h.shape[1] == 1 and n_channels > 1:
+            self.h = np.hstack([self.h] * n_channels)  # repeat mono IR for multi-channel input
+        if self.h.shape[1] > 1 and n_channels == 1:
+            self.h = self.h[:, np.random.randint(self.h.shape[1]), np.newaxis]  # randomly choose one IR channel
+
+        if self.parameters.wet_dry.value == 0.0:
+            return x
+        else:
+            # perform convolution to get wet signal
+            y = oaconvolve(x, self.h, mode='full', axes=0)
+
+            # cut out wet signal (compensating for the delay that the IR is introducing + predelay)
+            idx = np.argmax(np.max(np.abs(self.h), axis=1), axis=0)
+            idx += int(0.001 * self.parameters.pre_delay.value * self.parameters.sample_rate.value)
+
+            idx = np.clip(idx, 0, self.h.shape[0]-1)
+
+            y = y[idx:idx+x.shape[0], :]
+
+            # return weighted sum of dry and wet signal
+            return (1.0 - self.parameters.wet_dry.value) * x + self.parameters.wet_dry.value * y
+
+
+# %%%%%%%%%%%%%%%%%%%%%%%%%%%%% HAAS EFFECT %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
+def haas_process(x, delay, feedback, wet_channel):
+    """
+    Add Haas effect to audio.
+
+    Args:
+        x (Numpy array): input audio.
+        delay: Delay that we apply to one of the channels (in samples).
+        feedback: Feedback value.
+        wet_channel: Which channel we process (`left` or `right`).
+
+    Returns:
+        (Numpy array): Audio with Haas effect.
+    """
+    y = np.copy(x)
+    if wet_channel == 'left':
+        y[:, 0] += feedback * np.roll(x[:, 0], delay)
+    elif wet_channel == 'right':
+        y[:, 1] += feedback * np.roll(x[:, 1], delay)
+
+    return y
+
+
+class Haas(Processor):
+    """
+    Haas Effect Processor.
+
+    Randomly selects one channel and applies a short delay to it.
+
+    Important: This audio effect uses `np.roll` to perform the shift of one channel. Hence, you should use
+               `config['AUGMENTER_PADDING']` to ignore the first samples of the augmented audio sequence.
+
+    Processor parameters:
+        sample_rate (int): Current (assumed) sample rate of the audio.
+        delay (float)
+        feedback (float)
+        wet_channel (string)
+    """
+
+    def __init__(self, sample_rates, delay_range=(-0.040, 0.040), name='Haas', parameters=None, **kwargs):
+        """
+        Initialize processor.
+
+        Args:
+            sample_rates (list of ints): Sample rates of input audio.
+            delay_range (tuple of floats): minimum/maximum delay in milliseconds for Haas effect.
+            name (str): Name of processor.
+            parameters (parameter_list): Parameters for this processor.
+            **kwargs (dict): Keyword arguments that are passed to the Processor class.
+        """
+        super().__init__(name=name, parameters=parameters, block_size=None, **kwargs)
+
+        if not parameters:
+            self.parameters = ParameterList()
+            self.parameters.add(Parameter('sample_rate', 44100, 'string', options=sample_rates))
+            self.parameters.add(Parameter('delay', delay_range[1], 'float', units='ms',
+                                          minimum=delay_range[0], maximum=delay_range[1]))
+            self.parameters.add(Parameter('feedback', 0.35, 'float', minimum=0.33, maximum=0.66))
+            self.parameters.add(Parameter('wet_channel', 'left', 'string', options=['left', 'right']))
+        else:
+            self.parameters = parameters
+
+    def process(self, x):
+        """
+        Process audio.
+
+        Args:
+            x (Numpy array): input audio of size `n_samples x n_channels`.
+
+        Returns:
+            (Numpy array): audio with Haas effect of size `n_samples x n_channels`.
+        """
+        assert x.shape[1] == 1 or x.shape[1] == 2, 'Haas effect only works with monaural or stereo audio.'
+
+        if x.shape[1] < 2:
+            x = np.repeat(x, 2, axis=1)
+
+        y = haas_process(x, int(self.parameters.delay.value * self.parameters.sample_rate.value),
+                         self.parameters.feedback.value, self.parameters.wet_channel.value)
+
+        return y
+
+
+# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%% PANNER %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
+class Panner(Processor):
+    """
+    Simple stereo panner (adjusting amplitude and delay).
+
+    If input is mono, output is stereo.
+    Original edited from https://github.com/csteinmetz1/pymixconsole/blob/master/pymixconsole/processors/panner.py
+
+    Important: This audio effect uses `np.roll` to perform the shift of one channel. Hence, you should use
+               `config['AUGMENTER_PADDING']` to ignore the first samples of the augmented audio sequence.
+
+    Processor parameters:
+        sample_rate (int): Current (assumed) sample rate of the audio.
+        pan (float): Panning angle. Can take values in [0, 1] where `0` corresponds to fully panned to left
+            and `1` corresponds to fully panned to the right.
+        pan_law (str): Pan law to be used for amplitude panning. Can be '-4.5dB', 'linear' or 'constant_power'.
+        pan_mode (str): Scheme that is used for panning. Can be 'amplitude', 'delay' or 'both'.
+        pan_maxdelay (float): Maximum delay if we pan a source fully to the left/right.
+            For example `2. / 343.` is the maximum delay if the microphones are 2 meters apart.
+    """
+
+    def __init__(self, sample_rates, maxdelay_range=(0, 2. / 343.0), name='Panner', parameters=None, **kwargs):
+        """
+        Initialize processor.
+
+        Args:
+            sample_rates (list of ints): Sample rates of input audio.
+            maxdelay_range (tuple of floats): minimum/maximum delay for panning effect. `2. / 343.` corresponds to
+                the maximum delay if we have a microphone array where the microphones are 2 meters apart.
+            name (str): Name of processor.
+            parameters (parameter_list): Parameters for this processor.
+            **kwargs (dict): Keyword arguments that are passed to the Processor class.
+        """
+        # default processor class constructor
+        super().__init__(name=name, parameters=parameters, block_size=None, **kwargs)
+
+        if not parameters:
+            self.parameters = ParameterList()
+            self.parameters.add(Parameter('sample_rate', 44100, 'string', options=sample_rates))
+            self.parameters.add(Parameter('pan', 0.5, 'float', minimum=0.1, maximum=0.9))
+            self.parameters.add(Parameter('pan_law', '-4.5dB', 'string',
+                                          options=['-4.5dB', 'linear', 'constant_power']))
+            self.parameters.add(Parameter('pan_mode', 'amplitude', 'string',
+                                          options=['amplitude', 'delay', 'both']))
+            self.parameters.add(Parameter('pan_maxdelay', (maxdelay_range[0] + maxdelay_range[1]) / 2., 'float',
+                                          units='ms', minimum=maxdelay_range[0], maximum=maxdelay_range[1]))
+        else:
+            self.parameters = parameters
+
+        # setup the coefficents based on default params
+        self.update()
+
+    def _calculate_pan_coefficents(self):
+        """
+        Calculate panning coefficients from the chosen pan law.
+
+        Based on the set pan law determine the gain value
+        to apply for the left and right channel to achieve panning effect.
+        This operates on the assumption that the input channel is mono.
+        The output data will be stereo at the moment, but could be expanded
+        to a higher channel count format.
+        The panning value is in the range [0, 1], where
+        0 means the signal is panned completely to the left, and
+        1 means the signal is apanned copletely to the right.
+
+        Raises:
+            ValueError: `self.parameters.pan_law` is not supported.
+        """
+        self.gains = np.zeros(2, dtype=self.dtype)
+
+        # first scale the linear [0, 1] to [0, pi/2]
+        theta = self.parameters.pan.value * (np.pi/2)
+
+        if self.parameters.pan_law.value == 'linear':
+            self.gains[0] = ((np.pi/2) - theta) * (2/np.pi)
+            self.gains[1] = theta * (2/np.pi)
+        elif self.parameters.pan_law.value == 'constant_power':
+            self.gains[0] = np.cos(theta)
+            self.gains[1] = np.sin(theta)
+        elif self.parameters.pan_law.value == '-4.5dB':
+            self.gains[0] = np.sqrt(((np.pi/2) - theta) * (2/np.pi) * np.cos(theta))
+            self.gains[1] = np.sqrt(theta * (2/np.pi) * np.sin(theta))
+        else:
+            raise ValueError(f'Invalid pan_law {self.parameters.pan_law.value}.')
+
+    def _calculate_pan_delay(self):
+        """Calculate delay for the chosen pan angle."""
+        self.shifts = np.zeros(2, dtype=np.int32)
+
+        # compute overall shift that we need between the two channels
+        pan_maxdelay_samples = self.parameters.pan_maxdelay.value * self.parameters.sample_rate.value
+        shift = 2 * int(pan_maxdelay_samples * np.abs(self.parameters.pan.value - 0.5))
+
+        if self.parameters.pan.value < 0.5:
+            # panning to the left
+            self.shifts[0] = -shift // 2
+            self.shifts[1] = shift - shift // 2
+        else:
+            self.shifts[0] = shift - shift // 2
+            self.shifts[1] = -shift // 2
+
+    def process(self, x):
+        """
+        Process audio.
+
+        Args:
+            x (Numpy array): input audio of size `n_samples x n_channels`.
+
+        Returns:
+            (Numpy array): panned audio of size `n_samples x n_channels`.
+        """
+        assert x.shape[1] == 1 or x.shape[1] == 2, 'Panner only works with monaural or stereo audio.'
+
+        # convert to stereo if signal is monaural
+        if x.shape[1] < 2:
+            x = np.repeat(x, 2, axis=1)
+
+        y = np.copy(x)
+        if self.parameters.pan_mode.value in ['delay', 'both']:
+            y[:, 0] = np.roll(y[:, 0], self.shifts[0])
+            y[:, 1] = np.roll(y[:, 1], self.shifts[1])
+
+        if self.parameters.pan_mode.value in ['amplitude', 'both']:
+            y *= self.gains
+
+        return y
+
+    def update(self, parameter_name=None):
+        """
+        Update processor after randomization of parameters.
+
+        Args:
+            parameter_name (str): Parameter whose value has changed.
+        """
+        self._calculate_pan_coefficents()
+        self._calculate_pan_delay()
+
+
+# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% GAIN %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
+class Gain(Processor):
+    """
+    Gain Processor.
+
+    Applies gain in dB and optionally inverts polarity.
+
+    Processor parameters:
+        gain (float): Gain that should be applied (dB scale).
+        invert (bool): If True, then we also invert the waveform (all channels jointly).
+    """
+
+    def __init__(self, name='Gain', parameters=None, **kwargs):
+        """
+        Initialize processor.
+
+        Args:
+            name (str): Name of processor.
+            parameters (parameter_list): Parameters for this processor.
+            **kwargs (dict): Keyword arguments that are passed to the Processor class.
+        """
+        super().__init__(name, parameters=parameters, block_size=None, **kwargs)
+
+        if not parameters:
+            self.parameters = ParameterList()
+            self.parameters.add(Parameter('gain', 1.0, 'float', units='dB', minimum=-6.0, maximum=6.0))
+            self.parameters.add(Parameter('invert', False, 'bool'))
+        else:
+            self.parameters = parameters
+
+    def process(self, x):
+        """
+        Process audio.
+
+        Args:
+            x (Numpy array): input audio of size `n_samples x n_channels`.
+
+        Returns:
+            (Numpy array): gain-augmented audio of size `n_samples x n_channels`.
+        """
+        gain = 10 ** (self.parameters.gain.value / 20.)
+        if self.parameters.invert.value:
+            gain = -gain
+        return gain * x
+
+
+# %%%%%%%%%%%%%%%%%%%%%%% SIMPLE CHANNEL SWAP %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
+class SwapChannels(Processor):
+    """
+    Swap channels in multi-channel audio.
+
+    Processor parameters:
+        index (int) Selects the permutation that we are using.
+            Please note that "no permutation" is one of the permutations in `self.permutations` at index `0`.
+            Hence, this effect should be applied with probability "1.".
+    """
+
+    def __init__(self, n_channels, name='SwapChannels', parameters=None, **kwargs):
+        """
+        Initialize processor.
+
+        Args:
+            n_channels (int): Number of channels in audio that we want to process.
+            name (str): Name of processor.
+            parameters (parameter_list): Parameters for this processor.
+            **kwargs (dict): Keyword arguments that are passed to the Processor class.
+        """
+        super().__init__(name=name, parameters=parameters, block_size=None, **kwargs)
+
+        self.permutations = tuple(permutations(range(n_channels), n_channels))
+
+        if not parameters:
+            self.parameters = ParameterList()
+            self.parameters.add(Parameter('index', 0, 'int', minimum=0, maximum=len(self.permutations)))
+        else:
+            self.parameters = parameters
+
+    def process(self, x):
+        """
+        Process audio.
+
+        Args:
+            x (Numpy array): input audio of size `n_samples x n_channels`.
+
+        Returns:
+            (Numpy array): channel-swapped audio of size `n_samples x n_channels`.
+        """
+        return x[:, self.permutations[self.parameters.index.value]]
+
+
+# %%%%%%%%%%%%%%%%%%%%%%% Monauralize %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
+class Monauralize(Processor):
+    """
+    Monauralizes audio (i.e., removes spatial information).
+
+    Process parameters:
+        seed_channel (int): channel that we use for overwriting the others.
+            `-1` refers to using the mean over all channels.
+    """
+
+    def __init__(self, n_channels, name='Monauralize', parameters=None, **kwargs):
+        """
+        Initialize processor.
+
+        Args:
+            n_channels (int): Number of channels in audio that we want to process.
+            name (str): Name of processor.
+            parameters (parameter_list): Parameters for this processor.
+            **kwargs (dict): Keyword arguments that are passed to the Processor class.
+        """
+        super().__init__(name=name, parameters=parameters, block_size=None, **kwargs)
+
+        if not parameters:
+            self.parameters = ParameterList()
+            self.parameters.add(Parameter('seed_channel', 0, 'int', minimum=-1, maximum=n_channels))
+        else:
+            self.parameters = parameters
+
+    def process(self, x):
+        """
+        Process audio.
+
+        Args:
+            x (Numpy array): input audio of size `n_samples x n_channels`.
+
+        Returns:
+            (Numpy array): monauralized audio of size `n_samples x n_channels`.
+        """
+        tile_reps = (1, x.shape[1])
+        if self.parameters.seed_channel.value == -1:
+            # use average as seed
+            return np.tile(np.mean(x, axis=1, keepdims=True), tile_reps)
+        else:
+            # use one channel as seed
+            return np.tile(x[:, [self.parameters.seed_channel.value]], tile_reps)
+
+
+# %%%%%%%%%%%%%%%%%%%%%%% ChunkShuffle %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
+class ChunkShuffle(Processor):
+    """
+    Split audio into chunks and randomly re-arrange it.
+
+    Process parameters:
+        n_chunks (int): number of chunks into which we split
+    """
+
+    def __init__(self, name='ChunkShuffle', parameters=None, **kwargs):
+        """
+        Initialize processor.
+
+        Args:
+            name (str): Name of processor.
+            parameters (parameter_list): Parameters for this processor.
+            **kwargs (dict): Keyword arguments that are passed to the Processor class.
+        """
+        super().__init__(name=name, parameters=parameters, block_size=None, **kwargs)
+
+        if not parameters:
+            self.parameters = ParameterList()
+            self.parameters.add(Parameter('n_chunks', 2, 'int', minimum=2, maximum=6))
+        else:
+            self.parameters = parameters
+
+    def process(self, x):
+        """
+        Process audio.
+
+        Args:
+            x (Numpy array): input audio of size `n_samples x n_channels`.
+
+        Returns:
+            (Numpy array): chunk-shuffled audio of size `n_samples x n_channels`.
+        """
+        # determine random chunk boundaries
+        boundaries = np.sort(np.random.randint(x.shape[0], size=self.parameters.n_chunks.value-1))
+
+        # create chunks
+        chunks = np.split(x, boundaries)
+
+        # shuffle them
+        np.random.shuffle(chunks)
+
+        # return chunks in new random order
+        return np.concatenate(chunks)
+
+
+# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%% PITCH SHIFT %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
+class PitchShift(Processor):
+    """
+    Simple pitch shifter using SoX and soxbindings (https://github.com/pseeth/soxbindings).
+
+    Processor parameters:
+        sample_rate (int): Current (assumed) sample rate of the audio.
+        steps (float): Pitch shift as positive/negative semitones
+        quick (bool): If True, this effect will run faster but with lower sound quality.
+    """
+
+    def __init__(self, sample_rates, fix_length=True, name='PitchShift', parameters=None, **kwargs):
+        """
+        Initialize processor.
+
+        Args:
+            sample_rates (list of ints): Sample rates of input audio.
+            fix_length (bool): If True, then output has same length as input.
+            name (str): Name of processor.
+            parameters (parameter_list): Parameters for this processor.
+            **kwargs (dict): Keyword arguments that are passed to the Processor class.
+        """
+        super().__init__(name=name, parameters=parameters, block_size=None, **kwargs)
+
+        if not parameters:
+            self.parameters = ParameterList()
+            self.parameters.add(Parameter('sample_rate', 44100, 'string', options=sample_rates))
+            self.parameters.add(Parameter('steps', 0.0, 'float', minimum=-6., maximum=6.))
+            self.parameters.add(Parameter('quick', False, 'bool'))
+        else:
+            self.parameters = parameters
+
+        self.fix_length = fix_length
+        self.clips = False
+
+    def process(self, x):
+        """
+        Process audio.
+
+        Args:
+            x (Numpy array): input audio of size `n_samples x n_channels`.
+
+        Returns:
+            (Numpy array): pitch-shifted audio of size `n_samples x n_channels`.
+        """
+        if self.parameters.steps.value == 0.0:
+            y = x
+        else:
+            scale = np.max(np.abs(x))
+            if scale > MAX_SOX_PROCESSING_PEAK:
+                clips = True
+                x = x * (MAX_SOX_PROCESSING_PEAK / scale)
+            else:
+                clips = False
+
+            tfm = sox.Transformer()
+            tfm.pitch(self.parameters.steps.value, quick=bool(self.parameters.quick.value))
+            y = tfm.build_array(input_array=x, sample_rate_in=self.parameters.sample_rate.value).astype(np.float32)
+
+            if clips:
+                y *= scale / MAX_SOX_PROCESSING_PEAK  # rescale output to original scale
+
+        if self.fix_length:
+            n_samples_input = x.shape[0]
+            n_samples_output = y.shape[0]
+            if n_samples_input < n_samples_output:
+                idx1 = (n_samples_output - n_samples_input) // 2
+                idx2 = idx1 + n_samples_input
+                y = y[idx1:idx2]
+            elif n_samples_input > n_samples_output:
+                n_pad = n_samples_input - n_samples_output
+                y = np.pad(y, ((n_pad//2, n_pad - n_pad//2), (0, 0)))
+
+        return y
+
+
+# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%% TIME STRETCH %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
+class TimeStretch(Processor):
+    """
+    Simple time stretcher using SoX and soxbindings (https://github.com/pseeth/soxbindings).
+
+    Processor parameters:
+        sample_rate (int): Current (assumed) sample rate of the audio.
+        factor (float): Time stretch factor.
+        quick (bool): If True, this effect will run faster but with lower sound quality.
+        stretch_type (str): Algorithm used for stretching (`tempo` or `stretch`).
+        audio_type (str): Sets which time segments are most optmial when finding
+            the best overlapping points for time stretching.
+    """
+
+    def __init__(self, sample_rates, fix_length=True, name='TimeStretch', parameters=None, **kwargs):
+        """
+        Initialize processor.
+
+        Args:
+            sample_rates (list of ints): Sample rates of input audio.
+            fix_length (bool): If True, then output has same length as input.
+            name (str): Name of processor.
+            parameters (parameter_list): Parameters for this processor.
+            **kwargs (dict): Keyword arguments that are passed to the Processor class.
+        """
+        super().__init__(name=name, parameters=parameters, block_size=None, **kwargs)
+
+        if not parameters:
+            self.parameters = ParameterList()
+            self.parameters.add(Parameter('sample_rate', 44100, 'string', options=sample_rates))
+            self.parameters.add(Parameter('factor', 1.0, 'float', minimum=1/1.33, maximum=1.33))
+            self.parameters.add(Parameter('quick', False, 'bool'))
+            self.parameters.add(Parameter('stretch_type', 'tempo', 'string', options=['tempo', 'stretch']))
+            self.parameters.add(Parameter('audio_type', 'l', 'string', options=['m', 's', 'l']))
+        else:
+            self.parameters = parameters
+
+        self.fix_length = fix_length
+
+    def process(self, x):
+        """
+        Process audio.
+
+        Args:
+            x (Numpy array): input audio of size `n_samples x n_channels`.
+
+        Returns:
+            (Numpy array): time-stretched audio of size `n_samples x n_channels`.
+        """
+        if self.parameters.factor.value == 1.0:
+            y = x
+        else:
+            scale = np.max(np.abs(x))
+            if scale > MAX_SOX_PROCESSING_PEAK:
+                clips = True
+                x = x * (MAX_SOX_PROCESSING_PEAK / scale)
+            else:
+                clips = False
+
+            tfm = sox.Transformer()
+            if self.parameters.stretch_type.value == 'stretch':
+                tfm.stretch(self.parameters.factor.value)
+            elif self.parameters.stretch_type.value == 'tempo':
+                tfm.tempo(1/self.parameters.factor.value,
+                          audio_type=self.parameters.audio_type.value,
+                          quick=bool(self.parameters.quick.value))
+            y = tfm.build_array(input_array=x, sample_rate_in=self.parameters.sample_rate.value).astype(np.float32)
+
+            if clips:
+                y *= scale / MAX_SOX_PROCESSING_PEAK  # rescale output to original scale
+
+        if self.fix_length:
+            n_samples_input = x.shape[0]
+            n_samples_output = y.shape[0]
+            if n_samples_input < n_samples_output:
+                idx1 = (n_samples_output - n_samples_input) // 2
+                idx2 = idx1 + n_samples_input
+                y = y[idx1:idx2]
+            elif n_samples_input > n_samples_output:
+                n_pad = n_samples_input - n_samples_output
+                y = np.pad(y, ((n_pad//2, n_pad - n_pad//2), (0, 0)))
+
+        return y
+
+
+# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%% PLAYBACK SPEED %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
+class PlaybackSpeed(Processor):
+    """
+    Simple playback speed effect using SoX and soxbindings (https://github.com/pseeth/soxbindings).
+
+    Processor parameters:
+        sample_rate (int): Current (assumed) sample rate of the audio.
+        factor (float): Playback speed factor.
+    """
+
+    def __init__(self, sample_rates, fix_length=True, name='PlaybackSpeed', parameters=None, **kwargs):
+        """
+        Initialize processor.
+
+        Args:
+            sample_rates (list of ints): Sample rates of input audio.
+            fix_length (bool): If True, then output has same length as input.
+            name (str): Name of processor.
+            parameters (parameter_list): Parameters for this processor.
+            **kwargs (dict): Keyword arguments that are passed to the Processor class.
+        """
+        super().__init__(name=name, parameters=parameters, block_size=None, **kwargs)
+
+        if not parameters:
+            self.parameters = ParameterList()
+            self.parameters.add(Parameter('sample_rate', 44100, 'string', options=sample_rates))
+            self.parameters.add(Parameter('factor', 1.0, 'float', minimum=1./1.33, maximum=1.33))
+        else:
+            self.parameters = parameters
+
+        self.fix_length = fix_length
+
+    def process(self, x):
+        """
+        Process audio.
+
+        Args:
+            x (Numpy array): input audio of size `n_samples x n_channels`.
+
+        Returns:
+            (Numpy array): resampled audio of size `n_samples x n_channels`.
+        """
+        if self.parameters.factor.value == 1.0:
+            y = x
+        else:
+            scale = np.max(np.abs(x))
+            if scale > MAX_SOX_PROCESSING_PEAK:
+                clips = True
+                x = x * (MAX_SOX_PROCESSING_PEAK / scale)
+            else:
+                clips = False
+
+            tfm = sox.Transformer()
+            tfm.speed(self.parameters.factor.value)
+            y = tfm.build_array(input_array=x, sample_rate_in=self.parameters.sample_rate.value).astype(np.float32)
+
+            if clips:
+                y *= scale / MAX_SOX_PROCESSING_PEAK  # rescale output to original scale
+
+        if self.fix_length:
+            n_samples_input = x.shape[0]
+            n_samples_output = y.shape[0]
+            if n_samples_input < n_samples_output:
+                idx1 = (n_samples_output - n_samples_input) // 2
+                idx2 = idx1 + n_samples_input
+                y = y[idx1:idx2]
+            elif n_samples_input > n_samples_output:
+                n_pad = n_samples_input - n_samples_output
+                y = np.pad(y, ((n_pad//2, n_pad - n_pad//2), (0, 0)))
+
+        return y
+
+
+# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%% BEND %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
+class Bend(Processor):
+    """
+    Simple bend effect using SoX and soxbindings (https://github.com/pseeth/soxbindings).
+
+    Processor parameters:
+        sample_rate (int): Current (assumed) sample rate of the audio.
+        n_bends (int): Number of segments or intervals to pitch shift
+    """
+
+    def __init__(self, sample_rates, pitch_range=(-600, 600), fix_length=True, name='Bend', parameters=None, **kwargs):
+        """
+        Initialize processor.
+
+        Args:
+            sample_rates (list of ints): Sample rates of input audio.
+            pitch_range (tuple of ints): min and max pitch bending ranges in cents
+            fix_length (bool): If True, then output has same length as input.
+            name (str): Name of processor.
+            parameters (parameter_list): Parameters for this processor.
+            **kwargs (dict): Keyword arguments that are passed to the Processor class.
+        """
+        super().__init__(name=name, parameters=parameters, block_size=None, **kwargs)
+
+        if not parameters:
+            self.parameters = ParameterList()
+            self.parameters.add(Parameter('sample_rate', 44100, 'string', options=sample_rates))
+            self.parameters.add(Parameter('n_bends', 2, 'int', minimum=2, maximum=10))
+        else:
+            self.parameters = parameters
+        self.pitch_range_min, self.pitch_range_max = pitch_range
+
+    def process(self, x):
+        """
+        Process audio.
+
+        Args:
+            x (Numpy array): input audio of size `n_samples x n_channels`.
+
+        Returns:
+            (Numpy array): pitch-bended audio of size `n_samples x n_channels`.
+        """
+        n_bends = self.parameters.n_bends.value
+        max_length = x.shape[0] / self.parameters.sample_rate.value
+
+        # Generates random non-overlapping segments
+        delta = 1. / self.parameters.sample_rate.value
+        boundaries = np.sort(delta + np.random.rand(n_bends-1) * (max_length - delta))
+
+        start, end = np.zeros(n_bends), np.zeros(n_bends)
+        start[0] = delta
+        for i, b in enumerate(boundaries):
+            end[i] = b
+            start[i+1] = b
+        end[-1] = max_length
+
+        # randomly sample pitch-shifts in cents
+        cents = np.random.randint(self.pitch_range_min, self.pitch_range_max+1, n_bends)
+
+        # remove segment if cent value is zero or start == end (as SoX does not allow such values)
+        idx_keep = np.logical_and(cents != 0, start != end)
+        n_bends, start, end, cents = sum(idx_keep), start[idx_keep], end[idx_keep], cents[idx_keep]
+
+        scale = np.max(np.abs(x))
+        if scale > MAX_SOX_PROCESSING_PEAK:
+            clips = True
+            x = x * (MAX_SOX_PROCESSING_PEAK / scale)
+        else:
+            clips = False
+
+        tfm = sox.Transformer()
+        tfm.bend(n_bends=int(n_bends), start_times=list(start), end_times=list(end), cents=list(cents))
+        y = tfm.build_array(input_array=x, sample_rate_in=self.parameters.sample_rate.value).astype(np.float32)
+
+        if clips:
+            y *= scale / MAX_SOX_PROCESSING_PEAK  # rescale output to original scale
+
+        return y
+
+
+# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% MASTERING %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
+class Mastering(Processor):
+    """
+    Mastering Processor.
+
+    Models different mastering effects, i.e., ways how we deal with sample values outside [-1, 1].
+
+    Processor parameters:
+        method (str): Method that should be applied. Can take the following values:
+            `none`: Do nothing. E.g., useful if training a network for VST plugin where DAW can also handle peaks
+                larger/smaller than `+1`/`-1`.
+            `scale`: Scale to max-abs peak to avoid clipping.
+            `hardclip`: Apply hardclipping to [-1, 1].
+            `softclip`: Apply softclipping to [-1, 1] using tanh.
+    """
+
+    def __init__(self, name='Mastering', parameters=None, maximum_amplitude=1.0, **kwargs):
+        """
+        Initialize processor.
+
+        Args:
+            name (str): Name of processor.
+            parameters (parameter_list): Parameters for this processor.
+            maximum_amplitude (float): Maximum amplitude after applying the mastering processor.
+            **kwargs (dict): Keyword arguments that are passed to the Processor class.
+        """
+        super().__init__(name, parameters=parameters, block_size=None, **kwargs)
+
+        if not parameters:
+            self.parameters = ParameterList()
+            self.parameters.add(Parameter('method', 'none', 'string',
+                                          options=['none', 'scale', 'hardclip', 'softclip']))
+        else:
+            self.parameters = parameters
+
+        self.maximum_amplitude = maximum_amplitude
+
+    def process(self, x):
+        """
+        Process audio.
+
+        Args:
+            x (Numpy array): input audio of size `n_samples x n_channels`.
+
+        Returns:
+            (Numpy array): mastered audio of size `n_samples x n_channels`.
+        """
+        y = np.copy(x)
+        if self.parameters.method.value == 'none':
+            return y
+        elif self.parameters.method.value == 'scale':
+            maxabs = np.maximum(self.maximum_amplitude, np.max(np.abs(y)))
+            return y / (maxabs / self.maximum_amplitude)
+        elif self.parameters.method.value == 'hardclip':
+            return np.clip(y, -self.maximum_amplitude, self.maximum_amplitude)
+        elif self.parameters.method.value == 'softclip':
+            return self.maximum_amplitude * np.tanh(y)
+
+
+# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%% MP3 COMPRESSION %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
+class MP3Compression(Processor):
+    """
+    Models the effect of an mp3 compression using LAME.
+
+    LAME adds some delay at the beginning - we therefore truncate the first samples
+
+    Processor parameters:
+        sample_rate (int): Current (assumed) sample rate of the audio.
+        quality (int): in the range from 2 ... 7 (2 = highest, 7 = fastest)
+        bitrate (int): supported bitrates are
+            [8, 16, 24, 32, 40, 48, 56, 64, 80, 96, 112, 128, 144, 160, 192, 224, 256, 320]
+    """
+
+    def __init__(self, sample_rates, fix_length=True, name='MP3Compression', parameters=None, **kwargs):
+        """
+        Initialize processor.
+
+        Args:
+            sample_rates (list of ints): Sample rates of input audio.
+            fix_length (bool): If True, then output has same length as input.
+            name (str): Name of processor.
+            parameters (parameter_list): Parameters for this processor.
+            **kwargs (dict): Keyword arguments that are passed to the Processor class.
+        """
+        super().__init__(name=name, parameters=parameters, block_size=None, **kwargs)
+
+        if not parameters:
+            supported_bitrates = [8, 16, 24, 32, 40, 48, 56, 64, 80, 96, 112, 128, 144, 160, 192, 224, 256, 320]
+            self.parameters = ParameterList()
+            self.parameters.add(Parameter('sample_rate', 44100, 'string', options=sample_rates))
+            self.parameters.add(Parameter('quality', 2, 'string', options=[2, 3, 4, 5, 6, 7]))
+            self.parameters.add(Parameter('bitrate', 96, 'string', options=supported_bitrates))
+        else:
+            self.parameters = parameters
+
+        self.fix_length = fix_length
+
+    def process(self, x):
+        """
+        Process audio.
+
+        Args:
+            x (Numpy array): input audio of size `n_samples x n_channels`.
+
+        Returns:
+            (Numpy array): MP3 compressed audio of size `n_samples x n_channels`.
+        """
+        # scale if max-abs peak is larger than `MAX_SOX_PROCESSING_PEAK`
+        scale = np.max(np.abs(x))
+        if scale > MAX_SOX_PROCESSING_PEAK:
+            clips = True
+            x = x * (MAX_SOX_PROCESSING_PEAK / scale)
+        else:
+            clips = False
+
+        # convert to int16
+        x_int = (x * np.iinfo(np.int16).max).astype(np.int16)
+
+        encoder = lameenc.Encoder()
+        encoder.set_bit_rate(self.parameters.bitrate.value)
+        encoder.set_in_sample_rate(self.parameters.sample_rate.value)
+        encoder.set_channels(x.shape[1])
+        encoder.set_quality(self.parameters.quality.value)
+        encoder.silence()
+
+        mp3_data = encoder.encode(x_int.tobytes())
+        mp3_data += encoder.flush()
+
+        memory_file = io.BytesIO(initial_bytes=mp3_data)
+        y, fs = sf.read(memory_file, always_2d=True, dtype=np.float32)
+
+        # resample to original sampling rate if LAME changed it
+        if fs != self.parameters.sample_rate.value:
+            tfm = sox.Transformer()
+            tfm.speed(fs / self.parameters.sample_rate.value)
+            y = tfm.build_array(input_array=y, sample_rate_in=fs).astype(np.float32)
+
+        if clips:
+            y *= scale / MAX_SOX_PROCESSING_PEAK  # rescale output to original scale
+
+        if self.fix_length:
+            # LAME adds some samples at the beginning - remove them here
+            n_samples_input = x.shape[0]
+            n_samples_output = y.shape[0]
+            if n_samples_input < n_samples_output:
+                y = y[-n_samples_input:, :]
+
+        return y
+
+
+def __main__():
+    """
+    Main function for testing the audio effects.
+
+    examples:
+
+    config['AUGMENTER_CHAIN'] = AugmentationChain([(ConvolutionalReverb(impulse_responses=impulse_responses,
+                                                                    sample_rates=config['ACCEPTED_SAMPLING_RATES']), 0.5),
+                                               (Gain(), 0.5),
+                                               (Haas(sample_rates=config['ACCEPTED_SAMPLING_RATES']), 0.5),
+                                               (Panner(sample_rates=config['ACCEPTED_SAMPLING_RATES']), 0.5),
+                                               (SwapChannels(n_channels=config['N_CHANNELS']), 1.),
+                                               (Monauralize(n_channels=config['N_CHANNELS']), 0.5)],
+                                              shuffle=True)
+
+    config['AUGMENTER_CHAIN'] = AugmentationChain([(Gain(), 1.),
+                                               (SwapChannels(n_channels=config['N_CHANNELS']), 1.),
+                                               (Monauralize(n_channels=config['N_CHANNELS']), 0.5)], shuffle=True)
+    """
+    pass
+

utils/data_utils.py

@@ -0,0 +1,190 @@
+import soundfile as sf
+import torchaudio
+
+def read_wav_segment(file_path, start=None, end=None, dtype="float32"):
+    """
+    Reads a specific segment from a .wav file efficiently.
+
+    Args:
+        file_path (str): Path to the .wav file.
+        start (int): Start frame index.
+        end (int): End frame index.
+
+    Returns:
+        numpy.ndarray: Audio data for the specified segment.
+        int: Sample rate of the audio file.
+    """
+    # Open the .wav file
+    if start is None or end is None:
+        data, samplerate = sf.read(file_path, dtype=dtype)
+    else:
+        with sf.SoundFile(file_path) as audio_file:
+            # Read only the required frames
+            audio_file.seek(start)
+            data = audio_file.read(frames=end-start, dtype=dtype)
+            samplerate = audio_file.samplerate
+
+    return data, samplerate
+
+def get_audio_length(file_path):
+    """
+    Retrieves the length of an audio file in seconds and frames.
+
+    Args:
+        file_path (str): Path to the audio file.
+
+    Returns:
+        float: Length of the audio file in seconds.
+        int: Total number of frames in the audio file.
+        int: Sample rate of the audio file.
+    """
+    with sf.SoundFile(file_path) as audio_file:
+        total_frames = len(audio_file)  # Total number of frames
+        samplerate = audio_file.samplerate  # Sample rate
+        duration = total_frames / samplerate  # Duration in seconds
+
+    return duration, total_frames, samplerate
+
+def taxonomy2track(input_class, num_instr=8):
+
+    assert num_instr==8, "num_instr should be 8 for this function, the rest is not implemented yet"
+
+    if input_class is None:
+        return 'unknown'  
+    if num_instr == 8:
+        mapping = {0000: 'other', 1100: 'drums', 1200: 'drums', 1300: 'other', 2000: 'bass', 3000: 'guitar', 4100: 'piano', 4200: 'piano', 4300: 'piano', 4400: 'other', 4500: 'other', 4600: 'other', 4700: 'other', 4900: 'other', 5000: 'brass', 6100: 'strings', 6210: 'brass', 6220: 'brass', 8100: 'guitar', 8200: 'brass', 9000: 'vocals'}
+    else:
+        raise NotImplementedError()
+    
+    code_length = len(str(input_class))
+    if code_length < 4:
+        #pad zeros to the right to make it 4 digits
+        input_class = int(str(input_class) + "0" * (4 - code_length))
+ 
+    class_str = str(input_class)
+    for i in range(len(class_str), 0, -1):
+        general_class = int(class_str[:i] + "0" * (len(class_str) - i))
+        if general_class in mapping:
+            return mapping[general_class]  
+       
+    try:
+        raise ValueError(f"No mapping found for input class {input_class} with num_instr {num_instr}")
+    except ValueError as e:
+        print(f"Error: {e}")
+        return "other"  # Return a default value if no mapping is found
+
+import torch
+def efficient_roll(x, shift, dims=-1):
+    """
+    Efficiently roll tensor elements along a dimension without creating a full copy.
+    
+    Args:
+        x: Input tensor
+        shift: Number of places to roll (negative for left roll)
+        dim: Dimension along which to roll
+    
+    Returns:
+        Rolled tensor view where possible, minimal copy where necessary
+    """
+    if shift == 0:
+        return x
+    
+    # Get the size of the dimension
+    dim_size = x.size(dims)
+    
+    # Handle shift larger than dimension size
+    shift = shift % dim_size
+    if shift < 0:
+        shift += dim_size
+    
+    # Create indices for the roll
+    indices = torch.cat([torch.arange(dim_size-shift, dim_size), 
+                         torch.arange(0, dim_size-shift)])
+    
+    # Use index_select for the roll
+    return torch.index_select(x, dims, indices)
+
+#import loudness
+import numpy as np
+
+def apply_loud_normalization(x, lufs=-23, sample_rate=44100,device=None):
+    """
+    x shaPe: (batch_size, channels, time)
+    """
+
+    in_shape= x.shape
+    if x.ndim != 3:
+        x=x.view(-1, in_shape[-2], in_shape[-1])  # Ensure x is 3D
+
+    B, C, T = x.shape
+
+    if device is None:
+        device = x.device
+
+    x_out = torch.zeros_like(x)
+    #for b in range(B):
+    #    x_i=x[b].cpu().numpy().T
+    #    lufs_in=loudness.integrated_loudness(x_i, sample_rate)
+
+    #    delta_loudness= lufs - lufs_in
+    #    gain=np.power(10, delta_loudness / 20)  # Convert dB to linear gain
+
+    #    x_out[b] = torch.tensor(x_i.T * gain, device=device)
+
+    x=x.view(B* C,1, T)  # Ensure x is 3D
+
+    loudness=torchaudio.functional.loudness(x+1e-5, sample_rate=sample_rate)
+    delta_loudness = lufs - loudness
+    gain= torch.pow(10, delta_loudness / 20)  # Convert dB to linear gain
+    if gain.isnan().any():
+        print("NaN detected in gain, setting to -30 dB")
+        gain = torch.nan_to_num(gain, nan=-30.0)
+
+    x_out = x * gain.view(B * C, 1, 1)  # Apply gain to each channel
+
+    
+    x_out = x_out.view(in_shape)
+
+    return x_out
+
+
+
+from utils.feature_extractors.dsp_features import compute_log_rms_gated_max, compute_crest_factor, compute_stereo_width, compute_stereo_imbalance, compute_log_spread
+
+def apply_RMS_normalization(x, RMS_norm=-25, device=None, use_gate=False):
+        if device is None:
+            device = x.device
+
+        RMS= torch.tensor(RMS_norm, device=device).view(1, 1, 1).repeat(x.shape[0],1,1)  # Use fixed RMS for evaluation
+
+        x_RMS_ref=20*torch.log10(torch.sqrt(torch.mean(x**2, dim=(-1), keepdim=True).mean(dim=-2, keepdim=True)))
+        if use_gate:
+            x_RMS = compute_log_rms_gated_max(x).unsqueeze(-1)
+        else:
+            x_RMS=20*torch.log10(torch.sqrt(torch.mean(x**2, dim=(-1), keepdim=True).mean(dim=-2, keepdim=True)))
+        
+
+        gain= RMS - x_RMS
+        gain_linear = 10 ** (gain / 20 + 1e-6)  # Convert dB gain to linear scale, adding a small value to avoid division by zero
+        x=x* gain_linear
+
+        return x
+
+
+import pyloudnorm as pyln   
+
+def loudness_normalize(audio, target_loudness=-23.0, sample_rate=44100):
+    """
+    Normalize the loudness of the audio to a target level.
+    """
+
+    pylnmeter = pyln.Meter(sample_rate)  # Create a meter for 44100 Hz sampling rate
+
+    audio= np.array(audio, dtype=np.float32).T
+    loudness = pylnmeter.integrated_loudness(audio)
+
+
+    # loudness normalize audio to -12 dB LUFS
+    loudness_normalized_audio = pyln.normalize.loudness(audio, loudness, -14.0)
+
+    return  torch.tensor(loudness_normalized_audio.T, dtype=torch.float32)

utils/dnnlib/__init__.py

@@ -0,0 +1,8 @@
+# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# This work is licensed under a Creative Commons
+# Attribution-NonCommercial-ShareAlike 4.0 International License.
+# You should have received a copy of the license along with this
+# work. If not, see http://creativecommons.org/licenses/by-nc-sa/4.0/
+
+from .util import EasyDict, make_cache_dir_path, call_func_by_name

utils/dnnlib/util.py

@@ -0,0 +1,491 @@
+# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# This work is licensed under a Creative Commons
+# Attribution-NonCommercial-ShareAlike 4.0 International License.
+# You should have received a copy of the license along with this
+# work. If not, see http://creativecommons.org/licenses/by-nc-sa/4.0/
+
+"""Miscellaneous utility classes and functions."""
+
+import ctypes
+import fnmatch
+import importlib
+import inspect
+import numpy as np
+import os
+import shutil
+import sys
+import types
+import io
+import pickle
+import re
+import requests
+import html
+import hashlib
+import glob
+import tempfile
+import urllib
+import urllib.request
+import uuid
+
+from distutils.util import strtobool
+from typing import Any, List, Tuple, Union, Optional
+
+
+# Util classes
+# ------------------------------------------------------------------------------------------
+
+
+class EasyDict(dict):
+    """Convenience class that behaves like a dict but allows access with the attribute syntax."""
+
+    def __getattr__(self, name: str) -> Any:
+        try:
+            return self[name]
+        except KeyError:
+            raise AttributeError(name)
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        self[name] = value
+
+    def __delattr__(self, name: str) -> None:
+        del self[name]
+
+
+class Logger(object):
+    """Redirect stderr to stdout, optionally print stdout to a file, and optionally force flushing on both stdout and the file."""
+
+    def __init__(self, file_name: Optional[str] = None, file_mode: str = "w", should_flush: bool = True):
+        self.file = None
+
+        if file_name is not None:
+            self.file = open(file_name, file_mode)
+
+        self.should_flush = should_flush
+        self.stdout = sys.stdout
+        self.stderr = sys.stderr
+
+        sys.stdout = self
+        sys.stderr = self
+
+    def __enter__(self) -> "Logger":
+        return self
+
+    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
+        self.close()
+
+    def write(self, text: Union[str, bytes]) -> None:
+        """Write text to stdout (and a file) and optionally flush."""
+        if isinstance(text, bytes):
+            text = text.decode()
+        if len(text) == 0: # workaround for a bug in VSCode debugger: sys.stdout.write(''); sys.stdout.flush() => crash
+            return
+
+        if self.file is not None:
+            self.file.write(text)
+
+        self.stdout.write(text)
+
+        if self.should_flush:
+            self.flush()
+
+    def flush(self) -> None:
+        """Flush written text to both stdout and a file, if open."""
+        if self.file is not None:
+            self.file.flush()
+
+        self.stdout.flush()
+
+    def close(self) -> None:
+        """Flush, close possible files, and remove stdout/stderr mirroring."""
+        self.flush()
+
+        # if using multiple loggers, prevent closing in wrong order
+        if sys.stdout is self:
+            sys.stdout = self.stdout
+        if sys.stderr is self:
+            sys.stderr = self.stderr
+
+        if self.file is not None:
+            self.file.close()
+            self.file = None
+
+
+# Cache directories
+# ------------------------------------------------------------------------------------------
+
+_dnnlib_cache_dir = None
+
+def set_cache_dir(path: str) -> None:
+    global _dnnlib_cache_dir
+    _dnnlib_cache_dir = path
+
+def make_cache_dir_path(*paths: str) -> str:
+    if _dnnlib_cache_dir is not None:
+        return os.path.join(_dnnlib_cache_dir, *paths)
+    if 'DNNLIB_CACHE_DIR' in os.environ:
+        return os.path.join(os.environ['DNNLIB_CACHE_DIR'], *paths)
+    if 'HOME' in os.environ:
+        return os.path.join(os.environ['HOME'], '.cache', 'dnnlib', *paths)
+    if 'USERPROFILE' in os.environ:
+        return os.path.join(os.environ['USERPROFILE'], '.cache', 'dnnlib', *paths)
+    return os.path.join(tempfile.gettempdir(), '.cache', 'dnnlib', *paths)
+
+# Small util functions
+# ------------------------------------------------------------------------------------------
+
+
+def format_time(seconds: Union[int, float]) -> str:
+    """Convert the seconds to human readable string with days, hours, minutes and seconds."""
+    s = int(np.rint(seconds))
+
+    if s < 60:
+        return "{0}s".format(s)
+    elif s < 60 * 60:
+        return "{0}m {1:02}s".format(s // 60, s % 60)
+    elif s < 24 * 60 * 60:
+        return "{0}h {1:02}m {2:02}s".format(s // (60 * 60), (s // 60) % 60, s % 60)
+    else:
+        return "{0}d {1:02}h {2:02}m".format(s // (24 * 60 * 60), (s // (60 * 60)) % 24, (s // 60) % 60)
+
+
+def format_time_brief(seconds: Union[int, float]) -> str:
+    """Convert the seconds to human readable string with days, hours, minutes and seconds."""
+    s = int(np.rint(seconds))
+
+    if s < 60:
+        return "{0}s".format(s)
+    elif s < 60 * 60:
+        return "{0}m {1:02}s".format(s // 60, s % 60)
+    elif s < 24 * 60 * 60:
+        return "{0}h {1:02}m".format(s // (60 * 60), (s // 60) % 60)
+    else:
+        return "{0}d {1:02}h".format(s // (24 * 60 * 60), (s // (60 * 60)) % 24)
+
+
+def ask_yes_no(question: str) -> bool:
+    """Ask the user the question until the user inputs a valid answer."""
+    while True:
+        try:
+            print("{0} [y/n]".format(question))
+            return strtobool(input().lower())
+        except ValueError:
+            pass
+
+
+def tuple_product(t: Tuple) -> Any:
+    """Calculate the product of the tuple elements."""
+    result = 1
+
+    for v in t:
+        result *= v
+
+    return result
+
+
+_str_to_ctype = {
+    "uint8": ctypes.c_ubyte,
+    "uint16": ctypes.c_uint16,
+    "uint32": ctypes.c_uint32,
+    "uint64": ctypes.c_uint64,
+    "int8": ctypes.c_byte,
+    "int16": ctypes.c_int16,
+    "int32": ctypes.c_int32,
+    "int64": ctypes.c_int64,
+    "float32": ctypes.c_float,
+    "float64": ctypes.c_double
+}
+
+
+def get_dtype_and_ctype(type_obj: Any) -> Tuple[np.dtype, Any]:
+    """Given a type name string (or an object having a __name__ attribute), return matching Numpy and ctypes types that have the same size in bytes."""
+    type_str = None
+
+    if isinstance(type_obj, str):
+        type_str = type_obj
+    elif hasattr(type_obj, "__name__"):
+        type_str = type_obj.__name__
+    elif hasattr(type_obj, "name"):
+        type_str = type_obj.name
+    else:
+        raise RuntimeError("Cannot infer type name from input")
+
+    assert type_str in _str_to_ctype.keys()
+
+    my_dtype = np.dtype(type_str)
+    my_ctype = _str_to_ctype[type_str]
+
+    assert my_dtype.itemsize == ctypes.sizeof(my_ctype)
+
+    return my_dtype, my_ctype
+
+
+def is_pickleable(obj: Any) -> bool:
+    try:
+        with io.BytesIO() as stream:
+            pickle.dump(obj, stream)
+        return True
+    except:
+        return False
+
+
+# Functionality to import modules/objects by name, and call functions by name
+# ------------------------------------------------------------------------------------------
+
+def get_module_from_obj_name(obj_name: str) -> Tuple[types.ModuleType, str]:
+    """Searches for the underlying module behind the name to some python object.
+    Returns the module and the object name (original name with module part removed)."""
+
+    # allow convenience shorthands, substitute them by full names
+    obj_name = re.sub("^np.", "numpy.", obj_name)
+    obj_name = re.sub("^tf.", "tensorflow.", obj_name)
+
+    # list alternatives for (module_name, local_obj_name)
+    parts = obj_name.split(".")
+    name_pairs = [(".".join(parts[:i]), ".".join(parts[i:])) for i in range(len(parts), 0, -1)]
+
+    # try each alternative in turn
+    for module_name, local_obj_name in name_pairs:
+        try:
+            module = importlib.import_module(module_name) # may raise ImportError
+            get_obj_from_module(module, local_obj_name) # may raise AttributeError
+            return module, local_obj_name
+        except:
+            pass
+
+    # maybe some of the modules themselves contain errors?
+    for module_name, _local_obj_name in name_pairs:
+        try:
+            importlib.import_module(module_name) # may raise ImportError
+        except ImportError:
+            if not str(sys.exc_info()[1]).startswith("No module named '" + module_name + "'"):
+                raise
+
+    # maybe the requested attribute is missing?
+    for module_name, local_obj_name in name_pairs:
+        try:
+            module = importlib.import_module(module_name) # may raise ImportError
+            get_obj_from_module(module, local_obj_name) # may raise AttributeError
+        except ImportError:
+            pass
+
+    # we are out of luck, but we have no idea why
+    raise ImportError(obj_name)
+
+
+def get_obj_from_module(module: types.ModuleType, obj_name: str) -> Any:
+    """Traverses the object name and returns the last (rightmost) python object."""
+    if obj_name == '':
+        return module
+    obj = module
+    for part in obj_name.split("."):
+        obj = getattr(obj, part)
+    return obj
+
+
+def get_obj_by_name(name: str) -> Any:
+    """Finds the python object with the given name."""
+    module, obj_name = get_module_from_obj_name(name)
+    return get_obj_from_module(module, obj_name)
+
+
+def call_func_by_name(*args, func_name: str = None, **kwargs) -> Any:
+    """Finds the python object with the given name and calls it as a function."""
+    assert func_name is not None
+    func_obj = get_obj_by_name(func_name)
+    assert callable(func_obj)
+    return func_obj(*args, **kwargs)
+
+
+def construct_class_by_name(*args, class_name: str = None, **kwargs) -> Any:
+    """Finds the python class with the given name and constructs it with the given arguments."""
+    return call_func_by_name(*args, func_name=class_name, **kwargs)
+
+
+def get_module_dir_by_obj_name(obj_name: str) -> str:
+    """Get the directory path of the module containing the given object name."""
+    module, _ = get_module_from_obj_name(obj_name)
+    return os.path.dirname(inspect.getfile(module))
+
+
+def is_top_level_function(obj: Any) -> bool:
+    """Determine whether the given object is a top-level function, i.e., defined at module scope using 'def'."""
+    return callable(obj) and obj.__name__ in sys.modules[obj.__module__].__dict__
+
+
+def get_top_level_function_name(obj: Any) -> str:
+    """Return the fully-qualified name of a top-level function."""
+    assert is_top_level_function(obj)
+    module = obj.__module__
+    if module == '__main__':
+        module = os.path.splitext(os.path.basename(sys.modules[module].__file__))[0]
+    return module + "." + obj.__name__
+
+
+# File system helpers
+# ------------------------------------------------------------------------------------------
+
+def list_dir_recursively_with_ignore(dir_path: str, ignores: List[str] = None, add_base_to_relative: bool = False) -> List[Tuple[str, str]]:
+    """List all files recursively in a given directory while ignoring given file and directory names.
+    Returns list of tuples containing both absolute and relative paths."""
+    assert os.path.isdir(dir_path)
+    base_name = os.path.basename(os.path.normpath(dir_path))
+
+    if ignores is None:
+        ignores = []
+
+    result = []
+
+    for root, dirs, files in os.walk(dir_path, topdown=True):
+        for ignore_ in ignores:
+            dirs_to_remove = [d for d in dirs if fnmatch.fnmatch(d, ignore_)]
+
+            # dirs need to be edited in-place
+            for d in dirs_to_remove:
+                dirs.remove(d)
+
+            files = [f for f in files if not fnmatch.fnmatch(f, ignore_)]
+
+        absolute_paths = [os.path.join(root, f) for f in files]
+        relative_paths = [os.path.relpath(p, dir_path) for p in absolute_paths]
+
+        if add_base_to_relative:
+            relative_paths = [os.path.join(base_name, p) for p in relative_paths]
+
+        assert len(absolute_paths) == len(relative_paths)
+        result += zip(absolute_paths, relative_paths)
+
+    return result
+
+
+def copy_files_and_create_dirs(files: List[Tuple[str, str]]) -> None:
+    """Takes in a list of tuples of (src, dst) paths and copies files.
+    Will create all necessary directories."""
+    for file in files:
+        target_dir_name = os.path.dirname(file[1])
+
+        # will create all intermediate-level directories
+        if not os.path.exists(target_dir_name):
+            os.makedirs(target_dir_name)
+
+        shutil.copyfile(file[0], file[1])
+
+
+# URL helpers
+# ------------------------------------------------------------------------------------------
+
+def is_url(obj: Any, allow_file_urls: bool = False) -> bool:
+    """Determine whether the given object is a valid URL string."""
+    if not isinstance(obj, str) or not "://" in obj:
+        return False
+    if allow_file_urls and obj.startswith('file://'):
+        return True
+    try:
+        res = requests.compat.urlparse(obj)
+        if not res.scheme or not res.netloc or not "." in res.netloc:
+            return False
+        res = requests.compat.urlparse(requests.compat.urljoin(obj, "/"))
+        if not res.scheme or not res.netloc or not "." in res.netloc:
+            return False
+    except:
+        return False
+    return True
+
+
+def open_url(url: str, cache_dir: str = None, num_attempts: int = 10, verbose: bool = True, return_filename: bool = False, cache: bool = True) -> Any:
+    """Download the given URL and return a binary-mode file object to access the data."""
+    assert num_attempts >= 1
+    assert not (return_filename and (not cache))
+
+    # Doesn't look like an URL scheme so interpret it as a local filename.
+    if not re.match('^[a-z]+://', url):
+        return url if return_filename else open(url, "rb")
+
+    # Handle file URLs.  This code handles unusual file:// patterns that
+    # arise on Windows:
+    #
+    # file:///c:/foo.txt
+    #
+    # which would translate to a local '/c:/foo.txt' filename that's
+    # invalid.  Drop the forward slash for such pathnames.
+    #
+    # If you touch this code path, you should test it on both Linux and
+    # Windows.
+    #
+    # Some internet resources suggest using urllib.request.url2pathname() but
+    # but that converts forward slashes to backslashes and this causes
+    # its own set of problems.
+    if url.startswith('file://'):
+        filename = urllib.parse.urlparse(url).path
+        if re.match(r'^/[a-zA-Z]:', filename):
+            filename = filename[1:]
+        return filename if return_filename else open(filename, "rb")
+
+    assert is_url(url)
+
+    # Lookup from cache.
+    if cache_dir is None:
+        cache_dir = make_cache_dir_path('downloads')
+
+    url_md5 = hashlib.md5(url.encode("utf-8")).hexdigest()
+    if cache:
+        cache_files = glob.glob(os.path.join(cache_dir, url_md5 + "_*"))
+        if len(cache_files) == 1:
+            filename = cache_files[0]
+            return filename if return_filename else open(filename, "rb")
+
+    # Download.
+    url_name = None
+    url_data = None
+    with requests.Session() as session:
+        if verbose:
+            print("Downloading %s ..." % url, end="", flush=True)
+        for attempts_left in reversed(range(num_attempts)):
+            try:
+                with session.get(url) as res:
+                    res.raise_for_status()
+                    if len(res.content) == 0:
+                        raise IOError("No data received")
+
+                    if len(res.content) < 8192:
+                        content_str = res.content.decode("utf-8")
+                        if "download_warning" in res.headers.get("Set-Cookie", ""):
+                            links = [html.unescape(link) for link in content_str.split('"') if "export=download" in link]
+                            if len(links) == 1:
+                                url = requests.compat.urljoin(url, links[0])
+                                raise IOError("Google Drive virus checker nag")
+                        if "Google Drive - Quota exceeded" in content_str:
+                            raise IOError("Google Drive download quota exceeded -- please try again later")
+
+                    match = re.search(r'filename="([^"]*)"', res.headers.get("Content-Disposition", ""))
+                    url_name = match[1] if match else url
+                    url_data = res.content
+                    if verbose:
+                        print(" done")
+                    break
+            except KeyboardInterrupt:
+                raise
+            except:
+                if not attempts_left:
+                    if verbose:
+                        print(" failed")
+                    raise
+                if verbose:
+                    print(".", end="", flush=True)
+
+    # Save to cache.
+    if cache:
+        safe_name = re.sub(r"[^0-9a-zA-Z-._]", "_", url_name)
+        safe_name = safe_name[:min(len(safe_name), 128)]
+        cache_file = os.path.join(cache_dir, url_md5 + "_" + safe_name)
+        temp_file = os.path.join(cache_dir, "tmp_" + uuid.uuid4().hex + "_" + url_md5 + "_" + safe_name)
+        os.makedirs(cache_dir, exist_ok=True)
+        with open(temp_file, "wb") as f:
+            f.write(url_data)
+        os.replace(temp_file, cache_file) # atomic
+        if return_filename:
+            return cache_file
+
+    # Return data as file object.
+    assert not return_filename
+    return io.BytesIO(url_data)

utils/evaluation/KAD_evaluate.py

@@ -0,0 +1,307 @@
+
+import pandas as pd
+
+import torch
+import os
+import omegaconf
+import numpy as np
+import glob
+import pyloudnorm as pyln
+import sys
+#move to the parent directory to import datasets
+sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
+
+from datasets.eval_benchmark import load_audio
+from utils.evaluation.dist_metrics import KADFeatures
+
+# see device 1
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+
+path_test_set="/scratch/elec/t412-asp/automix/MDX_TM_benchmark"
+
+set="MDX_TM"
+
+features=[
+    "AFxRep",
+    "FxEncoder",
+    "FxEncoder++",
+    "CLAP",
+]
+
+
+KAD_metrics={}
+
+KAD_args = { 
+  "do_PCA_figure": False,
+  "do_TSNE_figure": False,
+  "kernel": "gaussian", #kernel to use for the KAD metric
+}
+
+KAD_args = omegaconf.OmegaConf.create(KAD_args)
+
+for feature in features:
+    if feature=="AFxRep":
+        AFxRep_args = {
+              "distance_type": "cosine",  # not used
+              "ckpt_path": "/scratch/work/molinee2/projects/project_mfm_eloi/src/tmp/afx-rep.ckpt"
+        }
+        AFxRep_args = omegaconf.OmegaConf.create(AFxRep_args)
+        KAD_metrics[feature] = KADFeatures(type="AFxRep", sample_rate=44100, AFxRep_args=AFxRep_args, KAD_args=KAD_args)
+    elif feature=="FxEncoder":
+        fx_encoder_args = {
+            "distance_type": "cosine",  # not used
+            "ckpt_path": "/scratch/work/molinee2/projects/project_mfm_eloi/src/utils/feature_extractors/ckpt/fxenc_default.pt"
+        }
+        fx_encoder_args = omegaconf.OmegaConf.create(fx_encoder_args)
+        KAD_metrics[feature] = KADFeatures(type="fx_encoder", sample_rate=44100, fx_encoder_args=fx_encoder_args, KAD_args=KAD_args)
+    elif feature=="FxEncoder++":
+        fx_encoder_plusplus_args = {
+            "distance_type": "cosine",  # not used
+            "ckpt_path": "/scratch/work/molinee2/projects/project_mfm_eloi/src_clean/checkpoints/fxenc_plusplus_default.pt"
+        }
+        fx_encoder_plusplus_args = omegaconf.OmegaConf.create(fx_encoder_plusplus_args)
+        KAD_metrics[feature] = KADFeatures(type="fx_encoder_++", sample_rate=44100, fx_encoder_plusplus_args=fx_encoder_plusplus_args, KAD_args=KAD_args)
+
+    elif feature=="CLAP":
+        clap_args = {
+            "ckpt_path": "/scratch/work/molinee2/projects/project_mfm_eloi/src_clean/checkpoints/music_audioset_epoch_15_esc_90.14.patched.pt",
+            "distance_type": "cosine",
+            "normalize": True,  # if True, the features will be normalized
+            "use_adaptor": False,  # if True, the features will be adapted to the CLAP space
+            "adaptor_checkpoint":None,
+            "adaptor_type":None,
+            "add_noise": False,  #   if True, the features will be augmented with orthogonal noise
+            "noise_sigma": 0  # sigma of the orthogonal noise to
+        }
+
+        clap_args = omegaconf.OmegaConf.create(clap_args)
+        KAD_metrics[feature] = KADFeatures(type="CLAP", sample_rate=44100, CLAP_args=clap_args, KAD_args=KAD_args)
+
+    elif feature=="bark":
+        bark_args = {
+            "distance_type": "cosine",  # not used
+            "normalize": True,  # if True, the features will be normalized
+        }
+        bark_args = omegaconf.OmegaConf.create(bark_args)
+        KAD_metrics[feature] = KADFeatures(type="bark", sample_rate=44100, bark_args=bark_args, KAD_args=KAD_args, normalize=True)
+
+    else:
+        raise ValueError(f"Unknown feature: {feature}")
+
+
+def get_ref_file_path(dir):
+    ref_mixture = os.path.join(dir,"mix", "mix.wav")
+    return ref_mixture
+
+pylnmeter = pyln.Meter(44100)  # Create a meter for 44100 Hz sampling rate
+
+def loudness_normalize(audio, target_loudness=-23.0):
+    """
+    Normalize the loudness of the audio to a target level.
+    """
+    audio= np.array(audio, dtype=np.float32).T
+    loudness = pylnmeter.integrated_loudness(audio)
+
+    # loudness normalize audio to -12 dB LUFS
+    loudness_normalized_audio = pyln.normalize.loudness(audio, loudness, -14.0)
+
+    return  torch.tensor(loudness_normalized_audio.T, dtype=torch.float32)
+
+song_dirs= sorted(glob.glob(os.path.join(path_test_set, "*")))
+
+reference_dict = {}
+
+for song_dir in song_dirs:
+
+    song_name = os.path.basename(song_dir)
+    segment_subdirs = sorted(glob.glob(os.path.join(song_dir, "*")))
+
+    song_id=os.path.basename(song_dir)
+
+    for segment_dir in segment_subdirs:
+
+        segment= os.path.basename(segment_dir)
+        id= f"{song_id}_{segment}"
+
+        ref_file_path= get_ref_file_path(segment_dir)
+        mix_ref, fs=load_audio(str(ref_file_path), stereo=True)
+        assert fs==44100, "Expected sampling rate of 44100 Hz"
+        mix_ref = loudness_normalize(mix_ref)
+        reference_dict[id] = mix_ref
+
+
+# create dataframe colums are features, rows are methods
+dataframe= pd.DataFrame(columns=["method"] +list(KAD_metrics.keys()))
+
+filename_results="results_eval/results_KAD"+f"_{set}_test.csv"
+
+#methods=[ "fxnorm_automix_S_Lb", "fxnorm_automix_L_Lb",]
+#methods=[ "proposed_random_churn","S3_random_churn", "S4_random_churn", "only_rms_S4_random_churn",  "only_rms_random_churn" ]
+#methods=[ "proposed_public", "only_rms_public" ]j
+#methods=[ "S1internal_S2public"]
+#methods=[ "S1internal_S2publicv3", "S1publicv3_S2publicv3",  "S1publicv3_S2internal"]
+#methods=[ "S1public_S2publicv3"]  
+#methods=[ "S1public_S2publicv3",  "S1public_S2publicv3_oracle",   "S1public_S2publicv3_rms"]  
+#methods=["mst_oracle",  "equal_loudness", "diff_baseline",  "proposed_oracle", "proposed_random", "only_rms_random" , "proposed_centroid_close", "proposed_centroid_far", "only_rms_centroid_close", "only_rms_centroid_far"]
+#methods=["S4_random"]
+#methods=["onlyrms_3108"]
+#methods=["publicv3_2"]
+#methods=["mst_oracle",  "equal_loudness", "diff_baseline",  "proposed_oracle", "proposed_random_3108_v2", "prop_random_indep_3108", "S4_random", "S3_random", "WUN_4instr", "publicv3", "DMC_14instr", "fxnorm_automix_S_Lb", "fxnorm_automix_L_Lb"]
+methods=[ "tencydb_test", "publicv3_test"]
+#methods=["prop_random_indep", "only_rms_indep"]
+#methods=["WUN_4instr"]
+#methods=["diff_baseline"]
+#methods=["internal_2708", "equal_loudness", "internal_2708_oracle"]
+#methods=["WUN_4instr", "WUN_14instr", "DMC_14instr"]
+
+
+for method in methods:
+
+        if method=="mst_oracle":
+            def get_pred_file_path(dir):
+                pred_mixture= os.path.join(dir,"mst_oracle", "mixture_output.wav")
+                return pred_mixture
+        elif method=="mst_oracle_multi":
+           def get_pred_file_path(dir):
+               pred_mixture= os.path.join(dir,"mst_oracle_multi", "mixture_output.wav")
+               return pred_mixture
+        elif method=="equal_loudness":
+           def get_pred_file_path(dir):
+               pred_mixture= os.path.join(dir,"anchor_equal_loudness", "mix.wav")
+               return pred_mixture
+        elif method=="diff_baseline":
+           def get_pred_file_path(dir):
+               pred_mixture= os.path.join(dir,"diff_baseline", "pred_mixture.wav")
+               return pred_mixture
+        elif method=="fxnorm_automix_S_Lb":
+           def get_pred_file_path(dir):
+               pred_mixture= os.path.join(dir,"fxnorm_automix_S_Lb_v2", "mixture_output.wav")
+               return pred_mixture
+        elif method=="fxnorm_automix_L_Lb":
+           def get_pred_file_path(dir):
+               pred_mixture= os.path.join(dir,"fxnorm_automix_L_Lb_v2", "mixture_output.wav")
+               return pred_mixture
+        elif method=="proposed_random":
+           def get_pred_file_path(dir):
+               pred_mixture= os.path.join(dir,"stylediffpipeline_S9v6_MF3wetv6_MDX_TM_benchmark_cfg_1_T30", "random.wav")
+               return pred_mixture
+        elif method=="proposed_random_3108_v2":
+           def get_pred_file_path(dir):
+               pred_mixture= os.path.join(dir,"stylediffpipeline_S9v6_MF3wetv6_MDX_TM_benchmark_cfg_1_T30_3108_v2", "random.wav")
+               return pred_mixture
+        elif method=="onlyrms_3108":
+           def get_pred_file_path(dir):
+               pred_mixture= os.path.join(dir,"stylediffpipeline_S9v6_MF3wetv6_MDX_TM_benchmark_cfg_1_T30_3108_v2", "only_rms_random.wav")
+               return pred_mixture
+        elif method=="proposed_random_3108_v2_oracle":
+           def get_pred_file_path(dir):
+               pred_mixture= os.path.join(dir,"stylediffpipeline_S9v6_MF3wetv6_MDX_TM_benchmark_cfg_1_T30_3108_v2_oracle", "random.wav")
+               return pred_mixture
+        elif method=="only_rms_random":
+           def get_pred_file_path(dir):
+               pred_mixture= os.path.join(dir,"stylediffpipeline_S9v6_MF3wetv6_MDX_TM_benchmark_cfg_1_T30", "only_rms_random.wav")
+               return pred_mixture
+        elif method=="S3_random":
+           def get_pred_file_path(dir):
+               pred_mixture= os.path.join(dir,"stylediffpipeline_S3v6_MF3wetv6_MDX_TM_benchmark_cfg_1_T30", "random.wav")
+               return pred_mixture
+        elif method=="S4_3108":
+           def get_pred_file_path(dir):
+               pred_mixture= os.path.join(dir,"stylediffpipeline_S4v6_MF3wetv6_MDX_TM_benchmark_cfg_1_T30_3108_v2", "random.wav")
+               return pred_mixture
+        elif method=="publicv3":
+           def get_pred_file_path(dir):
+               pred_mixture= os.path.join(dir,"stylediffpipeline_publicv3_publicv3_MDX_TM_benchmark_cfg_1_T30_publicv3", "random.wav")
+               return pred_mixture
+        elif method=="publicv3_test":
+           def get_pred_file_path(dir):
+               pred_mixture= os.path.join(dir,"public_public_MDX_TM_benchmark_public_test_Sep29", "random.wav")
+               return pred_mixture
+        elif method=="tencydb_test":
+           def get_pred_file_path(dir):
+               pred_mixture= os.path.join(dir,"internal_TencyDB_internal_TencyMastering_MDX_TM_benchmark_internal_test_Sep29", "random.wav")
+               return pred_mixture
+        elif method=="prop_random_indep_3108":
+           def get_pred_file_path(dir):
+               pred_mixture= os.path.join(dir,"stylediffpipeline_S9v6_MF3wetv6_MDX_TM_benchmark_cfg_1_T30_3108_v2_independent", "random.wav")
+               return pred_mixture
+        elif method=="WUN_4instr":
+            def get_pred_file_path(dir):
+                pred_mixture= os.path.join(dir,"WUN_4instr", "pred_mixture.wav")
+                return pred_mixture
+        elif method=="DMC_14instr":
+            def get_pred_file_path(dir):
+                pred_mixture= os.path.join(dir,"DMC_14instr", "pred_mixture.wav")
+                return pred_mixture
+        elif method=="only_rms_public":
+            def get_pred_file_path(dir):
+               pred_mixture= os.path.join(dir,"ours_public", "mixture_processed_onlyrms.wav")
+               return pred_mixture
+        elif method=="proposed_oracle":
+           def get_pred_file_path(dir):
+               pred_mixture= os.path.join(dir,"oracle_proposed", "mix.wav")
+               return pred_mixture
+        elif method=="publicv3_oracle":
+           def get_pred_file_path(dir):
+               pred_mixture= os.path.join(dir,"oracle_proposed_public", "mix.wav")
+               return pred_mixture
+        else:
+           raise ValueError(f"Unknown method: {method}")
+    
+    
+        method_dict = {}
+    
+        for song_dir in song_dirs:
+    
+            song_name = os.path.basename(song_dir)
+            segment_subdirs = sorted(glob.glob(os.path.join(song_dir, "*")))
+        
+            song_id=os.path.basename(song_dir)
+        
+            for segment_dir in segment_subdirs:
+        
+                segment= os.path.basename(segment_dir)
+                id= f"{song_id}_{segment}"
+        
+                pred_file_path= get_pred_file_path(segment_dir)
+                pred_mixture, fs=load_audio(str(pred_file_path), stereo=True)
+                assert fs==44100, "Expected sampling rate of 44100 Hz"
+                pred_mixture = loudness_normalize(pred_mixture)
+                assert not np.isnan(pred_mixture).any(), f"NaN values found in predicted mixture for {id}"
+                method_dict[id] = pred_mixture
+    
+    
+        #create new empty row with key as "method"
+        dataframe.loc[method] = [None] * (len(KAD_metrics)+1)
+    
+        for feature, metric in KAD_metrics.items():
+    
+            KAD_distance, dict_output=metric.compute(reference_dict, method_dict, None)
+
+            print(f"KAD distance for method {method} and feature {feature}: {KAD_distance}")
+        
+            #write the KAD distance to the dataframe
+            dataframe.at[method, "method"] = method
+            dataframe.at[method, feature] = KAD_distance
+        
+            for output_key, output_value in dict_output.items():
+                if "figure" in output_key:
+                    output_value.savefig(f"results_eval/{output_key}_{set}_{method}_{feature}.png")
+                    print(f"Saved figure: {output_key}_{set}_{method}_{feature}.png")
+    
+    
+    
+        dataframe.to_csv(filename_results, index=False)
+        print(f"Results saved to {filename_results}")
+
+        #except Exception as e:
+
+        #    print(f"Error processing method {method}: {e}")
+
+        #    dataframe.to_csv(filename_results, index=False)
+        #    print(f"Probably incompleted Results saved to {filename_results}")
+        #    continue
+
+
+
+

utils/evaluation/dist_metrics.py

@@ -0,0 +1,971 @@
+
+from numpy.lib.scimath import sqrt as scisqrt
+from scipy import linalg
+import numpy as np  
+import torchaudio
+import os
+from importlib import import_module
+import torch
+
+
+from utils.feature_extractors.load_features import load_AFxRep, load_fx_encoder, load_fx_encoder_plusplus, load_CLAP
+
+from utils.log import make_PCA_figure
+
+SCALE_FACTOR = 100
+
+def calc_kernel_audio_distance(
+    x: torch.Tensor,
+    y: torch.Tensor,
+    device: str,
+    bandwidth=None,
+    kernel='gaussian',
+    precision=torch.float32,
+    eps=1e-8
+) -> torch.Tensor:
+    """
+    Compute the Kernel Audio Distance (KAD) between two samples using PyTorch.
+
+    Args:
+        x: The first set of embeddings of shape (m, embedding_dim).
+        y: The second set of embeddings of shape (n, embedding_dim).
+        cache_dirs: Directories to cache kernel statistics.
+        bandwidth: The bandwidth value for the Gaussian RBF kernel.
+        kernel: Kernel function to use ('gaussian', 'iq', 'imq').
+        precision: Type setting for matrix calculation precision.
+        eps: Small value to prevent division by zero.
+
+    Returns:
+        The KAD between x and y embedding sets.
+    """
+    # Ensure x and y are of the correct precision
+    x = x.to(dtype=precision, device=device)
+    y = y.to(dtype=precision, device=device)
+
+    assert bandwidth is not None, "Bandwidth must be provided for KAD calculation"
+
+    m, n = x.shape[0], y.shape[0]
+    
+    # Define kernel functions
+    gamma = 1 / (2 * bandwidth**2 + eps)
+    if kernel == 'gaussian':    # Gaussian Kernel
+        kernel = lambda a: torch.exp(-gamma * a)
+    elif kernel == 'iq':        # Inverse Quadratic Kernel
+        kernel = lambda a: 1 / (1 + gamma * a)
+    elif kernel == 'imq':       # Inverse Multiquadric Kernel
+        kernel = lambda a: 1 / torch.sqrt(1 + gamma * a)
+    else:
+        raise ValueError("Invalid kernel type. Valid kernels: 'gaussian', 'iq', 'imq'")
+    
+    # Load x kernel statistics
+    xx = x @ x.T
+    x_sqnorms = torch.diagonal(xx)
+    d2_xx = x_sqnorms.unsqueeze(1) + x_sqnorms.unsqueeze(0) - 2 * xx # shape (m, m)
+            
+    k_xx = kernel(d2_xx)
+    k_xx = k_xx - torch.diag(torch.diagonal(k_xx))
+    k_xx_mean = k_xx.sum() / (m * (m - 1))
+    
+    # Load y kernel statistics
+    yy = y @ y.T
+    y_sqnorms = torch.diagonal(yy)
+    d2_yy = y_sqnorms.unsqueeze(1) + y_sqnorms.unsqueeze(0) - 2 * yy # shape (n, n)
+
+    k_yy = kernel(d2_yy)
+    k_yy = k_yy - torch.diag(torch.diagonal(k_yy))
+    k_yy_mean = k_yy.sum() / (n * (n - 1))
+    
+    # Compute kernel statistics for xy
+    xy = x @ y.T
+    d2_xy = x_sqnorms.unsqueeze(1) + y_sqnorms.unsqueeze(0) - 2 * xy # shape (m, n)
+    k_xy = kernel(d2_xy)
+    k_xy_mean = k_xy.mean()
+    
+    # Compute MMD
+    result = k_xx_mean + k_yy_mean - 2 * k_xy_mean
+
+    return result * SCALE_FACTOR
+
+
+def median_pairwise_distance(x, subsample=None):
+    """
+    Compute the median pairwise distance of an embedding set.
+    
+    Args:
+    x: torch.Tensor of shape (n_samples, embedding_dim)
+    subsample: int, number of random pairs to consider (optional)
+    
+    Returns:
+    The median pairwise distance between points in x.
+    """
+    x = torch.tensor(x, dtype=torch.float32)
+    n_samples = x.shape[0]
+    
+    if subsample is not None and subsample < n_samples * (n_samples - 1) / 2:
+        # Randomly select pairs of indices
+        idx1 = torch.randint(0, n_samples, (subsample,))
+        idx2 = torch.randint(0, n_samples, (subsample,))
+        
+        # Ensure idx1 != idx2
+        mask = idx1 == idx2
+        idx2[mask] = (idx2[mask] + 1) % n_samples
+        
+        # Compute distances for selected pairs
+        distances = torch.sqrt(torch.sum((x[idx1] - x[idx2])**2, dim=1))
+    else:
+        # Compute all pairwise distances
+        distances = torch.pdist(x)
+        
+    return torch.median(distances).item()
+
+class DistMetric:
+    """
+    Base class for pairwise metrics.
+    
+    This class should be subclassed to implement specific pairwise metrics.
+    """
+    def __init__(self,type, sample_rate,*args, **kwargs):
+        """
+        Initialize the PairwiseMetric instance.
+        
+        Args:
+            *args: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+        """
+        self.type=type
+        self.sample_rate=sample_rate
+
+        if self.type == "fx_encoder":
+            self.model_args= kwargs.get("fx_encoder_args", None)
+
+            assert self.model_args is not None, "model_args must be provided for fx_encoder type"
+
+            self.distance_type=self.model_args.distance_type
+
+            self.feat_extractor = load_fx_encoder(self.model_args, self.device)
+
+            #self.feat_extractor = load_effects_encoder(ckpt_path=ckpt_path).to(self.device)
+
+        elif self.type == "fx_encoder_++":
+            self.model_args= kwargs.get("fx_encoder_plusplus_args", None)
+
+            assert self.model_args is not None, "model_args must be provided for fx_encoder_plusplus type"
+            
+            print(self.model_args)
+
+            self.distance_type=self.model_args.distance_type
+
+            self.feat_extractor = load_fx_encoder_plusplus(self.model_args, self.device)
+
+            #self.feat_extractor = load_effects_encoder(ckpt_path=ckpt_path).to(self.device)
+        
+        elif self.type== "AFxRep-mid" or self.type== "AFxRep-side" or self.type== "AFxRep":
+
+            self.model_args= kwargs.get("AFxRep_args", None)
+
+            assert self.model_args is not None, "model_args must be provided for AFxRep type"
+
+            self.distance_type=self.model_args.distance_type
+
+            feat_extractor = load_AFxRep(self.model_args, self.device, sample_rate=self.sample_rate)
+
+            if self.type == "AFxRep-mid":
+                def feat_extractor_mid(x):
+
+                    features= feat_extractor(x)
+
+                    #print("features shape:", features.shape)
+                    #divide by 2 to get mid and side features
+
+                    feat_mid, feat_side = features.chunk(2, dim=-1)
+
+                    return feat_mid
+
+                self.feat_extractor = feat_extractor_mid
+            
+            elif self.type == "AFxRep-side":
+                def feat_extractor_side(x):
+
+                    features= feat_extractor(x)
+
+                    #divide by 2 to get mid and side features
+
+                    feat_mid, feat_side = features.chunk(2, dim=-1)
+
+                    return feat_side
+
+                self.feat_extractor = feat_extractor_side
+            else:
+                self.feat_extractor = feat_extractor
+
+        elif self.type == "CLAP":
+            self.model_args= kwargs.get("CLAP_args", None)
+            assert self.model_args is not None, "model_args must be provided for CLAP type"
+
+            CLAP_fn = load_CLAP(self.model_args, self.device)
+            self.feat_extractor=lambda x: CLAP_fn(x, None)
+
+        else:
+            raise ValueError(f"Unknown type: {self.type}. Supported types: fx_encoder, fx_encoder_plusplus, AFxRep-mid, AFxRep-side, AFxRep")
+
+
+    def compute(self, *args, **kwargs):
+        raise NotImplementedError("Subclasses should implement this method.")
+
+    def extract_features(self, y, y_hat, x=None):
+
+        y=torch.tensor(y).permute(1,0).unsqueeze(0).to(self.device)
+        y_hat=torch.tensor(y_hat).permute(1,0).unsqueeze(0).to(self.device)
+
+        with torch.no_grad():
+            feat_y= self.feat_extractor(y)
+            feat_y_hat= self.feat_extractor(y_hat)
+
+        if x is not None:
+            x=torch.tensor(x).permute(1,0).unsqueeze(0).to(self.device)
+            with torch.no_grad():
+                feat_x= self.feat_extractor(x)
+        else:
+            feat_x=None
+
+
+        return feat_y, feat_y_hat, feat_x
+
+
+    def do_TSNE_figure(self, dict_features_y, dict_features_y_hat, dict_features_x=None, fit_mode="target", dict_cluster=None):
+        """
+        Perform PCA on the features and create a figure.
+        
+        Args:
+            dict_features_y (dict): Dictionary containing features for the first set.
+            dict_features_y_hat (dict): Dictionary containing features for the second set.
+            
+        Returns:
+            fig: The created figure.
+        """
+        import matplotlib.pyplot as plt
+        from sklearn.decomposition import PCA
+        from sklearn.manifold import TSNE   
+
+        y_values = list(dict_features_y.values())
+        y_values = torch.cat(y_values, dim=0)
+
+        y_hat_values = list(dict_features_y_hat.values())
+        y_hat_values = torch.cat(y_hat_values, dim=0)
+
+
+        if dict_cluster is not None:
+            clusters= list(dict_cluster.values())
+            clusters = [c.unsqueeze(0) if c.dim() == 0 else c for c in clusters]
+            clusters = torch.cat(clusters, dim=0)
+
+            #check different clusters (0,1,2,3...)
+            assert len(torch.unique(clusters)) == 2, "Only two clusters are supported for PCA visualization"
+            C0= clusters == 0
+            C1= clusters == 1
+
+        #manage Nans in y_values
+        y_values=torch.nan_to_num(y_values, nan=0)
+        y_hat_values=torch.nan_to_num(y_hat_values, nan=0)
+
+        #if self.pca is None:
+        self.tsne =TSNE(n_components=2, perplexity=30)
+
+        combined_data = torch.cat([y_values, y_hat_values], dim=0).cpu().numpy()
+        combined_result = self.tsne.fit_transform(combined_data)
+        
+        # Split the results back
+        n_y = y_values.shape[0]
+        tsne_result = combined_result[:n_y]
+        tsne_result_hat = combined_result[n_y:]
+
+
+        if dict_cluster is not None:
+           data_dict = {
+               "y_C0": tsne_result[C0],
+               "y_hat_C0": tsne_result_hat[C0],
+               "y_C1": tsne_result[C1],
+               "y_hat_C1": tsne_result_hat[C1]
+           }
+        else:
+           data_dict = {
+               "y": tsne_result,
+               "y_hat": tsne_result_hat
+           }
+
+
+
+        fig= make_PCA_figure(data_dict, title=self.type + " TSNE")
+
+
+        return fig
+
+
+    def do_PCA_figure(self, dict_features_y, dict_features_y_hat, dict_features_x=None, fit_mode="target", dict_cluster=None):
+        """
+        Perform PCA on the features and create a figure.
+        
+        Args:
+            dict_features_y (dict): Dictionary containing features for the first set.
+            dict_features_y_hat (dict): Dictionary containing features for the second set.
+            
+        Returns:
+            fig: The created figure.
+        """
+        import matplotlib.pyplot as plt
+        from sklearn.decomposition import PCA
+
+        y_values = list(dict_features_y.values())
+        y_values = torch.cat(y_values, dim=0)
+
+        y_hat_values = list(dict_features_y_hat.values())
+        y_hat_values = torch.cat(y_hat_values, dim=0)
+
+
+        if dict_cluster is not None:
+            clusters= list(dict_cluster.values())
+            clusters = [c.unsqueeze(0) if c.dim() == 0 else c for c in clusters]
+            clusters = torch.cat(clusters, dim=0)
+
+            #check different clusters (0,1,2,3...)
+            assert len(torch.unique(clusters)) == 2, "Only two clusters are supported for PCA visualization"
+            C0= clusters == 0
+            C1= clusters == 1
+
+        #manage Nans in y_values
+        y_values=torch.nan_to_num(y_values, nan=0)
+        y_hat_values=torch.nan_to_num(y_hat_values, nan=0)
+
+        if self.pca is None:
+            self.pca = PCA(n_components=2)
+            if fit_mode == "target":
+                pca_result = self.pca.fit_transform(y_values.cpu().numpy())
+            elif fit_mode == "all":
+                self.pca = self.pca.fit(torch.cat([y_values,y_hat_values], dim=0).cpu().numpy())
+                pca_result= self.pca.transform(y_values.cpu().numpy())
+        else:
+            pca_result = self.pca.transform(y_values.cpu().numpy())
+
+
+        #project y_hat values into the same PCA space
+        pca_result_hat = self.pca.transform(y_hat_values.cpu().numpy())
+
+        if dict_cluster is not None:
+           data_dict = {
+               "y_C0": pca_result[C0],
+               "y_hat_C0": pca_result_hat[C0],
+               "y_C1": pca_result[C1],
+               "y_hat_C1": pca_result_hat[C1]
+           }
+        else:
+           data_dict = {
+               "y": pca_result,
+               "y_hat": pca_result_hat
+           }
+
+        if dict_features_x is not None:
+            x_values = list(dict_features_x.values())
+            x_values = torch.cat(x_values, dim=0)
+
+            pca_result_x = self.pca.transform(x_values.cpu().numpy())
+
+            data_dict["x"] = pca_result_x
+
+
+
+        fig= make_PCA_figure(data_dict, title=self.type + " PCA")
+
+
+        return fig
+
+
+
+
+class KADFeatures(DistMetric):
+    """
+    Class for computing the pairwise spectral metric.
+    
+    This class inherits from PairwiseMetric and implements the compute method
+    to calculate the pairwise spectral metric.
+    """
+    def __init__(self,
+        type=None,
+        sample_rate=44100,
+        KAD_args=None,
+        classwise=False,
+                  *args, **kwargs):
+        """
+        Initialize the PairwiseSpectral instance.
+        
+        Args:
+            *args: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+        """
+        self.type = type
+
+        self.classwise = classwise
+
+        self.device=torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+        self.KAD_args = KAD_args
+
+        assert self.KAD_args is not None, "FAD_args must be provided"
+
+        if self.KAD_args.do_PCA_figure:
+            self.pca = None
+
+        self.bandwidth = None
+
+        super().__init__(self.type, sample_rate, *args, **kwargs)
+    
+    def calculate_KAD_distance_classwise(self, dict_features_y, dict_features_y_hat, dict_cluster):
+
+        if dict_cluster is not None:
+            clusters= list(dict_cluster.values())
+            clusters = [c.unsqueeze(0) if c.dim() == 0 else c for c in clusters]
+            clusters = torch.cat(clusters, dim=0)
+
+            #check different clusters (0,1,2,3...)
+            assert len(torch.unique(clusters)) == 2, "Only two clusters are supported for PCA visualization"
+            C0= clusters == 0
+            C1= clusters == 1
+
+        y= torch.cat(list(dict_features_y.values()), dim=0)
+        y_hat= torch.cat(list(dict_features_y_hat.values()), dim=0)
+
+        y_C0= y[C0]
+        y_hat_C0= y_hat[C0]
+
+        with torch.no_grad():
+            KAD_C0=calc_kernel_audio_distance(y_C0, y_hat_C0, device=self.device, bandwidth=self.bandwidth, kernel=self.KAD_args.kernel, precision=torch.float32)
+
+        y_C1= y[C1]
+        y_hat_C1= y_hat[C1]
+        with torch.no_grad():
+            KAD_C1=calc_kernel_audio_distance(y_C1, y_hat_C1, device=self.device, bandwidth=self.bandwidth, kernel=self.KAD_args.kernel, precision=torch.float32)
+        
+        KAD= (KAD_C0 + KAD_C1) / 2
+
+        return KAD.cpu().item()
+    def calculate_KAD_distance(self, dict_features_y, dict_features_y_hat):
+        y= torch.cat(list(dict_features_y.values()), dim=0)
+        y_hat= torch.cat(list(dict_features_y_hat.values()), dim=0)
+        with torch.no_grad():
+            KAD=calc_kernel_audio_distance(y, y_hat, device=self.device, bandwidth=self.bandwidth, kernel=self.KAD_args.kernel, precision=torch.float32)
+        return KAD.cpu().item()
+
+    def calculate_bandwidth(self, dict_features_y):
+        print("Calculating bandwidth...")
+        y= torch.cat(list(dict_features_y.values()), dim=0)
+        print("y shape:", y.shape)
+        self.bandwidth = median_pairwise_distance(y)
+        print("Bandwidth:", self.bandwidth)
+
+
+
+    def compute(self, dict_y, dict_y_hat, dict_x, dict_p_hat=None, dict_cluster=None, normalize=False, *args, **kwargs):
+        """
+        Compute the pairwise spectral metric.
+        
+        Args:
+            *args: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+            
+        Returns:
+            The computed pairwise spectral metric.
+        """
+
+        print("Computing KAD distance...")
+
+
+        if self.classwise:
+            assert dict_cluster is not None, "dict_cluster must be provided if classwise is True"
+
+        dict_features_y={}
+        dict_features_y_hat={}
+        #dict_features_x={}
+
+        if dict_y_hat is None:
+            print("computing KAD with style embeddings")
+            assert dict_p_hat is not None, "dict_p_hat must be provided if dict_y_hat is None"
+
+            for key in dict_y.keys():
+                y= dict_y[key]
+                #x= dict_x[key]
+
+                embed= dict_p_hat[key]
+                embed=torch.tensor(embed).to(self.device).unsqueeze(0)
+
+                embed_mid, embed_side = torch.chunk(embed, 2, dim=-1)
+
+                if self.type== "AFxRep-mid":
+                    p_hat= embed_mid
+                elif self.type== "AFxRep-side":
+                    p_hat= embed_side
+                elif self.type== "AFxRep":
+                    p_hat= embed
+    
+                #if x.shape[-2] == 1:
+                #    x = x.repeat( 2, 1)
+    
+                #assert x.shape == y.shape, f"Shape mismatch for key {key}: {x.shape} vs {y.shape}"
+    
+                c, d=y.shape
+                #assert b==1, f"Expected batch size of 1, got {b} for key {key}"
+    
+                assert c==2, f"Expected 2 channels, got {c} for key {key}"
+    
+                y=y.T
+                #x=x.T
+
+                y=torch.tensor(y).permute(1,0).unsqueeze(0).to(self.device)
+                #x=torch.tensor(x).permute(1,0).unsqueeze(0).to(self.device)
+
+                with torch.no_grad():
+                    feat_y= self.feat_extractor(y)
+                    #feat_x= self.feat_extractor(x)
+    
+                assert p_hat.shape == feat_y.shape, f"Shape mismatch for key {key}: {p_hat.shape} vs {feat_y.shape}"
+                #assert p_hat.shape == feat_x.shape, f"Shape mismatch for key {key}: {p_hat.shape} vs {feat_x.shape}"
+
+                dict_features_y[key] = feat_y
+                #dict_features_x[key] = feat_x
+                dict_features_y_hat[key] = p_hat
+
+        else:
+            for key in dict_y.keys():
+                y= dict_y[key]
+                #x= dict_x[key]
+                y_hat= dict_y_hat[key]
+    
+                #if x.shape[-2] == 1:
+                #    x = x.repeat( 2, 1)
+    
+                assert y.shape == y_hat.shape, f"Shape mismatch for key {key}: {y.shape} vs {y_hat.shape}"
+                #assert x.shape == y.shape, f"Shape mismatch for key {key}: {x.shape} vs {y.shape}"
+    
+                c, d=y.shape
+                #assert b==1, f"Expected batch size of 1, got {b} for key {key}"
+    
+                assert c==2, f"Expected 2 channels, got {c} for key {key}"
+    
+                y=y.T
+                y_hat=y_hat.T
+                #x=x.T
+    
+                y=torch.tensor(y).permute(1,0).unsqueeze(0).to(self.device)
+                y_hat=torch.tensor(y_hat).permute(1,0).unsqueeze(0).to(self.device)
+                #x=torch.tensor(x).permute(1,0).unsqueeze(0).to(self.device)
+        
+                with torch.no_grad():
+                    #try:
+                    feat_y= self.feat_extractor(y)
+                    feat_y_hat= self.feat_extractor(y_hat)
+                    #except Exception as e:
+                    #print(f"Error extracting features for key {key}: {e}")
+                    #    continue
+                    #feat_x= self.feat_extractor(x)
+
+                dict_features_y[key] = feat_y
+                dict_features_y_hat[key] = feat_y_hat
+                #dict_features_x[key] = feat_x
+            
+        if normalize:
+            #compute normalization statistics wrt to reference features
+            y_values = list(dict_features_y.values())
+            y_values = torch.cat(y_values, dim=0)
+            mean= y_values.mean(dim=0, keepdim=True)
+            std= y_values.std(dim=0, keepdim=True)
+
+            #normalize features
+            for key in dict_features_y.keys():
+                dict_features_y[key] = (dict_features_y[key] - mean) / (std + 1e-8)
+                dict_features_y_hat[key] = (dict_features_y_hat[key] - mean) / (std + 1e-8)
+
+
+        dict_output = {}
+        if self.KAD_args.do_PCA_figure:
+            fig=self.do_PCA_figure(dict_features_y, dict_features_y_hat, fit_mode=self.KAD_args.PCA_fit_mode, dict_cluster=dict_cluster)
+            key= self.type+ "_PCA_figure"
+            dict_output = {key: fig}
+        
+        if self.KAD_args.do_TSNE_figure:
+            fig=self.do_TSNE_figure(dict_features_y, dict_features_y_hat, fit_mode="all", dict_cluster=dict_cluster)
+            key= self.type+ "_TSNE_figure"
+            dict_output[key] = fig
+        
+
+        # Compute mean features
+
+
+
+        # Compute the distance between the mean features
+        #FAD_distance = self.calculate_FAD_distance(y_mean, y_cov, y_hat_mean, y_hat_cov)
+        if self.bandwidth is None:
+            self.calculate_bandwidth(dict_features_y)
+        
+        if self.classwise:
+            #calculate KAD distance for each class
+            KAD_distance = self.calculate_KAD_distance_classwise(dict_features_y, dict_features_y_hat, dict_cluster)
+        else:
+
+            KAD_distance = self.calculate_KAD_distance(dict_features_y, dict_features_y_hat)
+
+
+        return  KAD_distance, dict_output
+
+class FADFeatures(DistMetric):
+    """
+    Class for computing the pairwise spectral metric.
+    
+    This class inherits from PairwiseMetric and implements the compute method
+    to calculate the pairwise spectral metric.
+    """
+    def __init__(self,
+        type=None,
+        sample_rate=44100,
+        FAD_args=None,
+                  *args, **kwargs):
+        """
+        Initialize the PairwiseSpectral instance.
+        
+        Args:
+            *args: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+        """
+        self.type = type
+
+        self.device=torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+        self.FAD_args = FAD_args
+
+        assert self.FAD_args is not None, "FAD_args must be provided"
+
+        if self.FAD_args.do_PCA_figure:
+            self.pca = None
+
+        super().__init__(self.type, sample_rate, *args, **kwargs)
+    
+    def calculate_FAD_distance(self, mu1, sigma1, mu2, sigma2, eps=1e-6):
+
+        """Numpy implementation of the Frechet Distance.
+        The Frechet distance between two multivariate Gaussians X_1 ~ N(mu_1, C_1)
+        and X_2 ~ N(mu_2, C_2) is
+                d^2 = ||mu_1 - mu_2||^2 + Tr(C_1 + C_2 - 2*sqrt(C_1*C_2)).
+    
+        Stable version by Dougal J. Sutherland.
+    
+        Params:
+        -- mu1   : Numpy array containing the activations of a layer of the
+                   inception net (like returned by the function 'get_predictions')
+                   for generated samples.
+        -- mu2   : The sample mean over activations, precalculated on an
+                   representative data set.
+        -- sigma1: The covariance matrix over activations for generated samples.
+        -- sigma2: The covariance matrix over activations, precalculated on an
+                   representative data set.
+    
+        Returns:
+        --   : The Frechet Distance.
+        """
+        mu1=mu1.detach().cpu().numpy()
+        mu2=mu2.detach().cpu().numpy()
+
+        sigma1=sigma1.detach().cpu().numpy()
+        sigma2=sigma2.detach().cpu().numpy()
+
+    
+        mu1 = np.atleast_1d(mu1)
+        mu2 = np.atleast_1d(mu2)
+    
+        sigma1 = np.atleast_2d(sigma1)
+        sigma2 = np.atleast_2d(sigma2)
+    
+        assert (
+            mu1.shape == mu2.shape
+        ), "Training and test mean vectors have different lengths"
+        assert (
+            sigma1.shape == sigma2.shape
+        ), "Training and test covariances have different dimensions"
+    
+        diff = mu1 - mu2
+    
+        # Product might be almost singular
+        covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False)
+        if not np.isfinite(covmean).all():
+            msg = (
+                "fid calculation produces singular product; "
+                "adding %s to diagonal of cov estimates"
+            ) % eps
+            offset = np.eye(sigma1.shape[0]) * eps
+            covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset))
+    
+        # Numerical error might give slight imaginary component
+        if np.iscomplexobj(covmean):
+            if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3):
+                m = np.max(np.abs(covmean.imag))
+                raise ValueError("Imaginary component {}".format(m))
+            covmean = covmean.real
+    
+        tr_covmean = np.trace(covmean)
+    
+        return diff.dot(diff) + np.trace(sigma1) + np.trace(sigma2) - 2 * tr_covmean
+
+
+    def calculate_emb_statistics(self, features_dicto):
+        """
+        Calculate the mean and standard deviation of the features.
+        
+        Args:
+            features_dicto (dict): Dictionary containing features for each key.
+            
+        Returns:
+            mean_features (torch.Tensor): Mean of the features.
+            std_features (torch.Tensor): Standard deviation of the features.
+        """
+
+        all_features = torch.cat(list(features_dicto.values()), dim=0)
+
+
+        #mean
+        #mean_features = all_features.mean(dim=0)
+
+        #cov
+        #cov_features = torch.cov(all_features.T)
+        mean_features = np.mean(all_features.cpu().numpy(), axis=0)
+
+        cov_features= np.cov(all_features.cpu().numpy(), rowvar=False)
+
+
+
+        return torch.tensor(mean_features), torch.tensor(cov_features)
+    
+
+    def do_PCA_figure(self, dict_features_y, dict_features_y_hat, dict_features_x=None):
+        """
+        Perform PCA on the features and create a figure.
+        
+        Args:
+            dict_features_y (dict): Dictionary containing features for the first set.
+            dict_features_y_hat (dict): Dictionary containing features for the second set.
+            
+        Returns:
+            fig: The created figure.
+        """
+        import matplotlib.pyplot as plt
+        from sklearn.decomposition import PCA
+
+        y_values = list(dict_features_y.values())
+        y_values = torch.cat(y_values, dim=0)
+
+
+
+        if self.pca is None:
+            self.pca = PCA(n_components=2)
+            pca_result = self.pca.fit_transform(y_values.cpu().numpy())
+        else:
+            pca_result = self.pca.transform(y_values.cpu().numpy())
+
+
+        y_hat_values = list(dict_features_y_hat.values())
+        y_hat_values = torch.cat(y_hat_values, dim=0)
+
+        #project y_hat values into the same PCA space
+        pca_result_hat = self.pca.transform(y_hat_values.cpu().numpy())
+
+        data_dict = {
+            "y": pca_result,
+            "y_hat": pca_result_hat
+        }
+
+        if dict_features_x is not None:
+            x_values = list(dict_features_x.values())
+            x_values = torch.cat(x_values, dim=0)
+
+            pca_result_x = self.pca.transform(x_values.cpu().numpy())
+
+            data_dict["x"] = pca_result_x
+
+
+
+        fig= make_PCA_figure(data_dict, title=self.type + " PCA")
+
+
+        return fig
+
+
+
+
+
+    def compute(self, dict_y, dict_y_hat, dict_x, dict_p_hat=None, dict_cluster=None,   *args, **kwargs):
+        """
+        Compute the pairwise spectral metric.
+        
+        Args:
+            *args: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+            
+        Returns:
+            The computed pairwise spectral metric.
+        """
+
+        print("Computing FAD distance...")
+
+
+
+        dict_features_y={}
+        dict_features_y_hat={}
+        #dict_features_x={}
+
+        if dict_y_hat is None:
+            print("computing FAD with style embeddings")
+            assert dict_p_hat is not None, "dict_p_hat must be provided if dict_y_hat is None"
+
+            for key in dict_y.keys():
+                y= dict_y[key]
+                #x= dict_x[key]
+
+                embed= dict_p_hat[key]
+                embed=torch.tensor(embed).to(self.device).unsqueeze(0)
+
+                print("embed shape:", embed.shape)
+                embed_mid, embed_side = torch.chunk(embed, 2, dim=-1)
+
+                if self.type== "AFxRep-mid":
+                    p_hat= embed_mid
+                elif self.type== "AFxRep-side":
+                    p_hat= embed_side
+                elif self.type== "AFxRep":
+                    p_hat= embed
+    
+                #if x.shape[-2] == 1:
+                #    x = x.repeat( 2, 1)
+    
+                #assert x.shape == y.shape, f"Shape mismatch for key {key}: {x.shape} vs {y.shape}"
+    
+                c, d=y.shape
+                #assert b==1, f"Expected batch size of 1, got {b} for key {key}"
+    
+                assert c==2, f"Expected 2 channels, got {c} for key {key}"
+    
+                y=y.T
+                #x=x.T
+
+                y=torch.tensor(y).permute(1,0).unsqueeze(0).to(self.device)
+                #x=torch.tensor(x).permute(1,0).unsqueeze(0).to(self.device)
+
+                with torch.no_grad():
+                    feat_y= self.feat_extractor(y)
+                    #feat_x= self.feat_extractor(x)
+    
+                assert p_hat.shape == feat_y.shape, f"Shape mismatch for key {key}: {p_hat.shape} vs {feat_y.shape}"
+                #assert p_hat.shape == feat_x.shape, f"Shape mismatch for key {key}: {p_hat.shape} vs {feat_x.shape}"
+
+                dict_features_y[key] = feat_y
+                #dict_features_x[key] = feat_x
+                dict_features_y_hat[key] = p_hat
+
+        else:
+            for key in dict_y.keys():
+                y= dict_y[key]
+                #x= dict_x[key]
+                y_hat= dict_y_hat[key]
+    
+                #if x.shape[-2] == 1:
+                #    x = x.repeat( 2, 1)
+    
+                assert y.shape == y_hat.shape, f"Shape mismatch for key {key}: {y.shape} vs {y_hat.shape}"
+                #assert x.shape == y.shape, f"Shape mismatch for key {key}: {x.shape} vs {y.shape}"
+    
+                c, d=y.shape
+                #assert b==1, f"Expected batch size of 1, got {b} for key {key}"
+    
+                assert c==2, f"Expected 2 channels, got {c} for key {key}"
+    
+                y=y.T
+                y_hat=y_hat.T
+                #x=x.T
+    
+                y=torch.tensor(y).permute(1,0).unsqueeze(0).to(self.device)
+                y_hat=torch.tensor(y_hat).permute(1,0).unsqueeze(0).to(self.device)
+                #x=torch.tensor(x).permute(1,0).unsqueeze(0).to(self.device)
+        
+                with torch.no_grad():
+                    feat_y= self.feat_extractor(y)
+                    feat_y_hat= self.feat_extractor(y_hat)
+                    #feat_x= self.feat_extractor(x)
+
+                dict_features_y[key] = feat_y
+                dict_features_y_hat[key] = feat_y_hat
+                #dict_features_x[key] = feat_x
+            
+        dict_output = {}
+        if self.FAD_args.do_PCA_figure:
+            fig=self.do_PCA_figure(dict_features_y, dict_features_y_hat, fit_mode="target")
+            key= self.type+ "_PCA_figure"
+            dict_output[key] = fig
+
+        if self.FAD_args.do_TSNE_figure:
+            fig=self.do_TSNE_figure(dict_features_y, dict_features_y_hat, fit_mode="target")
+            key= self.type+ "_TSNE_figure"
+            dict_output[key] = fig
+
+        # Compute mean features
+
+        y_mean, y_cov = self.calculate_emb_statistics(dict_features_y)
+
+        y_hat_mean, y_hat_cov = self.calculate_emb_statistics(dict_features_y_hat)
+
+
+        # Compute the distance between the mean features
+        #FAD_distance = self.calculate_FAD_distance(y_mean, y_cov, y_hat_mean, y_hat_cov)
+        FAD_distance = self.calculate_FAD_distance(y_mean, y_cov, y_hat_mean, y_hat_cov)
+
+
+        return  FAD_distance, dict_output
+
+
+def metric_factory(metric_name, sample_rate, *args, **kwargs):
+    """
+    Factory function to create a metric function based on the metric name.
+    
+    Args:
+        metric_name (str): The name of the metric to create.
+        *args: Variable length argument list.
+        **kwargs: Arbitrary keyword arguments.
+        
+    Returns:
+        An instance of a class that implements the metric function.
+    """
+    if metric_name == "fad-fx_encoder":
+        return FADFeatures(*args, **kwargs, type="fx_encoder", sample_rate=sample_rate )
+    elif metric_name == "fad-AFxRep":
+        return FADFeatures(*args, **kwargs, type="AFxRep", sample_rate=sample_rate)
+    elif metric_name == "fad-AFxRep-mid":
+        return FADFeatures(*args, **kwargs, type="AFxRep-side", sample_rate=sample_rate)
+    elif metric_name == "fad-AFxRep-side":
+        return FADFeatures(*args, **kwargs, type="AFxRep-mid", sample_rate=sample_rate)
+    elif metric_name == "kad-fx_encoder":
+        return KADFeatures(*args, **kwargs, type="fx_encoder", sample_rate=sample_rate)
+    elif metric_name == "kad-AFxRep":
+        return KADFeatures(*args, **kwargs, type="AFxRep", sample_rate=sample_rate)
+    elif metric_name == "kad-AFxRep-mid":
+        return KADFeatures(*args, **kwargs, type="AFxRep-side", sample_rate=sample_rate)
+    elif metric_name == "kad-AFxRep-side":
+        return KADFeatures(*args, **kwargs, type="AFxRep-mid", sample_rate=sample_rate)
+    elif metric_name == "kad-class-fx_encoder":
+        return KADFeatures(*args, **kwargs, type="fx_encoder", sample_rate=sample_rate, classwise=True)
+    elif metric_name == "kad-class-AFxRep":
+        return KADFeatures(*args, **kwargs, type="AFxRep", sample_rate=sample_rate, classwise=True)
+    elif metric_name == "kad-class-AFxRep-mid":
+        return KADFeatures(*args, **kwargs, type="AFxRep-side", sample_rate=sample_rate, classwise=True)
+    elif metric_name == "kad-class-AFxRep-side":
+        return KADFeatures(*args, **kwargs, type="AFxRep-mid", sample_rate=sample_rate, classwise=True)
+    else:
+        raise ValueError(f"Unknown metric: {metric_name}")
+
+# Example usage:
+#metric_instance = metric_factory("pairwise-spectral")
+#```

utils/evaluation/dist_metrics_multitrack.py

@@ -0,0 +1,619 @@
+
+import math
+from numpy.lib.scimath import sqrt as scisqrt
+from scipy import linalg
+import numpy as np  
+import torchaudio
+import os
+from importlib import import_module
+import torch
+
+
+from utils.feature_extractors.load_features import load_AFxRep, load_fx_encoder, load_fx_encoder_plusplus, load_MERT, load_CLAP, load_fx_encoder_plusplus_2048
+
+from utils.log import make_PCA_figure
+
+SCALE_FACTOR = 100
+
+def calc_kernel_audio_distance(
+    x: torch.Tensor,
+    y: torch.Tensor,
+    device: str,
+    bandwidth=None,
+    kernel='gaussian',
+    precision=torch.float32,
+    eps=1e-8
+) -> torch.Tensor:
+    """
+    Compute the Kernel Audio Distance (KAD) between two samples using PyTorch.
+
+    Args:
+        x: The first set of embeddings of shape (m, embedding_dim).
+        y: The second set of embeddings of shape (n, embedding_dim).
+        cache_dirs: Directories to cache kernel statistics.
+        bandwidth: The bandwidth value for the Gaussian RBF kernel.
+        kernel: Kernel function to use ('gaussian', 'iq', 'imq').
+        precision: Type setting for matrix calculation precision.
+        eps: Small value to prevent division by zero.
+
+    Returns:
+        The KAD between x and y embedding sets.
+    """
+    # Ensure x and y are of the correct precision
+    x = x.to(dtype=precision, device=device)
+    y = y.to(dtype=precision, device=device)
+
+    assert bandwidth is not None, "Bandwidth must be provided for KAD calculation"
+
+    m, n = x.shape[0], y.shape[0]
+    
+    # Define kernel functions
+    gamma = 1 / (2 * bandwidth**2 + eps)
+    if kernel == 'gaussian':    # Gaussian Kernel
+        kernel = lambda a: torch.exp(-gamma * a)
+    elif kernel == 'iq':        # Inverse Quadratic Kernel
+        kernel = lambda a: 1 / (1 + gamma * a)
+    elif kernel == 'imq':       # Inverse Multiquadric Kernel
+        kernel = lambda a: 1 / torch.sqrt(1 + gamma * a)
+    else:
+        raise ValueError("Invalid kernel type. Valid kernels: 'gaussian', 'iq', 'imq'")
+    
+    # Load x kernel statistics
+    xx = x @ x.T
+    x_sqnorms = torch.diagonal(xx)
+    d2_xx = x_sqnorms.unsqueeze(1) + x_sqnorms.unsqueeze(0) - 2 * xx # shape (m, m)
+            
+    k_xx = kernel(d2_xx)
+    k_xx = k_xx - torch.diag(torch.diagonal(k_xx))
+    k_xx_mean = k_xx.sum() / (m * (m - 1))
+    
+    # Load y kernel statistics
+    yy = y @ y.T
+    y_sqnorms = torch.diagonal(yy)
+    d2_yy = y_sqnorms.unsqueeze(1) + y_sqnorms.unsqueeze(0) - 2 * yy # shape (n, n)
+
+    k_yy = kernel(d2_yy)
+    k_yy = k_yy - torch.diag(torch.diagonal(k_yy))
+    k_yy_mean = k_yy.sum() / (n * (n - 1))
+    
+    # Compute kernel statistics for xy
+    xy = x @ y.T
+    d2_xy = x_sqnorms.unsqueeze(1) + y_sqnorms.unsqueeze(0) - 2 * xy # shape (m, n)
+    k_xy = kernel(d2_xy)
+    k_xy_mean = k_xy.mean()
+    
+    # Compute MMD
+    result = k_xx_mean + k_yy_mean - 2 * k_xy_mean
+
+    return result * SCALE_FACTOR
+
+
+def median_pairwise_distance(x, subsample=None):
+    """
+    Compute the median pairwise distance of an embedding set.
+    
+    Args:
+    x: torch.Tensor of shape (n_samples, embedding_dim)
+    subsample: int, number of random pairs to consider (optional)
+    
+    Returns:
+    The median pairwise distance between points in x.
+    """
+    x = torch.tensor(x, dtype=torch.float32)
+    n_samples = x.shape[0]
+    
+    if subsample is not None and subsample < n_samples * (n_samples - 1) / 2:
+        # Randomly select pairs of indices
+        idx1 = torch.randint(0, n_samples, (subsample,))
+        idx2 = torch.randint(0, n_samples, (subsample,))
+        
+        # Ensure idx1 != idx2
+        mask = idx1 == idx2
+        idx2[mask] = (idx2[mask] + 1) % n_samples
+        
+        # Compute distances for selected pairs
+        distances = torch.sqrt(torch.sum((x[idx1] - x[idx2])**2, dim=1))
+    else:
+        # Compute all pairwise distances
+        distances = torch.pdist(x)
+        
+    return torch.median(distances).item()
+
+class DistMetric:
+    """
+    Base class for pairwise metrics.
+    
+    This class should be subclassed to implement specific pairwise metrics.
+    """
+    def __init__(self,type, sample_rate,*args, **kwargs):
+        """
+        Initialize the PairwiseMetric instance.
+        
+        Args:
+            *args: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+        """
+        self.type=type
+        self.sample_rate=sample_rate
+
+        self.taxonomy_ref= kwargs.get("taxonomy_ref", None)
+
+        if self.type == "DynamicFeatures":
+            from utils.feature_extractors.AF_features_embedding import AF_fourier_embedding
+            AFembedding= AF_fourier_embedding(device=self.device)
+
+            def feat_extracfor_fn(x):
+                z, _ = AFembedding.encode(x)
+                return z
+
+            self.feat_extractor = feat_extracfor_fn
+        elif self.type == "FxEncoder++":
+            self.model_args= kwargs.get("fx_encoder_plusplus_args", None)
+            assert self.model_args is not None, "model_args must be provided for fxencAFv2-fxenc++ type"
+
+            self.distance_type=self.model_args.distance_type
+            fxencoder = load_fx_encoder_plusplus_2048(self.model_args, self.device)
+
+            def feat_extractor_fn(x):
+                z= fxencoder(x)
+                z=torch.nn.functional.normalize(z, dim=-1, p=2)  # normalize to unit variance
+                return z
+            
+            self.feat_extractor = feat_extractor_fn
+
+        else:
+            raise ValueError(f"Unknown type: {self.type}. Supported types: fx_encoder, fx_encoder_plusplus, AFxRep-mid, AFxRep-side, AFxRep")
+
+
+    def compute(self, *args, **kwargs):
+        raise NotImplementedError("Subclasses should implement this method.")
+
+    def extract_features(self, y, y_hat, x=None):
+
+        y=torch.tensor(y).permute(1,0).unsqueeze(0).to(self.device)
+        y_hat=torch.tensor(y_hat).permute(1,0).unsqueeze(0).to(self.device)
+
+        with torch.no_grad():
+            feat_y= self.feat_extractor(y)
+            feat_y_hat= self.feat_extractor(y_hat)
+
+        if x is not None:
+            x=torch.tensor(x).permute(1,0).unsqueeze(0).to(self.device)
+            with torch.no_grad():
+                feat_x= self.feat_extractor(x)
+        else:
+            feat_x=None
+
+
+        return feat_y, feat_y_hat, feat_x
+
+
+    def do_TSNE_figure(self, dict_features_y, dict_features_y_hat,  dict_taxonomy=None):
+        """
+        Perform PCA on the features and create a figure.
+        
+        Args:
+            dict_features_y (dict): Dictionary containing features for the first set.
+            dict_features_y_hat (dict): Dictionary containing features for the second set.
+            
+        Returns:
+            fig: The created figure.
+        """
+        import matplotlib.pyplot as plt
+        from sklearn.manifold import TSNE   
+
+
+        y_values = {}
+        y_hat_values = {}
+        for key, track_name in self.taxonomy_ref.items():
+
+
+            y_values[track_name] = []
+            for k in range(len(dict_features_y.keys())):
+                index = [i for i, tax_id in enumerate(dict_taxonomy[k]) if tax_id == key]
+                print("taxonomy",k, dict_taxonomy[k])
+                if index:  # If any matches were found
+                    # Select the corresponding rows from dict_features_y[k]
+                    selected_features = dict_features_y[k][index]
+                    y_values[track_name].append(selected_features)
+
+
+            if not y_values[track_name]:
+                print(f"Warning: No features found for track {track_name} with key {key}. Skipping this track.")
+                #remove the track from the dictionaries
+                y_values.pop(track_name, None)
+                continue
+
+            y_values[track_name]=torch.cat(y_values[track_name], dim=0)
+            y_values[track_name]=torch.nan_to_num(y_values[track_name], nan=0)
+
+            #print(clusters[track_name])
+
+            y_hat_values[track_name] = []
+
+            for k in range(len(dict_features_y_hat.keys())):
+                index = [i for i, tax_id in enumerate(dict_taxonomy[k]) if tax_id == key]
+                if index:  # If any matches were found
+                    # Select the corresponding rows from dict_features_y[k]
+                    selected_features = dict_features_y_hat[k][index]
+                    y_hat_values[track_name].append(selected_features)
+                #else:
+                #    y_hat_values[track_name].append(None)  # Append zeros if no match found
+
+
+            y_hat_values[track_name]=torch.cat(y_hat_values[track_name], dim=0)
+            y_hat_values[track_name]=torch.nan_to_num(y_hat_values[track_name], nan=0)
+
+            #print(clusters_hat[track_name])
+            
+        if self.pca is None:
+            self.pca= {}
+        
+
+        tsne_result = {}
+        tsne_result_hat = {}
+        data_dict = {}
+        figs= {}
+
+        for k in y_values.keys():
+            print("Processing track:", k)
+            self.tsne =TSNE(n_components=2, perplexity=30)
+            combined_data = torch.cat([y_values[k], y_hat_values[k]], dim=0).cpu().numpy()
+            combined_result = self.tsne.fit_transform(combined_data)
+
+            n_y = y_values[k].shape[0]
+            tsne_result[k] = combined_result[:n_y]
+            tsne_result_hat[k] = combined_result[n_y:]
+
+
+            data_dict[k] = {
+                   "y": tsne_result[k],
+                   "y_hat": tsne_result_hat[k] } 
+
+            figs[k]= make_PCA_figure(data_dict[k], title=self.type + " TSNE; track: "+k)
+
+
+        return figs
+
+    def do_PCA_figure(self, dict_features_y, dict_features_y_hat, fit_mode="target", dict_cluster=None, dict_taxonomy=None):
+        """
+        Perform PCA on the features and create a figure.
+        
+        Args:
+            dict_features_y (dict): Dictionary containing features for the first set.
+            dict_features_y_hat (dict): Dictionary containing features for the second set.
+            
+        Returns:
+            fig: The created figure.
+        """
+        import matplotlib.pyplot as plt
+        from sklearn.decomposition import PCA
+
+
+        y_values = {}
+        y_hat_values = {}
+        if dict_cluster is not None:
+            clusters={}
+            clusters_hat= {}
+
+        for key, track_name in self.taxonomy_ref.items():
+
+
+            y_values[track_name] = []
+            if dict_cluster is not None:
+                clusters[track_name] = []
+            for k in range(len(dict_features_y.keys())):
+                print("taxonomy",k, dict_taxonomy[k])
+                index = [i for i, tax_id in enumerate(dict_taxonomy[k]) if tax_id == key]
+                if index:  # If any matches were found
+                    # Select the corresponding rows from dict_features_y[k]
+                    selected_features = dict_features_y[k][index]
+                    y_values[track_name].append(selected_features)
+
+                    if dict_cluster is not None:
+                        clusters[track_name].append(dict_cluster[k].unsqueeze(0))
+
+            #check if y_values[track_name] is empty
+            if not y_values[track_name]:
+                print(f"Warning: No features found for track {track_name} with key {key}. Skipping this track.")
+                #remove the track from the dictionaries
+                y_values.pop(track_name, None)
+                continue
+
+            y_values[track_name]=torch.cat(y_values[track_name], dim=0)
+            y_values[track_name]=torch.nan_to_num(y_values[track_name], nan=0)
+
+            if dict_cluster is not None:
+                clusters[track_name]=torch.cat(clusters[track_name], dim=0)
+
+            y_hat_values[track_name] = []
+            if dict_cluster is not None:    
+                clusters_hat[track_name] = []
+
+            for k in range(len(dict_features_y_hat.keys())):
+                index = [i for i, tax_id in enumerate(dict_taxonomy[k]) if tax_id == key]
+                if index:  # If any matches were found
+                    # Select the corresponding rows from dict_features_y[k]
+                    selected_features = dict_features_y_hat[k][index]
+                    y_hat_values[track_name].append(selected_features)
+                    if dict_cluster is not None:
+                        clusters_hat[track_name].append(dict_cluster[k].unsqueeze(0))
+
+            y_hat_values[track_name]=torch.cat(y_hat_values[track_name], dim=0)
+            y_hat_values[track_name]=torch.nan_to_num(y_hat_values[track_name], nan=0)
+
+            if dict_cluster is not None:
+                clusters_hat[track_name]=torch.cat(clusters_hat[track_name], dim=0)
+            
+
+        if self.pca is None:
+            self.pca= {}
+
+        pca_result = {}
+        pca_result_hat = {}
+        data_dict = {}
+        figs= {}
+
+        for k in y_values.keys():
+            print("Processing track:", k)
+            #check if pca already exists for this key
+            if k not in self.pca.keys():
+                if fit_mode == "target":
+                    print("Fitting PCA for target values only track:", k)
+                    self.pca[k] = PCA(n_components=2).fit(y_values[k].cpu().numpy())
+                elif fit_mode == "all":
+                    print("Fitting PCA for all values track:", k)
+                    self.pca[k] = PCA(n_components=2).fit(torch.cat([y_values[k], y_hat_values[k]], dim=0).cpu().numpy())
+                else:
+                    raise ValueError(f"Unknown fit_mode: {fit_mode}. Supported modes: target, all")
+
+            pca_result[k]=self.pca[k].transform(y_values[k].cpu().numpy())
+            pca_result_hat[k] = self.pca[k].transform(y_hat_values[k].cpu().numpy())
+
+
+            if dict_cluster is not None:
+               data_dict[k] = {
+                   "y_C0": pca_result[k][clusters[k] == 0],
+                   "y_hat_C0": pca_result_hat[k][clusters_hat[k] == 0],
+                   "y_C1": pca_result[k][clusters[k] == 1],
+                   "y_hat_C1": pca_result_hat[k][clusters_hat[k] == 1]
+               }
+            else:
+               data_dict[k] = {
+                   "y": pca_result[k],
+                   "y_hat": pca_result_hat[k]
+               }
+
+
+            figs[k]= make_PCA_figure(data_dict[k], title=self.type + " PCA; track: "+k)
+
+
+        return figs
+
+
+class KADFeatures(DistMetric):
+    """
+    Class for computing the pairwise spectral metric.
+    
+    This class inherits from PairwiseMetric and implements the compute method
+    to calculate the pairwise spectral metric.
+    """
+    def __init__(self,
+        type=None,
+        sample_rate=44100,
+        KAD_args=None,
+        *args, **kwargs):
+        """
+        Initialize the PairwiseSpectral instance.
+        
+        Args:
+            *args: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+        """
+        self.type = type
+
+        self.device=torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        self.KAD_args = KAD_args
+
+        assert self.KAD_args is not None, "FAD_args must be provided"
+
+        if self.KAD_args.do_PCA_figure:
+            self.pca = None
+
+        self.bandwidth = None
+
+        super().__init__(self.type, sample_rate, *args, **kwargs)
+
+    def calculate_KAD_distance(self, dict_features_y, dict_features_y_hat, dict_taxonomy=None):
+
+        y_values = {}
+        y_hat_values = {}
+        for key, track_name in self.taxonomy_ref.items():
+            y_values[track_name] = []
+            for k in range(len(dict_features_y.keys())):
+                index = [i for i, tax_id in enumerate(dict_taxonomy[k]) if tax_id == key]
+                if index:  # If any matches were found
+                    # Select the corresponding rows from dict_features_y[k]
+                    selected_features = dict_features_y[k][index]
+                    y_values[track_name].append(selected_features)
+                #else:
+                #    y_values[track_name].append(None)  # Append zeros if no match found
+            if not y_values[track_name]:
+                print(f"Warning: No features found for track {track_name} with key {key}. Skipping this track.")
+                #remove the track from the dictionaries
+                y_values.pop(track_name, None)
+                continue
+
+            y_values[track_name]=torch.cat(y_values[track_name], dim=0)
+            y_values[track_name]=torch.nan_to_num(y_values[track_name], nan=0)
+
+            y_hat_values[track_name] = []
+            for k in range(len(dict_features_y_hat.keys())):
+                index = [i for i, tax_id in enumerate(dict_taxonomy[k]) if tax_id == key]
+                if index:  # If any matches were found
+                    # Select the corresponding rows from dict_features_y[k]
+                    selected_features = dict_features_y_hat[k][index]
+                    y_hat_values[track_name].append(selected_features)
+                #else:
+                #    y_hat_values[track_name].append(None)  # Append zeros if no match found
+
+            y_hat_values[track_name]=torch.cat(y_hat_values[track_name], dim=0)
+            y_hat_values[track_name]=torch.nan_to_num(y_hat_values[track_name], nan=0)
+            
+
+        KAD={}
+        weight={}
+        total_examples= 0
+        for key in y_values.keys():
+            y= y_values[key]
+            y_hat= y_hat_values[key]
+            with torch.no_grad():
+                KAD[key]=calc_kernel_audio_distance(y, y_hat, device=self.device, bandwidth=self.bandwidth, kernel=self.KAD_args.kernel, precision=torch.float32)
+
+            weight[key]= y.shape[0]
+            total_examples += weight[key]
+
+        KAD= sum([KAD[k] * weight[k] for k in KAD.keys()]) / total_examples
+
+        return KAD.cpu().item()
+
+    def calculate_bandwidth(self, dict_features_y):
+
+        print("Calculating bandwidth...")
+        y= torch.cat(list(dict_features_y.values()), dim=0)
+        print("y shape:", y.shape)
+        self.bandwidth = median_pairwise_distance(y)
+        print("Bandwidth:", self.bandwidth)
+
+
+
+    def compute(self, dict_y, dict_y_hat, dict_taxonomy=None, dict_p_hat=None,  *args, **kwargs):
+        """
+        Compute the pairwise spectral metric.
+        
+        Args:
+            *args: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+            
+        Returns:
+            The computed pairwise spectral metric.
+        """
+
+        print("Computing KAD distance...")
+
+        dict_features_y={}
+        dict_features_y_hat={}
+
+        if dict_y_hat is None and dict_p_hat is not None:
+
+            for key in dict_y.keys():
+                y= dict_y[key]
+
+                embed= dict_p_hat[key]
+                embed=torch.tensor(embed).to(self.device)
+
+                embed=embed*math.sqrt(embed.shape[-1])  # Scale the embedding
+
+                embed_fxenc=embed[...,:2048]/ math.sqrt(2048)  # Scale the first 128 dimensions
+                embed_AF=embed[...,2048:]/ math.sqrt(64)  # Scale the last 128 dimensions
+
+
+                if self.type == "FxEncoder++":
+                    p_hat= embed_fxenc
+
+                elif self.type == "DynamicFeatures":
+                    p_hat= embed_AF
+    
+                n, c, d=y.shape
+    
+                assert c==2, f"Expected 2 channels, got {c} for key {key}"
+
+                y=torch.tensor(y).to(self.device)
+
+                with torch.no_grad():
+                    feat_y= self.feat_extractor(y)
+    
+                assert p_hat.shape == feat_y.shape, f"Shape mismatch for key {key}: {p_hat.shape} vs {feat_y.shape}"
+
+                dict_features_y[key] = feat_y
+                dict_features_y_hat[key] = p_hat
+
+        else:
+            for key in dict_y.keys():
+                y= dict_y[key]
+                y_hat= dict_y_hat[key]
+    
+                assert y.shape == y_hat.shape, f"Shape mismatch for key {key}: {y.shape} vs {y_hat.shape}"
+    
+                n, c, d=y.shape
+                if dict_taxonomy is not None:
+                    taxonomy= dict_taxonomy[key]
+                    assert len(taxonomy) == n, f"Taxonomy length mismatch for key {key}: {len(taxonomy)} vs {n}"
+    
+                assert c==2, f"Expected 2 channels, got {c} for key {key}"
+    
+                y=torch.tensor(y).to(self.device)
+                y_hat=torch.tensor(y_hat).to(self.device)
+    
+                with torch.no_grad():
+                    feat_y= self.feat_extractor(y)
+                    feat_y_hat= self.feat_extractor(y_hat)
+
+                dict_features_y[key] = feat_y
+                dict_features_y_hat[key] = feat_y_hat
+            
+        dict_output = {}
+        if self.KAD_args.do_PCA_figure:
+            dict_figs=self.do_PCA_figure(dict_features_y, dict_features_y_hat, dict_taxonomy=dict_taxonomy, fit_mode=self.KAD_args.PCA_fit_mode)
+            #returns a dictionary with figures for each taxonomy key
+            #modify the key to include the type
+
+            for k, v in dict_figs.items():
+                #modify the key to include the type
+                key= self.type+ "_PCA_figure_" + k
+                dict_output[key] = v
+
+        if self.KAD_args.do_TSNE_figure:
+            dict_figs=self.do_TSNE_figure(dict_features_y, dict_features_y_hat, dict_taxonomy=dict_taxonomy)
+            for  k, fig in dict_figs.items():
+                key= self.type+ "_TSNE_figure"+k
+                dict_output[key] = fig
+        
+
+
+        # Compute the distance between the mean features
+        #FAD_distance = self.calculate_FAD_distance(y_mean, y_cov, y_hat_mean, y_hat_cov)
+        if self.bandwidth is None:
+            self.calculate_bandwidth(dict_features_y)
+        
+
+        KAD_distance = self.calculate_KAD_distance(dict_features_y, dict_features_y_hat, dict_taxonomy)
+
+
+        return  KAD_distance, dict_output
+
+
+
+def metric_factory(metric_name, sample_rate, *args, **kwargs):
+    """
+    Factory function to create a metric function based on the metric name.
+    
+    Args:
+        metric_name (str): The name of the metric to create.
+        *args: Variable length argument list.
+        **kwargs: Arbitrary keyword arguments.
+        
+    Returns:
+        An instance of a class that implements the metric function.
+    """
+    if metric_name == "kad-FxEncoder++":
+        return KADFeatures(*args, **kwargs, type="FxEncoder++", sample_rate=sample_rate)
+    elif metric_name == "kad-DynamicFeatures":
+        return KADFeatures(*args, **kwargs, type="DynamicFeatures", sample_rate=sample_rate)
+    else:
+        raise ValueError(f"Unknown metric: {metric_name}")
+
+# Example usage:
+#metric_instance = metric_factory("pairwise-spectral")
+#```

utils/evaluation/pairwise_metrics.py

@@ -0,0 +1,958 @@
+
+import math
+import torchaudio
+import os
+from importlib import import_module
+import yaml
+import torch
+from utils.feature_extractors.load_features import load_AFxRep, load_fx_encoder, load_fx_encoder_plusplus, load_fx_encoder_plusplus_2048
+from utils.MSS_loss import MultiScale_Spectral_Loss_MidSide_DDSP
+
+class PairwiseMetric:
+    """
+    Base class for pairwise metrics.
+    
+    This class should be subclassed to implement specific pairwise metrics.
+    """
+    def __init__(self, *args, **kwargs):
+        """
+        Initialize the PairwiseMetric instance.
+        
+        Args:
+            *args: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+        """
+        pass
+
+    def compute(self, *args, **kwargs):
+        raise NotImplementedError("Subclasses should implement this method.")
+
+class PairwiseFeatures(PairwiseMetric):
+    """
+    Class for computing the pairwise spectral metric.
+    
+    This class inherits from PairwiseMetric and implements the compute method
+    to calculate the pairwise spectral metric.
+    """
+    def __init__(self,
+        type=None,
+        sample_rate=44100,
+                  *args, **kwargs):
+        """
+        Initialize the PairwiseSpectral instance.
+        
+        Args:
+            *args: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+        """
+        self.type = type
+        self.sample_rate = sample_rate
+
+        self.device=torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        if self.type == "fx_encoder":
+            self.model_args= kwargs.get("fx_encoder_args", None)
+
+            assert self.model_args is not None, "model_args must be provided for fx_encoder type"
+
+            self.distance_type=self.model_args.distance_type
+
+            self.feat_extractor = load_fx_encoder(self.model_args, self.device)
+
+
+            #self.feat_extractor = load_effects_encoder(ckpt_path=ckpt_path).to(self.device)
+        elif self.type == "fxenc++":
+            self.model_args= kwargs.get("fx_encoder_plusplus_args", None)
+
+            assert self.model_args is not None, "model_args must be provided for fx_encoder type"
+
+            self.distance_type=self.model_args.distance_type
+
+            self.feat_extractor = load_fx_encoder_plusplus(self.model_args, self.device)
+
+
+            #self.feat_extractor = load_effects_encoder(ckpt_path=ckpt_path).to(self.device)
+        
+        elif self.type == "logrms":
+            self.model_args= kwargs.get("logrms_args", None)
+
+            assert self.model_args is not None, "model_args must be provided for logrms type"
+            self.distance_type=self.model_args.distance_type
+
+            from utils.feature_extractors.dsp_features import compute_log_rms
+
+            self.feat_extractor = lambda x: compute_log_rms(x)
+
+        elif self.type == "crestfactor":
+            self.model_args= kwargs.get("crestfactor_args", None)
+            assert self.model_args is not None, "model_args must be provided for crestfactor type"
+
+            self.distance_type=self.model_args.distance_type
+
+            from utils.feature_extractors.dsp_features import compute_crest_factor
+
+            self.feat_extractor = lambda x: compute_crest_factor(x)
+        elif self.type == "logspread":
+            self.model_args= kwargs.get("logspread_args", None)
+            assert self.model_args is not None, "model_args must be provided for logspread type"
+
+            self.distance_type=self.model_args.distance_type
+
+            from utils.feature_extractors.dsp_features import compute_log_spread
+
+            self.feat_extractor = lambda x: compute_log_spread(x).view(-1, 1)
+        elif self.type == "stereowidth":
+            self.model_args= kwargs.get("stereowidth_args", None)
+            assert self.model_args is not None, "model_args must be provided for stereowidth type"
+
+            self.distance_type=self.model_args.distance_type
+
+            from utils.feature_extractors.dsp_features import compute_stereo_width
+            self.feat_extractor = lambda x: compute_stereo_width(x).view(-1, 1)
+
+        elif self.type == "stereoimbalance":
+            self.model_args= kwargs.get("stereoimbalance_args", None)
+            assert self.model_args is not None, "model_args must be provided for stereoimbalance type"  
+            self.distance_type=self.model_args.distance_type
+            from utils.feature_extractors.dsp_features import compute_stereo_imbalance
+            self.feat_extractor = lambda x: compute_stereo_imbalance(x).view(-1, 1)
+        
+        elif self.type== "AFxRep-mid" or self.type== "AFxRep-side" or self.type== "AFxRep":
+
+            self.model_args= kwargs.get("AFxRep_args", None)
+
+            assert self.model_args is not None, "model_args must be provided for AFxRep type"
+
+            self.distance_type=self.model_args.distance_type
+
+            feat_extractor = load_AFxRep(self.model_args, self.device)
+
+            if self.type == "AFxRep-mid":
+                def feat_extractor_mid(x):
+
+                    features= feat_extractor(x)
+
+                    #divide by 2 to get mid and side features
+
+                    feat_mid, feat_side = features.chunk(2, dim=-1)
+
+                    return feat_mid
+
+                self.feat_extractor = feat_extractor_mid
+            
+            elif self.type == "AFxRep-side":
+                def feat_extractor_side(x):
+
+                    features= feat_extractor(x)
+
+                    #divide by 2 to get mid and side features
+
+                    feat_mid, feat_side = features.chunk(2, dim=-1)
+
+                    return feat_side
+
+                self.feat_extractor = feat_extractor_side
+            else:
+                self.feat_extractor = feat_extractor
+
+
+        super().__init__(*args, **kwargs)
+    
+    def compute_feature_distance(self, y, y_hat, sample_rate, type):    
+
+        y=torch.tensor(y).permute(1,0).unsqueeze(0).to(self.device)
+        y_hat=torch.tensor(y_hat).permute(1,0).unsqueeze(0).to(self.device)
+
+        with torch.no_grad():
+            feat_y= self.feat_extractor(y)
+            feat_y_hat= self.feat_extractor(y_hat)
+
+        if self.distance_type == "cosine":
+            cos_dist= 1- torch.cosine_similarity(feat_y_hat, feat_y, dim=1)
+
+
+            return {"distance": cos_dist.mean().item()}
+        elif self.distance_type == "l1":
+            l1_dist = torch.abs(feat_y_hat - feat_y).mean(dim=1)
+
+            return {"distance": l1_dist.mean().item()}
+
+
+    def compute(self, dict_y, dict_y_hat, dict_x,   *args, **kwargs):
+        """
+        Compute the pairwise spectral metric.
+        
+        Args:
+            *args: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+            
+        Returns:
+            The computed pairwise spectral metric.
+        """
+
+
+
+        dict_features={}
+
+        for key in dict_y.keys():
+            y= dict_y[key]
+            y_hat= dict_y_hat[key]
+
+            assert y.shape == y_hat.shape, f"Shape mismatch for key {key}: {y.shape} vs {y_hat.shape}"
+
+            c, d=y.shape
+            #assert b==1, f"Expected batch size of 1, got {b} for key {key}"
+
+            assert c==2, f"Expected 2 channels, got {c} for key {key}"
+
+            y=y.T
+            y_hat=y_hat.T
+
+
+            if self.type == "spectral":
+                from evaluation.automix_evaluation import compute_spectral_features 
+                dict_features_out= compute_spectral_features(y_hat, y ,self.sample_rate)
+                dict_features[key] = dict_features_out['mean_mape_spectral']
+            elif self.type=="panning":
+                from evaluation.automix_evaluation import compute_panning_features 
+                dict_features_out = compute_panning_features(y_hat, y, self.sample_rate)
+                dict_features[key] = dict_features_out['mean_mape_panning']
+            elif self.type=="loudness":
+                from evaluation.automix_evaluation import compute_loudness_features 
+                dict_features_out = compute_loudness_features(y_hat, y, self.sample_rate)
+                dict_features[key] = dict_features_out['mean_mape_loudness']
+            elif self.type=="dynamic":
+                from evaluation.automix_evaluation import compute_dynamic_features 
+                dict_features_out = compute_dynamic_features(y_hat, y, self.sample_rate)
+                dict_features[key] = dict_features_out['mean_mape_dynamic']
+            elif self.type=="fx_encoder":
+                dict_features_out = self.compute_feature_distance(y, y_hat, self.sample_rate, type=self.type)
+                dict_features[key] = dict_features_out["distance"]
+            elif self.type=="AFxRep":
+                dict_features_out = self.compute_feature_distance(y, y_hat, self.sample_rate, type=self.type)
+                dict_features[key] = dict_features_out["distance"]
+            elif self.type=="AFxRep-mid":
+                dict_features_out = self.compute_feature_distance(y, y_hat, self.sample_rate, type=self.type)
+                dict_features[key] = dict_features_out["distance"]
+            elif self.type=="AFxRep-side":
+                dict_features_out = self.compute_feature_distance(y, y_hat, self.sample_rate, type=self.type)
+                dict_features[key] = dict_features_out["distance"]
+            elif self.type=="fxenc++":
+                dict_features_out = self.compute_feature_distance(y, y_hat, self.sample_rate, type=self.type)
+                dict_features[key] = dict_features_out["distance"]
+            elif self.type=="logrms":
+                dict_features_out = self.compute_feature_distance(y, y_hat, self.sample_rate, type=self.type)
+                dict_features[key] = dict_features_out["distance"]
+            elif self.type=="crestfactor":
+                dict_features_out = self.compute_feature_distance(y, y_hat, self.sample_rate, type=self.type)
+                dict_features[key] = dict_features_out["distance"]
+            elif self.type=="logspread":
+                dict_features_out = self.compute_feature_distance(y, y_hat, self.sample_rate, type=self.type)
+                dict_features[key] = dict_features_out["distance"]
+            elif self.type=="stereowidth":
+                dict_features_out = self.compute_feature_distance(y, y_hat, self.sample_rate, type=self.type)
+                dict_features[key] = dict_features_out["distance"]
+            elif self.type=="stereoimbalance":
+                dict_features_out = self.compute_feature_distance(y, y_hat, self.sample_rate, type=self.type)
+                dict_features[key] = dict_features_out["distance"]
+            
+
+        # Compute the mean of the features across all keys
+
+        mean_features = sum(dict_features.values()) / len(dict_features)
+
+        return  mean_features, {}
+
+class PairwiseStyleMultitrackFeatures(PairwiseMetric):
+    """
+    Class for computing the pairwise spectral metric.
+    
+    This class inherits from PairwiseMetric and implements the compute method
+    to calculate the pairwise spectral metric.
+    """
+    def __init__(self,
+        type=None,
+        sample_rate=44100,
+                  *args, **kwargs):
+        """
+        Initialize the PairwiseSpectral instance.
+        
+        Args:
+            *args: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+        """
+        self.type = type
+        self.sample_rate = sample_rate
+
+        self.device=torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        if self.type == "fx_encoder":
+            raise NotImplementedError("Style features for fx_encoder not implemented yet")
+            self.model_args= kwargs.get("fx_encoder_args", None)
+
+            assert self.model_args is not None, "model_args must be provided for fx_encoder type"
+
+            self.distance_type=self.model_args.distance_type
+
+            self.feat_extractor = load_fx_encoder(self.model_args, self.device)
+
+
+            #self.feat_extractor = load_effects_encoder(ckpt_path=ckpt_path).to(self.device)
+        
+        elif self.type== "AFxRep-mid" or self.type== "AFxRep-side" or self.type== "AFxRep":
+
+            self.model_args= kwargs.get("AFxRep_args", None)
+
+            assert self.model_args is not None, "model_args must be provided for AFxRep type"
+
+            self.distance_type=self.model_args.distance_type
+
+            feat_extractor = load_AFxRep(self.model_args, self.device)
+
+            if self.type == "AFxRep-mid":
+                def feat_extractor_mid(x):
+
+                    features= feat_extractor(x)
+
+                    #divide by 2 to get mid and side features
+
+                    feat_mid, feat_side = features.chunk(2, dim=-1)
+
+                    return feat_mid
+
+                self.feat_extractor = feat_extractor_mid
+            
+            elif self.type == "AFxRep-side":
+                def feat_extractor_side(x):
+
+                    features= feat_extractor(x)
+
+                    #divide by 2 to get mid and side features
+
+                    feat_mid, feat_side = features.chunk(2, dim=-1)
+
+                    return feat_side
+
+                self.feat_extractor = feat_extractor_side
+            else:
+                self.feat_extractor = feat_extractor
+
+        elif self.type == "fxenc2048AFv3-fxenc++":
+
+            self.model_args= kwargs.get("fx_encoder_plusplus_args", None)
+            assert self.model_args is not None, "model_args must be provided for fxencAFv2-fxenc++ type"
+
+            self.distance_type= self.model_args.distance_type
+
+            fxencoder = load_fx_encoder_plusplus_2048(self.model_args, self.device)
+
+            def feat_extractor_fn(x):
+                z= fxencoder(x)
+                z=torch.nn.functional.normalize(z, dim=-1, p=2)  # normalize to unit variance
+                return z
+            
+            self.feat_extractor = feat_extractor_fn
+
+        elif self.type == "fxenc2048AFv3-AF":
+
+            from utils.AF_features_embedding_v2 import AF_fourier_embedding
+            AFembedding= AF_fourier_embedding(device=self.device)
+
+            self.distance_type="cosine"
+
+            def feat_extracfor_fn(x):
+                z, _ = AFembedding.encode(x)
+                return z
+
+            self.feat_extractor = feat_extracfor_fn
+        
+
+        super().__init__(*args, **kwargs)
+    
+    def compute_feature_distance(self, y, p_hat, sample_rate, type):    
+
+        y=torch.tensor(y).permute(1,0).unsqueeze(0).to(self.device)
+        #y_hat=torch.tensor(y_hat).permute(1,0).unsqueeze(0).to(self.device)
+
+        with torch.no_grad():
+            feat_y= self.feat_extractor(y)
+            #feat_y_hat= self.feat_extractor(y_hat)
+        
+        assert p_hat.shape == feat_y.shape, f"Shape mismatch: p_hat {p_hat.shape} vs feat_y {feat_y.shape}"
+
+        if self.distance_type == "cosine":
+            cos_dist= 1- torch.cosine_similarity(p_hat, feat_y, dim=1)
+
+            return {"distance": cos_dist.mean().item()}
+
+
+
+    def compute(self, dict_y, dict_y_hat, dict_x, dict_p_hat=None,   *args, **kwargs):
+        """
+        Compute the pairwise spectral metric.
+        
+        Args:
+            *args: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+            
+        Returns:
+            The computed pairwise spectral metric.
+        """
+
+        dict_features={}
+        
+        index=0
+
+        for key in dict_y.keys():
+            y= dict_y[key]
+
+            embed= dict_p_hat[key]
+            embed=torch.tensor(embed).to(self.device).unsqueeze(0)
+
+            print(f"Processing key: {key}, embed shape: {embed.shape}")
+            if "AFxRep" in self.type:
+                embed_mid, embed_side = torch.chunk(embed, 2, dim=-1)
+
+                if self.type== "AFxRep-mid":
+                        p_hat= embed_mid
+                elif self.type== "AFxRep-side":
+                        p_hat= embed_side
+                elif self.type== "AFxRep":
+                        p_hat= embed
+            elif "fxenc2048AFv3" in self.type:
+
+                embed=embed*math.sqrt(embed.shape[-1])  # Scale the embedding
+                embed_fxenc=embed[...,:2048]/ math.sqrt(2048)  # Scale the first 128 dimensions
+                embed_AF=embed[...,2048:]/ math.sqrt(64)  # Scale the last 128 dimensions
+
+                if self.type == "fxenc2048AFv3-fxenc++":
+                        p_hat= embed_fxenc
+                elif self.type == "fxenc2048AFv3-AF":
+                        p_hat= embed_AF
+    
+
+            n, c, d=y.shape
+
+            #assert b==1, f"Expected batch size of 1, got {b} for key {key}"
+
+            assert c==2, f"Expected 2 channels, got {c} for key {key}"
+
+            for i in range(n):
+                y_i=y[i].T
+
+                if self.type=="fx_encoder":
+                    raise NotImplementedError("Style features for fx_encoder not implemented yet")
+                    dict_features_out = self.compute_feature_distance(y, p_hat, self.sample_rate, type=self.type)
+                    dict_features[key] = dict_features_out["distance"]
+                elif self.type=="AFxRep":
+                    dict_features_out = self.compute_feature_distance(y_i, p_hat[i], self.sample_rate, type=self.type)
+                    dict_features[index] = dict_features_out["distance"]
+                elif self.type=="AFxRep-mid":
+                    dict_features_out = self.compute_feature_distance(y_i, p_hat[i], self.sample_rate, type=self.type)
+                    dict_features[index] = dict_features_out["distance"]
+                elif self.type=="AFxRep-side":
+                    dict_features_out = self.compute_feature_distance(y_i, p_hat[i], self.sample_rate, type=self.type)
+                    dict_features[index] = dict_features_out["distance"]
+                elif self.type=="fxenc2048AFv3-fxenc++":
+                    dict_features_out = self.compute_feature_distance(y_i, p_hat[i], self.sample_rate, type=self.type)
+                    dict_features[index] = dict_features_out["distance"]
+                elif self.type=="fxenc2048AFv3-AF":
+                    dict_features_out = self.compute_feature_distance(y_i, p_hat, self.sample_rate, type=self.type)
+                    dict_features[index] = dict_features_out["distance"]
+            
+
+        mean_features = sum(dict_features.values()) / len(dict_features)
+
+        return  mean_features, {}
+class PairwiseStyleFeatures(PairwiseMetric):
+    """
+    Class for computing the pairwise spectral metric.
+    
+    This class inherits from PairwiseMetric and implements the compute method
+    to calculate the pairwise spectral metric.
+    """
+    def __init__(self,
+        type=None,
+        sample_rate=44100,
+                  *args, **kwargs):
+        """
+        Initialize the PairwiseSpectral instance.
+        
+        Args:
+            *args: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+        """
+        self.type = type
+        self.sample_rate = sample_rate
+
+        self.device=torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        if self.type == "fx_encoder":
+            raise NotImplementedError("Style features for fx_encoder not implemented yet")
+            self.model_args= kwargs.get("fx_encoder_args", None)
+
+            assert self.model_args is not None, "model_args must be provided for fx_encoder type"
+
+            self.distance_type=self.model_args.distance_type
+
+            self.feat_extractor = load_fx_encoder(self.model_args, self.device)
+
+
+            #self.feat_extractor = load_effects_encoder(ckpt_path=ckpt_path).to(self.device)
+        
+        elif self.type== "AFxRep-mid" or self.type== "AFxRep-side" or self.type== "AFxRep":
+
+            self.model_args= kwargs.get("AFxRep_args", None)
+
+            assert self.model_args is not None, "model_args must be provided for AFxRep type"
+
+            self.distance_type=self.model_args.distance_type
+
+            feat_extractor = load_AFxRep(self.model_args, self.device)
+
+            if self.type == "AFxRep-mid":
+                def feat_extractor_mid(x):
+
+                    features= feat_extractor(x)
+
+                    #divide by 2 to get mid and side features
+
+                    feat_mid, feat_side = features.chunk(2, dim=-1)
+
+                    return feat_mid
+
+                self.feat_extractor = feat_extractor_mid
+            
+            elif self.type == "AFxRep-side":
+                def feat_extractor_side(x):
+
+                    features= feat_extractor(x)
+
+                    #divide by 2 to get mid and side features
+
+                    feat_mid, feat_side = features.chunk(2, dim=-1)
+
+                    return feat_side
+
+                self.feat_extractor = feat_extractor_side
+            else:
+                self.feat_extractor = feat_extractor
+
+        elif self.type == "fxenc2048AFv3-fxenc++":
+
+            self.model_args= kwargs.get("fx_encoder_plusplus_args", None)
+            assert self.model_args is not None, "model_args must be provided for fxencAFv2-fxenc++ type"
+
+            self.distance_type= self.model_args.distance_type
+
+            fxencoder = load_fx_encoder_plusplus_2048(self.model_args, self.device)
+
+            def feat_extractor_fn(x):
+                z= fxencoder(x)
+                z=torch.nn.functional.normalize(z, dim=-1, p=2)  # normalize to unit variance
+                return z
+            
+            self.feat_extractor = feat_extractor_fn
+
+        elif self.type == "fxenc2048AFv3-AF":
+
+            from utils.AF_features_embedding_v2 import AF_fourier_embedding
+            AFembedding= AF_fourier_embedding(device=self.device)
+
+            self.distance_type="cosine"
+
+            def feat_extracfor_fn(x):
+                z, _ = AFembedding.encode(x)
+                return z
+
+            self.feat_extractor = feat_extracfor_fn
+        
+
+        super().__init__(*args, **kwargs)
+    
+    def compute_feature_distance(self, y, p_hat, sample_rate, type):    
+
+        y=torch.tensor(y).permute(1,0).unsqueeze(0).to(self.device)
+        #y_hat=torch.tensor(y_hat).permute(1,0).unsqueeze(0).to(self.device)
+
+        with torch.no_grad():
+            feat_y= self.feat_extractor(y)
+            #feat_y_hat= self.feat_extractor(y_hat)
+        
+        assert p_hat.shape == feat_y.shape, f"Shape mismatch: p_hat {p_hat.shape} vs feat_y {feat_y.shape}"
+
+        if self.distance_type == "cosine":
+            cos_dist= 1- torch.cosine_similarity(p_hat, feat_y, dim=1)
+
+            return {"distance": cos_dist.mean().item()}
+
+
+
+    def compute(self, dict_y, dict_y_hat, dict_x, dict_p_hat=None,   *args, **kwargs):
+        """
+        Compute the pairwise spectral metric.
+        
+        Args:
+            *args: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+            
+        Returns:
+            The computed pairwise spectral metric.
+        """
+
+        dict_features={}
+
+        for key in dict_y.keys():
+            y= dict_y[key]
+
+            embed= dict_p_hat[key]
+            embed=torch.tensor(embed).to(self.device).unsqueeze(0)
+
+            print(f"Processing key: {key}, embed shape: {embed.shape}")
+            if "AFxRep" in self.type:
+                embed_mid, embed_side = torch.chunk(embed, 2, dim=-1)
+
+                if self.type== "AFxRep-mid":
+                        p_hat= embed_mid
+                elif self.type== "AFxRep-side":
+                        p_hat= embed_side
+                elif self.type== "AFxRep":
+                        p_hat= embed
+            elif "fxenc2048AFv3" in self.type:
+
+                embed=embed*math.sqrt(embed.shape[-1])  # Scale the embedding
+                embed_fxenc=embed[...,:2048]/ math.sqrt(2048)  # Scale the first 128 dimensions
+                embed_AF=embed[...,2048:]/ math.sqrt(64)  # Scale the last 128 dimensions
+
+                if self.type == "fxenc2048AFv3-fxenc++":
+                        p_hat= embed_fxenc
+                elif self.type == "fxenc2048AFv3-AF":
+                        p_hat= embed_AF
+    
+
+            c, d=y.shape
+            #assert b==1, f"Expected batch size of 1, got {b} for key {key}"
+
+            assert c==2, f"Expected 2 channels, got {c} for key {key}"
+
+            y=y.T
+            #y_hat=y_hat.T
+
+            if self.type=="fx_encoder":
+                raise NotImplementedError("Style features for fx_encoder not implemented yet")
+                dict_features_out = self.compute_feature_distance(y, p_hat, self.sample_rate, type=self.type)
+                dict_features[key] = dict_features_out["distance"]
+            elif self.type=="AFxRep":
+                dict_features_out = self.compute_feature_distance(y, p_hat, self.sample_rate, type=self.type)
+                dict_features[key] = dict_features_out["distance"]
+            elif self.type=="AFxRep-mid":
+                dict_features_out = self.compute_feature_distance(y, p_hat, self.sample_rate, type=self.type)
+                dict_features[key] = dict_features_out["distance"]
+            elif self.type=="AFxRep-side":
+                dict_features_out = self.compute_feature_distance(y, p_hat, self.sample_rate, type=self.type)
+                dict_features[key] = dict_features_out["distance"]
+            elif self.type=="fxenc2048AFv3-fxenc++":
+                dict_features_out = self.compute_feature_distance(y, p_hat, self.sample_rate, type=self.type)
+                dict_features[key] = dict_features_out["distance"]
+            elif self.type=="fxenc2048AFv3-AF":
+                dict_features_out = self.compute_feature_distance(y, p_hat, self.sample_rate, type=self.type)
+                dict_features[key] = dict_features_out["distance"]
+            
+            
+
+        # Compute the mean of the features across all keys
+
+        mean_features = sum(dict_features.values()) / len(dict_features)
+
+        return  mean_features, {}
+
+class PairwiseIMMSS(PairwiseMetric):
+    """
+    Class for computing the pairwise LDR metric.
+    
+    This class inherits from PairwiseMetric and implements the compute method
+    to calculate the pairwise LDR metric.
+    """
+    def __init__(self, mode=None, *args, **kwargs):
+        """
+        Initialize the PairwiseLDR instance.
+        
+        Args:
+            *args: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+        """
+        super().__init__(*args, **kwargs)
+
+        assert mode is not None, "Mode must be specified for PairwiseLDR"
+
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")   
+        if mode == "mss-lr":
+            multi_scale_spectral_ori = MultiScale_Spectral_Loss_MidSide_DDSP(mode='ori', eps=1e-6, device=device)
+            self.metric= lambda y_hat, y: multi_scale_spectral_ori(y_hat, y)
+        elif mode == "mss-ms":
+            multi_scale_spectral_midside = MultiScale_Spectral_Loss_MidSide_DDSP(mode='midside', eps=1e-6, device=device)
+            self.metric= lambda y_hat, y: multi_scale_spectral_midside(y_hat, y)
+
+    def compute(self, dict_y, dict_y_hat, dict_x,   *args, **kwargs):
+        """
+        Compute the pairwise spectral metric.
+        
+        Args:
+            *args: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+            
+        Returns:
+            The computed pairwise spectral metric.
+        """
+
+        dict_metrics={}
+
+        for key in dict_y.keys():
+            y= dict_y[key]
+            y_hat= dict_y_hat[key]
+
+            assert y.shape == y_hat.shape, f"Shape mismatch for key {key}: {y.shape} vs {y_hat.shape}"
+
+            c, d=y.shape
+            #assert b==1, f"Expected batch size of 1, got {b} for key {key}"
+
+            assert c==2, f"Expected 2 channels, got {c} for key {key}"
+
+            y=torch.from_numpy(y).cuda().unsqueeze(0)
+            y_hat=torch.from_numpy(y_hat).cuda().unsqueeze(0)
+
+            metric=self.metric(y_hat, y)
+
+            dict_metrics[key] = metric.item()
+
+        mean_features = sum(dict_metrics.values()) / len(dict_metrics)
+
+        return  mean_features, {}
+class PairwiseLDR(PairwiseMetric):
+    """
+    Class for computing the pairwise LDR metric.
+    
+    This class inherits from PairwiseMetric and implements the compute method
+    to calculate the pairwise LDR metric.
+    """
+    def __init__(self, mode=None, *args, **kwargs):
+        """
+        Initialize the PairwiseLDR instance.
+        
+        Args:
+            *args: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+        """
+        super().__init__(*args, **kwargs)
+
+        assert mode is not None, "Mode must be specified for PairwiseLDR"
+
+        if mode == "mldr-lr":
+            from evaluation.ldr import MLDRLoss
+            raw_metric= MLDRLoss(
+                sr=44100,
+                s_taus=[50, 100],
+                l_taus=[1000, 2000],
+            ).cuda()
+            self.metric=lambda y_hat, y: 0.5*raw_metric(y_hat, y)
+        elif mode == "mldr-ms":
+            from evaluation.ldr import MLDRLoss
+            raw_metric= MLDRLoss(
+                sr=44100,
+                s_taus=[50, 100],
+                l_taus=[1000, 2000],
+                mid_side=True
+            ).cuda()
+            self.metric=lambda y_hat, y: 0.25*raw_metric(y_hat, y)
+
+    def compute(self, dict_y, dict_y_hat, dict_x,   *args, **kwargs):
+        """
+        Compute the pairwise spectral metric.
+        
+        Args:
+            *args: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+            
+        Returns:
+            The computed pairwise spectral metric.
+        """
+
+        dict_metrics={}
+
+        for key in dict_y.keys():
+            y= dict_y[key]
+            y_hat= dict_y_hat[key]
+
+            assert y.shape == y_hat.shape, f"Shape mismatch for key {key}: {y.shape} vs {y_hat.shape}"
+
+            c, d=y.shape
+            #assert b==1, f"Expected batch size of 1, got {b} for key {key}"
+
+            assert c==2, f"Expected 2 channels, got {c} for key {key}"
+
+            y=y.T
+            y_hat=y_hat.T
+            
+            y=torch.from_numpy(y).cuda().unsqueeze(0)
+            y_hat=torch.from_numpy(y_hat).cuda().unsqueeze(0)
+
+            metric=self.metric(y_hat, y)
+
+            dict_metrics[key] = metric.item()
+
+        mean_features = sum(dict_metrics.values()) / len(dict_metrics)
+
+        return  mean_features, {}
+
+
+class PairwiseAuraloss(PairwiseMetric):
+    """
+    Class for computing the pairwise MSS metric.
+    
+    This class inherits from PairwiseMetric and implements the compute method
+    to calculate the pairwise MSS metric.
+    """
+    def __init__(self, mode=None, *args, **kwargs):
+        """
+        Initialize the PairwiseMSS instance.
+        
+        Args:
+            *args: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+        """
+        super().__init__(*args, **kwargs)
+
+        assert mode is not None, "Mode must be specified for PairwiseMSS"
+
+        if mode == "mss-lr":
+            from auraloss.freq import MultiResolutionSTFTLoss
+            raw_metric=MultiResolutionSTFTLoss(
+                [128, 512, 2048],
+                [32, 128, 512],
+                [128, 512, 2048],
+                sample_rate=44100,
+                perceptual_weighting=True,
+            ).cuda()
+            self.metric=lambda y_hat, y: raw_metric(y_hat, y)
+        elif mode == "mss-ms":
+            from auraloss.freq import  SumAndDifferenceSTFTLoss
+            raw_metric=SumAndDifferenceSTFTLoss(
+            [128, 512, 2048],
+            [32, 128, 512],
+            [128, 512, 2048],
+            sample_rate=44100,
+            perceptual_weighting=True,
+            ).cuda()
+            self.metric=lambda y_hat, y: 0.5*raw_metric(y_hat, y)
+
+    def compute(self, dict_y, dict_y_hat, dict_x,   *args, **kwargs):
+        """
+        Compute the pairwise spectral metric.
+        
+        Args:
+            *args: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+            
+        Returns:
+            The computed pairwise spectral metric.
+        """
+
+        dict_metrics={}
+
+        for key in dict_y.keys():
+            y= dict_y[key]
+            y_hat= dict_y_hat[key]
+
+            assert y.shape == y_hat.shape, f"Shape mismatch for key {key}: {y.shape} vs {y_hat.shape}"
+
+            c, d=y.shape
+            #assert b==1, f"Expected batch size of 1, got {b} for key {key}"
+
+            assert c==2, f"Expected 2 channels, got {c} for key {key}"
+
+            #y=y.T
+            #y_hat=y_hat.T
+            
+            y=torch.from_numpy(y).cuda().unsqueeze(0)
+            y_hat=torch.from_numpy(y_hat).cuda().unsqueeze(0)
+
+            metric=self.metric(y_hat, y)
+
+            dict_metrics[key] = metric.item()
+
+        mean_features = sum(dict_metrics.values()) / len(dict_metrics)
+
+        return  mean_features, {}
+
+def metric_factory(metric_name, sample_rate, *args, **kwargs):
+    """
+    Factory function to create a metric function based on the metric name.
+    
+    Args:
+        metric_name (str): The name of the metric to create.
+        *args: Variable length argument list.
+        **kwargs: Arbitrary keyword arguments.
+        
+    Returns:
+        An instance of a class that implements the metric function.
+    """
+    if metric_name == "pairwise-spectral":
+        return PairwiseFeatures(*args, **kwargs, type="spectral", sample_rate=sample_rate)
+    elif metric_name == "pairwise-panning":
+        return PairwiseFeatures(*args, **kwargs, type="panning", sample_rate=sample_rate)
+    elif metric_name == "pairwise-loudness":
+        return PairwiseFeatures(*args, **kwargs, type="loudness", sample_rate=sample_rate)
+    elif metric_name == "pairwise-dynamic":
+        return PairwiseFeatures(*args, **kwargs, type="dynamic", sample_rate=sample_rate)
+    elif metric_name == "pairwise-mss-lr":
+        return PairwiseAuraloss(mode="mss-lr",*args, **kwargs)
+    elif metric_name == "pairwise-mss-ms":
+        return PairwiseAuraloss(mode="mss-ms",*args, **kwargs)
+    elif metric_name == "pairwise-IM-mss-lr":
+        return PairwiseIMMSS(mode="mss-lr",*args, **kwargs)
+    elif metric_name == "pairwise-IM-mss-ms":
+        return PairwiseIMMSS(mode="mss-ms",*args, **kwargs)
+    elif metric_name == "pairwise-mldr-lr":
+        return PairwiseLDR(mode="mldr-lr",*args, **kwargs)
+    elif metric_name == "pairwise-mldr-ms":
+        return PairwiseLDR(mode="mldr-ms",*args, **kwargs)
+    elif metric_name == "pairwise-fx_encoder":
+        return PairwiseFeatures(*args, **kwargs, type="fx_encoder", sample_rate=sample_rate, model_args=kwargs.get('fx_encoder_args', None))
+    elif metric_name == "pairwise-AFxRep":
+        return PairwiseFeatures(*args, **kwargs, type="AFxRep", sample_rate=sample_rate, model_args=kwargs.get('AFxRep_args', None))
+    elif metric_name == "pairwise-AFxRep-mid":
+        return PairwiseFeatures(*args, **kwargs, type="AFxRep-mid", sample_rate=sample_rate, model_args=kwargs.get('AFxRep_args', None))
+    elif metric_name == "pairwise-AFxRep-side":
+        return PairwiseFeatures(*args, **kwargs, type="AFxRep-side", sample_rate=sample_rate, model_args=kwargs.get('AFxRep_args', None))
+    elif metric_name == "pairwise-style-AFxRep":
+        return PairwiseStyleFeatures(*args, **kwargs, type="AFxRep", sample_rate=sample_rate, model_args=kwargs.get('AFxRep_args', None))
+    elif metric_name == "pairwise-style-AFxRep-mid":
+        return PairwiseStyleFeatures(*args, **kwargs, type="AFxRep-mid", sample_rate=sample_rate, model_args=kwargs.get('AFxRep_args', None))
+    elif metric_name == "pairwise-style-AFxRep-side":
+        return PairwiseStyleFeatures(*args, **kwargs, type="AFxRep-side", sample_rate=sample_rate, model_args=kwargs.get('AFxRep_args', None))
+    elif metric_name == "pairwise-style-fxenc2048AFv3-fxenc++":
+        return PairwiseStyleFeatures(*args, **kwargs, type="fxenc2048AFv3-fxenc++", sample_rate=sample_rate, model_args=kwargs.get('AFxRep_args', None))
+    elif metric_name == "pairwise-style-fxenc2048AFv3-AF":
+        return PairwiseStyleFeatures(*args, **kwargs, type="fxenc2048AFv3-AF", sample_rate=sample_rate, model_args=kwargs.get('AFxRep_args', None))
+    elif metric_name == "pairwise-style-multitrack-fxenc2048AFv3-fxenc++":
+        return PairwiseStyleMultitrackFeatures(*args, **kwargs, type="fxenc2048AFv3-fxenc++", sample_rate=sample_rate, model_args=kwargs.get('AFxRep_args', None))
+    elif metric_name == "pairwise-style-multitrack-fxenc2048AFv3-AF":
+        return PairwiseStyleMultitrackFeatures(*args, **kwargs, type="fxenc2048AFv3-AF", sample_rate=sample_rate, model_args=kwargs.get('AFxRep_args', None))
+    elif metric_name == "pairwise-fxenc++":
+        return PairwiseFeatures(*args, **kwargs, type="fxenc++", sample_rate=sample_rate, model_args=kwargs.get('fx_encoder_plusplus_args', None))
+    elif metric_name == "pairwise-logrms":
+        return PairwiseFeatures(*args, **kwargs, type="logrms", sample_rate=sample_rate, model_args=kwargs.get('logrms_args', None))
+    elif metric_name == "pairwise-crestfactor":
+        return PairwiseFeatures(*args, **kwargs, type="crestfactor", sample_rate=sample_rate, model_args=kwargs.get('crestfactor_args', None))
+    elif metric_name == "pairwise-logspread":
+        return PairwiseFeatures(*args, **kwargs, type="logspread", sample_rate=sample_rate, model_args=kwargs.get('logspread_args', None))
+    elif metric_name == "pairwise-stereowidth":
+        return PairwiseFeatures(*args, **kwargs, type="stereowidth", sample_rate=sample_rate, model_args=kwargs.get('stereowidth_args', None))
+    elif metric_name == "pairwise-stereoimbalance":
+        return PairwiseFeatures(*args, **kwargs, type="stereoimbalance", sample_rate=sample_rate, model_args=kwargs.get('stereoimbalance_args', None))
+    else:
+        raise ValueError(f"Unknown metric: {metric_name}")
+
+# Example usage:
+#metric_instance = metric_factory("pairwise-spectral")
+#```

utils/feature_extractors/AF_features_embedding.py

@@ -0,0 +1,166 @@
+import math
+from utils.feature_extractors.dsp_features import compute_log_rms_gated_max, compute_crest_factor, compute_stereo_width, compute_stereo_imbalance, compute_log_spread
+import torch
+
+class AF_fourier_embedding:
+    def __init__(self,
+                input_dim=8,
+                output_dim=64,
+                sigma=0.2,
+                log_rms_shift=-26.5, #calculated as the mean from the dataset
+                log_rms_scale=7.0, #calculated as the std from the dataset
+                crest_shift=16.7, #calculated as the mean from the dataset
+                crest_scale=6.3,
+                log_spread_shift=-20.0, #calculated as the mean from the dataset
+                log_spread_scale=20.0, #calculated as the std from the dataset
+                stereo_width_shift=0.28,
+                stereo_width_scale=0.39,
+                stereo_imbalance_shift=0.0,
+                stereo_imbalance_scale=0.35,
+                device="cpu"
+                ):
+        """
+        Deterministic Fourier feature transformer using fixed cosine-based projection
+        """
+
+        self.device = device
+        # Ensure output_dim is even and >= 2 * input_dim
+        self.output_dim = max(input_dim * 2, output_dim)
+        if self.output_dim % 2 != 0:
+            self.output_dim += 1
+        
+        self.input_dim = input_dim
+        self.sigma = sigma
+        
+        # Create deterministic projection matrix
+        self.projection = self._create_deterministic_projection(input_dim, self.output_dim // 2, sigma)
+        self.projection = self.projection.to(self.device)
+        
+        # Normalization factor
+        self.scale_factor = math.sqrt(2.0 / self.output_dim)
+
+        self.log_rms_shift = log_rms_shift
+        self.log_rms_scale = log_rms_scale
+        self.crest_shift = crest_shift
+        self.crest_scale = crest_scale
+        self.log_spread_shift = log_spread_shift
+        self.log_spread_scale = log_spread_scale
+        self.stereo_width_shift = stereo_width_shift
+        self.stereo_width_scale = stereo_width_scale
+        self.stereo_imbalance_shift = stereo_imbalance_shift
+        self.stereo_imbalance_scale = stereo_imbalance_scale
+    
+    def _create_deterministic_projection(self, input_dim, proj_dim, sigma):
+        """
+        Create a deterministic projection matrix using a cosine basis
+        """
+        # Cosine-based matrix (like DCT type-II)
+        projection = torch.zeros(input_dim, proj_dim)
+        for i in range(input_dim):
+            for j in range(proj_dim):
+                projection[i, j] = math.cos(math.pi * (i + 0.5) * (j + 1) / proj_dim)
+        
+        return projection * sigma
+    
+    def encode(self, x):
+
+        log_rms=compute_log_rms_gated_max(x)
+        crest_factor= compute_crest_factor(x)
+        log_spread= compute_log_spread(x)
+        stereo_width= compute_stereo_width(x)
+        stereo_imbalance= compute_stereo_imbalance(x)
+
+
+        log_rms_std, crest_factor_std, log_spread_std, stereo_width_std, stereo_imbalance_std = self.standardize_features(
+            log_rms, crest_factor, log_spread, stereo_width, stereo_imbalance
+        )
+
+        embedding= self.transform(
+            log_rms_std, crest_factor_std, log_spread_std, stereo_width_std, stereo_imbalance_std
+        )
+
+
+        return embedding, (log_rms, crest_factor, log_spread, stereo_width, stereo_imbalance)
+    
+    def decode(self, fourier_features):
+        """
+        Invert Fourier features back to original feature space
+        (approximate due to phase-only reconstruction)
+        """
+        reconstructed = self.inverse_transform(fourier_features)
+        
+        # Reshape back to original feature dimensions
+        log_rms= reconstructed[:,0:2]
+        crest_factor = reconstructed[:,2:4]
+        log_spread= reconstructed[:,4:6]
+        stereo_width = reconstructed[:,6:7]
+        stereo_imbalance = reconstructed[:,7:8]
+
+        log_rms, crest_factor, log_spread,stereo_width, stereo_imbalance = self.destandardize_features(
+            log_rms, crest_factor, log_spread, stereo_width, stereo_imbalance
+        )
+        
+        return log_rms, crest_factor, log_spread, stereo_width, stereo_imbalance
+
+    def standardize_features(self, log_rms, crest_factor, log_spread, stereo_width, stereo_imbalance):
+        """
+        Standardize features using pre-computed mean and std
+        """
+        log_rms = (log_rms - self.log_rms_shift) / self.log_rms_scale
+        crest_factor = (crest_factor - self.crest_shift) / self.crest_scale
+        log_spread = (log_spread - self.log_spread_shift) / self.log_spread_scale
+        stereo_width = (stereo_width - self.stereo_width_shift) / self.stereo_width_scale
+        stereo_imbalance = (stereo_imbalance - self.stereo_imbalance_shift) / self.stereo_imbalance_scale
+        
+        return log_rms, crest_factor, log_spread, stereo_width, stereo_imbalance
+    
+    def destandardize_features(self, log_rms, crest_factor, log_spread, stereo_width, stereo_imbalance):
+        """
+        Reverse standardization to get back to original feature space
+        """
+        log_rms = log_rms * self.log_rms_scale + self.log_rms_shift
+        crest_factor = crest_factor * self.crest_scale + self.crest_shift
+        log_spread = log_spread * self.log_spread_scale + self.log_spread_shift
+        stereo_width = stereo_width * self.stereo_width_scale + self.stereo_width_shift
+        stereo_imbalance = stereo_imbalance * self.stereo_imbalance_scale + self.stereo_imbalance_shift
+        
+        return log_rms, crest_factor, log_spread, stereo_width, stereo_imbalance
+        
+    def transform(self, log_rms, crest_factor,log_spread, stereo_width, stereo_imbalance):
+        """
+        Transform features using the stored projection matrix
+        """
+
+        flat_features=torch.cat([log_rms, crest_factor, log_spread, stereo_width.unsqueeze(-1), stereo_imbalance.unsqueeze(-1)], dim=-1)
+        
+        # Project and transform
+        projected = flat_features @ self.projection
+        cos_features = torch.cos(projected)
+        sin_features = torch.sin(projected)
+        
+        # Concatenate and normalize
+        return torch.cat([cos_features, sin_features], dim=-1) * self.scale_factor
+    
+    def inverse_transform(self, fourier_features):
+        """
+        Invert Fourier features back to original feature space
+        (approximate due to phase-only reconstruction)
+        """
+        # Split into cosine and sine components
+        feature_dim = fourier_features.shape[-1] // 2
+        cos_features = fourier_features[:, :feature_dim]
+        sin_features = fourier_features[:, feature_dim:]
+        
+        # Denormalize
+        cos_features = cos_features / self.scale_factor
+        sin_features = sin_features / self.scale_factor
+        
+        # Compute phase angles
+        phases = torch.atan2(sin_features, cos_features)
+        
+        # Use pseudo-inverse for approximate inversion
+        projection_pinv = torch.pinverse(self.projection)
+        reconstructed = phases @ projection_pinv
+        
+        return reconstructed
+

utils/feature_extractors/audio_features.py

@@ -0,0 +1,76 @@
+import numpy as np
+import librosa
+
+def compute_spectrogram(file_path, set_limit=False, limit_seconds=10):
+    """Compute the spectrogram of an audio file."""
+    y, sr = librosa.load(file_path, sr=None)
+    if set_limit:
+        y = y[:int(sr*limit_seconds)]
+    S = librosa.stft(y)
+    S_db = librosa.amplitude_to_db(np.abs(S), ref=np.max)
+    return S_db, sr, y
+
+def smooth_curve(data, window_size=30):
+    """Smooth a curve using a moving average."""
+    return np.convolve(data, np.ones(window_size)/window_size, mode='same')
+    
+def compute_crest_factor(y, sr, frame_length=2048, hop_length=512):
+    # compute crest factor by windowing
+    crest_factor_total = []
+    for i in range(0, len(y), hop_length):
+        frame = y[i:i+frame_length]
+        peak = np.max(np.abs(frame))
+        rms = np.sqrt(np.mean(frame**2))
+        cur_crest_factor = peak / (rms + 1e-6)
+        crest_factor_total.append(cur_crest_factor)
+    crest_factor = np.asarray(crest_factor_total)
+    return crest_factor
+
+def compute_audio_features(y, sr, smooth_window_size=1, average_time=True):
+    # features to compute: RMS energy, crest factor, dynamic spread, spectral centroid, spectral contrast, spectral flatness, spectral balance, and spectral bandwidth
+    y_mono = y.mean(axis=0)
+    spectral_centroid = smooth_curve(librosa.feature.spectral_centroid(y=y_mono, sr=sr)[0], window_size=smooth_window_size)
+    spectral_contrast = smooth_curve(np.mean(librosa.feature.spectral_contrast(y=y_mono, sr=sr), axis=0), window_size=smooth_window_size)
+    spectral_flatness = smooth_curve(librosa.feature.spectral_flatness(y=y_mono)[0], window_size=smooth_window_size)
+    rms_energy = smooth_curve(librosa.feature.rms(y=y_mono)[0], window_size=smooth_window_size)
+    # Crest Factor: Peak-to-RMS ratio
+    crest_factor = smooth_curve(compute_crest_factor(y_mono, sr, frame_length=2048, hop_length=512), window_size=smooth_window_size)
+    # Dynamic Spread
+    dynamic_spread = np.std(rms_energy)
+    # Spectral Balance (Low, Mid, High Frequency Energy Ratios)
+    spec = np.abs(librosa.stft(y_mono))**2
+    freqs = librosa.fft_frequencies(sr=sr)
+    low_energy = np.sum(spec[freqs < 200])
+    mid_energy = np.sum(spec[(freqs >= 200) & (freqs < 2000)])
+    high_energy = np.sum(spec[freqs >= 2000])
+    total_energy = low_energy + mid_energy + high_energy
+    spectral_balance = np.array([low_energy / total_energy, mid_energy / total_energy, high_energy / total_energy])
+    spectral_bandwidth = smooth_curve(librosa.feature.spectral_bandwidth(y=y_mono, sr=sr)[0], window_size=smooth_window_size)
+    # Stereo Width + Mid-Side Ratio
+    y_mid = (y[:len(y)//2] + y[len(y)//2:]) / 2
+    y_side = (y[:len(y)//2] - y[len(y)//2:]) / 2
+    stereo_width = np.mean(np.abs(y_side)) / (np.mean(np.abs(y_mid)) + 1e-6)
+    mid_side_ratio = np.mean(np.abs(y_mid)) / (np.mean(np.abs(y_side)) + 1e-6)
+    # average across time
+    if average_time:
+        spectral_centroid = np.mean(spectral_centroid)
+        spectral_flatness = np.mean(spectral_flatness)
+        spectral_contrast = np.mean(spectral_contrast)
+        spectral_flatness = np.mean(spectral_flatness)
+        rms_energy = np.mean(rms_energy)
+        crest_factor = np.mean(crest_factor)
+        spectral_bandwidth = np.mean(spectral_bandwidth)
+
+    # return as a dictionary
+    return {'spectral_centroid': spectral_centroid,
+            'spectral_contrast': spectral_contrast,
+            'spectral_flatness': spectral_flatness,
+            'rms_energy': rms_energy,
+            'crest_factor': crest_factor,
+            'dynamic_spread': dynamic_spread,
+            'spectral_balance': spectral_balance,
+            'spectral_bandwidth': spectral_bandwidth,
+            'stereo_width': stereo_width,
+            'mid_side_ratio': mid_side_ratio}
+    
+

utils/feature_extractors/dsp_features.py

@@ -0,0 +1,141 @@
+import torch
+
+def compute_rms(x: torch.Tensor, **kwargs):
+    """Compute root mean square energy.
+
+    Args:
+        x: (bs, 1, seq_len)
+
+    Returns:
+        rms: (bs, )
+    """
+    rms = torch.sqrt(torch.mean(x**2, dim=-1).clamp(min=1e-8))
+    return rms
+
+def compute_log_rms_gated_max(x: torch.Tensor, sample_rate=44100, **kwargs):
+    """Compute gated log RMS energy.
+
+    Frames the signal in 400 ms windows with 75% overlap, computes RMS,
+    discards frames with RMS < -60 dBFS, and averages the log-RMS.
+
+    If all frames in a given (batch, channel) are below -60 dBFS,
+    returns -60 for that entry.
+
+    Args:
+        x: Tensor of shape (bs, c, seq_len)
+
+    Returns:
+        log_rms: Tensor of shape (bs, c)
+    """
+    seg_size = int(sample_rate * 0.4)
+    hop_size = int(sample_rate * 0.1)
+
+    # (bs, c, num_frames, seg_size)
+    B, C, L = x.size()
+
+    assert C==1 or C==2
+
+    x_frames = x.unfold(2, seg_size, hop_size)  # (bs, c, num_frames, seg_size)
+
+    # RMS over last dimension (seg_size)
+    rms = torch.sqrt((x_frames ** 2).mean(dim=-1))  # (bs, c, num_frames)
+
+    # dB conversion
+    rms_db = 20 * torch.log10(rms.clamp(min=1e-8))  # (bs, c, num_frames)
+    #print("rms db shape", rms_db.shape)
+
+    #take the maximum RMS across all frames
+    rms_max = rms_db.max(dim=2)[0]  # (bs, c)
+    #print(f"RMS max shape: {rms_max.shape}")
+
+    return rms_max
+
+
+def compute_log_rms(x: torch.Tensor, **kwargs):
+    """Compute root mean square energy.
+
+    Args:
+        x: (bs, 1, seq_len)
+
+    Returns:
+        rms: (bs, )
+    """
+    rms=compute_rms(x)
+    return 20 * torch.log10(rms.clamp(min=1e-8))
+
+def compute_crest_factor(x: torch.Tensor, **kwargs):
+    """Compute crest factor as ratio of peak to rms energy in dB.
+
+    Args:
+        x: (bs, 2, seq_len)
+
+    """
+    num = torch.max(torch.abs(x), dim=-1)[0]
+    den = compute_rms(x).clamp(min=1e-8)
+    cf = 20 * torch.log10((num / den).clamp(min=1e-8))
+    return cf
+
+def compute_log_spread(x: torch.Tensor, **kwargs):
+    """Compute log spread as the mean difference between log magnitude of samples and log RMS.
+    
+    Args:
+        x: (bs, 1, seq_len)
+        
+    Returns:
+        log_spread: (bs, )
+    """
+    # Compute log RMS
+    log_rms = compute_log_rms(x).unsqueeze(-1)  # (bs, 1, 1)
+    
+    # Compute log magnitude of each sample
+    log_magnitude = 20 * torch.log10(torch.abs(x).clamp(min=1e-8))  # (bs, 1, seq_len)
+    
+    # Compute the difference and take the mean
+    log_spread = torch.mean(log_magnitude - log_rms, dim=-1).squeeze(1)  # (bs, )
+    
+    return log_spread
+
+
+def compute_stereo_width(x: torch.Tensor, **kwargs):
+    """Compute stereo width as ratio of energy in sum and difference signals.
+
+    Args:
+        x: (bs, 2, seq_len)
+
+    """
+    bs, chs, seq_len = x.size()
+
+    assert chs == 2, "Input must be stereo"
+
+    # compute sum and diff of stereo channels
+    x_sum = x[:, 0, :] + x[:, 1, :]
+    x_diff = x[:, 0, :] - x[:, 1, :]
+
+    # compute power of sum and diff
+    sum_energy = torch.mean(x_sum**2, dim=-1)
+    diff_energy = torch.mean(x_diff**2, dim=-1)
+
+    # compute stereo width as ratio
+    stereo_width = diff_energy / sum_energy.clamp(min=1e-8)
+
+    return stereo_width
+
+
+def compute_stereo_imbalance(x: torch.Tensor, **kwargs):
+    """Compute stereo imbalance as ratio of energy in left and right channels.
+
+    Args:
+        x: (bs, 2, seq_len)
+
+    Returns:
+        stereo_imbalance: (bs, )
+
+    """
+    left_energy = torch.mean(x[:, 0, :] ** 2, dim=-1)
+    right_energy = torch.mean(x[:, 1, :] ** 2, dim=-1)
+
+    stereo_imbalance = (right_energy - left_energy) / (
+        right_energy + left_energy
+    ).clamp(min=1e-8)
+
+    return stereo_imbalance

utils/feature_extractors/fx_encoder_plus_plus.py

@@ -0,0 +1,82 @@
+import torch
+from utils.fxencoder_plusplus import FxEncoderPlusPlus
+from utils.fxencoder_plusplus.model import get_model_path
+
+
+def load_model(model_name="default", model_path=None, device="cuda", auto_download=True, cache_dir=None):
+    """
+    Load FxEncoderPlusPlus model.
+    
+    Args:
+        model_name: Name of pretrained model ('default', 'musdb', 'medleydb')
+        model_path: Custom checkpoint path. If provided, ignores model_name
+        device: Device to load model on ('cuda' or 'cpu')
+        auto_download: Automatically download if model not found
+        cache_dir: Custom cache directory for downloaded models
+        
+    Returns:
+        Loaded FxEncoderPlusPlus model
+        
+    Examples:
+        # Load default base model
+        model = load_model()
+        
+        # Load musdb model
+        model = load_model(model_name="musdb")
+        
+        # Load medleydb model
+        model = load_model(model_name="medleydb")
+        
+        # Load custom checkpoint
+        model = load_model(model_path="/path/to/custom.pt")
+        
+        # List available models
+        list_available_models()
+    """
+    # Handle device
+    if device == "cuda" and not torch.cuda.is_available():
+        print("CUDA not available, using CPU")
+        device = "cpu"
+    
+    
+    # Determine model path
+    if model_path is None:
+        if auto_download:
+            model_path = get_model_path(model_name, cache_dir=cache_dir)
+        else:
+            raise ValueError("model_path is None and auto_download is False")
+    
+    # Create model instance with specified device
+    model = FxEncoderPlusPlus(
+        embed_dim=2048, 
+        audio_clap_module=False, 
+        text_clap_module=False,
+        extractor_module=False,
+        device=device
+    )
+    
+    # Load checkpoint
+    checkpoint = torch.load(model_path, map_location=device, weights_only=False)
+    
+    if "epoch" in checkpoint:
+        # resuming a train checkpoint w/ epoch and optimizer state
+        start_epoch = checkpoint["epoch"]
+        sd = checkpoint["state_dict"]
+        if next(iter(sd.items()))[0].startswith("module"):
+            sd = {k[len("module."):]: v for k, v in sd.items()}
+        model.load_state_dict(sd, strict=False )
+        print(f"Loaded checkpoint from epoch {start_epoch}")
+    else:
+        # loading a bare (model only) checkpoint for fine-tune or evaluation
+        model.load_state_dict(checkpoint)
+        print("Loaded model checkpoint")
+    
+    model.to(device)
+    model.eval()
+    
+    # Freeze parameters for inference
+    for param in model.parameters():
+        param.requires_grad = False
+    
+    print(f"Model loaded successfully on {device}")
+    return model

utils/feature_extractors/load_features.py

@@ -0,0 +1,260 @@
+import os
+import math
+import sys
+import torchaudio
+from importlib import import_module
+import torch
+import yaml
+import torch.nn.functional as F
+
+def load_fx_encoder_plusplus(model_args, device, *args, **kwargs):
+    from utils.feature_extractors.fx_encoder_plus_plus import load_model 
+
+    assert model_args is not None, "model_args must be provided for fx_encoder type"
+
+    ckpt_path=model_args.ckpt_path
+
+    model=load_model(
+        model_path=ckpt_path,
+        device=device,
+    )
+
+    def effects_encoder_fn(x):
+        assert x.ndim == 3, f"Input tensor x must be 2D, got {x.ndim}D"
+        assert x.shape[1] == 2, f"Input tensor x must have 2 channels, got {x.shape[1]} channels"
+
+        emb=model.get_fx_embedding(x)
+        #l2 normalize the embeddings
+        emb = torch.nn.functional.normalize(emb, p=2, dim=-1)
+
+        return emb
+        
+
+    return lambda x: effects_encoder_fn(x)
+
+def load_fx_encoder_plusplus_2048(model_args, device, *args, **kwargs):
+    from utils.feature_extractors.fx_encoder_plus_plus import load_model 
+
+    assert model_args is not None, "model_args must be provided for fx_encoder type"
+
+    ckpt_path=model_args.ckpt_path
+
+    model=load_model(
+        model_path=ckpt_path,
+        device=device,
+    )
+
+    def effects_encoder_fn(x):
+        assert x.ndim == 3, f"Input tensor x must be 2D, got {x.ndim}D"
+        assert x.shape[1] == 2, f"Input tensor x must have 2 channels, got {x.shape[1]} channels"
+
+        emb=model.fx_encoder(x)
+        emb=emb["embedding"]  # Extract the embedding from the dictionary
+
+        return emb
+        
+
+    return lambda x: effects_encoder_fn(x)
+
+def add_isotropic_noise(z, sigma=0.1):
+        """
+        z: [..., D] normalized embeddings (e.g., from CLAP or a regressor)
+        sigma: scale of noise to inject
+        Returns: z with orthogonal Gaussian noise added
+        """
+        n=torch.randn_like(z)  # isotropic noise
+        z_noisy = F.normalize(z + sigma * n, dim=-1)
+        return z_noisy
+
+
+
+
+def load_CLAP(model_args, device, *args, **kwargs):
+
+    #original_path = sys.path.copy()
+    from utils.laion_clap.hook import CLAP_Module
+    model= CLAP_Module(enable_fusion=False, amodel= 'HTSAT-base')
+    #sys.path = original_path
+
+
+    print("checkpoint",model_args.ckpt_path)
+    #print current sys.path
+    print("sys.path", sys.path)
+    model.load_ckpt(model_args.ckpt_path)
+    model.to(device)
+
+    normalize = model_args.normalize
+
+    if model_args.use_adaptor:
+        if model_args.adaptor_type == "MLP_CLAP_regressor":
+            from networks.MLP_CLAP_regressor import MLP_CLAP_regressor
+            adaptor=MLP_CLAP_regressor()
+            ckpt=torch.load(model_args.adaptor_checkpoint, map_location=device, weights_only=False)
+            adaptor.load_state_dict(ckpt["network"], strict=True)
+            adaptor.to(device)
+
+        
+
+
+    def clap_fn(x, type=None):
+        B, C, T = x.shape
+        if C > 1:
+            x= x.mean(dim=1, keepdim=True)  # Convert to mono if stereo
+
+        with torch.no_grad():
+            x=torchaudio.functional.resample(x, orig_freq=44100, new_freq=48000)
+            x= x.squeeze(1)  # Remove channel dimension for CLAP
+            emb=model.get_audio_embedding_from_data(x,use_tensor=True)
+
+            if type is not None:
+                if type == "wet":
+                    #print("wet mode")
+                    if model_args.use_adaptor:
+                        emb= adaptor(emb)  # Apply the adaptor if specified
+    
+            if model_args.add_noise:
+                emb= torch.nn.functional.normalize(emb, p=2, dim=-1)  # Normalize before adding noise
+                emb = add_isotropic_noise(emb, sigma=model_args.noise_sigma)
+
+            # Normalize the embeddings
+            if normalize:
+                emb = torch.nn.functional.normalize(emb, p=2, dim=-1)
+
+            return emb
+
+    return lambda x, type: clap_fn(x, type=type)
+
+
+             
+def load_fx_encoder(model_args, device, *args, **kwargs):
+    """
+    Load the FX Encoder model.
+    
+    Args:
+        model_args: Arguments for the FX Encoder model.
+        device: Device to load the model on (CPU or GPU).
+        
+    Returns:
+        a function that extracts features from audio.
+    """
+    assert model_args is not None, "model_args must be provided for fx_encoder type"
+
+    ckpt_path=model_args.ckpt_path
+
+    #from utils.feature_extractors.fx_encoder import load_effects_encoder
+    from utils.feature_extractors.networks import Effects_Encoder
+
+    def reload_weights(model, ckpt_path, device):
+        checkpoint = torch.load(ckpt_path, map_location=device)
+    
+        from collections import OrderedDict
+        new_state_dict = OrderedDict()
+        for k, v in checkpoint["model"].items():
+            name = k[7:] # remove `module.`
+            new_state_dict[name] = v
+        model.load_state_dict(new_state_dict, strict=False)
+
+
+    try:
+        with open(os.path.join('.','utils','feature_extractors', 'networks', 'configs.yaml'), 'r') as f:
+            configs = yaml.full_load(f)
+    except:
+        with open(model_args.config_file, 'r') as f:
+            configs = yaml.full_load(f)
+
+    cfg_enc = configs['Effects_Encoder']['default']
+
+    effects_encoder = Effects_Encoder(cfg_enc)
+    reload_weights(effects_encoder, ckpt_path, device)
+    effects_encoder.to(device)
+    effects_encoder.eval()
+
+    def effects_encoder_fn(x):
+        emb=effects_encoder(x)
+        #l2 normalize the embeddings
+        emb = torch.nn.functional.normalize(emb, p=2, dim=-1)
+        return emb
+        
+
+    return lambda x, *args: effects_encoder_fn(x)
+
+def load_AFxRep(model_args, device, sample_rate=44100, peak_scaling=True, *args, **kwargs):
+
+    assert model_args is not None, "model_args must be provided for AFxRep type"
+
+    ckpt_path=model_args.ckpt_path
+
+    config_path = os.path.join(os.path.dirname(ckpt_path), "config.yaml")
+
+    with open(config_path) as f:
+        config = yaml.safe_load(f)
+    
+    encoder_configs = config["model"]["init_args"]["encoder"]
+
+    module_path, class_name = encoder_configs["class_path"].rsplit(".", 1)
+    module_path = module_path.replace("lcap", "utils.st_ito")
+
+    module = import_module(module_path)
+
+    model = getattr(module, class_name)(**encoder_configs["init_args"])
+
+    checkpoint = torch.load(ckpt_path, map_location="cpu")
+
+    # load state dicts
+    state_dict = {}
+    for k, v in checkpoint["state_dict"].items():
+        if k.startswith("encoder"):
+            state_dict[k.replace("encoder.", "", 1)] = v
+
+    model.load_state_dict(state_dict)
+
+    model.eval()
+
+    model.to(device)
+
+    def wrapper_fn(x, sample_rate):
+
+        x=x.to(device)
+
+        #x=torch.transpose(x,-1,-2)
+
+        if sample_rate != 48000:
+            x=torchaudio.functional.resample(x, sample_rate, 48000)
+
+        bs= x.shape[0]
+        #peak normalization. I do it because this is what ST-ITO get_param_embeds does. Not sure if it is good that this representation is invariant to gain
+        if peak_scaling:
+            x_max=[]
+            for batch_idx in range(bs):
+                #x[batch_idx, ...] /= x[batch_idx, ...].abs().max().clamp(1e-8)
+                x_max.append( x[batch_idx, ...].abs().max().clamp(1e-8) )
+            
+            if x.ndim == 3:
+                x_max=torch.stack(x_max, dim=0).view(bs, 1, 1)
+            elif x.ndim == 2:
+                x_max=torch.stack(x_max, dim=0).view(bs, 1)
+    
+            x=x/ x_max
+
+        mid_embeddings, side_embeddings = model(x)
+
+        # check for nan
+        if torch.isnan(mid_embeddings).any():
+            print("Warning: NaNs found in mid_embeddings")
+            mid_embeddings = torch.nan_to_num(mid_embeddings)
+        elif torch.isnan(side_embeddings).any():
+            print("Warning: NaNs found in side_embeddings")
+            side_embeddings = torch.nan_to_num(side_embeddings)
+
+        mid_embeddings = torch.nn.functional.normalize(mid_embeddings, p=2, dim=-1)
+        side_embeddings = torch.nn.functional.normalize(side_embeddings, p=2, dim=-1)
+        
+        embeddings_all= torch.cat([mid_embeddings, side_embeddings], dim=-1)
+
+        return embeddings_all
+
+    feat_extractor = lambda x, *args: wrapper_fn(x, sample_rate=sample_rate)
+
+    return feat_extractor
+
+

utils/feature_extractors/networks/__init__.py

@@ -0,0 +1,2 @@
+from .architectures import *
+from .network_utils import *

utils/feature_extractors/networks/architectures.py

@@ -0,0 +1,163 @@
+""" 
+    Implementation of neural networks used in the task 'Music Mastering Style Transfer'
+        - 'Effects Encoder'
+        - 'Mastering Style Transfer'
+        - 'Differentiable Mastering Style Transfer'
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.nn.init as init
+#import dasp_pytorch
+
+import os
+import sys
+import time
+
+# compute receptive field
+def compute_receptive_field(kernels, strides, dilations):
+    rf = 0
+    for i in range(len(kernels)):
+        rf += rf * strides[i] + (kernels[i]-strides[i]) * dilations[i]
+    return rf
+
+
+# Encoder of music effects for contrastive learning of music effects
+class Effects_Encoder(nn.Module):
+    def __init__(self, config):
+        super(Effects_Encoder, self).__init__()
+        # input is stereo channeled audio
+        config["channels"].insert(0, 2)
+
+        # encoder layers
+        encoder = []
+        for i in range(len(config["kernels"])):
+            if config["conv_block"]=='res':
+                encoder.append(Res_ConvBlock(dimension=1, \
+                                                in_channels=config["channels"][i], \
+                                                out_channels=config["channels"][i+1], \
+                                                kernel_size=config["kernels"][i], \
+                                                stride=config["strides"][i], \
+                                                padding="SAME", \
+                                                dilation=config["dilation"][i], \
+                                                norm=config["norm"], \
+                                                activation=config["activation"], \
+                                                last_activation=config["activation"]))
+        self.encoder = nn.Sequential(*encoder)
+
+        # pooling method
+        self.glob_pool = nn.AdaptiveAvgPool1d(1)
+
+
+    # network forward operation
+    def forward(self, input):
+        enc_output = self.encoder(input)
+        glob_pooled = self.glob_pool(enc_output).squeeze(-1)
+
+        # outputs c feature
+        return glob_pooled
+
+
+# Residual Block
+    # the input is added after the first convolutional layer, retaining its original channel size
+    # therefore, the second convolutional layer's output channel may differ
+class Res_ConvBlock(nn.Module):
+    def __init__(self, dimension, \
+                        in_channels, out_channels, \
+                        kernel_size, \
+                        stride=1, padding="SAME", \
+                        dilation=1, \
+                        bias=True, \
+                        norm="batch", \
+                        activation="relu", last_activation="relu", \
+                        mode="conv"):
+        super(Res_ConvBlock, self).__init__()
+
+        if dimension==1:
+            self.conv1 = Conv1d_layer(in_channels, in_channels, kernel_size, padding=padding, dilation=dilation, bias=bias, norm=norm, activation=activation)
+            self.conv2 = Conv1d_layer(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, bias=bias, norm=norm, activation=last_activation, mode=mode)
+
+    def forward(self, input):
+        c1_out = self.conv1(input) + input
+        c2_out = self.conv2(c1_out)
+        return c2_out
+
+
+# 1-dimensional convolutional layer
+# in the order of conv -> norm -> activation
+class Conv1d_layer(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, \
+                                    stride=1, \
+                                    padding="SAME", dilation=1, bias=True, \
+                                    norm="batch", activation="relu", \
+                                    mode="conv"):
+        super(Conv1d_layer, self).__init__()
+        
+        self.conv1d = nn.Sequential()
+
+        ''' padding '''
+        if mode=="deconv":
+            padding = int(dilation * (kernel_size-1) / 2)
+            out_padding = 0 if stride==1 else 1
+        elif mode=="conv" or "alias_free" in mode:
+            if padding == "SAME":
+                pad = int((kernel_size-1) * dilation)
+                l_pad = int(pad//2)
+                r_pad = pad - l_pad
+                padding_area = (l_pad, r_pad)
+            elif padding == "VALID":
+                padding_area = (0, 0)
+            else:
+                pass
+
+        ''' convolutional layer '''
+        if mode=="deconv":
+            self.conv1d.add_module("deconv1d", nn.ConvTranspose1d(in_channels, out_channels, kernel_size, \
+                                                            stride=stride, padding=padding, output_padding=out_padding, \
+                                                            dilation=dilation, \
+                                                            bias=bias))
+        elif mode=="conv":
+            self.conv1d.add_module(f"{mode}1d_pad", nn.ReflectionPad1d(padding_area))
+            self.conv1d.add_module(f"{mode}1d", nn.Conv1d(in_channels, out_channels, kernel_size, \
+                                                            stride=stride, padding=0, \
+                                                            dilation=dilation, \
+                                                            bias=bias))
+        elif "alias_free" in mode:
+            if "up" in mode:
+                up_factor = stride * 2
+                down_factor = 2
+            elif "down" in mode:
+                up_factor = 2
+                down_factor = stride * 2
+            else:
+                raise ValueError("choose alias-free method : 'up' or 'down'")
+            # procedure : conv -> upsample -> lrelu -> low-pass filter -> downsample
+            # the torchaudio.transforms.Resample's default resampling_method is 'sinc_interpolation' which performs low-pass filter during the process
+            # details at https://pytorch.org/audio/stable/transforms.html
+            self.conv1d.add_module(f"{mode}1d_pad", nn.ReflectionPad1d(padding_area))
+            self.conv1d.add_module(f"{mode}1d", nn.Conv1d(in_channels, out_channels, kernel_size, \
+                                                            stride=1, padding=0, \
+                                                            dilation=dilation, \
+                                                            bias=bias))
+            self.conv1d.add_module(f"{mode}upsample", torchaudio.transforms.Resample(orig_freq=1, new_freq=up_factor))
+            self.conv1d.add_module(f"{mode}lrelu", nn.LeakyReLU())
+            self.conv1d.add_module(f"{mode}downsample", torchaudio.transforms.Resample(orig_freq=down_factor, new_freq=1))
+
+        ''' normalization '''
+        if norm=="batch":
+            self.conv1d.add_module("batch_norm", nn.BatchNorm1d(out_channels))
+            # self.conv1d.add_module("batch_norm", nn.SyncBatchNorm(out_channels))
+
+        ''' activation '''
+        if 'alias_free' not in mode:
+            if activation=="relu":
+                self.conv1d.add_module("relu", nn.ReLU())
+            elif activation=="lrelu":
+                self.conv1d.add_module("lrelu", nn.LeakyReLU())
+
+
+    def forward(self, input):
+        # input shape should be : batch x channel x height x width
+        output = self.conv1d(input)
+        return output
+

utils/feature_extractors/networks/configs.yaml

@@ -0,0 +1,14 @@
+# model architecture configurations
+
+# Music Effects Encoder
+Effects_Encoder:
+
+  default:
+    channels: [16, 32, 64, 128, 256, 256, 512, 512, 1024, 1024, 2048, 2048]
+    kernels: [25, 25, 15, 15, 10, 10, 10, 10, 5, 5, 5, 5]
+    strides: [4, 4, 2, 2, 2, 2, 2, 2, 2, 2, 1, 1]
+    dilation: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]
+    bias: True
+    norm: 'batch'
+    conv_block: 'res'
+    activation: "relu"

utils/feature_extractors/networks/network_utils.py

@@ -0,0 +1,254 @@
+"""
+    Utility File
+        containing functions for neural networks
+"""
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.nn.init as init
+import torch
+import torchaudio
+
+
+
+# 2-dimensional convolutional layer
+# in the order of conv -> norm -> activation
+class Conv2d_layer(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, \
+                                    stride=1, \
+                                    padding="SAME", dilation=(1,1), bias=True, \
+                                    norm="batch", activation="relu", \
+                                    mode="conv"):
+        super(Conv2d_layer, self).__init__()
+
+        self.conv2d = nn.Sequential()
+
+        if isinstance(kernel_size, int):
+            kernel_size = [kernel_size, kernel_size]
+        if isinstance(stride, int):
+            stride = [stride, stride]
+        if isinstance(dilation, int):
+            dilation = [dilation, dilation]
+
+        ''' padding '''
+        if mode=="deconv":
+            padding = tuple(int((current_kernel - 1)/2) for current_kernel in kernel_size)
+            out_padding = tuple(0 if current_stride == 1 else 1 for current_stride in stride)
+        elif mode=="conv":
+            if padding == "SAME":
+                f_pad = int((kernel_size[0]-1) * dilation[0])
+                t_pad = int((kernel_size[1]-1) * dilation[1])
+                t_l_pad = int(t_pad//2)
+                t_r_pad = t_pad - t_l_pad
+                f_l_pad = int(f_pad//2)
+                f_r_pad = f_pad - f_l_pad
+                padding_area = (t_l_pad, t_r_pad, f_l_pad, f_r_pad)
+            elif padding == "VALID":
+                padding = 0
+            else:
+                pass
+
+        ''' convolutional layer '''
+        if mode=="deconv":
+            self.conv2d.add_module("deconv2d", nn.ConvTranspose2d(in_channels, out_channels, \
+                                                            (kernel_size[0], kernel_size[1]), \
+                                                            stride=stride, \
+                                                            padding=padding, output_padding=out_padding, \
+                                                            dilation=dilation, \
+                                                            bias=bias))
+        elif mode=="conv":
+            self.conv2d.add_module(f"{mode}2d_pad", nn.ReflectionPad2d(padding_area))
+            self.conv2d.add_module(f"{mode}2d", nn.Conv2d(in_channels, out_channels, \
+                                                            (kernel_size[0], kernel_size[1]), \
+                                                            stride=stride, \
+                                                            padding=0, \
+                                                            dilation=dilation, \
+                                                            bias=bias))
+        
+        ''' normalization '''
+        if norm=="batch":
+            self.conv2d.add_module("batch_norm", nn.BatchNorm2d(out_channels))
+
+        ''' activation '''
+        if activation=="relu":
+            self.conv2d.add_module("relu", nn.ReLU())
+        elif activation=="lrelu":
+            self.conv2d.add_module("lrelu", nn.LeakyReLU())
+
+
+    def forward(self, input):
+        # input shape should be : batch x channel x height x width
+        output = self.conv2d(input)
+        return output
+
+
+
+# 1-dimensional convolutional layer
+# in the order of conv -> norm -> activation
+class Conv1d_layer(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, \
+                                    stride=1, \
+                                    padding="SAME", dilation=1, bias=True, \
+                                    norm="batch", activation="relu", \
+                                    mode="conv"):
+        super(Conv1d_layer, self).__init__()
+        
+        self.conv1d = nn.Sequential()
+
+        ''' padding '''
+        if mode=="deconv":
+            padding = int(dilation * (kernel_size-1) / 2)
+            out_padding = 0 if stride==1 else 1
+        elif mode=="conv" or "alias_free" in mode:
+            if padding == "SAME":
+                pad = int((kernel_size-1) * dilation)
+                l_pad = int(pad//2)
+                r_pad = pad - l_pad
+                padding_area = (l_pad, r_pad)
+            elif padding == "VALID":
+                padding_area = (0, 0)
+            else:
+                pass
+
+        ''' convolutional layer '''
+        if mode=="deconv":
+            self.conv1d.add_module("deconv1d", nn.ConvTranspose1d(in_channels, out_channels, kernel_size, \
+                                                            stride=stride, padding=padding, output_padding=out_padding, \
+                                                            dilation=dilation, \
+                                                            bias=bias))
+        elif mode=="conv":
+            self.conv1d.add_module(f"{mode}1d_pad", nn.ReflectionPad1d(padding_area))
+            self.conv1d.add_module(f"{mode}1d", nn.Conv1d(in_channels, out_channels, kernel_size, \
+                                                            stride=stride, padding=0, \
+                                                            dilation=dilation, \
+                                                            bias=bias))
+        elif "alias_free" in mode:
+            if "up" in mode:
+                up_factor = stride * 2
+                down_factor = 2
+            elif "down" in mode:
+                up_factor = 2
+                down_factor = stride * 2
+            else:
+                raise ValueError("choose alias-free method : 'up' or 'down'")
+            # procedure : conv -> upsample -> lrelu -> low-pass filter -> downsample
+            # the torchaudio.transforms.Resample's default resampling_method is 'sinc_interpolation' which performs low-pass filter during the process
+            # details at https://pytorch.org/audio/stable/transforms.html
+            self.conv1d.add_module(f"{mode}1d_pad", nn.ReflectionPad1d(padding_area))
+            self.conv1d.add_module(f"{mode}1d", nn.Conv1d(in_channels, out_channels, kernel_size, \
+                                                            stride=1, padding=0, \
+                                                            dilation=dilation, \
+                                                            bias=bias))
+            self.conv1d.add_module(f"{mode}upsample", torchaudio.transforms.Resample(orig_freq=1, new_freq=up_factor))
+            self.conv1d.add_module(f"{mode}lrelu", nn.LeakyReLU())
+            self.conv1d.add_module(f"{mode}downsample", torchaudio.transforms.Resample(orig_freq=down_factor, new_freq=1))
+
+        ''' normalization '''
+        if norm=="batch":
+            self.conv1d.add_module("batch_norm", nn.BatchNorm1d(out_channels))
+            # self.conv1d.add_module("batch_norm", nn.SyncBatchNorm(out_channels))
+
+        ''' activation '''
+        if 'alias_free' not in mode:
+            if activation=="relu":
+                self.conv1d.add_module("relu", nn.ReLU())
+            elif activation=="lrelu":
+                self.conv1d.add_module("lrelu", nn.LeakyReLU())
+
+
+    def forward(self, input):
+        # input shape should be : batch x channel x height x width
+        output = self.conv1d(input)
+        return output
+
+
+
+# Residual Block
+    # the input is added after the first convolutional layer, retaining its original channel size
+    # therefore, the second convolutional layer's output channel may differ
+class Res_ConvBlock(nn.Module):
+    def __init__(self, dimension, \
+                        in_channels, out_channels, \
+                        kernel_size, \
+                        stride=1, padding="SAME", \
+                        dilation=1, \
+                        bias=True, \
+                        norm="batch", \
+                        activation="relu", last_activation="relu", \
+                        mode="conv"):
+        super(Res_ConvBlock, self).__init__()
+
+        if dimension==1:
+            self.conv1 = Conv1d_layer(in_channels, in_channels, kernel_size, padding=padding, dilation=dilation, bias=bias, norm=norm, activation=activation)
+            self.conv2 = Conv1d_layer(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, bias=bias, norm=norm, activation=last_activation, mode=mode)
+        elif dimension==2:
+            self.conv1 = Conv2d_layer(in_channels, in_channels, kernel_size, padding=padding, dilation=dilation, bias=bias, norm=norm, activation=activation)
+            self.conv2 = Conv2d_layer(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, bias=bias, norm=norm, activation=last_activation, mode=mode)
+
+
+    def forward(self, input):
+        c1_out = self.conv1(input) + input
+        c2_out = self.conv2(c1_out)
+        return c2_out
+
+
+
+# Convoluaionl Block
+    # consists of multiple (number of layer_num) convolutional layers
+    # only the final convoluational layer outputs the desired 'out_channels'
+class ConvBlock(nn.Module):
+    def __init__(self, dimension, layer_num, \
+                        in_channels, out_channels, \
+                        kernel_size, \
+                        stride=1, padding="SAME", \
+                        dilation=1, \
+                        bias=True, \
+                        norm="batch", \
+                        activation="relu", last_activation="relu", \
+                        mode="conv"):
+        super(ConvBlock, self).__init__()
+
+        conv_block = []
+        if dimension==1:
+            for i in range(layer_num-1):
+                conv_block.append(Conv1d_layer(in_channels, in_channels, kernel_size, padding=padding, dilation=dilation, bias=bias, norm=norm, activation=activation))
+            conv_block.append(Conv1d_layer(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, bias=bias, norm=norm, activation=last_activation, mode=mode))
+        elif dimension==2:
+            for i in range(layer_num-1):
+                conv_block.append(Conv2d_layer(in_channels, in_channels, kernel_size, padding=padding, dilation=dilation, bias=bias, norm=norm, activation=activation))
+            conv_block.append(Conv2d_layer(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, bias=bias, norm=norm, activation=last_activation, mode=mode))
+        self.conv_block = nn.Sequential(*conv_block)
+
+
+    def forward(self, input):
+        return self.conv_block(input)
+
+
+# Feature-wise Linear Modulation
+class FiLM(nn.Module):
+    def __init__(self, condition_len=2048, feature_len=1024):
+        super(FiLM, self).__init__()
+        self.film_fc = nn.Linear(condition_len, feature_len*2)
+        self.feat_len = feature_len
+
+
+    def forward(self, feature, condition, sefa=None):
+        # SeFA
+        if sefa:
+            weight = self.film_fc.weight.T
+            weight = weight / torch.linalg.norm((weight+1e-07), dim=0, keepdims=True)
+            eigen_values, eigen_vectors = torch.eig(torch.matmul(weight, weight.T), eigenvectors=True)
+
+            ####### custom parameters #######
+            chosen_eig_idx = sefa[0]
+            alpha = eigen_values[chosen_eig_idx][0] * sefa[1]
+            #################################
+
+            An = eigen_vectors[chosen_eig_idx].repeat(condition.shape[0], 1)
+            alpha_An = alpha * An
+
+            condition += alpha_An
+
+        film_factor = self.film_fc(condition).unsqueeze(-1)
+        r, b = torch.split(film_factor, self.feat_len, dim=1)
+        return r*feature + b
+

utils/feature_extractors/networks/pytorch_utils.py

@@ -0,0 +1,256 @@
+"""
+From: https://github.com/qiuqiangkong/audioset_tagging_cnn/blob/master/pytorch/pytorch_utils.py
+
+Copyright (c) 2018-2020 Qiuqiang Kong
+"""
+import numpy as np
+import time
+import torch
+import torch.nn as nn
+
+
+def move_data_to_device(x, device):
+    if 'float' in str(x.dtype):
+        x = torch.Tensor(x)
+    elif 'int' in str(x.dtype):
+        x = torch.LongTensor(x)
+    else:
+        return x
+
+    return x.to(device)
+
+
+def do_mixup(x, mixup_lambda):
+    """Mixup x of even indexes (0, 2, 4, ...) with x of odd indexes 
+    (1, 3, 5, ...).
+
+    Args:
+      x: (batch_size * 2, ...)
+      mixup_lambda: (batch_size * 2,)
+
+    Returns:
+      out: (batch_size, ...)
+    """
+    out = (x[0 :: 2].transpose(0, -1) * mixup_lambda[0 :: 2] + \
+        x[1 :: 2].transpose(0, -1) * mixup_lambda[1 :: 2]).transpose(0, -1)
+    return out
+    
+
+def append_to_dict(dict, key, value):
+    if key in dict.keys():
+        dict[key].append(value)
+    else:
+        dict[key] = [value]
+
+
+def forward(model, generator, return_input=False, 
+    return_target=False):
+    """Forward data to a model.
+    
+    Args: 
+      model: object
+      generator: object
+      return_input: bool
+      return_target: bool
+
+    Returns:
+      audio_name: (audios_num,)
+      clipwise_output: (audios_num, classes_num)
+      (ifexist) segmentwise_output: (audios_num, segments_num, classes_num)
+      (ifexist) framewise_output: (audios_num, frames_num, classes_num)
+      (optional) return_input: (audios_num, segment_samples)
+      (optional) return_target: (audios_num, classes_num)
+    """
+    output_dict = {}
+    device = next(model.parameters()).device
+    time1 = time.time()
+
+    # Forward data to a model in mini-batches
+    for n, batch_data_dict in enumerate(generator):
+        print(n)
+        batch_waveform = move_data_to_device(batch_data_dict['waveform'], device)
+        
+        with torch.no_grad():
+            model.eval()
+            batch_output = model(batch_waveform)
+
+        append_to_dict(output_dict, 'audio_name', batch_data_dict['audio_name'])
+
+        append_to_dict(output_dict, 'clipwise_output', 
+            batch_output['clipwise_output'].data.cpu().numpy())
+
+        if 'segmentwise_output' in batch_output.keys():
+            append_to_dict(output_dict, 'segmentwise_output', 
+                batch_output['segmentwise_output'].data.cpu().numpy())
+
+        if 'framewise_output' in batch_output.keys():
+            append_to_dict(output_dict, 'framewise_output', 
+                batch_output['framewise_output'].data.cpu().numpy())
+            
+        if return_input:
+            append_to_dict(output_dict, 'waveform', batch_data_dict['waveform'])
+            
+        if return_target:
+            if 'target' in batch_data_dict.keys():
+                append_to_dict(output_dict, 'target', batch_data_dict['target'])
+
+        if n % 10 == 0:
+            print(' --- Inference time: {:.3f} s / 10 iterations ---'.format(
+                time.time() - time1))
+            time1 = time.time()
+
+    for key in output_dict.keys():
+        output_dict[key] = np.concatenate(output_dict[key], axis=0)
+
+    return output_dict
+
+
+def interpolate(x, ratio):
+    """Interpolate data in time domain. This is used to compensate the 
+    resolution reduction in downsampling of a CNN.
+    
+    Args:
+      x: (batch_size, time_steps, classes_num)
+      ratio: int, ratio to interpolate
+
+    Returns:
+      upsampled: (batch_size, time_steps * ratio, classes_num)
+    """
+    (batch_size, time_steps, classes_num) = x.shape
+    upsampled = x[:, :, None, :].repeat(1, 1, ratio, 1)
+    upsampled = upsampled.reshape(batch_size, time_steps * ratio, classes_num)
+    return upsampled
+
+
+def pad_framewise_output(framewise_output, frames_num):
+    """Pad framewise_output to the same length as input frames. The pad value 
+    is the same as the value of the last frame.
+
+    Args:
+      framewise_output: (batch_size, frames_num, classes_num)
+      frames_num: int, number of frames to pad
+
+    Outputs:
+      output: (batch_size, frames_num, classes_num)
+    """
+    pad = framewise_output[:, -1 :, :].repeat(1, frames_num - framewise_output.shape[1], 1)
+    """tensor for padding"""
+
+    output = torch.cat((framewise_output, pad), dim=1)
+    """(batch_size, frames_num, classes_num)"""
+
+    return output
+
+
+def count_parameters(model):
+    return sum(p.numel() for p in model.parameters() if p.requires_grad)
+
+
+def count_flops(model, audio_length):
+    """Count flops. Code modified from others' implementation.
+    """
+    multiply_adds = True
+    list_conv2d=[]
+    def conv2d_hook(self, input, output):
+        batch_size, input_channels, input_height, input_width = input[0].size()
+        output_channels, output_height, output_width = output[0].size()
+ 
+        kernel_ops = self.kernel_size[0] * self.kernel_size[1] * (self.in_channels / self.groups) * (2 if multiply_adds else 1)
+        bias_ops = 1 if self.bias is not None else 0
+ 
+        params = output_channels * (kernel_ops + bias_ops)
+        flops = batch_size * params * output_height * output_width
+ 
+        list_conv2d.append(flops)
+
+    list_conv1d=[]
+    def conv1d_hook(self, input, output):
+        batch_size, input_channels, input_length = input[0].size()
+        output_channels, output_length = output[0].size()
+ 
+        kernel_ops = self.kernel_size[0] * (self.in_channels / self.groups) * (2 if multiply_adds else 1)
+        bias_ops = 1 if self.bias is not None else 0
+ 
+        params = output_channels * (kernel_ops + bias_ops)
+        flops = batch_size * params * output_length
+ 
+        list_conv1d.append(flops)
+ 
+    list_linear=[] 
+    def linear_hook(self, input, output):
+        batch_size = input[0].size(0) if input[0].dim() == 2 else 1
+ 
+        weight_ops = self.weight.nelement() * (2 if multiply_adds else 1)
+        bias_ops = self.bias.nelement()
+ 
+        flops = batch_size * (weight_ops + bias_ops)
+        list_linear.append(flops)
+ 
+    list_bn=[] 
+    def bn_hook(self, input, output):
+        list_bn.append(input[0].nelement() * 2)
+ 
+    list_relu=[] 
+    def relu_hook(self, input, output):
+        list_relu.append(input[0].nelement() * 2)
+ 
+    list_pooling2d=[]
+    def pooling2d_hook(self, input, output):
+        batch_size, input_channels, input_height, input_width = input[0].size()
+        output_channels, output_height, output_width = output[0].size()
+ 
+        kernel_ops = self.kernel_size * self.kernel_size
+        bias_ops = 0
+        params = output_channels * (kernel_ops + bias_ops)
+        flops = batch_size * params * output_height * output_width
+ 
+        list_pooling2d.append(flops)
+
+    list_pooling1d=[]
+    def pooling1d_hook(self, input, output):
+        batch_size, input_channels, input_length = input[0].size()
+        output_channels, output_length = output[0].size()
+ 
+        kernel_ops = self.kernel_size[0]
+        bias_ops = 0
+        
+        params = output_channels * (kernel_ops + bias_ops)
+        flops = batch_size * params * output_length
+ 
+        list_pooling2d.append(flops)
+ 
+    def foo(net):
+        childrens = list(net.children())
+        if not childrens:
+            if isinstance(net, nn.Conv2d):
+                net.register_forward_hook(conv2d_hook)
+            elif isinstance(net, nn.Conv1d):
+                net.register_forward_hook(conv1d_hook)
+            elif isinstance(net, nn.Linear):
+                net.register_forward_hook(linear_hook)
+            elif isinstance(net, nn.BatchNorm2d) or isinstance(net, nn.BatchNorm1d):
+                net.register_forward_hook(bn_hook)
+            elif isinstance(net, nn.ReLU):
+                net.register_forward_hook(relu_hook)
+            elif isinstance(net, nn.AvgPool2d) or isinstance(net, nn.MaxPool2d):
+                net.register_forward_hook(pooling2d_hook)
+            elif isinstance(net, nn.AvgPool1d) or isinstance(net, nn.MaxPool1d):
+                net.register_forward_hook(pooling1d_hook)
+            else:
+                print('Warning: flop of module {} is not counted!'.format(net))
+            return
+        for c in childrens:
+            foo(c)
+
+    # Register hook
+    foo(model)
+    
+    device = device = next(model.parameters()).device
+    input = torch.rand(1, audio_length).to(device)
+
+    out = model(input)
+ 
+    total_flops = sum(list_conv2d) + sum(list_conv1d) + sum(list_linear) + \
+        sum(list_bn) + sum(list_relu) + sum(list_pooling2d) + sum(list_pooling1d)
+    
+    return total_flops

utils/fx_normalization/features.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1271839aa88a321deb76fb3e4f3059eeee26cf64c4071cd116e1e4a69d896d74
+size 1574572

utils/fx_normalization/fxnorm.py

@@ -0,0 +1,154 @@
+import torch
+import numpy as np
+
+from utils.training_utils import  Gauss_smooth_vectorized,  prepare_smooth_filter
+
+def T602logmag(t60, sample_rate=44100, hop_length=512):
+    return 6.908 / (t60 * (sample_rate / hop_length))  # Convert T60 to delta log magnitude
+
+from utils.data_utils import apply_RMS_normalization
+
+
+class FxNormAug:
+
+    def __init__(self,
+                sample_rate=44100,  # Sample rate of the audio
+                device="cuda" if torch.cuda.is_available() else "cpu",
+                mode="train",  # Mode can be "train" or "eval"
+                seed=42,
+                features_path="features_tency1_4instr_v4.npy",  # Path to the features file
+                ):
+
+        torch.random.manual_seed(seed)
+
+        #the path is in the same directory as this file
+        self.features_path = features_path
+
+        self.sample_rate = sample_rate
+        self.device = device
+        self.train_setup()
+
+        self.EQ_normalize_setup()  # Initialize EQ normalization function
+
+
+    def train_setup(self):
+        #hardcoded setuo used for training
+
+        self.RMS_norm=-25  # Target RMS level in dB, used for normalization
+
+    def EQ_normalize_setup(self ):
+
+        features_mean = np.load(self.features_path, allow_pickle='TRUE')[()]
+
+        target_cuves_original= {
+            "vocals": torch.tensor(features_mean["eq"]["vocals"]).to(torch.float32).to(self.device),
+            "drums": torch.tensor(features_mean["eq"]["drums"]).to(torch.float32).to(self.device),
+            "bass": torch.tensor(features_mean["eq"]["bass"]).to(torch.float32).to(self.device),
+            "other": torch.tensor(features_mean["eq"]["other"]).to(torch.float32).to(self.device),
+        }
+
+
+        nfft=4096 # FFT size hardcoded
+        nfft_orig = 65536  # FFT size for the smoothing filter
+
+        win_length=2048  # Window length hardcoded
+        hop_length=1024  # Hop length hardcoded
+
+        window = torch.sqrt(torch.hann_window(win_length, device=self.device))
+        window_energy = window.pow(2).sum().sqrt()  # Energy of the window
+
+        freqs = torch.fft.rfftfreq(nfft, d=1.0).to(self.device)
+        freqs_Hz=torch.fft.rfftfreq(nfft, d=1.0).to(self.device) * self.sample_rate
+
+        smooth_filter = prepare_smooth_filter(freqs_Hz, Noct=3).to(self.device)  # Prepare the smoothing filter
+
+        freqs_Hz_orig=torch.fft.rfftfreq(nfft_orig, d=1.0).to(self.device) * self.sample_rate
+        smooth_filter_orig = prepare_smooth_filter(freqs_Hz_orig, Noct=3).to(self.device)  # Prepare the smoothing filter
+
+        def downsample_curve(x):
+            return torch.nn.functional.interpolate(
+                x.unsqueeze(0).unsqueeze(0),
+                size=(nfft // 2 + 1,),
+                mode='linear',
+                align_corners=False
+            ).squeeze(0).squeeze(0)
+
+        target_curves = {
+            "vocals": downsample_curve(Gauss_smooth_vectorized(target_cuves_original["vocals"], freqs_Hz_orig, Noct=3, smooth_filter=smooth_filter_orig)),
+            "drums": downsample_curve(Gauss_smooth_vectorized(target_cuves_original["drums"], freqs_Hz_orig, Noct=3, smooth_filter=smooth_filter_orig)),
+            "bass": downsample_curve(Gauss_smooth_vectorized(target_cuves_original["bass"], freqs_Hz_orig, Noct=3, smooth_filter=smooth_filter_orig)),
+            "other": downsample_curve(Gauss_smooth_vectorized(target_cuves_original["other"], freqs_Hz_orig, Noct=3, smooth_filter=smooth_filter_orig)),
+        }
+
+
+        def EQ_normalize_fn(x):
+
+            shape= x.shape
+
+            target_curves_tensor=torch.zeros((shape[0], nfft // 2 + 1), device=self.device, dtype=torch.float32)
+
+            for i in range(shape[0]):
+                track_class = "other"
+
+                assert track_class in target_curves, f"track_class {track_class} not found in target_curves"
+                target_curves_tensor[i] = target_curves[track_class] 
+
+            x=x.view(-1, shape[-1])
+            #ensure x.shape[-1] is divisible by hop_length
+            if x.shape[-1] % hop_length != 0:
+                # Pad the input to make it divisible by hop_length
+                pad_length = hop_length - (x.shape[-1] % hop_length)
+                x = torch.nn.functional.pad(x, (0, pad_length), mode='constant', value=0)
+            X=torch.stft(x, n_fft=nfft, hop_length=hop_length, win_length=win_length, window=window, return_complex=True)/ window_energy
+            X_pow=X.abs().pow(2)
+            X_mean= torch.sqrt(X_pow.mean(dim=-1, keepdim=False))  # Mean power spectrum
+
+
+
+            ratio= target_curves_tensor  / (X_mean + 1e-6)
+
+            ratio = torch.clamp(ratio, max=10.0**(20.0/20.0))
+
+
+            ratio_smooth = Gauss_smooth_vectorized(ratio, freqs_Hz, Noct=3, smooth_filter=smooth_filter)
+
+
+
+            X= X * ratio_smooth.unsqueeze(-1)
+
+            X_unnormalized=X* window_energy
+
+            x_reconstructed = torch.istft(X_unnormalized, 
+                             n_fft=nfft, 
+                             hop_length=hop_length, 
+                             win_length=win_length, 
+                             window=window,
+                             return_complex=False)  # Set to True if you want complex outpu
+
+            #remove the padding if it was added
+            if x_reconstructed.shape[-1] > shape[-1]:
+                x_reconstructed = x_reconstructed[..., :shape[-1]]
+            x_reconstructed = x_reconstructed.view(shape)
+            return x_reconstructed
+        
+        self.EQ_normalize = EQ_normalize_fn
+
+
+    def __call__(self, x, use_gate=False, RMS=None):
+
+        B, C, T = x.shape
+        if C > 1:
+            x = x.mean(dim=1, keepdim=True)
+
+
+        x= apply_RMS_normalization(x, self.RMS_norm , use_gate=use_gate)  # Apply RMS normalization to the input
+
+        x=self.EQ_normalize(x)
+
+
+
+        x= apply_RMS_normalization(x, self.RMS_norm, use_gate=use_gate)
+        assert not torch.isnan(x).any(), "NaN detected in x after EQ normalization"
+
+
+        return x

utils/fx_normalization/fxnorm_v2_public.py

@@ -0,0 +1,142 @@
+import torch
+import numpy as np
+
+from utils.training_utils import  Gauss_smooth_vectorized,  prepare_smooth_filter
+
+def T602logmag(t60, sample_rate=44100, hop_length=512):
+    return 6.908 / (t60 * (sample_rate / hop_length))  # Convert T60 to delta log magnitude
+
+from utils.data_utils import apply_RMS_normalization
+
+
+class FxNormAug:
+
+    def __init__(self,
+                sample_rate=44100,  # Sample rate of the audio
+                device="cuda" if torch.cuda.is_available() else "cpu",
+                mode="train",  # Mode can be "train" or "eval"
+                seed=42,
+                features_path="features_tency1_4instr_v4.npy",  # Path to the features file
+                ):
+
+        torch.random.manual_seed(seed)
+
+        #the path is in the same directory as this file
+        self.features_path = features_path
+
+        self.sample_rate = sample_rate
+        self.device = device
+
+        self.EQ_normalize_setup()  # Initialize EQ normalization function
+
+
+
+    def EQ_normalize_setup(self ):
+
+        features_mean = np.load(self.features_path, allow_pickle='TRUE')[()]
+
+        target_cuves_original= {
+            "vocals": torch.tensor(features_mean["eq"]["vocals"]).to(torch.float32).to(self.device),
+            "drums": torch.tensor(features_mean["eq"]["drums"]).to(torch.float32).to(self.device),
+            "bass": torch.tensor(features_mean["eq"]["bass"]).to(torch.float32).to(self.device),
+            "other": torch.tensor(features_mean["eq"]["other"]).to(torch.float32).to(self.device),
+        }
+
+
+        nfft=4096 # FFT size hardcoded
+        nfft_orig = 65536  # FFT size for the smoothing filter
+
+        win_length=2048  # Window length hardcoded
+        hop_length=1024  # Hop length hardcoded
+
+        window = torch.sqrt(torch.hann_window(win_length, device=self.device))
+        window_energy = window.pow(2).sum().sqrt()  # Energy of the window
+
+        freqs = torch.fft.rfftfreq(nfft, d=1.0).to(self.device)
+        freqs_Hz=torch.fft.rfftfreq(nfft, d=1.0).to(self.device) * self.sample_rate
+
+        smooth_filter = prepare_smooth_filter(freqs_Hz, Noct=3).to(self.device)  # Prepare the smoothing filter
+
+        freqs_Hz_orig=torch.fft.rfftfreq(nfft_orig, d=1.0).to(self.device) * self.sample_rate
+        smooth_filter_orig = prepare_smooth_filter(freqs_Hz_orig, Noct=3).to(self.device)  # Prepare the smoothing filter
+
+        def downsample_curve(x):
+            return torch.nn.functional.interpolate(
+                x.unsqueeze(0).unsqueeze(0),
+                size=(nfft // 2 + 1,),
+                mode='linear',
+                align_corners=False
+            ).squeeze(0).squeeze(0)
+
+        target_curves = {
+            "vocals": downsample_curve(Gauss_smooth_vectorized(target_cuves_original["vocals"], freqs_Hz_orig, Noct=3, smooth_filter=smooth_filter_orig)),
+            "drums": downsample_curve(Gauss_smooth_vectorized(target_cuves_original["drums"], freqs_Hz_orig, Noct=3, smooth_filter=smooth_filter_orig)),
+            "bass": downsample_curve(Gauss_smooth_vectorized(target_cuves_original["bass"], freqs_Hz_orig, Noct=3, smooth_filter=smooth_filter_orig)),
+            "other": downsample_curve(Gauss_smooth_vectorized(target_cuves_original["other"], freqs_Hz_orig, Noct=3, smooth_filter=smooth_filter_orig)),
+        }
+
+
+        def EQ_normalize_fn(x):
+
+            shape= x.shape
+
+            target_curves_tensor=torch.zeros((shape[0], nfft // 2 + 1), device=self.device, dtype=torch.float32)
+
+            for i in range(shape[0]):
+                track_class = "other"
+
+                assert track_class in target_curves, f"track_class {track_class} not found in target_curves"
+                target_curves_tensor[i] = target_curves[track_class] 
+
+            x=x.view(-1, shape[-1])
+            #ensure x.shape[-1] is divisible by hop_length
+            if x.shape[-1] % hop_length != 0:
+                # Pad the input to make it divisible by hop_length
+                pad_length = hop_length - (x.shape[-1] % hop_length)
+                x = torch.nn.functional.pad(x, (0, pad_length), mode='constant', value=0)
+            X=torch.stft(x, n_fft=nfft, hop_length=hop_length, win_length=win_length, window=window, return_complex=True)/ window_energy
+            X_pow=X.abs().pow(2)
+            X_mean= torch.sqrt(X_pow.mean(dim=-1, keepdim=False))  # Mean power spectrum
+
+            ratio= target_curves_tensor  / (X_mean + 1e-6)
+
+            ratio = torch.clamp(ratio, max=10.0**(40.0/20.0))
+
+
+            ratio_smooth = Gauss_smooth_vectorized(ratio, freqs_Hz, Noct=3, smooth_filter=smooth_filter)
+
+
+            X= X * ratio_smooth.unsqueeze(-1)
+
+            X_unnormalized=X* window_energy
+
+            x_reconstructed = torch.istft(X_unnormalized, 
+                             n_fft=nfft, 
+                             hop_length=hop_length, 
+                             win_length=win_length, 
+                             window=window,
+                             return_complex=False)  # Set to True if you want complex outpu
+
+            #remove the padding if it was added
+            if x_reconstructed.shape[-1] > shape[-1]:
+                x_reconstructed = x_reconstructed[..., :shape[-1]]
+            x_reconstructed = x_reconstructed.view(shape)
+            return x_reconstructed
+        
+        self.EQ_normalize = EQ_normalize_fn
+
+
+    def __call__(self, x, use_gate=False, RMS=-25):
+
+        B, C, T = x.shape
+        if C > 1:
+            x = x.mean(dim=1, keepdim=True)
+
+        x=x/x.max()
+
+        x=self.EQ_normalize(x)
+
+        x= apply_RMS_normalization(x, RMS, use_gate=use_gate)
+        assert not torch.isnan(x).any(), "NaN detected in x after EQ normalization"
+
+        return x

utils/fxencoder_plusplus/__init__.py

@@ -0,0 +1,12 @@
+from .model import FxEncoderPlusPlus, load_model, load_default_model
+
+
+__version__ = "0.1.0"
+
+# This defines what gets imported when someone does "from fxencoder_plusplus import *"
+__all__ = [
+    'FxEncoderPlusPlus',
+    'load_model',
+    'load_default_model'
+]
+

utils/fxencoder_plusplus/model.py

@@ -0,0 +1,676 @@
+import torch
+#import laion_clap
+import logging
+import sys
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchlibrosa.stft import Spectrogram, LogmelFilterBank
+from collections import OrderedDict
+from dataclasses import dataclass
+import numpy as np
+import os
+from pathlib import Path
+from huggingface_hub import hf_hub_download
+
+def init_layer(layer):
+    """Initialize a Linear or Convolutional layer. """
+    nn.init.xavier_uniform_(layer.weight)
+
+    if hasattr(layer, 'bias'):
+        if layer.bias is not None:
+            layer.bias.data.fill_(0.)
+
+def init_bn(bn):
+    """Initialize a Batchnorm layer. """
+    bn.bias.data.fill_(0.)
+    bn.weight.data.fill_(1.)
+
+class ConvBlock(nn.Module):
+    def __init__(self, in_channels, out_channels):
+
+        super(ConvBlock, self).__init__()
+
+        self.conv1 = nn.Conv2d(in_channels=in_channels,
+                              out_channels=out_channels,
+                              kernel_size=(3, 3), stride=(1, 1),
+                              padding=(1, 1), bias=False)
+
+        self.conv2 = nn.Conv2d(in_channels=out_channels,
+                              out_channels=out_channels,
+                              kernel_size=(3, 3), stride=(1, 1),
+                              padding=(1, 1), bias=False)
+
+        self.bn1 = nn.BatchNorm2d(out_channels)
+        self.bn2 = nn.BatchNorm2d(out_channels)
+
+        self.init_weight()
+
+    def init_weight(self):
+        init_layer(self.conv1)
+        init_layer(self.conv2)
+        init_bn(self.bn1)
+        init_bn(self.bn2)
+
+
+    def forward(self, input, pool_size=(2, 2), pool_type='avg'):
+
+        x = input
+        x = F.relu_(self.bn1(self.conv1(x)))
+        x = F.relu_(self.bn2(self.conv2(x)))
+        if pool_type == 'max':
+            x = F.max_pool2d(x, kernel_size=pool_size)
+        elif pool_type == 'avg':
+            x = F.avg_pool2d(x, kernel_size=pool_size)
+        elif pool_type == 'avg+max':
+            x1 = F.avg_pool2d(x, kernel_size=pool_size)
+            x2 = F.max_pool2d(x, kernel_size=pool_size)
+            x = x1 + x2
+        else:
+            raise Exception('Incorrect argument!')
+
+        return x
+
+class CLAP_AUDIO_ENCODER(torch.nn.Module):
+    def __init__(self, pretrained: bool = True, frozen: bool = False) -> None:
+        super().__init__()
+        self.pretrained = pretrained
+        self.frozen = frozen
+
+        # load the model
+        self.encoder = laion_clap.CLAP_Module(enable_fusion=False, amodel= 'HTSAT-tiny', tmodel='roberta')
+        if self.pretrained:
+            self.encoder.load_ckpt()  # download the default pretrained checkpoint.
+
+        self.embed_dim = 512
+
+    def forward(self, x: torch.Tensor):
+        if self.frozen:
+            with torch.no_grad():
+                embed = self.encoder.get_audio_embedding_from_data(
+                    x=x, use_tensor=True
+                )
+        else:
+            embed = self.encoder.get_audio_embedding_from_data(
+                    x=x, use_tensor=True
+                )
+
+        return embed
+            
+class CLAP_TEXT_ENCODER(torch.nn.Module):
+    def __init__(self, pretrained: bool = True, frozen: bool = False) -> None:
+        super().__init__()
+        self.pretrained = pretrained
+        self.frozen = frozen
+
+        # load the model
+        self.encoder = laion_clap.CLAP_Module(enable_fusion=False, amodel= 'HTSAT-tiny', tmodel='roberta')
+        if self.pretrained:
+            self.encoder.load_ckpt()  # download the default pretrained checkpoint.
+
+        self.embed_dim = 512
+
+    def forward(self, x):
+        if self.frozen:
+            with torch.no_grad():
+                embed = self.encoder.get_text_embedding(
+                    x=x, use_tensor=True
+                )
+        else:
+            embed = self.encoder.get_text_embedding(
+                    x=x, use_tensor=True
+                )
+
+        return embed 
+    
+# > ================ Proposed =================== <
+class MixtureFxEncoder(nn.Module):
+    def __init__(self, sample_rate, window_size, hop_size, mel_bins, fmin, fmax, enable_fusion=False, fusion_type='None'):
+        super().__init__()
+        self.enable_fusion = enable_fusion
+        self.fusion_type = fusion_type
+
+        window = "hann"
+        center = True
+        pad_mode = "reflect"
+        ref = 1.0
+        amin = 1e-10
+        top_db = None
+        self.input_norm = "minmax"
+        # Spectrogram extractor
+        self.spectrogram_extractor = Spectrogram(
+            n_fft=window_size,
+            hop_length=hop_size,
+            win_length=window_size,
+            window=window,
+            center=center,
+            pad_mode=pad_mode,
+            freeze_parameters=True,
+        )
+
+        # Logmel feature extractor
+        self.logmel_extractor = LogmelFilterBank(
+            sr=sample_rate,
+            n_fft=window_size,
+            n_mels=mel_bins,
+            fmin=fmin,
+            fmax=fmax,
+            ref=ref,
+            amin=amin,
+            top_db=top_db,
+            freeze_parameters=True,
+        )
+
+        self.bn0 = nn.BatchNorm2d(64)
+
+        self.conv_block1 = ConvBlock(in_channels=2, out_channels=64)
+        self.conv_block2 = ConvBlock(in_channels=64, out_channels=128)
+        self.conv_block3 = ConvBlock(in_channels=128, out_channels=256)
+        self.conv_block4 = ConvBlock(in_channels=256, out_channels=512)
+        self.conv_block5 = ConvBlock(in_channels=512, out_channels=1024)
+        self.conv_block6 = ConvBlock(in_channels=1024, out_channels=2048)
+
+        self.fc_1 = nn.Linear(2048, 2048, bias=True)
+
+        self.init_weight()
+        
+    def init_weight(self):
+        init_bn(self.bn0)
+        init_layer(self.fc_1)
+
+    def forward(self, x):
+        """
+        Input: (batch_size, 2, data_length)
+        """
+        batch_size, chs, seq_len = x.size()
+
+        # move to batch dim
+        x = x.view(batch_size * chs, seq_len)
+        
+        # extract logmel features
+        x = self.spectrogram_extractor(x)  # (batch_size, 1, time_steps, freq_bins)
+        x = self.logmel_extractor(x)  # (batch_size, 1, time_steps, mel_bins)
+
+        if self.input_norm == "batchnorm":
+            # this normalizes over mel bins which is problematic for equalization
+            x = x.transpose(1, 3)
+            x = self.bn0(x)
+            x = x.transpose(1, 3)
+        elif self.input_norm == "minmax":
+            x = x.clamp(-80, 40.0)  # clamp the logmels between -80 and 40
+            x = (x + 80) / 120  # normalize the logmels between 0 and 1
+            x = (x * 2) - 1  # normalize the logmels between -1 and 1
+        elif self.input_norm == "none":
+            pass
+        else:
+            raise ValueError(f"Invalid input_norm: {self.input_norm}")
+
+        x = x.view(batch_size, chs, x.size(-2), x.size(-1))
+        
+        x = self.conv_block1(x, pool_size=(2, 2), pool_type='avg')
+        x = F.dropout(x, p=0.2, training=self.training)
+        x = self.conv_block2(x, pool_size=(2, 2), pool_type='avg')
+        x = F.dropout(x, p=0.2, training=self.training)
+        x = self.conv_block3(x, pool_size=(2, 2), pool_type='avg')
+        x = F.dropout(x, p=0.2, training=self.training)
+        x = self.conv_block4(x, pool_size=(2, 2), pool_type='avg')
+        x = F.dropout(x, p=0.2, training=self.training)
+        x = self.conv_block5(x, pool_size=(2, 2), pool_type='avg')
+        x = F.dropout(x, p=0.2, training=self.training)
+        x = self.conv_block6(x, pool_size=(1, 1), pool_type='avg')
+        x = F.dropout(x, p=0.2, training=self.training)
+        x = torch.mean(x, dim=3)
+        
+
+        (x1, _) = torch.max(x, dim=2)
+        x2 = torch.mean(x, dim=2)
+        x = x1 + x2
+        x = F.relu_(self.fc_1(x))
+        embedding = x
+
+        output_dict = {
+            'embedding': embedding, 
+        }
+
+        return output_dict
+
+def create_MixtureFxEncoder():
+    model = MixtureFxEncoder(
+        sample_rate = 44100, #audio_cfg.sample_rate,
+        window_size = 2048, #audio_cfg.window_size,
+        hop_size = 512, #audio_cfg.hop_size,
+        mel_bins = 64, #audio_cfg.mel_bins,
+        fmin = 50, #audio_cfg.fmin,
+        fmax = 18000, #audio_cfg.fmax,
+    )
+    return model
+
+class MLPLayers(nn.Module):
+    def __init__(self, units=[512, 512, 512], nonlin=nn.ReLU(), dropout=0.1):
+        super(MLPLayers, self).__init__()
+        self.nonlin = nonlin
+        self.dropout = dropout
+
+        sequence = []
+        for u0, u1 in zip(units[:-1], units[1:]):
+            sequence.append(nn.Linear(u0, u1))
+            sequence.append(self.nonlin)
+            sequence.append(nn.Dropout(self.dropout))
+        sequence = sequence[:-2]
+
+        self.sequential = nn.Sequential(*sequence)
+
+    def forward(self, X):
+        X = self.sequential(X)
+        return X
+
+class BernoulliDynamicDropout(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.p_min = 0.75
+        self.p_max = 0.95
+    
+    def get_random_dropout_rate(self):
+        return torch.empty(1).uniform_(self.p_min, self.p_max).item()
+    
+    def forward(self, x):
+        if self.training:
+            p = self.get_random_dropout_rate()
+            mask = torch.bernoulli(torch.full_like(x, 1-p))
+            return x * mask / (1 - p)
+        return x
+
+class AudioExtracter(nn.Module):
+    def __init__(self, fx_embedding_dim=128, clap_embedding_dim=512):
+        super().__init__()
+        
+        # Simple fusion network
+        self.fusion = nn.Sequential(
+            nn.Linear(fx_embedding_dim+clap_embedding_dim, 128),
+            nn.LeakyReLU(0.1),
+            nn.Linear(128, 128),
+            nn.LeakyReLU(0.1),
+            nn.Linear(128, 128),
+        )
+    
+    def forward(self, mixture_emb, query_emb):
+        # Concatenate and project
+        x = torch.cat([mixture_emb, query_emb], dim=-1)  # [B, 2D]
+        stem_emb = self.fusion(x)  # [B, D]
+        
+        return stem_emb
+
+@dataclass
+class AudioCfg:
+    model_type: str = "PANN"
+    model_name: str = "Cnn14"
+    sample_rate: int = 44100
+    # Param
+    audio_length: int = 1024
+    window_size: int = 1024
+    hop_size: int = 1024
+    fmin: int = 50
+    fmax: int = 14000
+    mel_bins: int = 64
+    clip_samples: int = 441000
+    class_num: int = 527
+    condition_dim: int = 512
+
+class FxEncoderPlusPlus(nn.Module): 
+    def __init__(
+        self,
+        embed_dim: int = 2048,
+        mixture_cfg: AudioCfg = None,
+        enable_fusion: bool = False,
+        fusion_type: str = 'None',
+        joint_embed_shape: int = 128,
+        mlp_act: str = 'relu',
+        audio_clap_module: bool = True,
+        text_clap_module: bool = False,
+        extractor_module: bool = True,
+        device: str = "cpu",
+    ):
+        super().__init__()
+        
+        self.mixture_cfg = mixture_cfg
+        self.enable_fusion = enable_fusion
+        self.fusion_type = fusion_type
+        self.joint_embed_shape = joint_embed_shape
+        self.mlp_act = mlp_act
+        self.device = device
+
+        if mlp_act == 'relu':
+            mlp_act_layer = nn.ReLU()
+        elif mlp_act == 'gelu':
+            mlp_act_layer = nn.GELU()
+        else:
+            raise NotImplementedError
+
+        # > ========================= FX Encoder ========================= <
+        self.fx_encoder = create_MixtureFxEncoder()
+        self.fx_encoder_transform = MLPLayers(units=[self.joint_embed_shape, self.joint_embed_shape, self.joint_embed_shape], dropout=0.1)
+        
+        self.fx_encoder_projection = nn.Sequential(
+            nn.Linear(embed_dim, self.joint_embed_shape),
+            mlp_act_layer,
+            nn.Linear(self.joint_embed_shape, self.joint_embed_shape)
+        )
+        
+        self.logit_scale_m = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+        self.logit_scale_t = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+        
+        if audio_clap_module:
+            # Freeze all layers
+            # print("Loading CLAP Audio Model")
+            self.audio_clap_model = CLAP_AUDIO_ENCODER(pretrained=True, frozen=True)
+            self.audio_clap_model.to(device)
+            for param in self.audio_clap_model.parameters():
+                param.requires_grad = False
+            self.clap_dropout = BernoulliDynamicDropout()
+        if text_clap_module:
+            # Freeze all layers
+            # print("Loading CLAP Text Model")
+            self.text_clap_model = CLAP_TEXT_ENCODER(pretrained=True, frozen=True)
+            self.text_clap_model.to(device)
+            for param in self.text_clap_model.parameters():
+                param.requires_grad = False
+            
+        if extractor_module:
+            # extractor 
+            self.extractor = AudioExtracter()
+
+        self.use_audio_clap_module = audio_clap_module
+        self.use_text_clap_module = text_clap_module 
+        self.use_extractor_module = extractor_module
+
+    def get_fx_embedding(self, x):
+        fx_emb = self.fx_encoder(x) 
+        fx_emb = self.fx_encoder_projection(fx_emb["embedding"])
+        fx_emb = F.normalize(fx_emb, dim=-1)
+        return fx_emb
+    
+    def get_fx_embedding_by_audio_query(self, x, audio_query):
+        # mixture fx embedding 
+        fx_mixture_emb = self.fx_encoder(x) 
+        fx_mixture_emb = self.fx_encoder_projection(fx_mixture_emb["embedding"])
+        fx_mixture_emb = F.normalize(fx_mixture_emb, dim=-1)
+        
+        # stem fx embedding 
+        query_content_embeded = self.audio_clap_model(torch.mean(audio_query, dim=1))
+        fx_stem_emb = self.extractor(fx_mixture_emb, query_content_embeded)
+        fx_stem_emb = F.normalize(fx_stem_emb, dim=-1)
+        return fx_mixture_emb, fx_stem_emb
+    
+    def get_fx_embedding_by_text_query(self, x, text_query):
+        # mixture fx embedding 
+        fx_mixture_emb = self.fx_encoder(x) 
+        fx_mixture_emb = self.fx_encoder_projection(fx_mixture_emb["embedding"])
+        fx_mixture_emb = F.normalize(fx_mixture_emb, dim=-1)
+
+        # stem fx embedding 
+        query_embeded = self.text_clap_model(text_query)
+        fx_stem_emb = self.extractor(fx_mixture_emb, query_embeded)
+        fx_stem_emb = F.normalize(fx_stem_emb, dim=-1)
+        return fx_mixture_emb, fx_stem_emb
+    
+    def forward(
+        self, 
+        mixture_a, 
+        mixture_b, 
+        stem_a, 
+        query_stem, 
+        device = None
+    ):
+
+        if device is None:
+            if mixture_a is not None:
+                device = mixture_a.device
+            elif mixture_b is not None:
+                device = mixture_b.device
+        if mixture_a is None and mixture_b is None:
+            # a hack to get the logit scale
+            return self.logit_scale_m.exp(), self.logit_scale_t.exp()
+        
+        # ======== Global ========
+        mixture_a_features = self.fx_encoder_projection(
+            self.fx_encoder(mixture_a)["embedding"]
+        )
+        mixture_a_features = F.normalize(mixture_a_features, dim=-1)
+        
+        mixture_b_features = self.fx_encoder_projection(
+            self.fx_encoder(mixture_b)["embedding"]
+        )
+        mixture_b_features = F.normalize(mixture_b_features, dim=-1)
+
+        mixture_a_features_mlp = self.fx_encoder_transform(mixture_a_features)
+        mixture_b_features_mlp = self.fx_encoder_transform(mixture_b_features)
+        
+        # ======= Local ========
+        stem_a_features = self.fx_encoder_projection(
+            self.fx_encoder(stem_a)["embedding"]
+        )
+        stem_a_features = F.normalize(stem_a_features, dim=-1)
+        
+        if self.use_audio_clap_module and self.use_extractor_module:
+            query_stem_content_embeded = self.clap_dropout(
+                self.audio_clap_model(
+                    torch.mean(query_stem, dim=1)
+                )
+            )
+            extracted_stem_a_features = self.extractor(mixture_a_features, query_stem_content_embeded)
+            extracted_stem_a_features = F.normalize(extracted_stem_a_features, dim=-1)
+        elif self.use_text_clap_module and self.use_extractor_module:
+            query_stem_content_embeded = self.text_clap_model(query_stem)
+            extracted_stem_a_features = self.extractor(mixture_a_features, query_stem_content_embeded)
+            extracted_stem_a_features = F.normalize(extracted_stem_a_features, dim=-1)
+        
+        return (
+            mixture_a_features, # global
+            mixture_b_features, # global 
+            stem_a_features, # local 
+            extracted_stem_a_features, # local 
+            mixture_a_features_mlp,
+            mixture_b_features_mlp,
+            self.logit_scale_m.exp(),
+            self.logit_scale_t.exp(),
+        )
+
+    def get_logit_scale(self):
+        return self.logit_scale_m.exp(), self.logit_scale_t.exp()
+
+# def load_model(model_path, device):
+    
+#     model = FxEncoderPlusPlus(
+#         embed_dim = 2048, 
+#         audio_clap_module = True, 
+#         extractor_module = True
+#     )
+
+#     # load model
+#     checkpoint = torch.load(model_path, map_location=device, weights_only=False)
+#     if "epoch" in checkpoint:
+#         # resuming a train checkpoint w/ epoch and optimizer state
+#         start_epoch = checkpoint["epoch"]
+#         sd = checkpoint["state_dict"]
+#         if next(iter(sd.items()))[0].startswith(
+#                 "module"
+#         ):
+#             sd = {k[len("module."):]: v for k, v in sd.items()}
+#         model.load_state_dict(sd)
+#         logging.info(
+#             f"=> resuming checkpoint '{model_path}' (epoch {start_epoch})"
+#         )
+#     else:
+#         # loading a bare (model only) checkpoint for fine-tune or evaluation
+#         model.load_state_dict(checkpoint)
+#         start_epoch = 0
+    
+#     model.to(device)
+#     model.eval()
+#     for param in model.parameters():
+#         param.requires_grad = False
+#     return model
+
+# Define available models
+MODEL_REGISTRY = {
+    "default": {
+        "repo_id": "yytung/fxencoder-plusplus",
+        "filename": "fxenc_plusplus_default.pt",
+        "description": "Default model",
+    },
+    # "musdb": {
+    #     "repo_id": "yytung/fxencoder-plusplus", 
+    #     "filename": "fxenc_plusplus_musdb.pt",
+    #     "description": "Fx-Encoder++ trained on musdb",
+    # },
+    # "medleydb": {
+    #     "repo_id": "yytung/fxencoder-plusplus",
+    #     "filename": "fxenc_plusplus_medleydb.pt", 
+    #     "description": "Fx-Encoder++ trained on medleydb",
+    # },
+}
+
+def get_model_path(model_name="default", cache_dir=None, force_download=False):
+    """
+    Download or retrieve the path to a pretrained model.
+    
+    Args:
+        model_name: Name of the model variant ('default', 'musdb', 'medleydb')
+        cache_dir: Custom cache directory. If None, uses ~/.cache/fxencoder_plusplus
+        force_download: Force re-download even if file exists
+        
+    Returns:
+        Path to the model file
+    """
+    if model_name not in MODEL_REGISTRY:
+        available = ", ".join(MODEL_REGISTRY.keys())
+        raise ValueError(f"Unknown model: {model_name}. Available models: {available}")
+    
+    if cache_dir is None:
+        cache_dir = Path.home() / ".cache" / "fxencoder_plusplus"
+    else:
+        cache_dir = Path(cache_dir)
+    
+    cache_dir.mkdir(parents=True, exist_ok=True)
+    
+    model_info = MODEL_REGISTRY[model_name]
+    model_path = cache_dir / model_info["filename"]
+    
+    # Check if already downloaded
+    if model_path.exists() and not force_download:
+        print(f"Using cached model: {model_path}")
+        return str(model_path)
+    
+    print(f"Description: {model_info['description']}")
+    
+    # Download from Hugging Face
+    downloaded_path = hf_hub_download(
+        repo_id=model_info["repo_id"],
+        filename=model_info["filename"],
+        cache_dir=str(cache_dir),
+        force_download=force_download
+    )
+    
+    print(f"Model downloaded successfully to: {downloaded_path}")
+    return downloaded_path
+
+def list_available_models():
+    """List all available pretrained models."""
+    print("Available FxEncoder++ models:")
+    print("-" * 50)
+    for name, info in MODEL_REGISTRY.items():
+        print(f"  {name}:")
+        print(f"    - Description: {info['description']}")
+    print("-" * 50)
+
+def load_model(model_name="default", model_path=None, device="cuda", auto_download=True, cache_dir=None):
+    """
+    Load FxEncoderPlusPlus model.
+    
+    Args:
+        model_name: Name of pretrained model ('default', 'musdb', 'medleydb')
+        model_path: Custom checkpoint path. If provided, ignores model_name
+        device: Device to load model on ('cuda' or 'cpu')
+        auto_download: Automatically download if model not found
+        cache_dir: Custom cache directory for downloaded models
+        
+    Returns:
+        Loaded FxEncoderPlusPlus model
+        
+    Examples:
+        # Load default base model
+        model = load_model()
+        
+        # Load musdb model
+        model = load_model(model_name="musdb")
+        
+        # Load medleydb model
+        model = load_model(model_name="medleydb")
+        
+        # Load custom checkpoint
+        model = load_model(model_path="/path/to/custom.pt")
+        
+        # List available models
+        list_available_models()
+    """
+    # Handle device
+    if device == "cuda" and not torch.cuda.is_available():
+        print("CUDA not available, using CPU")
+        device = "cpu"
+    
+    
+    # Determine model path
+    if model_path is None:
+        if auto_download:
+            model_path = get_model_path(model_name, cache_dir=cache_dir)
+        else:
+            raise ValueError("model_path is None and auto_download is False")
+    
+    # Create model instance with specified device
+    model = FxEncoderPlusPlus(
+        embed_dim=2048, 
+        audio_clap_module=True, 
+        text_clap_module=True,
+        extractor_module=True,
+        device=device
+    )
+    
+    # Load checkpoint
+    checkpoint = torch.load(model_path, map_location=device, weights_only=False)
+    
+    if "epoch" in checkpoint:
+        # resuming a train checkpoint w/ epoch and optimizer state
+        start_epoch = checkpoint["epoch"]
+        sd = checkpoint["state_dict"]
+        if next(iter(sd.items()))[0].startswith("module"):
+            sd = {k[len("module."):]: v for k, v in sd.items()}
+        model.load_state_dict(sd)
+        print(f"Loaded checkpoint from epoch {start_epoch}")
+    else:
+        # loading a bare (model only) checkpoint for fine-tune or evaluation
+        model.load_state_dict(checkpoint)
+        print("Loaded model checkpoint")
+    
+    model.to(device)
+    model.eval()
+    
+    # Freeze parameters for inference
+    for param in model.parameters():
+        param.requires_grad = False
+    
+    print(f"Model loaded successfully on {device}")
+    return model
+
+# Convenience functions for specific models
+def load_default_model(device="cuda", **kwargs):
+    """Load the default FxEncoder++ model."""
+    return load_model(model_name="default", device=device, **kwargs)
+
+# def load_musdb_model(device="cuda", **kwargs):
+#     """Load the musdb FxEncoder++ model."""
+#     return load_model(model_name="musdb", device=device, **kwargs)
+
+# def load_medleydb_model(device="cuda", **kwargs):
+#     """Load the medleydb FxEncoder++ model."""
+#     return load_model(model_name="medleydb", device=device, **kwargs)
+

utils/laion_clap/clap_module/__init__.py

@@ -0,0 +1,8 @@
+from .factory import list_models, create_model, create_model_and_transforms, add_model_config
+from .loss import ClipLoss, gather_features, LPLoss, lp_gather_features, LPMetrics
+from .model import CLAP, CLAPTextCfg, CLAPVisionCfg, CLAPAudioCfp, convert_weights_to_fp16, trace_model
+from .openai import load_openai_model, list_openai_models
+from .pretrained import list_pretrained, list_pretrained_tag_models, list_pretrained_model_tags,\
+    get_pretrained_url, download_pretrained
+from .tokenizer import SimpleTokenizer, tokenize
+from .transform import image_transform

utils/laion_clap/clap_module/bert.py

@@ -0,0 +1,32 @@
+from transformers import BertTokenizer, BertModel
+tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
+model = BertModel.from_pretrained("bert-base-uncased")
+text = "Replace me by any text you'd like."
+
+def bert_embeddings(text):
+    # text = "Replace me by any text you'd like."
+    encoded_input = tokenizer(text, return_tensors='pt')
+    output = model(**encoded_input)
+    return output
+    
+from transformers import RobertaTokenizer, RobertaModel
+
+tokenizer = RobertaTokenizer.from_pretrained('roberta-base')
+model = RobertaModel.from_pretrained('roberta-base')
+text = "Replace me by any text you'd like."
+def Roberta_embeddings(text):
+    # text = "Replace me by any text you'd like."
+    encoded_input = tokenizer(text, return_tensors='pt')
+    output = model(**encoded_input)
+    return output
+
+from transformers import BartTokenizer, BartModel
+
+tokenizer = BartTokenizer.from_pretrained('facebook/bart-base')
+model = BartModel.from_pretrained('facebook/bart-base')
+text = "Replace me by any text you'd like."
+def bart_embeddings(text):
+    # text = "Replace me by any text you'd like."
+    encoded_input = tokenizer(text, return_tensors='pt')
+    output = model(**encoded_input)
+    return output

utils/laion_clap/clap_module/bpe_simple_vocab_16e6.txt.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:924691ac288e54409236115652ad4aa250f48203de50a9e4722a6ecd48d6804a
+size 1356917