utils/training_utils.py

import math
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


def resample_batch(audio, fs, fs_target, length_target=None):

        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
        #print(fs_target)
        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)
                    elif f_s==22050:
                        pass
                    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
                #B,C,L=audio.shape
                #proc_batch=torch.zeros((B,C,L), device=device)
                proc_batch=torch.zeros((B,length_target), device=device)
                #print("debigging resample batch")
                #print(audio.shape,fs.shape)
                #for i, (a, f_s) in enumerate(zip(audio, fs.tolist())): #I hope this shit wll not slow down everythingh
                for i, (a, f_s) in enumerate(zip(audio, fs)): #I hope this shit wll not slow down everythingh
                    #print(i,a.shape,f_s)
                    if f_s==44100:
                        #resample by 2
                        pass
                    elif f_s==22050:
                        a=torchaudio.functional.resample(a, 1,2)
                    elif f_s==48000:
                        a=torchaudio.functional.resample(a, 160,147)
                    elif f_s==96000:
                        a=torchaudio.functional.resample(a, 320, 147)
                    else:
                        print("WARNING, strange fs", f_s)
           

                    proc_batch[i]=a[...,0:length_target] 
                return proc_batch
        else:
            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():
                 print("resampling 48000 to 16000", length_target, audio.shape)
                 #approcimate resamppleint
                 audio=torchaudio.functional.resample(audio, 48000,fs_target)
                 print(audio.shape)
                 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:
            print("assuming the network was saved from a DDP model, and I forgot to add .module to the keys")
            #verify that the keys are the same but with .module removed
            if network is not None:
                network_state_dict = network.state_dict()
                for key in list(state_dict['network'].keys()):
                    print("checking", key)
                    if key.startswith('module.'):
                        new_key = key.replace('module.', '')
                        state_dict['network'][new_key] = state_dict['network'].pop(key)

                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'])

            print("loaded state dict with .module removed from the keys")

            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)
        #try:
        # this is for the dmae1d mddel, assuming there is only one network
        if network is not None:
            network.load_state_dict(state_dict, strict=True)
        if ema is not None:
            ema.load_state_dict(state_dict, strict=True)
        return True

        #except Exception as e:
        #    if log:
        #        print(e)

        return False

            
def unnormalize(x,stds, args):
        #unnormalize the STN separated audio
        new_std=args.exp.normalization.target_std
        if new_std=="sigma_data":
            new_std=args.diff_params.sigma_data
        x=stds*x/(new_std+1e-8)
        return x
def normalize( xS, xT, xN, args, return_std=False):
        #normalize the STN separated audio
        if args.exp.normalization.mode=="None":
            pass
        elif args.exp.normalization.mode=="residual_noise":
            #normalize the residual noise

            std=xN.std(dim=-1, keepdim=True).mean(dim=1, keepdim=True)
            new_std=args.exp.normalization.target_std

            if new_std=="sigma_data":
                new_std=args.diff_params.sigma_data

            #print(std, new_std)

            xN=new_std*xN/(std+1e-8)
            #print(xN.std(dim=-1, keepdim=True))
            xS=new_std*xS/(std+1e-8)
            xT=new_std*xT/(std+1e-8)
        elif args.exp.normalization.mode=="residual_noise_batch":
            #normalize the residual noise per batch
            #get the std of the entire batch
            std=xN.std(dim=(0,1,2),unbiased=True, keepdim=False)
        
            new_std=args.exp.normalization.target_std

            if new_std=="sigma_data":
                new_std=args.diff_params.sigma_data

            #print(std, new_std)

            xN=new_std*xN/(std+1e-8)
            #print(xN.std(dim=-1, keepdim=True).mean(dim=1, keepdim=True))
            xS=new_std*xS/(std+1e-8)
            xT=new_std*xT/(std+1e-8)

        elif args.exp.normalization.mode=="all":
            std=(xN+xS+xT).std(dim=-1, keepdim=True).mean(dim=1, keepdim=True)
            new_std=args.exp.normalization.target_std
            if new_std=="sigma_data":
                new_std=args.diff_params.sigma_data
            xN=new_std*xN/(std+1e-8)
            xS=new_std*xS/(std+1e-8)
            xT=new_std*xT/(std+1e-8)
            #print("std",xN.std(dim=-1, keepdim=True).mean(dim=1, keepdim=True))
        else:
            print("normalization mode not recognized")
            pass

        try:
            if return_std:
                return xS, xT, xN, std
        except Exception as e:
            print(e)
            print("warning!, std cannot be returned")
            pass

        return xS, xT, xN


