Upload 2 files

inference.py

@@ -10,11 +10,14 @@ import os
 import argparse
 import torch
 import torchaudio
+import torch.nn as nn
 import librosa
 from models.stfts import mag_phase_stft, mag_phase_istft
 from models.generator_SEMamba_time_d4 import SEMamba
 from utils.util import load_config, pad_or_trim_to_match
 
+RELU = nn.ReLU()
+
 def get_filepaths(directory, file_type=None):
     file_paths = []  # List which will store all of the full filepaths.
     # Walk the tree.
@@ -75,6 +78,10 @@ def inference(args, device):
                 addeps=False
             )
             amp_g, pha_g, _ = SE_model(noisy_mag, noisy_pha)
+            # To remove "strange sweep artifact"
+            mag = torch.expm1(RELU(amp_g)) # [1, F, T]
+            zero_portion = torch.sum(mag==0, 1)/mag.shape[1]
+            amp_g[:,:,(zero_portion>0.5)[0]] = 0
 
             audio_g = mag_phase_istft(amp_g, pha_g, n_fft_scaled, hop_size_scaled, win_size_scaled, compress_factor)
             audio_g = pad_or_trim_to_match(noisy_wav.detach(), audio_g, pad_value=1e-8)  # Align lengths using epsilon padding

inference_chunk.py

@@ -10,12 +10,14 @@ import os
 import argparse
 import torch
 import torchaudio
+import torch.nn as nn
 import librosa
 import math
 from models.stfts import mag_phase_stft, mag_phase_istft
 from models.generator_SEMamba_time_d4 import SEMamba
 from utils.util import load_config, pad_or_trim_to_match
 
+RELU = nn.ReLU()
 
 def get_filepaths(directory, file_type=None):
     file_paths = []  # List which will store all of the full filepaths.
@@ -89,6 +91,10 @@ def inference(args, device):
                         addeps=False
                     )
                     amp_g, pha_g, _ = SE_model(noisy_mag, noisy_pha)
+                    # To remove "strange sweep artifact"
+                    mag = torch.expm1(RELU(amp_g)) # [1, F, T]
+                    zero_portion = torch.sum(mag==0, 1)/mag.shape[1]
+                    amp_g[:,:,(zero_portion>0.5)[0]] = 0
 
                     audio_g = mag_phase_istft(amp_g, pha_g, n_fft_scaled, hop_size_scaled, win_size_scaled, compress_factor)
                     audio_g = pad_or_trim_to_match(noisy_wav_chunk.detach(), audio_g, pad_value=1e-8)  # Align lengths using epsilon padding