initial deploy

.gitignore

@@ -0,0 +1,5 @@
+__pycache__/
+*.pyc
+.DS_Store
+.venv/
+audio-inpainting-diffusion-repo/

README.md

@@ -1,13 +1,12 @@
 ---
 title: Audio Inpainting Diffusion
-emoji: 💻
-colorFrom: indigo
-colorTo: pink
+emoji: 🩹
+colorFrom: blue
+colorTo: purple
 sdk: gradio
-sdk_version: 6.20.0
-python_version: '3.13'
+sdk_version: 5.28.0
+python_version: '3.11'
 app_file: app.py
 pinned: false
+license: mit
 ---
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

SOURCES.md

@@ -0,0 +1,4 @@
+# Sources — audio-inpainting-diffusion
+
+- Source repo: https://github.com/eloimoliner/audio-inpainting-diffusion
+- Paper: https://arxiv.org/html/2305.15266

app.py

@@ -0,0 +1,291 @@
+try:
+    import spaces
+except ImportError:
+    # keep @spaces.GPU usable as a no-op; ZeroGPU requires this exact name.
+    class spaces:
+        class GPU:
+            def __init__(self, func=None, duration=60):
+                self.func = func
+
+            def __call__(self, *args, **kwargs):
+                if self.func is not None:
+                    return self.func(*args, **kwargs)
+                func = args[0]
+                return func
+
+import sys
+sys.stdout.reconfigure(line_buffering=True)
+
+import tempfile
+import threading
+
+import gradio as gr
+import numpy as np
+import soundfile as sf
+import torch
+import torchaudio
+from huggingface_hub import hf_hub_download
+from omegaconf import OmegaConf
+from pyharp import ModelCard, build_endpoint
+
+# cqt_nsgt_pytorch clips an int-dtype array against np.inf with out=, which
+# numpy's ufunc casting rules reject on any numpy version. Compute without
+# `out` and write the result back instead of failing outright.
+_np_clip = np.clip
+
+
+def _clip_int_out_safe(a, a_min, a_max, out=None, **kwargs):
+    if out is not None:
+        out[...] = _np_clip(a, a_min, a_max).astype(out.dtype)
+        return out
+    return _np_clip(a, a_min, a_max, **kwargs)
+
+
+np.clip = _clip_int_out_safe
+
+from diff_params.edm import EDM
+from networks.unet_cqt_oct_with_projattention_adaLN_2 import Unet_CQT_oct_with_attention
+from testing.edm_sampler_inpainting import Sampler
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+CHECKPOINT_REPO = "Eloimoliner/audio-inpainting-diffusion"
+
+# Same network configs as conf/network/paper_1912_unet_cqt_oct_attention_adaLN_2.yaml
+# and conf/network/paper_1912_unet_cqt_oct_attention_44k_2.yaml.
+CHECKPOINTS = {
+    "MAESTRO (solo piano)": {
+        "filename": "maestro_22k_8s-750000.pt",
+        "sample_rate": 22050,
+        "audio_len": 184184,
+        "network": {
+            "use_fencoding": False,
+            "use_norm": True,
+            "depth": 7,
+            "emb_dim": 256,
+            "Ns": [64, 96, 96, 128, 128, 256, 256],
+            "attention_layers": [0, 0, 0, 0, 1, 1, 1, 1],
+            "Ss": [2, 2, 2, 2, 2, 2, 2],
+            "num_dils": [2, 3, 4, 5, 6, 7, 7],
+            "cqt": {"window": "kaiser", "beta": 1, "num_octs": 7, "bins_per_oct": 64},
+            "bottleneck_type": "res_dil_convs",
+            "num_bottleneck_layers": 1,
+            "attention_dict": {
+                "num_heads": 8, "attn_dropout": 0.0, "bias_qkv": False, "N": 0,
+                "rel_pos_num_buckets": 32, "rel_pos_max_distance": 64,
+                "use_rel_pos": False, "Nproj": 8,
+            },
+        },
+    },
+    "MusicNet (mixed classical instrumentation)": {
+        "filename": "musicnet_44k_4s-560000.pt",
+        "sample_rate": 44100,
+        "audio_len": 184184,
+        "network": {
+            "use_fencoding": False,
+            "use_norm": True,
+            "depth": 8,
+            "emb_dim": 256,
+            "Ns": [64, 64, 96, 96, 128, 128, 256, 256],
+            "attention_layers": [0, 0, 0, 0, 0, 1, 1, 1, 1],
+            "Ss": [2, 2, 2, 2, 2, 2, 2],
+            "num_dils": [2, 3, 4, 5, 6, 7, 8, 8],
+            "cqt": {"window": "kaiser", "beta": 1, "num_octs": 8, "bins_per_oct": 64},
+            "bottleneck_type": "res_dil_convs",
+            "num_bottleneck_layers": 1,
+            "attention_dict": {
+                "num_heads": 8, "attn_dropout": 0.0, "bias_qkv": False, "N": 0,
+                "rel_pos_num_buckets": 32, "rel_pos_max_distance": 64,
+                "use_rel_pos": False, "Nproj": 8,
+            },
+        },
+    },
+}
+
+# conf/diff_params/edm.yaml -- the training-time noise schedule, read by EDM.__init__.
+BASE_DIFF_PARAMS = {
+    "sigma_data": 0.063,
+    "sigma_min": 1e-5,
+    "sigma_max": 10,
+    "P_mean": -1.2,
+    "P_std": 1.2,
+    "ro": 13,
+    "ro_train": 10,
+    "Schurn": 5,
+    "Snoise": 1,
+    "Stmin": 0,
+    "Stmax": 50,
+    "aweighting": {"use_aweighting": False, "ntaps": 101},
+}
+
+# conf/tester/inpainting_tester.yaml -- sampling-time overrides, read by Sampler.
+# This is the same config the repo's own notebooks/demo_inpainting_spectrogram.ipynb
+# uses for both checkpoints (it only swaps exp= and network=).
+TESTER_DIFF_PARAMS = {
+    "same_as_training": False,
+    "sigma_data": 0.063,
+    "sigma_min": 1e-4,
+    "sigma_max": 1,
+    "P_mean": -1.2,
+    "P_std": 1.2,
+    "ro": 13,
+    "ro_train": 13,
+    "Schurn": 10,
+    "Snoise": 1.0,
+    "Stmin": 0,
+    "Stmax": 50,
+}
+
+_model_cache = {}
+_model_locks = {name: threading.Lock() for name in CHECKPOINTS}
+
+
+def build_args(checkpoint_name, steps, guidance_strength):
+    """Assemble the config tree the network, EDM and Sampler classes expect, mirroring the repo's hydra config composition for one checkpoint."""
+    ckpt_cfg = CHECKPOINTS[checkpoint_name]
+    return OmegaConf.create({
+        "exp": {
+            "sample_rate": ckpt_cfg["sample_rate"],
+            "audio_len": ckpt_cfg["audio_len"],
+        },
+        "network": ckpt_cfg["network"],
+        "diff_params": BASE_DIFF_PARAMS,
+        "tester": {
+            "T": steps,
+            "order": 2,
+            "filter_out_cqt_DC_Nyq": True,
+            "posterior_sampling": {"xi": guidance_strength, "norm": 2, "smoothl1_beta": 1},
+            "data_consistency": {"use": True, "type": "always", "smooth": True, "hann_size": 50},
+            "diff_params": TESTER_DIFF_PARAMS,
+        },
+    })
+
+
+def get_network(checkpoint_name):
+    """Build the network directly on DEVICE and load its checkpoint, caching per checkpoint name.
+
+    Building fresh on DEVICE (rather than building on CPU and calling .to(DEVICE)
+    later) matters here: the network's CQT transform allocates its kernels for a
+    specific device at construction time, so a later .to() would leave it stale.
+    """
+    with _model_locks[checkpoint_name]:
+        if checkpoint_name not in _model_cache:
+            ckpt_cfg = CHECKPOINTS[checkpoint_name]
+            args = build_args(checkpoint_name, steps=35, guidance_strength=0.25)
+            network = Unet_CQT_oct_with_attention(args, DEVICE).to(DEVICE)
+            ckpt_path = hf_hub_download(repo_id=CHECKPOINT_REPO, filename=ckpt_cfg["filename"])
+            state_dict = torch.load(ckpt_path, map_location=DEVICE, weights_only=False)
+            network.load_state_dict(state_dict["ema"])
+            network.eval()
+            _model_cache[checkpoint_name] = network
+    return _model_cache[checkpoint_name]
+
+
+def prepare_segment(audio, sr_in, target_sr, target_len):
+    """Resample to the checkpoint's rate, then center-crop or zero-pad to its fixed segment length.
+
+    Adapted from Tester.resample_audio in testing/tester_inpainting.py, extended
+    to pad shorter clips: the repo only ever sees pre-cropped test-set segments,
+    but an interactive demo can be handed audio of any length.
+    """
+    if sr_in != target_sr:
+        audio = torchaudio.functional.resample(audio, sr_in, target_sr)
+    n = audio.shape[-1]
+    if n > target_len:
+        start = (n - target_len) // 2
+        audio = audio[..., start:start + target_len]
+    elif n < target_len:
+        audio = torch.nn.functional.pad(audio, (0, target_len - n))
+    return audio
+
+
+def build_gap_mask(audio_len, gap_ms, position_ms, sample_rate, device):
+    """Binary mask with a single gap, matching Tester.prepare_mask's 'long' mode. position_ms of -1 centers the gap, per the repo's own 'start_gap_idx: None' convention."""
+    gap = int(gap_ms * sample_rate / 1000)
+    gap = max(1, min(gap, audio_len - 1))
+    if position_ms < 0:
+        start = audio_len // 2 - gap // 2
+    else:
+        start = int(position_ms * sample_rate / 1000)
+        start = max(0, min(start, audio_len - gap))
+    mask = torch.ones((1, audio_len), device=device)
+    mask[..., start:start + gap] = 0
+    return mask
+
+
+@spaces.GPU
+@torch.inference_mode()
+def process_fn(input_audio_path, checkpoint_choice, gap_length_ms, gap_position_ms, steps, guidance_strength):
+    """Fill a gap in the input audio using the selected diffusion model."""
+    ckpt_cfg = CHECKPOINTS[checkpoint_choice]
+    sample_rate = ckpt_cfg["sample_rate"]
+    audio_len = ckpt_cfg["audio_len"]
+
+    audio_np, sr_in = sf.read(input_audio_path, dtype="float32", always_2d=False)
+    if audio_np.ndim > 1:
+        audio_np = audio_np.mean(axis=-1)  # the model is mono-only
+    audio = torch.from_numpy(audio_np).unsqueeze(0).to(DEVICE)
+
+    segment = prepare_segment(audio, sr_in, sample_rate, audio_len)
+    mask = build_gap_mask(audio_len, gap_length_ms, gap_position_ms, sample_rate, DEVICE)
+
+    network = get_network(checkpoint_choice)
+    args = build_args(checkpoint_choice, steps, guidance_strength)
+    diff_params = EDM(args)
+    sampler = Sampler(network, diff_params, args)
+
+    y_masked = segment * mask
+    pred = sampler.predict_inpainting(y_masked, mask)
+
+    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
+        out_path = f.name
+    sf.write(out_path, pred.squeeze(0).cpu().numpy(), sample_rate)
+    return out_path
+
+
+model_card = ModelCard(
+    name="Diffusion-Based Audio Inpainting",
+    description=(
+        "Fills a gap in an audio clip using a diffusion model trained on either solo "
+        "piano (MAESTRO) or mixed classical instrumentation (MusicNet). The input is "
+        "resampled to the chosen model's rate and center-cropped or zero-padded to its "
+        "fixed segment length (about 8.4s for MAESTRO, 4.2s for MusicNet) before the gap "
+        "is filled."
+    ),
+    author="Eloi Moliner, Vesa Välimäki",
+    tags=["audio-inpainting", "diffusion", "restoration"],
+)
+
+
+with gr.Blocks() as demo:
+    input_components = [
+        gr.Audio(type="filepath", label="Input Audio").harp_required(True),
+        gr.Dropdown(
+            choices=list(CHECKPOINTS.keys()),
+            value="MAESTRO (solo piano)",
+            label="Model",
+            info="Which checkpoint to inpaint with.",
+        ),
+        gr.Slider(minimum=50, maximum=3000, step=50, value=1500, label="Gap Length (ms)",
+                  info="Length of the gap to fill (default: 1500ms, per repo config)"),
+        gr.Number(value=-1, label="Gap Position (ms)",
+                  info="Where the gap starts, from the beginning of the processed segment. -1 = centered (default, per repo config)"),
+        gr.Slider(minimum=10, maximum=150, step=1, value=35, label="Processing Steps",
+                  info="Diffusion sampling steps (default: T=35, per repo config)"),
+        gr.Slider(minimum=0.0, maximum=1.0, step=0.05, value=0.25, label="Reconstruction Guidance",
+                  info="How strongly the output is pulled toward the surrounding audio (default: ξ=0.25, per repo config)"),
+    ]
+    output_components = [
+        gr.Audio(type="filepath", label="Output Audio").set_info(
+            "The processed segment with the gap filled in."
+        ),
+    ]
+
+    build_endpoint(
+        model_card=model_card,
+        input_components=input_components,
+        output_components=output_components,
+        process_fn=process_fn,
+    )
+
+if __name__ == "__main__":
+    demo.queue().launch(pwa=True)