#def find_time_offset(x: torch.Tensor, y: torch.Tensor, sign_shift=False):


#def find_time_offset(y: torch.Tensor, x: torch.Tensor, margin=3000, check_sign_flip=False):
#    x = x.double()
#    y = y.double()
#    N = x.size(-1)
#    M = y.size(-1)
#    #print("x",x.shape)
#    #print("y",y.shape)
#    #print(N,M)

#    X = torch.fft.rfft(x, n=N + M - 1)

#    Y = torch.fft.rfft(y, n=N + M - 1)
#    print("X",X.shape, "Y",Y.shape, "x",x.shape, "y",y.shape)

#    corr = torch.fft.irfft(X.conj() * Y)
#    print("corr",corr.shape)

#    shifts = torch.argmax(corr, dim=-1) - x.shape[-1]

#    return shifts


def align_batch(y, x, sample_rate):

    x = x.double()
    y = y.double()
    N = x.size(-1)
    M = y.size(-1)
    y_device=y.device
    x_device=x.device
    y_dtype=y.dtype
    x_dtype=x.dtype

    X = torch.fft.rfft(x, n=N + M - 1)
    X_flipped=torch.fft.rfft(x*-1, n=N + M - 1)

    Y = torch.fft.rfft(y, n=N + M - 1)

    corr = torch.fft.irfft(X.conj() * Y)
    corr_flipped = torch.fft.irfft(X_flipped.conj() * Y)

    corr=corr.sum(dim=1)
    corr_flipped=corr_flipped.sum(dim=1)

    #print("corr",corr.shape)
    shifts = torch.argmax(corr, dim=-1)

    shifts_flipped = torch.argmax(corr_flipped, dim=-1)
#    corr_values_flipped= corr_flipped[...,shifts_flipped]
    #shifts= torch.where(shifts >= N, shifts - N - M + 1, shifts)
    #shifts_flipped= torch.where(shifts_flipped >= N, shifts_flipped - N - M + 1, shifts_flipped)

    shifts=shifts.to(torch.int64)
            
    result=[]
    for i in range(len(shifts)):
        corr_value=corr[i, shifts[i]]
        corr_value_flipped=corr_flipped[i, shifts_flipped[i]]
        if corr_value < corr_value_flipped:
            shift=shifts_flipped[i].item()
            if shift >=N:
                shift = shift - N - M + 1
            result.append(torch.roll(x[i]*-1, shifts=shift, dims=-1) )
        else:
            shift=shifts[i].item()
            if shift >=N:
                shift = shift - N - M + 1
            result.append(torch.roll(x[i], shifts=shift, dims=-1) )


    x= torch.stack(result)
    #x = torch.stack([torch.roll(x[i], shifts=shifts[i].item(), dims=-1) for i in range(x.shape[0])])

    return y.to(y_device).to(y_dtype), x.to(x_device).to(x_dtype)

def get_pink_noise_magnitude(freqs, device='cpu'):
    a=torch.ones_like(freqs, device=device)
    return a / torch.sqrt(torch.clamp(freqs, min=1e-6))  # Avoid division by zero

