cleanup (~25k lines) + torchlibrosa in requirements.txt

inference/validator_FxProcessor.py

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

requirements.txt

@@ -11,6 +11,7 @@ pyloudnorm
 einops
 packaging
 librosa
+torchlibrosa
 transformers
 wget
 huggingface-hub

utils/MSS_loss.py

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

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

@@ -1,1729 +0,0 @@
-"""
-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/dnnlib/__init__.py

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

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

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

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

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

@@ -1,958 +0,0 @@
-
-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/audio_features.py

@@ -1,76 +0,0 @@
-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/load_features.py

@@ -1,37 +1,8 @@
-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 
 
@@ -124,137 +95,3 @@ def load_CLAP(model_args, device, *args, **kwargs):
 
     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/pytorch_utils.py

@@ -1,256 +0,0 @@
-"""
-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/fxnorm.py

@@ -1,154 +0,0 @@
-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/laion_clap/clap_module/__init__.py

@@ -1,8 +1,6 @@
 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

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

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

utils/laion_clap/clap_module/linear_probe.py

@@ -1,63 +0,0 @@
-import numpy as np
-import torch.nn.functional as F
-from torch import nn
-from .model import MLPLayers
-
-
-class LinearProbe(nn.Module):
-    def __init__(self, model, mlp, freeze, in_ch, out_ch, act=None):
-        """
-        Args:
-            model: nn.Module
-            mlp: bool, if True, then use the MLP layer as the linear probe module
-            freeze: bool, if Ture, then freeze all the CLAP model's layers when training the linear probe
-            in_ch: int, the output channel from CLAP model
-            out_ch: int, the output channel from linear probe (class_num)
-            act: torch.nn.functional, the activation function before the loss function
-        """
-        super().__init__()
-        in_ch = 512
-        self.clap_model = model
-        self.clap_model.text_branch = None  # to save memory
-        self.freeze = freeze
-        if mlp:
-            self.lp_layer = MLPLayers(units=[in_ch, in_ch * 2, out_ch])
-        else:
-            self.lp_layer = nn.Linear(in_ch, out_ch)
-
-        if self.freeze:
-            for param in self.clap_model.parameters():
-                param.requires_grad = False
-
-        if act == 'None':
-            self.act = None
-        elif act == 'relu':
-            self.act = nn.ReLU()
-        elif act == 'elu':
-            self.act = nn.ELU()
-        elif act == 'prelu':
-            self.act = nn.PReLU(num_parameters=in_ch)
-        elif act == 'softmax':
-            self.act = nn.Softmax(dim=-1)
-        elif act == 'sigmoid':
-            self.act = nn.Sigmoid()
-
-    def forward(self, x, mix_lambda=None, device=None):
-        """
-        Args:
-            x: waveform, torch.tensor [batch, t_samples] / batch of mel_spec and longer list
-            mix_lambda: torch.tensor [batch], the mixup lambda
-        Returns:
-            class_prob: torch.tensor [batch, class_num]
-
-        """
-        # batchnorm cancel grandient
-        if self.freeze:
-            self.clap_model.eval()
-
-        x = self.clap_model.audio_projection(
-            self.clap_model.audio_branch(x, mixup_lambda=mix_lambda, device=device)["embedding"])
-        out = self.lp_layer(x)
-        if self.act is not None:
-            out = self.act(out)
-        return out

utils/laion_clap/clap_module/loss.py

@@ -1,307 +0,0 @@
-from multiprocessing.sharedctypes import Value
-import torch
-import torch.distributed.nn
-from torch import distributed as dist, nn as nn
-from torch.nn import functional as F
-import numpy as np
-from sklearn.metrics import average_precision_score, roc_auc_score, accuracy_score 
-
-try:
-    import horovod.torch as hvd
-except ImportError:
-    hvd = None
-
-
-def gather_features(
-        audio_features,
-        text_features,
-        audio_features_mlp=None, 
-        text_features_mlp=None,
-        local_loss=False,
-        gather_with_grad=False,
-        rank=0,
-        world_size=1,
-        use_horovod=False,
-        mlp_loss=False
-):
-    if use_horovod:
-        assert hvd is not None, 'Please install horovod'
-        if gather_with_grad:
-            all_audio_features = hvd.allgather(audio_features)
-            all_text_features = hvd.allgather(text_features)
-            if mlp_loss:
-                all_audio_features_mlp = hvd.allgather(audio_features_mlp)
-                all_text_features_mlp = hvd.allgather(text_features_mlp)
-        else:
-            with torch.no_grad():
-                all_audio_features = hvd.allgather(audio_features)
-                all_text_features = hvd.allgather(text_features)
-                if mlp_loss:
-                    all_audio_features_mlp = hvd.allgather(audio_features_mlp)
-                    all_text_features_mlp = hvd.allgather(text_features_mlp)
-            if not local_loss:
-                # ensure grads for local rank when all_* features don't have a gradient
-                gathered_audio_features = list(all_audio_features.chunk(world_size, dim=0))
-                gathered_text_features = list(all_text_features.chunk(world_size, dim=0))
-                gathered_audio_features[rank] = audio_features
-                gathered_text_features[rank] = text_features
-                all_audio_features = torch.cat(gathered_audio_features, dim=0)
-                all_text_features = torch.cat(gathered_text_features, dim=0)
-                if mlp_loss:
-                    gathered_audio_features_mlp = list(all_audio_features_mlp.chunk(world_size, dim=0))
-                    gathered_text_features_mlp = list(all_text_features_mlp.chunk(world_size, dim=0))
-                    gathered_audio_features_mlp[rank] = audio_features_mlp
-                    gathered_text_features_mlp[rank] = text_features_mlp
-                    all_audio_features_mlp = torch.cat(gathered_audio_features_mlp, dim=0)
-                    all_text_features_mlp = torch.cat(gathered_text_features_mlp, dim=0)
-    else:
-        # We gather tensors from all gpus
-        if gather_with_grad:
-            all_audio_features = torch.cat(torch.distributed.nn.all_gather(audio_features), dim=0)
-            all_text_features = torch.cat(torch.distributed.nn.all_gather(text_features), dim=0)
-            if mlp_loss:
-                all_audio_features_mlp = torch.cat(torch.distributed.nn.all_gather(audio_features_mlp), dim=0)
-                all_text_features_mlp = torch.cat(torch.distributed.nn.all_gather(text_features_mlp), dim=0)
-        else:
-            gathered_audio_features = [torch.zeros_like(audio_features) for _ in range(world_size)]
-            gathered_text_features = [torch.zeros_like(text_features) for _ in range(world_size)]
-            dist.all_gather(gathered_audio_features, audio_features)
-            dist.all_gather(gathered_text_features, text_features)
-            if mlp_loss:
-                gathered_audio_features_mlp = [torch.zeros_like(audio_features_mlp) for _ in range(world_size)]
-                gathered_text_features_mlp = [torch.zeros_like(text_features_mlp) for _ in range(world_size)]
-                dist.all_gather(gathered_audio_features_mlp, audio_features_mlp)
-                dist.all_gather(gathered_text_features_mlp, text_features_mlp)
-            if not local_loss:
-                # ensure grads for local rank when all_* features don't have a gradient
-                gathered_audio_features[rank] = audio_features
-                gathered_text_features[rank] = text_features
-                if mlp_loss:
-                    gathered_audio_features_mlp[rank] = audio_features_mlp
-                    gathered_text_features_mlp[rank] = text_features_mlp
-
-            all_audio_features = torch.cat(gathered_audio_features, dim=0)
-            all_text_features = torch.cat(gathered_text_features, dim=0)
-            if mlp_loss:
-                all_audio_features_mlp = torch.cat(gathered_audio_features_mlp, dim=0)
-                all_text_features_mlp = torch.cat(gathered_text_features_mlp, dim=0)
-    if mlp_loss:
-        return all_audio_features, all_text_features, all_audio_features_mlp, all_text_features_mlp
-    else:
-        return all_audio_features, all_text_features
-
-class ClipLoss(nn.Module):
-
-    def __init__(
-            self,
-            local_loss=False,
-            gather_with_grad=False,
-            cache_labels=False,
-            rank=0,
-            world_size=1,
-            use_horovod=False,
-            mlp_loss=False,
-            weight_loss_kappa=0,
-    ):
-        super().__init__()
-        self.local_loss = local_loss
-        self.gather_with_grad = gather_with_grad
-        self.cache_labels = cache_labels
-        self.rank = rank
-        self.world_size = world_size
-        self.use_horovod = use_horovod
-        self.mlp_loss = mlp_loss
-        self.weighted_loss = bool(weight_loss_kappa!=0)
-        self.weight_loss_kappa = weight_loss_kappa
-        # cache state
-        self.prev_num_logits = 0
-        self.labels = {}
-
-    def forward(self, audio_features, text_features, logit_scale_a, logit_scale_t=None, audio_features_mlp=None, text_features_mlp=None):
-        device = audio_features.device
-        if self.mlp_loss:
-            if self.world_size > 1:
-                all_audio_features, all_text_features, all_audio_features_mlp, all_text_features_mlp = gather_features(
-                    audio_features=audio_features,text_features=text_features,
-                    audio_features_mlp=audio_features_mlp,text_features_mlp=text_features_mlp,
-                    local_loss=self.local_loss,gather_with_grad=self.gather_with_grad,
-                    rank=self.rank,world_size=self.world_size,use_horovod=self.use_horovod,
-                    mlp_loss=self.mlp_loss
-                )
-                if self.local_loss:
-                    a_logits_per_audio = logit_scale_a * audio_features @ all_text_features_mlp.T
-                    a_logits_per_text = logit_scale_a * text_features_mlp @ all_audio_features.T
-                    t_logits_per_audio = logit_scale_t * audio_features_mlp @ all_text_features.T
-                    t_logits_per_text = logit_scale_t * text_features @ all_audio_features_mlp.T
-                else:
-                    a_logits_per_audio = logit_scale_a * all_audio_features @ all_text_features_mlp.T
-                    a_logits_per_text = a_logits_per_audio.T
-                    t_logits_per_audio = logit_scale_t * all_audio_features_mlp @ all_text_features.T
-                    t_logits_per_text = t_logits_per_audio.T
-            else:
-                a_logits_per_audio = logit_scale_a * audio_features @ text_features_mlp.T
-                a_logits_per_text = logit_scale_a * text_features_mlp @ audio_features.T
-                t_logits_per_audio = logit_scale_t * audio_features_mlp @ text_features.T
-                t_logits_per_text = logit_scale_t * text_features @ audio_features_mlp.T
-
-            # calculated ground-truth and cache if enabled
-            num_logits = a_logits_per_audio.shape[0]
-            if self.prev_num_logits != num_logits or device not in self.labels:
-                labels = torch.arange(num_logits, device=device, dtype=torch.long)
-                if self.world_size > 1 and self.local_loss:
-                    labels = labels + num_logits * self.rank
-                if self.cache_labels:
-                    self.labels[device] = labels
-                    self.prev_num_logits = num_logits
-            else:
-                labels = self.labels[device]
-
-            if not self.weighted_loss:
-                total_loss = (
-                    F.cross_entropy(a_logits_per_audio, labels) +
-                    F.cross_entropy(a_logits_per_text, labels) + 
-                    F.cross_entropy(t_logits_per_audio, labels) +
-                    F.cross_entropy(t_logits_per_text, labels) 
-                    ) / 4
-            else:
-                audio_weight = (audio_features@audio_features.T).detach()
-                audio_weight = (torch.exp(torch.sum(audio_weight, axis=1)/(self.weight_loss_kappa*len(audio_weight)))).detach()
-                text_weight = (text_features@text_features.T).detach()
-                text_weight = (torch.exp(torch.sum(text_weight, axis=1)/(self.weight_loss_kappa*len(text_features)))).detach()
-                total_loss = (
-                    F.cross_entropy(a_logits_per_audio, labels, weight=audio_weight) +
-                    F.cross_entropy(a_logits_per_text, labels, weight=audio_weight) + 
-                    F.cross_entropy(t_logits_per_audio, labels, weight=text_weight) +
-                    F.cross_entropy(t_logits_per_text, labels, weight=text_weight) 
-                    ) / 4
-        else:
-            if self.world_size > 1:
-                all_audio_features, all_text_features = gather_features(
-                    audio_features=audio_features,text_features=text_features,
-                    local_loss=self.local_loss,gather_with_grad=self.gather_with_grad,
-                    rank=self.rank,world_size=self.world_size,use_horovod=self.use_horovod,
-                    mlp_loss=self.mlp_loss
-                )
-
-                if self.local_loss:
-                    logits_per_audio = logit_scale_a * audio_features @ all_text_features.T
-                    logits_per_text = logit_scale_a * text_features @ all_audio_features.T
-                else:
-                    logits_per_audio = logit_scale_a * all_audio_features @ all_text_features.T
-                    logits_per_text = logits_per_audio.T
-            else:
-                logits_per_audio = logit_scale_a * audio_features @ text_features.T
-                logits_per_text = logit_scale_a * text_features @ audio_features.T
-
-            # calculated ground-truth and cache if enabled
-            num_logits = logits_per_audio.shape[0]
-            if self.prev_num_logits != num_logits or device not in self.labels:
-                labels = torch.arange(num_logits, device=device, dtype=torch.long)
-                if self.world_size > 1 and self.local_loss:
-                    labels = labels + num_logits * self.rank
-                if self.cache_labels:
-                    self.labels[device] = labels
-                    self.prev_num_logits = num_logits
-            else:
-                labels = self.labels[device]
-            if not self.weighted_loss:
-                total_loss = (
-                    F.cross_entropy(logits_per_audio, labels) +
-                    F.cross_entropy(logits_per_text, labels)
-                    ) / 2
-            else:
-                audio_weight = (all_audio_features@all_audio_features.T).detach()
-                audio_weight = (torch.exp(torch.sum(audio_weight, axis=1)/(self.weight_loss_kappa*len(all_audio_features)))).detach()
-                text_weight = (all_text_features@all_text_features.T).detach()
-                text_weight = (torch.exp(torch.sum(text_weight, axis=1)/(self.weight_loss_kappa*len(all_text_features)))).detach()
-                total_loss = (
-                    F.cross_entropy(logits_per_audio, labels, weight=text_weight) +
-                    F.cross_entropy(logits_per_text, labels, weight=audio_weight)
-                    ) / 2
-        return total_loss
-
-def lp_gather_features(
-        pred,
-        target,
-        world_size=1,
-        use_horovod=False
-):
-    if use_horovod:
-        assert hvd is not None, 'Please install horovod'
-        with torch.no_grad():
-            all_preds = hvd.allgather(pred)
-            all_targets = hvd.allgath(target)
-    else:
-        gathered_preds = [torch.zeros_like(pred) for _ in range(world_size)]
-        gathered_targets = [torch.zeros_like(target) for _ in range(world_size)]
-
-        dist.all_gather(gathered_preds, pred)
-        dist.all_gather(gathered_targets, target)
-        all_preds = torch.cat(gathered_preds, dim=0)
-        all_targets = torch.cat(gathered_targets, dim=0)
-
-    return all_preds, all_targets
-
-
-def get_map(pred, target):
-    pred = torch.sigmoid(pred).numpy()
-    target = target.numpy()
-    return np.mean(average_precision_score(target, pred, average=None))
-
-def get_acc(pred, target):
-    pred = torch.argmax(pred,1).numpy()
-    target = torch.argmax(target,1).numpy()
-    return accuracy_score(target, pred)
-
-def get_mauc(pred, target):
-    pred = torch.sigmoid(pred).numpy()
-    target = target.numpy()
-    return np.mean(roc_auc_score(target, pred, average=None))
-
-
-class LPMetrics(object):
-    def __init__(self, metric_names = ['map','acc','mauc']):
-        self.metrics = []
-        for name in metric_names:
-            self.metrics.append(self.get_metric(name))
-        self.metric_names = metric_names
-
-    def get_metric(self,name):
-        if name == 'map':
-            return get_map
-        elif name == 'acc':
-            return get_acc
-        elif name == 'mauc':
-            return get_mauc
-        else:
-            raise ValueError(f'the metric should be at least one of [map, acc, mauc]')
-
-    def evaluate_mertics(self, pred, target):
-        metric_dict = {}
-        for i in range(len(self.metric_names)):
-            metric_dict[self.metric_names[i]] = self.metrics[i](pred, target)
-        return metric_dict
-
-
-def calc_celoss(pred, target):
-    target = torch.argmax(target, 1).long()
-    return nn.CrossEntropyLoss()(pred, target)
-
-
-class LPLoss(nn.Module):
-
-    def __init__(self, loss_name):
-        super().__init__()
-        if loss_name == 'bce':
-            self.loss_func = nn.BCEWithLogitsLoss()
-        elif loss_name == 'ce':
-            self.loss_func = calc_celoss
-        elif loss_name == 'mse':
-            self.loss_func = nn.MSELoss()
-        else:
-            raise ValueError(f'the loss func should be at least one of [bce, ce, mse]')
-
-    def forward(self, pred, target):
-        loss = self.loss_func(pred, target)
-        return loss
-        

utils/laion_clap/clap_module/tokenizer.py

@@ -1,180 +0,0 @@
-""" CLIP tokenizer
-
-Copied from https://github.com/openai/CLIP. Originally MIT License, Copyright (c) 2021 OpenAI.
-"""
-import gzip
-import html
-import os
-from functools import lru_cache
-from typing import Union, List
-
-import ftfy
-import regex as re
-import torch
-
-
-@lru_cache()
-def default_bpe():
-    return os.path.join(os.path.dirname(os.path.abspath(__file__)), "bpe_simple_vocab_16e6.txt.gz")
-
-
-@lru_cache()
-def bytes_to_unicode():
-    """
-    Returns list of utf-8 byte and a corresponding list of unicode strings.
-    The reversible bpe codes work on unicode strings.
-    This means you need a large # of unicode characters in your vocab if you want to avoid UNKs.
-    When you're at something like a 10B token dataset you end up needing around 5K for decent coverage.
-    This is a signficant percentage of your normal, say, 32K bpe vocab.
-    To avoid that, we want lookup tables between utf-8 bytes and unicode strings.
-    And avoids mapping to whitespace/control characters the bpe code barfs on.
-    """
-    bs = list(range(ord("!"), ord("~")+1))+list(range(ord("¡"), ord("¬")+1))+list(range(ord("®"), ord("ÿ")+1))
-    cs = bs[:]
-    n = 0
-    for b in range(2**8):
-        if b not in bs:
-            bs.append(b)
-            cs.append(2**8+n)
-            n += 1
-    cs = [chr(n) for n in cs]
-    return dict(zip(bs, cs))
-
-
-def get_pairs(word):
-    """Return set of symbol pairs in a word.
-    Word is represented as tuple of symbols (symbols being variable-length strings).
-    """
-    pairs = set()
-    prev_char = word[0]
-    for char in word[1:]:
-        pairs.add((prev_char, char))
-        prev_char = char
-    return pairs
-
-
-def basic_clean(text):
-    text = ftfy.fix_text(text)
-    text = html.unescape(html.unescape(text))
-    return text.strip()
-
-
-def whitespace_clean(text):
-    text = re.sub(r'\s+', ' ', text)
-    text = text.strip()
-    return text
-
-
-class SimpleTokenizer(object):
-    def __init__(self, bpe_path: str = default_bpe(), special_tokens=None):
-        self.byte_encoder = bytes_to_unicode()
-        self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}
-        merges = gzip.open(bpe_path).read().decode("utf-8").split('\n')
-        merges = merges[1:49152-256-2+1]
-        merges = [tuple(merge.split()) for merge in merges]
-        vocab = list(bytes_to_unicode().values())
-        vocab = vocab + [v+'</w>' for v in vocab]
-        for merge in merges:
-            vocab.append(''.join(merge))
-        if not special_tokens:
-            special_tokens = ['<start_of_text>', '<end_of_text>']
-        else:
-            special_tokens = ['<start_of_text>', '<end_of_text>'] + special_tokens
-        vocab.extend(special_tokens)
-        self.encoder = dict(zip(vocab, range(len(vocab))))
-        self.decoder = {v: k for k, v in self.encoder.items()}
-        self.bpe_ranks = dict(zip(merges, range(len(merges))))
-        self.cache = {t:t for t in special_tokens}
-        special = "|".join(special_tokens)
-        self.pat = re.compile(special + r"""|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+""", re.IGNORECASE)
-
-        self.vocab_size = len(self.encoder)
-        self.all_special_ids = [self.encoder[t] for t in special_tokens]
-
-    def bpe(self, token):
-        if token in self.cache:
-            return self.cache[token]
-        word = tuple(token[:-1]) + ( token[-1] + '</w>',)
-        pairs = get_pairs(word)
-
-        if not pairs:
-            return token+'</w>'
-
-        while True:
-            bigram = min(pairs, key = lambda pair: self.bpe_ranks.get(pair, float('inf')))
-            if bigram not in self.bpe_ranks:
-                break
-            first, second = bigram
-            new_word = []
-            i = 0
-            while i < len(word):
-                try:
-                    j = word.index(first, i)
-                    new_word.extend(word[i:j])
-                    i = j
-                except:
-                    new_word.extend(word[i:])
-                    break
-
-                if word[i] == first and i < len(word)-1 and word[i+1] == second:
-                    new_word.append(first+second)
-                    i += 2
-                else:
-                    new_word.append(word[i])
-                    i += 1
-            new_word = tuple(new_word)
-            word = new_word
-            if len(word) == 1:
-                break
-            else:
-                pairs = get_pairs(word)
-        word = ' '.join(word)
-        self.cache[token] = word
-        return word
-
-    def encode(self, text):
-        bpe_tokens = []
-        text = whitespace_clean(basic_clean(text)).lower()
-        for token in re.findall(self.pat, text):
-            token = ''.join(self.byte_encoder[b] for b in token.encode('utf-8'))
-            bpe_tokens.extend(self.encoder[bpe_token] for bpe_token in self.bpe(token).split(' '))
-        return bpe_tokens
-
-    def decode(self, tokens):
-        text = ''.join([self.decoder[token] for token in tokens])
-        text = bytearray([self.byte_decoder[c] for c in text]).decode('utf-8', errors="replace").replace('</w>', ' ')
-        return text
-
-
-_tokenizer = SimpleTokenizer()
-
-
-def tokenize(texts: Union[str, List[str]], context_length: int = 77) -> torch.LongTensor:
-    """
-    Returns the tokenized representation of given input string(s)
-
-    Parameters
-    ----------
-    texts : Union[str, List[str]]
-        An input string or a list of input strings to tokenize
-    context_length : int
-        The context length to use; all CLIP models use 77 as the context length
-
-    Returns
-    -------
-    A two-dimensional tensor containing the resulting tokens, shape = [number of input strings, context_length]
-    """
-    if isinstance(texts, str):
-        texts = [texts]
-
-    sot_token = _tokenizer.encoder["<start_of_text>"]
-    eot_token = _tokenizer.encoder["<end_of_text>"]
-    all_tokens = [[sot_token] + _tokenizer.encode(text) + [eot_token] for text in texts]
-    result = torch.zeros(len(all_tokens), context_length, dtype=torch.long)
-
-    for i, tokens in enumerate(all_tokens):
-        if len(tokens) > context_length:
-            tokens = tokens[:context_length]  # Truncate
-        result[i, :len(tokens)] = torch.tensor(tokens)
-
-    return result

utils/laion_clap/clap_module/utils.py

@@ -1,62 +1,6 @@
-import numpy as np
 import torch
 from torch import nn as nn
 from torchvision.ops.misc import FrozenBatchNorm2d
-import logging
-import h5py
-from tqdm import tqdm
-import random
-import json
-import os
-import pathlib
-
-# TODO: (yusong) this not a good place to store those information and does not scale. Need to be fixed later.
-dataset_split = {
-    "audiocaps": ["train", "valid", "test"],
-    "audioset": ["balanced_train", "unbalanced_train", "eval"],
-    "BBCSoundEffects": ["train", "test"],
-    "Clotho": ["train", "test", "valid"],
-    "free_to_use_sounds": ["train", "test"],
-    "paramount_motion": ["train", "test"],
-    "sonniss_game_effects": ["train", "test"],
-    "wesoundeffects": ["train", "test"],
-    "MACS": ["train", "test"],
-    "freesound": ["train", "test"],
-    "FSD50K": ["train", "test", "valid"],
-    "fsd50k_class_label": ["train", "test", "valid"],
-    "esc50": ["train", "test"],
-    "ESC50_1": ["train", "test"],
-    "ESC50_2": ["train", "test"],
-    "ESC50_3": ["train", "test"],
-    "ESC50_4": ["train", "test"],
-    "ESC50_5": ["train", "test"],
-    "audiostock": ["train", "test"],
-    "freesound_no_overlap_noesc50": ["train", "test"],
-    "epidemic_sound_effects": ["train", "test"],
-    "VGGSound": ["train", "test"],
-    "urbansound8k_class_label": ["train", "test"],
-    "audioset_t5": ["balanced_train", "unbalanced_train", "eval"],
-    "audioset_t5_debiased": ["balanced_train", "unbalanced_train", "eval"],
-    "epidemic_sound_effects_t5": ["train", "test"],
-    "epidemic_sound_effects_t5_debiased": ["train", "test"],
-    "WavText5K": ["train", "test"],
-    "esc50_no_overlap": ["train", "test"],
-    "usd8k_no_overlap": ["train", "test"],
-    "fsd50k_200_class_label": ["train", "test", "valid"],
-    "fma_full": ["train", "test"],
-    "Genius": ["train", "test"],
-    "Jamendo": ["train", "test"],
-    "juno": ["train", "test"],
-    "CMU_Arctic": ["train", "test"],
-    "ravdess": ["train", "test"],
-    "Europarl-st": ["train", "test"],
-    "common_voice": ["train", "test"],
-    "Jamendo_16bit": ["train", "test"],
-    "genius_16bit_128": ["train", "test"],
-    "juno_16bit": ["train", "test"],
-    "fma_full_16bit_128": ["train", "test"],
-    "GTZAN": ["train", "test"],
-    }
 
 
 def freeze_batch_norm_2d(module, module_match={}, name=""):
@@ -100,99 +44,6 @@ def freeze_batch_norm_2d(module, module_match={}, name=""):
     return res
 
 
-def exist(dataset_name, dataset_type):
-    """
-    Check if dataset exists
-    """
-    if dataset_type in dataset_split[dataset_name]:
-        return True
-    else:
-        return False
-
-
-def get_tar_path_from_dataset_name(
-        dataset_names,
-        dataset_types,
-        islocal,
-        dataset_path,
-        proportion=1,
-        full_dataset=None
-):
-    """
-    Get tar path from dataset name and type
-    """
-    output = []
-    for n in dataset_names:
-        if full_dataset is not None and n in full_dataset:
-            current_dataset_types = dataset_split[n]
-        else:
-            current_dataset_types = dataset_types
-        for s in current_dataset_types:
-            tmp = []
-            if islocal:
-                sizefilepath_ = f"{dataset_path}/{n}/{s}/sizes.json"
-                if not os.path.exists(sizefilepath_):
-                    sizefilepath_ = f"./json_files/{n}/{s}/sizes.json"
-            else:
-                sizefilepath_ = f"./json_files/{n}/{s}/sizes.json"
-            if not os.path.exists(sizefilepath_):
-                continue
-            sizes = json.load(open(sizefilepath_, "r"))
-            for k in sizes.keys():
-                if islocal:
-                    tmp.append(f"{dataset_path}/{n}/{s}/{k}")
-                else:
-                    tmp.append(
-                        f"pipe:aws s3 --cli-connect-timeout 0 cp s3://s-laion-audio/webdataset_tar/{n}/{s}/{k} -"
-                    )
-            if proportion != 1:
-                tmp = random.sample(tmp, int(proportion * len(tmp)))
-            output.append(tmp)
-    return sum(output, [])
-
-
-def get_tar_path_from_txts(txt_path, islocal, proportion=1):
-    """
-    Get tar path from txt path
-    """
-    if isinstance(txt_path, (list, tuple)):
-        return sum(
-            [
-                get_tar_path_from_txts(
-                    txt_path[i], islocal=islocal, proportion=proportion
-                )
-                for i in range(len(txt_path))
-            ],
-            [],
-        )
-    if isinstance(txt_path, str):
-        with open(txt_path) as f:
-            lines = f.readlines()
-        if islocal:
-            lines = [
-                lines[i]
-                .split("\n")[0]
-                .replace("pipe:aws s3 cp s3://s-laion-audio/", "/mnt/audio_clip/")
-                for i in range(len(lines))
-            ]
-        else:
-            lines = [
-                lines[i].split("\n")[0].replace(".tar", ".tar -")
-                for i in range(len(lines))
-            ]
-        if proportion != 1:
-            print("Sampling tars with proportion of {}".format(proportion))
-            lines = random.sample(lines, int(proportion * len(lines)))
-        return lines
-
-
-def get_mix_lambda(mixup_alpha, batch_size):
-    mixup_lambdas = [
-        np.random.beta(mixup_alpha, mixup_alpha, 1)[0] for _ in range(batch_size)
-    ]
-    return np.array(mixup_lambdas).astype(np.float32)
-
-
 def do_mixup(x, mixup_lambda):
     """
     Args:
@@ -240,150 +91,3 @@ def pad_framewise_output(framewise_output, frames_num):
 
     output = torch.cat((framewise_output, pad), dim=1)
     """(batch_size, frames_num, classes_num)"""
-
-
-def process_ipc(index_path, classes_num, filename):
-    # load data
-    logging.info("Load Data...............")
-    ipc = [[] for _ in range(classes_num)]
-    with h5py.File(index_path, "r") as f:
-        for i in tqdm(range(len(f["target"]))):
-            t_class = np.where(f["target"][i])[0]
-            for t in t_class:
-                ipc[t].append(i)
-    print(ipc)
-    np.save(filename, ipc)
-    logging.info("Load Data Succeed...............")
-
-
-def save_to_dict(s, o_={}):
-    sp = s.split(": ")
-    o_.update({sp[0]: float(sp[1])})
-    return o_
-
-
-def get_data_from_log(txt_path):
-    """
-    Output dictionary from out.txt log file
-    """
-    with open(txt_path) as f:
-        lines = f.readlines()
-    val_data = {}
-    train_data = {}
-    train_losses = []
-    train_losses_epoch = []
-    for i in range(len(lines)):
-        if "| INFO |" in lines[i]:
-            if "Eval Epoch" in lines[i]:
-                if "val_loss" in lines[i]:
-                    # float(regex.sub("", lines[310].split("	")[-1]).replace(" ", ""))
-                    line = lines[i].split("Eval Epoch: ")[-1]
-                    num_epoch = int(line.split("	")[0].split(" ")[0])
-                    d = {
-                        line.split("	")[0]
-                        .split(" ")[1]
-                        .replace(":", ""): float(line.split("	")[0].split(" ")[-1])
-                    }
-                    for i in range(1, len(line.split("	"))):
-                        d = save_to_dict(line.split("	")[i], d)
-                    val_data[num_epoch] = d
-            elif "Train Epoch" in lines[i]:
-                num_epoch = int(lines[i].split("Train Epoch: ")[1][0])
-                loss = float(lines[i].split("Loss: ")[-1].split(" (")[0])
-                train_losses.append(loss)
-                train_losses_epoch.append(num_epoch)
-    for i in range(len(train_losses)):
-        train_data[i] = {
-            "num_epoch": train_losses_epoch[i],
-            "train_loss": train_losses[i],
-        }
-    return train_data, val_data
-
-
-def save_p(obj, filename):
-    import pickle
-
-    try:
-        from deepdiff import DeepDiff
-    except:
-        os.system("pip install deepdiff")
-        from deepdiff import DeepDiff
-    with open(filename, "wb") as file:
-        pickle.dump(obj, file, protocol=pickle.HIGHEST_PROTOCOL)  # highest protocol
-    with open(filename, "rb") as file:
-        z = pickle.load(file)
-    assert (
-            DeepDiff(obj, z, ignore_string_case=True) == {}
-    ), "there is something wrong with the saving process"
-    return
-
-
-def load_p(filename):
-    import pickle
-
-    with open(filename, "rb") as file:
-        z = pickle.load(file)
-    return z
-
-
-def save_json(data, name="data.json"):
-    import json
-    with open(name, 'w') as fp:
-        json.dump(data, fp)
-    return
-
-
-def load_json(name):
-    import json
-    with open(name, 'r') as fp:
-        data = json.load(fp)
-    return data
-
-
-from multiprocessing import Process, Manager
-from multiprocessing import Process, Value, Array
-from ctypes import c_wchar
-
-
-def load_class_label(path):
-    # https://stackoverflow.com/questions/48004243/how-to-share-large-read-only-dictionary-list-across-processes-in-multiprocessing
-    # https://stackoverflow.com/questions/45693949/storing-strings-in-a-multiprocessing-sharedctypes-array
-    out = None
-    if path is not None:
-        if pathlib.Path(path).suffix in [".pkl", ".pickle"]:
-            out = load_p(path)
-        elif pathlib.Path(path).suffix in [".json", ".txt"]:
-            out = load_json(path)
-        elif pathlib.Path(path).suffix in [".npy", ".npz"]:
-            out = np.load(path)
-        elif pathlib.Path(path).suffix in [".csv"]:
-            import pandas as pd
-            out = pd.read_csv(path)
-    return out
-    # if out is None:
-    #     return None
-    # else:
-    #     key = Array(c_wchar, '\n'.join(list(out.keys())), lock=False)
-    #     val = Array('i', out.values(), lock=False)
-    #     return (key, val)
-
-
-from torch import optim
-
-
-def get_optimizer(params, lr, betas, eps, momentum, optimizer_name):
-    if optimizer_name.lower() == "adamw":
-        optimizer = optim.AdamW(
-            params, lr=lr, betas=betas, eps=eps
-        )
-    elif optimizer_name.lower() == "sgd":
-        optimizer = optim.SGD(
-            params, lr=lr, momentum=momentum
-        )
-    elif optimizer_name.lower() == "adam":
-        optimizer = optim.Adam(
-            params, lr=lr, betas=betas, eps=eps
-        )
-    else:
-        raise ValueError("optimizer name is not correct")
-    return optimizer

utils/laion_clap/clap_module/version.py

@@ -1 +0,0 @@
-__version__ = '0.2.1'

utils/laion_clap/evaluate/eval_dcase.py

@@ -1,150 +0,0 @@
-import torch
-import torch.nn.functional as F
-import torch.backends.cudnn as cudnn
-from open_clip import create_model
-from open_clip import tokenize
-import glob
-import json
-import librosa
-from tqdm import tqdm
-import numpy as np
-import os
-from laion_clap.training.params import parse_args
-
-
-def get_output_from_single_audio(audio, text, model, device):
-
-    # audio_embedding = model.audio_infer(audio, hopsize=5 * 48000, key="embedding", device=device)['embedding']
-    # if audio_embedding.ndim > 1:
-    #     audio_embedding = audio_embedding.mean(dim=0, keepdim=True)
-    # else:
-    #     audio_embedding = audio_embedding.unsqueeze(0)
-    audio_features = model(audio, None, device)
-    audio_features = F.normalize(audio_features, dim=-1)
-    text_features = model(None, text, device=device)
-    text_features = F.normalize(text_features, dim=-1)
-
-    # CHANGE: before normalize or after
-    audio_features_mlp = model.audio_transform(audio_features)
-    text_features_mlp = model.text_transform(text_features)
-    return audio_features, text_features, audio_features_mlp, text_features_mlp, model.logit_scale_a.exp(), model.logit_scale_t.exp()
-
-
-def get_metrics(text_to_audio_logits):
-    metrics = {}
-
-    # repeat ground truth 5 times because Clotho has 5 text for 1 audio
-    ground_truth = torch.repeat_interleave(torch.arange(len(text_features) // 5), 5).view(-1, 1)
-
-    ranking = torch.argsort(text_to_audio_logits, descending=True)
-    preds = torch.where(ranking == ground_truth)[1]  # (yusong) this line is slow because it uses single thread
-    preds = preds.detach().cpu().numpy()
-    metrics[f"mean_rank"] = preds.mean() + 1
-    metrics[f"median_rank"] = np.floor(np.median(preds)) + 1
-    for k in [1, 5, 10]:
-        metrics[f"R@{k}"] = np.mean(preds < k)
-    # map@10
-    metrics[f"mAP@10"] = np.mean(np.where(preds < 10, 1 / (preds + 1), 0.0))
-    return metrics
-
-
-if __name__ == '__main__':
-    args = parse_args()
-
-    model_path = args.pretrained
-
-    clotho_test_preprocessed_dir = "/fsx/yusong/clotho_test_set/test"
-
-    cudnn.benchmark = True
-    cudnn.deterministic = False
-
-    audio_features_ensemble_all = []
-    text_features_ensemble_all = []
-    audio_features_mlp_ensemble_all = []
-    text_features_mlp_ensemble_all = []
-    logit_scale_a_ensemble_all = []
-    logit_scale_t_ensemble_all = []
-
-
-    device = torch.device('cuda')
-    model, clap_model_cfg = create_model(
-        args.amodel,
-        args.tmodel,
-        args.pretrained,
-        precision=args.precision,
-        device=device,
-        jit=args.torchscript,
-        force_quick_gelu=args.force_quick_gelu,
-        openai_model_cache_dir=os.path.expanduser(args.openai_model_cache_dir),
-        skip_params=False
-    )
-
-    # load model
-    checkpoint = torch.load(model_path, map_location=device)
-    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)
-    else:
-        # loading a bare (model only) checkpoint for fine-tune or evaluation
-        model.load_state_dict(checkpoint)
-
-    model.to(device)
-    model.eval()
-    for param in model.parameters():
-        param.requires_grad = False
-
-    # take every 5th file because clotho has 5 texts for 1 audio
-    test_file_list = sorted(glob.glob(f"{clotho_test_preprocessed_dir}/*.flac"))
-
-    audio_features_all = []
-    text_features_all = []
-    audio_features_mlp_all = []
-    text_features_mlp_all = []
-    logit_scale_a_all = []
-    logit_scale_t_all = []
-
-    with torch.no_grad():
-        for file_path in tqdm(test_file_list):
-            json_path = file_path.replace(".flac", ".json")
-            with open(json_path, "r") as f:
-                json_data = json.load(f)
-            audio, sr = librosa.load(file_path, sr=48000, mono=True)
-            audio = torch.from_numpy(audio).to(device)
-            audio = {'waveform': audio.unsqueeze(0), 'sample_rate': sr}
-            text = json_data["text"]
-
-            if args.tmodel == "transformer":
-                from open_clip import tokenize
-                text = tokenize(text)
-            else:
-                from laion_clap.training.data import tokenizer
-                text = tokenizer(text, tmodel=args.tmodel)  # 5 texts for each audio
-
-            audio_features, text_features, audio_features_mlp, text_features_mlp, logit_scale_a, logit_scale_t = \
-                get_output_from_single_audio(audio, text, model, device)
-
-            audio_features_all.append(audio_features.detach().cpu())
-            text_features_all.append(text_features.detach().cpu())
-            audio_features_mlp_all.append(audio_features_mlp.detach().cpu())
-            text_features_mlp_all.append(text_features_mlp.detach().cpu())
-            logit_scale_a_all.append(logit_scale_a.detach().cpu())
-            logit_scale_t_all.append(logit_scale_t.detach().cpu())
-
-    audio_features = torch.cat(audio_features_all)
-    text_features = torch.cat(text_features_all)
-    logit_scale_a = logit_scale_a_all[0]
-
-    logits_per_audio = (logit_scale_a * audio_features @ text_features.t()).detach().cpu()
-    logits_per_text = logits_per_audio.t().detach().cpu()
-
-    metrics = get_metrics(
-        logits_per_text
-    )
-
-    print(metrics)

utils/laion_clap/evaluate/eval_linear_probe.py

@@ -1,515 +0,0 @@
-'''
-Evalute the linear probe performance on different checkpoints
-'''
-import logging
-import os
-import random
-from datetime import datetime
-import copy
-import numpy as np
-import torch
-import torch.backends.cudnn as cudnn
-from torch.cuda.amp import GradScaler
-import glob
-
-try:
-    import wandb
-except ImportError:
-    wandb = None
-
-try:
-    import torch.utils.tensorboard as tensorboard
-except ImportError:
-    tensorboard = None
-
-try:
-    import horovod.torch as hvd
-except ImportError:
-    hvd = None
-
-from clap_module import create_model_and_transforms, trace_model, create_model
-from training.data import get_data
-from training.params import parse_args
-from training.distributed import is_master, init_distributed_device, world_info_from_env
-from training.logger import setup_logging
-from training.scheduler import cosine_lr
-from training.lp_main import config_lp_optimizer
-from training.lp_train import train_one_epoch, evaluate
-from clap_module.utils import get_tar_path_from_dataset_name, dataset_split
-from clap_module.utils import load_p, load_class_label
-from clap_module.linear_probe import LinearProbe
-
-def maintain_ckpts(args, startidx, all_idx_len):
-    for i in reversed(range(startidx, all_idx_len)):
-        if os.path.exists(os.path.join(args.checkpoint_path, f"epoch_top_{i}.pt")):
-            os.rename(
-                os.path.join(args.checkpoint_path, f"epoch_top_{i}.pt"),
-                os.path.join(args.checkpoint_path, f"epoch_top_{i+1}.pt"),
-            )
-    if os.path.exists(
-        os.path.join(args.checkpoint_path, f"epoch_top_{all_idx_len}.pt")
-    ):
-        os.remove(os.path.join(args.checkpoint_path, f"epoch_top_{all_idx_len}.pt"))
-    return
-
-
-def update_top_k_performance(
-    new_metrics_inputs, current_top_k_ckpt_metrics, args, ckpt, bignumbetter=True, pretrain_epoch=0
-):
-    """
-    Record the top-k performance of the current epoch.
-    current_top_k_metrics is a dictionary of the form: {1: top_1_ckpt_measure, 2: top_2_ckpt_measure, ...}
-    """
-    if isinstance(new_metrics_inputs, (list, tuple)):
-        new_metrics_inputs = np.mean(new_metrics_inputs)
-        return update_top_k_performance(
-            new_metrics_inputs,
-            current_top_k_ckpt_metrics,
-            args=args,
-            ckpt=ckpt,
-            bignumbetter=bignumbetter,
-            pretrain_epoch=pretrain_epoch
-        )
-    elif isinstance(new_metrics_inputs, dict):
-        new_metrics_inputs = np.mean(list(new_metrics_inputs.values()))
-        return update_top_k_performance(
-            new_metrics_inputs,
-            current_top_k_ckpt_metrics,
-            args=args,
-            ckpt=ckpt,
-            bignumbetter=bignumbetter,
-            pretrain_epoch=pretrain_epoch
-        )
-    elif isinstance(new_metrics_inputs, (float, int)):
-        update_flag = {k: False for k in current_top_k_ckpt_metrics.keys()}
-        sorted_keys = sorted(current_top_k_ckpt_metrics.keys())
-        sorted_values = sorted(
-            current_top_k_ckpt_metrics.values(), reverse=bignumbetter
-        )
-        sorted_values_ = copy.deepcopy(sorted_values)
-        sorted_values.append(new_metrics_inputs)
-        sorted_values = sorted(sorted_values, reverse=bignumbetter)
-        sorted_values = sorted_values[:-1]
-
-        if sorted_values == sorted_values_:
-            return current_top_k_ckpt_metrics, new_metrics_inputs
-        else:
-            for i in range(len(sorted_keys)):
-                if current_top_k_ckpt_metrics[sorted_keys[i]] != sorted_values[i]:
-                    current_top_k_ckpt_metrics[sorted_keys[i]] = sorted_values[i]
-                    update_flag[sorted_keys[i]] = True
-            for i in range(len(update_flag)):
-                if update_flag[i]:
-                    maintain_ckpts(args, i, len(sorted_keys))
-                    torch.save(
-                        ckpt,
-                        os.path.join(args.checkpoint_path, f"pretrain_epoch_{pretrain_epoch}_lp_epoch_top_{i}.pt"),
-                    )
-                    break
-            return current_top_k_ckpt_metrics, new_metrics_inputs
-
-
-# def updateifNone(a, b):
-#     a = b if None else a
-#     return a
-
-
-def is_pretrained_params(n):
-    return (
-        n.startswith("clap_model.transformer")
-        or n in ["clap_model.positional_embedding", "clap_model.text_projection"]
-        or n.startswith("clap_model.token_embedding")
-        or n.startswith("clap_model.ln_final")
-        or n.startswith("clap_model.logit_scale_t")
-    )
-
-
-def random_seed(seed=42, rank=0):
-    torch.manual_seed(seed + rank)
-    np.random.seed(seed + rank)
-    random.seed(seed + rank)
-
-def main():
-    args = parse_args()
-    # sanitize model name for filesystem / uri use, easier if we don't use / in name as a rule?
-    args.amodel = args.amodel.replace("/", "-")
-    
-    pretrained_ckpts = sorted(glob.glob(os.path.join(args.pretrained, "*.pt")), key=os.path.getmtime)
-
-    if args.name is None:
-        args.name = "-".join(
-            [
-                datetime.now().strftime("%Y_%m_%d-%H_%M_%S"),
-                f"linear_probe"
-                f"model_{args.amodel}",
-                f"lr_{args.lr}",
-                f"b_{args.batch_size}",
-                f"j_{args.workers}",
-                f"p_{args.precision}",
-            ]
-        )
-
-    # discover initial world args early so we can log properly
-    args.distributed = False
-    args.local_rank, args.rank, args.world_size = world_info_from_env()
-
-    if args.remotedata and is_master(args):
-        for dataset_name in args.datasetnames:
-            for split in dataset_split[dataset_name]:
-                if not os.path.exists(f"./json_files/{dataset_name}/{split}"):
-                    os.makedirs(f"./json_files/{dataset_name}/{split}")
-                os.system(
-                    f"aws s3 cp s3://s-laion-audio/webdataset_tar/{dataset_name}/{split}/sizes.json ./json_files/{dataset_name}/{split}/sizes.json"
-                )
-    args.log_path = None
-    if is_master(args, local=args.log_local):
-        log_base_path = os.path.join(args.logs, args.name)
-        os.makedirs(log_base_path, exist_ok=True)
-        log_filename = f"out-{args.rank}" if args.log_local else "out.log"
-        args.log_path = os.path.join(log_base_path, log_filename)
-
-        # avoid log dir in same name:
-        postfix = 0
-        while os.path.exists(args.log_path):
-            postfix += 1
-            log_base_path_new = log_base_path+'-'+str(postfix)
-            os.makedirs(log_base_path_new, exist_ok=True)
-            log_filename = f"out-{args.rank}" if args.log_local else "out.log"
-            args.log_path = os.path.join(log_base_path_new, log_filename)
-            # print(
-            #     "Error. Experiment already exists. Use --name {} to specify a new experiment."
-            # )
-            # return -1
-
-    # Set logger
-    args.log_level = logging.DEBUG if args.debug else logging.INFO
-    setup_logging(args.log_path, args.log_level)
-
-    # fully initialize distributed device environment
-    device = init_distributed_device(args)
-
-    args.wandb = "wandb" in args.report_to or "all" in args.report_to
-    args.tensorboard = "tensorboard" in args.report_to or "all" in args.report_to
-    if is_master(args):
-        args.tensorboard_path = (
-            os.path.join(args.logs, args.name, "tensorboard")
-            if args.tensorboard
-            else ""
-        )
-        args.checkpoint_path = os.path.join(args.logs, args.name, "checkpoints")
-        for dirname in [args.tensorboard_path, args.checkpoint_path]:
-            if dirname:
-                os.makedirs(dirname, exist_ok=True)
-    else:
-        args.tensorboard_path = ""
-        args.checkpoint_path = ""
-
-    if args.copy_codebase:
-        copy_codebase(args)
-
-    assert args.precision in ["amp", "fp16", "fp32"]
-    if args.precision == "fp16":
-        logging.warning(
-            "It is recommended to use AMP mixed-precision instead of FP16. "
-            "FP16 support needs further verification and tuning, especially for train."
-        )
-
-    if args.horovod:
-        logging.info(
-            f"Running in horovod mode with multiple processes / nodes. Device: {args.device}."
-            f"Process (global: {args.rank}, local {args.local_rank}), total {args.world_size}."
-        )
-    elif args.distributed:
-        logging.info(
-            f"Running in distributed mode with multiple processes. Device: {args.device}."
-            f"Process (global: {args.rank}, local {args.local_rank}), total {args.world_size}."
-        )
-    else:
-        logging.info(f"Running with a single process. Device {args.device}.")
-
-    logging.info(f'openai cache dir: {os.path.expanduser(args.openai_model_cache_dir)}')
-
-    # determine if this worker should save logs and checkpoints. only do so if it is rank == 0
-    args.save_logs = args.logs and args.logs.lower() != "none" and is_master(args)
-    writer = None
-    if args.save_logs and args.tensorboard:
-        assert tensorboard is not None, "Please install tensorboard."
-        writer = tensorboard.SummaryWriter(args.tensorboard_path)
-
-    if args.wandb and is_master(args):
-        assert wandb is not None, "Please install wandb."
-        logging.debug("Starting wandb.")
-        # you will have to configure this for your project!
-        wandb.init(
-            project="clap",
-            notes=args.wandb_notes,
-            name=args.wandb_notes,
-            tags=[],
-            config=vars(args),
-        )
-        logging.debug("Finished loading wandb.")
-
-    for idx, f in enumerate(pretrained_ckpts):
-        logging.info(f"pretrained on {f}")
-        args.pretrained = f
-        ckpt = torch.load(f, map_location='cpu')
-        pretrain_epoch = 0
-        if 'epoch' in ckpt:
-            pretrain_epoch = ckpt['epoch']
-        # train 
-        best_metrics = lp_main(args, device, writer, pretrain_epoch, idx)
-        
-        if args.wandb and is_master(args):
-            assert wandb is not None, "Please install wandb."
-            for name, val in best_metrics.items():
-                wandb.log({f"val/summary/{name}": val, "epoch": pretrain_epoch})
-
-    if args.wandb and is_master(args):
-        wandb.finish()
-
-def update_metric(best_metric, new_metric):
-    for key in new_metric:
-        if key not in best_metric:
-            best_metric[key] = new_metric[key]
-        else:
-            best_metric[key] = max(best_metric[key], new_metric[key])
-    return best_metric
-
-def lp_main(args, device, writer, pretrain_epoch, idx):
-
-    random.seed(args.seed)
-    torch.manual_seed(args.seed)
-    torch.cuda.manual_seed(args.seed)
-    torch.cuda.manual_seed_all(args.seed)
-    np.random.seed(args.seed)
-    args.class_index_dict = load_class_label(args.class_label_path)
-
-
-    # Create CLAP model
-    clap_model, clap_model_cfg = create_model(
-        args.amodel,
-        args.tmodel,
-        args.pretrained,
-        precision=args.precision,
-        device=device,
-        jit=args.torchscript,
-        force_quick_gelu=args.force_quick_gelu,
-        openai_model_cache_dir=os.path.expanduser(args.openai_model_cache_dir),
-        skip_params=False,
-        enable_fusion=args.enable_fusion,
-        fusion_type=args.fusion_type
-    )
-    
-    args.lp_out_ch = len(list(args.class_index_dict.keys()))
-    # Linear Probe 
-    if idx == 0:
-        logging.info(f"linear probe using mlp: {args.lp_mlp}")
-        logging.info(f"linear probe using freeze: {args.lp_freeze}")
-        logging.info(f"linear probe act layer: {args.lp_act}")
-        logging.info(f"linear probe out ch: {args.lp_out_ch}")
-        logging.info(f"linear probe learning rate (if applicable): {args.lp_lr}")
-        logging.info(f"linear probe loss func: {args.lp_loss}")
-        logging.info(f"linear probe lp_metrics: {args.lp_metrics}")
-
-    model = LinearProbe(
-        clap_model, 
-        mlp=args.lp_mlp, freeze=args.lp_freeze, 
-        in_ch=512, out_ch=args.lp_out_ch,
-        act=args.lp_act
-    ) # in_ch is fixed (i.e., 512)
-    model = model.to(device)
-    
-    if args.horovod:
-        with torch.no_grad():
-            for param in model.parameters():
-                param.set_(param.contiguous())
-
-    if args.trace:
-        model = trace_model(model, batch_size=args.batch_size, device=device)
-
-    if is_master(args) and idx == 0:
-        logging.info("Linear Probe CLAP Model:")
-        logging.info(f"{str(clap_model)}")
-        logging.info("Params:")
-        params_file = os.path.join(args.logs, args.name, "params.txt")
-        with open(params_file, "w") as f:
-            for name in sorted(vars(args)):
-                val = getattr(args, name)
-                logging.info(f"  {name}: {val}")
-                f.write(f"{name}: {val}\n")
-    
-
-    if args.distributed and not args.horovod:
-        if args.use_bn_sync:
-            model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
-        ddp_args = {}
-        if args.ddp_static_graph:
-            # this doesn't exist in older PyTorch, arg only added if enabled
-            ddp_args["static_graph"] = True
-        model = torch.nn.parallel.DistributedDataParallel(
-            model, device_ids=[device], find_unused_parameters=True, **ddp_args
-        )
-
-    data = get_data(args, clap_model_cfg)
-    assert len(data), "At least one train or eval dataset must be specified."
-    if args.trace:
-        assert "train" not in data, "Cannot train with traced model"
-
-    optimizer, scheduler, text_freeze_parameters = config_lp_optimizer(model, data, args)
-
-    scaler = GradScaler() if args.precision == "amp" else None
-
-    # optionally resume from a checkpoint
-    start_epoch = 0
-    if args.resume is not None:
-        if os.path.isfile(args.resume):
-            checkpoint = torch.load(args.resume, map_location=device)
-            if "epoch" in checkpoint:
-                # resuming a train checkpoint w/ epoch and optimizer state
-                start_epoch = checkpoint["epoch"]
-                sd = checkpoint["state_dict"]
-                if not args.distributed and next(iter(sd.items()))[0].startswith(
-                    "module"
-                ):
-                    sd = {k[len("module.") :]: v for k, v in sd.items()}
-                model.load_state_dict(sd)
-                if args.split_opt:
-                    if optimizer is not None:
-                        for k, o_ in optimizer.items():
-                            o_.load_state_dict(checkpoint[k + "_" + "optimizer"])
-                if optimizer is not None:
-                    optimizer.load_state_dict(checkpoint["optimizer"])
-                if scaler is not None and "scaler" in checkpoint:
-                    scaler.load_state_dict(checkpoint["scaler"])
-                logging.info(
-                    f"=> resuming checkpoint '{args.resume}' (epoch {start_epoch})"
-                )
-            else:
-                # loading a bare (model only) checkpoint for fine-tune or evaluation
-                model.load_state_dict(checkpoint)
-                logging.info(
-                    f"=> loaded checkpoint '{args.resume}' (epoch {start_epoch})"
-                )
-            if args.freeze_text:
-                print("Freeze Text!!!!")
-                for k in text_freeze_parameters:
-                    k.requires_grad = False
-        else:
-            logging.info("=> no checkpoint found at '{}'".format(args.resume))
-
-    cudnn.benchmark = True
-    cudnn.deterministic = False
-
-    if args.wandb and is_master(args):
-        args.train_sz = data["train"].dataloader.num_samples
-        if args.val_data is not None:
-            args.val_sz = data["val"].dataloader.num_samples
-        if args.debug:
-            wandb.watch(model, log="all")
-        if idx == 0:
-            wandb.save(params_file)
-
-    best_metrics = {}
-
-    if "train" not in data:
-        metric = evaluate(model, data, start_epoch, args, writer, extra_suffix="_pe@" + str(pretrain_epoch))
-        if is_master(args):
-            best_metrics = update_metric(best_metrics, metric)
-        return
-    elif start_epoch == 0 and "val" in data and not args.no_eval:
-        metric = evaluate(model, data, 0, args, writer, extra_suffix="_pe@" + str(pretrain_epoch))
-        if is_master(args):
-            best_metrics = update_metric(best_metrics, metric)
-    if args.save_top_performance:
-        current_top_k_ckpt_metrics = {
-            i: 0 for i in range(args.save_top_performance)
-        }  # initialize the top-k metric for ckpts to 0
-
-    for epoch in range(start_epoch, args.epochs):
-        # freeze the text param after (include) args.freeze_text_after, this is -1 by default
-        if epoch == args.freeze_text_after:
-            print("Text pretrained parameters are freezed since this epoch.")
-            for k in text_freeze_parameters:
-                k.requires_grad = False
-        if is_master(args):
-            logging.info(f"Start epoch {epoch}")
-
-        train_one_epoch(model, data, epoch, optimizer, scaler, scheduler, args, writer, extra_suffix="_pe@" + str(pretrain_epoch))
-        completed_epoch = epoch + 1
-
-        if any(v in data for v in ("val", "imagenet-val", "imagenet-v2")) and not args.no_eval:
-            metric = evaluate(model, data, completed_epoch, args, writer, extra_suffix="_pe@" + str(pretrain_epoch))
-            if is_master(args):
-                best_metrics = update_metric(best_metrics, metric)
-            if args.save_top_performance:
-                top_k_dataset = args.top_k_checkpoint_select_dataset
-                top_k_metric = args.top_k_checkpoint_select_metric
-                filtered_metrics = [
-                    v
-                    for k, v in metric.items()
-                    if top_k_metric in k and top_k_dataset in k
-                ]  # check all R@10 metrics (all dataset) and use it to update the ckpt
-        # Saving checkpoints.
-        if args.save_logs:
-            opt_dict = {
-                    k + "_" + "optimizer": v.state_dict() for k, v in optimizer.items()
-                }
-            checkpoint_dict = {
-                "epoch": completed_epoch,
-                "pretrain_epoch": pretrain_epoch,
-                "name": args.name,
-                "state_dict": model.state_dict(),
-            }
-            checkpoint_dict.update(opt_dict)
-            if scaler is not None:
-                checkpoint_dict["scaler"] = scaler.state_dict()
-
-            if completed_epoch == args.epochs or (
-                args.save_frequency > 0 and (completed_epoch % args.save_frequency) == 0
-            ):
-                torch.save(
-                    checkpoint_dict,
-                    os.path.join(args.checkpoint_path, f"pretrain_epoch_{pretrain_epoch}_lp_epoch_{completed_epoch}.pt"),
-                )
-            if args.save_most_recent:
-                torch.save(
-                    checkpoint_dict,
-                    os.path.join(args.checkpoint_path, f"pretrain_epoch_{pretrain_epoch}_lp_epoch_latest.pt"),
-                )
-            if args.save_top_performance and not args.no_eval:
-                update_top_k_performance(
-                    filtered_metrics,
-                    current_top_k_ckpt_metrics,
-                    args,
-                    checkpoint_dict,
-                    bignumbetter=True,
-                    pretrain_epoch=pretrain_epoch
-                )
-    del clap_model
-    return best_metrics
-
-
-def copy_codebase(args):
-    from shutil import copytree, ignore_patterns
-
-    new_code_path = os.path.join(args.logs, args.name, "code")
-    if os.path.exists(new_code_path):
-        print(
-            f"Error. Experiment already exists at {new_code_path}. Use --name to specify a new experiment."
-        )
-        return -1
-    print(f"Copying codebase to {new_code_path}")
-    current_code_path = os.path.realpath(__file__)
-    for _ in range(3):
-        current_code_path = os.path.dirname(current_code_path)
-    copytree(
-        current_code_path, new_code_path, ignore=ignore_patterns("log", "logs", "wandb")
-    )
-    print("Done copying code.")
-    return 1
-
-
-if __name__ == "__main__":
-    main()
-
-

utils/laion_clap/evaluate/eval_retrieval.py

@@ -1,192 +0,0 @@
-import os.path
-import glob
-import random
-import numpy as np
-import logging
-import wandb
-import torch
-import torch.backends.cudnn as cudnn
-from laion_clap import create_model
-from laion_clap.training.logger import setup_logging
-from laion_clap.training.data import get_data
-from laion_clap.training.train import evaluate
-from laion_clap.utils import get_tar_path_from_dataset_name, dataset_split
-from laion_clap.training.params import parse_args
-
-
-def find_params_value(file, key):
-    # find value of params in params_file
-    with open(file, 'r') as f:
-        for line in f:
-            if key + ': ' in line:
-                return line.split(': ')[1].strip()
-    return None
-
-
-if __name__ == '__main__':
-    # (yusong) repeated run might have different metric results.
-    # This is because we randomly select crop 10s for each audio.
-    args = parse_args()
-
-    if os.path.isdir(args.pretrained):
-        log_dir = os.path.dirname(args.pretrained)
-    else:
-        log_dir = os.path.dirname(os.path.dirname(args.pretrained))
-
-    args.log_level = logging.DEBUG if args.debug else logging.INFO
-    log_path = os.path.join(log_dir, 'out.log')
-    setup_logging(log_path, args.log_level)
-    params_file = os.path.join(log_dir, 'params.txt')
-
-    seed = 3407
-    random.seed(seed)
-    torch.manual_seed(seed)
-    torch.cuda.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)
-    np.random.seed(seed)
-
-    cudnn.benchmark = True
-    cudnn.deterministic = False
-    pretrained = 'openai'
-    amodel = find_params_value(params_file, 'amodel')
-    tmodel = find_params_value(params_file, 'tmodel')
-
-    if amodel is None or tmodel is None:
-        raise ValueError('model type not found in params file')
-
-    # set up dummy values for args
-    args.parallel_eval = False
-    args.rank = 0
-    args.local_rank = 0
-    args.world_size = 1
-    args.val_frequency = 1
-    args.epochs = 1
-    args.precision = 'fp32'
-    args.save_logs = True
-    args.wandb = True
-    args.class_index_dict = None
-
-    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
-    args.device = device
-
-    if args.remotedata:
-        for dataset_name in args.datasetnames:
-            for split in dataset_split[dataset_name]:
-                if not os.path.exists(f"./json_files/{dataset_name}/{split}"):
-                    os.makedirs(f"./json_files/{dataset_name}/{split}")
-                os.system(
-                    f"aws s3 cp s3://s-laion-audio/webdataset_tar/{dataset_name}/{split}/sizes.json ./json_files/{dataset_name}/{split}/sizes.json"
-                )
-
-    if args.datasetinfos is None:
-        args.datasetinfos = ["train", "unbalanced_train", "balanced_train"]
-    if args.dataset_type == "webdataset":
-        args.train_data = get_tar_path_from_dataset_name(
-            args.datasetnames,
-            args.datasetinfos,
-            islocal=not args.remotedata,
-            proportion=args.dataset_proportion,
-            dataset_path=args.datasetpath,
-        )
-        args.val_data = get_tar_path_from_dataset_name(
-            args.datasetnames,
-            ["valid", "test", "eval"],
-            islocal=not args.remotedata,
-            proportion=1,
-            dataset_path=args.datasetpath,
-        )
-    model, model_cfg = create_model(
-        amodel,
-        tmodel,
-        pretrained,
-        precision='fp32',
-        device=device,
-        jit=False,
-        force_quick_gelu=False,
-        openai_model_cache_dir=os.path.expanduser(args.openai_model_cache_dir),
-        skip_params=False,
-        enable_fusion=args.enable_fusion,
-        fusion_type=args.fusion_type
-    )  # a hack to get model_cfg
-
-    data = get_data(args, model_cfg=model_cfg)  # (yusong): hack: no model_cfg needed to get data
-
-    writer = None  # if use tensorboard, initalize writer here
-
-    if args.wandb:
-        assert wandb is not None, "Please install wandb."
-
-        # # find the line with "wandb_notes" and get the value
-        # wandb_notes = find_params_value(params_file, 'wandb_notes')
-        # if wandb_notes is None:
-        #     print(f'wandb_notes not found in params file: {params_file}, set to timestamp.')
-        #     wandb_notes = f'experiment_{time.strftime("%Y%m%d-%H%M%S")}'
-        # wandb_notes = wandb_notes + '-eval-retrieval'
-        wandb_notes = args.wandb_notes
-
-        logging.debug("Starting wandb.")
-        args.train_sz = data["train"].dataloader.num_samples
-        if args.val_data is not None:
-            args.val_sz = data["val"].dataloader.num_samples
-        # you will have to configure this for your project!
-        if args.wandb_id is not None:
-            wandb.init(
-                project="clap",
-                id=args.wandb_id,
-                resume=True
-            )
-        else:
-            wandb.init(
-                project="clap",
-                notes=wandb_notes,
-                name=wandb_notes,
-                tags=[],
-                config=vars(args),
-            )
-        logging.debug("Finished loading wandb.")
-
-    if os.path.isdir(args.pretrained):
-        all_model_checkpoints = sorted(glob.glob(os.path.join(log_dir, 'checkpoints', '*.pt')), key=os.path.getmtime)
-    else:
-        all_model_checkpoints = [args.pretrained]
-    for model_path in all_model_checkpoints:
-        args.checkpoint_path = os.path.dirname(model_path)
-        model, model_cfg = create_model(
-            amodel,
-            tmodel,
-            pretrained,
-            precision='fp32',
-            device=device,
-            jit=False,
-            force_quick_gelu=False,
-            openai_model_cache_dir=os.path.expanduser(args.openai_model_cache_dir),
-            skip_params=False,
-            enable_fusion=args.enable_fusion,
-            fusion_type=args.fusion_type
-        )
-
-        # load model
-        checkpoint = torch.load(model_path, map_location=device)
-        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
-
-        evaluate(model, data, start_epoch, args, writer)

utils/laion_clap/evaluate/eval_retrieval_main.py

@@ -1,257 +0,0 @@
-import os.path
-import glob
-import random
-import numpy as np
-import logging
-import wandb
-import torch
-import torch.nn.functional as F
-import torch.backends.cudnn as cudnn
-from clap_module import create_model
-from clap_module import tokenize
-from training.logger import setup_logging
-from training.data import get_data
-from training.train import evaluate
-from clap_module.utils import get_tar_path_from_dataset_name, dataset_split
-from training.params import parse_args
-
-
-def find_params_value(file, key):
-    # find value of params in params_file
-    with open(file, 'r') as f:
-        for line in f:
-            if key + ': ' in line:
-                return line.split(': ')[1].strip()
-    return None
-
-
-def evaluate_zeroshot(model, data, start_epoch, args, writer):
-    dataloader = data["val"].dataloader
-    metrics = {}
-    device = torch.device(args.device)
-    model.eval()
-    metrics.update({"epoch": start_epoch})
-
-    all_audio_features = []
-    all_class_labels = []
-    with torch.no_grad():
-        for i, batch in enumerate(dataloader):
-            audios = batch  # contains mel_spec, wavform, and longer list
-            audio_features = model(audios, None, device)
-            audio_features = F.normalize(audio_features, dim=-1)
-            all_audio_features.append(audio_features.detach().cpu())
-            all_class_labels.append(torch.argmax(batch["class_label"], 1).long())
-        all_audio_features = torch.cat(all_audio_features, dim=0)
-        all_class_labels = torch.cat(all_class_labels, dim=0)
-        metrics["num_samples"] = all_audio_features.shape[0]
-
-        # get text features
-        all_texts = ["This is a sound of " + t for t in args.class_index_dict.keys()]
-        # (yusong): a hack, can make it better
-        if args.tmodel == "transformer":
-            from clap_module.tokenizer import tokenize
-            all_texts = tokenize(all_texts)
-        else:
-            from training.data import tokenizer
-            all_texts = tokenizer(all_texts)
-        all_text_features = model(None, all_texts, device)
-        all_text_features = F.normalize(all_text_features, dim=-1).detach().cpu()
-
-        # compute similarity
-        logit_scale_a, logit_scale_t = model(None, None, device)
-        logit_scale_a = logit_scale_a.cpu()
-
-        logits_per_audio = (logit_scale_a * all_audio_features @ all_text_features.t()).detach().cpu()
-        logits_per_text = logits_per_audio.t().detach().cpu()
-
-        ground_truth = all_class_labels.view(-1, 1)
-        logit = logits_per_audio
-
-        ranking = torch.argsort(logit, descending=True)
-        preds = torch.where(ranking == ground_truth)[1]  # (yusong) this line is slow because it uses single thread
-        preds = preds.detach().cpu().numpy()
-        metrics[f"{args.datasetnames[0]}_mean_rank"] = preds.mean() + 1
-        metrics[f"{args.datasetnames[0]}_median_rank"] = np.floor(np.median(preds)) + 1
-        for k in [1, 5, 10]:
-            metrics[f"{args.datasetnames[0]}_R@{k}"] = np.mean(preds < k)
-        # map@10
-        metrics[f"{args.datasetnames[0]}_mAP@10"] = np.mean(np.where(preds < 10, 1 / (preds + 1), 0.0))
-
-        logging.info(
-            f"Eval Epoch: {start_epoch} "
-            + "\t".join([f"{k}: {round(v, 4):.4f}" for k, v in metrics.items()])
-        )
-
-        if args.wandb:
-            assert wandb is not None, "Please install wandb."
-            for name, val in metrics.items():
-                wandb.log({f"val/{name}": val, "epoch": start_epoch})
-
-
-if __name__ == '__main__':
-    # (yusong) repeated run might have different metric results.
-    # This is because we randomly select crop 10s for each audio.
-    args = parse_args()
-
-    if os.path.isdir(args.pretrained):
-        log_dir = os.path.dirname(args.pretrained)
-    else:
-        log_dir = os.path.dirname(os.path.dirname(args.pretrained))
-
-    args.log_level = logging.DEBUG if args.debug else logging.INFO
-    log_path = os.path.join(log_dir, 'out.log')
-    setup_logging(log_path, args.log_level)
-    params_file = os.path.join(log_dir, 'params.txt')
-
-    seed = 3407
-    random.seed(seed)
-    torch.manual_seed(seed)
-    torch.cuda.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)
-    np.random.seed(seed)
-
-    cudnn.benchmark = True
-    cudnn.deterministic = False
-    pretrained = 'openai'
-    amodel = find_params_value(params_file, 'amodel')
-    tmodel = find_params_value(params_file, 'tmodel')
-
-    if amodel is None or tmodel is None:
-        raise ValueError('model type not found in params file')
-
-    # set up dummy values for args
-    args.parallel_eval = False
-    args.rank = 0
-    args.local_rank = 0
-    args.world_size = 1
-    args.val_frequency = 1
-    args.epochs = 1
-    args.precision = 'fp32'
-    args.save_logs = True
-    args.wandb = args.report_to == 'wandb'
-    args.class_index_dict = None
-
-    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
-    args.device = device
-
-    if args.remotedata:
-        for dataset_name in args.datasetnames:
-            for split in dataset_split[dataset_name]:
-                if not os.path.exists(f"./json_files/{dataset_name}/{split}"):
-                    os.makedirs(f"./json_files/{dataset_name}/{split}")
-                os.system(
-                    f"aws s3 cp s3://s-laion-audio/webdataset_tar/{dataset_name}/{split}/sizes.json ./json_files/{dataset_name}/{split}/sizes.json"
-                )
-
-    if args.datasetinfos is None:
-        args.datasetinfos = ["train", "unbalanced_train", "balanced_train"]
-    if args.dataset_type == "webdataset":
-        args.train_data = get_tar_path_from_dataset_name(
-            args.datasetnames,
-            args.datasetinfos,
-            islocal=not args.remotedata,
-            proportion=args.dataset_proportion,
-            dataset_path=args.datasetpath,
-        )
-        args.val_data = get_tar_path_from_dataset_name(
-            args.datasetnames,
-            ["valid", "test", "eval"],
-            islocal=not args.remotedata,
-            proportion=1,
-            dataset_path=args.datasetpath,
-        )
-    model, model_cfg = create_model(
-        amodel,
-        tmodel,
-        pretrained,
-        precision='fp32',
-        device=device,
-        jit=False,
-        force_quick_gelu=False,
-        openai_model_cache_dir=os.path.expanduser(args.openai_model_cache_dir),
-        skip_params=False,
-        enable_fusion=args.enable_fusion,
-        fusion_type=args.fusion_type
-    )  # a hack to get model_cfg
-
-    data = get_data(args, model_cfg=model_cfg)  # (yusong): hack: no model_cfg needed to get data
-
-    writer = None  # if use tensorboard, initalize writer here
-
-    if args.wandb:
-        assert wandb is not None, "Please install wandb."
-
-        # # find the line with "wandb_notes" and get the value
-        # wandb_notes = find_params_value(params_file, 'wandb_notes')
-        # if wandb_notes is None:
-        #     print(f'wandb_notes not found in params file: {params_file}, set to timestamp.')
-        #     wandb_notes = f'experiment_{time.strftime("%Y%m%d-%H%M%S")}'
-        # wandb_notes = wandb_notes + '-eval-retrieval'
-        wandb_notes = args.wandb_notes
-
-        logging.debug("Starting wandb.")
-        args.train_sz = data["train"].dataloader.num_samples
-        if args.val_data is not None:
-            args.val_sz = data["val"].dataloader.num_samples
-        # you will have to configure this for your project!
-        if args.wandb_id is not None:
-            wandb.init(
-                project="clap",
-                id=args.wandb_id,
-                resume=True
-            )
-        else:
-            wandb.init(
-                project="clap",
-                notes=wandb_notes,
-                name=wandb_notes,
-                tags=[],
-                config=vars(args),
-            )
-        logging.debug("Finished loading wandb.")
-
-    if os.path.isdir(args.pretrained):
-        all_model_checkpoints = sorted(glob.glob(os.path.join(log_dir, 'checkpoints', '*.pt')), key=os.path.getmtime)
-    else:
-        all_model_checkpoints = [args.pretrained]
-    for model_path in all_model_checkpoints:
-        args.checkpoint_path = os.path.dirname(model_path)
-        model, model_cfg = create_model(
-            amodel,
-            tmodel,
-            pretrained,
-            precision='fp32',
-            device=device,
-            jit=False,
-            force_quick_gelu=False,
-            openai_model_cache_dir=os.path.expanduser(args.openai_model_cache_dir),
-            skip_params=False,
-            enable_fusion=args.enable_fusion,
-            fusion_type=args.fusion_type
-        )
-
-        # load model
-        checkpoint = torch.load(model_path, map_location=device)
-        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
-
-        evaluate_zeroshot(model, data, start_epoch, args, writer)

utils/laion_clap/evaluate/eval_zeroshot_classification.py

@@ -1,261 +0,0 @@
-import os.path
-import glob
-import random
-import numpy as np
-import logging
-import wandb
-import torch
-import torch.nn.functional as F
-import torch.backends.cudnn as cudnn
-from clap_module import create_model
-from clap_module import tokenize
-from training.logger import setup_logging
-from training.data import get_data
-from training.train import evaluate
-from clap_module.utils import get_tar_path_from_dataset_name, dataset_split
-from training.params import parse_args
-
-
-def find_params_value(file, key):
-    # find value of params in params_file
-    with open(file, 'r') as f:
-        for line in f:
-            if key + ': ' in line:
-                return line.split(': ')[1].strip()
-    return None
-
-
-def evaluate_zeroshot(model, data, start_epoch, args, writer):
-    dataloader = data["val"].dataloader
-    metrics = {}
-    device = torch.device(args.device)
-    model.eval()
-    metrics.update({"epoch": start_epoch})
-
-    all_audio_features = []
-    all_class_labels = []
-    with torch.no_grad():
-        for i, batch in enumerate(dataloader):
-            audios = batch  # contains mel_spec, wavform, and longer list
-            audio_features = model(audios, None, device)
-            audio_features = F.normalize(audio_features, dim=-1)
-            all_audio_features.append(audio_features.detach().cpu())
-            all_class_labels.append(torch.argmax(batch["class_label"], 1).long())
-        all_audio_features = torch.cat(all_audio_features, dim=0)
-        all_class_labels = torch.cat(all_class_labels, dim=0)
-        metrics["num_samples"] = all_audio_features.shape[0]
-
-        # get text features
-        if args.val_dataset_names == ['GTZAN']:
-            all_texts = [f"This is a {t} song." for t in args.class_index_dict.keys()]
-        else:
-            all_texts = [f"This is a sound of {t}." for t in args.class_index_dict.keys()]
-        logging.info(f'class label prompts: {all_texts}')
-        # (yusong): a hack, can make it better
-        if args.tmodel == "transformer":
-            from clap_module.tokenizer import tokenize
-            all_texts = tokenize(all_texts)
-        else:
-            from training.data import tokenizer
-            all_texts = tokenizer(all_texts)
-        all_text_features = model(None, all_texts, device)
-        all_text_features = F.normalize(all_text_features, dim=-1).detach().cpu()
-
-        # compute similarity
-        logit_scale_a, logit_scale_t = model(None, None, device)
-        logit_scale_a = logit_scale_a.cpu()
-
-        logits_per_audio = (logit_scale_a * all_audio_features @ all_text_features.t()).detach().cpu()
-        logits_per_text = logits_per_audio.t().detach().cpu()
-
-        ground_truth = all_class_labels.view(-1, 1)
-        logit = logits_per_audio
-
-        ranking = torch.argsort(logit, descending=True)
-        preds = torch.where(ranking == ground_truth)[1]  # (yusong) this line is slow because it uses single thread
-        preds = preds.detach().cpu().numpy()
-        metrics[f"{args.datasetnames[0]}_mean_rank"] = preds.mean() + 1
-        metrics[f"{args.datasetnames[0]}_median_rank"] = np.floor(np.median(preds)) + 1
-        for k in [1, 5, 10]:
-            metrics[f"{args.datasetnames[0]}_R@{k}"] = np.mean(preds < k)
-        # map@10
-        metrics[f"{args.datasetnames[0]}_mAP@10"] = np.mean(np.where(preds < 10, 1 / (preds + 1), 0.0))
-
-        logging.info(
-            f"Eval Epoch: {start_epoch} "
-            + "\t".join([f"{k}: {round(v, 4):.4f}" for k, v in metrics.items()])
-        )
-
-        if args.wandb:
-            assert wandb is not None, "Please install wandb."
-            for name, val in metrics.items():
-                wandb.log({f"val/{name}": val, "epoch": start_epoch})
-
-
-if __name__ == '__main__':
-    # (yusong) repeated run might have different metric results.
-    # This is because we randomly select crop 10s for each audio.
-    args = parse_args()
-
-    if os.path.isdir(args.pretrained):
-        log_dir = os.path.dirname(args.pretrained)
-    else:
-        log_dir = os.path.dirname(os.path.dirname(args.pretrained))
-
-    args.log_level = logging.DEBUG if args.debug else logging.INFO
-    log_path = os.path.join(log_dir, 'out.log')
-    setup_logging(log_path, args.log_level)
-    params_file = os.path.join(log_dir, 'params.txt')
-
-    seed = 3407
-    random.seed(seed)
-    torch.manual_seed(seed)
-    torch.cuda.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)
-    np.random.seed(seed)
-
-    cudnn.benchmark = True
-    cudnn.deterministic = False
-    pretrained = 'openai'
-    amodel = find_params_value(params_file, 'amodel')
-    tmodel = find_params_value(params_file, 'tmodel')
-
-    if amodel is None or tmodel is None:
-        raise ValueError('model type not found in params file')
-
-    # set up dummy values for args
-    args.parallel_eval = False
-    args.rank = 0
-    args.local_rank = 0
-    args.world_size = 1
-    args.val_frequency = 1
-    args.epochs = 1
-    args.precision = 'fp32'
-    args.save_logs = True
-    args.wandb = args.report_to == 'wandb'
-    args.class_index_dict = None
-
-    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
-    args.device = device
-
-    if args.remotedata:
-        for dataset_name in args.datasetnames:
-            for split in dataset_split[dataset_name]:
-                if not os.path.exists(f"./json_files/{dataset_name}/{split}"):
-                    os.makedirs(f"./json_files/{dataset_name}/{split}")
-                os.system(
-                    f"aws s3 cp s3://s-laion-audio/webdataset_tar/{dataset_name}/{split}/sizes.json ./json_files/{dataset_name}/{split}/sizes.json"
-                )
-
-    if args.datasetinfos is None:
-        args.datasetinfos = ["train", "unbalanced_train", "balanced_train"]
-    if args.dataset_type == "webdataset":
-        args.train_data = get_tar_path_from_dataset_name(
-            args.datasetnames,
-            args.datasetinfos,
-            islocal=not args.remotedata,
-            proportion=args.dataset_proportion,
-            dataset_path=args.datasetpath,
-        )
-        args.val_data = get_tar_path_from_dataset_name(
-            args.datasetnames,
-            ["valid", "test", "eval"],
-            islocal=not args.remotedata,
-            proportion=1,
-            dataset_path=args.datasetpath,
-        )
-    model, model_cfg = create_model(
-        amodel,
-        tmodel,
-        pretrained,
-        precision='fp32',
-        device=device,
-        jit=False,
-        force_quick_gelu=False,
-        openai_model_cache_dir=os.path.expanduser(args.openai_model_cache_dir),
-        skip_params=False,
-        enable_fusion=args.enable_fusion,
-        fusion_type=args.fusion_type
-    )  # a hack to get model_cfg
-
-    data = get_data(args, model_cfg=model_cfg)  # (yusong): hack: no model_cfg needed to get data
-
-    writer = None  # if use tensorboard, initalize writer here
-
-    if args.wandb:
-        assert wandb is not None, "Please install wandb."
-
-        # # find the line with "wandb_notes" and get the value
-        # wandb_notes = find_params_value(params_file, 'wandb_notes')
-        # if wandb_notes is None:
-        #     print(f'wandb_notes not found in params file: {params_file}, set to timestamp.')
-        #     wandb_notes = f'experiment_{time.strftime("%Y%m%d-%H%M%S")}'
-        # wandb_notes = wandb_notes + '-eval-retrieval'
-        wandb_notes = args.wandb_notes
-
-        logging.debug("Starting wandb.")
-        args.train_sz = data["train"].dataloader.num_samples
-        if args.val_data is not None:
-            args.val_sz = data["val"].dataloader.num_samples
-        # you will have to configure this for your project!
-        if args.wandb_id is not None:
-            wandb.init(
-                project="clap",
-                id=args.wandb_id,
-                resume=True
-            )
-        else:
-            wandb.init(
-                project="clap",
-                notes=wandb_notes,
-                name=wandb_notes,
-                tags=[],
-                config=vars(args),
-            )
-        logging.debug("Finished loading wandb.")
-
-    if os.path.isdir(args.pretrained):
-        all_model_checkpoints = sorted(glob.glob(os.path.join(log_dir, 'checkpoints', '*.pt')), key=os.path.getmtime)
-    else:
-        all_model_checkpoints = [args.pretrained]
-    for model_path in all_model_checkpoints:
-        args.checkpoint_path = os.path.dirname(model_path)
-        model, model_cfg = create_model(
-            amodel,
-            tmodel,
-            pretrained,
-            precision='fp32',
-            device=device,
-            jit=False,
-            force_quick_gelu=False,
-            openai_model_cache_dir=os.path.expanduser(args.openai_model_cache_dir),
-            skip_params=False,
-            enable_fusion=args.enable_fusion,
-            fusion_type=args.fusion_type
-        )
-
-        # load model
-        checkpoint = torch.load(model_path, map_location=device)
-        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
-
-        evaluate_zeroshot(model, data, start_epoch, args, writer)

utils/laion_clap/training/data.py

@@ -1,94 +1,15 @@
-import ast
-import json
-import logging
-import math
-import os
-import random
-import h5py
-from dataclasses import dataclass
-import braceexpand
+from contextlib import suppress
+
 import numpy as np
-import pandas as pd
 import torch
 import torch.nn.functional as F
-import torchvision.datasets as datasets
 import torchvision.transforms
-import webdataset as wds
-from PIL import Image
-from torch.utils.data import Dataset, DataLoader, SubsetRandomSampler
-from torch.utils.data.distributed import DistributedSampler
-from functools import partial
-from pathlib import Path
-import wget
-import tempfile
-import copy
-from contextlib import suppress
-
-from utils.laion_clap.clap_module.utils import get_tar_path_from_dataset_name, dataset_split
-from utils.laion_clap.clap_module.utils import load_p, load_class_label
-from utils.laion_clap.clap_module import tokenize as clip_tokenizer
-from transformers import BertTokenizer
-from transformers import RobertaTokenizer
-from transformers import BartTokenizer
-
-try:
-    import horovod.torch as hvd
-except ImportError:
-    hvd = None
 
 try:
     import torchaudio
 except ImportError:
     torchaudio = None
 
-bert_tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
-roberta_tokenizer = RobertaTokenizer.from_pretrained("roberta-base")
-bart_tokenizer = BartTokenizer.from_pretrained("facebook/bart-base")
-
-def tokenizer(text, tmodel="roberta", max_length=77):
-    """tokenizer for different models
-    tmodel is default to roberta as it is the best model for our task
-    max_length is default to 77 from the OpenAI CLIP parameters
-    We assume text to be a single string, but it can also be a list of strings
-    """
-    if tmodel == "transformer":
-        return clip_tokenizer(text).squeeze(0)
-
-    elif tmodel == "bert":
-        result = bert_tokenizer(
-            text,
-            padding="max_length",
-            truncation=True,
-            max_length=max_length,
-            return_tensors="pt",
-        )
-        return {k: v.squeeze(0) for k, v in result.items()}
-
-    elif tmodel == "roberta":
-        result = roberta_tokenizer(
-            text,
-            padding="max_length",
-            truncation=True,
-            max_length=max_length,
-            return_tensors="pt",
-        )
-        return {k: v.squeeze(0) for k, v in result.items()}
-
-    elif tmodel == "bart":
-        result = bart_tokenizer(
-            text,
-            padding="max_length",
-            truncation=True,
-            max_length=max_length,
-            return_tensors="pt",
-        )
-        return {k: v.squeeze(0) for k, v in result.items()}
-
-
-# initizlied the audioset map
-_AUDIOSET_MAP_PATH = os.path.join(Path(__file__).parent, "audioset_textmap.npy")
-_AUDIOSET_MAP = np.load(_AUDIOSET_MAP_PATH, allow_pickle=True)
-
 
 def int16_to_float32(x):
     return (x / 32767.0).astype(np.float32)
@@ -99,267 +20,6 @@ def float32_to_int16(x):
     return (x * 32767.).astype(np.int16)
 
 
-def int16_to_float32_torch(x):
-    return (x / 32767.0).type(torch.float32)
-
-
-def float32_to_int16_torch(x):
-    x = torch.clamp(x, min=-1., max=1.)
-    return (x * 32767.).type(torch.int16)
-
-
-# For Toy Dataset
-class ToyDataset(Dataset):
-    def __init__(self, index_path, ipc, config, eval_mode=False):
-        """Toy Dataset for testing the audioset input with text labels
-        Parameters
-        ----------
-            index_path: str
-                the link to the h5 file of each audio
-            idc: str
-                the link to the npy file, the number of samples in each class
-            config: dict
-                the audio cfg file
-           eval_model (bool): to indicate if the dataset is a testing dataset
-        """
-        self.audio_cfg = config["audio_cfg"]
-        self.text_cfg = config["text_cfg"]
-        self.fp = h5py.File(index_path, "r")
-        self.ipc = np.load(ipc, allow_pickle=True)
-        self.total_size = len(self.fp["audio_name"])
-        self.classes_num = self.audio_cfg["class_num"]
-        self.eval_mode = eval_mode
-
-        if not eval_mode:
-            self.generate_queue()
-        else:
-            self.queue = []
-            for i in range(self.total_size):
-                target = self.fp["target"][i]
-                if np.sum(target) > 0:
-                    self.queue.append(i)
-            self.total_size = len(self.queue)
-        logging.info("total dataset size: %d" % (self.total_size))
-        logging.info("class num: %d" % (self.classes_num))
-
-    def time_shifting(self, x):
-        frame_num = len(x)
-        shift_len = random.randint(0, frame_num - 1)
-        new_sample = np.concatenate([x[shift_len:], x[:shift_len]], axis=0)
-        return new_sample
-
-    def generate_queue(self):
-        self.queue = []
-        while len(self.queue) < self.total_size:
-            class_set = [*range(self.classes_num)]
-            random.shuffle(class_set)
-            self.queue += [
-                self.ipc[d][random.randint(0, len(self.ipc[d]) - 1)] for d in class_set
-            ]
-        self.queue = self.queue[: self.total_size]
-
-        logging.info("queue regenerated:%s" % (self.queue[-5:]))
-
-    def crop_wav(self, x):
-        crop_size = self.audio_cfg["crop_size"]
-        crop_pos = random.randint(0, len(x) - crop_size - 1)
-        return x[crop_pos: crop_pos + crop_size]
-
-    def prompt_text(self, target):
-        events = _AUDIOSET_MAP[np.where(target > 0)]
-        event_text = "The sounds of " + ", ".join(events[:-1]) + " and " + events[-1]
-        text = tokenizer(event_text)[0]
-        return text
-
-    def __getitem__(self, index):
-        """Load waveform, text, and target of an audio clip
-
-        Parameters
-        ----------
-            index: int
-                the index number
-        Return
-        ------
-            output: dict {
-                "hdf5_path": str,
-                "index_in_hdf5": int,
-                "audio_name": str,
-                "waveform": list (audio_length,),
-                "target": list (class_num, ),
-                "text": torch.tensor (context_length,)
-            }
-                the output dictionary
-        """
-        s_index = self.queue[index]
-
-        audio_name = self.fp["audio_name"][s_index].decode()
-        # Hardcode here CHANGE
-        hdf5_path = (
-            self.fp["hdf5_path"][s_index]
-            .decode()
-            .replace(
-                "../workspace",
-                "/home/la/kechen/Research/ke_zsasp/workspace",
-            )
-        )
-        r_idx = self.fp["index_in_hdf5"][s_index]
-        target = self.fp["target"][s_index].astype(np.float32)
-        text = self.prompt_text(target)
-        with h5py.File(hdf5_path, "r") as f:
-            waveform = int16_to_float32(f["waveform"][r_idx])[
-                       : self.audio_cfg["clip_samples"]
-                       ]
-        assert (
-                len(waveform) == self.audio_cfg["clip_samples"]
-        ), "The sample length is not match"
-        # Time shift
-        # if (self.config.enable_time_shift) and (not self.eval_mode):
-        #     waveform = self.time_shifting(waveform)
-        # # Label Enhance
-        # if (self.config.crop_size is not None) and (not self.eval_mode):
-        #     waveform = self.crop_wav(waveform)
-        # # the label enhance rate is fixed 0.5
-        # if (self.config.enable_label_enhance) and (not self.eval_mode) and random.random() < 0.5:
-        #     kidx = np.where(target)[0]
-        #     for k in kidx:
-        #         for add_key in self.class_map[k][1]:
-        #             target[add_key] = 1.0
-        #         if len(self.class_map[k][2]) > 0:
-        #             add_key = random.choice(self.class_map[k][2])
-        #             target[add_key] = 1.0
-
-        # missing the text input
-        mel_spec = get_mel(torch.from_numpy(waveform), self.audio_cfg)[None, :, :]
-        mel_spec = torch.cat([mel_spec, mel_spec.clone(), mel_spec.clone(), mel_spec.clone()], dim=0).cpu().numpy()
-        longer = random.choice([True, False])
-        if longer == False:
-            mel_spec[1:, :, :] = 0.0
-        data_dict = {
-            "hdf5_path": hdf5_path,
-            "index_in_hdf5": r_idx,
-            "audio_name": audio_name,
-            "waveform": waveform,
-            "class_label": target,
-            "text": text,
-            "longer": longer,
-            "mel_fusion": mel_spec
-        }
-        return data_dict
-
-    def __len__(self):
-        return self.total_size
-
-@dataclass
-class DataInfo:
-    dataloader: DataLoader
-    sampler: DistributedSampler
-
-
-def get_dataset_size(shards, sizefilepath_=None, is_local=True):
-    if isinstance(shards, list):
-        size_list = []
-        for s in shards:
-            size_list.append(
-                get_dataset_size(s, sizefilepath_=sizefilepath_, is_local=is_local)[0]
-            )
-    else:
-        if not is_local:
-            for n in dataset_split.keys():
-                if n in shards.split("/"):
-                    break
-            for s in dataset_split[n]:
-                if s in shards.split("/"):
-                    break
-            sizefilepath_ = f"./json_files/{n}/{s}/sizes.json"
-        shards_list = list(braceexpand.braceexpand(shards))
-        dir_path = os.path.dirname(shards)
-        if sizefilepath_ is not None:
-            sizes = json.load(open(sizefilepath_, "r"))
-            total_size = sum(
-                [
-                    int(sizes[os.path.basename(shard.replace(".tar -", ".tar"))])
-                    for shard in shards_list
-                ]
-            )
-        else:
-            sizes_filename = os.path.join(dir_path, "sizes.json")
-            len_filename = os.path.join(dir_path, "__len__")
-            if os.path.exists(sizes_filename):
-                sizes = json.load(open(sizes_filename, "r"))
-                total_size = sum(
-                    [int(sizes[os.path.basename(shard)]) for shard in shards_list]
-                )
-            elif os.path.exists(len_filename):
-                # FIXME this used to be eval(open(...)) but that seemed rather unsafe
-                total_size = ast.literal_eval(open(len_filename, "r").read())
-            else:
-                raise Exception(
-                    f"Cannot find sizes file for dataset {shards}. Please specify the path to the file."
-                )
-                # total_size = None  # num samples undefined
-                # some common dataset sizes (at time of authors last download)
-                # cc3m-train: 2905954
-                # cc12m: 10968539
-                # LAION-400m: 407332084
-        num_shards = len(shards_list)
-    if isinstance(shards, list):
-        return sum(size_list), len(shards)
-    else:
-        return total_size, num_shards
-
-
-def count_samples(dataloader):
-    os.environ["WDS_EPOCH"] = "0"
-    n_elements, n_batches = 0, 0
-    for images, texts in dataloader:
-        n_batches += 1
-        n_elements += len(images)
-        assert len(images) == len(texts)
-    return n_elements, n_batches
-
-
-def log_and_continue(exn):
-    """Call in an exception handler to ignore any exception, isssue a warning, and continue."""
-    logging.warning(f"Handling webdataset error ({repr(exn)}). Ignoring.")
-    return True
-
-
-_SHARD_SHUFFLE_SIZE = 2000
-_SHARD_SHUFFLE_INITIAL = 500
-_SAMPLE_SHUFFLE_SIZE = 5000
-_SAMPLE_SHUFFLE_INITIAL = 1000
-
-
-def sample_prop(sizefile, inputs, proportion, is_local=True):
-    """
-    Sample a proportion of the data.
-    """
-    file_path_dict = {
-        os.path.split(inputs[i])[1]: os.path.split(inputs[i])[0]
-        for i in range(len(inputs))
-    }
-    sampled_filepath_dict = {}
-    sampled_size_dict = {}
-    if not is_local:
-        if os.path.exists("sizes.json"):
-            os.remove("sizes.json")
-        wget.download(sizefile, "sizes.json")
-        sizefile = "sizes.json"
-    with open(sizefile, "r", encoding="UTF-8") as f:
-        load_dict = json.load(f)
-    L = int(len(file_path_dict) * proportion)
-    subkeys = random.sample(file_path_dict.keys(), L)
-    for k in subkeys:
-        sampled_size_dict[k] = load_dict[k]
-        sampled_filepath_dict[k] = file_path_dict[k]
-    return (
-        sum(sampled_size_dict.values()),
-        L,
-        [os.path.join(v, k) for k, v in sampled_filepath_dict.items()],
-        sampled_size_dict,
-    )
-
-
 def get_mel(audio_data, audio_cfg):
     # mel shape: (n_mels, T)
     mel_tf = torchaudio.transforms.MelSpectrogram(
@@ -376,24 +36,8 @@ def get_mel(audio_data, audio_cfg):
         f_min=audio_cfg['fmin'],
         f_max=audio_cfg['fmax']
     ).to(audio_data.device)
-    
+
     mel = mel_tf(audio_data)
-    # Align to librosa:
-    # librosa_melspec = librosa.feature.melspectrogram(
-    #     waveform,
-    #     sr=audio_cfg['sample_rate'],
-    #     n_fft=audio_cfg['window_size'],
-    #     hop_length=audio_cfg['hop_size'],
-    #     win_length=audio_cfg['window_size'],
-    #     center=True,
-    #     pad_mode="reflect",
-    #     power=2.0,
-    #     n_mels=audio_cfg['mel_bins'],
-    #     norm=None,
-    #     htk=True,
-    #     f_min=audio_cfg['fmin'],
-    #     f_max=audio_cfg['fmax']
-    # )
     # we use log mel spectrogram as input
     mel = torchaudio.transforms.AmplitudeToDB(top_db=None)(mel)
     return mel.T  # (T, n_mels)
@@ -431,10 +75,6 @@ def get_audio_features(sample, audio_data, max_len, data_truncating, data_fillin
                     longer = torch.tensor([False])
                 else:
                     ranges = np.array_split(list(range(0, total_frames - chunk_frames + 1)), 3)
-                    # print('total_frames-chunk_frames:', total_frames-chunk_frames,
-                    #       'len(audio_data):', len(audio_data),
-                    #       'chunk_frames:', chunk_frames,
-                    #       'total_frames:', total_frames)
                     if len(ranges[1]) == 0:
                         # if the audio is too short, we just use the first chunk
                         ranges[1] = [0]
@@ -452,7 +92,6 @@ def get_audio_features(sample, audio_data, max_len, data_truncating, data_fillin
 
                     # shrink the mel
                     mel_shrink = torchvision.transforms.Resize(size=[chunk_frames, audio_cfg['mel_bins']])(mel[None])[0]
-                    # logging.info(f"mel_shrink.shape: {mel_shrink.shape}")
 
                     # stack
                     mel_fusion = torch.stack([mel_shrink, mel_chunk_front, mel_chunk_middle, mel_chunk_back], dim=0)
@@ -472,8 +111,6 @@ def get_audio_features(sample, audio_data, max_len, data_truncating, data_fillin
                 if data_filling == "repeatpad":
                     n_repeat = int(max_len / len(audio_data))
                     audio_data = audio_data.repeat(n_repeat)
-                    # audio_data = audio_data.unsqueeze(0).unsqueeze(0).unsqueeze(0)
-                    # audio_data = F.interpolate(audio_data,size=max_len,mode="bicubic")[0,0,0]
                     audio_data = F.pad(
                         audio_data,
                         (0, max_len - len(audio_data)),
@@ -504,392 +141,3 @@ def get_audio_features(sample, audio_data, max_len, data_truncating, data_fillin
     sample["waveform"] = audio_data
 
     return sample
-
-
-def select_text(json_dict_raw, text_augment_selection):
-    # For selecting augmented text from dataset
-    if text_augment_selection is None or text_augment_selection == "none":
-        texts = json_dict_raw["text"]
-    elif text_augment_selection == "all":
-        if "text_augment_all" in json_dict_raw.keys():
-            texts = json_dict_raw["text_augment_all"]
-        else:
-            texts = json_dict_raw["text"]
-    elif text_augment_selection == "augment_only":
-        if "text_augment_all" in json_dict_raw.keys():
-            if json_dict_raw["text_augment_t5"] is None:
-                texts = json_dict_raw["text"]
-            else:
-                texts = json_dict_raw["text_augment_t5"]
-        else:
-            texts = json_dict_raw["text"]
-    else:
-        raise NotImplementedError(
-            f"text_augment_selection {text_augment_selection} not implemented"
-        )
-    return texts
-
-
-def preprocess_single(
-        sample,
-        audio_ext,
-        text_ext,
-        max_len,
-        audio_cfg,
-        tmodel,
-        class_index_dict,
-        data_filling,
-        data_truncating,
-        text_augment_selection,
-):
-    """
-    Preprocess a single sample for wdsdataloader.
-    """
-    audio_data, orig_sr = sample[audio_ext]
-    audio_data = int16_to_float32_torch(float32_to_int16_torch(audio_data[0]))
-
-    sample = get_audio_features(sample, audio_data, max_len, data_truncating, data_filling, audio_cfg)
-    del sample[audio_ext]
-
-    json_dict_raw = sample[text_ext]
-
-    texts = select_text(json_dict_raw, text_augment_selection)
-    sample["full_text"] = texts
-
-    if isinstance(texts, list) and isinstance(texts[0], str) and len(texts) > 1:
-        texts = random.choice(texts)
-    sample["raw_text"] = texts
-    sample["text"] = tokenizer(texts, tmodel=tmodel)  # text shape: [num_token]
-    if class_index_dict is not None:
-        # https://stackoverflow.com/questions/48004243/how-to-share-large-read-only-dictionary-list-across-processes-in-multiprocessing
-        # https://stackoverflow.com/questions/45693949/storing-strings-in-a-multiprocessing-sharedctypes-array
-
-        # in case the re-written version is wrong, here is the old version:
-        # sample["class_label"] = np.zeros(len(class_index_dict.keys()))
-        # for x in json_dict_raw["tag"]:
-        #     sample["class_label"][class_index_dict[x]] = 1
-        # sample["class_label"] = torch.tensor(sample["class_label"]).float()
-
-        class_labels = np.zeros(len(class_index_dict))
-        class_labels[np.in1d(list(class_index_dict.keys()), json_dict_raw["tag"])] = 1
-        sample["class_label"] = torch.tensor(class_labels).float()
-
-    del sample[text_ext]
-    sample["audio_name"] = sample["__key__"].split("/")[-1] + "." + audio_ext
-    sample["text_name"] = sample["__key__"].split("/")[-1] + "." + text_ext
-    sample["audio_orig_sr"] = orig_sr
-    return sample
-
-
-def collate_fn_with_preprocess(batch,
-                               audio_ext,
-                               text_ext,
-                               max_len,
-                               audio_cfg,
-                               args,
-                               ):
-    """
-    Collate function for wdsdataloader.
-    batch: a list of dict, each dict is a sample
-    """
-
-    class_index_dict = copy.deepcopy(args.class_index_dict)  # To avoid deadlock in multiprocessing
-    data_filling = args.data_filling
-    data_truncating = args.data_truncating
-    text_augment_selection = args.text_augment_selection
-    tmodel = args.tmodel
-
-    # concatenate values in each dictionary. if it is a tensor, concatenate. if it is a list, extend.
-    data_preprocessed = []
-
-    for sample in batch:
-        data_preprocessed.append(
-            preprocess_single(sample, audio_ext, text_ext, max_len, audio_cfg, tmodel, class_index_dict, data_filling,
-                              data_truncating, text_augment_selection))
-
-    batch_dict = {}
-    for k in data_preprocessed[0].keys():
-        if isinstance(data_preprocessed[0][k], dict):  # dealwith bert tokenizer output
-            batch_dict[k] = {}
-            for kk in data_preprocessed[0][k].keys():
-                tmp = []
-                for i in range(len(data_preprocessed)):
-                    tmp.append(data_preprocessed[i][k][kk])
-                batch_dict[k][kk] = torch.vstack(tmp)
-        elif isinstance(data_preprocessed[0][k], torch.Tensor):
-            batch_dict[k] = torch.stack([sample[k] for sample in data_preprocessed])
-        elif isinstance(data_preprocessed[0][k], np.ndarray):
-            batch_dict[k] = torch.tensor(np.stack([sample[k] for sample in data_preprocessed]))
-        else:
-            batch_dict[k] = [sample[k] for sample in data_preprocessed]
-    del data_preprocessed
-    return batch_dict
-
-
-def get_wds_dataset(
-        args,
-        model_cfg,
-        is_train,
-        audio_ext="flac",
-        text_ext="json",
-        max_len=480000,
-        proportion=1.0,
-        sizefilepath_=None,
-        is_local=None,
-):
-    """
-    Get a dataset for wdsdataloader.
-    """
-    if is_local is None and (not args.remotedata is None):
-        is_local = not args.remotedata
-
-    input_shards = args.train_data if is_train else args.val_data
-    assert input_shards is not None
-
-    if not sizefilepath_ is None:
-        sizefilepath = sizefilepath_
-    else:
-        sizefilepath = os.path.join(os.path.dirname(input_shards[0]), "sizes.json")
-
-    if proportion != 1.0:
-        num_samples, num_shards, input_shards, _ = sample_prop(
-            sizefilepath, input_shards, proportion, is_local=is_local
-        )
-    else:
-        num_samples, num_shards = get_dataset_size(
-            input_shards, sizefilepath_=sizefilepath_, is_local=is_local
-        )
-
-    if not num_samples:
-        if is_train:
-            num_samples = args.train_num_samples
-            if not num_samples:
-                raise RuntimeError(
-                    "Currently, number of dataset samples must be specified for training dataset. "
-                    "Please specify via `--train-num-samples` if no dataset length info present."
-                )
-        else:
-            num_samples = (
-                    args.val_num_samples or 0
-            )  # eval will just exhaust the iterator if not specified
-
-    pipeline = [wds.SimpleShardList(input_shards)]
-    # at this point we have an iterator over all the shards
-    # TODO: (yusong): add a if statement of distributed. If not, we don't need to split_by_node
-    if is_train or args.parallel_eval:
-        pipeline.extend(
-            [
-                wds.detshuffle(
-                    bufsize=_SHARD_SHUFFLE_SIZE,
-                    initial=_SHARD_SHUFFLE_INITIAL,
-                    seed=args.seed,
-                ),
-                wds.split_by_node,
-                wds.split_by_worker,
-                # at this point, we have an iterator over the shards assigned to each worker at each node
-                wds.tarfile_to_samples(handler=log_and_continue),
-                wds.shuffle(
-                    bufsize=_SAMPLE_SHUFFLE_SIZE,
-                    initial=_SAMPLE_SHUFFLE_INITIAL,
-                    rng=random.Random(args.seed),
-                ),
-                # wds.repeatedly,  # FIXME determine if this is beneficial
-            ]
-        )
-    else:
-        pipeline.extend(
-            [
-                wds.split_by_worker,
-                # at this point, we have an iterator over the shards assigned to each worker
-                wds.tarfile_to_samples(handler=log_and_continue),
-            ]
-        )
-
-    pipeline.append(
-        wds.decode(wds.torch_audio),
-    )
-
-    pipeline.append(
-        wds.batched(
-            args.batch_size,
-            partial=not (is_train or args.parallel_eval),
-            collation_fn=partial(collate_fn_with_preprocess,
-                                 audio_ext=audio_ext,
-                                 text_ext=text_ext,
-                                 max_len=max_len,
-                                 audio_cfg=model_cfg['audio_cfg'],
-                                 args=args,
-                                 ),
-
-        )
-    )
-
-    dataset = wds.DataPipeline(*pipeline)
-    if is_train or args.parallel_eval:
-        # (yusong): Currently parallel evaluation will be not precise as we are repeat the last few samples.
-        # (yusong): See comments below.
-        # roll over and repeat a few samples to get same number of full batches on each node
-        global_batch_size = args.batch_size * args.world_size
-        num_batches = math.ceil(num_samples / global_batch_size)
-        num_workers = max(1, args.workers)
-        num_worker_batches = math.ceil(
-            num_batches / num_workers
-        )  # per dataloader worker
-        num_batches = num_worker_batches * num_workers
-        num_samples = num_batches * global_batch_size
-        dataset = dataset.with_epoch(
-            num_worker_batches
-        )  # each worker is iterating over this
-    else:
-        # last batches are partial, eval is done on single (master) node
-        num_batches = math.ceil(num_samples / args.batch_size)
-
-    kwargs = {}
-    if args.horovod:  # multi-node training on summit
-        kwargs["multiprocessing_context"] = "forkserver"
-
-    if is_train:
-        if args.prefetch_factor:
-            prefetch_factor = args.prefetch_factor
-        else:
-            prefetch_factor = max(2, args.batch_size // args.workers)
-    else:
-        prefetch_factor = 2
-
-    dataloader = wds.WebLoader(
-        dataset,
-        batch_size=None,
-        shuffle=False,
-        num_workers=args.workers,
-        pin_memory=True,
-        prefetch_factor=prefetch_factor,
-        **kwargs
-    )
-
-    # FIXME not clear which approach is better, with_epoch before vs after dataloader?
-    # hoping to resolve via https://github.com/webdataset/webdataset/issues/169
-    # if is_train:
-    #     # roll over and repeat a few samples to get same number of full batches on each node
-    #     global_batch_size = args.batch_size * args.world_size
-    #     num_batches = math.ceil(num_samples / global_batch_size)
-    #     num_workers = max(1, args.workers)
-    #     num_batches = math.ceil(num_batches / num_workers) * num_workers
-    #     num_samples = num_batches * global_batch_size
-    #     dataloader = dataloader.with_epoch(num_batches)
-    # else:
-    #     # last batches are partial, eval is done on single (master) node
-    #     num_batches = math.ceil(num_samples / args.batch_size)
-
-    # add meta-data to dataloader instance for convenience
-    dataloader.num_batches = num_batches
-    dataloader.num_samples = num_samples
-
-    return DataInfo(dataloader, None)
-
-
-def wds_batch_list2dict(
-        batch,
-        keys=[
-            "__url__",
-            "__key__",
-            "waveform",
-            "text",
-            "raw_text",
-            "audio_name",
-            "text_name",
-            "audio_orig_sr",
-        ],
-):
-    """
-    Return a dictionary of the batch, with keys as the names of the fields.
-    """
-    assert len(keys) == len(
-        batch
-    ), "batch must have same number of keys as keys argument"
-    return {keys[i]: batch[i] for i in range(len(batch))}
-
-
-
-def get_toy_dataset(args, model_cfg, is_train):
-    index_path = args.train_data if is_train else args.val_data
-    ipc_path = args.train_ipc if is_train else args.val_ipc
-    assert index_path and ipc_path
-    eval_mode = not is_train
-    dataset = ToyDataset(index_path, ipc_path, model_cfg, eval_mode=eval_mode)
-
-    num_samples = len(dataset)
-    sampler = (
-        DistributedSampler(dataset, shuffle=False)
-        if args.distributed and is_train
-        else None
-    )
-
-    dataloader = DataLoader(
-        dataset,
-        batch_size=args.batch_size,
-        shuffle=False,
-        num_workers=args.workers,
-        sampler=sampler,
-        drop_last=is_train,
-    )
-    dataloader.num_samples = num_samples
-    dataloader.num_batches = len(dataloader)
-
-    return DataInfo(dataloader, sampler)
-
-
-def get_dataset_fn(dataset_type):
-    if dataset_type == "webdataset":
-        return get_wds_dataset
-    elif dataset_type == "toy":
-        return get_toy_dataset
-    else:
-        raise ValueError(f"Unsupported dataset type: {dataset_type}")
-
-
-def get_data(args, model_cfg):
-    data = {}
-
-    args.class_index_dict = load_class_label(args.class_label_path)
-
-    if args.datasetinfos is None:
-        args.datasetinfos = ["train", "unbalanced_train", "balanced_train"]
-    if args.dataset_type == "webdataset":
-        args.train_data = get_tar_path_from_dataset_name(
-            args.datasetnames,
-            args.datasetinfos,
-            islocal=not args.remotedata,
-            proportion=args.dataset_proportion,
-            dataset_path=args.datasetpath,
-            full_dataset=args.full_train_dataset,
-        )
-
-        if args.full_train_dataset is None:
-            args.full_train_dataset = []
-        if args.exclude_eval_dataset is None:
-            args.exclude_eval_dataset = []
-        excluded_eval_datasets = args.full_train_dataset + args.exclude_eval_dataset
-
-        val_dataset_names = [n for n in args.datasetnames if n not in excluded_eval_datasets] \
-            if excluded_eval_datasets else args.datasetnames
-        args.val_dataset_names = val_dataset_names
-        args.val_data = get_tar_path_from_dataset_name(
-            val_dataset_names,
-            ["valid", "test", "eval"],
-            islocal=not args.remotedata,
-            proportion=1,
-            dataset_path=args.datasetpath,
-            full_dataset=None,
-        )
-
-    if args.train_data:
-        data["train"] = get_dataset_fn(args.dataset_type)(
-            args, model_cfg, is_train=True
-        )
-
-    if args.val_data:
-        data["val"] = get_dataset_fn(args.dataset_type)(
-            args, model_cfg, is_train=False
-        )
-
-    return data

utils/laion_clap/training/distributed.py

@@ -1,139 +0,0 @@
-import os
-
-import torch
-import socket
-
-try:
-    import horovod.torch as hvd
-except ImportError:
-    hvd = None
-
-
-def is_global_master(args):
-    return args.rank == 0
-
-
-def is_local_master(args):
-    return args.local_rank == 0
-
-
-def is_master(args, local=False):
-    return is_local_master(args) if local else is_global_master(args)
-
-
-def is_using_horovod():
-    # NOTE w/ horovod run, OMPI vars should be set, but w/ SLURM PMI vars will be set
-    # Differentiating between horovod and DDP use via SLURM may not be possible, so horovod arg still required...
-    ompi_vars = ["OMPI_COMM_WORLD_RANK", "OMPI_COMM_WORLD_SIZE"]
-    pmi_vars = ["PMI_RANK", "PMI_SIZE"]
-    if all([var in os.environ for var in ompi_vars]) or all([var in os.environ for var in pmi_vars]):
-        return True
-    else:
-        return False
-
-
-def is_using_distributed():
-    if 'WORLD_SIZE' in os.environ:
-        return int(os.environ['WORLD_SIZE']) > 1
-    if 'SLURM_NTASKS' in os.environ:
-        return int(os.environ['SLURM_NTASKS']) > 1
-    return False
-
-
-def world_info_from_env():
-    local_rank = 0
-    for v in ('SLURM_LOCALID', 'MPI_LOCALRANKID', 'OMPI_COMM_WORLD_LOCAL_RANK', 'LOCAL_RANK'):
-        if v in os.environ:
-            local_rank = int(os.environ[v])
-            break
-    global_rank = 0
-    for v in ('SLURM_PROCID', 'PMI_RANK', 'OMPI_COMM_WORLD_RANK', 'RANK'):
-        if v in os.environ:
-            global_rank = int(os.environ[v])
-            break
-    world_size = 1
-    for v in ('SLURM_NTASKS', 'PMI_SIZE', 'OMPI_COMM_WORLD_SIZE', 'WORLD_SIZE'):
-        if v in os.environ:
-            world_size = int(os.environ[v])
-            break
-
-    return local_rank, global_rank, world_size
-
-
-def init_distributed_device(args):
-    # Distributed training = training on more than one GPU.
-    # Works in both single and multi-node scenarios.
-    args.distributed = False
-    args.world_size = 1
-    args.rank = 0  # global rank
-    args.local_rank = 0
-    if args.horovod:
-        assert hvd is not None, "Horovod is not installed"
-        hvd.init()
-        world_size = int(os.environ['OMPI_COMM_WORLD_SIZE'])
-        world_rank = int(os.environ['OMPI_COMM_WORLD_RANK'])
-        local_rank = int(os.environ['OMPI_COMM_WORLD_LOCAL_RANK'])
-        args.local_rank = local_rank
-        args.rank = world_rank
-        args.world_size = world_size
-        # args.local_rank = int(hvd.local_rank())
-        # args.rank = hvd.rank()
-        # args.world_size = hvd.size()
-        args.distributed = True
-        os.environ['LOCAL_RANK'] = str(args.local_rank)
-        os.environ['RANK'] = str(args.rank)
-        os.environ['WORLD_SIZE'] = str(args.world_size)
-        print(f"Distributed training: local_rank={args.local_rank}, "
-              f"rank={args.rank}, world_size={args.world_size}, "
-              f"hostname={socket.gethostname()}, pid={os.getpid()}")
-    elif is_using_distributed():
-        if 'SLURM_PROCID' in os.environ:
-            # DDP via SLURM
-            args.local_rank, args.rank, args.world_size = world_info_from_env()
-            # SLURM var -> torch.distributed vars in case needed
-            os.environ['LOCAL_RANK'] = str(args.local_rank)
-            os.environ['RANK'] = str(args.rank)
-            os.environ['WORLD_SIZE'] = str(args.world_size)
-            torch.distributed.init_process_group(
-                backend=args.dist_backend,
-                init_method=args.dist_url,
-                world_size=args.world_size,
-                rank=args.rank,
-            )
-        elif 'OMPI_COMM_WORLD_SIZE' in os.environ: # using Summit cluster
-            world_size = int(os.environ['OMPI_COMM_WORLD_SIZE'])
-            world_rank = int(os.environ['OMPI_COMM_WORLD_RANK'])
-            local_rank = int(os.environ['OMPI_COMM_WORLD_LOCAL_RANK'])
-            args.local_rank = local_rank
-            args.rank = world_rank
-            args.world_size = world_size
-            torch.distributed.init_process_group(
-                backend=args.dist_backend,
-                init_method=args.dist_url,
-                world_size=args.world_size,
-                rank=args.rank,
-            )
-        else:
-            # DDP via torchrun, torch.distributed.launch
-            args.local_rank, _, _ = world_info_from_env()
-            torch.distributed.init_process_group(
-                backend=args.dist_backend,
-                init_method=args.dist_url)
-            args.world_size = torch.distributed.get_world_size()
-            args.rank = torch.distributed.get_rank()
-        args.distributed = True
-        print(f"Distributed training: local_rank={args.local_rank}, "
-              f"rank={args.rank}, world_size={args.world_size}, "
-              f"hostname={socket.gethostname()}, pid={os.getpid()}")
-
-    if torch.cuda.is_available():
-        if args.distributed and not args.no_set_device_rank:
-            device = 'cuda:%d' % args.local_rank
-        else:
-            device = 'cuda:0'
-        torch.cuda.set_device(device)
-    else:
-        device = 'cpu'
-    args.device = device
-    device = torch.device(device)
-    return device

utils/laion_clap/training/imagenet_zeroshot_data.py

@@ -1,254 +0,0 @@
-# NOTE: This script is currently not supported for CLAP.
-
-imagenet_classnames = ["tench", "goldfish", "great white shark", "tiger shark", "hammerhead shark", "electric ray",
-                        "stingray", "rooster", "hen", "ostrich", "brambling", "goldfinch", "house finch", "junco",
-                        "indigo bunting", "American robin", "bulbul", "jay", "magpie", "chickadee", "American dipper",
-                        "kite (bird of prey)", "bald eagle", "vulture", "great grey owl", "fire salamander",
-                        "smooth newt", "newt", "spotted salamander", "axolotl", "American bullfrog", "tree frog",
-                        "tailed frog", "loggerhead sea turtle", "leatherback sea turtle", "mud turtle", "terrapin",
-                        "box turtle", "banded gecko", "green iguana", "Carolina anole",
-                        "desert grassland whiptail lizard", "agama", "frilled-necked lizard", "alligator lizard",
-                        "Gila monster", "European green lizard", "chameleon", "Komodo dragon", "Nile crocodile",
-                        "American alligator", "triceratops", "worm snake", "ring-necked snake",
-                        "eastern hog-nosed snake", "smooth green snake", "kingsnake", "garter snake", "water snake",
-                        "vine snake", "night snake", "boa constrictor", "African rock python", "Indian cobra",
-                        "green mamba", "sea snake", "Saharan horned viper", "eastern diamondback rattlesnake",
-                        "sidewinder rattlesnake", "trilobite", "harvestman", "scorpion", "yellow garden spider",
-                        "barn spider", "European garden spider", "southern black widow", "tarantula", "wolf spider",
-                        "tick", "centipede", "black grouse", "ptarmigan", "ruffed grouse", "prairie grouse", "peafowl",
-                        "quail", "partridge", "african grey parrot", "macaw", "sulphur-crested cockatoo", "lorikeet",
-                        "coucal", "bee eater", "hornbill", "hummingbird", "jacamar", "toucan", "duck",
-                        "red-breasted merganser", "goose", "black swan", "tusker", "echidna", "platypus", "wallaby",
-                        "koala", "wombat", "jellyfish", "sea anemone", "brain coral", "flatworm", "nematode", "conch",
-                        "snail", "slug", "sea slug", "chiton", "chambered nautilus", "Dungeness crab", "rock crab",
-                        "fiddler crab", "red king crab", "American lobster", "spiny lobster", "crayfish", "hermit crab",
-                        "isopod", "white stork", "black stork", "spoonbill", "flamingo", "little blue heron",
-                        "great egret", "bittern bird", "crane bird", "limpkin", "common gallinule", "American coot",
-                        "bustard", "ruddy turnstone", "dunlin", "common redshank", "dowitcher", "oystercatcher",
-                        "pelican", "king penguin", "albatross", "grey whale", "killer whale", "dugong", "sea lion",
-                        "Chihuahua", "Japanese Chin", "Maltese", "Pekingese", "Shih Tzu", "King Charles Spaniel",
-                        "Papillon", "toy terrier", "Rhodesian Ridgeback", "Afghan Hound", "Basset Hound", "Beagle",
-                        "Bloodhound", "Bluetick Coonhound", "Black and Tan Coonhound", "Treeing Walker Coonhound",
-                        "English foxhound", "Redbone Coonhound", "borzoi", "Irish Wolfhound", "Italian Greyhound",
-                        "Whippet", "Ibizan Hound", "Norwegian Elkhound", "Otterhound", "Saluki", "Scottish Deerhound",
-                        "Weimaraner", "Staffordshire Bull Terrier", "American Staffordshire Terrier",
-                        "Bedlington Terrier", "Border Terrier", "Kerry Blue Terrier", "Irish Terrier",
-                        "Norfolk Terrier", "Norwich Terrier", "Yorkshire Terrier", "Wire Fox Terrier",
-                        "Lakeland Terrier", "Sealyham Terrier", "Airedale Terrier", "Cairn Terrier",
-                        "Australian Terrier", "Dandie Dinmont Terrier", "Boston Terrier", "Miniature Schnauzer",
-                        "Giant Schnauzer", "Standard Schnauzer", "Scottish Terrier", "Tibetan Terrier",
-                        "Australian Silky Terrier", "Soft-coated Wheaten Terrier", "West Highland White Terrier",
-                        "Lhasa Apso", "Flat-Coated Retriever", "Curly-coated Retriever", "Golden Retriever",
-                        "Labrador Retriever", "Chesapeake Bay Retriever", "German Shorthaired Pointer", "Vizsla",
-                        "English Setter", "Irish Setter", "Gordon Setter", "Brittany dog", "Clumber Spaniel",
-                        "English Springer Spaniel", "Welsh Springer Spaniel", "Cocker Spaniel", "Sussex Spaniel",
-                        "Irish Water Spaniel", "Kuvasz", "Schipperke", "Groenendael dog", "Malinois", "Briard",
-                        "Australian Kelpie", "Komondor", "Old English Sheepdog", "Shetland Sheepdog", "collie",
-                        "Border Collie", "Bouvier des Flandres dog", "Rottweiler", "German Shepherd Dog", "Dobermann",
-                        "Miniature Pinscher", "Greater Swiss Mountain Dog", "Bernese Mountain Dog",
-                        "Appenzeller Sennenhund", "Entlebucher Sennenhund", "Boxer", "Bullmastiff", "Tibetan Mastiff",
-                        "French Bulldog", "Great Dane", "St. Bernard", "husky", "Alaskan Malamute", "Siberian Husky",
-                        "Dalmatian", "Affenpinscher", "Basenji", "pug", "Leonberger", "Newfoundland dog",
-                        "Great Pyrenees dog", "Samoyed", "Pomeranian", "Chow Chow", "Keeshond", "brussels griffon",
-                        "Pembroke Welsh Corgi", "Cardigan Welsh Corgi", "Toy Poodle", "Miniature Poodle",
-                        "Standard Poodle", "Mexican hairless dog (xoloitzcuintli)", "grey wolf", "Alaskan tundra wolf",
-                        "red wolf or maned wolf", "coyote", "dingo", "dhole", "African wild dog", "hyena", "red fox",
-                        "kit fox", "Arctic fox", "grey fox", "tabby cat", "tiger cat", "Persian cat", "Siamese cat",
-                        "Egyptian Mau", "cougar", "lynx", "leopard", "snow leopard", "jaguar", "lion", "tiger",
-                        "cheetah", "brown bear", "American black bear", "polar bear", "sloth bear", "mongoose",
-                        "meerkat", "tiger beetle", "ladybug", "ground beetle", "longhorn beetle", "leaf beetle",
-                        "dung beetle", "rhinoceros beetle", "weevil", "fly", "bee", "ant", "grasshopper",
-                        "cricket insect", "stick insect", "cockroach", "praying mantis", "cicada", "leafhopper",
-                        "lacewing", "dragonfly", "damselfly", "red admiral butterfly", "ringlet butterfly",
-                        "monarch butterfly", "small white butterfly", "sulphur butterfly", "gossamer-winged butterfly",
-                        "starfish", "sea urchin", "sea cucumber", "cottontail rabbit", "hare", "Angora rabbit",
-                        "hamster", "porcupine", "fox squirrel", "marmot", "beaver", "guinea pig", "common sorrel horse",
-                        "zebra", "pig", "wild boar", "warthog", "hippopotamus", "ox", "water buffalo", "bison",
-                        "ram (adult male sheep)", "bighorn sheep", "Alpine ibex", "hartebeest", "impala (antelope)",
-                        "gazelle", "arabian camel", "llama", "weasel", "mink", "European polecat",
-                        "black-footed ferret", "otter", "skunk", "badger", "armadillo", "three-toed sloth", "orangutan",
-                        "gorilla", "chimpanzee", "gibbon", "siamang", "guenon", "patas monkey", "baboon", "macaque",
-                        "langur", "black-and-white colobus", "proboscis monkey", "marmoset", "white-headed capuchin",
-                        "howler monkey", "titi monkey", "Geoffroy's spider monkey", "common squirrel monkey",
-                        "ring-tailed lemur", "indri", "Asian elephant", "African bush elephant", "red panda",
-                        "giant panda", "snoek fish", "eel", "silver salmon", "rock beauty fish", "clownfish",
-                        "sturgeon", "gar fish", "lionfish", "pufferfish", "abacus", "abaya", "academic gown",
-                        "accordion", "acoustic guitar", "aircraft carrier", "airliner", "airship", "altar", "ambulance",
-                        "amphibious vehicle", "analog clock", "apiary", "apron", "trash can", "assault rifle",
-                        "backpack", "bakery", "balance beam", "balloon", "ballpoint pen", "Band-Aid", "banjo",
-                        "baluster / handrail", "barbell", "barber chair", "barbershop", "barn", "barometer", "barrel",
-                        "wheelbarrow", "baseball", "basketball", "bassinet", "bassoon", "swimming cap", "bath towel",
-                        "bathtub", "station wagon", "lighthouse", "beaker", "military hat (bearskin or shako)",
-                        "beer bottle", "beer glass", "bell tower", "baby bib", "tandem bicycle", "bikini",
-                        "ring binder", "binoculars", "birdhouse", "boathouse", "bobsleigh", "bolo tie", "poke bonnet",
-                        "bookcase", "bookstore", "bottle cap", "hunting bow", "bow tie", "brass memorial plaque", "bra",
-                        "breakwater", "breastplate", "broom", "bucket", "buckle", "bulletproof vest",
-                        "high-speed train", "butcher shop", "taxicab", "cauldron", "candle", "cannon", "canoe",
-                        "can opener", "cardigan", "car mirror", "carousel", "tool kit", "cardboard box / carton",
-                        "car wheel", "automated teller machine", "cassette", "cassette player", "castle", "catamaran",
-                        "CD player", "cello", "mobile phone", "chain", "chain-link fence", "chain mail", "chainsaw",
-                        "storage chest", "chiffonier", "bell or wind chime", "china cabinet", "Christmas stocking",
-                        "church", "movie theater", "cleaver", "cliff dwelling", "cloak", "clogs", "cocktail shaker",
-                        "coffee mug", "coffeemaker", "spiral or coil", "combination lock", "computer keyboard",
-                        "candy store", "container ship", "convertible", "corkscrew", "cornet", "cowboy boot",
-                        "cowboy hat", "cradle", "construction crane", "crash helmet", "crate", "infant bed",
-                        "Crock Pot", "croquet ball", "crutch", "cuirass", "dam", "desk", "desktop computer",
-                        "rotary dial telephone", "diaper", "digital clock", "digital watch", "dining table",
-                        "dishcloth", "dishwasher", "disc brake", "dock", "dog sled", "dome", "doormat", "drilling rig",
-                        "drum", "drumstick", "dumbbell", "Dutch oven", "electric fan", "electric guitar",
-                        "electric locomotive", "entertainment center", "envelope", "espresso machine", "face powder",
-                        "feather boa", "filing cabinet", "fireboat", "fire truck", "fire screen", "flagpole", "flute",
-                        "folding chair", "football helmet", "forklift", "fountain", "fountain pen", "four-poster bed",
-                        "freight car", "French horn", "frying pan", "fur coat", "garbage truck",
-                        "gas mask or respirator", "gas pump", "goblet", "go-kart", "golf ball", "golf cart", "gondola",
-                        "gong", "gown", "grand piano", "greenhouse", "radiator grille", "grocery store", "guillotine",
-                        "hair clip", "hair spray", "half-track", "hammer", "hamper", "hair dryer", "hand-held computer",
-                        "handkerchief", "hard disk drive", "harmonica", "harp", "combine harvester", "hatchet",
-                        "holster", "home theater", "honeycomb", "hook", "hoop skirt", "gymnastic horizontal bar",
-                        "horse-drawn vehicle", "hourglass", "iPod", "clothes iron", "carved pumpkin", "jeans", "jeep",
-                        "T-shirt", "jigsaw puzzle", "rickshaw", "joystick", "kimono", "knee pad", "knot", "lab coat",
-                        "ladle", "lampshade", "laptop computer", "lawn mower", "lens cap", "letter opener", "library",
-                        "lifeboat", "lighter", "limousine", "ocean liner", "lipstick", "slip-on shoe", "lotion",
-                        "music speaker", "loupe magnifying glass", "sawmill", "magnetic compass", "messenger bag",
-                        "mailbox", "tights", "one-piece bathing suit", "manhole cover", "maraca", "marimba", "mask",
-                        "matchstick", "maypole", "maze", "measuring cup", "medicine cabinet", "megalith", "microphone",
-                        "microwave oven", "military uniform", "milk can", "minibus", "miniskirt", "minivan", "missile",
-                        "mitten", "mixing bowl", "mobile home", "ford model t", "modem", "monastery", "monitor",
-                        "moped", "mortar and pestle", "graduation cap", "mosque", "mosquito net", "vespa",
-                        "mountain bike", "tent", "computer mouse", "mousetrap", "moving van", "muzzle", "metal nail",
-                        "neck brace", "necklace", "baby pacifier", "notebook computer", "obelisk", "oboe", "ocarina",
-                        "odometer", "oil filter", "pipe organ", "oscilloscope", "overskirt", "bullock cart",
-                        "oxygen mask", "product packet / packaging", "paddle", "paddle wheel", "padlock", "paintbrush",
-                        "pajamas", "palace", "pan flute", "paper towel", "parachute", "parallel bars", "park bench",
-                        "parking meter", "railroad car", "patio", "payphone", "pedestal", "pencil case",
-                        "pencil sharpener", "perfume", "Petri dish", "photocopier", "plectrum", "Pickelhaube",
-                        "picket fence", "pickup truck", "pier", "piggy bank", "pill bottle", "pillow", "ping-pong ball",
-                        "pinwheel", "pirate ship", "drink pitcher", "block plane", "planetarium", "plastic bag",
-                        "plate rack", "farm plow", "plunger", "Polaroid camera", "pole", "police van", "poncho",
-                        "pool table", "soda bottle", "plant pot", "potter's wheel", "power drill", "prayer rug",
-                        "printer", "prison", "missile", "projector", "hockey puck", "punching bag", "purse", "quill",
-                        "quilt", "race car", "racket", "radiator", "radio", "radio telescope", "rain barrel",
-                        "recreational vehicle", "fishing casting reel", "reflex camera", "refrigerator",
-                        "remote control", "restaurant", "revolver", "rifle", "rocking chair", "rotisserie", "eraser",
-                        "rugby ball", "ruler measuring stick", "sneaker", "safe", "safety pin", "salt shaker", "sandal",
-                        "sarong", "saxophone", "scabbard", "weighing scale", "school bus", "schooner", "scoreboard",
-                        "CRT monitor", "screw", "screwdriver", "seat belt", "sewing machine", "shield", "shoe store",
-                        "shoji screen / room divider", "shopping basket", "shopping cart", "shovel", "shower cap",
-                        "shower curtain", "ski", "balaclava ski mask", "sleeping bag", "slide rule", "sliding door",
-                        "slot machine", "snorkel", "snowmobile", "snowplow", "soap dispenser", "soccer ball", "sock",
-                        "solar thermal collector", "sombrero", "soup bowl", "keyboard space bar", "space heater",
-                        "space shuttle", "spatula", "motorboat", "spider web", "spindle", "sports car", "spotlight",
-                        "stage", "steam locomotive", "through arch bridge", "steel drum", "stethoscope", "scarf",
-                        "stone wall", "stopwatch", "stove", "strainer", "tram", "stretcher", "couch", "stupa",
-                        "submarine", "suit", "sundial", "sunglasses", "sunglasses", "sunscreen", "suspension bridge",
-                        "mop", "sweatshirt", "swim trunks / shorts", "swing", "electrical switch", "syringe",
-                        "table lamp", "tank", "tape player", "teapot", "teddy bear", "television", "tennis ball",
-                        "thatched roof", "front curtain", "thimble", "threshing machine", "throne", "tile roof",
-                        "toaster", "tobacco shop", "toilet seat", "torch", "totem pole", "tow truck", "toy store",
-                        "tractor", "semi-trailer truck", "tray", "trench coat", "tricycle", "trimaran", "tripod",
-                        "triumphal arch", "trolleybus", "trombone", "hot tub", "turnstile", "typewriter keyboard",
-                        "umbrella", "unicycle", "upright piano", "vacuum cleaner", "vase", "vaulted or arched ceiling",
-                        "velvet fabric", "vending machine", "vestment", "viaduct", "violin", "volleyball",
-                        "waffle iron", "wall clock", "wallet", "wardrobe", "military aircraft", "sink",
-                        "washing machine", "water bottle", "water jug", "water tower", "whiskey jug", "whistle",
-                        "hair wig", "window screen", "window shade", "Windsor tie", "wine bottle", "airplane wing",
-                        "wok", "wooden spoon", "wool", "split-rail fence", "shipwreck", "sailboat", "yurt", "website",
-                        "comic book", "crossword", "traffic or street sign", "traffic light", "dust jacket", "menu",
-                        "plate", "guacamole", "consomme", "hot pot", "trifle", "ice cream", "popsicle", "baguette",
-                        "bagel", "pretzel", "cheeseburger", "hot dog", "mashed potatoes", "cabbage", "broccoli",
-                        "cauliflower", "zucchini", "spaghetti squash", "acorn squash", "butternut squash", "cucumber",
-                        "artichoke", "bell pepper", "cardoon", "mushroom", "Granny Smith apple", "strawberry", "orange",
-                        "lemon", "fig", "pineapple", "banana", "jackfruit", "cherimoya (custard apple)", "pomegranate",
-                        "hay", "carbonara", "chocolate syrup", "dough", "meatloaf", "pizza", "pot pie", "burrito",
-                        "red wine", "espresso", "tea cup", "eggnog", "mountain", "bubble", "cliff", "coral reef",
-                        "geyser", "lakeshore", "promontory", "sandbar", "beach", "valley", "volcano", "baseball player",
-                        "bridegroom", "scuba diver", "rapeseed", "daisy", "yellow lady's slipper", "corn", "acorn",
-                        "rose hip", "horse chestnut seed", "coral fungus", "agaric", "gyromitra", "stinkhorn mushroom",
-                        "earth star fungus", "hen of the woods mushroom", "bolete", "corn cob", "toilet paper"]
-
-
-
-
-
-openai_imagenet_template = [
-    lambda c: f'a bad photo of a {c}.',
-    lambda c: f'a photo of many {c}.',
-    lambda c: f'a sculpture of a {c}.',
-    lambda c: f'a photo of the hard to see {c}.',
-    lambda c: f'a low resolution photo of the {c}.',
-    lambda c: f'a rendering of a {c}.',
-    lambda c: f'graffiti of a {c}.',
-    lambda c: f'a bad photo of the {c}.',
-    lambda c: f'a cropped photo of the {c}.',
-    lambda c: f'a tattoo of a {c}.',
-    lambda c: f'the embroidered {c}.',
-    lambda c: f'a photo of a hard to see {c}.',
-    lambda c: f'a bright photo of a {c}.',
-    lambda c: f'a photo of a clean {c}.',
-    lambda c: f'a photo of a dirty {c}.',
-    lambda c: f'a dark photo of the {c}.',
-    lambda c: f'a drawing of a {c}.',
-    lambda c: f'a photo of my {c}.',
-    lambda c: f'the plastic {c}.',
-    lambda c: f'a photo of the cool {c}.',
-    lambda c: f'a close-up photo of a {c}.',
-    lambda c: f'a black and white photo of the {c}.',
-    lambda c: f'a painting of the {c}.',
-    lambda c: f'a painting of a {c}.',
-    lambda c: f'a pixelated photo of the {c}.',
-    lambda c: f'a sculpture of the {c}.',
-    lambda c: f'a bright photo of the {c}.',
-    lambda c: f'a cropped photo of a {c}.',
-    lambda c: f'a plastic {c}.',
-    lambda c: f'a photo of the dirty {c}.',
-    lambda c: f'a jpeg corrupted photo of a {c}.',
-    lambda c: f'a blurry photo of the {c}.',
-    lambda c: f'a photo of the {c}.',
-    lambda c: f'a good photo of the {c}.',
-    lambda c: f'a rendering of the {c}.',
-    lambda c: f'a {c} in a video game.',
-    lambda c: f'a photo of one {c}.',
-    lambda c: f'a doodle of a {c}.',
-    lambda c: f'a close-up photo of the {c}.',
-    lambda c: f'a photo of a {c}.',
-    lambda c: f'the origami {c}.',
-    lambda c: f'the {c} in a video game.',
-    lambda c: f'a sketch of a {c}.',
-    lambda c: f'a doodle of the {c}.',
-    lambda c: f'a origami {c}.',
-    lambda c: f'a low resolution photo of a {c}.',
-    lambda c: f'the toy {c}.',
-    lambda c: f'a rendition of the {c}.',
-    lambda c: f'a photo of the clean {c}.',
-    lambda c: f'a photo of a large {c}.',
-    lambda c: f'a rendition of a {c}.',
-    lambda c: f'a photo of a nice {c}.',
-    lambda c: f'a photo of a weird {c}.',
-    lambda c: f'a blurry photo of a {c}.',
-    lambda c: f'a cartoon {c}.',
-    lambda c: f'art of a {c}.',
-    lambda c: f'a sketch of the {c}.',
-    lambda c: f'a embroidered {c}.',
-    lambda c: f'a pixelated photo of a {c}.',
-    lambda c: f'itap of the {c}.',
-    lambda c: f'a jpeg corrupted photo of the {c}.',
-    lambda c: f'a good photo of a {c}.',
-    lambda c: f'a plushie {c}.',
-    lambda c: f'a photo of the nice {c}.',
-    lambda c: f'a photo of the small {c}.',
-    lambda c: f'a photo of the weird {c}.',
-    lambda c: f'the cartoon {c}.',
-    lambda c: f'art of the {c}.',
-    lambda c: f'a drawing of the {c}.',
-    lambda c: f'a photo of the large {c}.',
-    lambda c: f'a black and white photo of a {c}.',
-    lambda c: f'the plushie {c}.',
-    lambda c: f'a dark photo of a {c}.',
-    lambda c: f'itap of a {c}.',
-    lambda c: f'graffiti of the {c}.',
-    lambda c: f'a toy {c}.',
-    lambda c: f'itap of my {c}.',
-    lambda c: f'a photo of a cool {c}.',
-    lambda c: f'a photo of a small {c}.',
-    lambda c: f'a tattoo of the {c}.',
-]

utils/laion_clap/training/infer_demo.py

@@ -1,92 +0,0 @@
-import torch
-import librosa
-from clap_module import create_model
-from training.data import get_audio_features
-from training.data import int16_to_float32, float32_to_int16
-from transformers import RobertaTokenizer
-
-tokenize = RobertaTokenizer.from_pretrained('roberta-base')
-def tokenizer(text):
-    result = tokenize(
-        text,
-        padding="max_length",
-        truncation=True,
-        max_length=77,
-        return_tensors="pt",
-    )
-    return {k: v.squeeze(0) for k, v in result.items()}
-
-def infer_text():
-    device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
-    precision = 'fp32'
-    amodel = 'HTSAT-tiny' # or 'PANN-14'
-    tmodel = 'roberta' # the best text encoder in our training
-    enable_fusion = True # False if you do not want to use the fusion model
-    fusion_type = 'aff_2d'
-    pretrained = "/home/la/kechen/Research/KE_CLAP/ckpt/fusion_best.pt" # the checkpoint name, the unfusion model can also be loaded.
-
-    model, model_cfg = create_model(
-        amodel,
-        tmodel,
-        pretrained,
-        precision=precision,
-        device=device,
-        enable_fusion=enable_fusion,
-        fusion_type=fusion_type
-    )
-    # load the text, can be a list (i.e. batch size)
-    text_data = ["I love the contrastive learning", "I love the pretrain model"] 
-    # tokenize for roberta, if you want to tokenize for another text encoder, please refer to data.py#L43-90 
-    text_data = tokenizer(text_data)
-    model.eval()
-    text_embed = model.get_text_embedding(text_data)
-    text_embed = text_embed.detach().cpu().numpy()
-    print(text_embed)
-    print(text_embed.shape)
-
-def infer_audio():
-    
-    device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
-    precision = 'fp32'
-    amodel = 'HTSAT-tiny' # or 'PANN-14'
-    tmodel = 'roberta' # the best text encoder in our training
-    enable_fusion = True # False if you do not want to use the fusion model
-    fusion_type = 'aff_2d'
-    pretrained = "/home/la/kechen/Research/KE_CLAP/ckpt/fusion_best.pt" # the checkpoint name, the unfusion model can also be loaded.
-
-    model, model_cfg = create_model(
-        amodel,
-        tmodel,
-        pretrained,
-        precision=precision,
-        device=device,
-        enable_fusion=enable_fusion,
-        fusion_type=fusion_type
-    )
-
-    # load the waveform of the shape (T,), should resample to 48000
-    audio_waveform, sr = librosa.load('/home/la/kechen/Research/KE_CLAP/ckpt/test_clap_short.wav', sr=48000) 
-    # quantize
-    audio_waveform = int16_to_float32(float32_to_int16(audio_waveform))
-    audio_waveform = torch.from_numpy(audio_waveform).float()
-    audio_dict = {}
-
-    # the 'fusion' truncate mode can be changed to 'rand_trunc' if run in unfusion mode
-    audio_dict = get_audio_features(
-        audio_dict, audio_waveform, 480000, 
-        data_truncating='fusion', 
-        data_filling='repeatpad',
-        audio_cfg=model_cfg['audio_cfg']
-    )
-    model.eval()
-    # can send a list to the model, to process many audio tracks in one time (i.e. batch size)
-    audio_embed = model.get_audio_embedding([audio_dict])
-    audio_embed = audio_embed.detach().cpu().numpy()
-    print(audio_embed)
-    print(audio_embed.shape)
-    
-
-
-if __name__ == "__main__":
-    infer_text()
-    # infer_audio()

utils/laion_clap/training/logger.py

@@ -1,26 +0,0 @@
-import logging
-
-
-def setup_logging(log_file, level, include_host=False):
-    if include_host:
-        import socket
-        hostname = socket.gethostname()
-        formatter = logging.Formatter(
-            f'%(asctime)s |  {hostname} | %(levelname)s | %(message)s', datefmt='%Y-%m-%d,%H:%M:%S')
-    else:
-        formatter = logging.Formatter('%(asctime)s | %(levelname)s | %(message)s', datefmt='%Y-%m-%d,%H:%M:%S')
-
-    logging.root.setLevel(level)
-    loggers = [logging.getLogger(name) for name in logging.root.manager.loggerDict]
-    for logger in loggers:
-        logger.setLevel(level)
-
-    stream_handler = logging.StreamHandler()
-    stream_handler.setFormatter(formatter)
-    logging.root.addHandler(stream_handler)
-
-    if log_file:
-        file_handler = logging.FileHandler(filename=log_file)
-        file_handler.setFormatter(formatter)
-        logging.root.addHandler(file_handler)
-

utils/laion_clap/training/lp_main.py

@@ -1,643 +0,0 @@
-import logging
-import os
-import random
-from datetime import datetime
-import copy
-import numpy as np
-import torch
-import torch.backends.cudnn as cudnn
-from torch.cuda.amp import GradScaler
-import time
-
-try:
-    import wandb
-except ImportError:
-    wandb = None
-
-try:
-    import torch.utils.tensorboard as tensorboard
-except ImportError:
-    tensorboard = None
-
-try:
-    import horovod.torch as hvd
-except ImportError:
-    hvd = None
-
-from clap_module import create_model_and_transforms, trace_model, create_model
-from training.data import get_data
-from training.params import parse_args
-from training.distributed import is_master, init_distributed_device, world_info_from_env
-from training.logger import setup_logging
-from training.scheduler import cosine_lr
-from training.lp_train import train_one_epoch, evaluate
-from clap_module.utils import get_tar_path_from_dataset_name, dataset_split, get_optimizer
-from clap_module.utils import load_p, load_class_label
-from clap_module.linear_probe import LinearProbe
-
-
-def maintain_ckpts(args, startidx, all_idx_len):
-    for i in reversed(range(startidx, all_idx_len)):
-        if os.path.exists(os.path.join(args.checkpoint_path, f"epoch_top_{i}.pt")):
-            os.rename(
-                os.path.join(args.checkpoint_path, f"epoch_top_{i}.pt"),
-                os.path.join(args.checkpoint_path, f"epoch_top_{i+1}.pt"),
-            )
-    if os.path.exists(
-        os.path.join(args.checkpoint_path, f"epoch_top_{all_idx_len}.pt")
-    ):
-        os.remove(os.path.join(args.checkpoint_path, f"epoch_top_{all_idx_len}.pt"))
-    return
-
-
-def update_top_k_performance(
-    new_metrics_inputs, current_top_k_ckpt_metrics, args, ckpt, bignumbetter=True
-):
-    """
-    Record the top-k performance of the current epoch.
-    current_top_k_metrics is a dictionary of the form: {1: top_1_ckpt_measure, 2: top_2_ckpt_measure, ...}
-    """
-    if isinstance(new_metrics_inputs, (list, tuple)):
-        new_metrics_inputs = np.mean(new_metrics_inputs)
-        return update_top_k_performance(
-            new_metrics_inputs,
-            current_top_k_ckpt_metrics,
-            args=args,
-            ckpt=ckpt,
-            bignumbetter=bignumbetter,
-        )
-    elif isinstance(new_metrics_inputs, dict):
-        new_metrics_inputs = np.mean(list(new_metrics_inputs.values()))
-        return update_top_k_performance(
-            new_metrics_inputs,
-            current_top_k_ckpt_metrics,
-            args=args,
-            ckpt=ckpt,
-            bignumbetter=bignumbetter,
-        )
-    elif isinstance(new_metrics_inputs, (float, int)):
-        update_flag = {k: False for k in current_top_k_ckpt_metrics.keys()}
-        sorted_keys = sorted(current_top_k_ckpt_metrics.keys())
-        sorted_values = sorted(
-            current_top_k_ckpt_metrics.values(), reverse=bignumbetter
-        )
-        sorted_values_ = copy.deepcopy(sorted_values)
-        sorted_values.append(new_metrics_inputs)
-        sorted_values = sorted(sorted_values, reverse=bignumbetter)
-        sorted_values = sorted_values[:-1]
-
-        if sorted_values == sorted_values_:
-            return current_top_k_ckpt_metrics, new_metrics_inputs
-        else:
-            for i in range(len(sorted_keys)):
-                if current_top_k_ckpt_metrics[sorted_keys[i]] != sorted_values[i]:
-                    current_top_k_ckpt_metrics[sorted_keys[i]] = sorted_values[i]
-                    update_flag[sorted_keys[i]] = True
-            for i in range(len(update_flag)):
-                if update_flag[i]:
-                    maintain_ckpts(args, i, len(sorted_keys))
-                    torch.save(
-                        ckpt,
-                        os.path.join(args.checkpoint_path, f"epoch_top_{i}.pt"),
-                    )
-                    break
-            return current_top_k_ckpt_metrics, new_metrics_inputs
-
-
-# def updateifNone(a, b):
-#     a = b if None else a
-#     return a
-
-
-def is_pretrained_params(n):
-    return (
-        n.startswith("clap_model.transformer")
-        or n in ["clap_model.positional_embedding", "clap_model.text_projection"]
-        or n.startswith("clap_model.token_embedding")
-        or n.startswith("clap_model.ln_final")
-        or n.startswith("clap_model.logit_scale_t")
-    )
-
-
-def random_seed(seed=42, rank=0):
-    torch.manual_seed(seed + rank)
-    np.random.seed(seed + rank)
-    random.seed(seed + rank)
-
-def config_lp_optimizer(model, data, args):
-    # set wd-related params to 0 if use adam optimizer
-    if args.optimizer == "adam":
-        args.wd = 0
-        args.wd_pretrained = 0
-        args.wd_new = 0
-
-    in_clap = (
-        lambda n, p: n.startswith("clap_model")
-    )
-
-    named_parameters = list(model.named_parameters())
-
-    optimizer = {}
-    scheduler = {}
-
-    # freeze text encoder
-    text_freeze_parameters = [
-        p
-        for n, p in named_parameters
-        if n.startswith("clap_model.transformer")
-        or n in ["clap_model.positional_embedding", "clap_model.text_projection"]
-        or n.startswith("clap_model.token_embedding")
-        or n.startswith("clap_model.ln_final")
-    ]
-
-    if args.freeze_text:
-        logging.info("Freeze Text!!!!")
-        for k in text_freeze_parameters:
-            k.requires_grad = False
-
-    if not args.lp_freeze:
-        exclude = (
-            lambda n, p: p.ndim < 2
-            or "bn" in n
-            or "ln" in n
-            or "bias" in n
-            or "logit_scale" in n
-        )
-        include = lambda n, p: not exclude(n, p)
-
-        # (yusong): we do not split the learning rate anymore
-        # p for n, p in named_parameters if in_clap(n,p) and exclude(n, p) and p.requires_grad
-        gain_or_bias_params = [
-            p for n, p in named_parameters if exclude(n, p) and p.requires_grad
-        ]
-        # rest_params = [p for n, p in named_parameters if in_clap(n,p) and include(n, p) and p.requires_grad]
-        rest_params = [p for n, p in named_parameters if include(n, p) and p.requires_grad]
-
-        if args.train_data is None:
-            optimizer = None
-            scheduler = None
-        else:
-            total_steps = data["train"].dataloader.num_batches * args.epochs
-
-            if args.split_opt:
-                for x in ["lr", "beta1", "beta2", "eps", "wd"]:
-                    for y in ["_new", "_pretrained"]:
-                        if getattr(args, x + y) is None:
-                            setattr(args, x + y, getattr(args, x))
-
-                gain_or_bias_pretrained_params = [
-                    p
-                    for n, p in named_parameters
-                    if (exclude(n, p) and p.requires_grad) and is_pretrained_params(n)
-                ]
-                rest_pretrained_params = [
-                    p
-                    for n, p in named_parameters
-                    if (include(n, p) and p.requires_grad) and is_pretrained_params(n)
-                ]
-                gain_or_bias_new_params = [
-                    p
-                    for n, p in named_parameters
-                    if (exclude(n, p) and p.requires_grad) and (not is_pretrained_params(n))
-                ]
-                rest_new_params = [
-                    p
-                    for n, p in named_parameters
-                    if (include(n, p) and p.requires_grad) and (not is_pretrained_params(n))
-                ]
-
-                pretrained_params_optimizer = get_optimizer(
-                    [
-                        {"params": gain_or_bias_pretrained_params, "weight_decay": 0.0},
-                        {
-                            "params": rest_pretrained_params,
-                            "weight_decay": args.wd_pretrained,
-                        },
-                    ],
-                    lr=args.lr_pretrained,
-                    betas=(args.beta1_pretrained, args.beta2_pretrained),
-                    eps=args.eps_pretrained,
-                    momentum=args.momentum_pretrained,
-                    optimizer_name=args.optimizer,
-                )
-                pretrained_params_scheduler = cosine_lr(
-                    pretrained_params_optimizer,
-                    args.lr_pretrained,
-                    args.warmup,
-                    total_steps,
-                )
-
-                new_params_optimizer = get_optimizer(
-                    [
-                        {"params": gain_or_bias_new_params, "weight_decay": 0.0},
-                        {"params": rest_new_params, "weight_decay": args.wd_new},
-                    ],
-                    lr=args.lr_new,
-                    betas=(args.beta1_new, args.beta2_new),
-                    eps=args.eps_new,
-                    momentum=args.momentum_new,
-                    optimizer_name=args.optimizer,
-                )
-                new_params_scheduler = cosine_lr(
-                    new_params_optimizer, args.lr_new, args.warmup, total_steps
-                )
-
-                optimizer["text"] = pretrained_params_optimizer
-                optimizer["audio"] = new_params_optimizer
-                scheduler["text"] = pretrained_params_scheduler
-                scheduler["audio"] = new_params_scheduler
-
-                if args.horovod:
-                    pretrained_params_optimizer = hvd.DistributedOptimizer(
-                        pretrained_params_optimizer,
-                        named_parameters=model.named_parameters(),
-                    )
-                    new_params_optimizer = hvd.DistributedOptimizer(
-                        new_params_optimizer, named_parameters=model.named_parameters()
-                    )
-                    hvd.broadcast_parameters(model.state_dict(), root_rank=0)
-                    hvd.broadcast_optimizer_state(pretrained_params_optimizer, root_rank=0)
-                    hvd.broadcast_optimizer_state(new_params_optimizer, root_rank=0)
-            else:
-
-                optimizer["clap"] = get_optimizer(
-                    [
-                        {"params": gain_or_bias_params, "weight_decay": 0.0},
-                        {"params": rest_params, "weight_decay": args.wd},
-                    ],
-                    lr=args.lr,
-                    betas=(args.beta1, args.beta2),
-                    eps=args.eps,
-                    momentum=args.momentum,
-                    optimizer_name=args.optimizer,
-            )
-                scheduler["clap"] = cosine_lr(optimizer["clap"], args.lr, args.warmup, total_steps)
-
-                if args.horovod:
-                    optimizer["clap"] = hvd.DistributedOptimizer(
-                        optimizer["clap"], named_parameters=model.named_parameters()
-                    )
-                    hvd.broadcast_parameters(model.state_dict(), root_rank=0)
-                    hvd.broadcast_optimizer_state(optimizer["clap"], root_rank=0)
-
-    # linear probe optimizer
-    else:
-        lp_params = [p for n, p in named_parameters if (not in_clap(n, p)) and p.requires_grad]
-        lp_optim = get_optimizer(lp_params, lr=args.lp_lr, betas=(args.beta1, args.beta2), eps=args.eps, momentum=0.9,
-                                 optimizer_name=args.optimizer)
-        optimizer["lp"] = lp_optim
-
-    return optimizer, scheduler, text_freeze_parameters
-
-
-def main():
-    args = parse_args()
-
-    time.sleep(args.sleep)
-
-    # sanitize model name for filesystem / uri use, easier if we don't use / in name as a rule?
-    args.amodel = args.amodel.replace("/", "-")
-    # download sizes.json file
-
-    # (yusong): the below two lines are for debug
-    # print("setting up faulthandler")
-    # faulthandler.register(10)
-
-    random.seed(args.seed)
-    torch.manual_seed(args.seed)
-    torch.cuda.manual_seed(args.seed)
-    torch.cuda.manual_seed_all(args.seed)
-    np.random.seed(args.seed)
-    args.class_index_dict = load_class_label(args.class_label_path)
-
-    # get the name of the experiments
-    if args.name is None:
-        args.name = "-".join(
-            [
-                datetime.now().strftime("%Y_%m_%d-%H_%M_%S"),
-                f"linear_probe"
-                f"model_{args.amodel}",
-                f"lr_{args.lr}",
-                f"b_{args.batch_size}",
-                f"j_{args.workers}",
-                f"p_{args.precision}",
-            ]
-        )
-
-    # discover initial world args early so we can log properly
-    args.distributed = False
-    args.local_rank, args.rank, args.world_size = world_info_from_env()
-
-    if args.remotedata and is_master(args):
-        for dataset_name in args.datasetnames:
-            for split in dataset_split[dataset_name]:
-                if not os.path.exists(f"./json_files/{dataset_name}/{split}"):
-                    os.makedirs(f"./json_files/{dataset_name}/{split}")
-                os.system(
-                    f"aws s3 cp s3://s-laion-audio/webdataset_tar/{dataset_name}/{split}/sizes.json ./json_files/{dataset_name}/{split}/sizes.json"
-                )
-
-    args.log_path = None
-    if is_master(args, local=args.log_local):
-        log_base_path = os.path.join(args.logs, args.name)
-        os.makedirs(log_base_path, exist_ok=True)
-        log_filename = f"out-{args.rank}" if args.log_local else "out.log"
-        args.log_path = os.path.join(log_base_path, log_filename)
-
-        # avoid log dir in same name:
-        postfix = 0
-        while os.path.exists(args.log_path):
-            postfix += 1
-            log_base_path_new = log_base_path+'-'+str(postfix)
-            os.makedirs(log_base_path_new, exist_ok=True)
-            log_filename = f"out-{args.rank}" if args.log_local else "out.log"
-            args.log_path = os.path.join(log_base_path_new, log_filename)
-            # print(
-            #     "Error. Experiment already exists. Use --name {} to specify a new experiment."
-            # )
-            # return -1
-
-    # Set logger
-    args.log_level = logging.DEBUG if args.debug else logging.INFO
-    setup_logging(args.log_path, args.log_level)
-
-    # fully initialize distributed device environment
-    device = init_distributed_device(args)
-
-    args.wandb = "wandb" in args.report_to or "all" in args.report_to
-    args.tensorboard = "tensorboard" in args.report_to or "all" in args.report_to
-    if is_master(args):
-        args.tensorboard_path = (
-            os.path.join(args.logs, args.name, "tensorboard")
-            if args.tensorboard
-            else ""
-        )
-        args.checkpoint_path = os.path.join(args.logs, args.name, "checkpoints")
-        for dirname in [args.tensorboard_path, args.checkpoint_path]:
-            if dirname:
-                os.makedirs(dirname, exist_ok=True)
-    else:
-        args.tensorboard_path = ""
-        args.checkpoint_path = ""
-
-    if args.copy_codebase:
-        copy_codebase(args)
-
-    assert args.precision in ["amp", "fp16", "fp32"]
-    if args.precision == "fp16":
-        logging.warning(
-            "It is recommended to use AMP mixed-precision instead of FP16. "
-            "FP16 support needs further verification and tuning, especially for train."
-        )
-
-    if args.horovod:
-        logging.info(
-            f"Running in horovod mode with multiple processes / nodes. Device: {args.device}."
-            f"Process (global: {args.rank}, local {args.local_rank}), total {args.world_size}."
-        )
-    elif args.distributed:
-        logging.info(
-            f"Running in distributed mode with multiple processes. Device: {args.device}."
-            f"Process (global: {args.rank}, local {args.local_rank}), total {args.world_size}."
-        )
-    else:
-        logging.info(f"Running with a single process. Device {args.device}.")
-
-    logging.info(f'openai cache dir: {os.path.expanduser(args.openai_model_cache_dir)}')
-
-    # Create CLAP model
-    clap_model, clap_model_cfg = create_model(
-        args.amodel,
-        args.tmodel,
-        args.pretrained,
-        precision=args.precision,
-        device=device,
-        jit=args.torchscript,
-        force_quick_gelu=args.force_quick_gelu,
-        openai_model_cache_dir=os.path.expanduser(args.openai_model_cache_dir),
-        skip_params=False,
-        pretrained_audio=args.pretrained_audio,
-        pretrained_text=args.pretrained_text,
-        enable_fusion=args.enable_fusion,
-        fusion_type=args.fusion_type
-    )
-
-    args.lp_out_ch = len(list(args.class_index_dict.keys()))
-    # Linear Probe 
-    logging.info(f"linear probe using mlp: {args.lp_mlp}")
-    logging.info(f"linear probe using freeze: {args.lp_freeze}")
-    logging.info(f"linear probe act layer: {args.lp_act}")
-    logging.info(f"linear probe out ch: {args.lp_out_ch}")
-    logging.info(f"linear probe learning rate (if applicable): {args.lp_lr}")
-    logging.info(f"linear probe loss func: {args.lp_loss}")
-    logging.info(f"linear probe lp_metrics: {args.lp_metrics}")
-
-    model = LinearProbe(
-        clap_model,
-        mlp=args.lp_mlp, freeze=args.lp_freeze,
-        in_ch=512, out_ch=args.lp_out_ch,
-        act=args.lp_act
-    ) # in_ch is fixed (i.e., 512)
-    model = model.to(device)
-
-    if args.horovod:
-        with torch.no_grad():
-            for param in model.parameters():
-                param.set_(param.contiguous())
-
-    if args.trace:
-        model = trace_model(model, batch_size=args.batch_size, device=device)
-
-    if is_master(args):
-        logging.info("Linear Probe CLAP Model:")
-        logging.info(f"{str(clap_model)}")
-        logging.info("Params:")
-        params_file = os.path.join(args.logs, args.name, "params.txt")
-        with open(params_file, "w") as f:
-            for name in sorted(vars(args)):
-                val = getattr(args, name)
-                logging.info(f"  {name}: {val}")
-                f.write(f"{name}: {val}\n")
-
-    if args.distributed and not args.horovod:
-        if args.use_bn_sync:
-            model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
-        ddp_args = {}
-        if args.ddp_static_graph:
-            # this doesn't exist in older PyTorch, arg only added if enabled
-            ddp_args["static_graph"] = True
-        model = torch.nn.parallel.DistributedDataParallel(
-            model, device_ids=[device], find_unused_parameters=True, **ddp_args
-        )
-
-    data = get_data(args, clap_model_cfg)
-    assert len(data), "At least one train or eval dataset must be specified."
-    if args.trace:
-        assert "train" not in data, "Cannot train with traced model"
-
-
-    optimizer, scheduler, text_freeze_parameters = config_lp_optimizer(model, data, args)
-
-
-    scaler = GradScaler() if args.precision == "amp" else None
-
-    # optionally resume from a checkpoint
-    start_epoch = 0
-    if args.resume is not None:
-        if os.path.isfile(args.resume):
-            checkpoint = torch.load(args.resume, map_location=device)
-            if "epoch" in checkpoint:
-                # resuming a train checkpoint w/ epoch and optimizer state
-                start_epoch = checkpoint["epoch"]
-                sd = checkpoint["state_dict"]
-                if not args.distributed and next(iter(sd.items()))[0].startswith(
-                    "module"
-                ):
-                    sd = {k[len("module.") :]: v for k, v in sd.items()}
-                model.load_state_dict(sd)
-                if args.split_opt:
-                    if optimizer is not None:
-                        for k, o_ in optimizer.items():
-                            o_.load_state_dict(checkpoint[k + "_" + "optimizer"])
-                if optimizer is not None:
-                    optimizer.load_state_dict(checkpoint["optimizer"])
-                if scaler is not None and "scaler" in checkpoint:
-                    scaler.load_state_dict(checkpoint["scaler"])
-                logging.info(
-                    f"=> resuming checkpoint '{args.resume}' (epoch {start_epoch})"
-                )
-            else:
-                # loading a bare (model only) checkpoint for fine-tune or evaluation
-                model.load_state_dict(checkpoint)
-                logging.info(
-                    f"=> loaded checkpoint '{args.resume}' (epoch {start_epoch})"
-                )
-            if args.freeze_text:
-                print("Freeze Text!!!!")
-                for k in text_freeze_parameters:
-                    k.requires_grad = False
-        else:
-            logging.info("=> no checkpoint found at '{}'".format(args.resume))
-
-    cudnn.benchmark = True
-    cudnn.deterministic = False
-
-    # determine if this worker should save logs and checkpoints. only do so if it is rank == 0
-    args.save_logs = args.logs and args.logs.lower() != "none" and is_master(args)
-    writer = None
-    if args.save_logs and args.tensorboard:
-        assert tensorboard is not None, "Please install tensorboard."
-        writer = tensorboard.SummaryWriter(args.tensorboard_path)
-
-    if args.wandb and is_master(args):
-        assert wandb is not None, "Please install wandb."
-        logging.debug("Starting wandb.")
-        args.train_sz = data["train"].dataloader.num_samples
-        if args.val_data is not None:
-            args.val_sz = data["val"].dataloader.num_samples
-        # you will have to configure this for your project!
-        wandb.init(
-            project="clap",
-            notes=args.wandb_notes,
-            name=args.wandb_notes,
-            tags=[],
-            config=vars(args),
-        )
-        if args.debug:
-            wandb.watch(model, log="all")
-        wandb.save(params_file)
-        logging.debug("Finished loading wandb.")
-
-    if "train" not in data:
-        evaluate(model, data, start_epoch, args, writer)
-        return
-    elif start_epoch == 0 and "val" in data and not args.no_eval:
-        evaluate(model, data, 0, args, writer)
-    if args.save_top_performance:
-        current_top_k_ckpt_metrics = {
-            i: 0 for i in range(args.save_top_performance)
-        }  # initialize the top-k metric for ckpts to 0
-
-    for epoch in range(start_epoch, args.epochs):
-        # freeze the text param after (include) args.freeze_text_after, this is -1 by default
-        if epoch == args.freeze_text_after:
-            print("Text pretrained parameters are freezed since this epoch.")
-            for k in text_freeze_parameters:
-                k.requires_grad = False
-        if is_master(args):
-            logging.info(f"Start epoch {epoch}")
-
-        train_one_epoch(model, data, epoch, optimizer, scaler, scheduler, args, writer)
-        completed_epoch = epoch + 1
-
-        if any(v in data for v in ("val", "imagenet-val", "imagenet-v2")) and not args.no_eval:
-            metrics = evaluate(model, data, completed_epoch, args, writer)
-            if args.save_top_performance:
-                top_k_dataset = args.top_k_checkpoint_select_dataset
-                top_k_metric = args.top_k_checkpoint_select_metric
-                filtered_metrics = [
-                    v
-                    for k, v in metrics.items()
-                    if top_k_metric in k and top_k_dataset in k
-                ]  # check all R@10 metrics (all dataset) and use it to update the ckpt
-        # Saving checkpoints.
-        if args.save_logs:
-            opt_dict = {
-                    k + "_" + "optimizer": v.state_dict() for k, v in optimizer.items()
-                }
-            checkpoint_dict = {
-                "epoch": completed_epoch,
-                "name": args.name,
-                "state_dict": model.state_dict(),
-            }
-            checkpoint_dict.update(opt_dict)
-            if scaler is not None:
-                checkpoint_dict["scaler"] = scaler.state_dict()
-
-            if completed_epoch == args.epochs or (
-                args.save_frequency > 0 and (completed_epoch % args.save_frequency) == 0
-            ):
-                torch.save(
-                    checkpoint_dict,
-                    os.path.join(args.checkpoint_path, f"epoch_{completed_epoch}.pt"),
-                )
-            if args.save_most_recent:
-                torch.save(
-                    checkpoint_dict,
-                    os.path.join(args.checkpoint_path, f"epoch_latest.pt"),
-                )
-            if args.save_top_performance and not args.no_eval:
-                update_top_k_performance(
-                    filtered_metrics,
-                    current_top_k_ckpt_metrics,
-                    args,
-                    checkpoint_dict,
-                    bignumbetter=True,
-                )
-
-    if args.wandb and is_master(args):
-        wandb.finish()
-
-
-def copy_codebase(args):
-    from shutil import copytree, ignore_patterns
-
-    new_code_path = os.path.join(args.logs, args.name, "code")
-    if os.path.exists(new_code_path):
-        print(
-            f"Error. Experiment already exists at {new_code_path}. Use --name to specify a new experiment."
-        )
-        return -1
-    print(f"Copying codebase to {new_code_path}")
-    current_code_path = os.path.realpath(__file__)
-    for _ in range(3):
-        current_code_path = os.path.dirname(current_code_path)
-    copytree(
-        current_code_path, new_code_path, ignore=ignore_patterns("log", "logs", "wandb")
-    )
-    print("Done copying code.")
-    return 1
-
-
-if __name__ == "__main__":
-    main()

utils/laion_clap/training/lp_train.py

@@ -1,292 +0,0 @@
-import json
-import logging
-import math
-import os
-import time
-from contextlib import suppress
-
-import numpy as np
-import torch
-import torch.nn.functional as F
-
-try:
-    import wandb
-except ImportError:
-    wandb = None
-
-from clap_module import LPLoss, LPMetrics, lp_gather_features
-from clap_module.utils import do_mixup, get_mix_lambda
-from .distributed import is_master
-from .zero_shot import zero_shot_eval
-
-
-class AverageMeter(object):
-    """Computes and stores the average and current value"""
-
-    def __init__(self):
-        self.reset()
-
-    def reset(self):
-        self.val = 0
-        self.avg = 0
-        self.sum = 0
-        self.count = 0
-
-    def update(self, val, n=1):
-        self.val = val
-        self.sum += val * n
-        self.count += n
-        self.avg = self.sum / self.count
-
-
-def unwrap_model(model):
-    if hasattr(model, "module"):
-        return model.module
-    else:
-        return model
-
-
-def train_one_epoch(
-        model, data, epoch, optimizer, scaler, scheduler, args, tb_writer=None, extra_suffix=""
-):
-    device = torch.device(args.device)
-    autocast = torch.cuda.amp.autocast if args.precision == "amp" else suppress
-    model.train()
-    loss = LPLoss(args.lp_loss)
-
-    dataloader, sampler = data["train"].dataloader, data["train"].sampler
-    if args.distributed and sampler is not None:
-        sampler.set_epoch(epoch)
-    num_batches_per_epoch = dataloader.num_batches
-    sample_digits = math.ceil(math.log(dataloader.num_samples + 1, 10))
-
-    # for toy dataset
-    if args.dataset_type == "toy":
-        dataloader.dataset.generate_queue()
-
-    loss_m = AverageMeter()
-    batch_time_m = AverageMeter()
-    data_time_m = AverageMeter()
-    end = time.time()
-
-    for i, batch in enumerate(dataloader):
-        step = num_batches_per_epoch * epoch + i
-
-        if isinstance(scheduler, dict):
-            for s in scheduler.values():
-                s(step)
-        else:
-            scheduler(step)
-
-        audio = batch # contains mel_spec, wavform, and longer list
-        class_label = batch['class_label']
-        # audio = audio.to(device=device, non_blocking=True)
-        class_label = class_label.to(device=device, non_blocking=True)
-
-        if args.mixup:
-            # https://github.com/RetroCirce/HTS-Audio-Transformer/blob/main/utils.py#L146
-            mix_lambda = torch.from_numpy(get_mix_lambda(0.5, len(audio["waveform"]))).to(device)
-            class_label = do_mixup(class_label, mix_lambda)
-        else:
-            mix_lambda = None
-
-        data_time_m.update(time.time() - end)
-        if isinstance(optimizer, dict):
-            for o_ in optimizer.values():
-                o_.zero_grad()
-        else:
-            optimizer.zero_grad()
-
-        with autocast():
-            pred = model(audio, mix_lambda=mix_lambda, device=device)
-            total_loss = loss(pred, class_label)
-
-        if isinstance(optimizer, dict):
-            if scaler is not None:
-                scaler.scale(total_loss).backward()
-                for o_ in optimizer.values():
-                    if args.horovod:
-                        o_.synchronize()
-                        scaler.unscale_(o_)
-                        with o_.skip_synchronize():
-                            scaler.step(o_)
-                    else:
-                        scaler.step(o_)
-                scaler.update()
-            else:
-                total_loss.backward()
-                for o_ in optimizer.values():
-                    o_.step()
-        else:
-            if scaler is not None:
-                scaler.scale(total_loss).backward()
-                if args.horovod:
-                    optimizer.synchronize()
-                    scaler.unscale_(optimizer)
-                    with optimizer.skip_synchronize():
-                        scaler.step(optimizer)
-                else:
-                    scaler.step(optimizer)
-                scaler.update()
-            else:
-                total_loss.backward()
-                optimizer.step()
-
-        # Note: we clamp to 4.6052 = ln(100), as in the original paper.
-        with torch.no_grad():
-            unwrap_model(model).clap_model.logit_scale_a.clamp_(0, math.log(100))
-            unwrap_model(model).clap_model.logit_scale_t.clamp_(0, math.log(100))
-
-        batch_time_m.update(time.time() - end)
-        end = time.time()
-        batch_count = i + 1
-
-        if is_master(args) and (i % 100 == 0 or batch_count == num_batches_per_epoch):
-            if isinstance(audio, dict):
-                batch_size = len(audio["waveform"])
-            else:
-                batch_size = len(audio)
-            num_samples = batch_count * batch_size * args.world_size
-            samples_per_epoch = dataloader.num_samples
-            percent_complete = 100.0 * batch_count / num_batches_per_epoch
-
-            # NOTE loss is coarsely sampled, just master node and per log update
-            loss_m.update(total_loss.item(), batch_size)
-            if isinstance(optimizer, dict):
-                logging.info(
-                    f"Train Epoch: {epoch} [{num_samples:>{sample_digits}}/{samples_per_epoch} ({percent_complete:.0f}%)] "
-                    f"Loss: {loss_m.val:#.5g} ({loss_m.avg:#.4g}) "
-                    f"Data (t): {data_time_m.avg:.3f} "
-                    f"Batch (t): {batch_time_m.avg:.3f} "
-                    f"LR: {[o_.param_groups[0]['lr'] for o_ in optimizer.values()]}"
-                )
-                log_data = {
-                    "loss": loss_m.val,
-                    "data_time": data_time_m.val,
-                    "batch_time": batch_time_m.val,
-                    "lr": [o_.param_groups[0]["lr"] for o_ in optimizer.values()],
-                }
-            else:
-                logging.info(
-                    f"Train Epoch: {epoch} [{num_samples:>{sample_digits}}/{samples_per_epoch} ({percent_complete:.0f}%)] "
-                    f"Loss: {loss_m.val:#.5g} ({loss_m.avg:#.4g}) "
-                    f"Data (t): {data_time_m.avg:.3f} "
-                    f"Batch (t): {batch_time_m.avg:.3f} "
-                    f"LR: {optimizer.param_groups[0]['lr']:5f} "
-                )
-
-                # Save train loss / etc. Using non avg meter values as loggers have their own smoothing
-                log_data = {
-                    "loss": loss_m.val,
-                    "data_time": data_time_m.val,
-                    "batch_time": batch_time_m.val,
-                    "lr": optimizer.param_groups[0]["lr"],
-                }
-            for name, val in log_data.items():
-                name = f"train{extra_suffix}/{name}"
-                if tb_writer is not None:
-                    tb_writer.add_scalar(name, val, step)
-                if args.wandb:
-                    assert wandb is not None, "Please install wandb."
-                    wandb.log({name: val, "step": step})
-
-            # resetting batch / data time meters per log window
-            batch_time_m.reset()
-            data_time_m.reset()
-    # end for
-
-def evaluate(model, data, epoch, args, tb_writer=None, extra_suffix=""):
-    metrics = {}
-    if not args.parallel_eval:
-        if not is_master(args):
-            return metrics
-    device = torch.device(args.device)
-    model.eval()
-
-    # CHANGE
-    # zero_shot_metrics = zero_shot_eval(model, data, epoch, args)
-    # metrics.update(zero_shot_metrics)
-    if is_master(args):
-        print('Evaluating...')
-        metric_names = args.lp_metrics.split(',')
-        eval_tool = LPMetrics(metric_names=metric_names)
-
-    autocast = torch.cuda.amp.autocast if args.precision == "amp" else suppress        
-    if "val" in data and (
-            args.val_frequency
-            and ((epoch % args.val_frequency) == 0 or epoch == args.epochs)
-    ):
-        if args.parallel_eval:
-            dataloader, sampler = data["val"].dataloader, data["val"].sampler
-            if args.distributed and sampler is not None:
-                sampler.set_epoch(epoch)
-            samples_per_val = dataloader.num_samples
-        else:
-            dataloader = data["val"].dataloader
-            num_samples = 0
-            samples_per_val = dataloader.num_samples
-            
-        eval_info = {
-            'pred': [],
-            'target': []
-        }
-        with torch.no_grad():
-            for i, batch in enumerate(dataloader):
-                audio = batch # contains mel_spec, wavform, and longer list
-                class_label = batch['class_label']
-                           
-                # audio = audio.to(device=device, non_blocking=True)
-                class_label = class_label.to(device=device, non_blocking=True)
-
-                with autocast():
-                    pred = model(audio, device=device)
-                    if args.parallel_eval:
-                        pred, class_label = lp_gather_features(pred, class_label, args.world_size, args.horovod)
-                    eval_info['pred'].append(pred)
-                    eval_info['target'].append(class_label)
-                    
-                num_samples += class_label.shape[0]
-
-                if (i % 100) == 0:  # and i != 0:
-                    logging.info(
-                        f"Eval Epoch: {epoch} [{num_samples} / {samples_per_val}]"
-                    )
-                    
-            if is_master(args):
-                eval_info['pred'] = torch.cat(eval_info['pred'], 0).cpu()
-                eval_info['target'] = torch.cat(eval_info['target'], 0).cpu()
-                metric_dict = eval_tool.evaluate_mertics(eval_info['pred'], eval_info['target'])
-                metrics.update(metric_dict)
-                if "epoch" not in metrics.keys():
-                    metrics.update({"epoch": epoch})
-
-    if is_master(args):
-        if not metrics:
-            return metrics
-
-        logging.info(
-            f"Eval Epoch: {epoch} "
-            + "\n".join(
-                [
-                    "\t".join([f"{m}: {round(metrics[m], 4):.4f}" ])
-                    for m in metrics
-                ]
-            )
-        )
-        if args.save_logs:
-            for name, val in metrics.items():
-                if tb_writer is not None:
-                    tb_writer.add_scalar(f"val{extra_suffix}/{name}", val, epoch)
-
-            with open(os.path.join(args.checkpoint_path, "results.jsonl"), "a+") as f:
-                f.write(json.dumps(metrics))
-                f.write("\n")
-
-        if args.wandb:
-            assert wandb is not None, "Please install wandb."
-            for name, val in metrics.items():
-                wandb.log({f"val{extra_suffix}/{name}": val, "epoch": epoch})
-
-        return metrics
-    else:
-        return metrics

utils/laion_clap/training/main.py

@@ -1,597 +0,0 @@
-import logging
-import os
-import random
-from datetime import datetime
-import copy
-import numpy as np
-import torch
-import torch.backends.cudnn as cudnn
-from torch.cuda.amp import GradScaler
-
-try:
-    import wandb
-except ImportError:
-    wandb = None
-
-try:
-    import torch.utils.tensorboard as tensorboard
-except ImportError:
-    tensorboard = None
-
-try:
-    import horovod.torch as hvd
-except ImportError:
-    hvd = None
-
-from clap_module import create_model_and_transforms, trace_model, create_model
-from training.data import get_data
-from training.distributed import is_master, init_distributed_device, world_info_from_env
-from training.logger import setup_logging
-from training.params import parse_args
-from training.scheduler import cosine_lr
-from training.train import train_one_epoch, evaluate
-from clap_module.utils import dataset_split, get_optimizer
-
-
-def maintain_ckpts(args, startidx, all_idx_len):
-    for i in reversed(range(startidx, all_idx_len)):
-        if os.path.exists(os.path.join(args.checkpoint_path, f"epoch_top_{i}.pt")):
-            os.rename(
-                os.path.join(args.checkpoint_path, f"epoch_top_{i}.pt"),
-                os.path.join(args.checkpoint_path, f"epoch_top_{i+1}.pt"),
-            )
-    if os.path.exists(
-        os.path.join(args.checkpoint_path, f"epoch_top_{all_idx_len}.pt")
-    ):
-        os.remove(os.path.join(args.checkpoint_path, f"epoch_top_{all_idx_len}.pt"))
-    return
-
-
-def update_top_k_performance(
-    new_metrics_inputs, current_top_k_ckpt_metrics, args, ckpt, bignumbetter=True
-):
-    """
-    Record the top-k performance of the current epoch.
-    current_top_k_metrics is a dictionary of the form: {1: top_1_ckpt_measure, 2: top_2_ckpt_measure, ...}
-    """
-    if isinstance(new_metrics_inputs, (list, tuple)):
-        new_metrics_inputs = np.mean(new_metrics_inputs)
-        return update_top_k_performance(
-            new_metrics_inputs,
-            current_top_k_ckpt_metrics,
-            args=args,
-            ckpt=ckpt,
-            bignumbetter=bignumbetter,
-        )
-    elif isinstance(new_metrics_inputs, dict):
-        new_metrics_inputs = np.mean(list(new_metrics_inputs.values()))
-        return update_top_k_performance(
-            new_metrics_inputs,
-            current_top_k_ckpt_metrics,
-            args=args,
-            ckpt=ckpt,
-            bignumbetter=bignumbetter,
-        )
-    elif isinstance(new_metrics_inputs, (float, int)):
-        update_flag = {k: False for k in current_top_k_ckpt_metrics.keys()}
-        sorted_keys = sorted(current_top_k_ckpt_metrics.keys())
-        sorted_values = sorted(
-            current_top_k_ckpt_metrics.values(), reverse=bignumbetter
-        )
-        sorted_values_ = copy.deepcopy(sorted_values)
-        sorted_values.append(new_metrics_inputs)
-        sorted_values = sorted(sorted_values, reverse=bignumbetter)
-        sorted_values = sorted_values[:-1]
-
-        if sorted_values == sorted_values_:
-            return current_top_k_ckpt_metrics, new_metrics_inputs
-        else:
-            for i in range(len(sorted_keys)):
-                if current_top_k_ckpt_metrics[sorted_keys[i]] != sorted_values[i]:
-                    current_top_k_ckpt_metrics[sorted_keys[i]] = sorted_values[i]
-                    update_flag[sorted_keys[i]] = True
-            for i in range(len(update_flag)):
-                if update_flag[i]:
-                    maintain_ckpts(args, i, len(sorted_keys))
-                    torch.save(
-                        ckpt,
-                        os.path.join(args.checkpoint_path, f"epoch_top_{i}.pt"),
-                    )
-                    break
-            return current_top_k_ckpt_metrics, new_metrics_inputs
-
-
-# def updateifNone(a, b):
-#     a = b if None else a
-#     return a
-
-
-def is_pretrained_params(n):
-    return (
-        n.startswith("transformer")
-        or n in ["positional_embedding", "text_projection"]
-        or n.startswith("token_embedding")
-        or n.startswith("ln_final")
-        or n.startswith("logit_scale_t")
-    )
-
-
-def random_seed(seed=42, rank=0):
-    torch.manual_seed(seed + rank)
-    np.random.seed(seed + rank)
-    random.seed(seed + rank)
-
-
-def main():
-    args = parse_args()
-    # sanitize model name for filesystem / uri use, easier if we don't use / in name as a rule?
-    args.amodel = args.amodel.replace("/", "-")
-    # download sizes.json file
-
-    # (yusong): the below two lines are for debug
-    # print("setting up faulthandler")
-    # faulthandler.register(10)
-
-    random.seed(args.seed)
-    torch.manual_seed(args.seed)
-    torch.cuda.manual_seed(args.seed)
-    torch.cuda.manual_seed_all(args.seed)
-    np.random.seed(args.seed)
-    if args.tmodel == "bert" or args.tmodel == "roberta" or args.tmodel == "bart":
-        assert (
-            args.pretrained == "" or args.pretrained is None
-        ), "bert/roberta/bart text encoder does not support pretrained models."
-
-    # get the name of the experiments
-    if args.name is None:
-        args.name = "-".join(
-            [
-                datetime.now().strftime("%Y_%m_%d-%H_%M_%S"),
-                f"model_{args.amodel}",
-                f"lr_{args.lr}",
-                f"b_{args.batch_size}",
-                f"j_{args.workers}",
-                f"p_{args.precision}",
-            ]
-        )
-
-    # discover initial world args early so we can log properly
-    args.distributed = False
-    args.local_rank, args.rank, args.world_size = world_info_from_env()
-
-    if args.remotedata and is_master(args):
-        for dataset_name in args.datasetnames:
-            for split in dataset_split[dataset_name]:
-                if not os.path.exists(f"./json_files/{dataset_name}/{split}"):
-                    os.makedirs(f"./json_files/{dataset_name}/{split}")
-                os.system(
-                    f"aws s3 cp s3://s-laion-audio/webdataset_tar/{dataset_name}/{split}/sizes.json ./json_files/{dataset_name}/{split}/sizes.json"
-                )
-
-    args.log_path = None
-    if is_master(args, local=args.log_local):
-        log_base_path = os.path.join(args.logs, args.name)
-        os.makedirs(log_base_path, exist_ok=True)
-        log_filename = f"out-{args.rank}" if args.log_local else "out.log"
-        args.log_path = os.path.join(log_base_path, log_filename)
-        if os.path.exists(args.log_path):
-            print(
-                "Error. Experiment already exists. Use --name {} to specify a new experiment."
-            )
-            return -1
-
-    # Set logger
-    args.log_level = logging.DEBUG if args.debug else logging.INFO
-    setup_logging(args.log_path, args.log_level)
-
-    # fully initialize distributed device environment
-    device = init_distributed_device(args)
-
-    args.wandb = "wandb" in args.report_to or "all" in args.report_to
-    args.tensorboard = "tensorboard" in args.report_to or "all" in args.report_to
-    if is_master(args):
-        args.tensorboard_path = (
-            os.path.join(args.logs, args.name, "tensorboard")
-            if args.tensorboard
-            else ""
-        )
-        args.checkpoint_path = os.path.join(args.logs, args.name, "checkpoints")
-        for dirname in [args.tensorboard_path, args.checkpoint_path]:
-            if dirname:
-                os.makedirs(dirname, exist_ok=True)
-    else:
-        args.tensorboard_path = ""
-        args.checkpoint_path = ""
-
-    if args.copy_codebase:
-        copy_codebase(args)
-
-    assert args.precision in ["amp", "fp16", "fp32"]
-    if args.precision == "fp16":
-        logging.warning(
-            "It is recommended to use fp32 mixed-precision instead of FP16 and AMP in this model. "
-            "They will cause NaN loss and NaN gradients. "
-            "FP16 and AMP support needs further verification and tuning, especially for train."
-        )
-
-    if args.horovod:
-        logging.info(
-            f"Running in horovod mode with multiple processes / nodes. Device: {args.device}."
-            f"Process (global: {args.rank}, local {args.local_rank}), total {args.world_size}."
-        )
-    elif args.distributed:
-        logging.info(
-            f"Running in distributed mode with multiple processes. Device: {args.device}."
-            f"Process (global: {args.rank}, local {args.local_rank}), total {args.world_size}."
-        )
-    else:
-        logging.info(f"Running with a single process. Device {args.device}.")
-
-    logging.info(f"openai cache dir: {os.path.expanduser(args.openai_model_cache_dir)}")
-
-    model, model_cfg = create_model(
-        args.amodel,
-        args.tmodel,
-        args.pretrained,
-        precision=args.precision,
-        device=device,
-        jit=args.torchscript,
-        force_quick_gelu=args.force_quick_gelu,
-        openai_model_cache_dir=os.path.expanduser(args.openai_model_cache_dir),
-        skip_params=True,
-        pretrained_audio=args.pretrained_audio,
-        pretrained_text=args.pretrained_text,
-        enable_fusion=args.enable_fusion,
-        fusion_type=args.fusion_type
-    )
-
-    if args.horovod:
-        with torch.no_grad():
-            for param in model.parameters():
-                param.set_(param.contiguous())
-
-    if args.trace:
-        model = trace_model(model, batch_size=args.batch_size, device=device)
-
-    if is_master(args):
-        logging.info("Model:")
-        logging.info(f"{str(model)}")
-        logging.info("Params:")
-        params_file = os.path.join(args.logs, args.name, "params.txt")
-        with open(params_file, "w") as f:
-            for name in sorted(vars(args)):
-                val = getattr(args, name)
-                logging.info(f"  {name}: {val}")
-                f.write(f"{name}: {val}\n")
-
-    if args.distributed and not args.horovod:
-        if args.use_bn_sync:
-            model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
-        ddp_args = {}
-        if args.ddp_static_graph:
-            # this doesn't exist in older PyTorch, arg only added if enabled
-            ddp_args["static_graph"] = True
-        model = torch.nn.parallel.DistributedDataParallel(
-            model, device_ids=[device], find_unused_parameters=True, **ddp_args
-        )
-
-    data = get_data(args, model_cfg)
-    assert len(data), "At least one train or eval dataset must be specified."
-    if args.trace:
-        assert "train" not in data, "Cannot train with traced model"
-
-    exclude = (
-        lambda n, p: p.ndim < 2
-        or "bn" in n
-        or "ln" in n
-        or "bias" in n
-        or "logit_scale" in n
-    )
-    include = lambda n, p: not exclude(n, p)
-
-    named_parameters = list(model.named_parameters())
-
-    # freeze text encoder
-    text_freeze_parameters = [
-        p
-        for n, p in named_parameters
-        if 'text_branch' in n
-    ]
-
-    if args.freeze_text:
-        print("Freeze Text!!!!")
-        for k in text_freeze_parameters:
-            k.requires_grad = False
-
-    gain_or_bias_params = [
-        p for n, p in named_parameters if exclude(n, p) and p.requires_grad
-    ]
-    rest_params = [p for n, p in named_parameters if include(n, p) and p.requires_grad]
-
-    # set wd-related params to 0 if use adam optimizer
-    if args.optimizer == "adam":
-        args.wd = 0
-        args.wd_pretrained = 0
-        args.wd_new = 0
-
-    if args.train_data is None:
-        optimizer = None
-        scheduler = None
-    else:
-        total_steps = data["train"].dataloader.num_batches * args.epochs
-
-        if args.split_opt:
-            for x in ["lr", "beta1", "beta2", "eps", "wd"]:
-                for y in ["_new", "_pretrained"]:
-                    if getattr(args, x + y) is None:
-                        setattr(args, x + y, getattr(args, x))
-
-            gain_or_bias_pretrained_params = [
-                p
-                for n, p in named_parameters
-                if (exclude(n, p) and p.requires_grad) and is_pretrained_params(n)
-            ]
-            rest_pretrained_params = [
-                p
-                for n, p in named_parameters
-                if (include(n, p) and p.requires_grad) and is_pretrained_params(n)
-            ]
-            gain_or_bias_new_params = [
-                p
-                for n, p in named_parameters
-                if (exclude(n, p) and p.requires_grad) and (not is_pretrained_params(n))
-            ]
-            rest_new_params = [
-                p
-                for n, p in named_parameters
-                if (include(n, p) and p.requires_grad) and (not is_pretrained_params(n))
-            ]
-            pretrained_params_optimizer = get_optimizer(
-                [
-                    {"params": gain_or_bias_pretrained_params, "weight_decay": 0.0},
-                    {
-                        "params": rest_pretrained_params,
-                        "weight_decay": args.wd_pretrained,
-                    },
-                ],
-                lr=args.lr_pretrained,
-                betas=(args.beta1_pretrained, args.beta2_pretrained),
-                eps=args.eps_pretrained,
-                momentum=args.momentum_pretrained,
-                optimizer_name=args.optimizer,
-                )
-            pretrained_params_scheduler = cosine_lr(
-                pretrained_params_optimizer,
-                args.lr_pretrained,
-                args.warmup,
-                total_steps,
-            )
-            new_params_optimizer = get_optimizer(
-                [
-                    {"params": gain_or_bias_new_params, "weight_decay": 0.0},
-                    {"params": rest_new_params, "weight_decay": args.wd_new},
-                ],
-                lr=args.lr_new,
-                betas=(args.beta1_new, args.beta2_new),
-                eps=args.eps_new,
-                momentum=args.momentum_new,
-                optimizer_name=args.optimizer,
-                )
-
-            new_params_scheduler = cosine_lr(
-                new_params_optimizer, args.lr_new, args.warmup, total_steps
-            )
-
-            optimizer = {
-                "pretrained": pretrained_params_optimizer,
-                "new": new_params_optimizer,
-            }
-            scheduler = {
-                "pretrained": pretrained_params_scheduler,
-                "new": new_params_scheduler,
-            }
-
-            if args.horovod:
-                pretrained_params_optimizer = hvd.DistributedOptimizer(
-                    pretrained_params_optimizer,
-                    named_parameters=model.named_parameters(),
-                )
-                new_params_optimizer = hvd.DistributedOptimizer(
-                    new_params_optimizer, named_parameters=model.named_parameters()
-                )
-                hvd.broadcast_parameters(model.state_dict(), root_rank=0)
-                hvd.broadcast_optimizer_state(pretrained_params_optimizer, root_rank=0)
-                hvd.broadcast_optimizer_state(new_params_optimizer, root_rank=0)
-        else:
-            optimizer = get_optimizer(
-                [
-                    {"params": gain_or_bias_params, "weight_decay": 0.0},
-                    {"params": rest_params, "weight_decay": args.wd},
-                ],
-                lr=args.lr,
-                betas=(args.beta1, args.beta2),
-                eps=args.eps,
-                momentum=args.momentum,
-                optimizer_name=args.optimizer,
-            )
-
-            scheduler = cosine_lr(optimizer, args.lr, args.warmup, total_steps)
-
-            if args.horovod:
-                optimizer = hvd.DistributedOptimizer(
-                    optimizer, named_parameters=model.named_parameters()
-                )
-                hvd.broadcast_parameters(model.state_dict(), root_rank=0)
-                hvd.broadcast_optimizer_state(optimizer, root_rank=0)
-
-    scaler = GradScaler() if args.precision == "amp" else None
-
-    # optionally resume from a checkpoint
-    start_epoch = 0
-    if args.resume is not None:
-        if os.path.isfile(args.resume):
-            checkpoint = torch.load(args.resume, map_location=device)
-            if "epoch" in checkpoint:
-                # resuming a train checkpoint w/ epoch and optimizer state
-                start_epoch = checkpoint["epoch"]
-                sd = checkpoint["state_dict"]
-                if not args.distributed and next(iter(sd.items()))[0].startswith(
-                    "module"
-                ):
-                    sd = {k[len("module.") :]: v for k, v in sd.items()}
-                model.load_state_dict(sd)
-                if args.split_opt:
-                    if optimizer is not None:
-                        for k, o_ in optimizer.items():
-                            o_.load_state_dict(checkpoint[k + "_" + "optimizer"])
-                if optimizer is not None:
-                    optimizer.load_state_dict(checkpoint["optimizer"])
-                if scaler is not None and "scaler" in checkpoint:
-                    scaler.load_state_dict(checkpoint["scaler"])
-                logging.info(
-                    f"=> resuming checkpoint '{args.resume}' (epoch {start_epoch})"
-                )
-            else:
-                # loading a bare (model only) checkpoint for fine-tune or evaluation
-                model.load_state_dict(checkpoint)
-                logging.info(
-                    f"=> loaded checkpoint '{args.resume}' (epoch {start_epoch})"
-                )
-            if args.freeze_text:
-                print("Freeze Text!!!!")
-                for k in text_freeze_parameters:
-                    k.requires_grad = False
-        else:
-            logging.info("=> no checkpoint found at '{}'".format(args.resume))
-
-    cudnn.benchmark = True
-    cudnn.deterministic = False
-
-    # determine if this worker should save logs and checkpoints. only do so if it is rank == 0
-    args.save_logs = args.logs and args.logs.lower() != "none" and is_master(args)
-    writer = None
-    if args.save_logs and args.tensorboard:
-        assert tensorboard is not None, "Please install tensorboard."
-        writer = tensorboard.SummaryWriter(args.tensorboard_path)
-
-    if args.wandb and is_master(args):
-        assert wandb is not None, "Please install wandb."
-        logging.debug("Starting wandb.")
-        args.train_sz = data["train"].dataloader.num_samples
-        if args.val_data is not None:
-            args.val_sz = data["val"].dataloader.num_samples
-        # you will have to configure this for your project!
-        wandb.init(
-            entity="clap",
-            project="clap",
-            notes=args.wandb_notes,
-            name=args.wandb_notes,
-            tags=[],
-            config=vars(args),
-        )
-        if args.debug:
-            wandb.watch(model, log="all")
-        wandb.save(params_file)
-        logging.debug("Finished loading wandb.")
-
-    if "train" not in data:
-        evaluate(model, data, start_epoch, args, writer)
-        return
-    elif start_epoch == 0 and "val" in data and not args.no_eval:
-        evaluate(model, data, 0, args, writer)
-        #  print(f'rank {args.rank}, Start First Evaluation')#  (yusong): for debug
-    if args.save_top_performance:
-        current_top_k_ckpt_metrics = {
-            i: 0 for i in range(args.save_top_performance)
-        }  # initialize the top-k metric for ckpts to 0
-
-    #  print(f'rank {args.rank}, Start Training') #  (yusong): for debug
-    for epoch in range(start_epoch, args.epochs):
-        # freeze the text param after (include) args.freeze_text_after, this is -1 by default
-        if epoch == args.freeze_text_after:
-            print("Text pretrained parameters are freezed since this epoch.")
-            for k in text_freeze_parameters:
-                k.requires_grad = False
-        if is_master(args):
-            logging.info(f"Start epoch {epoch}")
-
-        train_one_epoch(model, data, epoch, optimizer, scaler, scheduler, args, writer)
-        completed_epoch = epoch + 1
-
-        if (
-            any(v in data for v in ("val", "imagenet-val", "imagenet-v2"))
-            and not args.no_eval
-        ):
-            metrics = evaluate(model, data, completed_epoch, args, writer)
-            if args.save_top_performance:
-                top_k_dataset = args.top_k_checkpoint_select_dataset
-                top_k_metric = args.top_k_checkpoint_select_metric
-                filtered_metrics = [
-                    v
-                    for k, v in metrics.items()
-                    if top_k_metric in k and top_k_dataset in k
-                ]  # check all R@10 metrics (all dataset) and use it to update the ckpt
-        # Saving checkpoints.
-        if args.save_logs:
-            if args.split_opt:
-                opt_dict = {
-                    k + "_" + "optimizer": v.state_dict() for k, v in optimizer.items()
-                }
-            else:
-                opt_dict = {"optimizer": optimizer.state_dict()}
-            checkpoint_dict = {
-                "epoch": completed_epoch,
-                "name": args.name,
-                "state_dict": model.state_dict(),
-            }
-            checkpoint_dict.update(opt_dict)
-            if scaler is not None:
-                checkpoint_dict["scaler"] = scaler.state_dict()
-
-            if completed_epoch == args.epochs or (
-                args.save_frequency > 0 and (completed_epoch % args.save_frequency) == 0
-            ):
-                torch.save(
-                    checkpoint_dict,
-                    os.path.join(args.checkpoint_path, f"epoch_{completed_epoch}.pt"),
-                )
-            if args.save_most_recent:
-                torch.save(
-                    checkpoint_dict,
-                    os.path.join(args.checkpoint_path, f"epoch_latest.pt"),
-                )
-            if args.save_top_performance and not args.no_eval:
-                update_top_k_performance(
-                    filtered_metrics,
-                    current_top_k_ckpt_metrics,
-                    args,
-                    checkpoint_dict,
-                    bignumbetter=True,
-                )
-
-    if args.wandb and is_master(args):
-        wandb.finish()
-
-
-def copy_codebase(args):
-    from shutil import copytree, ignore_patterns
-
-    new_code_path = os.path.join(args.logs, args.name, "code")
-    if os.path.exists(new_code_path):
-        print(
-            f"Error. Experiment already exists at {new_code_path}. Use --name to specify a new experiment."
-        )
-        return -1
-    print(f"Copying codebase to {new_code_path}")
-    current_code_path = os.path.realpath(__file__)
-    for _ in range(3):
-        current_code_path = os.path.dirname(current_code_path)
-    copytree(
-        current_code_path, new_code_path, ignore=ignore_patterns("log", "logs", "wandb")
-    )
-    print("Done copying code.")
-    return 1
-
-
-if __name__ == "__main__":
-    main()

utils/laion_clap/training/params.py

@@ -1,567 +0,0 @@
-import argparse
-
-
-def get_default_params(model_name):
-    # Params from paper (https://arxiv.org/pdf/2103.00020.pdf)
-    model_name = model_name.lower()
-    if "vit" in model_name:
-        return {"lr": 5.0e-4, "beta1": 0.9, "beta2": 0.98, "eps": 1.0e-6}
-    else:
-        return {"lr": 5.0e-4, "beta1": 0.9, "beta2": 0.999, "eps": 1.0e-8}
-
-
-def parse_args():
-    parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "--train-data",
-        type=str,
-        default=None,
-        help="Path to h5 filewith training data",
-    )
-    parser.add_argument(
-        "--val-data",
-        type=str,
-        default=None,
-        help="Path to h5 file with validation data",
-    )
-    parser.add_argument(
-        "--freeze-text",
-        default=False,
-        action="store_true",
-        help="if you need to freeze the text encoder, make this True",
-    )
-    parser.add_argument(
-        "--freeze-text-after",
-        type=int,
-        default=-1,
-        help="if you need to freeze the text encoder after (include) epoch x, set this param to x. Set -1 to disable it",
-    )
-    parser.add_argument(
-        "--train-ipc",
-        type=str,
-        default=None,
-        help="Path to npy file of the number of instance per class in training data",
-    )
-    parser.add_argument(
-        "--val-ipc",
-        type=str,
-        default=None,
-        help="Path to npy file of the number of instance per class in validation data",
-    )
-    parser.add_argument(
-        "--train-num-samples",
-        type=int,
-        default=None,
-        help="Number of samples in dataset. Required for webdataset if not available in info file.",
-    )
-    parser.add_argument(
-        "--val-num-samples",
-        type=int,
-        default=None,
-        help="Number of samples in dataset. Useful for webdataset if not available in info file.",
-    )
-    parser.add_argument(
-        "--dataset-type",
-        choices=["webdataset", "csv", "auto", "toy"],
-        default="auto",
-        help="Which type of dataset to process.",
-    )
-    parser.add_argument(
-        "--csv-separator",
-        type=str,
-        default="\t",
-        help="For csv-like datasets, which separator to use.",
-    )
-    parser.add_argument(
-        "--csv-img-key",
-        type=str,
-        default="filepath",
-        help="For csv-like datasets, the name of the key for the image paths.",
-    )
-    parser.add_argument(
-        "--csv-caption-key",
-        type=str,
-        default="title",
-        help="For csv-like datasets, the name of the key for the captions.",
-    )
-    parser.add_argument(
-        "--imagenet-val",
-        type=str,
-        default=None,
-        help="Path to imagenet val set for conducting zero shot evaluation.",
-    )
-    parser.add_argument(
-        "--imagenet-v2",
-        type=str,
-        default=None,
-        help="Path to imagenet v2 for conducting zero shot evaluation.",
-    )
-    parser.add_argument(
-        "--datasetnames",
-        nargs="+",
-        default=None,
-        help="If loading webdataset, spedify the dataset names to load. Can be some of these: Clotho, audioset, audiocaps, BBCSoundEffects",
-    )
-    parser.add_argument(
-        "--full-train-dataset",
-        nargs="+",
-        default=None,
-        help="Which dataset will be trained with all the subsets. (train+test)",
-    )
-    parser.add_argument(
-        "--exclude-eval-dataset",
-        nargs="+",
-        default=None,
-        help="Which dataset will be excluded with evaluation",
-    )
-    parser.add_argument(
-        "--datasetinfos",
-        nargs="+",
-        default=None,
-        help="If loading webdataset, spedify the dataset types to load. Can be some of these: train, test, valid, unbalanced_train, balanced_train, eval",
-    )
-    parser.add_argument(
-        "--dataset-proportion",
-        type=float,
-        default=1.0,
-        help="How much proportion of dataset we want to train.",
-    )
-    parser.add_argument(
-        "--remotedata",
-        default=False,
-        action="store_true",
-        help="if the dataset is remote, set this flag",
-    )
-    parser.add_argument(
-        "--class-label-path",
-        type=str,
-        default=None,
-        help="The path of the class label pickle or csv.",
-    )
-    parser.add_argument(
-        "--datasetpath",
-        type=str,
-        default="/mnt/audio_clip/webdataset_tar",
-        help="The path to the dataset",
-    )
-    parser.add_argument(
-        "--logs",
-        type=str,
-        default="./logs/",
-        help="Where to store tensorboard logs. Use None to avoid storing logs.",
-    )
-    parser.add_argument(
-        "--log-local",
-        action="store_true",
-        default=False,
-        help="log files on local master, otherwise global master only.",
-    )
-    parser.add_argument(
-        "--name",
-        type=str,
-        default=None,
-        help="Optional identifier for the experiment when storing logs. Otherwise use current time.",
-    )
-    parser.add_argument(
-        "--workers", type=int, default=1, help="Number of workers per GPU."
-    )
-    parser.add_argument(
-        "--batch-size", type=int, default=64, help="Batch size per GPU."
-    )
-    parser.add_argument(
-        "--epochs", type=int, default=32, help="Number of epochs to train for."
-    )
-    parser.add_argument("--lr", type=float, default=None, help="Learning rate.")
-    parser.add_argument("--beta1", type=float, default=None, help="Adam beta 1.")
-    parser.add_argument("--beta2", type=float, default=None, help="Adam beta 2.")
-    parser.add_argument("--eps", type=float, default=None, help="Adam epsilon.")
-    parser.add_argument("--momentum", type=float, default=None, help="SGD epsilon.")
-    parser.add_argument("--wd", type=float, default=0.2, help="Weight decay.")
-
-    parser.add_argument(
-        "--split-opt",
-        action="store_true",
-        default=False,
-        help="Use this flag to skip the learning rate decay.",
-    )
-    parser.add_argument(
-        "--lr-pretrained", type=float, default=None, help="Learning rate for text."
-    )
-    parser.add_argument(
-        "--beta1-pretrained", type=float, default=None, help="Adam beta 1 for text."
-    )
-    parser.add_argument(
-        "--beta2-pretrained", type=float, default=None, help="Adam beta 2 for text."
-    )
-    parser.add_argument(
-        "--eps-pretrained", type=float, default=None, help="Adam epsilon for text."
-    )
-    parser.add_argument(
-        "--wd-pretrained", type=float, default=0.2, help="Weight decay for text."
-    )
-    parser.add_argument(
-        "--momentum-pretrained", type=float, default=0.9, help="Momentum for text."
-    )
-    parser.add_argument(
-        "--lr-new", type=float, default=None, help="Learning rate for audio."
-    )
-    parser.add_argument(
-        "--beta1-new", type=float, default=None, help="Adam beta 1 for audio."
-    )
-    parser.add_argument(
-        "--beta2-new", type=float, default=None, help="Adam beta 2 for audio."
-    )
-    parser.add_argument(
-        "--eps-new", type=float, default=None, help="Adam epsilon for audio."
-    )
-    parser.add_argument(
-        "--wd-new", type=float, default=0.2, help="Weight decay for audio."
-    )
-    parser.add_argument(
-        "--momentum-new", type=float, default=0.9, help="Momentum for audio."
-    )
-    parser.add_argument(
-        "--warmup", type=int, default=10000, help="Number of steps to warmup for."
-    )
-    parser.add_argument(
-        "--use-bn-sync",
-        default=False,
-        action="store_true",
-        help="Whether to use batch norm sync.",
-    )
-    parser.add_argument(
-        "--skip-scheduler",
-        action="store_true",
-        default=False,
-        help="Use this flag to skip the learning rate decay.",
-    )
-    parser.add_argument(
-        "--save-frequency", type=int, default=1, help="How often to save checkpoints."
-    )
-    parser.add_argument(
-        "--save-top-performance",
-        type=int,
-        default=0,
-        help="Save the top x performance weights if the value >0",
-    )
-    parser.add_argument(
-        "--save-most-recent",
-        action="store_true",
-        default=False,
-        help="Always save the most recent model trained to epoch_latest.pt.",
-    )
-    parser.add_argument(
-        "--zeroshot-frequency", type=int, default=2, help="How often to run zero shot."
-    )
-    parser.add_argument(
-        "--val-frequency",
-        type=int,
-        default=1,
-        help="How often to run evaluation with val data.",
-    )
-    parser.add_argument(
-        "--resume",
-        default=None,
-        type=str,
-        help="path to latest checkpoint (default: none)",
-    )
-    parser.add_argument(
-        "--precision",
-        choices=["amp", "fp16", "fp32"],
-        default="amp",
-        help="Floating point precision.",
-    )
-    parser.add_argument(
-        "--amodel",
-        type=str,
-        default="RN50",
-        help="Name of the audio backbone to use.",
-    )
-    parser.add_argument(
-        "--tmodel",
-        type=str,
-        default="transformer",
-        help="Name of the text backbone to use. Can be [transformer, bert, roberta, bart]",
-    )
-    parser.add_argument(
-        "--pretrained-audio",
-        default="",
-        type=str,
-        help="Use a pretrained audio model weights for the audio encoder of CLAP",
-    )
-    parser.add_argument(
-        "--pretrained-text",
-        default="",
-        type=str,
-        help="Use a pretrained text model weights for the text encoder of CLAP",
-    )
-    parser.add_argument(
-        "--pretrained",
-        default="",
-        type=str,
-        help="Use a pretrained CLIP model weights with the specified tag or file path.",
-    )
-    parser.add_argument(
-        "--pretrained-image",
-        default=False,
-        action="store_true",
-        help="Load imagenet pretrained weights for image tower backbone if available.",
-    )
-    parser.add_argument(
-        "--lock-image",
-        default=False,
-        action="store_true",
-        help="Lock full image tower by disabling gradients.",
-    )
-    parser.add_argument(
-        "--lock-image-unlocked-groups",
-        type=int,
-        default=0,
-        help="Leave last n image tower layer groups unlocked.",
-    )
-    parser.add_argument(
-        "--lock-image-freeze-bn-stats",
-        default=False,
-        action="store_true",
-        help="Freeze BatchNorm running stats in image tower for any locked layers.",
-    )
-    parser.add_argument(
-        "--local-loss",
-        default=False,
-        action="store_true",
-        help="calculate loss w/ local features @ global (instead of realizing full global @ global matrix)",
-    )
-    parser.add_argument(
-        "--gather-with-grad",
-        default=False,
-        action="store_true",
-        help="enable full distributed gradient for feature gather",
-    )
-    parser.add_argument(
-        "--force-quick-gelu",
-        default=False,
-        action="store_true",
-        help="Force use of QuickGELU activation for non-OpenAI transformer models.",
-    )
-    parser.add_argument(
-        "--torchscript",
-        default=False,
-        action="store_true",
-        help="torch.jit.script the model, also uses jit version of OpenAI models if pretrained=='openai'",
-    )
-    parser.add_argument(
-        "--trace",
-        default=False,
-        action="store_true",
-        help="torch.jit.trace the model for inference / eval only",
-    )
-    # arguments for distributed training
-    parser.add_argument(
-        "--dist-url",
-        default="env://",
-        type=str,
-        help="url used to set up distributed training",
-    )
-    parser.add_argument(
-        "--dist-backend", default="nccl", type=str, help="distributed backend"
-    )
-    parser.add_argument(
-        "--report-to",
-        default="",
-        type=str,
-        help="Options are ['wandb', 'tensorboard', 'wandb,tensorboard']",
-    )
-    parser.add_argument(
-        "--wandb-notes", default="", type=str, help="Notes if logging with wandb"
-    )
-    parser.add_argument(
-        "--C", type=float, default=3.16, help="inverse regularizer for logistic reg."
-    )
-    parser.add_argument(
-        "--debug",
-        default=False,
-        action="store_true",
-        help="If true, more information is logged.",
-    )
-    parser.add_argument(
-        "--copy-codebase",
-        default=False,
-        action="store_true",
-        help="If true, we copy the entire base on the log diretory, and execute from there.",
-    )
-    parser.add_argument(
-        "--horovod",
-        default=False,
-        action="store_true",
-        help="Use horovod for distributed training.",
-    )
-    parser.add_argument(
-        "--ddp-static-graph",
-        default=False,
-        action="store_true",
-        help="Enable static graph optimization for DDP in PyTorch >= 1.11.",
-    )
-    parser.add_argument(
-        "--no-set-device-rank",
-        default=False,
-        action="store_true",
-        help="Don't set device index from local rank (when CUDA_VISIBLE_DEVICES restricted to one per proc).",
-    )
-    parser.add_argument("--seed", type=int, default=4242, help="Default random seed.")
-
-    parser.add_argument(
-        "--top-k-checkpoint-select-dataset",
-        type=str,
-        default="all",
-        help="The dataset of selecting top-k checkpoint.",
-    )
-
-    # @R10, @R@5, @R1, mAP@10
-    parser.add_argument(
-        "--top-k-checkpoint-select-metric",
-        type=str,
-        default="_R@10",
-        help="The metric for selecting top-k checkpoint.",
-    )
-    parser.add_argument(
-        "--openai-model-cache-dir",
-        type=str,
-        default="~/.cache/clip",
-        help="Directory to download OpenAI models.",
-    )
-    parser.add_argument(
-        "--optimizer",
-        type=str,
-        default="adamw",
-        help="can be AdamW or SGD",
-    )
-    parser.add_argument(
-        "--parallel-eval",
-        default=False,
-        action="store_true",
-        help="Eval in parallel (multi-GPU, multi-node).",
-    )
-
-    parser.add_argument(
-        "--no-eval",
-        default=False,
-        action="store_true",
-        help="Training without evaluation.",
-    )
-
-    parser.add_argument(
-        "--lp-mlp",
-        default=False,
-        action="store_true",
-        help="Linear Probe using MLP layer or not.",
-    )
-
-    parser.add_argument(
-        "--lp-freeze",
-        default=False,
-        action="store_true",
-        help="Linear Probe using Freeze CLAP or not",
-    )
-
-    parser.add_argument(
-        "--lp-act",
-        default="None",
-        type=str,
-        help="Options are ['relu','elu','prelu','softmax','sigmoid']",
-    )
-
-    parser.add_argument(
-        "--lp-loss", type=str, default="bce", help="Loss func of Linear Probe."
-    )
-
-    parser.add_argument(
-        "--lp-metrics",
-        type=str,
-        default="map,mauc,acc",
-        help="Metrics of Linear Probe.",
-    )
-
-    parser.add_argument(
-        "--lp-lr", type=float, default=1e-4, help="learning rate of linear probe"
-    )
-    parser.add_argument(
-        "--kappa", type=float, default=0,
-        help="the kappa in the weighted contrastive loss, default is to turn off the weighted contrastive loss"
-    )
-
-    parser.add_argument(
-        "--data-filling",
-        type=str,
-        default="pad",
-        help="type of data filling when the audio length is shorter than the max length."
-             "Can be one of the following: repeat, repeatpad, pad",
-    )
-    parser.add_argument(
-        "--data-truncating",
-        type=str,
-        default="rand_trunc",
-        help="type of data truncation when the audio length is longer than the max length."
-             "Can be one of the following: rand_trunc, fusion",
-    )
-
-    parser.add_argument(
-        "--clap-mlploss",
-        default=False,
-        action="store_true",
-        help="Using MLP loss for CLAP model or not",
-    )
-
-    parser.add_argument(
-        "--wandb-id",
-        type=str,
-        default=None,
-        help="the id of wandb experiment to restore.",
-    )
-
-    parser.add_argument(
-        "--sleep", type=float, default=0, help="sleep n seconds before start training"
-    )
-
-    # variable length processing
-    parser.add_argument(
-        "--enable-fusion",
-        default=False,
-        action="store_true",
-        help="Enable feature funsion for variable-length data",
-    )
-
-    parser.add_argument(
-        "--fusion-type",
-        type=str,
-        default='None',
-        help="Type is among ['channel_map', 'daf_1d','aff_1d','iaff_1d','daf_2d','aff_2d','iaff_2d']",
-    )
-
-    parser.add_argument(
-        "--mixup",
-        default=False,
-        action="store_true",
-        help="Enable mixup in finetuning training.",
-    )
-    parser.add_argument(
-        "--text-augment-selection",
-        type=str,
-        default=None,
-        help="For selecting levels of augmented text. Type is among ['all', 'augment_only', 'none']",
-    )
-    parser.add_argument(
-        "--prefetch-factor",
-        type=int,
-        default=None,
-        help="The prefetch factor for dataloader. Larger value will use more memory and CPU but faster.",
-    )
-
-    args = parser.parse_args()
-
-    # If some params are not passed, we use the default values based on model name.
-    default_params = get_default_params(args.amodel)
-    for name, val in default_params.items():
-        if getattr(args, name) is None:
-            setattr(args, name, val)
-
-    return args

utils/laion_clap/training/scheduler.py

@@ -1,23 +0,0 @@
-import numpy as np
-
-
-def assign_learning_rate(optimizer, new_lr):
-    for param_group in optimizer.param_groups:
-        param_group["lr"] = new_lr
-
-
-def _warmup_lr(base_lr, warmup_length, step):
-    return base_lr * (step + 1) / warmup_length
-
-
-def cosine_lr(optimizer, base_lr, warmup_length, steps):
-    def _lr_adjuster(step):
-        if step < warmup_length:
-            lr = _warmup_lr(base_lr, warmup_length, step)
-        else:
-            e = step - warmup_length
-            es = steps - warmup_length
-            lr = 0.5 * (1 + np.cos(np.pi * e / es)) * base_lr
-        assign_learning_rate(optimizer, lr)
-        return lr
-    return _lr_adjuster

utils/laion_clap/training/train.py

@@ -1,781 +0,0 @@
-import json
-import logging
-import math
-import os
-import time
-from contextlib import suppress
-
-import numpy as np
-import torch
-import torch.nn.functional as F
-
-try:
-    import wandb
-except ImportError:
-    wandb = None
-
-from clap_module import ClipLoss, gather_features
-from .distributed import is_master
-
-
-class AverageMeter(object):
-    """Computes and stores the average and current value"""
-
-    def __init__(self):
-        self.reset()
-
-    def reset(self):
-        self.val = 0
-        self.avg = 0
-        self.sum = 0
-        self.count = 0
-
-    def update(self, val, n=1):
-        self.val = val
-        self.sum += val * n
-        self.count += n
-        self.avg = self.sum / self.count
-
-
-def unwrap_model(model):
-    if hasattr(model, "module"):
-        return model.module
-    else:
-        return model
-
-
-def train_one_epoch(
-        model, data, epoch, optimizer, scaler, scheduler, args, tb_writer=None
-):
-    device = torch.device(args.device)
-    autocast = torch.cuda.amp.autocast if args.precision == "amp" else suppress
-    model.train()
-    loss = ClipLoss(
-        local_loss=args.local_loss,
-        gather_with_grad=args.gather_with_grad,
-        cache_labels=True,
-        rank=args.rank,
-        world_size=args.world_size,
-        use_horovod=args.horovod,
-        mlp_loss=args.clap_mlploss,
-        weight_loss_kappa=args.kappa,
-    )
-
-    dataloader, sampler = data["train"].dataloader, data["train"].sampler
-    if args.distributed and sampler is not None:
-        sampler.set_epoch(epoch)
-    num_batches_per_epoch = dataloader.num_batches
-    sample_digits = math.ceil(math.log(dataloader.num_samples + 1, 10))
-
-    # for toy dataset
-    if args.dataset_type == "toy":
-        dataloader.dataset.generate_queue()
-
-    loss_m = AverageMeter()
-    batch_time_m = AverageMeter()
-    data_time_m = AverageMeter()
-    end = time.time()
-
-    for i, batch in enumerate(dataloader):
-        # logging.info(f"batch {i} of {num_batches_per_epoch}")
-        step = num_batches_per_epoch * epoch + i
-        if isinstance(scheduler, dict):
-            for s in scheduler.values():
-                s(step)
-        else:
-            scheduler(step)
-        audios = batch  # contains mel_spec, wavform, and longer list
-        texts = batch['text']
-        # audios = audios.to(device=device, non_blocking=True)
-        # texts = texts.to(device=device, non_blocking=True)
-
-        data_time_m.update(time.time() - end)
-        if isinstance(optimizer, dict):
-            for o_ in optimizer.values():
-                o_.zero_grad()
-        else:
-            optimizer.zero_grad()
-
-        with autocast():
-            (
-                audio_features,
-                text_features,
-                audio_features_mlp,
-                text_features_mlp,
-                logit_scale_a,
-                logit_scale_t,
-            ) = model(audios, texts, device)
-
-            if args.clap_mlploss:
-                total_loss = loss(
-                    audio_features=audio_features,
-                    text_features=text_features,
-                    logit_scale_a=logit_scale_a,
-                    logit_scale_t=logit_scale_t,
-                    audio_features_mlp=audio_features_mlp,
-                    text_features_mlp=text_features_mlp
-                )
-            else:
-                total_loss = loss(
-                    audio_features=audio_features,
-                    text_features=text_features,
-                    logit_scale_a=logit_scale_a
-                )
-        if isinstance(optimizer, dict):
-            if scaler is not None:
-                scaler.scale(total_loss).backward()
-                for o_ in optimizer.values():
-                    if args.horovod:
-                        o_.synchronize()
-                        scaler.unscale_(o_)
-                        with o_.skip_synchronize():
-                            scaler.step(o_)
-                    else:
-                        scaler.step(o_)
-                scaler.update()
-            else:
-                total_loss.backward()
-                for o_ in optimizer.values():
-                    o_.step()
-        else:
-            if scaler is not None:
-                scaler.scale(total_loss).backward()
-                if args.horovod:
-                    optimizer.synchronize()
-                    scaler.unscale_(optimizer)
-                    with optimizer.skip_synchronize():
-                        scaler.step(optimizer)
-                else:
-                    scaler.step(optimizer)
-                scaler.update()
-            else:
-                total_loss.backward()
-                optimizer.step()
-
-        # Note: we clamp to 4.6052 = ln(100), as in the original paper.
-        with torch.no_grad():
-            unwrap_model(model).logit_scale_a.clamp_(0, math.log(100))
-            if args.clap_mlploss:
-                unwrap_model(model).logit_scale_t.clamp_(0, math.log(100))
-
-        batch_time_m.update(time.time() - end)
-        end = time.time()
-        batch_count = i + 1
-        if is_master(args) and (i % 100 == 0 or batch_count == num_batches_per_epoch):
-            if isinstance(audios, dict):
-                batch_size = len(audios["waveform"])
-            else:
-                batch_size = len(audios)
-            num_samples = batch_count * batch_size * args.world_size
-            samples_per_epoch = dataloader.num_samples
-            percent_complete = 100.0 * batch_count / num_batches_per_epoch
-
-            # NOTE loss is coarsely sampled, just master node and per log update
-            loss_m.update(total_loss.item(), batch_size)
-            logit_scale_scalar_a = logit_scale_a.item()
-            logit_scale_scalar_t = logit_scale_t.item()
-            if isinstance(optimizer, dict):
-                if args.clap_mlploss:
-                    logging.info(
-                        f"Train Epoch: {epoch} [{num_samples:>{sample_digits}}/{samples_per_epoch} ({percent_complete:.0f}%)] "
-                        f"Loss: {loss_m.val:#.5g} ({loss_m.avg:#.4g}) "
-                        f"Data (t): {data_time_m.avg:.3f} "
-                        f"Batch (t): {batch_time_m.avg:.3f} "
-                        f"LR: {[o_.param_groups[0]['lr'] for o_ in optimizer.values()]} "
-                        f"Logit Scale Audio: {logit_scale_scalar_a:.3f}"
-                        f"Logit Scale Text: {logit_scale_scalar_t:.3f}"
-                    )
-                    log_data = {
-                        "loss": loss_m.val,
-                        "data_time": data_time_m.val,
-                        "batch_time": batch_time_m.val,
-                        "scale_audio": logit_scale_scalar_a,
-                        "scale_text": logit_scale_scalar_t,
-                        "lr": [o_.param_groups[0]["lr"] for o_ in optimizer.values()],
-                    }
-                else:
-                    logging.info(
-                        f"Train Epoch: {epoch} [{num_samples:>{sample_digits}}/{samples_per_epoch} ({percent_complete:.0f}%)] "
-                        f"Loss: {loss_m.val:#.5g} ({loss_m.avg:#.4g}) "
-                        f"Data (t): {data_time_m.avg:.3f} "
-                        f"Batch (t): {batch_time_m.avg:.3f} "
-                        f"LR: {[o_.param_groups[0]['lr'] for o_ in optimizer.values()]} "
-                        f"Logit Scale Audio: {logit_scale_scalar_a:.3f}"
-                    )
-                    log_data = {
-                        "loss": loss_m.val,
-                        "data_time": data_time_m.val,
-                        "batch_time": batch_time_m.val,
-                        "scale_audio": logit_scale_scalar_a,
-                        "lr": [o_.param_groups[0]["lr"] for o_ in optimizer.values()],
-                    }
-
-            else:
-                if args.clap_mlploss:
-                    logging.info(
-                        f"Train Epoch: {epoch} [{num_samples:>{sample_digits}}/{samples_per_epoch} ({percent_complete:.0f}%)] "
-                        f"Loss: {loss_m.val:#.5g} ({loss_m.avg:#.4g}) "
-                        f"Data (t): {data_time_m.avg:.3f} "
-                        f"Batch (t): {batch_time_m.avg:.3f} "
-                        f"LR: {optimizer.param_groups[0]['lr']:5f} "
-                        f"Logit Scale Audio: {logit_scale_scalar_a:.3f}"
-                        f"Logit Scale Text: {logit_scale_scalar_t:.3f}"
-                    )
-
-                    # Save train loss / etc. Using non avg meter values as loggers have their own smoothing
-                    log_data = {
-                        "loss": loss_m.val,
-                        "data_time": data_time_m.val,
-                        "batch_time": batch_time_m.val,
-                        "scale_audio": logit_scale_scalar_a,
-                        "scale_text": logit_scale_scalar_t,
-                        "lr": optimizer.param_groups[0]["lr"],
-                    }
-                else:
-                    logging.info(
-                        f"Train Epoch: {epoch} [{num_samples:>{sample_digits}}/{samples_per_epoch} ({percent_complete:.0f}%)] "
-                        f"Loss: {loss_m.val:#.5g} ({loss_m.avg:#.4g}) "
-                        f"Data (t): {data_time_m.avg:.3f} "
-                        f"Batch (t): {batch_time_m.avg:.3f} "
-                        f"LR: {optimizer.param_groups[0]['lr']:5f} "
-                        f"Logit Scale Audio: {logit_scale_scalar_a:.3f}"
-                    )
-
-                    # Save train loss / etc. Using non avg meter values as loggers have their own smoothing
-                    log_data = {
-                        "loss": loss_m.val,
-                        "data_time": data_time_m.val,
-                        "batch_time": batch_time_m.val,
-                        "scale_audio": logit_scale_scalar_a,
-                        "lr": optimizer.param_groups[0]["lr"],
-                    }
-            for name, val in log_data.items():
-                name = "train/" + name
-                if tb_writer is not None:
-                    tb_writer.add_scalar(name, val, step)
-                if args.wandb:
-                    assert wandb is not None, "Please install wandb."
-                    wandb.log({name: val, "step": step})
-
-            # resetting batch / data time meters per log window
-            batch_time_m.reset()
-            data_time_m.reset()
-    # end for
-
-
-def evaluate(model, data, epoch, args, tb_writer=None):
-    metrics = {}
-    if not args.parallel_eval:
-        if not is_master(args):
-            return metrics
-    device = torch.device(args.device)
-    model.eval()
-
-    # CHANGE
-    # zero_shot_metrics = zero_shot_eval(model, data, epoch, args)
-    # metrics.update(zero_shot_metrics)
-    if is_master(args):
-        print('Evaluating...')
-    autocast = torch.cuda.amp.autocast if args.precision == "amp" else suppress
-    if args.val_dataset_names == ['Clotho', 'audiocaps']:
-        # if only clotho and audiocaps are used, then we will use a different evaluation function.
-        # This is because in the Clotho and audiocaps valid and test set, there are 5 text for 1 audio.
-        if args.parallel_eval:
-            # (yusong): just a hack here. Don't use parallel eval when evaluating only clotho and audiocaps.
-            raise NotImplementedError("Parallel evaluation not supported for eval only Clotho and audiocaps.")
-        val_metrics_per_dataset = evaluate_clotho_audiocaps(model, data, epoch, args, autocast, device, tb_writer)
-        for m in val_metrics_per_dataset.values():
-            metrics.update(m)
-        if "epoch" not in metrics.keys():
-            metrics.update({"epoch": epoch})
-        metrics = select_top_metric_clotho_audiocaps(metrics, val_metrics_per_dataset, args)
-    elif "val" in data and (
-            args.val_frequency
-            and ((epoch % args.val_frequency) == 0 or epoch == args.epochs)
-    ):
-        dataloader = data["val"].dataloader
-        num_samples = 0
-        samples_per_val = dataloader.num_samples
-
-        # FIXME this does not scale past small eval datasets
-        # all_audio_features @ all_text_features will blow up memory and compute very quickly
-        eval_info = {}
-        if args.clap_mlploss:
-            eval_info["all"] = {
-                "cumulative_loss": 0.0,
-                "num_samples": 0,
-                "all_audio_features": [],
-                "all_text_features": [],
-                "all_audio_features_mlp": [],
-                "all_text_features_mlp": []
-            }  # cumulative_loss = 0.0
-        else:
-            eval_info["all"] = {
-                "cumulative_loss": 0.0,
-                "num_samples": 0,
-                "all_audio_features": [],
-                "all_text_features": []
-            }  # cumu
-        # all_audio_features, all_text_features, all_audio_features_mlp, all_text_features_mlp = [], [], [], []
-        with torch.no_grad():
-            for i, batch in enumerate(dataloader):
-                audios = batch  # contains mel_spec, wavform, and longer list
-                texts = batch['text']
-                # audios = audios.to(device=device, non_blocking=True)
-
-                all_names = list(set(["-".join(b.split("/")[-3:-1]) for b in batch['__url__']]))
-                for name in all_names:
-                    if name not in eval_info.keys():
-                        if args.clap_mlploss:
-                            eval_info[name] = {
-                                "cumulative_loss": 0.0,
-                                "num_samples": 0,
-                                "all_audio_features": [],
-                                "all_text_features": [],
-                                "all_audio_features_mlp": [],
-                                "all_text_features_mlp": [],
-                            }
-                        else:
-                            eval_info[name] = {
-                                "cumulative_loss": 0.0,
-                                "num_samples": 0,
-                                "all_audio_features": [],
-                                "all_text_features": []
-                            }
-                with autocast():
-                    (
-                        audio_features,
-                        text_features,
-                        audio_features_mlp,
-                        text_features_mlp,
-                        logit_scale_a,
-                        logit_scale_t,
-                    ) = model(audios, texts, device)
-
-                    if args.parallel_eval:
-                        # multi-GPU eval
-                        if args.clap_mlploss:
-                            (
-                                audio_features,
-                                text_features,
-                                audio_features_mlp,
-                                text_features_mlp,
-                            ) = gather_features(
-                                audio_features=audio_features,
-                                text_features=text_features,
-                                audio_features_mlp=audio_features_mlp,
-                                text_features_mlp=text_features_mlp,
-                                local_loss=False,
-                                gather_with_grad=False,
-                                rank=args.rank,
-                                world_size=args.world_size,
-                                use_horovod=args.horovod,
-                                mlp_loss=args.clap_mlploss
-                            )
-                        else:
-                            (
-                                audio_features,
-                                text_features,
-                            ) = gather_features(
-                                audio_features=audio_features,
-                                text_features=text_features,
-                                local_loss=False,
-                                gather_with_grad=False,
-                                rank=args.rank,
-                                world_size=args.world_size,
-                                use_horovod=args.horovod,
-                                mlp_loss=args.clap_mlploss
-                            )
-
-                    if is_master(args):
-                        num_samples += audio_features.shape[0]
-                        for n in [*all_names, "all"]:
-                            if n == "all":
-                                eval_info[n]["all_audio_features"].append(
-                                    audio_features.cpu()
-                                )
-                                eval_info[n]["all_text_features"].append(
-                                    text_features.cpu()
-                                )
-                                if args.clap_mlploss:
-                                    eval_info[n]["all_audio_features_mlp"].append(
-                                        audio_features_mlp.cpu()
-                                    )
-                                    eval_info[n]["all_text_features_mlp"].append(
-                                        text_features_mlp.cpu()
-                                    )
-                            else:
-                                idx = np.where(
-                                    np.array(
-                                        ["-".join(b.split("/")[-3:-1]) for b in batch['__url__']]
-                                    )
-                                    == n
-                                )[0]
-                                eval_info[n]["all_audio_features"].append(
-                                    audio_features.cpu().index_select(
-                                        0, torch.tensor(idx).long()
-                                    )
-                                )
-                                eval_info[n]["all_text_features"].append(
-                                    text_features.cpu().index_select(
-                                        0, torch.tensor(idx).long()
-                                    )
-                                )
-                                if args.clap_mlploss:
-                                    eval_info[n]["all_audio_features_mlp"].append(
-                                        audio_features_mlp.cpu().index_select(
-                                            0, torch.tensor(idx).long()
-                                        )
-                                    )
-                                    eval_info[n]["all_text_features_mlp"].append(
-                                        text_features_mlp.cpu().index_select(
-                                            0, torch.tensor(idx).long()
-                                        )
-                                    )
-                        #  print(f'eval step {i}') #  (yusong): for debug
-
-                # cumulative_loss += total_loss * batch_size
-                # num_samples += batch_size
-                if is_master(args) and (i % 100) == 0:  # and i != 0:
-                    logging.info(
-                        f"Eval Epoch: {epoch} [{num_samples} / {samples_per_val}]"
-                    )
-            if is_master(args):
-                val_metrics_per_dataset = {}
-                for n in eval_info.keys():
-                    if args.clap_mlploss:
-                        metrics_single_dataset = get_metrics(
-                            audio_features=torch.cat(eval_info[n]["all_audio_features"]),
-                            text_features=torch.cat(eval_info[n]["all_text_features"]),
-                            logit_scale_a=logit_scale_a.cpu(),
-                            audio_features_mlp=torch.cat(
-                                eval_info[n]["all_audio_features_mlp"]
-                            ),
-                            text_features_mlp=torch.cat(eval_info[n]["all_text_features_mlp"]),
-                            logit_scale_t=logit_scale_t.cpu(),
-                            mlp_loss=args.clap_mlploss
-                        )
-                    else:
-                        metrics_single_dataset = get_metrics(
-                            audio_features=torch.cat(eval_info[n]["all_audio_features"]),
-                            text_features=torch.cat(eval_info[n]["all_text_features"]),
-                            logit_scale_a=logit_scale_a.cpu(),
-                            mlp_loss=args.clap_mlploss
-                        )
-                    val_metrics_per_dataset[n] = {
-                        n + "/" + k: v for k, v in metrics_single_dataset.items()
-                    }
-                    metrics.update(val_metrics_per_dataset[n])
-                    if "epoch" not in metrics.keys():
-                        metrics.update({"epoch": epoch})
-    if is_master(args):
-        if not metrics:
-            return metrics
-
-        logging.info(
-            f"Eval Epoch: {epoch} "
-            + "\n".join(
-                [
-                    "\t".join([f"{k}: {round(v, 4):.4f}" for k, v in m.items()])
-                    for m in val_metrics_per_dataset.values()
-                ]
-            )
-        )
-
-        if args.save_logs:
-            for name, val in metrics.items():
-                if tb_writer is not None:
-                    tb_writer.add_scalar(f"val/{name}", val, epoch)
-
-            with open(os.path.join(args.checkpoint_path, "results.jsonl"), "a+") as f:
-                f.write(json.dumps(metrics))
-                f.write("\n")
-
-        if args.wandb:
-            assert wandb is not None, "Please install wandb."
-            for name, val in metrics.items():
-                wandb.log({f"val/{name}": val, "epoch": epoch})
-
-        return metrics
-    else:
-        return metrics
-
-
-def get_metrics(
-        audio_features,
-        text_features,
-        logit_scale_a,
-        audio_features_mlp=None,
-        text_features_mlp=None,
-        logit_scale_t=None,
-        mlp_loss=False
-):
-    metrics = {}
-    if mlp_loss:
-        # Set up audio to text & text to audio similary matrice
-        a_logits_per_audio = (
-            (logit_scale_a * audio_features @ text_features_mlp.t()).detach().cpu()
-        )
-        a_logits_per_text = a_logits_per_audio.t().detach().cpu()
-        t_logits_per_audio = (
-            (logit_scale_t * audio_features_mlp @ text_features.t()).detach().cpu()
-        )
-        t_logits_per_text = t_logits_per_audio.t().detach().cpu()
-
-        labels = torch.arange(audio_features.shape[0]).long()
-        # Change the loss from two terms into four terms with 2x2 combined CE loss
-        total_loss = (
-                             F.cross_entropy(a_logits_per_audio, labels)
-                             + F.cross_entropy(a_logits_per_text, labels)
-                             + F.cross_entropy(t_logits_per_audio, labels)
-                             + F.cross_entropy(t_logits_per_text, labels)
-                     ) / 4
-
-        metrics[f"cumulative_loss"] = total_loss.item()
-        metrics[f"num_samples"] = audio_features.shape[0]
-
-        logits = {
-            "audio_to_text": (a_logits_per_audio + t_logits_per_audio) / 2,
-            "text_to_audio": (a_logits_per_text + t_logits_per_text) / 2,
-        }
-        ground_truth = torch.arange(len(text_features)).view(-1, 1)
-
-    else:
-        # print("text_features", text_features)
-        # print("text_features.shape", text_features.shape)
-        logits_per_audio = (logit_scale_a * audio_features @ text_features.t()).detach().cpu()
-        logits_per_text = logits_per_audio.t().detach().cpu()
-
-        labels = torch.arange(audio_features.shape[0]).long()
-        # Change the loss from two terms into four terms with 2x2 combined CE loss
-        total_loss = (
-                             F.cross_entropy(logits_per_audio, labels)
-                             + F.cross_entropy(logits_per_text, labels)
-                     ) / 2
-
-        metrics[f"cumulative_loss"] = total_loss.item()
-        metrics[f"num_samples"] = audio_features.shape[0]
-
-        logits = {"audio_to_text": logits_per_audio, "text_to_audio": logits_per_text}
-
-        ground_truth = torch.arange(len(text_features)).view(-1, 1)
-
-    for name, logit in logits.items():
-        ranking = torch.argsort(logit, descending=True)
-        preds = torch.where(ranking == ground_truth)[1]  # (yusong) this line is slow because it uses single thread
-        preds = preds.detach().cpu().numpy()
-        metrics[f"{name}_mean_rank"] = preds.mean() + 1
-        metrics[f"{name}_median_rank"] = np.floor(np.median(preds)) + 1
-        for k in [1, 5, 10]:
-            metrics[f"{name}_R@{k}"] = np.mean(preds < k)
-        # map@10
-        metrics[f"{name}_mAP@10"] = np.mean(np.where(preds < 10, 1 / (preds + 1), 0.0))
-
-    return metrics
-
-
-def evaluate_clotho_audiocaps(
-        model, data, epoch, args, autocast, device, tb_writer=None
-):
-    """
-    Adapted from https://github.com/XinhaoMei/audio-text_retrieval/blob/main/tools/utils.py.
-    1. for text-to-audio retrieval, do 5 times and average the results
-    2. for R@1, R@5, R@10 in audio-to-text retrieval, take the best rank among 5 text
-    3. for map@10 in audio-to-text retrieval:
-        3.1: sort the rank of 5 text
-        3.2: exclude the rank >=10 (0-index)
-        3.3: compute the map regarding the remaining ranks: np.mean(np.arange(1, len(ranks)+1) / ranks).
-        (3.3) That is, take the top ranks of 5 text that is < 10, and assign the descending number as ground truth.
-        (3.3) E.g.: the ground truth of first rank of the 5 text should be 1, the second rank should be 2, etc.
-    """
-    # TODO: (yusong) only support single GPU evaluation and only support non-mlp case for now.
-    dataloader = data["val"].dataloader
-    with torch.no_grad():
-        eval_info = {}
-        for i, batch in enumerate(dataloader):
-            audios = batch  # contains mel_spec, wavform, and longer list
-
-            # each item in the list has 5 texts
-            if args.tmodel == "transformer":
-                from clap_module import tokenize
-                texts = [tokenize(t) for t in batch['full_text']]
-                texts = torch.cat(texts)
-            else:
-                from .data import tokenizer
-                texts = [tokenizer(t, tmodel=args.tmodel) for t in batch['full_text']]  # 5 texts for each audio
-                texts = {k: torch.cat([t[k] for t in texts]) for k in texts[0].keys()}  # 5 x batch
-
-            # audios = audios.to(device=device, non_blocking=True)
-
-            # batch['__url__'] contains the path to the data tar this sample is from
-            # So, b.split("/")[-3:-1] will get you '<dataset_name>-<dataset-split>'
-            all_names = list(set(["-".join(b.split("/")[-3:-1]) for b in batch['__url__']]))
-            for name in all_names:
-                if name not in eval_info.keys():
-                    # we will not use mlp outputs even if args.clap_mlploss=True
-                    eval_info[name] = {
-                        "cumulative_loss": 0.0,
-                        "num_samples": 0,
-                        "all_audio_features": [],
-                        "all_text_features": []
-                    }
-            with autocast():
-                audio_features = model(audios, None, device)
-                text_features = model(None, texts, device)
-                audio_features = F.normalize(audio_features, dim=-1)
-                text_features = F.normalize(text_features, dim=-1)
-
-                all_names = list(set(["-".join(b.split("/")[-3:-1]) for b in batch['__url__']]))
-                for n in all_names:
-                    idx = np.where(
-                        np.array(
-                            ["-".join(b.split("/")[-3:-1]) for b in batch['__url__']]
-                        )
-                        == n
-                    )[0]
-                    eval_info[n]["all_audio_features"].append(
-                        audio_features.cpu().index_select(
-                            0, torch.tensor(idx).long()
-                        )
-                    )
-                    # (yusong) please double-check. This is for selecting 5 text features at once.
-                    # because idx is a list of indices in size of num_samples,
-                    # and text_features is a tensor of size (5*num_samples, dim)
-                    # so we need to select 5 consecutive indices at once for a single index in idx.
-                    eval_info[n]["all_text_features"].append(
-                        text_features.cpu().reshape([-1, 5, text_features.shape[1]]).index_select(
-                            0, torch.tensor(idx).long()
-                        ).reshape([-1, text_features.shape[1]])
-                    )
-
-        val_metrics_all = {}
-
-        for n in eval_info.keys():
-            logit_scale_a, logit_scale_t = model(None, None, device)
-            logit_scale_a = logit_scale_a.cpu()
-
-            audio_features = torch.cat(eval_info[n]["all_audio_features"], dim=0)
-            text_features = torch.cat(eval_info[n]["all_text_features"], dim=0)
-
-            logits_per_audio = (logit_scale_a * audio_features @ text_features.t()).detach().cpu()
-            logits_per_text = logits_per_audio.t().detach().cpu()
-
-            # logits_per_audio shape: [num_samples, num_samples*5]
-            # logits_per_text shape: [num_samples*5, num_samples]
-
-            logging.info(f"dataset {n}, logits_per_audio shape: {logits_per_audio.shape}, "
-                         f"logits_per_text shape: {logits_per_text.shape}")
-
-            metrics = {}
-            num_samples = audio_features.shape[0]
-            metrics[f"num_samples"] = num_samples
-
-            # (yusong) the following code is very important, please double-check:
-            # logits_per_audio.reshape(num_samples, num_samples, 5)[:, :, d]
-            # logits_per_text.reshape(num_samples, 5, num_samples)[:, d, :]
-            # Those two are retrieving one of the 5 text for each audio.
-            labels = torch.arange(audio_features.shape[0]).long()
-            audio_to_text_loss = [
-                F.cross_entropy(
-                    logits_per_audio.reshape(num_samples, num_samples, 5)[:, :, d], labels) for d in range(5)
-            ]
-            text_to_audio_loss = [
-                F.cross_entropy(
-                    logits_per_text.reshape(num_samples, 5, num_samples)[:, d, :], labels) for d in range(5)
-            ]
-            total_loss = (
-                                 np.mean(audio_to_text_loss) + np.mean(text_to_audio_loss)
-                         ) / 2
-
-            metrics[f"cumulative_loss"] = total_loss.item()
-
-            # text to audio: do 5 times
-            pred_text = []
-            for d in range(5):
-                logit = logits_per_text.reshape(num_samples, 5, num_samples)[:, d, :]
-                ground_truth = torch.arange(len(logit)).view(-1, 1)
-                ranking = torch.argsort(logit, descending=True)  # [num_samples, num_samples]
-                preds = torch.where(ranking == ground_truth)[1]
-                pred_text.append(preds.detach().cpu().numpy())
-            pred_text_concat = np.concatenate(pred_text, axis=0)  # [5*num_samples]
-            metrics[f"text_to_audio_mean_rank"] = pred_text_concat.mean() + 1
-            metrics[f"text_to_audio_median_rank"] = np.floor(np.median(pred_text_concat)) + 1
-            for k in [1, 5, 10]:
-                metrics[f"text_to_audio_R@{k}"] = np.mean(pred_text_concat < k)
-            # map@10
-            metrics[f"text_to_audio_mAP@10"] = np.mean(np.where(pred_text_concat < 10, 1 / (pred_text_concat + 1), 0.0))
-
-            # audio to text: take the best result
-            # for audio to text map 10, sort and assign descending ground truth.
-            # see https://github.com/XinhaoMei/audio-text_retrieval/blob/main/tools/utils.py#L103
-            # map@10
-            map_all = []
-            pred_audio_all = []
-            for d in range(num_samples):
-                # logits_per_audio: [num_samples, num_samples*5]
-                logit_single = logits_per_audio[d, :]  # [5*num_samples]
-                # Ground-truth index: [d*5, d*5+1, d*5+2, d*5+3, d*5+4]
-                ranking = torch.argsort(logit_single, descending=True)  # [5*num_samples]
-                # ranking: the index of first match, second match, ...
-                ground_truth = torch.arange(d * 5, d * 5 + 5)[None]
-                all_pred = torch.where(torch.stack([ranking] * 5) == ground_truth.view(-1, 1))[1]
-                min_pred = torch.min(all_pred)
-                pred_audio_all.append(min_pred.detach().cpu().numpy())
-                all_pred_filter = all_pred[all_pred < 10].detach().cpu().numpy()
-                # /5 because we have 5 text, so it means for the text rank >=10 we count as 0.
-                map_single = np.sum((np.arange(1, len(all_pred_filter) + 1) / (all_pred_filter + 1))) / 5
-                map_all.append(map_single)
-            metrics[f"audio_to_text_mAP@10"] = np.mean(map_all)
-            for k in [1, 5, 10]:
-                metrics[f"audio_to_text_R@{k}"] = np.mean(np.array(pred_audio_all) < k)
-
-            val_metrics_all[n] = {
-                n + "/" + k: v for k, v in metrics.items()
-            }
-    return val_metrics_all
-
-
-def calculate_selection_performance_clotho_audiocaps(val_metrics_per_dataset):
-    """
-    Calculate performance for Clotho+AudioCaps for model selection.
-    """
-    selection_performance_all = []
-    for n in val_metrics_per_dataset.keys():
-        selection_performance = (val_metrics_per_dataset[n][f"{n}/audio_to_text_mAP@10"] +
-                                 val_metrics_per_dataset[n][f"{n}/text_to_audio_mAP@10"]) / 2
-        selection_performance_all.append(selection_performance)
-    return np.mean(selection_performance_all)
-
-
-def select_top_metric_clotho_audiocaps(metrics, val_metrics_per_dataset, args):
-    # val_metrics_per_dataset: dict, key: dataset name, value: dict, key: metric name, value: metric value
-    # metrics: dict, key: metric name, value: metric value
-    # Hack: use args to save the top performance
-    if not hasattr(args, "top_selection_performance"):
-        selection_performance = calculate_selection_performance_clotho_audiocaps(val_metrics_per_dataset)
-        # TODO: write the if and else together
-        metric_update = {}
-        for n in val_metrics_per_dataset.keys():
-            for k in val_metrics_per_dataset[n].keys():
-                metric_update[k.split('/')[0] + '-top' + '/' + k.split('/')[1]] = val_metrics_per_dataset[n][k]
-        metric_update['top_selection_performance'] = selection_performance
-        metric_update['top-selection-epoch'] = metrics['epoch']
-        metrics.update(metric_update)
-        args.top_metric = metric_update
-        args.top_selection_performance = selection_performance
-    else:
-        selection_performance_new = calculate_selection_performance_clotho_audiocaps(val_metrics_per_dataset)
-        selection_performance_old = args.top_selection_performance
-        if selection_performance_new > selection_performance_old:
-            metric_update = {}
-            for n in val_metrics_per_dataset.keys():
-                for k in val_metrics_per_dataset[n].keys():
-                    metric_update[k.split('/')[0] + '-top' + '/' + k.split('/')[1]] = val_metrics_per_dataset[n][k]
-            metric_update['top_selection_performance'] = selection_performance_new
-            metric_update['top-selection-epoch'] = metrics['epoch']
-            metrics.update(metric_update)
-            args.top_metric = metric_update
-            args.top_selection_performance = selection_performance_new
-        else:
-            metrics.update(args.top_metric)
-    return metrics

utils/laion_clap/training/zero_shot.py

@@ -1,90 +0,0 @@
-# NOTE: This script is currently not supported for CLAP.
-import logging
-from contextlib import suppress
-
-import torch
-import torch.nn.functional as F
-from tqdm import tqdm
-
-from clap_module import tokenize
-from .imagenet_zeroshot_data import imagenet_classnames, openai_imagenet_template
-
-
-def zero_shot_classifier(model, classnames, templates, args):
-    with torch.no_grad():
-        zeroshot_weights = []
-        for classname in tqdm(classnames):
-            texts = [template(classname) for template in templates]  # format with class
-            texts = tokenize(texts).to(args.device)  # tokenize
-            if args.distributed and not args.horovod:
-                class_embeddings = model.module.encode_text(texts)
-            else:
-                class_embeddings = model.encode_text(texts)
-            class_embedding = F.normalize(class_embeddings, dim=-1).mean(dim=0)
-            class_embedding /= class_embedding.norm()
-            zeroshot_weights.append(class_embedding)
-        zeroshot_weights = torch.stack(zeroshot_weights, dim=1).to(args.device)
-    return zeroshot_weights
-
-
-def accuracy(output, target, topk=(1,)):
-    pred = output.topk(max(topk), 1, True, True)[1].t()
-    correct = pred.eq(target.view(1, -1).expand_as(pred))
-    return [float(correct[:k].reshape(-1).float().sum(0, keepdim=True).cpu().numpy()) for k in topk]
-
-
-def run(model, classifier, dataloader, args):
-    autocast = torch.cuda.amp.autocast if args.precision == 'amp' else suppress
-    with torch.no_grad():
-        top1, top5, n = 0., 0., 0.
-        for images, target in tqdm(dataloader, unit_scale=args.batch_size):
-            images = images.to(args.device)
-            target = target.to(args.device)
-
-            with autocast():
-                # predict
-                if args.distributed and not args.horovod:
-                    image_features = model.module.encode_image(images)
-                else:
-                    image_features = model.encode_image(images)
-                image_features = F.normalize(image_features, dim=-1)
-                logits = 100. * image_features @ classifier
-
-            # measure accuracy
-            acc1, acc5 = accuracy(logits, target, topk=(1, 5))
-            top1 += acc1
-            top5 += acc5
-            n += images.size(0)
-
-    top1 = (top1 / n)
-    top5 = (top5 / n)
-    return top1, top5
-
-
-def zero_shot_eval(model, data, epoch, args):
-    if 'imagenet-val' not in data and 'imagenet-v2' not in data:
-        return {}
-    if args.zeroshot_frequency == 0:
-        return {}
-    if (epoch % args.zeroshot_frequency) != 0 and epoch != args.epochs:
-        return {}
-
-    logging.info('Starting zero-shot imagenet.')
-
-    logging.info('Building zero-shot classifier')
-    classifier = zero_shot_classifier(model, imagenet_classnames, openai_imagenet_template, args)
-
-    logging.info('Using classifier')
-    results = {}
-    if 'imagenet-val' in data:
-        top1, top5 = run(model, classifier, data['imagenet-val'].dataloader, args)
-        results['imagenet-zeroshot-val-top1'] = top1
-        results['imagenet-zeroshot-val-top5'] = top5
-    if 'imagenet-v2' in data:
-        top1, top5 = run(model, classifier, data['imagenet-v2'].dataloader, args)
-        results['imagenetv2-zeroshot-val-top1'] = top1
-        results['imagenetv2-zeroshot-val-top5'] = top5
-
-    logging.info('Finished zero-shot imagenet.')
-
-    return results

utils/laion_clap/unit_test.py

@@ -1,75 +0,0 @@
-"""
-Contrastive Language-Audio Pretraining Model from LAION
---------------------------------------------------------
-Paper: https://arxiv.org/abs/2211.06687
-Authors (equal contributions): Ke Chen, Yusong Wu, Tianyu Zhang, Yuchen Hui
-Support: LAION
-"""
-
-import numpy as np
-import librosa
-import torch
-import laion_clap
-
-# quantization
-def int16_to_float32(x):
-    return (x / 32767.0).astype(np.float32)
-
-
-def float32_to_int16(x):
-    x = np.clip(x, a_min=-1., a_max=1.)
-    return (x * 32767.).astype(np.int16)
-
-model = laion_clap.CLAP_Module(enable_fusion=False)
-model.load_ckpt()
-
-# Directly get audio embeddings from audio files
-audio_file = [
-    '/home/la/kechen/Research/KE_CLAP/ckpt/test_clap_short.wav',
-    '/home/la/kechen/Research/KE_CLAP/ckpt/test_clap_long.wav'
-]
-audio_embed = model.get_audio_embedding_from_filelist(x = audio_file, use_tensor=False)
-print(audio_embed[:,-20:])
-print(audio_embed.shape)
-
-# Get audio embeddings from audio data
-audio_data, _ = librosa.load('/home/la/kechen/Research/KE_CLAP/ckpt/test_clap_short.wav', sr=48000) # sample rate should be 48000
-audio_data = audio_data.reshape(1, -1) # Make it (1,T) or (N,T)
-audio_embed = model.get_audio_embedding_from_data(x = audio_data, use_tensor=False)
-print(audio_embed[:,-20:])
-print(audio_embed.shape)
-
-# Directly get audio embeddings from audio files, but return torch tensor
-audio_file = [
-    '/home/la/kechen/Research/KE_CLAP/ckpt/test_clap_short.wav',
-    '/home/la/kechen/Research/KE_CLAP/ckpt/test_clap_long.wav'
-]
-audio_embed = model.get_audio_embedding_from_filelist(x = audio_file, use_tensor=True)
-print(audio_embed[:,-20:])
-print(audio_embed.shape)
-
-# Get audio embeddings from audio data
-audio_data, _ = librosa.load('/home/la/kechen/Research/KE_CLAP/ckpt/test_clap_short.wav', sr=48000) # sample rate should be 48000
-audio_data = audio_data.reshape(1, -1) # Make it (1,T) or (N,T)
-audio_data = torch.from_numpy(int16_to_float32(float32_to_int16(audio_data))).float() # quantize before send it in to the model
-audio_embed = model.get_audio_embedding_from_data(x = audio_data, use_tensor=True)
-print(audio_embed[:,-20:])
-print(audio_embed.shape)
-
-# Get text embedings from texts:
-text_data = ["I love the contrastive learning", "I love the pretrain model"] 
-text_embed = model.get_text_embedding(text_data)
-print(text_embed)
-print(text_embed.shape)
-
-# Get text embedings from texts, but return torch tensor:
-text_data = ["I love the contrastive learning", "I love the pretrain model"] 
-text_embed = model.get_text_embedding(text_data, use_tensor=True)
-print(text_embed)
-print(text_embed.shape)
-
-
-
-
-
-

utils/log.py

@@ -1,976 +0,0 @@
-import os
-import torch 
-import time
-import numpy as np
-#import torchaudio
-import plotly.express as px
-import soundfile as sf
-#import plotly.graph_objects as go
-import pandas as pd
-import plotly
-import plotly.graph_objects as go
-import matplotlib.pyplot as plt
-
-#from notebooks.shared_log_filter import frequencies
-
-"""
-Logging related functions that I wrote for my own use
-This is quite a mess, but I'm too lazy to clean it up
-"""
-#from src.CQT_nsgt import CQT_cpx
-
-def do_stft(noisy, clean=None, win_size=2048, hop_size=512, device="cpu", DC=True):
-    """
-        applies the stft, this an ugly old function, but I'm using it for logging and I'm to lazy to modify it
-    """
-    
-    #window_fn = tf.signal.hamming_window
-
-    #win_size=args.stft.win_size
-    #hop_size=args.stft.hop_size
-    window=torch.hamming_window(window_length=win_size)
-    window=window.to(noisy.device)
-    noisy=torch.cat((noisy, torch.zeros(noisy.shape[0],win_size).to(noisy.device)),1)
-    stft_signal_noisy=torch.stft(noisy, win_size, hop_length=hop_size,window=window,center=False,return_complex=False)
-    stft_signal_noisy=stft_signal_noisy.permute(0,3,2,1)
-    #stft_signal_noisy=tf.signal.stft(noisy,frame_length=win_size, window_fn=window_fn, frame_step=hop_size)
-    #stft_noisy_stacked=tf.stack( values=[tf.math.real(stft_signal_noisy), tf.math.imag(stft_signal_noisy)], axis=-1)
-    
-    if clean!=None:
-
-       # stft_signal_clean=tf.signal.stft(clean,frame_length=win_size, window_fn=window_fn, frame_step=hop_size)
-        clean=torch.cat((clean, torch.zeros(clean.shape[0],win_size).to(device)),1)
-        stft_signal_clean=torch.stft(clean, win_size, hop_length=hop_size,window=window, center=False,return_complex=False)
-        stft_signal_clean=stft_signal_clean.permute(0,3,2,1)
-        #stft_clean_stacked=tf.stack( values=[tf.math.real(stft_signal_clean), tf.math.imag(stft_signal_clean)], axis=-1)
-
-
-        if DC:
-            return stft_signal_noisy, stft_signal_clean
-        else:
-            return stft_signal_noisy[...,1:], stft_signal_clean[...,1:]
-    else:
-
-        if DC:
-            return stft_signal_noisy
-        else:
-            return stft_signal_noisy[...,1:]
-
-def plot_norms(path, normsscores, normsguides, t, name):
-    values=t.cpu().numpy()
-     
-    df=pd.DataFrame.from_dict(
-                {"sigma": values[0:-1], "score": normsscores.cpu().numpy(), "guidance": normsguides.cpu().numpy()}
-                )
-    fig= px.line(df, x="sigma", y=["score", "guidance"],log_x=True,  log_y=True, markers=True)
-
-    path_to_plotly_html = path+"/"+name+".html"
-    
-    fig.write_html(path_to_plotly_html, auto_play = False)
-    return fig
-    
-
-def plot_spectral_analysis_sampling( avgspecNF,avgspecDEN, ts, fs=22050, nfft=1024):
-    T,F=avgspecNF.shape
-    f=torch.arange(0, F)*fs/nfft
-    f=f.unsqueeze(1).repeat(1,T).view(-1)
-    avgspecNF=avgspecNF.permute(1,0).reshape(-1)
-    avgspecDEN=avgspecDEN.permute(1,0).reshape(-1)
-    ts=ts.cpu().numpy().tolist()
-    ts=513*ts
-    #ts=np.repeat(ts,513)
-    print(f.shape, avgspecDEN.shape, avgspecNF.shape, len(ts))
-    df=pd.DataFrame.from_dict(
-        {"f": f,  "noisy": avgspecNF, "denoised":  avgspecDEN, "sigma":ts
-        }
-        )
-
-    fig= px.line(df, x="f", y=["noisy", "denoised"],  animation_frame="sigma", log_x=False,log_y=False,  markers=False)
-    return fig
-
-def plot_spectral_analysis(avgspecY, avgspecNF,avgspecDEN, ts, fs=22050, nfft=1024):
-    T,F=avgspecNF.shape
-    f=torch.arange(0, F)*fs/nfft
-    f=f.unsqueeze(1).repeat(1,T).view(-1)
-    avgspecY=avgspecY.squeeze(0).unsqueeze(1).repeat(1,T).view(-1)
-    avgspecNF=avgspecNF.permute(1,0).reshape(-1)
-    avgspecDEN=avgspecDEN.permute(1,0).reshape(-1)
-    ts=ts.cpu().numpy().tolist()
-    ts=513*ts
-    #ts=np.repeat(ts,513)
-    print(f.shape, avgspecY.shape, avgspecNF.shape, len(ts))
-    df=pd.DataFrame.from_dict(
-        {"f": f, "y": avgspecY, "noisy": avgspecNF, "denoised":  avgspecDEN, "sigma":ts
-        }
-        )
-
-    fig= px.line(df, x="f", y=["y", "noisy", "denoised"],  animation_frame="sigma", log_x=False,log_y=False,  markers=False)
-    return fig
-
-def plot_loss_by_sigma_test_snr(average_snr, average_snr_out, t):
-    #write a fancy plot to log in wandb
-    values=np.array(t)
-    df=pd.DataFrame.from_dict(
-                {"sigma": values, "SNR": average_snr, "SNR_denoised": average_snr_out
-                }
-                )
-
-    fig= px.line(df, x="sigma", y=["SNR", "SNR_denoised"],log_x=True,  markers=True)
-    return fig
-    
-
-
-
-# Create and style traces
-
-def plot_loss_by_sigma(sigma_means, sigma_stds, sigma_bins, log_scale=True):
-    df=pd.DataFrame.from_dict(
-                {"sigma": sigma_bins, "loss": sigma_means, "std": sigma_stds
-                }
-                )
-
-    fig= error_line('bar', data_frame=df, x="sigma", y="loss", error_y="std", log_x=log_scale,  markers=True, range_y=[0, 2])
-    
-    return fig
-
-def plot_loss_by_sigma_and_freq(sigma_freq_means, sigma_freq_stds, sigma_bins, freq_bins):
-    df=pd.DataFrame.from_dict(
-                {"sigma": np.tile(sigma_bins, len(freq_bins))}) 
-    names=[]
-    means=[]
-    stds=[]
-    for i in range(len(freq_bins)):
-        name=[str(freq_bins[i])+"Hz" for j in range(len(sigma_bins))]
-        names.extend(name)
-        means.extend(sigma_freq_means[i])
-        stds.extend(sigma_freq_stds[i])
-    df['freq']=names
-    df['means']=means
-    df['stds']=stds
-    print(df)
-
-
-    #basically, we want to plot the loss as a function of sigma and freq. We do it by plotting a line for each freq.
-    fig= error_line('bar', data_frame=df, x="sigma", y="means", error_y="stds", color="freq", log_x=True,  markers=True, range_y=[0, 2])
-    
-    return fig
-
-
-
-def plot_melspectrogram(X, refr=1):
-    X=X.squeeze(1) #??
-    X=X.cpu().numpy()
-    if refr==None:
-        refr=np.max(np.abs(X))+1e-8
-     
-    S_db = 10*np.log10(np.abs(X)/refr)
-
-    S_db=np.transpose(S_db, (0,2,1))
-    S_db=np.flip(S_db, axis=1)
-
-    for i in range(X.shape[0]): #iterate over batch size, shity way of ploting all the batched spectrograms
-        o=S_db[i]
-        if i==0:
-             res=o
-        else:
-             res=np.concatenate((res,o), axis=1)
-      
-    fig=px.imshow( res,  zmin=-40, zmax=20)
-
-    fig.update_layout(coloraxis_showscale=False)
-
-    return fig
-
-def plot_cpxspectrogram(X):
-    X=X.squeeze(1)
-    X=X.cpu().numpy()
-    fig=px.imshow(X, facet_col=3, animation_frame=0)
-    fig.update_layout(coloraxis_showscale=False)
-    return fig
-def print_cuda_memory():
-    t = torch.cuda.get_device_properties(0).total_memory
-    r = torch.cuda.memory_reserved(0)
-    a = torch.cuda.memory_allocated(0)
-    f = r-a  # free inside reservedk
-    print("memrylog",t,r,a,f)
-
-
-
-def plot_mag_spectrogram(X, refr=None, path=None,name="spec"):
-    X=X.cpu().numpy()
-    if refr==None:
-        refr=np.max(np.abs(X))+1e-8
-     
-    S_db = 10*np.log10(np.abs(X)/refr)
-    S_db=np.flip(S_db, axis=1)
-
-    for i in range(X.shape[0]): #iterate over batch size, shity way of ploting all the batched spectrograms
-        o=S_db[i]
-        if i==0:
-             res=o
-        else:
-             res=np.concatenate((res,o), axis=1)
-      
-    fig = px.imshow(res,  zmin=-40, zmax=20)
-
-    fig.update_layout(coloraxis_showscale=False)
-
-    path_to_plotly_png = path+"/"+name+".png"
-    plotly.io.write_image(fig, path_to_plotly_png)
-    
-    return fig
-
-
-def plot_spectrogram(X, refr=None, cplx=True, freqs=None, time=None):
-    X=X.squeeze(1)
-    X=X.cpu().numpy()
-    if cplx:
-        X=np.sqrt(X[:,:,:,0]**2 + X[:,:,:,1]**2)
-
-    if refr==None:
-            refr=np.max(np.abs(X))+1e-8
-     
-    S_db = 10*np.log10(np.abs(X)/refr)
-
-    S_db=np.transpose(S_db, (0,2,1))
-    S_db=np.flip(S_db, axis=1)
-
-    for i in range(X.shape[0]): #iterate over batch size, shity way of ploting all the batched spectrograms
-        o=S_db[i]
-        if i==0:
-             res=o
-        else:
-             res=np.concatenate((res,o), axis=1)
-    #freqs ia a tensor of shape (F,)
-    #time is a tensor of shape (T,)
-    #add the frequency and time axis correctly
-
-      
-    fig=px.imshow( res,  zmin=-40, zmax=20)
-    fig.update_layout(coloraxis_showscale=False)
-
-    return fig
-
-def write_audio_file(x, sr, string: str, path='tmp', stereo=False, normalize=False):
-    if normalize:
-        x=x/torch.max(torch.abs(x))
-    if not(os.path.exists(path)): 
-        os.makedirs(path)
-      
-    path=os.path.join(path,string+".wav")
-    if stereo:
-        '''
-        x has shape (B,2,T)
-        '''
-        x=x.permute(0,2,1) #B,T,2
-        x=x.flatten(0,1) #B*T,2
-        x=x.cpu().numpy()
-        #if np.abs(np.max(x))>=1:
-        #    #normalize to avoid clipping
-        #    x=x/np.abs(np.max(x))
-    else:
-        x=x.flatten()
-        x=x.unsqueeze(1)
-        x=x.cpu().numpy()
-        #if np.abs(np.max(x))>=1:
-        #    #normalize to avoid clipping
-        #    x=x/np.abs(np.max(x))
-    sf.write(path,x,sr)
-    return path
-
-def plot_trajectories(x,sr, sigma_data):
-    """
-    x: (B,T) Batch of trajectories
-    sr: sampling rate
-    args: args object
-    """
-    x=x.cpu()
-    times=np.arange(x.shape[1])/sr
-
-    fig=plot_batch_of_lines(x/sigma_data, times, log_x=False)
-    return fig
-
-def plot_trajectories_fft(x, sr, sigma_data):
-    """
-    x: (B,T) Batch of trajectories
-    sr: sampling rate
-    args: args object
-
-    Apply fft to the trajectories and plot the magnitude of the fft
-    """
-    x=x.cpu()
-
-    times=np.arange(x.shape[1])/sr
-    x=x/sigma_data
-    x_fft=np.fft.rfft(x, axis=1)
-    x_fft=np.abs(x_fft)
-    freqs=np.fft.rfftfreq(x.shape[1], d=1/sr)
-    #now do it on db scale
-    x_fft_db=10*np.log10(x_fft)
-    fig=plot_batch_of_lines(x_fft_db, freqs, log_x=False)
-    return fig
-    return fig
-
-
-def plot_cpxCQT_from_raw_audio(x, args, refr=None ):
-    #shape of input spectrogram:  (     ,T,F, )
-    fmax=args.sample_rate/2
-    fmin=fmax/(2**args.cqt.numocts)
-    fbins=int(args.cqt.binsoct*args.cqt.numocts) 
-    device=x.device
-    CQTransform=CQT_cpx(fmin,fbins, args.sample_rate, args.audio_len, device=device, split_0_nyq=False)
-
-    x=x
-    X=CQTransform.fwd(x)
-    return plot_cpxspectrogram(X)
-
-def plot_CQT_from_raw_audio(x, args, refr=None ):
-    #shape of input spectrogram:  (     ,T,F, )
-    fmax=args.sample_rate/2
-    fmin=fmax/(2**args.cqt.numocts)
-    fbins=int(args.cqt.binsoct*args.cqt.numocts) 
-    device=x.device
-    CQTransform=CQT_cpx(fmin,fbins, args.sample_rate, args.audio_len, device=device, split_0_nyq=False)
-
-    refr=3
-    x=x
-    X=CQTransform.fwd(x)
-    return plot_spectrogram(X, refr)
-
-def get_spectrogram_from_raw_audio(x, stft, refr=1):
-    X=do_stft(x, win_size=stft.win_size, hop_size=stft.hop_size)
-    X=X.permute(0,2,3,1)
-
-    X=X.squeeze(1)
-    #X=X.cpu().numpy()
-    X=torch.sqrt(X[:,:,:,0]**2 + X[:,:,:,1]**2)
-
-    #if refr==None:
-    #    refr=np.max(np.abs(X))+1e-8
-     
-    S_db = 10*torch.log10(torch.abs(X)/refr)
-
-    S_db=S_db.permute(0,2,1)
-    #np.transpose(S_db, (0,2,1))
-    S_db=torch.flip(S_db, [1])
-
-    for i in range(X.shape[0]): #iterate over batch size, shity way of ploting all the batched spectrograms
-        o=S_db[i]
-        if i==0:
-             res=o
-        else:
-             res=torch.cat((res,o), 1)
-    return res
-
-
-
-def diffusion_CQT_animation(path, x ,t,  args, refr=1, name="animation_diffusion", resample_factor=1 ):
-    """
-        Utility for creating an animation of the cqt diffusion process
-    """
-    #shape of input spectrograms:  (Nsteps,B,Time,Freq)
-    #print(noisy.shape)
-    Nsteps=x.shape[0]
-    numsteps=10
-    tt=torch.linspace(0, Nsteps-1, numsteps)
-    i_s=[]
-    allX=None
-    device=x.device
-
-    fmax=args.sample_rate/2
-    fmin=fmax/(2**args.cqt.numocts)
-    fbins=int(args.cqt.binsoct*args.cqt.numocts) 
-    CQTransform=CQT_cpx(fmin,fbins, args.sample_rate, args.audio_len, device=device, split_0_nyq=False)
-
-    for i in tt:
-        i=int(torch.floor(i))
-        i_s.append(i)
-        xx=x[i]
-        X=CQTransform.fwd(xx)
-
-        X=torch.sqrt(X[...,0]**2 + X[...,1]**2)
-
-        S_db = 10*torch.log10(torch.abs(X)/refr)
-        S_db = S_db[:,1:-1,1:-1]
-       
-        S_db=S_db.permute(0,2,1)
-        #np.transpose(S_db, (0,2,1))
-        S_db=torch.flip(S_db, [1])
-        S_db=S_db.unsqueeze(0)
-        S_db=torch.nn.functional.interpolate(S_db, size= (S_db.shape[2]//resample_factor, S_db.shape[3]//resample_factor), mode="bilinear")
-        S_db=S_db.squeeze(0)
-        S_db=S_db.cpu().numpy()
-        #S_db=S_db.squeeze(0)        
-
-        if i==0:
-             allX=S_db
-        else:
-             allX=np.concatenate((allX,S_db), 0)
-    
-    #fig=px.imshow(S_db, animation_frame=0, facet_col=3, binary_compression_level=0) #I need to incorporate t here!!!
-    fig=px.imshow(allX, animation_frame=0,  zmin=-40, zmax=10, binary_compression_level=0) #I need to incorporate t here!!!
-
-    fig.update_layout(coloraxis_showscale=False)
-    fig.update_xaxes(showticklabels=False)
-    fig.update_yaxes(showticklabels=False)
-    
-    t=t[i_s].cpu().numpy()
-
-    assert len(t)==len(fig.frames), print(len(t), len(fig.frames))
-    #print(fig)
-    for i, f in enumerate(fig.layout.sliders[0].steps):
-        #hacky way of changing the axis values
-        #f.args[0]=[str(t[i])]
-        f.label=str(t[i])
-
-    path_to_plotly_html = path+"/"+name+".html"
-    
-    fig.write_html(path_to_plotly_html, auto_play = False)
-
-
-    return fig
-
-def diffusion_joint_filter_animation(path, score, grads , f, t,  refr=1, name="animation_diffusion", NT=15 ):
-    '''
-    plot an animation of the reverse diffusion process of filters
-    args:
-        path: path to save the animation
-        x: filters (Nsteps, F)
-        f: frequencies (F)
-        t: timesteps (sigma)
-        name: name of the animation
-    '''
-    #shape of input spectrograms:  (Nsteps,B,Time,Freq)
-    #print(noisy.shape)
-    Nsteps=score.shape[0]
-    #numsteps=min(Nsteps,NT) #hardcoded, I'll probably need more!
-
-    #tt=torch.linspace(0, Nsteps-1, numsteps)
-    #i_s=[]
-    #for i in tt:
-    #    i=int(torch.floor(i))
-    #    i_s.append(i)
-      
-    print(f.shape, score.shape, grads.shape, t.shape)  #19, (100,19)
-    #fig=px.line(x=f.unsqueeze(0).numpy(),y=x, animation_frame=0) #I need to incorporate t here!!!
-    sigma=t
-    s=score.squeeze(1)# (100,19)
-    g=grads.squeeze(1)# (100,19)
-    f=f# (19,)
-    f=f.unsqueeze(0).expand(s.shape[0], -1).reshape(-1)
-    sigma=sigma.unsqueeze(-1).expand(-1, s.shape[1]).reshape(-1)
-    s=s.reshape(-1)
-    g=g.reshape(-1)
-    max_score=torch.max(s).item()
-    df=pd.DataFrame(
-        {
-            "f": f.cpu().numpy(),
-            "score": s.cpu().numpy(),
-            "grads": g.cpu().numpy(),
-            "sigma": sigma.cpu().numpy()
-        }
-    )
-    fig=px.line(df, x="f",y=["score","grads"], animation_frame="sigma", log_x=True) #I need to incorporate t here!!!
-    #set the range of th y axis
-    fig.update_yaxes(range=[-max_score,max_score])
-    #t=t[i_s].cpu().numpy()
-    #t=t.cpu().numpy()
-       
-    #assert len(t)==len(fig.frames), print(len(t), len(fig.frames))
-    #print(fig)
-    #for i, f in enumerate(fig.layout.sliders[0].steps):
-        #hacky way of changing the axis values
-        #f.args[0]=[str(t[i])]
-    #    f.label=str(t[i])
-
-    path_to_plotly_html = path+"/"+name+".html"
-    
-    fig.write_html(path_to_plotly_html, auto_play = False)
-
-    return fig
-
-def diffusion_filter_animation(path, x , f, t,  refr=1, name="animation_diffusion", NT=15 ):
-    '''
-    plot an animation of the reverse diffusion process of filters
-    args:
-        path: path to save the animation
-        x: filters (Nsteps, F)
-        f: frequencies (F)
-        t: timesteps (sigma)
-        name: name of the animation
-    '''
-    #shape of input spectrograms:  (Nsteps,B,Time,Freq)
-    #print(noisy.shape)
-    Nsteps=x.shape[0]
-    #numsteps=min(Nsteps,NT) #hardcoded, I'll probably need more!
-
-    #tt=torch.linspace(0, Nsteps-1, numsteps)
-    #i_s=[]
-    #for i in tt:
-    #    i=int(torch.floor(i))
-    #    i_s.append(i)
-      
-    print(f.shape, x.shape, t.shape)  #19, (100,19)
-    #fig=px.line(x=f.unsqueeze(0).numpy(),y=x, animation_frame=0) #I need to incorporate t here!!!
-    sigma=t
-    x=x.squeeze(1)# (100,19)
-    f=f# (19,)
-    f=f.unsqueeze(0).expand(x.shape[0], -1).reshape(-1)
-    sigma=sigma.unsqueeze(-1).expand(-1, x.shape[1]).reshape(-1)
-    x=x.reshape(-1)
-    df=pd.DataFrame(
-        {
-            "f": f.cpu().numpy(),
-            "x": x.cpu().numpy(),
-            "sigma": sigma.cpu().numpy()
-        }
-    )
-    fig=px.line(df, x="f",y="x", animation_frame="sigma", log_x=True) #I need to incorporate t here!!!
-    #t=t[i_s].cpu().numpy()
-    #t=t.cpu().numpy()
-       
-    #assert len(t)==len(fig.frames), print(len(t), len(fig.frames))
-    #print(fig)
-    #for i, f in enumerate(fig.layout.sliders[0].steps):
-        #hacky way of changing the axis values
-        #f.args[0]=[str(t[i])]
-    #    f.label=str(t[i])
-
-    path_to_plotly_html = path+"/"+name+".html"
-    
-    fig.write_html(path_to_plotly_html, auto_play = False)
-
-    return fig
-
-def diffusion_spec_animation(path, x ,t,  stft, refr=1, name="animation_diffusion",NT=15 ):
-    '''
-    plot an animation of the reverse diffusion process of filters
-    args:
-        path: path to save the animation
-        x: input audio (N,T)
-        t: timesteps (sigma)
-        name: name of the animation
-    '''
-    #print(noisy.shape)
-    Nsteps=x.shape[0]
-    numsteps=min(Nsteps,NT) #hardcoded, I'll probably need more!
-    tt=torch.linspace(0, Nsteps-1, numsteps)
-    i_s=[]
-    allX=None
-    for i in tt:
-        i=int(torch.floor(i))
-        i_s.append(i)
-        X=get_spectrogram_from_raw_audio(x[i],stft, refr)
-        X=X.unsqueeze(0)
-        if allX==None:
-             allX=X
-        else:
-             allX=torch.cat((allX,X), 0)
-
-      
-    allX=allX.cpu().numpy()
-    fig=px.imshow(allX, animation_frame=0,  zmin=-40, zmax=20) #I need to incorporate t here!!!
-
-    fig.update_layout(coloraxis_showscale=False)
-    
-    t=t[i_s].cpu().numpy()
-       
-    assert len(t)==len(fig.frames), print(len(t), len(fig.frames))
-    #print(fig)
-    for i, f in enumerate(fig.layout.sliders[0].steps):
-        #hacky way of changing the axis values
-        #f.args[0]=[str(t[i])]
-        f.label=str(t[i])
-
-    path_to_plotly_html = path+"/"+name+".html"
-    
-    fig.write_html(path_to_plotly_html, auto_play = False)
-
-
-    return fig
-
-def plot_spectrogram_from_raw_audio(x, stft, refr=None ):
-    #shape of input spectrogram:  (     ,T,F, )
-    refr=3
-    x=x
-    X=do_stft(x, win_size=stft.win_size, hop_size=stft.hop_size)
-    X=X.permute(0,2,3,1)
-    return plot_spectrogram(X, refr)
-
-def plot_spectrogram_from_cpxspec(X, refr=None):
-    #shape of input spectrogram:  (     ,T,F, )
-    return plot_spectrogram(X, refr)
-
-
-def error_line(error_y_mode='band', **kwargs):
-    print("hello error line", error_y_mode)
-    """Extension of `plotly.express.line` to use error bands."""
-    ERROR_MODES = {'bar','band','bars','bands',None}
-    if error_y_mode not in ERROR_MODES:
-        raise ValueError(f"'error_y_mode' must be one of {ERROR_MODES}, received {repr(error_y_mode)}.")
-    if error_y_mode in {'bar','bars',None}:
-        print("Using error bars", kwargs)
-        fig = px.line(**kwargs)
-    elif error_y_mode in {'band','bands'}:
-        if 'error_y' not in kwargs:
-            raise ValueError(f"If you provide argument 'error_y_mode' you must also provide 'error_y'.")
-        figure_with_error_bars = px.line(**kwargs)
-        fig = px.line(**{arg: val for arg,val in kwargs.items() if arg != 'error_y'})
-        for data in figure_with_error_bars.data:
-            x = list(data['x'])
-            y_upper = list(data['y'] + data['error_y']['array'])
-            y_lower = list(data['y'] - data['error_y']['array'] if data['error_y']['arrayminus'] is None else data['y'] - data['error_y']['arrayminus'])
-            color = f"rgba({tuple(int(data['line']['color'].lstrip('#')[i:i+2], 16) for i in (0, 2, 4))},.3)".replace('((','(').replace('),',',').replace(' ','')
-            fig.add_trace(
-                go.Scatter(
-                    x = x+x[::-1],
-                    y = y_upper+y_lower[::-1],
-                    fill = 'toself',
-                    fillcolor = color,
-                    line = dict(
-                        color = 'rgba(255,255,255,0)'
-                    ),
-                    hoverinfo = "skip",
-                    showlegend = False,
-                    legendgroup = data['legendgroup'],
-                    xaxis = data['xaxis'],
-                    yaxis = data['yaxis'],
-                )
-            )
-        # Reorder data as said here: https://stackoverflow.com/a/66854398/8849755
-        reordered_data = []
-        for i in range(int(len(fig.data)/2)):
-            reordered_data.append(fig.data[i+int(len(fig.data)/2)])
-            reordered_data.append(fig.data[i])
-        fig.data = tuple(reordered_data)
-    return fig
-
-
-def plot_filters( x, freqs):
-    '''
-    This function plots a batch of lines using plotly
-    args:
-        x: (B, F)
-    '''
-    fig = px.line(x=freqs, y=x[0,:], log_x=True)
-
-    for i in range(1,x.shape[0]):
-        fig.add_trace(go.Scatter(x=freqs, y=x[i,:]))
-    return fig
-
-def plot_batch_of_lines( x, freqs, log_x=True):
-    '''
-    This function plots a batch of lines using plotly
-    args:
-        x: (B, F)
-    '''
-    fig = px.line(x=freqs, y=x[0,:], log_x=log_x)
-
-    for i in range(1,x.shape[0]):
-        fig.add_trace(go.Scatter(x=freqs, y=x[i,:]))
-    return fig
-
-
-
-
-def plot_rir(rir, T_separation, fs, path=None, name=None, save=True):
-    '''
-    Plot the RIR and indicate the separation time between early and late reflections
-    '''
-    rir = rir.cpu()
-    RIR = torch.stft(rir, n_fft=512, hop_length=128, window=torch.hann_window(512).to(rir.device), return_complex=True)
-    
-    fig, ax = plt.subplots(1, 2)
-
-    ax[0].pcolormesh(
-        torch.linspace(0, RIR.size(-1)/128, RIR.size(-1)), 
-        torch.linspace(0, fs/2, RIR.size(-2)),
-        20*torch.log10(RIR.abs() + 1e-10))
-
-    ax[1].plot(
-        torch.linspace(0, rir.size(-1)/fs, rir.size(-1)), 
-        rir)
-    
-    if save:
-        plt.savefig(os.path.join(path, name+".png"))
-    else:
-        return fig
-
-
-def create_mag_phase_plot():
-    # Create subplots
-    from plotly.subplots import make_subplots
-
-    # fig = make_subplots(rows=2, cols=1, vertical_spacing=0.1,
-    #                     subplot_titles=('Magnitude Response', 'Phase Response'))
-    fig = make_subplots(specs=[[{"secondary_y": True}]])
-    # Update layout
-    fig.update_layout(height=600, width=800, title_text="Magnitude and Phase of Filter")
-    # fig.update_xaxes(type='log',title_text="Frequency [Hz])", row=2, col=1)
-    fig.update_xaxes(type='log', title_text="Frequency [Hz])")
-    # fig.update_yaxes(title_text="Magnitude [dB]", row=1, col=1)
-    # fig.update_yaxes(title_text="Phase [rad]", row=2, col=1)
-    # Set y-axes titles
-    fig.update_yaxes(title_text="<b>Magnitude [dB]</b>", secondary_y=False)
-    fig.update_yaxes(title_text="<b>Phase [rad]</b>", secondary_y=True)
-    return fig
-
-
-def plot_STFT(x, operator, type='mag_dB', min_freq=100):
-
-    # Ensure that shape is (B, F, T)
-    fs = operator.sample_rate
-    nfft = operator.n_fft
-    hop_size = operator.hop_length
-
-    print(x.device, operator.device)
-    X = operator.apply_stft(x.to(operator.device))
-
-    print(X.shape, x.shape, X.device)
-
-    # , X.device)
-    if type == 'mag':
-        x_spec = X.abs()
-    if type == 'mag_dB':
-        x_spec = 20*torch.log10(X.abs()+1e-7)
-    if type == 'phase':
-        x_spec = X.angle()
-    # if type == 'phase_unwrap':
-    #     x_spec = X.angle().grad()
-
-
-    if "mag" in type:
-        z_min=-100
-        z_max= 10
-    else:
-        z_min=-np.pi
-        z_max= np.pi
-
-
-
-
-    frequencies = torch.fft.rfftfreq(nfft) * fs
-    timestamps = torch.arange(X.shape[-1]) * (hop_size / fs)
-    fig = go.Figure(data=go.Heatmap(
-        z=x_spec.squeeze(0).cpu().numpy(),
-        x=timestamps.numpy(),
-        y=frequencies.numpy(),
-        colorscale='Magma',
-        zmax=z_max,
-        zmin=z_min
-    ))
-
-    fig.update_layout(
-        yaxis=dict(
-            type='log',
-            title='Frequency(Hz)',
-            range= [np.log10(min_freq), np.log10(frequencies[-1].item())],
-        ),
-        xaxis_title='Time(s)',
-        title='Magnitude'
-    )
-
-    #define min and max for frequency axis
-
-    return fig
-
-def lineplot(x=None, dict=None, xaxis="linear", yaxis="log"):
-    assert dict is not None
-
-    y=dict.items()
-    y=torch.stack([v for k,v in y])
-
-    labels = [k for k,v in dict.items()]
-
-
-
-    # Ensure y is a NumPy array for compatibility with Plotly
-    if isinstance(y, torch.Tensor):
-        y = y.detach().cpu().numpy()
-
-    if x is None:
-        x = np.arange(y.shape[-1])  # Ensure x is NumPy array too
-    elif isinstance(x, torch.Tensor):
-        x = x.detach().cpu().numpy()
-
-    # Validate labels
-    num_series = y.shape[0]
-    if labels is None:
-        labels = [f"Series {i}" for i in range(num_series)]
-    elif len(labels) != num_series:
-        raise ValueError(f"Expected {num_series} labels, but got {len(labels)}.")
-
-    # Convert to DataFrame for Plotly
-    df = pd.DataFrame(y.T, index=x, columns=labels)  # Use labels as column names
-    df.index.name = "x"
-    df = df.reset_index()
-
-    # Melt DataFrame for Plotly (long format)
-    df_melted = df.melt(id_vars=["x"], var_name="series", value_name="y")
-
-    fig = px.line(df_melted, x="x", y="y", color="series", labels={"y": "Value", "x": "t", "series": "Legend"})
-
-    # Set y-axis to log scale
-    if yaxis=="log":
-        fig.update_yaxes(type="log")
-        #fix range from 1e-4 to 10
-        fig.update_yaxes(range=[-4, 1])
-    if xaxis=="log":
-        fig.update_xaxes(type="log")
-
-
-    return fig
-
-def make_PCA_figure(data_dict, num_bins=20, title="PCA"):
-
-    fig = plt.figure(figsize=(6, 5), dpi=200)
-    gs = fig.add_gridspec(
-        2,
-        2,
-        width_ratios=(4, 1),
-        height_ratios=(1, 4),
-        left=0.1,
-        right=0.9,
-        bottom=0.1,
-        top=0.9,
-    )
-
-    # Create the Axes.
-    ax = fig.add_subplot(gs[1, 0])
-    ax_histx = fig.add_subplot(gs[0, 0], sharex=ax)
-    ax_histy = fig.add_subplot(gs[1, 1], sharey=ax)
-    # Draw the scatter plot and marginals.
-    # no labels
-    ax_histx.tick_params(axis="x", labelbottom=False)
-    ax_histy.tick_params(axis="y", labelleft=False)
-
-    xymax=0
-    xymin=999999
-    markers= ["o", "x", "D", "^", "v", "x", "p", "*", "h", "+", "s"]
-    colors=["#1f77b4", "#ff7f0e", "#2ca02c", "#d62728", "#9467bd", "#8c564b"]
-
-
-    for i,(k, v) in enumerate(data_dict.items()):
-        if v is None or len(v) == 0:
-            continue  # Skip empty or None values
-
-        # the scatter plot:
-        ax.scatter(
-            v[:, 0],
-            v[:, 1],
-            marker=markers[i % len(markers)],
-            c=colors[i % len(colors)],
-            linewidths=0.8,
-            s=24,
-            label=k,
-            alpha=0.7,
-        )
-
-        try:
-            xymax=max(xymax, np.max(v[:, :2]))
-            xymin=min(xymin, np.min(v[:, :2]))
-        except Exception as e:
-            print(f"Error processing data for {k}: {e}")
-            print("data:",v)
-            print(f"Data shape: {v.shape}")
-            continue
-
-    ax.legend(loc="upper left")
-    ax.set_xlabel("PC 1")
-    ax.set_ylabel("PC 2")
-
-    # now determine nice limits by hand:
-    data_range = xymax - xymin
-    binwidth = data_range / num_bins  # Calculate bin width based on the range and number of bins
-
-    lim = (int(xymax / binwidth) + 1) * binwidth
-    bins = np.arange(-lim, lim + binwidth, binwidth)
-
-    style = { "edgecolor": "black", "linewidth": 0.8, "alpha": 0.5 }
-
-
-
-    #linestyles=[":", "-", "--", "-.", ":", "-"]
-    facecolors=["#1f77b4", "#ff7f0e", "#2ca02c", "#d62728", "#9467bd", "#8c564b"]
-    hatches = ['/', '\\', '.', '+', '*', 'o', 'O', '.', '-', '|', 'x']
-
-    for i, (k, v) in enumerate(data_dict.items()):
-
-        hist_style = style.copy()
-        hist_style["hatch"] = hatches[i % len(hatches)]
-
-        ax_histx.hist(v[:, 0], bins=bins, density=False,  label=k, facecolor=facecolors[i], **hist_style)
-
-        ax_histy.hist(
-            v[:, 1],
-            bins=bins,
-            orientation="horizontal",
-            density=False,
-            facecolor=facecolors[i],
-            **hist_style
-        )
-
-    ax_histx.set_ylabel("Count")
-    ax_histy.set_xlabel("Count")
-
-    #ax_histx.legend(loc="upper right")
-
-    #add title
-    fig.suptitle(title, fontsize=16, fontweight='bold')
-
-    return fig
-
-def make_histogram_figure(data_dict, num_bins=20):
-
-    fig = plt.figure(figsize=(5, 3), dpi=100)
-
-    # Create the Axes.
-    #ax = fig.add_subplot(gs[1, 0])
-    #ax_histx = fig.add_subplot(gs[0, 0], sharex=ax)
-    #ax_histy = fig.add_subplot(gs[1, 1], sharey=ax)
-    # Draw the scatter plot and marginals.
-    # no labels
-
-    #get ax from the figure
-    ax = fig.add_subplot(111)
-
-    ax.tick_params(axis="x", labelbottom=True)
-
-
-    xymax=0
-    xymin=999999
-
-    for i,(k, v) in enumerate(data_dict.items()):
-
-        xymax=max(xymax, np.max(v))
-        xymin=min(xymin, np.min(v))
-
-
-    data_range = xymax - xymin  
-    binwidth = data_range / num_bins  # Calculate bin width based on the range and number of bins
-
-    lim_up = (int(xymax / binwidth) + 1) * binwidth+binwidth
-    lim_down= (int(xymin / binwidth) - 1) * binwidth-binwidth
-
-    bins = np.arange(lim_down, lim_up, binwidth)
-
-    style = { "edgecolor": "black", "linewidth": 0.8, "alpha": 0.5 }
-
-    #linestyles=[":", "-", "--", "-.", ":", "-"]
-    facecolors=["#1f77b4", "#ff7f0e", "#2ca02c", "#d62728", "#9467bd", "#8c564b"]
-    hatches = ['/', '\\', '.', '+', '*', 'o', 'O', '.', '-', '|', 'x']
-
-    for i, (k, v) in enumerate(data_dict.items()):
-
-        hist_style = style.copy()
-        hist_style["hatch"] = hatches[i % len(hatches)]
-
-        ax.hist(v[:], bins=bins, density=False,  label=k, facecolor=facecolors[i], **hist_style)
-
-    ax.set_ylabel("Count")
-
-    ax.legend(loc="upper right")
-
-
-    return fig
-

utils/st_ito/callbacks.py

@@ -1,164 +0,0 @@
-import os
-import wandb
-import torch
-import shutil
-import numpy as np
-import pytorch_lightning as pl
-import matplotlib.pyplot as plt
-
-from typing import Any
-from sklearn.metrics import confusion_matrix
-from pytorch_lightning.utilities.types import STEP_OUTPUT
-
-
-# trainer.logger.experiment.dir:  /import/c4dm-datasets-ext/wandb/run-20230826_192342-4zjz1inp/files
-
-
-class LogAudioCallback(pl.Callback):
-    def __init__(self) -> None:
-        super().__init__()
-
-    def on_validation_batch_end(
-        self,
-        trainer: pl.Trainer,
-        pl_module: pl.LightningModule,
-        outputs: STEP_OUTPUT,
-        batch: Any,
-        batch_idx: int,
-        dataloader_idx: int = 0,
-    ) -> None:
-        # Only log audio for the first batch
-        if batch_idx > 0:
-            return
-        # Log audio to wandb
-        else:
-            data_dict = outputs
-            silence_gap = 48000
-
-            # check the batch size
-            batch_size = data_dict["input_audio_A"].shape[0]
-
-            for bidx in range(min(4, batch_size)):
-                audio_tensors = []
-                for key in ["input_audio_A", "target_audio_A", "output_audio_A"]:
-                    audio_segment = data_dict[key][bidx, ...]
-                    chs, seq_len = audio_segment.shape
-
-                    if chs == 1:  # make it stereo
-                        audio_segment = audio_segment.repeat(2, 1)
-
-                    audio_tensors.append(audio_segment)
-
-                # check the length of the audio tensors
-                total_audio_length = sum(
-                    [audio_tensor.shape[-1] for audio_tensor in audio_tensors]
-                )
-
-                # create a single tensor to hold the audio
-                output_tensor = torch.zeros((total_audio_length + 2 * silence_gap, 2))
-
-                # copy the audio tensors into the single tensor
-                start_idx = 0
-                for audio_tensor in audio_tensors:
-                    end_idx = start_idx + audio_tensor.shape[-1]
-                    output_tensor[start_idx:end_idx, :] = audio_tensor.permute(1, 0)
-                    start_idx = end_idx + silence_gap
-
-                trainer.logger.experiment.log(
-                    {
-                        f"{bidx}-audio": wandb.Audio(
-                            output_tensor.cpu().numpy(), sample_rate=48000
-                        ),
-                    }
-                )
-
-
-class MoveConfigCallback(pl.Callback):
-    def __init__(self, project_name: str) -> None:
-        super().__init__()
-        self.project_name = project_name
-
-    def on_fit_start(self, trainer, pl_module):
-        full_run_dir = trainer.logger.experiment.dir
-        run_id = full_run_dir.split(os.sep)[-2].split("-")[-1]
-        src_dir = os.path.join(trainer.log_dir, "config.yaml")
-        dest_dir = os.path.join(
-            trainer.log_dir, self.project_name, run_id, "checkpoints", "config.yaml"
-        )
-
-        run_dir = os.path.join(trainer.log_dir, self.project_name, run_id)
-        ckpt_dir = os.path.join(run_dir, "checkpoints")
-        os.makedirs(run_dir, exist_ok=True)
-        os.makedirs(ckpt_dir, exist_ok=True)
-
-        shutil.copyfile(src_dir, dest_dir)
-
-
-class ConfusionMatrixCallback(pl.Callback):
-    def __init__(self) -> None:
-        super().__init__()
-        self.outputs = []
-
-    def on_validation_batch_end(
-        self,
-        trainer: pl.Trainer,
-        pl_module: pl.LightningModule,
-        outputs: STEP_OUTPUT,
-        batch: Any,
-        batch_idx: int,
-        dataloader_idx: int = 0,
-    ) -> None:
-        logits, labels = outputs
-        self.outputs.append(outputs)
-
-    def on_validation_epoch_end(
-        self, trainer: pl.Trainer, pl_module: pl.LightningModule
-    ) -> None:
-        labels = list(trainer.datamodule.val_dataset.plugins.keys())
-
-        # Concatenate logits and labels from all batches in the validation set
-        all_logits = torch.cat([output[0] for output in self.outputs], dim=0)
-        all_labels = torch.cat([output[1] for output in self.outputs], dim=0)
-
-        # Convert logits to predicted labels
-        predicted_labels = torch.argmax(all_logits, dim=-1)
-
-        # Compute confusion matrix
-        cm = confusion_matrix(all_labels.cpu().numpy(), predicted_labels.cpu().numpy())
-
-        # Plot confusion matrix using matplotlib
-        num_classes = cm.shape[0]
-        class_labels = [str(i) for i in labels]
-
-        if num_classes != len(class_labels):
-            print(
-                f"Number of classes ({num_classes}) does not match number of class labels ({len(class_labels)}). Skipping."
-            )
-            return
-
-        fig, ax = plt.subplots(figsize=(num_classes // 2, num_classes // 2))
-        plt.imshow(cm, interpolation="nearest", cmap=plt.cm.Blues)
-        plt.title("Confusion Matrix")
-        plt.colorbar()
-
-        tick_marks = np.arange(num_classes)
-        plt.xticks(tick_marks, class_labels, rotation=90)
-        ax.xaxis.set_tick_params(top=True)
-        plt.yticks(tick_marks, class_labels)
-
-        plt.xlabel("Predicted")
-        plt.ylabel("True")
-        plt.tight_layout()
-
-        for i in range(num_classes):
-            for j in range(num_classes):
-                plt.text(
-                    j, i, str(cm[i, j]), horizontalalignment="center", color="black"
-                )
-
-        # log confusion matrix to wandb
-        trainer.logger.experiment.log({"confusion_matrix": plt})
-        plt.close("all")
-
-        # reset outputs
-        self.outputs = []

utils/st_ito/config.yaml

@@ -1,151 +0,0 @@
-# pytorch_lightning==2.1.2
-seed_everything: 4000152388
-trainer:
-  accelerator: gpu
-  strategy: ddp_find_unused_parameters_true
-  devices: -1
-  num_nodes: 1
-  precision: 32
-  logger:
-    class_path: pytorch_lightning.loggers.WandbLogger
-    init_args:
-      name: null
-      save_dir: /import/c4dm-datasets-ext/lcap/param
-      version: null
-      offline: false
-      dir: null
-      id: null
-      anonymous: null
-      project: lcap-param
-      log_model: false
-      experiment: null
-      prefix: ''
-      checkpoint_name: null
-      job_type: null
-      config: null
-      entity: null
-      reinit: null
-      tags: null
-      group: null
-      notes: null
-      magic: null
-      config_exclude_keys: null
-      config_include_keys: null
-      mode: null
-      allow_val_change: null
-      resume: null
-      force: null
-      tensorboard: null
-      sync_tensorboard: null
-      monitor_gym: null
-      save_code: null
-      settings: null
-  callbacks:
-  - class_path: pytorch_lightning.callbacks.ModelCheckpoint
-    init_args:
-      dirpath: null
-      filename: '{epoch}-{step}'
-      monitor: val_loss
-      verbose: false
-      save_last: true
-      save_top_k: 1
-      save_weights_only: false
-      mode: min
-      auto_insert_metric_name: true
-      every_n_train_steps: null
-      train_time_interval: null
-      every_n_epochs: null
-      save_on_train_epoch_end: null
-      enable_version_counter: true
-  - class_path: pytorch_lightning.callbacks.ModelSummary
-    init_args:
-      max_depth: 2
-  - class_path: pytorch_lightning.callbacks.LearningRateMonitor
-    init_args:
-      logging_interval: null
-      log_momentum: false
-      log_weight_decay: false
-  - class_path: lcap.callbacks.MoveConfigCallback
-    init_args:
-      project_name: lcap-param
-  - class_path: lcap.callbacks.ConfusionMatrixCallback
-  fast_dev_run: false
-  max_epochs: 250
-  min_epochs: null
-  max_steps: -1
-  min_steps: null
-  max_time: null
-  limit_train_batches: null
-  limit_val_batches: null
-  limit_test_batches: null
-  limit_predict_batches: null
-  overfit_batches: 0.0
-  val_check_interval: null
-  check_val_every_n_epoch: 5
-  num_sanity_val_steps: 32
-  log_every_n_steps: 100
-  enable_checkpointing: null
-  enable_progress_bar: null
-  enable_model_summary: null
-  accumulate_grad_batches: 1
-  gradient_clip_val: null
-  gradient_clip_algorithm: null
-  deterministic: null
-  benchmark: true
-  inference_mode: true
-  use_distributed_sampler: true
-  profiler: null
-  detect_anomaly: false
-  barebones: false
-  plugins: null
-  sync_batchnorm: true
-  reload_dataloaders_every_n_epochs: 0
-  default_root_dir: /import/c4dm-datasets-ext/lcap/param
-ckpt_path: null
-model:
-  class_path: lcap.methods.param.ParameterEstimator
-  init_args:
-    encoder:
-      class_path: lcap.models.panns.Cnn14
-      init_args:
-        embed_dim: 512
-        sample_rate: 48000.0
-        window_size: 2048
-        hop_size: 1024
-        mel_bins: 128
-        fmin: 20.0
-        fmax: 20000.0
-        use_batchnorm: true
-        input_norm: minmax
-    lr: 0.0001
-    num_instances: 63
-    num_presets: 10
-    weight_decay: 0.0001
-    max_epochs: 250
-    embed_mode: concat
-    norm: L2
-data:
-  class_path: lcap.dataset.dataset_param.PluginTarfileDataModule
-  init_args:
-    train_tarfiles:
-    - /import/c4dm-datasets-ext/lcap-datasets/20k/musdb18-vst3-presets.tar
-    - /import/c4dm-datasets-ext/lcap-datasets/20k/mtg-jamendo-vst3-presets.tar
-    - /import/c4dm-datasets-ext/lcap-datasets/20k/ursing-vst3-presets.tar
-    - /import/c4dm-datasets-ext/lcap-datasets/20k/enst-drums-vst3-presets.tar
-    - /import/c4dm-datasets-ext/lcap-datasets/20k/fsd50k-vst3-presets.tar
-    - /import/c4dm-datasets-ext/lcap-datasets/20k/librispeech-vst3-presets.tar
-    - /import/c4dm-datasets-ext/lcap-datasets/20k/medley-solos-db-vst3-presets.tar
-    - /import/c4dm-datasets-ext/lcap-datasets/20k/guitarset-vst3-presets.tar
-    val_tarfiles:
-    - /import/c4dm-datasets-ext/lcap-datasets/20k/guitarset-vst3-presets.tar
-    plugin_json: plugins/valid_plugins.json
-    batch_size: 64
-    num_workers: 8
-    length: 262144
-    stereo: true
-    test_tarfiles: null
-    pin_memory: true
-    persistent_workers: true
-    num_train_examples: 140000
-    num_val_examples: 2000
-    num_test_examples: 2000

utils/st_ito/dataset/dataset_param.py

@@ -1,595 +0,0 @@
-import os
-import io
-import sys
-import glob
-import json
-import math
-import torch
-import tarfile
-import torchaudio
-import numpy as np
-import pytorch_lightning as pl
-
-from tqdm import tqdm
-from typing import List
-from torch.utils.data import Dataset, IterableDataset
-
-
-def tarfile_worker_init_fn(worker_id):
-    worker_info = torch.utils.data.get_worker_info()
-    dataset = worker_info.dataset  # the dataset copy in this worker process
-
-    # if there is more than one tar file in the dataset, then we need to
-    # select one tar file for this worker
-    print(f"Worker {worker_id} creating new tar file handles.")
-
-    # reopen the tarfile handles for this worker
-    dataset.tar_handles = []
-    for tar_file in dataset.tar_files:
-        tar_handle = tarfile.open(tar_file, "r:")
-        tar_handle.next()  # skip the first root directory
-        tar_handle.next()  # skip the first example directory
-        dataset.tar_handles.append(tar_handle)
-
-
-def torchaudio_decode(b):
-    x, sr = torchaudio.load(b, format="FLAC")
-    return x
-
-
-class PluginTarfileDataset(IterableDataset):
-    def __init__(
-        self,
-        tar_files: List[str],
-        plugin_json: str,
-        length: int = 262144,
-        stereo: bool = True,
-        num_examples: int = 20000,
-        label_files: List[str] = None,
-    ) -> None:
-        super().__init__()
-        """ Interable dataset that returns a single input/output pair from a tar file. 
-
-        When multiple tar files are provided, we open a file handle to each tar file. Then on each
-        iteration we select a tar file at random and then select a random example from that tar file.
-        We repeat this until each tar file has been read once. Then we start over again.
-
-        Args:
-            tar_files (List[str]): List of tar files to read from.
-            plugin_json (str): Path to json file containing plugin instance details.
-            length (int): Length of audio to return. Defaults to 262144.
-            stereo (bool): whether to return stereo audio. Defaults to True.
-                When set to False, the mean of the two channels is returned.
-            num_examples (int): Number of examples to return. Defaults to 20000.
-            label_files (List[str]): List of label files to use for classification. Defaults to None.
-                Note: they must be int he same order as the associated tar_files.
-
-        """
-        self.tar_files = tar_files
-        self.length = length
-        self.stereo = stereo
-        self.num_examples = num_examples
-        self.label_files = label_files
-
-        # first open plugin info
-        # this is used for indexing into the plugin list
-        with open(plugin_json, "r") as f:
-            self.plugins = json.load(f)
-
-        # get the tar file handles
-        self.tar_handles = []
-        for tar_file in tar_files:
-            tar_handle = tarfile.open(tar_file, "r:")
-            tar_handle.next()  # skip the first root directory
-            tar_handle.next()  # skip the first example directory
-            self.tar_handles.append(tar_handle)
-
-        # check if we have label files
-        if label_files is not None:
-            assert len(label_files) == len(tar_files)  # must be the same length
-            self.labels = []
-            for label_file in label_files:
-                results = torch.load(label_file)
-                self.labels.append(results)
-
-    def __len__(self):
-        return self.num_examples
-
-    def next_example(self):
-        while True:
-            # select tar file at random
-            tar_index = np.random.choice(len(self.tar_handles))
-
-            # get the next member (should be root directory of this example)
-            members = {}
-            done = False
-            handle_reset = False
-            while not done:
-                # get next member
-                member = self.tar_handles[tar_index].next()
-
-                if member is None:
-                    # we have reached the end of this tar file
-                    # so, we will close it and remove open tar handle
-
-                    # close this tar handle
-                    self.tar_handles[tar_index].close()
-
-                    # reopen
-                    tar_handle = tarfile.open(self.tar_files[tar_index], "r:")
-                    tar_handle.next()  # skip the first root directory
-                    tar_handle.next()  # skip the first example directory
-                    self.tar_handles[tar_index] = tar_handle  # replace handle
-                    print(f"Resetting tar handle {self.tar_files[tar_index]}.")
-                    handle_reset = True
-                    break
-
-                if member.isdir():
-                    done = True
-                    break
-
-                members[os.path.basename(member.name)] = member
-
-            if handle_reset:
-                continue
-
-            # get input file
-            input_file = self.tar_handles[tar_index].extractfile(members["input.flac"])
-            input = torchaudio_decode(input_file)
-
-            parent_dir = os.path.dirname(members["input.flac"].name).split("/")[-1]
-
-            # select at random from the processed files
-            processed_members = [
-                x for x in members.keys() if x != "input.flac" and ".json" not in x
-            ]
-            processed_member = np.random.choice(processed_members)
-            processed_fp = self.tar_handles[tar_index].extractfile(
-                members[processed_member]
-            )
-            output = torchaudio_decode(processed_fp)
-
-            # load the json file
-            json_member = members[processed_member.replace(".flac", ".json")]
-            json_fp = self.tar_handles[tar_index].extractfile(json_member)
-            details = json.load(json_fp)
-
-            # if using labels, get the classification info
-            if self.label_files is not None:
-                print(
-                    tar_index, members["input.flac"].name, self.label_files[tar_index]
-                )
-                label_info = self.labels[tar_index][parent_dir]
-                predicted_class_id = label_info["predicted_class_ids"]
-                predicted_label = label_info["predicted_label"]
-                logits = label_info["logits"].squeeze()
-            else:
-                logits = torch.empty(0)
-
-            # get the index of the plugin instance for classification
-            instance_index = self.plugins[details["instance"]]
-            instance_index -= 1  # zero index
-            instance_index = torch.tensor(instance_index).long()
-            preset_index = torch.tensor(details["preset"]).long()
-
-            # fit the proper length
-            if input.shape[-1] < self.length:
-                input = torch.cat(
-                    [
-                        input,
-                        torch.zeros([input.shape[0], self.length - input.shape[-1]]),
-                    ],
-                    dim=-1,
-                )
-            if output.shape[-1] < self.length:
-                output = torch.cat(
-                    [
-                        output,
-                        torch.zeros([output.shape[0], self.length - output.shape[-1]]),
-                    ],
-                    dim=-1,
-                )
-
-            # take a different random crop from each file
-            if input.shape[-1] > self.length:
-                start = torch.randint(0, (input.shape[-1] - self.length), [1])
-                input = input[..., start : start + self.length]
-
-            # take a random crop
-            if output.shape[-1] > self.length:
-                start = torch.randint(0, (output.shape[-1] - self.length), [1])
-                output = output[..., start : start + self.length]
-
-            # check if stereo
-            if self.stereo:
-                # convert mono to stereo
-                if input.shape[0] == 1:
-                    input = input.repeat(2, 1)
-                if output.shape[0] == 1:
-                    output = output.repeat(2, 1)
-            else:
-                # convert stereo to mono
-                if input.shape[0] == 2:
-                    input = torch.mean(input, dim=0, keepdim=True)
-                if output.shape[0] == 2:
-                    output = torch.mean(output, dim=0, keepdim=True)
-
-            # randomize the gain of input and output
-            input /= torch.max(torch.abs(input)).clamp(1e-8)
-            output /= torch.max(torch.abs(output)).clamp(1e-8)
-
-            # from 0 to -32 dB
-            gain_db = torch.rand(1) * -32.0
-            input *= 10.0 ** (gain_db / 20.0)
-
-            # from 0 to -32 dB
-            gain_db = torch.rand(1) * -32.0
-            output *= 10.0 ** (gain_db / 20.0)
-
-            # swap left and right channels at random
-            if torch.rand(1) > 0.5:
-                input = input.flip(0)
-                output = output.flip(0)
-
-            yield input, output, instance_index, preset_index, tar_index, logits
-
-    def __iter__(self):
-        return iter(self.next_example())
-
-
-class PluginTarfileDataModule(pl.LightningDataModule):
-    def __init__(
-        self,
-        train_tarfiles: List[str],
-        val_tarfiles: List[str],
-        plugin_json: str,
-        batch_size: int,
-        num_workers: int,
-        length: int = 262144,
-        stereo: bool = True,
-        test_tarfiles: List[str] = None,
-        train_labels: List[str] = None,
-        val_labels: List[str] = None,
-        test_labels: List[str] = None,
-        pin_memory: bool = False,
-        persistent_workers: bool = True,
-        num_train_examples: int = 20000,
-        num_val_examples: int = 2000,
-        num_test_examples: int = 2000,
-    ) -> None:
-        super().__init__()
-        """ Pytorch Lightning DataModule for loading data from tar files.
-        
-        Args:
-            train_tarfiles (List[str]): List of tar files to use for training.
-            val_tarfiles (List[str]): List of tar files to use for validation.
-            plugin_json (str): Path to json file containing plugin instance details.
-            batch_size (int): Batch size.
-            num_workers (int): Number of workers.
-            length (int): Length of audio to return. Defaults to 262144.
-            stereo (bool): whether to return stereo audio. Defaults to True.
-                When set to False, the mean of the two channels is returned.
-            test_tarfiles (List[str]): List of tar files to use for testing.
-            train_labels (List[str]): List of label files to use for training. Defaults to None.
-            val_labels (List[str]): List of label files to use for validation. Defaults to None.
-            test_labels (List[str]): List of label files to use for testing. Defaults to None.
-            pin_memory (bool): Whether to pin memory. Defaults to False.
-            persistent_workers (bool): Whether to use persistent workers. Defaults to True.
-            num_train_examples (int): Number of training examples to return. Defaults to 20000.
-            num_val_examples (int): Number of validation examples to return. Defaults to 2000.
-            num_test_examples (int): Number of testing examples to return. Defaults to 2000.
-        """
-        self.save_hyperparameters()
-
-    def setup(self, stage=None):
-        self.train_dataset = PluginTarfileDataset(
-            self.hparams.train_tarfiles,
-            self.hparams.plugin_json,
-            self.hparams.length,
-            self.hparams.stereo,
-            self.hparams.num_train_examples,
-            label_files=self.hparams.train_labels,
-        )
-
-        self.val_dataset = PluginTarfileDataset(
-            self.hparams.val_tarfiles,
-            self.hparams.plugin_json,
-            self.hparams.length,
-            self.hparams.stereo,
-            self.hparams.num_val_examples,
-            label_files=self.hparams.val_labels,
-        )
-
-        if self.hparams.test_tarfiles is not None:
-            self.test_dataset = PluginTarfileDataset(
-                self.hparams.test_tarfiles,
-                self.hparams.plugin_json,
-                self.hparams.length,
-                self.hparams.stereo,
-                self.hparams.num_test_examples,
-                label_files=self.hparams.test_labels,
-            )
-
-    def train_dataloader(self):
-        return torch.utils.data.DataLoader(
-            self.train_dataset,
-            batch_size=self.hparams.batch_size,
-            num_workers=self.hparams.num_workers,
-            drop_last=True,
-            pin_memory=self.hparams.pin_memory,
-            persistent_workers=self.hparams.persistent_workers,
-            worker_init_fn=tarfile_worker_init_fn,
-        )
-
-    def val_dataloader(self):
-        return torch.utils.data.DataLoader(
-            self.val_dataset,
-            batch_size=self.hparams.batch_size,
-            num_workers=self.hparams.num_workers,
-            pin_memory=self.hparams.pin_memory,
-            persistent_workers=self.hparams.persistent_workers,
-            worker_init_fn=tarfile_worker_init_fn,
-        )
-
-    def test_dataloader(self):
-        return torch.utils.data.DataLoader(
-            self.val_dataset,
-            batch_size=self.hparams.batch_size,
-            num_workers=self.hparams.num_workers,
-            persistent_workers=self.hparams.persistent_workers,
-            worker_init_fn=tarfile_worker_init_fn,
-        )
-
-
-# ------------------------------ File based dataset --------------------------------
-
-
-class PluginAudioFileWebDataset(Dataset):
-    def __init__(
-        self,
-        root_dirs: str,
-        plugin_json: str,
-        length: int = 524288,
-        stereo: bool = True,
-    ) -> None:
-        """
-
-        Args:
-            root_dirs (List[str]): List of directories containing dataset files.
-            plugin_json: Path to json file containing plugin instance details.
-            length (int): Length of audio to return. Defaults to 524288.
-            stereo (bool): whether to return stereo audio. Defaults to True.
-                When set to False, the mean of the two channels is returned.
-
-        """
-        super().__init__()
-        self.root_dirs = root_dirs
-        self.stereo = stereo
-        self.length = length
-
-        self.examples = []
-
-        # find all segment directories
-        for root_dir in root_dirs:
-            examples = []
-            # first, check if there is an examples.json file
-            examples_json = os.path.join(root_dir, "examples.json")
-            if os.path.exists(examples_json):
-                print(f"Found {examples_json}.")
-                with open(examples_json, "r") as f:
-                    examples = json.load(f)
-            else:
-                print(
-                    f"No examples.json file found. Crawling to find all examples in {root_dir}..."
-                )
-                segment_dirs = glob.glob(os.path.join(root_dir, "*"))
-                segment_dirs = sorted([x for x in segment_dirs if os.path.isdir(x)])
-
-                for segment_dir in tqdm(segment_dirs):
-                    # find the input file
-                    input_file = os.path.join(segment_dir, "input.flac")
-                    if not os.path.exists(input_file):
-                        continue
-
-                    # find all processed audio files
-                    processed_files = glob.glob(os.path.join(segment_dir, "*.flac"))
-
-                    for processed_file in processed_files:
-                        if processed_file == input_file:
-                            continue
-                        # find the appropriate json file
-                        json_file = processed_file.replace(".flac", ".json")
-
-                        # check if json file exists
-                        if not os.path.exists(json_file):
-                            print(f"Could not find {json_file}.")
-                            continue
-
-                        # test open json file
-                        # try:
-                        #    with open(json_file, "r") as f:
-                        #        details = json.load(f)
-                        # except Exception:
-                        #    print(f"Error opening {json_file}.")
-                        #    continue
-
-                        example = (input_file, processed_file, json_file)
-                        examples.append(example)
-
-            self.examples += examples
-            print(
-                f"Found {len(examples)} examples in {root_dir}. Total: {len(self.examples)}"
-            )
-
-        with open(plugin_json, "r") as f:
-            self.plugins = json.load(f)
-
-        # for testing, load one example into memory
-        # input_file, processed_file, json_file = self.examples[0]
-
-        # open the json file
-        # with open(json_file, "r") as f:
-        #    details = json.load(f)
-
-        # input, sr = torchaudio.load(input_file, backend="soundfile")
-        # output, sr = torchaudio.load(processed_file, backend="soundfile")
-
-        # self.input = input
-        # self.output = output
-        # self.details = details
-
-    def __len__(self):
-        return len(self.examples)
-
-    def __getitem__(self, idx):
-        """Get a single example from the dataset."""
-
-        if True:
-            input_file, processed_file, json_file = self.examples[idx]
-
-            # open the json file
-            with open(json_file, "r") as f:
-                details = json.load(f)
-
-            input, sr = torchaudio.load(input_file, backend="soundfile")
-            output, sr = torchaudio.load(processed_file, backend="soundfile")
-        else:
-            input = self.input
-            output = self.output
-            details = self.details
-
-        # check length of input and output
-        # if input.shape[-1] != output.shape[-1]:
-        #    raise RuntimeError(
-        #        f"Input and output shapes do not match: {input.shape} vs {output.shape}"
-        #    )
-
-        # fit the proper length
-        if input.shape[-1] < self.length:
-            input = torch.cat(
-                [input, torch.zeros([input.shape[0], self.length - input.shape[-1]])],
-                dim=-1,
-            )
-        if output.shape[-1] < self.length:
-            output = torch.cat(
-                [
-                    output,
-                    torch.zeros([output.shape[0], self.length - output.shape[-1]]),
-                ],
-                dim=-1,
-            )
-
-        # take a random crop
-        start = torch.randint(
-            0, (min(input.shape[-1], output.shape[-1]) - self.length), [1]
-        )
-        input = input[..., start : start + self.length]
-        output = output[..., start : start + self.length]
-
-        # note: now we take a crop from the same position in both files
-        # however, it might be a good augmentation to take a different
-        # random crop from each file
-
-        # get the index of the plugin instance for classification
-        instance_index = self.plugins[details["instance"]]
-        instance_index -= 1  # zero index
-        instance_index = torch.tensor(instance_index).long()
-        preset_index = torch.tensor(details["preset"]).long()
-
-        # check if stereo
-        if self.stereo:
-            # convert mono to stereo
-            if input.shape[0] == 1:
-                input = input.repeat(2, 1)
-            if output.shape[0] == 1:
-                output = output.repeat(2, 1)
-        else:
-            # convert stereo to mono
-            if input.shape[0] == 2:
-                input = torch.mean(input, dim=0, keepdim=True)
-            if output.shape[0] == 2:
-                output = torch.mean(output, dim=0, keepdim=True)
-
-        # randomize the gain of input and output
-        input /= torch.max(torch.abs(input)).clamp(1e-8)
-        output /= torch.max(torch.abs(output)).clamp(1e-8)
-
-        # from 0 to -32 dB
-        gain_db = torch.rand(1) * -32.0
-        input *= 10.0 ** (gain_db / 20.0)
-
-        # from 0 to -32 dB
-        gain_db = torch.rand(1) * -32.0
-        output *= 10.0 ** (gain_db / 20.0)
-
-        return input, output, instance_index, preset_index
-
-
-class PluginAudioFileWebDataModule(pl.LightningDataModule):
-    def __init__(
-        self,
-        train_root_dirs: List[str],
-        val_root_dirs: List[str],
-        plugin_json: str,
-        batch_size: int,
-        num_workers: int,
-        length: int = 262144,
-        test_root_dirs: List[str] = None,
-        pin_memory: bool = False,
-        persistent_workers: bool = False,
-        use_tarfile: bool = False,
-    ) -> None:
-        super().__init__()
-        self.save_hyperparameters()
-
-    def setup(self, stage=None):
-        if not self.hparams.use_tarfile:
-            self.train_dataset = PluginAudioFileWebDataset(
-                self.hparams.train_root_dirs,
-                self.hparams.plugin_json,
-                self.hparams.length,
-            )
-
-            self.val_dataset = PluginAudioFileWebDataset(
-                self.hparams.val_root_dirs,
-                self.hparams.plugin_json,
-                self.hparams.length,
-            )
-
-            if self.hparams.test_root_dirs is not None:
-                self.test_dataset = PluginAudioFileWebDataset(
-                    self.hparams.test_root_dirs,
-                    self.hparams.plugin_json,
-                    self.hparams.length,
-                )
-        else:
-            raise NotImplementedError
-
-    def train_dataloader(self):
-        return torch.utils.data.DataLoader(
-            self.train_dataset,
-            batch_size=self.hparams.batch_size,
-            shuffle=True,
-            num_workers=self.hparams.num_workers,
-            drop_last=True,
-            pin_memory=self.hparams.pin_memory,
-            persistent_workers=self.hparams.persistent_workers,
-            prefetch_factor=8,
-        )
-
-    def val_dataloader(self):
-        return torch.utils.data.DataLoader(
-            self.val_dataset,
-            batch_size=self.hparams.batch_size,
-            num_workers=self.hparams.num_workers,
-            pin_memory=self.hparams.pin_memory,
-            persistent_workers=self.hparams.persistent_workers,
-        )
-
-    def test_dataloader(self):
-        return torch.utils.data.DataLoader(
-            self.val_dataset,
-            batch_size=self.hparams.batch_size,
-            num_workers=self.hparams.num_workers,
-            persistent_workers=self.hparams.persistent_workers,
-        )

utils/st_ito/dataset/dataset_sim.py

@@ -1,320 +0,0 @@
-import os
-import glob
-import torch
-import torchaudio
-import pedalboard
-import numpy as np
-import pytorch_lightning as pl
-
-from tqdm import tqdm
-from typing import List, Optional, Tuple, Iterable
-from st_ito.effects import (
-    BasicChorus,
-    BasicCompressor,
-    BasicDelay,
-    BasicDistortion,
-    BasicReverb,
-)
-
-
-def worker_init_fn(worker_id):
-    worker_info = torch.utils.data.get_worker_info()
-    # re-init the vst_plugins for each worker
-    vst_plugins = load_plugins(worker_info.dataset.vst_dir)
-    worker_info.dataset.vst_plugins = vst_plugins
-
-
-def find_files_with_glob(
-    directory: str,
-):
-    """Find all files in the given directory with specific pattern."""
-    audio_filepaths = []
-    for ext in ["*.wav", "*.mp3", "*.ogg", "*.flac"]:
-        audio_filepaths += glob.glob(os.path.join(directory, "**", ext), recursive=True)
-
-    return audio_filepaths
-
-
-def find_files_with_extensions(
-    directory: str,
-    extensions: Tuple[str] = (".wav", ".mp3", ".ogg", ".flac"),
-):
-    """Find all files in the given directory with specific extensions."""
-    for entry in os.scandir(directory):
-        if entry.is_file() and entry.name.lower().endswith(extensions):
-            yield entry.path
-        elif entry.is_dir():
-            yield from find_files_with_extensions(entry.path, extensions)
-
-
-def find_files_lazy(
-    directory: str,
-    extensions: Tuple[str] = (".wav", ".mp3", ".ogg", ".flac"),
-):
-    """Lazily find files with specific extensions in the given directory."""
-    for root, dirs, files in os.walk(directory):
-        for file in files:
-            if file.lower().endswith(extensions):
-                yield os.path.join(root, file)
-
-
-def get_nonsilent_region(
-    audio_filepath: str,
-    total_num_frames: int,
-    length: int,
-    threshold: float = 1e-8,
-    max_tires: int = 100,
-):
-    silent = True
-    tries = 0
-    while silent:
-
-        if total_num_frames > length:
-            # load a random segment of the audio file
-            audio, sr = torchaudio.load(
-                audio_filepath,
-                num_frames=length,
-                frame_offset=np.random.randint(0, total_num_frames - length),
-                backend="soundfile",
-            )
-
-        else:
-            # load the entire audio file
-            audio, sr = torchaudio.load(
-                audio_filepath,
-                num_frames=total_num_frames,
-                backend="soundfile",
-            )
-
-            # repeat the audio to reach the desired length
-            audio = torch.cat([audio] * (length // total_num_frames + 1), axis=1)
-            audio = audio[:, :length]
-
-        # check for silence
-        energy = torch.mean(audio**2)
-        print(energy)
-        if energy > threshold:
-            silent = False
-        else:
-            silent = False  # TODO: remove this line
-
-        tries += 1
-
-        if tries > max_tires:
-            raise ValueError(
-                f"Could not find a non-silent region in the audio: {audio_filepath}."
-            )
-
-    return audio, sr
-
-
-def random_gain(audio: torch.Tensor):
-    # first peak normalize
-    audio = audio / torch.max(torch.abs(audio))
-    # then apply random gain
-    gain_db = np.random.uniform(-18, 0)
-    gain_lin = 10 ** (gain_db / 20)
-    audio = audio * gain_lin
-    return audio
-
-
-def conform_channels(audio_input: torch.Tensor, num_channels: int):
-    if audio_input.shape[0] == 1 and num_channels == 2:
-        audio_input = torch.cat((audio_input, audio_input), axis=0)
-    if audio_input.shape[0] == 2 and num_channels == 1:
-        audio_input = audio_input[0:1, :]
-    return audio_input
-
-
-def load_plugins(vst_dir: str, sample_rate: int = 48000):
-    # find all vst3 files
-    vst_filepaths = []
-    for ext in ["*.vst3"]:
-        vst_filepaths += glob.glob(os.path.join(vst_dir, ext), recursive=True)
-
-    # vst_filepaths = vst_filepaths[:2]
-
-    vst_filepaths = sorted(vst_filepaths)
-    print(f"Found {len(vst_filepaths)} VST3 files in {vst_dir}.")
-
-    vst_plugins = {}
-    print("Loading plugins...")
-    for vst_filepath in tqdm(vst_filepaths):
-        instance_name = os.path.basename(vst_filepath).replace(".vst3", "")
-        print(f"Loading {instance_name}...")
-        # try to load each plugin
-        try:
-            plugin = pedalboard.load_plugin(vst_filepath)
-        except Exception:
-            print(f"Error loading {vst_filepath}.")
-
-        # try to pass audio through the plugin
-        failed = True
-        try:
-            num_channels = 2
-            audio_input = np.random.randn(2, sample_rate * 5)
-            audio_output = plugin(audio_input, sample_rate)
-            failed = False
-        except Exception as e:
-            num_channels = 1
-            print(e)
-            print(f"Error processing audio with 2 channels.")
-
-        try:
-            audio_input = np.random.randn(1, sample_rate * 5)
-            audio_output = plugin(audio_input, sample_rate)
-            failed = False
-        except Exception as e:
-            num_channels = 2
-            print(e)
-            print(f"Failed processing audio with 1 channel.")
-
-        if failed:
-            continue
-
-        print(f"Loaded {os.path.basename(vst_filepath)} successfully.")
-
-        vst_plugins[instance_name] = {
-            "filepath": vst_filepath,
-            "instance": plugin,
-            "num_channels": num_channels,
-        }
-
-    for idx, (instance_name, vst_plugin) in enumerate(vst_plugins.items()):
-        print(idx + 1, instance_name)
-
-    return vst_plugins
-
-
-class PluginSimilarityDataset(torch.utils.data.Dataset):
-    def __init__(
-        self,
-        audio_filepaths: List[str],
-        vst_dir: str,
-        length: int = 262144,
-        num_examples_per_epoch: int = 10000,
-    ):
-        self.audio_filepaths = audio_filepaths
-        self.vst_dir = vst_dir
-        self.length = length
-        self.num_examples_per_epoch = num_examples_per_epoch
-        self.plugins = load_plugins(self.vst_dir)
-
-    def __len__(self):
-        return self.num_examples_per_epoch
-
-    def __getitem__(self, idx: int):
-        # sample the first filepath
-        filepath_a = np.random.choice(self.audio_filepaths)
-        # sample the second filepath
-        filepath_b = np.random.choice(self.audio_filepaths)
-
-        # load the audio files
-        num_frames_a = torchaudio.info(filepath_a).num_frames
-        audio_a, sr_a = get_nonsilent_region(filepath_a, num_frames_a, self.length)
-
-        num_frames_b = torchaudio.info(filepath_b).num_frames
-        audio_b, sr_b = get_nonsilent_region(filepath_b, num_frames_b, self.length)
-
-        # sample randomly a plugin from the dictionary
-        plugin_name = np.random.choice(list(self.plugins.keys()))
-        plugin = self.plugins[plugin_name]
-
-        # randomize the parameters of the plugin
-        for name, parameter in plugin["instance"].parameters.items():
-            if "bypass" in name.lower():
-                parameter.raw_value = 0
-            elif name in ["buffer_size_frames", "sample_rate_frames"]:
-                pass
-            else:  # set the parameter to the preset value
-                parameter.raw_value = np.random.rand()
-
-        audio_a = conform_channels(audio_a, plugin["num_channels"])
-        audio_b = conform_channels(audio_b, plugin["num_channels"])
-
-        # apply the effect to the two audio files
-        audio_a_out = torch.from_numpy(
-            plugin["instance"].process(audio_a.numpy(), sr_a)
-        )
-        audio_b_out = torch.from_numpy(
-            plugin["instance"].process(audio_b.numpy(), sr_b)
-        )
-
-        # peak normalize and then apply random gain
-        audio_a = random_gain(audio_a)
-        audio_b = random_gain(audio_b)
-        audio_a_out = random_gain(audio_a_out)
-        audio_b_out = random_gain(audio_b_out)
-
-        # ensure all outputs are stereo
-        audio_a = conform_channels(audio_a, 2)
-        audio_b = conform_channels(audio_b, 2)
-        audio_a_out = conform_channels(audio_a_out, 2)
-        audio_b_out = conform_channels(audio_b_out, 2)
-
-        return audio_a, audio_b, audio_a_out, audio_b_out
-
-
-class PluginSimilarityDataModule(pl.LightningDataModule):
-    def __init__(
-        self,
-        audio_dirs: List[str],
-        vst_dir: str,
-        length: int = 262144,
-        batch_size: int = 32,
-        num_train_examples: int = 10000,
-        num_val_examples: int = 1000,
-        num_workers: int = 4,
-        persistent_workers: bool = True,
-        pin_memory: bool = True,
-    ) -> None:
-        super().__init__()
-        self.save_hyperparameters()
-
-    def prepare_data(self):
-        # download, split, etc...
-        pass
-
-    def setup(self, stage: Optional[str] = None):
-        # find all audio files in the input directory based on file extensions
-        audio_filepaths = []
-        for audio_dir in self.hparams.audio_dirs:
-            audio_filepaths += find_files_with_glob(audio_dir)
-            print("Found", len(audio_filepaths), "audio files in", audio_dir)
-
-        # Assign train/val datasets for use in dataloaders
-        if stage == "fit" or stage is None:
-            self.plugin_similarity_train = PluginSimilarityDataset(
-                audio_filepaths,
-                self.hparams.vst_dir,
-                self.hparams.length,
-                self.hparams.num_train_examples,
-            )
-            self.plugin_similarity_val = PluginSimilarityDataset(
-                audio_filepaths,
-                self.hparams.vst_dir,
-                self.hparams.length,
-                self.hparams.num_val_examples,
-            )
-
-    def train_dataloader(self):
-        return torch.utils.data.DataLoader(
-            self.plugin_similarity_train,
-            batch_size=self.hparams.batch_size,
-            shuffle=True,
-            num_workers=self.hparams.num_workers,
-            persistent_workers=self.hparams.persistent_workers,
-            pin_memory=self.hparams.pin_memory,
-            # worker_init_fn=worker_init_fn,
-        )
-
-    def val_dataloader(self):
-        return torch.utils.data.DataLoader(
-            self.plugin_similarity_val,
-            batch_size=self.hparams.batch_size,
-            shuffle=False,
-            num_workers=self.hparams.num_workers,
-            persistent_workers=self.hparams.persistent_workers,
-            pin_memory=self.hparams.pin_memory,
-            # worker_init_fn=worker_init_fn,
-        )

utils/st_ito/dataset/dataset_style.py

@@ -1,351 +0,0 @@
-import os
-import glob
-import json
-import torch
-import tarfile
-import torchaudio
-import numpy as np
-import pytorch_lightning as pl
-
-from typing import List, Tuple, Dict, Optional
-from st_ito.dataset.utils import tarfile_worker_init_fn, torchaudio_decode
-
-
-class StyleTransferTarfileDataset(torch.utils.data.IterableDataset):
-    def __init__(
-        self,
-        tar_files: List[str],
-        length: int,
-        num_examples: int = 10000,
-        audio_format: str = "wav",
-        input_only: bool = False,
-    ):
-        super().__init__()
-        self.tar_files = tar_files
-        self.length = length
-        self.num_examples = num_examples
-        self.audio_format = audio_format
-        self.input_only = input_only
-
-        # get the tar file handles
-        self.tar_handles = []
-        for tar_file in tar_files:
-            tar_handle = tarfile.open(tar_file, "r:")
-            tar_handle.next()  # skip the first root directory
-            tar_handle.next()  # skip the first example directory
-            self.tar_handles.append(tar_handle)
-
-    def __len__(self):
-        return self.num_examples
-
-    def next_example(self):
-        while True:
-            # select tar file at random
-            tar_index = np.random.choice(len(self.tar_handles))
-
-            # get the next member (should be root directory of this example)
-            members = {}
-            done = False
-            handle_reset = False
-            while not done:
-                # get next member
-                member = self.tar_handles[tar_index].next()
-
-                if member is None:
-                    # we have reached the end of this tar file
-                    # so, we will close it and remove open tar handle
-
-                    # close this tar handle
-                    self.tar_handles[tar_index].close()
-
-                    # reopen
-                    tar_handle = tarfile.open(self.tar_files[tar_index], "r:")
-                    tar_handle.next()  # skip the first root directory
-                    tar_handle.next()  # skip the first example directory
-                    self.tar_handles[tar_index] = tar_handle  # replace handle
-                    print(f"Resetting tar handle {self.tar_files[tar_index]}.")
-                    handle_reset = True
-                    break
-
-                if member.isdir():
-                    done = True
-                    break
-
-                members[os.path.basename(member.name)] = member
-
-            if handle_reset:
-                continue
-
-            # get input file
-            input_file = self.tar_handles[tar_index].extractfile(
-                members[f"input.{self.audio_format}"]
-            )
-            audio_input = torchaudio_decode(input_file)
-
-            if self.input_only:
-
-                num_frames = audio_input.shape[1]
-                start_frame = torch.randint(0, num_frames - self.length, (1,)).item()
-                audio_input = audio_input[:, start_frame : start_frame + self.length]
-
-                # we will generate the output in the system
-                audio_output = torch.empty_like(audio_input)
-                params = torch.empty(1)
-
-            else:
-                # get the output file
-                output_file = self.tar_handles[tar_index].extractfile(
-                    members[f"output.{self.audio_format}"]
-                )
-                audio_output = torchaudio_decode(output_file)
-
-                # get the params file
-                params_file = self.tar_handles[tar_index].extractfile(
-                    members["params.json"]
-                )
-                params = json.load(params_file)
-                # convert params to tensor
-                params = torch.tensor(
-                    [val for val in params.values()], dtype=torch.float32
-                )
-
-                # crop the audio to the desired length
-                num_frames = audio_input.shape[1]
-                start_frame = torch.randint(0, num_frames - self.length, (1,)).item()
-                audio_input = audio_input[:, start_frame : start_frame + self.length]
-                audio_output = audio_output[:, start_frame : start_frame + self.length]
-
-                # check for silence in the output
-                input_energy = torch.mean(audio_input**2, dim=1)
-                output_energy = torch.mean(audio_output**2, dim=1)
-
-                if output_energy < 1e-6 or input_energy < 1e-6:
-                    # skip this example
-                    continue
-
-            yield audio_input, audio_output, params
-
-    def __iter__(self):
-        return iter(self.next_example())
-
-
-class StyleTransferTarfileDataModule(pl.LightningDataModule):
-    def __init__(
-        self,
-        train_tarfiles: List[str],
-        val_tarfiles: List[str],
-        batch_size: int,
-        num_workers: int,
-        input_only: bool = False,
-        length: int = 262144,
-        test_tarfiles: List[str] = None,
-        pin_memory: bool = False,
-        persistent_workers: bool = True,
-        num_train_examples: int = 20000,
-        num_val_examples: int = 2000,
-        num_test_examples: int = 2000,
-        audio_format: str = "wav",
-    ) -> None:
-        super().__init__()
-        """ Pytorch Lightning DataModule for loading data from tar files.
-        
-        Args:
-            train_tarfiles (List[str]): List of tar files to use for training.
-            val_tarfiles (List[str]): List of tar files to use for validation.
-            batch_size (int): Batch size.
-            num_workers (int): Number of workers.
-            input_only (bool): Whether to use only the input audio. Defaults to False.
-            length (int): Length of audio to return. Defaults to 262144.
-            test_tarfiles (List[str]): List of tar files to use for testing.
-            pin_memory (bool): Whether to pin memory. Defaults to False.
-            persistent_workers (bool): Whether to use persistent workers. Defaults to True.
-            num_train_examples (int): Number of training examples to return. Defaults to 20000.
-            num_val_examples (int): Number of validation examples to return. Defaults to 2000.
-            num_test_examples (int): Number of testing examples to return. Defaults to 2000.
-            audio_format (str): Audio format to use. Defaults to "wav".
-        """
-        self.save_hyperparameters()
-
-    def setup(self, stage=None):
-        self.train_dataset = StyleTransferTarfileDataset(
-            self.hparams.train_tarfiles,
-            self.hparams.length,
-            self.hparams.num_train_examples,
-            self.hparams.audio_format,
-            self.hparams.input_only,
-        )
-
-        self.val_dataset = StyleTransferTarfileDataset(
-            self.hparams.val_tarfiles,
-            self.hparams.length,
-            self.hparams.num_val_examples,
-            self.hparams.audio_format,
-            self.hparams.input_only,
-        )
-
-        if self.hparams.test_tarfiles is not None:
-            self.test_dataset = StyleTransferTarfileDataset(
-                self.hparams.test_tarfiles,
-                self.hparams.length,
-                self.hparams.num_test_examples,
-                self.hparams.audio_format,
-                self.hparams.input_only,
-            )
-
-    def train_dataloader(self):
-        return torch.utils.data.DataLoader(
-            self.train_dataset,
-            batch_size=self.hparams.batch_size,
-            num_workers=self.hparams.num_workers,
-            drop_last=True,
-            pin_memory=self.hparams.pin_memory,
-            persistent_workers=self.hparams.persistent_workers,
-            worker_init_fn=tarfile_worker_init_fn,
-        )
-
-    def val_dataloader(self):
-        return torch.utils.data.DataLoader(
-            self.val_dataset,
-            batch_size=self.hparams.batch_size,
-            num_workers=self.hparams.num_workers,
-            pin_memory=self.hparams.pin_memory,
-            persistent_workers=self.hparams.persistent_workers,
-            worker_init_fn=tarfile_worker_init_fn,
-        )
-
-    def test_dataloader(self):
-        return torch.utils.data.DataLoader(
-            self.val_dataset,
-            batch_size=self.hparams.batch_size,
-            num_workers=self.hparams.num_workers,
-            persistent_workers=self.hparams.persistent_workers,
-            worker_init_fn=tarfile_worker_init_fn,
-        )
-
-
-# ------------------------------ File based dataset --------------------------------
-
-
-class StyleTransferDataset(torch.utils.data.Dataset):
-    def __init__(
-        self,
-        root_dirs: List[str],
-        length: int,
-    ):
-        super().__init__()
-        self.root_dirs = root_dirs
-        self.length = length
-
-        self.examples = []
-        for root_dir in root_dirs:
-            # get all examples in the root dir
-            print("Searching for examples in", root_dir)
-            examples = glob.glob(os.path.join(root_dir, "*"))
-            self.examples.extend(examples)
-
-        # for testing use only the first 100 examples
-        # self.examples = self.examples[:100]
-        # repeat those examples 10 times
-        # self.examples = self.examples * 100
-
-        print(f"Found {len(self.examples)} examples.")
-
-    def __len__(self):
-        return len(self.examples)
-
-    def __getitem__(self, idx):
-        example = self.examples[idx]
-
-        input_filepath = os.path.join(example, "input.wav")
-        output_filepath = os.path.join(example, "output.wav")
-        params_filepath = os.path.join(example, "params.json")
-
-        # check the length of the audio
-        md = torchaudio.info(input_filepath, backend="soundfile")
-        num_frames = md.num_frames
-
-        # pick random start frame
-        start_frame = torch.randint(0, num_frames - self.length, (1,)).item()
-        # start_frame = 0  # for now, just use the first frame
-
-        # load the audio
-        input_audio, sr = torchaudio.load(
-            input_filepath,
-            frame_offset=start_frame,
-            num_frames=self.length,
-            backend="soundfile",
-        )
-        output_audio, sr = torchaudio.load(
-            output_filepath,
-            frame_offset=start_frame,
-            num_frames=self.length,
-            backend="soundfile",
-        )
-
-        # load the params
-        with open(params_filepath, "r") as fp:
-            params = json.load(fp)
-
-        # convert params to tensor
-        params = torch.tensor([val for val in params.values()], dtype=torch.float32)
-
-        return input_audio, output_audio, params
-
-
-class StyleTransferDataModule(pl.LightningDataModule):
-    def __init__(
-        self,
-        train_root_dirs: List[str],
-        val_root_dirs: List[str],
-        batch_size: int,
-        num_workers: int,
-        length: int = 262144,
-        test_root_dirs: List[str] = None,
-        pin_memory: bool = False,
-        persistent_workers: bool = False,
-    ) -> None:
-        super().__init__()
-        self.save_hyperparameters()
-
-    def setup(self, stage=None):
-        self.train_dataset = StyleTransferDataset(
-            self.hparams.train_root_dirs,
-            self.hparams.length,
-        )
-        self.val_dataset = StyleTransferDataset(
-            self.hparams.val_root_dirs,
-            self.hparams.length,
-        )
-        if self.hparams.test_root_dirs is not None:
-            self.test_dataset = StyleTransferDataset(
-                self.hparams.test_root_dirs,
-                self.hparams.length,
-            )
-
-    def train_dataloader(self):
-        return torch.utils.data.DataLoader(
-            self.train_dataset,
-            batch_size=self.hparams.batch_size,
-            shuffle=True,
-            num_workers=self.hparams.num_workers,
-            drop_last=True,
-            pin_memory=self.hparams.pin_memory,
-            persistent_workers=self.hparams.persistent_workers,
-        )
-
-    def val_dataloader(self):
-        return torch.utils.data.DataLoader(
-            self.val_dataset,
-            batch_size=self.hparams.batch_size,
-            num_workers=self.hparams.num_workers,
-            pin_memory=self.hparams.pin_memory,
-            persistent_workers=self.hparams.persistent_workers,
-        )
-
-    def test_dataloader(self):
-        return torch.utils.data.DataLoader(
-            self.val_dataset,
-            batch_size=self.hparams.batch_size,
-            num_workers=self.hparams.num_workers,
-            persistent_workers=self.hparams.persistent_workers,
-        )

utils/st_ito/dataset/utils.py

@@ -1,25 +0,0 @@
-import torch
-import tarfile
-import torchaudio
-
-
-def tarfile_worker_init_fn(worker_id):
-    worker_info = torch.utils.data.get_worker_info()
-    dataset = worker_info.dataset  # the dataset copy in this worker process
-
-    # if there is more than one tar file in the dataset, then we need to
-    # select one tar file for this worker
-    print(f"Worker {worker_id} creating new tar file handles.")
-
-    # reopen the tarfile handles for this worker
-    dataset.tar_handles = []
-    for tar_file in dataset.tar_files:
-        tar_handle = tarfile.open(tar_file, "r:")
-        tar_handle.next()  # skip the first root directory
-        tar_handle.next()  # skip the first example directory
-        dataset.tar_handles.append(tar_handle)
-
-
-def torchaudio_decode(b):
-    x, sr = torchaudio.load(b, format="FLAC")
-    return x