diff_params/edm.py

@@ -0,0 +1,194 @@
+import torch
+import numpy as np
+
+import utils.training_utils as utils
+
+
+class EDM():
+    """
+        Definition of most of the diffusion parameterization, following ( Karras et al., "Elucidating...", 2022)
+    """
+
+    def __init__(self, args):
+        """
+        Args:
+            args (dictionary): hydra arguments
+            sigma_data (float): 
+        """
+        self.args=args
+        self.sigma_min = args.diff_params.sigma_min
+        self.sigma_max =args.diff_params.sigma_max
+        self.P_mean=args.diff_params.P_mean
+        self.P_std=args.diff_params.P_std
+        self.ro=args.diff_params.ro
+        self.ro_train=args.diff_params.ro_train
+        self.sigma_data=args.diff_params.sigma_data #depends on the training data!! precalculated variance of the dataset
+        #parameters stochastic sampling
+        self.Schurn=args.diff_params.Schurn
+        self.Stmin=args.diff_params.Stmin
+        self.Stmax=args.diff_params.Stmax
+        self.Snoise=args.diff_params.Snoise
+
+        #perceptual filter
+        if self.args.diff_params.aweighting.use_aweighting:
+            self.AW=utils.FIRFilter(filter_type="aw", fs=args.exp.sample_rate, ntaps=self.args.diff_params.aweighting.ntaps)
+
+       
+
+    def get_gamma(self, t): 
+        """
+        Get the parameter gamma that defines the stochasticity of the sampler
+        Args
+            t (Tensor): shape: (N_steps, ) Tensor of timesteps, from which we will compute gamma
+        """
+        N=t.shape[0]
+        gamma=torch.zeros(t.shape).to(t.device)
+        
+        #If desired, only apply stochasticity between a certain range of noises Stmin is 0 by default and Stmax is a huge number by default. (Unless these parameters are specified, this does nothing)
+        indexes=torch.logical_and(t>self.Stmin , t<self.Stmax)
+         
+        #We use Schurn=5 as the default in our experiments
+        gamma[indexes]=gamma[indexes]+torch.min(torch.Tensor([self.Schurn/N, 2**(1/2) -1]))
+        
+        return gamma
+
+    def create_schedule(self,nb_steps):
+        """
+        Define the schedule of timesteps
+        Args:
+           nb_steps (int): Number of discretized steps
+        """
+        i=torch.arange(0,nb_steps+1)
+        t=(self.sigma_max**(1/self.ro) +i/(nb_steps-1) *(self.sigma_min**(1/self.ro) - self.sigma_max**(1/self.ro)))**self.ro
+        t[-1]=0
+        return t
+
+
+    def sample_ptrain(self,N):
+        """
+        For training, getting t as a normal distribution, folowing Karras et al. 
+        I'm not using this
+        Args:
+            N (int): batch size
+        """
+        lnsigma=np.random.randn(N)*self.P_std +self.P_mean
+        return np.clip(np.exp(lnsigma),self.sigma_min, self.sigma_max) #not sure if clipping here is necessary, but makes sense to me
+    
+    def sample_ptrain_safe(self,N):
+        """
+        For training, getting  t according to the same criteria as sampling
+        Args:
+            N (int): batch size
+        """
+        a=torch.rand(N)
+        t=(self.sigma_max**(1/self.ro_train) +a *(self.sigma_min**(1/self.ro_train) - self.sigma_max**(1/self.ro_train)))**self.ro_train
+        return t
+
+    def sample_prior(self,shape,sigma):
+        """
+        Just sample some gaussian noise, nothing more
+        Args:
+            shape (tuple): shape of the noise to sample, something like (B,T)
+            sigma (float): noise level of the noise
+        """
+        n=torch.randn(shape).to(sigma.device)*sigma
+        return n
+
+    def cskip(self, sigma):
+        """
+        Just one of the preconditioning parameters
+        Args:
+            sigma (float): noise level (equal to timestep is sigma=t, which is our default)
+        
+        """
+        return self.sigma_data**2 *(sigma**2+self.sigma_data**2)**-1
+
+    def cout(self,sigma ):
+        """
+        Just one of the preconditioning parameters
+        Args:
+            sigma (float): noise level (equal to timestep is sigma=t, which is our default)
+        """
+        return sigma*self.sigma_data* (self.sigma_data**2+sigma**2)**(-0.5)
+
+    def cin(self, sigma):
+        """
+        Just one of the preconditioning parameters
+        Args:
+            sigma (float): noise level (equal to timestep is sigma=t, which is our default)
+        """
+        return (self.sigma_data**2+sigma**2)**(-0.5)
+
+    def cnoise(self,sigma ):
+        """
+        preconditioning of the noise embedding
+        Args:
+            sigma (float): noise level (equal to timestep is sigma=t, which is our default)
+        """
+        return (1/4)*torch.log(sigma)
+
+    def lambda_w(self,sigma):
+        return (sigma*self.sigma_data)**(-2) * (self.sigma_data**2+sigma**2)
+        
+    def denoiser(self, xn , net, sigma):
+        """
+        This method does the whole denoising step, which implies applying the model and the preconditioning
+        Args:
+            x (Tensor): shape: (B,T) Intermediate noisy latent to denoise
+            model (nn.Module): Model of the denoiser
+            sigma (float): noise level (equal to timestep is sigma=t, which is our default)
+        """
+        if len(sigma.shape)==1:
+            sigma=sigma.unsqueeze(-1)
+        cskip=self.cskip(sigma)
+        cout=self.cout(sigma)
+        cin=self.cin(sigma)
+        cnoise=self.cnoise(sigma)
+
+        return cskip * xn +cout*net(cin*xn, cnoise)  #this will crash because of broadcasting problems, debug later!
+
+    def prepare_train_preconditioning(self, x, sigma):
+        #weight=self.lambda_w(sigma)
+        #Is calling the denoiser here a good idea? Maybe it would be better to apply directly the preconditioning as in the paper, even though Karras et al seem to do it this way in their code
+        print(x.shape)
+        noise=self.sample_prior(x.shape,sigma)
+
+        cskip=self.cskip(sigma)
+        cout=self.cout(sigma)
+        cin=self.cin(sigma)
+        cnoise=self.cnoise(sigma)
+
+        target=(1/cout)*(x-cskip*(x+noise))
+
+        return cin*(x+noise), target, cnoise
+
+
+    def loss_fn(self, net, x):
+        """
+        Loss function, which is the mean squared error between the denoised latent and the clean latent
+        Args:
+            net (nn.Module): Model of the denoiser
+            x (Tensor): shape: (B,T) Intermediate noisy latent to denoise
+            sigma (float): noise level (equal to timestep is sigma=t, which is our default)
+        """
+        sigma=self.sample_ptrain_safe(x.shape[0]).unsqueeze(-1).to(x.device)
+
+        input, target, cnoise= self.prepare_train_preconditioning(x, sigma)
+        estimate=net(input,cnoise)
+        
+        error=(estimate-target)
+
+        try:
+            #this will only happen if the model is cqt-based, if it crashes it is normal
+            if self.args.net.use_cqt_DC_correction:
+                error=net.CQTransform.apply_hpf_DC(error) #apply the DC correction to the error as we dont want to propagate the DC component of the error as the network is discarding it. It also applies for the nyquit frequency, but this is less critical.
+        except:
+            pass 
+
+        #APPLY A-WEIGHTING
+        if self.args.diff_params.aweighting.use_aweighting:
+            error=self.AW(error)
+
+        #here we have the chance to apply further emphasis to the error, as some kind of perceptual frequency weighting could be
+        return error**2, sigma
+

model.json

@@ -0,0 +1,28 @@
+{
+  "name": "audio-inpainting-diffusion",
+  "package_dir": ".",
+  "entry_point": "testing.edm_sampler_inpainting.Sampler",
+  "checkpoints": [
+    {
+      "label": "MAESTRO (solo piano)",
+      "repo": "Eloimoliner/audio-inpainting-diffusion",
+      "filename": "maestro_22k_8s-750000.pt",
+      "size_mb": 2984,
+      "exp_config": "conf/exp/maestro22k_8s.yaml",
+      "network_config": "conf/network/paper_1912_unet_cqt_oct_attention_adaLN_2.yaml",
+      "sample_rate": 22050,
+      "audio_len_samples": 184184
+    },
+    {
+      "label": "MusicNet (mixed classical instrumentation)",
+      "repo": "Eloimoliner/audio-inpainting-diffusion",
+      "filename": "musicnet_44k_4s-560000.pt",
+      "size_mb": 3884,
+      "exp_config": "conf/exp/musicnet44k_4s.yaml",
+      "network_config": "conf/network/paper_1912_unet_cqt_oct_attention_44k_2.yaml",
+      "sample_rate": 44100,
+      "audio_len_samples": 184184
+    }
+  ],
+  "notes": "Both checkpoints share the same tester/diff_params hyperparameters (conf/tester/inpainting_tester.yaml, conf/diff_params/edm.yaml) -- this is exactly the hydra.compose() config composition used in the repo's own notebooks/demo_inpainting_spectrogram.ipynb, which swaps only exp= and network= per checkpoint. Skipped the third HF checkpoint (22k_8s-850000.pt, the no-attention MAESTRO ablation variant) per user choice."
+}