def generate_pink_noise(shape, device='cpu'):
    """
    Generate pink noise with 1/f frequency scaling for a batch of stereo signals.
    
    Args:
        shape (tuple): Shape of the noise signal (B, 2, T).
        device (str): Device for tensor computation ('cpu' or 'cuda').
    
    Returns:
        torch.Tensor: Pink noise signal with shape (B, 2, T).
    """
    B, C, T = shape
    
    # Generate white noise
    white_noise = torch.randn(B, C, T, device=device)
    
    # Perform FFT to move to frequency domain
    fft = torch.fft.rfft(white_noise, dim=-1)
    
    # Generate frequency bins
    freqs = torch.fft.rfftfreq(T, d=1.0).to(device)
    
    # Scale the amplitude by 1/sqrt(frequency) to approximate pink noise
    # Avoid division by zero by clamping frequencies to a minimum value
    H= get_pink_noise_magnitude(freqs, device=device)
    fft*=H.unsqueeze(0).unsqueeze(0)
    
    # Perform inverse FFT to return to time domain
    pink_noise = torch.fft.irfft(fft, n=T, dim=-1)
    
    return pink_noise

def add_pink_noise(signal, snr_db):
    """
    Add pink noise to a signal based on the desired SNR.
    
    Args:
        signal (torch.Tensor): Original signal with shape (B, 2, T).
        snr_db (float): Desired Signal-to-Noise Ratio in decibels.
    
    Returns:
        torch.Tensor: Noisy signal with the specified SNR.
    """
    # Calculate signal power
    signal_power = torch.mean(signal ** 2, dim=(-1, -2), keepdim=True)

    
    # Calculate noise power based on desired SNR
    snr_linear = 10 ** (snr_db / 10)
    noise_power = signal_power / snr_linear.view(-1, 1, 1)  # Adjust shape for broadcasting
    
    # Generate pink noise
    pink_noise = generate_pink_noise(signal.shape, device=signal.device)
    
    # Scale pink noise to achieve the desired noise power
    noise_scaling_factor = torch.sqrt(noise_power / torch.mean(pink_noise ** 2, dim=-1, keepdim=True))

    scaled_noise = pink_noise * noise_scaling_factor
    
    # Add scaled noise to the original signal
    noisy_signal = signal + scaled_noise
    
    return noisy_signal


def Gauss_smooth(X,f,Noct=1, smooth_filter=None):
    """
    based on https://github.com/IoSR-Surrey/MatlabToolbox/blob/4bff1bb2da7c95de0ce2713e7c710a0afa70c705/%2Biosr/%2Bdsp/smoothSpectrum.m
    Smooths the magnitude spectrum X using a Gaussian filter.

    Args:
        X (torch.Tensor): Input spectrum to be smoothed, shape (B, N,).
        f (torch.Tensor): Frequency bins corresponding to the spectrum, shape (N,).
        Noct (int, optional): Number of octaves for smoothing. Default is 1
    Returns:
        torch.Tensor: Smoothed spectrum, same shape as X.
    """

    def gauss_f(f_x,F,Noct):
        sigma = (F/Noct)/np.pi
        g = torch.exp(-(((f_x-F)**2)/(2*(sigma**2))))
        g = g/torch.sum(g)
        return g

    shape=X.shape

    x_oct = X.clone().view(-1, shape[-1])  # Initialize smoothed output
    if Noct>0:
        for i in range(1,len(f)):
            g = gauss_f(f,f[i],Noct)
            g=g.to(X.device).view(1, -1)
            x_oct[...,i] = torch.sum(g*X)


        x_oct = x_oct.clamp(min=0)  # Ensure non-negative values

    return x_oct.view(shape)  # Reshape back to original dimensions

def prepare_smooth_filter(f, Noct=1):
     f_matrix = f.unsqueeze(0)  # Shape: [1, N]
     F_matrix = f.unsqueeze(1)  # Shape: [N, 1]
     
     # Calculate sigma for each center frequency
     sigma = (F_matrix / Noct) / np.pi  # Shape: [N, 1]
     
     # Calculate Gaussian weights for all frequency pairs at once
     g = torch.exp(-((f_matrix - F_matrix)**2) / (2 * (sigma**2)))  # Shape: [N, N]
     
     # Normalize each row to sum to 1
     g = g / g.sum(dim=1, keepdim=True)  # Shape: [N, N]
     
     # Move to the same device as X
     return g