networks/unet_cqt_oct_with_projattention_adaLN_2.py

@@ -0,0 +1,882 @@
+import torch.nn as nn
+import numpy as np
+import torch.nn.functional as F
+import math as m
+import torch
+#import torchaudio
+torch.pi = torch.acos(torch.zeros(1)).item() * 2 # which is 3.1415927410125732
+
+from cqt_nsgt_pytorch import CQT_nsgt
+import torchaudio
+import einops
+import math
+
+"""
+As similar as possible to the original CQTdiff architecture, but using the octave-base representation of the CQT
+This should be more memory efficient, and also more efficient in terms of computation, specially when using higher sampling rates.
+I am expecting similar performance to the original CQTdiff architecture, but faster. 
+Perhaps the fact that I am using powers of 2 for the time sizes is critical for transient reconstruction. I should thest CQT matrix model with powers of 2, this requires modifying the CQT_nsgt_pytorch.py file.
+"""
+def weight_init(shape, mode, fan_in, fan_out):
+    if mode == 'xavier_uniform': return np.sqrt(6 / (fan_in + fan_out)) * (torch.rand(*shape) * 2 - 1)
+    if mode == 'xavier_normal':  return np.sqrt(2 / (fan_in + fan_out)) * torch.randn(*shape)
+    if mode == 'kaiming_uniform': return np.sqrt(3 / fan_in) * (torch.rand(*shape) * 2 - 1)
+    if mode == 'kaiming_normal':  return np.sqrt(1 / fan_in) * torch.randn(*shape)
+    raise ValueError(f'Invalid init mode "{mode}"')
+
+class Linear(torch.nn.Module):
+    def __init__(self, in_features, out_features, bias=True, init_mode='kaiming_normal', init_weight=1, init_bias=0):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        init_kwargs = dict(mode=init_mode, fan_in=in_features, fan_out=out_features)
+        self.weight = torch.nn.Parameter(weight_init([out_features, in_features], **init_kwargs) * init_weight)
+        self.bias = torch.nn.Parameter(weight_init([out_features], **init_kwargs) * init_bias) if bias else None
+
+    def forward(self, x):
+        x = x @ self.weight.to(x.dtype).t()
+        if self.bias is not None:
+            x = x.add_(self.bias.to(x.dtype))
+        return x
+
+class Conv1d(torch.nn.Module):
+    def __init__(self,
+        in_channels, out_channels, kernel=1, bias=False, dilation=1,
+        init_mode='kaiming_normal', init_weight=1, init_bias=0,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.dilation = dilation
+        init_kwargs = dict(mode=init_mode, fan_in=in_channels*kernel, fan_out=out_channels*kernel)
+        self.weight = torch.nn.Parameter(weight_init([out_channels, in_channels, kernel], **init_kwargs) * init_weight) 
+        self.bias = torch.nn.Parameter(weight_init([out_channels], **init_kwargs) * init_bias) if bias else None
+
+    def forward(self, x):
+        w = self.weight.to(x.dtype) if self.weight is not None else None
+        b = self.bias.to(x.dtype) if self.bias is not None else None
+        w_pad = w.shape[-1] // 2 if w is not None else 0
+        #f_pad = (f.shape[-1] - 1) // 2 if f is not None else 0
+        #print(x.shape, w.shape)
+        if w is not None:
+                x = torch.nn.functional.conv1d(x, w, padding="same", dilation=self.dilation)
+        if b is not None:
+            x = x.add_(b.reshape(1, -1, 1))
+        return x
+class Conv2d(torch.nn.Module):
+    def __init__(self,
+        in_channels, out_channels, kernel=(1,1), bias=False, dilation=1,
+        init_mode='kaiming_normal', init_weight=1, init_bias=0,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.dilation = dilation
+        init_kwargs = dict(mode=init_mode, fan_in=in_channels*kernel[0]*kernel[1], fan_out=out_channels*kernel[0]*kernel[1])
+        self.weight = torch.nn.Parameter(weight_init([out_channels, in_channels, kernel[0], kernel[1]], **init_kwargs) * init_weight) 
+        self.bias = torch.nn.Parameter(weight_init([out_channels], **init_kwargs) * init_bias) if bias else None
+
+    def forward(self, x):
+        w = self.weight.to(x.dtype) if self.weight is not None else None
+        b = self.bias.to(x.dtype) if self.bias is not None else None
+        w_pad = w.shape[-1] // 2 if w is not None else 0
+        #f_pad = (f.shape[-1] - 1) // 2 if f is not None else 0
+        if w is not None:
+                x = torch.nn.functional.conv2d(x, w, padding="same", dilation=self.dilation)
+        if b is not None:
+            x = x.add_(b.reshape(1, -1, 1, 1))
+        return x
+
+class LayerScale(nn.Module):
+    """Layer scale from [Touvron et al 2021] (https://arxiv.org/pdf/2103.17239.pdf).
+    This rescales diagonaly residual outputs close to 0 initially, then learnt.
+    """
+
+    def __init__(self, channels: int, init: float = 1e-4, channel_last=True):
+        """
+        channel_last = False corresponds to (B, C, T) tensors
+        channel_last = True corresponds to (T, B, C) tensors
+        """
+        super().__init__()
+        self.channel_last = channel_last
+        self.scale = nn.Parameter(torch.zeros(channels, requires_grad=True))
+        self.scale.data[:] = init
+
+    def forward(self, x):
+        if self.channel_last:
+            return self.scale * x
+        else:
+            return self.scale[:, None] * x
+
+class BiasFreeLayerNorm(nn.Module):
+
+    def __init__(self, num_features, eps=1e-7):
+        super(BiasFreeLayerNorm, self).__init__()
+        self.gamma = nn.Parameter(torch.ones(1,1,num_features))
+        #self.beta = nn.Parameter(torch.zeros(1,num_features,1,1))
+        #self.beta = torch.zeros(1,num_features,1,1)
+        self.eps = eps
+
+    def forward(self, x):
+        N, T, C = x.size()
+        #x = x.view(N, self.num_groups ,-1,H,W)
+        #x=einops.rearrange(x, 'n t c -> n (t c)')
+        #mean = x.mean(-1, keepdim=True)
+        #var = x.var(-1, keepdim=True)
+
+        std=x.std(-1, keepdim=True) #reduce over channels and time
+        #var = x.var(-1, keepdim=True)
+
+        ## normalize
+        x = (x) / (std+self.eps)
+        # normalize
+        #x=einops.rearrange(x, 'n (t c) -> n t c', t=T)
+        #x = x.view(N,C,H,W)
+        return x * self.gamma
+
+class BiasFreeGroupNorm(nn.Module):
+
+    def __init__(self, num_features, num_groups=32, eps=1e-7):
+        super(BiasFreeGroupNorm, self).__init__()
+        self.gamma = nn.Parameter(torch.ones(1,num_features,1,1))
+        #self.beta = nn.Parameter(torch.zeros(1,num_features,1,1))
+        #self.beta = torch.zeros(1,num_features,1,1)
+        self.num_groups = num_groups
+        self.eps = eps
+
+    def forward(self, x):
+        N, C, F, T = x.size()
+        #x = x.view(N, self.num_groups ,-1,H,W)
+        gc=C//self.num_groups
+        x=einops.rearrange(x, 'n (g gc) f t -> n g (gc f t)', g=self.num_groups, gc=gc)
+        #mean = x.mean(-1, keepdim=True)
+        #var = x.var(-1, keepdim=True)
+
+        std=x.std(-1, keepdim=True) #reduce over channels and time
+        #var = x.var(-1, keepdim=True)
+
+        ## normalize
+        x = (x) / (std+self.eps)
+        # normalize
+        x=einops.rearrange(x, 'n g (gc f t) -> n (g gc) f t', g=self.num_groups, gc=gc, f=F, t=T)
+        #x = x.view(N,C,H,W)
+        return x * self.gamma
+
+
+
+class RFF_MLP_Block(nn.Module):
+    """
+        Encoder of the noise level embedding
+        Consists of:
+            -Random Fourier Feature embedding
+            -MLP
+    """
+    def __init__(self, emb_dim=512, rff_dim=32, init=None):
+        super().__init__()
+        self.RFF_freq = nn.Parameter(
+            16 * torch.randn([1, rff_dim]), requires_grad=False)
+        self.MLP = nn.ModuleList([
+            Linear(2*rff_dim, 128, **init),
+            Linear(128, 256, **init),
+            Linear(256, emb_dim, **init),
+        ])
+
+    def forward(self, sigma):
+        """
+        Arguments:
+          sigma:
+              (shape: [B, 1], dtype: float32)
+
+        Returns:
+          x: embedding of sigma
+              (shape: [B, 512], dtype: float32)
+        """
+        x = self._build_RFF_embedding(sigma)
+        for layer in self.MLP:
+            x = F.relu(layer(x))
+        return x
+
+    def _build_RFF_embedding(self, sigma):
+        """
+        Arguments:
+          sigma:
+              (shape: [B, 1], dtype: float32)
+        Returns:
+          table:
+              (shape: [B, 64], dtype: float32)
+        """
+        freqs = self.RFF_freq
+        table = 2 * np.pi * sigma * freqs
+        table = torch.cat([torch.sin(table), torch.cos(table)], dim=1)
+        return table
+
+class AddFreqEncodingRFF(nn.Module):
+    '''
+    [B, T, F, 2] => [B, T, F, 12]  
+    Generates frequency positional embeddings and concatenates them as 10 extra channels
+    This function is optimized for F=1025
+    '''
+    def __init__(self, f_dim, N):
+        super(AddFreqEncodingRFF, self).__init__()
+        self.N=N
+        self.RFF_freq = nn.Parameter(
+            16 * torch.randn([1, N]), requires_grad=False)
+
+
+        self.f_dim=f_dim #f_dim is fixed
+        embeddings=self.build_RFF_embedding()
+        self.embeddings=nn.Parameter(embeddings, requires_grad=False) 
+
+        
+    def build_RFF_embedding(self):
+        """
+        Returns:
+          table:
+              (shape: [C,F], dtype: float32)
+        """
+        freqs = self.RFF_freq
+        #freqs = freqs.to(device=torch.device("cuda"))
+        freqs=freqs.unsqueeze(-1) # [1, 32, 1]
+
+        self.n=torch.arange(start=0,end=self.f_dim)
+        self.n=self.n.unsqueeze(0).unsqueeze(0)  #[1,1,F]
+
+        table = 2 * np.pi * self.n * freqs
+
+        #print(freqs.shape, x.shape, table.shape)
+        table = torch.cat([torch.sin(table), torch.cos(table)], dim=1) #[1,32,F]
+
+        return table
+    
+
+    def forward(self, input_tensor):
+
+        #print(input_tensor.shape)
+        batch_size_tensor = input_tensor.shape[0]  # get batch size
+        time_dim = input_tensor.shape[-1]  # get time dimension
+
+        fembeddings_2 = torch.broadcast_to(self.embeddings, [batch_size_tensor, time_dim,self.N*2, self.f_dim])
+        fembeddings_2=fembeddings_2.permute(0,2,3,1)
+    
+        
+        #print(input_tensor.shape, fembeddings_2.shape)
+        return torch.cat((input_tensor,fembeddings_2),1)  
+
+
+class RelativePositionBias(nn.Module):
+    def __init__(self, num_buckets: int, max_distance: int, num_heads: int):
+        super().__init__()
+        self.num_buckets = num_buckets
+        self.max_distance = max_distance
+        self.num_heads = num_heads
+        self.relative_attention_bias = nn.Embedding(num_buckets, num_heads)
+
+    @staticmethod
+    def _relative_position_bucket(
+        relative_position, num_buckets: int, max_distance: int
+    ):
+        num_buckets //= 2
+        ret = (relative_position >= 0).to(torch.long) * num_buckets
+        n = torch.abs(relative_position)
+
+        max_exact = num_buckets // 2
+        is_small = n < max_exact
+
+        val_if_large = (
+            max_exact
+            + (
+                torch.log(n.float() / max_exact)
+                / math.log(max_distance / max_exact)
+                * (num_buckets - max_exact)
+            ).long()
+        )
+        val_if_large = torch.min(
+            val_if_large, torch.full_like(val_if_large, num_buckets - 1)
+        )
+
+        ret += torch.where(is_small, n, val_if_large)
+        return ret
+
+    def forward(self, num_queries: int, num_keys: int):
+        i, j, device = num_queries, num_keys, self.relative_attention_bias.weight.device
+        q_pos = torch.arange(j - i, j, dtype=torch.long, device=device)
+        k_pos = torch.arange(j, dtype=torch.long, device=device)
+        rel_pos = einops.rearrange(k_pos, "j -> 1 j") - einops.rearrange(q_pos, "i -> i 1")
+
+        relative_position_bucket = self._relative_position_bucket(
+            rel_pos, num_buckets=self.num_buckets, max_distance=self.max_distance
+        )
+
+        bias = self.relative_attention_bias(relative_position_bucket)
+        bias = einops.rearrange(bias, "m n h -> 1 h m n")
+        return bias
+
+class TimeAttentionBlock(nn.Module):
+    def __init__(self, Nin,attention_dict, init, init_zero, Fdim) -> None:
+        super().__init__()
+        #NA=attention_dict.N
+        self.attention_dict=attention_dict
+        self.Fdim=Fdim
+        N=attention_dict.num_heads*Fdim 
+        self.qk = Conv1d(N, N*2, bias=self.attention_dict.bias_qkv, **init )
+        self.proj_in=Conv2d(Nin, attention_dict.num_heads, (1,1), bias=False, **init)
+        self.proj_out=Conv2d(attention_dict.num_heads, Nin, (1,1), bias=False, **init)
+        #not sure if a bias is a good idea here
+        #self.v = Conv2d(N, N*2, (1,1), bias=False,**init )
+        #I think that as long as the main signal path layers are bias free, we should be safe from artifacts
+        #self.proj = Conv1d(NA, NA, 1, bias=False, **init)
+
+        self.scale=(N/self.attention_dict.num_heads)**-0.5
+        self.use_rel_pos = self.attention_dict.use_rel_pos
+        if self.use_rel_pos:
+            self.rel_pos = RelativePositionBias(
+                num_buckets=attention_dict.rel_pos_num_buckets,
+                max_distance=attention_dict.rel_pos_max_distance,
+                num_heads=attention_dict.num_heads,
+            )
+
+    def forward(self, x):
+        #shape of x is [batch, C,F, T]
+
+        #we need shape: [batch, heads, T, D]
+        #with heands on different (original) channels
+        #print(x.shape, self.Fdim)
+
+        x=self.proj_in(x) #reduce the C dimensionality
+
+        #print(x.shape, self.Fdim)
+        #normalize everyting (easy)
+
+        #split into heads
+        x=einops.rearrange(x, "b h f t -> b (h f) t")
+
+        v=einops.rearrange(x,"b (h f) t -> b h t f", f=self.Fdim) #identity layer for the values
+
+        qk=self.qk(x) #linear layer
+        #for now, f are features (all merged) but still represents frequency
+
+        qk=einops.rearrange(qk, "b (h d) t -> b h t d", h=self.attention_dict.num_heads)
+        q,k=qk.chunk(2,dim=-1)
+
+        #print("qk",q.shape, k.shape)
+        sim = torch.einsum("... n d, ... m d -> ... n m", q, k)
+        #print("sim",sim.shape)
+        sim = (sim + self.rel_pos(*sim.shape[-2:])) if self.use_rel_pos else sim
+        #print("sim",sim.shape)
+        sim = sim * self.scale
+        # Get attention matrix with softmax
+        attn = sim.softmax(dim=-1)
+        # Compute values
+        #print("attn",attn.shape, v.shape)
+        out = torch.einsum("... n m, ... m d -> ... n d", attn, v)
+
+        #print("out",out.shape)
+        out = einops.rearrange(out, "b h t f -> b h f t", f=self.Fdim)
+        #out = einops.rearrange(out, "b (h f) t -> b h f t", f=self.Fdim)
+
+        #reverse step
+        out=self.proj_out(out)
+
+        return out
+        
+class ResnetBlock(nn.Module):
+    def __init__(
+        self,
+        dim,
+        dim_out,
+        use_norm=True,
+        num_dils = 6,
+        bias=False,
+        kernel_size=(5,3),
+        emb_dim=512,
+        proj_place='before', #using 'after' in the decoder out blocks
+        init=None,
+        init_zero=None,
+        attention_dict=None,
+        Fdim=128, #number of frequency bins
+    ):
+        super().__init__()
+
+        self.bias=bias
+        self.use_norm=use_norm
+        self.num_dils=num_dils
+        self.proj_place=proj_place
+        self.Fdim=Fdim
+
+        if self.proj_place=='before':
+            #dim_out is the block dimension
+            N=dim_out
+        else:
+            #dim in is the block dimension
+            N=dim
+            self.proj_out = Conv2d(N, dim_out,   bias=bias, **init) if N!=dim_out else nn.Identity() #linear projection
+
+        self.res_conv = Conv2d(dim, dim_out, bias=bias, **init) if dim!= dim_out else nn.Identity() #linear projection
+        self.proj_in = Conv2d(dim, N,   bias=bias, **init) if dim!=N else nn.Identity()#linear projection
+
+
+
+        self.H=nn.ModuleList()
+        self.affine=nn.ModuleList()
+        self.gate=nn.ModuleList()
+        if self.use_norm:
+            self.norm=nn.ModuleList()
+
+        for i in range(self.num_dils):
+
+            if self.use_norm:
+                self.norm.append(BiasFreeGroupNorm(N,8))
+
+            self.affine.append(Linear(emb_dim, N, **init))
+            self.gate.append(Linear(emb_dim, N, **init_zero))
+            #self.H.append(Gated_residual_layer(dim_out, (5,3), (2**i,1), bias=bias)) #sometimes I changed this 1,5 to 3,5. be careful!!! (in exp 80 as far as I remember)
+            self.H.append(Conv2d(N,N,    
+                                    kernel=kernel_size,
+                                    dilation=(2**i,1),
+                                    bias=bias, **init)) #freq convolution (dilated) 
+
+        self.attention_dict=attention_dict
+        if self.attention_dict is not None:
+            #NA=self.attention_dict.N
+            self.norm2=BiasFreeGroupNorm(N,8)
+            self.affine2=Linear(emb_dim, N, **init)
+            self.gate2=Linear(emb_dim, N, **init_zero)
+            #self.norm2 = BiasFreeGroupNorm(N,8)
+            #self.proj_attn_in = Conv1d(N*Fdim, NA,   bias=bias, **init) if (N*Fdim)!=NA else nn.Identity()#linear projection
+            #self.proj_attn_out = Conv1d(NA, N*Fdim,   bias=bias, **init_zero) if NA!=(N*Fdim) else nn.Identity() #linear projection
+            ##the attention is applied time-wise, since channels times frequency is too much, we need to reduce the dimensionality using a linear projection
+            self.attn_block=TimeAttentionBlock(N,self.attention_dict, init,init_zero, self.Fdim)
+
+
+
+    def forward(self, input_x, sigma):
+        
+        x=input_x
+
+        x=self.proj_in(x)
+
+        if self.attention_dict is not None:
+            i_x=x
+
+            gamma=self.affine2(sigma)
+            scale=self.gate2(sigma)
+
+            x=self.norm2(x)
+            x=x*(gamma.unsqueeze(2).unsqueeze(3)+1) #no bias
+
+            x=self.attn_block(x)*scale.unsqueeze(2).unsqueeze(3)
+
+            #x=(x+i_x)
+            x=(x+i_x)/(2**0.5)
+
+        for norm, affine, gate, conv in zip(self.norm, self.affine, self.gate, self.H):
+            x0=x
+            if self.use_norm:
+                x=norm(x)
+            gamma =affine(sigma)
+            scale=gate(sigma)
+
+            x=x*(gamma.unsqueeze(2).unsqueeze(3)+1) #no bias
+
+
+            x=(x0+conv(F.gelu(x))*scale.unsqueeze(2).unsqueeze(3))/(2**0.5) 
+            #x=(x0+conv(F.gelu(x))*scale.unsqueeze(2).unsqueeze(3))
+        
+        #one residual connection here after the dilated convolutions
+
+
+        if self.proj_place=='after':
+            x=self.proj_out(x)
+
+        x=(x + self.res_conv(input_x))/(2**0.5)
+
+        return x
+
+
+class AttentionOp(torch.autograd.Function):
+
+    def forward(ctx, q, k):
+        w = torch.einsum('ncq,nck->nqk', q.to(torch.float32), (k / np.sqrt(k.shape[1])).to(torch.float32)).softmax(dim=2).to(q.dtype)
+        ctx.save_for_backward(q, k, w)
+        return w
+
+    def backward(ctx, dw):
+        q, k, w = ctx.saved_tensors
+        db = torch._softmax_backward_data(grad_output=dw.to(torch.float32), output=w.to(torch.float32), dim=2, input_dtype=torch.float32)
+        dq = torch.einsum('nck,nqk->ncq', k.to(torch.float32), db).to(q.dtype) / np.sqrt(k.shape[1])
+        dk = torch.einsum('ncq,nqk->nck', q.to(torch.float32), db).to(k.dtype) / np.sqrt(k.shape[1])
+        return dq, dk
+
+_kernels = {
+    'linear':
+        [1 / 8, 3 / 8, 3 / 8, 1 / 8],
+    'cubic': 
+        [-0.01171875, -0.03515625, 0.11328125, 0.43359375,
+        0.43359375, 0.11328125, -0.03515625, -0.01171875],
+    'lanczos3': 
+        [0.003689131001010537, 0.015056144446134567, -0.03399861603975296,
+        -0.066637322306633, 0.13550527393817902, 0.44638532400131226,
+        0.44638532400131226, 0.13550527393817902, -0.066637322306633,
+        -0.03399861603975296, 0.015056144446134567, 0.003689131001010537]
+}
+class UpDownResample(nn.Module):
+    def __init__(self,
+        up=False, 
+        down=False,
+        mode_resample="T", #T for time, F for freq, TF for both
+        resample_filter='cubic', 
+        pad_mode='reflect'
+        ):
+        super().__init__()
+        assert not (up and down) #you cannot upsample and downsample at the same time
+        assert up or down #you must upsample or downsample
+        self.down=down
+        self.up=up
+        if up or down:
+            #upsample block
+            self.pad_mode = pad_mode #I think reflect is a goof choice for padding
+            self.mode_resample=mode_resample
+            if mode_resample=="T":
+                kernel_1d = torch.tensor(_kernels[resample_filter], dtype=torch.float32)
+            elif mode_resample=="F":
+                #kerel shouuld be the same
+                kernel_1d = torch.tensor(_kernels[resample_filter], dtype=torch.float32)
+            else:
+                raise NotImplementedError("Only time upsampling is implemented")
+                #TODO implement freq upsampling and downsampling
+            self.pad = kernel_1d.shape[0] // 2 - 1
+            self.register_buffer('kernel', kernel_1d)
+    def forward(self, x):
+        shapeorig=x.shape
+        #x=x.view(x.shape[0],-1,x.shape[-1])
+        x=x.view(-1,x.shape[-2],x.shape[-1]) #I have the feeling the reshape makes everything consume too much memory. There is no need to have the channel dimension different than 1. I leave it like this because otherwise it requires a contiguous() call, but I should check if the memory gain / speed, would be significant.
+        if self.mode_resample=="F":
+            x=x.permute(0,2,1)#call contiguous() here?
+
+        #print("after view",x.shape)
+        if self.down:
+            x = F.pad(x, (self.pad,) * 2, self.pad_mode)
+        elif self.up:
+            x = F.pad(x, ((self.pad + 1) // 2,) * 2, self.pad_mode)
+
+        #print("after pad",x.shape)
+
+        weight = x.new_zeros([x.shape[1], x.shape[1], self.kernel.shape[0]])
+        #print("weight",weight.shape)
+        indices = torch.arange(x.shape[1], device=x.device)
+        #print("indices",indices.shape)
+        #weight = self.kernel.to(x.device).unsqueeze(0).unsqueeze(0).expand(x.shape[1], x.shape[1], -1)
+        #print("weight",weight.shape)
+        weight[indices, indices] = self.kernel.to(weight)
+        if self.down:
+            x_out= F.conv1d(x, weight, stride=2)
+        elif self.up:
+            x_out =F.conv_transpose1d(x, weight, stride=2, padding=self.pad * 2 + 1)
+
+        if self.mode_resample=="F":
+            x_out=x_out.permute(0,2,1).contiguous()
+            return x_out.view(shapeorig[0],-1,x_out.shape[-2], shapeorig[-1])
+        else:
+            return x_out.view(shapeorig[0],-1,shapeorig[2], x_out.shape[-1])
+
+
+class Unet_CQT_oct_with_attention(nn.Module):
+    """
+        Main U-Net model based on the CQT
+    """
+    def __init__(self, args, device):
+        """
+        Args:
+            args (dictionary): hydra dictionary
+            device: torch device ("cuda" or "cpu")
+        """
+        super(Unet_CQT_oct_with_attention, self).__init__()
+        self.args=args
+        self.depth=args.network.cqt.num_octs
+        #self.depth=args.network.inner_depth+self.args.network.cqt.num_octs
+        #assert self.depth==args.network.depth, "The depth of the network should be the sum of the inner depth and the number of octaves" #make sure we are aware of the depth of the network
+
+        init = dict(init_mode='kaiming_uniform', init_weight=np.sqrt(1/3)) #same as ADM, according to edm implementation
+        init_zero = dict(init_mode='kaiming_uniform', init_weight=1e-7) #I think it is safer to initialize the last layer with a small weight, rather than zero. Breaking symmetry and all that. 
+
+
+
+        self.emb_dim=args.network.emb_dim
+        self.embedding = RFF_MLP_Block(emb_dim=args.network.emb_dim, init=init)
+        self.use_norm=args.network.use_norm
+
+        #fmax=self.args.exp.sample_rate/2
+        #self.fmin=fmax/(2**self.args.cqt.numocts)
+        self.fbins=int(self.args.network.cqt.bins_per_oct*self.args.network.cqt.num_octs) 
+        self.device=device
+        self.bins_per_oct=self.args.network.cqt.bins_per_oct
+        self.num_octs=self.args.network.cqt.num_octs
+        #self.CQTransform=CQT_nsgt(self.args.network.cqt.num_octs,self.args.network.cqt.bins_per_oct, "oct",  self.args.exp.sample_rate, self.args.exp.audio_len, device=self.device)
+        if self.args.network.cqt.window=="kaiser":
+            win=("kaiser",self.args.network.cqt.beta)
+        else:
+            win=self.args.network.cqt.window
+
+        self.CQTransform=CQT_nsgt(self.args.network.cqt.num_octs, self.args.network.cqt.bins_per_oct, mode="oct",window=win,fs=self.args.exp.sample_rate, audio_len=self.args.exp.audio_len, dtype=torch.float32, device=self.device)
+
+
+        self.f_dim=self.fbins #assuming we have thrown away the DC component and the Nyquist frequency
+
+        self.use_fencoding=self.args.network.use_fencoding
+        if self.use_fencoding:
+            N_freq_encoding=32
+    
+            self.freq_encodings=nn.ModuleList([])
+            for i in range(self.num_octs):
+                self.freq_encodings.append(AddFreqEncodingRFF(self.bins_per_oct,N_freq_encoding))
+            Nin=2*N_freq_encoding+2
+        else:
+            Nin=2
+
+        #Encoder
+        self.Ns= self.args.network.Ns
+        self.Ss= self.args.network.Ss
+
+        self.num_dils= self.args.network.num_dils #intuition: less dilations for the first layers and more for the deeper layers
+        #self.inner_num_dils=self.args.network.inner_num_dils
+        
+        self.attention_dict=self.args.network.attention_dict
+        #self.attention_Ns=self.args.network.attention_Ns
+
+        self.downsamplerT=UpDownResample(down=True, mode_resample="T")
+        #self.downsamplerF=UpDownResample(down=True, mode_resample="F")
+        self.upsamplerT=UpDownResample(up=True, mode_resample="T")
+        #self.upsamplerF=UpDownResample(up=True, mode_resample="F")
+
+        self.downs=nn.ModuleList([])
+        self.middle=nn.ModuleList([])
+        self.ups=nn.ModuleList([])
+
+        self.attention_layers=self.args.network.attention_layers
+        #sth like [0,0,0,0,0,0,1,1]
+        
+        for i in range(self.num_octs):
+            if i==0:
+                dim_in=self.Ns[i]
+                dim_out=self.Ns[i]
+            else:
+                dim_in=self.Ns[i-1]
+                dim_out=self.Ns[i]
+            if self.attention_layers[i]:
+                print("Attention layer at (down) octave {}".format(i))
+                attn_dict=self.attention_dict
+                #attn_dict.N=self.attention_Ns[i]
+                #assert attn_dict.N > 0
+            else:
+                attn_dict=None
+
+            self.downs.append(
+                               nn.ModuleList([
+                                        ResnetBlock(Nin, dim_in, self.use_norm,num_dils=1, bias=False, kernel_size=(1,1), emb_dim=self.emb_dim, init=init, init_zero=init_zero),
+                                        Conv2d(2, dim_out, kernel=(5,3), bias=False, **init),
+                                        ResnetBlock(dim_in, dim_out, self.use_norm,num_dils=self.num_dils[i], bias=False , attention_dict=attn_dict, emb_dim=self.emb_dim, init=init, init_zero=init_zero, Fdim=(i+1)*self.bins_per_oct)
+                                        ]))
+
+        if self.args.network.bottleneck_type=="res_dil_convs":
+            for i in range(self.args.network.num_bottleneck_layers):
+                if self.attention_layers[-1]:
+                    attn_dict=self.attention_dict
+                    #attn_dict.N=self.attention_Ns[-1]
+                    #assert attn_dict.N > 0
+                else:
+                    attn_dict=None
+    
+                self.middle.append(nn.ModuleList([
+                                ResnetBlock(self.Ns[-1], 2, use_norm=self.use_norm,num_dils= 1,bias=False, kernel_size=(1,1), proj_place="after", emb_dim=self.emb_dim, init=init, init_zero=init_zero),
+                                ResnetBlock(self.Ns[-1], self.Ns[-1], self.use_norm, num_dils=self.num_dils[-1], bias=False, emb_dim=self.emb_dim,attention_dict=attn_dict, init=init, init_zero=init_zero,
+                                Fdim=(self.num_octs)*self.bins_per_oct)]))
+        else:
+            raise NotImplementedError("bottleneck type not implemented")
+                        
+
+
+        #self.pyr_up_proj_first=nn.Conv2d(dim_out, 2, (5,3), padding="same", padding_mode="zeros", bias=False)
+        
+        for i in range(self.num_octs-1,-1,-1):
+
+            if i==0:
+                dim_in=self.Ns[i]*2
+                dim_out=self.Ns[i]
+            else:
+                dim_in=self.Ns[i]*2
+                dim_out=self.Ns[i-1]
+
+            if self.attention_layers[i]:
+                print("Attention layer at (up) oct layer {}".format(i))
+                attn_dict=self.attention_dict
+                #attn_dict.N=self.attention_Ns[i]
+                #assert attn_dict.N > 0
+            else:
+                attn_dict=None
+
+            self.ups.append(nn.ModuleList(
+                                        [
+                                        ResnetBlock(dim_out, 2, use_norm=self.use_norm,num_dils= 1,bias=False, kernel_size=(1,1), proj_place="after", emb_dim=self.emb_dim, init=init, init_zero=init_zero),
+                                        ResnetBlock(dim_in, dim_out, use_norm=self.use_norm,num_dils= self.num_dils[i],attention_dict=attn_dict, bias=False, emb_dim=self.emb_dim, init=init, init_zero=init_zero, Fdim=(i+1)*self.bins_per_oct),
+                                        ]))
+
+
+
+        #self.cropconcat = CropConcatBlock()
+
+
+
+
+    def forward(self, inputs, sigma):
+        """
+        Args: 
+            inputs (Tensor):  Input signal in time-domsin, shape (B,T)
+            sigma (Tensor): noise levels,  shape (B,1)
+        Returns:
+            pred (Tensor): predicted signal in time-domain, shape (B,T)
+        """
+        #apply RFF embedding+MLP of the noise level
+        sigma = self.embedding(sigma)
+
+        
+        #apply CQT to the inputs
+        X_list =self.CQTransform.fwd(inputs.unsqueeze(1))
+        X_list_out=X_list
+
+        hs=[]
+        for i,modules in enumerate(self.downs):
+            #print("downsampler", i)
+            if i <=(self.num_octs-1):
+                C=X_list[-1-i]#get the corresponding CQT octave
+                C=C.squeeze(1)
+                C=torch.view_as_real(C)
+                C=C.permute(0,3,1,2).contiguous() # call contiguous() here?
+                if self.use_fencoding:
+                    #Cfreq=self.freq_encoding(C)
+                    C2=self.freq_encodings[i](C) #B, C + Nfreq*2, F,T
+                else:
+                    C2=C
+                    
+                init_block, pyr_down_proj, ResBlock=modules
+                C2=init_block(C2,sigma)
+            else:
+                pyr_down_proj, ResBlock=modules
+            
+            if i==0:
+                X=C2 #starting the main signal path
+                pyr=self.downsamplerT(C) #starting the auxiliary path
+            elif i<(self.num_octs-1):
+                pyr=torch.cat((self.downsamplerT(C),self.downsamplerT(pyr)),dim=2) #updating the auxiliary path
+                X=torch.cat((C2,X),dim=2) #updating the main signal path with the new octave
+            elif i==(self.num_octs-1):# last layer
+                #pyr=torch.cat((self.downsamplerF(C),self.downsamplerF(pyr)),dim=2) #updating the auxiliary path
+                pyr=torch.cat((C,pyr), dim=2) #no downsampling in the last layer
+                X=torch.cat((C2,X),dim=2) #updating the main signal path with the new octave
+            else: #last layer
+                pass
+                #pyr=pyr
+                #X=X
+
+            X=ResBlock(X, sigma)
+            hs.append(X)
+
+            #downsample the main signal path
+            #we do not need to downsample in the inner layer
+            if i<(self.num_octs-1): 
+                X=self.downsamplerT(X)
+                #apply the residual connection
+                #X=(X+pyr_down_proj(pyr))/(2**0.5) #I'll my need to put that inside a combiner block??
+            else: #last layer
+                #no downsampling in the last layer
+                pass
+
+            #apply the residual connection
+            X=(X+pyr_down_proj(pyr))/(2**0.5) #I'll my need to put that inside a combiner block??
+            #print("encoder ", i, X.shape, X.mean().item(), X.std().item())
+                
+
+        #middle layers
+        #print("bttleneck")
+        if self.args.network.bottleneck_type=="res_dil_convs":
+            for i in range(self.args.network.num_bottleneck_layers):
+                OutBlock, ResBlock =self.middle[i]
+                X=ResBlock(X, sigma)   
+                Xout=OutBlock(X,sigma)
+
+
+        for i,modules in enumerate(self.ups):
+            j=len(self.ups) -i-1
+            #print("upsampler", j)
+
+            OutBlock,  ResBlock=modules
+
+            skip=hs.pop()
+            X=torch.cat((X,skip),dim=1)
+            X=ResBlock(X, sigma)
+            
+            Xout=(Xout+OutBlock(X,sigma))/(2**0.5)
+
+
+            if j<=(self.num_octs-1):
+                X= X[:,:,self.bins_per_oct::,:]
+                Out, Xout= Xout[:,:,0:self.bins_per_oct,:], Xout[:,:,self.bins_per_oct::,:]
+                #pyr_out, pyr= pyr[:,:,0:self.bins_per_oct,:], pyr[:,:,self.bins_per_oct::,:]
+                #X_out=(pyr_up_proj(X_out)+pyr_out)/(2**0.5)
+
+                Out=Out.permute(0,2,3,1).contiguous() #call contiguous() here?
+                Out=torch.view_as_complex(Out)
+
+                #save output
+                X_list_out[i]=Out.unsqueeze(1)
+
+            elif j>(self.num_octs-1):
+                print("We should not be here")
+                pass
+
+            if j>0 and j<=(self.num_octs-1):
+                #pyr=self.upsampler(pyr) #call contiguous() here?
+                X=self.upsamplerT(X) #call contiguous() here?
+                Xout=self.upsamplerT(Xout) #call contiguous() here?
+
+        pred_time=self.CQTransform.bwd(X_list_out)
+        pred_time=pred_time.squeeze(1)
+        pred_time=pred_time[:,0:inputs.shape[-1]]
+        assert pred_time.shape==inputs.shape, "bad shapes"
+        return pred_time
+
+            
+
+
+class CropAddBlock(nn.Module):
+
+    def forward(self,down_layer, x,  **kwargs):
+        x1_shape = down_layer.shape
+        x2_shape = x.shape
+
+        #print(x1_shape,x2_shape)
+        height_diff = (x1_shape[2] - x2_shape[2]) // 2
+        width_diff = (x1_shape[3] - x2_shape[3]) // 2
+
+
+        down_layer_cropped = down_layer[:,
+                                        :,
+                                        height_diff: (x2_shape[2] + height_diff),
+                                        width_diff: (x2_shape[3] + width_diff),:]
+        x = torch.add(down_layer_cropped, x)
+        return x
+
+class CropConcatBlock(nn.Module):
+
+    def forward(self, down_layer, x, **kwargs):
+        x1_shape = down_layer.shape
+        x2_shape = x.shape
+
+        height_diff = (x1_shape[2] - x2_shape[2]) // 2
+        width_diff = (x1_shape[3] - x2_shape[3]) // 2
+        down_layer_cropped = down_layer[:,
+                                        :,
+                                        height_diff: (x2_shape[2] + height_diff),
+                                        width_diff: (x2_shape[3] + width_diff)]
+        x = torch.cat((down_layer_cropped, x),1)
+        return x
+

requirements.txt

@@ -0,0 +1,9 @@
+git+https://github.com/TEAMuP-dev/pyharp.git@develop
+cqt_nsgt_pytorch
+einops
+torch
+torchaudio
+omegaconf
+soundfile
+numpy
+scipy

testing/edm_sampler_inpainting.py

@@ -0,0 +1,367 @@
+from tqdm import tqdm
+import torch
+import torchaudio
+#import scipy.signal
+#import numpy as np
+#from utils.decSTN_pytorch import apply_STN_mask
+
+class Sampler():
+
+    def __init__(self, model, diff_params, args, rid=False):
+
+        self.model = model
+        self.diff_params = diff_params #same as training, useful if we need to apply a wrapper or something
+        self.args=args
+        if not(self.args.tester.diff_params.same_as_training):
+            self.update_diff_params()
+
+
+        self.order=self.args.tester.order
+        self.xi=self.args.tester.posterior_sampling.xi
+         #hyperparameter for the reconstruction guidance
+        self.data_consistency=self.args.tester.data_consistency.use and self.args.tester.data_consistency.type=="always"
+
+        self.data_consistency_end=self.args.tester.data_consistency.use and self.args.tester.data_consistency.type=="end"
+        if self.data_consistency or self.data_consistency_end:
+            if self.args.tester.data_consistency.smooth:
+                self.smooth=True
+            else:
+                self.smooth=False
+
+        #use reconstruction gudance without replacement
+        self.nb_steps=self.args.tester.T
+
+        #self.treshold_on_grads=args.tester.inference.max_thresh_grads
+        self.rid=rid #this is for logging, ignore for now
+
+        #try:
+        #    self.stereo=self.args.tester.stereo
+        #except:
+        #    self.stereo=False
+
+
+    def update_diff_params(self):
+        #the parameters for testing might not be necesarily the same as the ones used for training
+        self.diff_params.sigma_min=self.args.tester.diff_params.sigma_min
+        self.diff_params.sigma_max =self.args.tester.diff_params.sigma_max
+        self.diff_params.ro=self.args.tester.diff_params.ro
+        self.diff_params.sigma_data=self.args.tester.diff_params.sigma_data
+        #par.diff_params.meters stochastic sampling
+        self.diff_params.Schurn=self.args.tester.diff_params.Schurn
+        self.diff_params.Stmin=self.args.tester.diff_params.Stmin
+        self.diff_params.Stmax=self.args.tester.diff_params.Stmax
+        self.diff_params.Snoise=self.args.tester.diff_params.Snoise
+
+
+
+    def get_score_rec_guidance(self, x, y, t_i, degradation):
+
+        x.requires_grad_()
+        x_hat=self.diff_params.denoiser(x, self.model, t_i.unsqueeze(-1))
+
+        if self.args.tester.filter_out_cqt_DC_Nyq:
+            x_hat=self.model.CQTransform.apply_hpf_DC(x_hat)
+
+        den_rec= degradation(x_hat) 
+
+        if len(y.shape)==3:
+            dim=(1,2)
+        elif len(y.shape)==2:
+            dim=1
+
+        if self.args.tester.posterior_sampling.norm=="smoothl1":
+            norm=torch.nn.functional.smooth_l1_loss(y, den_rec, reduction='sum', beta=self.args.tester.posterior_sampling.smoothl1_beta)
+        else:
+            norm=torch.linalg.norm(y-den_rec,dim=dim, ord=self.args.tester.posterior_sampling.norm)
+
+        
+        rec_grads=torch.autograd.grad(outputs=norm,
+                                      inputs=x)
+
+        rec_grads=rec_grads[0]
+        
+        normguide=torch.linalg.norm(rec_grads)/self.args.exp.audio_len**0.5
+        
+        #normalize scaling
+        #s=self.xi/(normguide*t_i+1e-6)
+        s=t_i*self.xi/(normguide+1e-6)
+        
+        #optionally apply a treshold to the gradients
+        if False:
+            #pply tresholding to the gradients. It is a dirty trick but helps avoiding bad artifacts 
+            rec_grads=torch.clip(rec_grads, min=-self.treshold_on_grads, max=self.treshold_on_grads)
+        
+        if self.rid:
+            x_hat_old=x_hat.detach().clone()
+
+        x_hat=x_hat-s*rec_grads #apply gradients
+
+        if self.rid:
+            x_hat_old_2=x_hat.detach().clone()
+
+        if self.data_consistency:
+            x_hat=self.proj_convex_set(x_hat.detach())
+
+        score=(x_hat.detach()-x)/t_i**2
+
+        #apply scaled guidance to the score
+        #score=score-s*rec_grads
+        if self.rid:
+            #score, denoised esimate,  s*rec_grads, denoised_estimate minus gradients, x_hat after pocs
+            return score, x_hat_old, s*rec_grads, x_hat_old_2, x_hat
+        else:
+            return score
+
+    def get_score(self,x, y, t_i, degradation):
+        if y==None:
+            assert degradation==None
+            #unconditional sampling
+            with torch.no_grad():
+                #print("In sampling", x.shape, t_i.shape)
+                x_hat=self.diff_params.denoiser(x, self.model, t_i.unsqueeze(-1))
+                if self.args.tester.filter_out_cqt_DC_Nyq:
+                    x_hat=self.model.CQTransform.apply_hpf_DC(x_hat)
+                score=(x_hat-x)/t_i**2
+            return score
+        else:
+            if self.xi>0:
+                #apply rec. guidance
+                score=self.get_score_rec_guidance(x, y, t_i, degradation)
+    
+                #optionally apply replacement or consistency step
+                #if self.data_consistency:
+                #    #convert score to denoised estimate using Tweedie's formula
+                #    x_hat=score*t_i**2+x
+    
+                #    x_hat=self.proj_convex_set(x_hat)
+    
+                #    #convert back to score
+                #    score=(x_hat-x)/t_i**2
+    
+            else:
+                #denoised with replacement method
+                with torch.no_grad():
+                    x_hat=self.diff_params.denoiser(x, self.model, t_i.unsqueeze(-1))
+                    x_hat=self.proj_convex_set(x_hat.detach())
+
+                    score=(x_hat.detach()-x)/t_i**2
+                        
+                    #x_hat=self.data_consistency_step(x_hat,y, degradation)
+        
+                    #score=(x_hat-x)/t_i**2
+    
+            return score
+
+    def predict_unconditional(
+        self,
+        shape,  #observations (lowpssed signal) Tensor with shape ??
+        device
+    ):
+        self.y=None
+        self.degradation=None
+        return self.predict(shape, device)
+
+    def predict_resample(
+        self,
+        y,  #observations (lowpssed signal) Tensor with shape ??
+        shape,
+        degradation, #lambda function
+    ):
+        self.degradation=degradation 
+        self.y=y
+        #print(shape)
+        return self.predict(shape, y.device)
+
+
+
+
+    def predict(
+        self,
+        shape,  #observations (lowpssed signal) Tensor with shape ??
+        device, #lambda function
+    ):
+
+        if self.rid:
+            rid_xt=torch.zeros((self.nb_steps,shape[0], shape[1]))
+            rid_grads=torch.zeros((self.nb_steps,shape[0], shape[1]))
+            rid_denoised=torch.zeros((self.nb_steps,shape[0], shape[1]))
+            rid_grad_update=torch.zeros((self.nb_steps,shape[0], shape[1]))
+            rid_pocs=torch.zeros((self.nb_steps,shape[0], shape[1]))
+            rid_xt2=torch.zeros((self.nb_steps,shape[0], shape[1]))
+
+        #get the noise schedule
+        t = self.diff_params.create_schedule(self.nb_steps).to(device)
+        #sample from gaussian distribution with sigma_max variance
+        x = self.diff_params.sample_prior(shape,t[0]).to(device)
+
+        #parameter for langevin stochasticity, if Schurn is 0, gamma will be 0 to, so the sampler will be deterministic
+        gamma=self.diff_params.get_gamma(t).to(device)
+
+
+        for i in tqdm(range(0, self.nb_steps, 1)):
+            #print("sampling step ",i," from ",self.nb_steps)
+
+            if gamma[i]==0:
+                #deterministic sampling, do nothing
+                t_hat=t[i] 
+            else:
+                #stochastic sampling
+                #move timestep
+                t_hat=t[i]+gamma[i]*t[i] 
+                #sample noise, Snoise is 1 by default
+                epsilon=torch.randn(shape).to(device)*self.diff_params.Snoise
+                #add extra noise
+                x=x+((t_hat**2 - t[i]**2)**(1/2))*epsilon  #x_hat
+                del epsilon
+
+            if self.rid:
+                rid_xt[i]=x
+
+            score=self.get_score(x, self.y, t_hat, self.degradation)    
+            if self.rid:
+                score, x_hat1, grads, x_hat2, x_hat3=score
+                rid_denoised[i]=x_hat1
+                rid_grads[i]=grads
+                rid_grad_update[i]=x_hat2
+                rid_pocs[i]=x_hat3
+
+
+            #d=-t_hat*((denoised-x_hat)/t_hat**2)
+            d=-t_hat*score
+            
+            #apply second order correction
+            h=t[i+1]-t_hat
+
+
+            if t[i+1]!=0 and self.order==2:  #always except last step
+                #second order correction2
+                #h=t[i+1]-t_hat
+                t_prime=t[i+1]
+                x_prime=x+h*d
+                score=self.get_score(x_prime, self.y, t_prime, self.degradation)
+                if self.rid:
+                    score, x_hat1, grads, x_hat2, x_hat3=score
+
+                d_prime=-t_prime*score
+
+                x=(x+h*((1/2)*d +(1/2)*d_prime))
+
+            elif t[i+1]==0 or self.order==1: #first condition  is to avoid dividing by 0
+                #first order Euler step
+                x=x+h*d
+
+            if self.rid: 
+                rid_xt2[i]=x
+            
+        if self.data_consistency_end:
+            x=self.proj_convex_set(x)
+
+        if self.rid:
+            return x.detach(), rid_denoised.detach(), rid_grads.detach(), rid_grad_update.detach(), rid_pocs.detach(), rid_xt.detach(), rid_xt2.detach(), t.detach()
+        else:
+            return x.detach()
+
+    def apply_mask(self, x, mask=None):
+        
+        if mask is None:
+            mask=self.mask
+
+        return mask*x
+
+    def apply_spectral_mask(self, x):
+
+        if self.args.tester.spectrogram_inpainting.stft.window=="hann":
+            window=torch.hann_window(self.args.tester.spectrogram_inpainting.stft.win_length).to(x.device)
+        else:
+            raise NotImplementedError("Only hann window is implemented for now")
+
+        n_fft=self.args.tester.spectrogram_inpainting.stft.n_fft
+        #add padding to the signal
+        input_shape=x.shape
+        x=torch.nn.functional.pad(x, (0, n_fft-x.shape[-1]%n_fft), mode='constant', value=0)
+        x=torch.stft(x, n_fft, self.args.tester.spectrogram_inpainting.stft.hop_length, self.args.tester.spectrogram_inpainting.stft.win_length, window, return_complex=True) 
+        #X.shape=(B, F, T)
+        x=x*self.mask.unsqueeze(0)
+        #apply the inverse stft
+        x=torch.istft(x, n_fft, self.args.tester.spectrogram_inpainting.stft.hop_length, self.args.tester.spectrogram_inpainting.stft.win_length, window, return_complex=False)
+        #why is the size different? Because of the padding?
+        x=x[...,0:input_shape[-1]]
+
+        return x
+
+
+    #def proj_convex_set(self, x_hat, y, degradation):
+    #    """
+    #    Simple replacement method, used for inpainting and FIR bwe
+    #    """
+    #    #get reconstruction estimate
+    #    den_rec= degradation(x_hat)     
+    #    #apply replacment (valid for linear degradations)
+    #    return y+x_hat-den_rec 
+
+    def prepare_smooth_mask(self, mask, size=10):
+        hann=torch.hann_window(size*2)
+        hann_left=hann[0:size]
+        hann_right=hann[size::]
+        B,N=mask.shape
+        mask=mask[0]
+        prev=1
+        new_mask=mask.clone()
+        #print(hann.shape)
+        for i in range(len(mask)):
+            if mask[i] != prev:
+                #print(i, mask.shape, mask[i], prev)
+                #transition
+                if mask[i]==0:
+                   #gap encountered, apply hann right before
+                   new_mask[i-size:i]=hann_right
+                if mask[i]==1:
+                   #gap encountered, apply hann left after
+                   new_mask[i:i+size]=hann_left
+                #print(mask[i-2*size:i+2*size])
+                #print(new_mask[i-2*size:i+2*size])
+                
+            prev=mask[i]
+        return new_mask.unsqueeze(0).expand(B,-1)
+
+    def predict_inpainting(
+        self,
+        y_masked,
+        mask
+        ):
+        self.mask=mask.to(y_masked.device)
+
+        self.y=y_masked
+
+        self.degradation=lambda x: self.apply_mask(x)
+        if self.data_consistency or self.data_consistency_end:
+            if self.smooth:
+                smooth_mask=self.prepare_smooth_mask(mask, self.args.tester.data_consistency.hann_size)
+            else:
+                smooth_mask=mask
+
+            self.proj_convex_set= lambda x: smooth_mask*y_masked+(1-smooth_mask)*x #will this work? I am too scared about lambdas
+
+
+        return self.predict(self.y.shape, self.y.device)
+
+    def predict_spectrogram_inpainting(
+        self,
+        y_masked,
+        mask
+        ):
+        self.mask=mask.to(y_masked.device)
+        del mask
+
+        self.y=y_masked
+        del y_masked
+
+        self.degradation=lambda x: self.apply_spectral_mask(x)
+        if self.data_consistency or self.data_consistency_end:
+
+            self.proj_convex_set= lambda x: self.y+x-self.apply_spectral_mask(x) #If this fails, consider using a more appropiate projection
+
+        return self.predict(self.y.shape, self.y.device)
+
+
+

utils/training_utils.py

@@ -0,0 +1,384 @@
+
+import torch
+import torchaudio
+import numpy as np
+import scipy.signal
+class EMAWarmup:
+    """Implements an EMA warmup using an inverse decay schedule.
+    If inv_gamma=1 and power=1, implements a simple average. inv_gamma=1, power=2/3 are
+    good values for models you plan to train for a million or more steps (reaches decay
+    factor 0.999 at 31.6K steps, 0.9999 at 1M steps), inv_gamma=1, power=3/4 for models
+    you plan to train for less (reaches decay factor 0.999 at 10K steps, 0.9999 at
+    215.4k steps).
+    Args:
+        inv_gamma (float): Inverse multiplicative factor of EMA warmup. Default: 1.
+        power (float): Exponential factor of EMA warmup. Default: 1.
+        min_value (float): The minimum EMA decay rate. Default: 0.
+        max_value (float): The maximum EMA decay rate. Default: 1.
+        start_at (int): The epoch to start averaging at. Default: 0.
+        last_epoch (int): The index of last epoch. Default: 0.
+    """
+
+    def __init__(self, inv_gamma=1., power=1., min_value=0., max_value=1., start_at=0,
+                 last_epoch=0):
+        self.inv_gamma = inv_gamma
+        self.power = power
+        self.min_value = min_value
+        self.max_value = max_value
+        self.start_at = start_at
+        self.last_epoch = last_epoch
+
+    def state_dict(self):
+        """Returns the state of the class as a :class:`dict`."""
+        return dict(self.__dict__.items())
+
+    def load_state_dict(self, state_dict):
+        """Loads the class's state.
+        Args:
+            state_dict (dict): scaler state. Should be an object returned
+                from a call to :meth:`state_dict`.
+        """
+        self.__dict__.update(state_dict)
+
+    def get_value(self):
+        """Gets the current EMA decay rate."""
+        epoch = max(0, self.last_epoch - self.start_at)
+        value = 1 - (1 + epoch / self.inv_gamma) ** -self.power
+        return 0. if epoch < 0 else min(self.max_value, max(self.min_value, value))
+
+    def step(self):
+        """Updates the step count."""
+        self.last_epoch += 1
+
+
+#from https://github.com/csteinmetz1/auraloss/blob/main/auraloss/perceptual.py
+class FIRFilter(torch.nn.Module):
+    """FIR pre-emphasis filtering module.
+    Args:
+        filter_type (str): Shape of the desired FIR filter ("hp", "fd", "aw"). Default: "hp"
+        coef (float): Coefficient value for the filter tap (only applicable for "hp" and "fd"). Default: 0.85
+        ntaps (int): Number of FIR filter taps for constructing A-weighting filters. Default: 101
+        plot (bool): Plot the magnitude respond of the filter. Default: False
+    Based upon the perceptual loss pre-empahsis filters proposed by
+    [Wright & Välimäki, 2019](https://arxiv.org/abs/1911.08922).
+    A-weighting filter - "aw"
+    First-order highpass - "hp"
+    Folded differentiator - "fd"
+    Note that the default coefficeint value of 0.85 is optimized for
+    a sampling rate of 44.1 kHz, considering adjusting this value at differnt sampling rates.
+    """
+
+    def __init__(self, filter_type="hp", coef=0.85, fs=44100, ntaps=101, plot=False): 
+        """Initilize FIR pre-emphasis filtering module."""
+        super(FIRFilter, self).__init__()
+        self.filter_type = filter_type
+        self.coef = coef
+        self.fs = fs
+        self.ntaps = ntaps
+        self.plot = plot
+
+        if ntaps % 2 == 0:
+            raise ValueError(f"ntaps must be odd (ntaps={ntaps}).")
+
+        if filter_type == "hp":
+            self.fir = torch.nn.Conv1d(1, 1, kernel_size=3, bias=False, padding=1)
+            self.fir.weight.requires_grad = False
+            self.fir.weight.data = torch.tensor([1, -coef, 0]).view(1, 1, -1)
+        elif filter_type == "fd":
+            self.fir = torch.nn.Conv1d(1, 1, kernel_size=3, bias=False, padding=1)
+            self.fir.weight.requires_grad = False
+            self.fir.weight.data = torch.tensor([1, 0, -coef]).view(1, 1, -1)
+        elif filter_type == "aw":
+            # Definition of analog A-weighting filter according to IEC/CD 1672.
+            f1 = 20.598997
+            f2 = 107.65265
+            f3 = 737.86223
+            f4 = 12194.217
+            A1000 = 1.9997
+
+            NUMs = [(2 * np.pi * f4) ** 2 * (10 ** (A1000 / 20)), 0, 0, 0, 0]
+            DENs = np.polymul(
+                [1, 4 * np.pi * f4, (2 * np.pi * f4) ** 2],
+                [1, 4 * np.pi * f1, (2 * np.pi * f1) ** 2],
+            )
+            DENs = np.polymul(
+                np.polymul(DENs, [1, 2 * np.pi * f3]), [1, 2 * np.pi * f2]
+            )
+
+            # convert analog filter to digital filter
+            b, a = scipy.signal.bilinear(NUMs, DENs, fs=fs)
+
+            # compute the digital filter frequency response
+            w_iir, h_iir = scipy.signal.freqz(b, a, worN=512, fs=fs)
+
+            # then we fit to 101 tap FIR filter with least squares
+            taps = scipy.signal.firls(ntaps, w_iir, abs(h_iir), fs=fs)
+
+            # now implement this digital FIR filter as a Conv1d layer
+            self.fir = torch.nn.Conv1d(
+                1, 1, kernel_size=ntaps, bias=False, padding=ntaps // 2
+            )
+            self.fir.weight.requires_grad = False
+            self.fir.weight.data = torch.tensor(taps.astype("float32")).view(1, 1, -1)
+
+    def forward(self, error):
+        """Calculate forward propagation.
+        Args:
+            input (Tensor): Predicted signal (B, #channels, #samples).
+            target (Tensor): Groundtruth signal (B, #channels, #samples).
+        Returns:
+            Tensor: Filtered signal.
+        """
+        self.fir.weight.data=self.fir.weight.data.to(error.device)
+        error=error.unsqueeze(1)
+        error = torch.nn.functional.conv1d(
+            error, self.fir.weight.data, padding=self.ntaps // 2
+        )
+        error=error.squeeze(1)
+        return error
+
+def resample_batch(audio, fs, fs_target, length_target):
+
+        device=audio.device
+        dtype=audio.dtype
+        B=audio.shape[0]
+        #if possible resampe in a batched way
+        #check if all the fs are the same and equal to 44100
+        if fs_target==22050:
+            if (fs==44100).all():
+                 audio=torchaudio.functional.resample(audio, 2,1)
+                 return audio[:, 0:length_target] #trow away the last samples
+            elif (fs==48000).all():
+                 #approcimate resamppleint
+                 audio=torchaudio.functional.resample(audio, 160*2,147)
+                 return audio[:, 0:length_target]
+            else:
+                #if revious is unsuccesful bccause we have examples at 441000 and 48000 in the same batch,, just iterate over the batch
+                proc_batch=torch.zeros((B,length_target), device=device)
+                for i, (a, f_s) in enumerate(zip(audio, fs)): #I hope this shit wll not slow down everythingh
+                    if f_s==44100:
+                        #resample by 2
+                        a=torchaudio.functional.resample(a, 2,1)
+                    elif f_s==48000:
+                        a=torchaudio.functional.resample(a, 160*2,147)
+                    else:
+                        print("WARNING, strange fs", f_s)
+           
+                    proc_batch[i]=a[0:length_target]
+                    return proc_batch
+        elif fs_target==44100:
+            if (fs==44100).all():
+                 return audio[:, 0:length_target] #trow away the last samples
+            elif (fs==48000).all():
+                 #approcimate resamppleint
+                 audio=torchaudio.functional.resample(audio, 160,147)
+                 return audio[:, 0:length_target]
+            else:
+                #if revious is unsuccesful bccause we have examples at 441000 and 48000 in the same batch,, just iterate over the batch
+                proc_batch=torch.zeros((B,length_target), device=device)
+                for i, (a, f_s) in enumerate(zip(audio, fs)): #I hope this shit wll not slow down everythingh
+                    if f_s==44100:
+                        #resample by 2
+                        pass
+                    elif f_s==48000:
+                        a=torchaudio.functional.resample(a, 160,147)
+                    else:
+                        print("WARNING, strange fs", f_s)
+           
+                    proc_batch[i]=a[0:length_target] 
+                    return proc_batch
+        else:
+            print(" resampling to fs_target", fs_target)
+            if (fs==44100).all():
+                 audio=torchaudio.functional.resample(audio, 44100, fs_target)
+                 return audio[:, 0:length_target] #trow away the last samples
+            elif (fs==48000).all():
+                 #approcimate resamppleint
+                 audio=torchaudio.functional.resample(audio, 48000,fs_target)
+                 return audio[:, 0:length_target]
+            else:
+                #if revious is unsuccesful bccause we have examples at 441000 and 48000 in the same batch,, just iterate over the batch
+                proc_batch=torch.zeros((B,length_target), device=device)
+                for i, (a, f_s) in enumerate(zip(audio, fs)): #I hope this shit wll not slow down everythingh
+                    if f_s==44100:
+                        #resample by 2
+                        a=torchaudio.functional.resample(a, 44100,fs_target)
+                    elif f_s==48000:
+                        a=torchaudio.functional.resample(a, 48000,fs_target)
+                    else:
+                        print("WARNING, strange fs", f_s)
+           
+                    proc_batch[i]=a[0:length_target] 
+                    return proc_batch
+
+def load_state_dict( state_dict, network=None, ema=None, optimizer=None, log=True):
+        '''
+        utility for loading state dicts for different models. This function sequentially tries different strategies
+        args:
+            state_dict: the state dict to load
+        returns:
+            True if the state dict was loaded, False otherwise
+        Assuming the operations are don in_place, this function will not create a copy of the network and optimizer (I hope)
+        '''
+        #print(state_dict)
+        if log: print("Loading state dict")
+        if log:
+            print(state_dict.keys())
+        #if there
+        try:
+            if log: print("Attempt 1: trying with strict=True")
+            if network is not None:
+                network.load_state_dict(state_dict['network'])
+            if optimizer is not None:
+                optimizer.load_state_dict(state_dict['optimizer'])
+            if ema is not None:
+                ema.load_state_dict(state_dict['ema'])
+            return True
+        except Exception as e:
+            if log:
+                print("Could not load state dict")
+                print(e)
+        try:
+            if log: print("Attempt 2: trying with strict=False")
+            if network is not None:
+                network.load_state_dict(state_dict['network'], strict=False)
+            #we cannot load the optimizer in this setting
+            #self.optimizer.load_state_dict(state_dict['optimizer'], strict=False)
+            if ema is not None:
+                ema.load_state_dict(state_dict['ema'], strict=False)
+            return True
+        except Exception as e:
+            if log:
+                print("Could not load state dict")
+                print(e)
+                print("training from scratch")
+        try:
+            if log: print("Attempt 3: trying with strict=False,but making sure that the shapes are fine")
+            if ema is not None:
+                ema_state_dict = ema.state_dict()
+            if network is not None:
+                network_state_dict = network.state_dict()
+            i=0 
+            if network is not None:
+                for name, param in state_dict['network'].items():
+                    if log: print("checking",name) 
+                    if name in network_state_dict.keys():
+                        if network_state_dict[name].shape==param.shape:
+                                network_state_dict[name]=param
+                                if log:
+                                    print("assigning",name)
+                                i+=1
+            network.load_state_dict(network_state_dict)
+            if ema is not None:
+                for name, param in state_dict['ema'].items():
+                        if log: print("checking",name) 
+                        if name in ema_state_dict.keys():
+                            if ema_state_dict[name].shape==param.shape:
+                                ema_state_dict[name]=param
+                                if log:
+                                    print("assigning",name)
+                                i+=1
+     
+            ema.load_state_dict(ema_state_dict)
+     
+            if i==0:
+                if log: print("WARNING, no parameters were loaded")
+                raise Exception("No parameters were loaded")
+            elif i>0:
+                if log: print("loaded", i, "parameters")
+                return True
+
+        except Exception as e:
+            print(e)
+            print("the second strict=False failed")
+
+
+        try:
+            if log: print("Attempt 4: Assuming the naming is different, with the network and ema called 'state_dict'")
+            if network is not None:
+                network.load_state_dict(state_dict['state_dict'])
+            if ema is not None:
+                ema.load_state_dict(state_dict['state_dict'])
+        except Exception as e:
+            if log:
+                print("Could not load state dict")
+                print(e)
+                print("training from scratch")
+                print("It failed 3 times!! but not giving up")
+            #print the names of the parameters in self.network
+
+        try:
+            if log: print("Attempt 5: trying to load with different names, now model='model' and ema='ema_weights'")
+            if ema is not None:
+                dic_ema = {}
+                for (key, tensor) in zip(state_dict['model'].keys(), state_dict['ema_weights']):
+                    dic_ema[key] = tensor
+                    ema.load_state_dict(dic_ema)
+                return True
+        except Exception as e:
+            if log:
+                print(e)
+
+        try:
+            if log: print("Attempt 6: If there is something wrong with the name of the ema parameters, we can try to load them using the names of the parameters in the model")
+            if ema is not None:
+                dic_ema = {}
+                i=0
+                for (key, tensor) in zip(state_dict['model'].keys(), state_dict['model'].values()):
+                    if tensor.requires_grad:
+                        dic_ema[key]=state_dict['ema_weights'][i]
+                        i=i+1
+                    else:
+                        dic_ema[key]=tensor     
+                ema.load_state_dict(dic_ema)
+                return True
+        except Exception as e:
+            if log:
+                print(e)
+
+        #try:
+        #assign the parameters in state_dict to self.network using a for loop
+        print("Attempt 7: Trying to load the parameters one by one. This is for the dance diffusion model, looking for parameters starting with 'diffusion.' or 'diffusion_ema.'")
+        if ema is not None:
+            ema_state_dict = ema.state_dict()
+        if network is not None:
+            network_state_dict = ema.state_dict()
+        i=0 
+        if network is not None:
+            for name, param in state_dict['state_dict'].items():
+                print("checking",name) 
+                if name.startswith("diffusion."):
+                    i+=1
+                    name=name.replace("diffusion.","")
+                    if network_state_dict[name].shape==param.shape:
+                        #print(param.shape, network.state_dict()[name].shape)
+                        network_state_dict[name]=param
+                        #print("assigning",name)
+
+            network.load_state_dict(network_state_dict, strict=False)
+
+        if ema is not None:
+            for name, param in state_dict['state_dict'].items():
+                if name.startswith("diffusion_ema."): 
+                    i+=1
+                    name=name.replace("diffusion_ema.","")
+                    if ema_state_dict[name].shape==param.shape:
+                        if log:
+                                print(param.shape, ema.state_dict()[name].shape)
+                        ema_state_dict[name]=param
+
+            ema.load_state_dict(ema_state_dict, strict=False)
+
+        if i==0:
+            print("WARNING, no parameters were loaded")
+            raise Exception("No parameters were loaded")
+        elif i>0:
+            print("loaded", i, "parameters")
+            return True
+        #except Exception as e:
+        #    if log:
+        #        print(e)
+
+        return False
+
+