def Gauss_smooth_vectorized(X, f, Noct=1, smooth_filter=None):
    """
    based on https://github.com/IoSR-Surrey/MatlabToolbox/blob/4bff1bb2da7c95de0ce2713e7c710a0afa70c705/%2Biosr/%2Bdsp/smoothSpectrum.m
    Smooths the magnitude spectrum X using a Gaussian filter.

    Args:
        X (torch.Tensor): Input spectrum to be smoothed, shape (B, N,).
        f (torch.Tensor): Frequency bins corresponding to the spectrum, shape (N,).
        Noct (int, optional): Number of octaves for smoothing. Default is 1
    Returns:
        torch.Tensor: Smoothed spectrum, same shape as X.
    """
    shape = X.shape
    x_oct = X.clone().view(-1, shape[-1])  # Initialize smoothed output
    
    if Noct > 0:
        # Vectorized implementation
        # Create a matrix of all frequency pairs
        if smooth_filter is None:
            f_matrix = f.unsqueeze(0)  # Shape: [1, N]
            F_matrix = f.unsqueeze(1)  # Shape: [N, 1]
            
            # Calculate sigma for each center frequency
            sigma = (F_matrix / Noct) / np.pi  # Shape: [N, 1]
            
            # Calculate Gaussian weights for all frequency pairs at once
            g = torch.exp(-((f_matrix - F_matrix)**2) / (2 * (sigma**2)))  # Shape: [N, N]
            
            # Normalize each row to sum to 1
            g = g / g.sum(dim=1, keepdim=True)  # Shape: [N, N]
            
            # Move to the same device as X
        else:
            g=smooth_filter

        g = g.to(X.device)
        
        # Apply the filter to each batch element
        # Skip the first bin (i=0) as in the original loop
        x_oct_new = torch.matmul(x_oct, g[1:].T)  # Shape: [B, N-1]
        
        # Keep the first bin unchanged and update the rest
        x_oct[:, 1:] = x_oct_new
        
        # Ensure non-negative values
        x_oct = x_oct.clamp(min=0)
    
    return x_oct.view(shape)  # Reshape back to original dimensions

def create_music_mean_spectrum_curve(freqs_Hz, device='cpu'):
    """
    Create a target curve approximating the mean spectrum of music.
    
    Args:
        freqs_Hz (torch.Tensor or np.ndarray): Frequency bins (Hz).
        device (str): Device for tensor computation ('cpu' or 'cuda').
    
    Returns:
        torch.Tensor: Target curve for equalization.
    """
    import torch
    import numpy as np
    
    # Convert to numpy for easier manipulation if needed
    if isinstance(freqs_Hz, torch.Tensor):
        freqs_np = freqs_Hz.cpu().numpy()
    else:
        freqs_np = freqs_Hz
    
    # Create the curve in segments
    target = np.ones_like(freqs_np)
    
    # Define the segments
    f1 = 100    # First transition point
    f2 = 1000   # Second transition point
    f3 = 5000  # Third transition point
    
    # Calculate the curve for each segment
    for i, f in enumerate(freqs_np):
        if f < f1:
            # Below 100Hz: -3dB/octave roll-off
            target[i] = np.sqrt(f / f1)
        elif f <= f2:
            # 100Hz to 1kHz: Flat
            target[i] = 1.0
        elif f <= f3:
            # 1kHz to 10kHz: -3dB/octave roll-off
            target[i] = np.sqrt(f2 / f)
        else:
            # Above 10kHz: -6dB/octave roll-off
            # Calculate what the value would be at 10kHz using the -3dB/octave formula
            val_at_f3 = np.sqrt(f2 / f3)
            # Continue from that value with a -6dB/octave slope
            target[i] = val_at_f3 * (f3 / f)
    
    # Convert back to tensor
    target_curve = torch.tensor(target, device=device)
    target_curve[0]=target_curve[1]  # Ensure the first value is not zero to avoid division issues

    #decrease 20dB
    target_curve = target_curve * 0.1  # Scale down to -20dB
    
    return target_curve