Upload 14 files

.gitattributes

@@ -33,3 +33,7 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+enhanced_audio/offline_enhanced_mic_test.flac filter=lfs diff=lfs merge=lfs -text
+enhanced_audio/online_enhanced_mic_test.flac filter=lfs diff=lfs merge=lfs -text
+noisy_audio/mic_test.wav filter=lfs diff=lfs merge=lfs -text
+offline_enhanced_audio/mic_test.flac filter=lfs diff=lfs merge=lfs -text

enhanced_audio/offline_enhanced_mic_test.flac

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

enhanced_audio/online_enhanced_mic_test.flac

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

models/stfts.py

@@ -0,0 +1,95 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import torch
+import torch.nn as nn
+
+def decompress_signed_log1p(y):
+   return torch.sign(y) * (torch.expm1(torch.abs(y)))
+
+RELU = nn.ReLU()
+
+def mag_phase_stft(y, n_fft, hop_size, win_size, compress_factor=1.0, center=True, addeps=False):
+    """
+    Compute magnitude and phase using STFT.
+
+    Args:
+        y (torch.Tensor): Input audio signal.
+        n_fft (int): FFT size.
+        hop_size (int): Hop size.
+        win_size (int): Window size.
+        compress_factor (float, optional): Magnitude compression factor. Defaults to 1.0.
+        center (bool, optional): Whether to center the signal before padding. Defaults to True.
+        eps (bool, optional): Whether adding epsilon to magnitude and phase or not. Defaults to False. 
+
+    Returns:
+        tuple: Magnitude, phase, and complex representation of the STFT.
+    """
+    eps = 1e-10
+    hann_window = torch.hann_window(win_size).to(y.device)
+    stft_spec = torch.stft(
+                    y, n_fft, 
+                    hop_length=hop_size, 
+                    win_length=win_size, 
+                    window=hann_window,
+                    center=center, 
+                    pad_mode='reflect', 
+                    normalized=False, 
+                    return_complex=True)
+
+    if addeps==False:
+        mag = torch.abs(stft_spec)
+        pha = torch.angle(stft_spec)
+    else:
+        real_part = stft_spec.real
+        imag_part = stft_spec.imag
+        mag = torch.sqrt(real_part.pow(2) + imag_part.pow(2) + eps)
+        pha = torch.atan2(imag_part + eps, real_part + eps)
+    # Compress the magnitude
+    if compress_factor in ['log1p','relu_log1p', 'signed_log1p']:
+        mag = torch.log1p(mag)
+    else:
+        mag = torch.pow(mag, compress_factor)
+    com = torch.stack((mag * torch.cos(pha), mag * torch.sin(pha)), dim=-1)
+    return mag, pha, com
+
+
+def mag_phase_istft(mag, pha, n_fft, hop_size, win_size, compress_factor=1.0, center=True):
+    """
+    Inverse STFT to reconstruct the audio signal from magnitude and phase.
+
+    Args:
+        mag (torch.Tensor): Magnitude of the STFT.
+        pha (torch.Tensor): Phase of the STFT.
+        n_fft (int): FFT size.
+        hop_size (int): Hop size.
+        win_size (int): Window size.
+        compress_factor (float, optional): Magnitude compression factor. Defaults to 1.0.
+        center (bool, optional): Whether to center the signal before padding. Defaults to True.
+
+    Returns:
+        torch.Tensor: Reconstructed audio signal.
+    """
+    if compress_factor == 'log1p':
+        mag = torch.expm1(mag)
+    elif compress_factor == 'signed_log1p':
+        mag = decompress_signed_log1p(mag) 
+    elif compress_factor == 'relu_log1p':
+        mag = torch.expm1(RELU(mag))
+    else:
+        mag = torch.pow(RELU(mag), 1.0 / compress_factor)
+    com = torch.complex(mag * torch.cos(pha), mag * torch.sin(pha))
+    hann_window = torch.hann_window(win_size).to(com.device)
+    wav = torch.istft(
+                    com, 
+                    n_fft, 
+                    hop_length=hop_size, 
+                    win_length=win_size, 
+                    window=hann_window, 
+                    center=center)
+    return wav

models/streaming_codec_module_time_d1_input_ahead_sep_conv.py

@@ -0,0 +1,216 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import torch
+import torch.nn as nn
+import numpy as np
+from einops import rearrange
+import torch.nn.functional as F
+
+class ChannelLayerNorm(nn.Module):
+    """
+    LayerNorm over channels only. Normalizes over C only.
+    Input shape:  [B, C, T, F]
+    """
+    def __init__(self, num_channels):
+        super().__init__()
+        self.norm = nn.LayerNorm(num_channels)
+
+    def forward(self, x):
+        # [B, C, T, F] -> [B, T, F, C]
+        x = x.permute(0, 2, 3, 1)
+        # LayerNorm over C
+        x = self.norm(x)
+        # Back to [B, C, T, F]
+        x = x.permute(0, 3, 1, 2)
+        return x
+
+def get_causal_padding_2d(kernel_size, dilation=(1,1)):
+    """
+    Causal padding only along time axis.
+    Frequency axis uses symmetric padding.
+    """
+    pad_t =  kernel_size[0] * dilation[0] - dilation[0]  # all padding on left side
+    pad_f = (kernel_size[1] * dilation[1] - dilation[1]) // 2
+    return (pad_f, pad_f, pad_t, 0)  # (left, right, top, bottom)
+
+def get_causal_padding_2d_FT(kernel_size, dilation=(1,1)):
+    """
+    Causal padding only along time axis.
+    Frequency axis uses symmetric padding.
+    """
+    pad_t =  kernel_size[1] * dilation[1] - dilation[1]  # all padding on left side
+    pad_f = (kernel_size[0] * dilation[0] - dilation[0]) // 2
+    return (pad_t, 0, pad_f, pad_f)  # (left, right, top, bottom)
+
+class SPConvTranspose2d(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, padding, r=1):
+        super(SPConvTranspose2d, self).__init__()
+        self.pad = nn.ConstantPad2d(padding, value=0.)
+        self.out_channels = out_channels
+        self.r = r
+        self.conv = nn.Conv2d(in_channels, out_channels * r, kernel_size=kernel_size, padding=0)
+        
+
+    def forward(self, x):
+        x = self.pad(x)
+        out = self.conv(x)
+        batch_size, nchannels, H, W = out.shape
+        out = out.view((batch_size, self.r, nchannels // self.r, H, W))
+        out = out.permute(0, 2, 3, 4, 1)
+        out = out.contiguous().view((batch_size, nchannels // self.r, H, -1))
+        return out
+
+class DenseBlock(nn.Module):
+    """
+    DenseBlock module consisting of multiple convolutional layers with dilation.
+    """
+    def __init__(self, cfg, kernel_size=(3, 3), depth=4):
+        super(DenseBlock, self).__init__()
+        self.cfg = cfg
+        self.depth = depth
+        self.hid_feature = cfg['model_cfg']['hid_feature']
+        self.dense_block = nn.ModuleList()
+
+        for i in range(depth):
+            dil = 2 ** i 
+            pad = get_causal_padding_2d(kernel_size, (dil, 1))
+            dense_conv = nn.Sequential(
+                nn.ConstantPad2d(pad, 0.0),
+                nn.Conv2d(self.hid_feature * (i + 1), self.hid_feature, kernel_size, 
+                          dilation=(dil, 1), padding=0),
+                ChannelLayerNorm(self.hid_feature),
+                nn.PReLU(self.hid_feature)
+            )
+            self.dense_block.append(dense_conv)
+
+    def forward(self, x):
+        skip = x
+        for i in range(self.depth):
+            x = self.dense_block[i](skip)
+            skip = torch.cat([x, skip], dim=1)
+        return x
+
+class DenseEncoder(nn.Module):
+    """
+    DenseEncoder module consisting of initial convolution, dense block, and a final convolution.
+    """
+    def __init__(self, cfg):
+        super(DenseEncoder, self).__init__()
+        self.cfg = cfg
+        self.input_channel = cfg['model_cfg']['input_channel']
+        self.hid_feature = cfg['model_cfg']['hid_feature']
+
+        self.dense_conv_1_1 = nn.Sequential(
+            nn.Conv2d(self.input_channel, self.hid_feature, (3, 3)),
+            ChannelLayerNorm(self.hid_feature),
+            nn.PReLU(self.hid_feature)
+        )
+
+        self.dense_conv_1_2 = nn.Sequential(
+            nn.Conv2d(self.input_channel, self.hid_feature, (3, 3)),
+            ChannelLayerNorm(self.hid_feature),
+            nn.PReLU(self.hid_feature)
+        )
+
+        self.dense_conv_1_3 = nn.Sequential(
+            nn.Conv2d(self.input_channel, self.hid_feature, (3, 3)),
+            ChannelLayerNorm(self.hid_feature),
+            nn.PReLU(self.hid_feature)
+        )
+
+        self.dense_block = DenseBlock(cfg, depth=4)
+
+        self.dense_conv_2 = nn.Sequential(
+            nn.Conv2d(self.hid_feature, self.hid_feature, (1, 3), stride=(1, 2)),
+            ChannelLayerNorm(self.hid_feature),
+            nn.PReLU(self.hid_feature)
+        )
+
+    def forward(self, x, number_ahead):
+        x = F.pad(x, (1,1, 2-number_ahead, number_ahead), "constant", 0) # (pad_f, pad_f, pad_t, 0)  (left, right, top, bottom)
+        if number_ahead==0:
+            x = self.dense_conv_1_1(x)  # [batch, hid_feature, time, freq]
+        elif number_ahead==1:
+            x = self.dense_conv_1_2(x)  # [batch, hid_feature, time, freq]
+        elif number_ahead==2: 
+            x = self.dense_conv_1_3(x)  # [batch, hid_feature, time, freq]
+        else:
+            print('number_ahead not support!')
+
+        x = self.dense_block(x)   # [batch, hid_feature, time, freq]
+        x = self.dense_conv_2(x)  # [batch, hid_feature, time, freq//2]
+        return x
+
+class MagDecoder(nn.Module):
+    """
+    MagDecoder module for decoding magnitude information.
+    """
+    def __init__(self, cfg):
+        super(MagDecoder, self).__init__()
+        self.dense_block = DenseBlock(cfg, depth=4)
+        self.hid_feature = cfg['model_cfg']['hid_feature']
+        self.output_channel = cfg['model_cfg']['output_channel']
+        self.n_fft = cfg['stft_cfg']['n_fft']
+        self.beta = cfg['model_cfg']['beta']
+
+        self.up_conv1 = nn.Sequential(
+            SPConvTranspose2d(self.hid_feature, self.hid_feature, (1, 3), get_causal_padding_2d((1,3)), 2),
+            ChannelLayerNorm(self.hid_feature),
+            nn.PReLU(self.hid_feature)
+        )
+
+        self.up_conv2 = nn.Sequential(
+            SPConvTranspose2d(self.hid_feature, self.hid_feature, (1, 3), get_causal_padding_2d_FT((1,3)), 1),
+            ChannelLayerNorm(self.hid_feature),
+            nn.PReLU(self.hid_feature)
+        )
+
+        self.final_conv = nn.Conv2d(self.hid_feature, self.output_channel, (1, 1))
+        
+    def forward(self, x):
+        x = self.dense_block(x)
+        x = self.up_conv1(x)
+        x = self.up_conv2(x.permute(0,1,3,2)).permute(0,1,3,2)
+        x = self.final_conv(x)
+
+        return x
+
+class PhaseDecoder(nn.Module):
+    """
+    PhaseDecoder module for decoding phase information.
+    """
+    def __init__(self, cfg):
+        super(PhaseDecoder, self).__init__()
+        self.dense_block = DenseBlock(cfg, depth=4)
+        self.hid_feature = cfg['model_cfg']['hid_feature']
+        self.output_channel = cfg['model_cfg']['output_channel']
+
+        self.up_conv1 = nn.Sequential(
+            SPConvTranspose2d(self.hid_feature, self.hid_feature, (1, 3), get_causal_padding_2d((1,3)), 2),
+            ChannelLayerNorm(self.hid_feature),
+            nn.PReLU(self.hid_feature)
+        )
+
+        self.up_conv2 = nn.Sequential(
+            SPConvTranspose2d(self.hid_feature, self.hid_feature, (1, 3), get_causal_padding_2d_FT((1,3)), 1),
+            ChannelLayerNorm(self.hid_feature),
+            nn.PReLU(self.hid_feature)
+        )
+
+        self.phase_conv_r = nn.Conv2d(self.hid_feature, self.output_channel, (1, 1))
+        self.phase_conv_i = nn.Conv2d(self.hid_feature, self.output_channel, (1, 1))
+
+    def forward(self, x):
+        x = self.dense_block(x)
+        x = self.up_conv1(x)
+        x = self.up_conv2(x.permute(0,1,3,2)).permute(0,1,3,2)
+        x_r = self.phase_conv_r(x)
+        x_i = self.phase_conv_i(x)
+        x = torch.atan2(x_i, x_r)
+        return x

models/streaming_generator_SEMamba_time_d1_random_layer_ahead_sep_conv.py

@@ -0,0 +1,93 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+from huggingface_hub import PyTorchModelHubMixin
+from .streaming_mamba_block2_SEMamba import TFMambaBlock
+from .streaming_codec_module_time_d1_input_ahead_sep_conv import DenseEncoder, MagDecoder, PhaseDecoder
+
+class SEMamba_decoder_list(nn.Module, PyTorchModelHubMixin):
+    """
+    SEMamba model for speech enhancement using Mamba blocks.
+    
+    This model uses a dense encoder, multiple Mamba blocks, and separate magnitude
+    and phase decoders to process noisy magnitude and phase inputs.
+    """
+    def __init__(self, cfg):
+        """
+        Initialize the SEMamba model.
+        
+        Args:
+        - cfg: Configuration object containing model parameters.
+        """
+        super(SEMamba_decoder_list, self).__init__()
+        self.cfg = cfg
+        self.num_tscblocks = cfg['model_cfg']['num_tfmamba'] if cfg['model_cfg']['num_tfmamba'] is not None else 4  # default tfmamba: 4
+        self.mapping = cfg['model_cfg']['mapping'] if cfg['model_cfg']['mapping'] is not None else False  # default tfmamba: 4
+
+        # Initialize dense encoder
+        self.dense_encoder = DenseEncoder(cfg)
+
+        # Initialize Mamba blocks
+        self.TSMamba = nn.ModuleList([TFMambaBlock(cfg) for _ in range(self.num_tscblocks)])
+
+        # Initialize decoders
+        self.mask_decoder = MagDecoder(cfg)
+        self.phase_decoder = PhaseDecoder(cfg)
+
+        self.mask_decoder_list = nn.ModuleList([MagDecoder(cfg) for _ in range(self.num_tscblocks)])
+        self.phase_decoder_list = nn.ModuleList([PhaseDecoder(cfg) for _ in range(self.num_tscblocks)])
+
+    def forward(self, noisy_mag, noisy_pha, layer_use, number_ahead):
+        """
+        Forward pass for the SEMamba model.
+        
+        Args:
+        - noisy_mag (torch.Tensor): Noisy magnitude input tensor [B, F, T].
+        - noisy_pha (torch.Tensor): Noisy phase input tensor [B, F, T].
+        
+        Returns:
+        - denoised_mag (torch.Tensor): Denoised magnitude tensor [B, F, T].
+        - denoised_pha (torch.Tensor): Denoised phase tensor [B, F, T].
+        - denoised_com (torch.Tensor): Denoised complex tensor [B, F, T, 2].
+        """
+        # Reshape inputs
+        noisy_mag = rearrange(noisy_mag, 'b f t -> b t f').unsqueeze(1)  # [B, 1, T, F]
+        noisy_pha = rearrange(noisy_pha, 'b f t -> b t f').unsqueeze(1)  # [B, 1, T, F]
+
+        # Concatenate magnitude and phase inputs
+        x = torch.cat((noisy_mag, noisy_pha), dim=1)  # [B, 2, T, F]
+        
+        # Prevent unpredictable errors
+        B, C, T, F = x.shape
+        zeros = torch.zeros(B, C, T, 2, device=x.device)
+        x = torch.cat((x, zeros), dim=-1)
+        
+        # Encode input
+        x = self.dense_encoder(x, number_ahead)
+
+        # Apply Mamba blocks
+        for block in self.TSMamba[0:layer_use]:
+            x = block(x)
+
+        denoised_mag = rearrange(self.mask_decoder_list[layer_use-1](x), 'b c t f -> b f t c').squeeze(-1)
+        denoised_pha = rearrange(self.phase_decoder_list[layer_use-1](x), 'b c t f -> b f t c').squeeze(-1)
+        
+        # Prevent unpredictable errors
+        denoised_mag = denoised_mag[:, :F, :T]
+        denoised_pha = denoised_pha[:, :F, :T]
+                
+        # Combine denoised magnitude and phase into a complex representation
+        denoised_com = torch.stack(
+            (denoised_mag * torch.cos(denoised_pha), denoised_mag * torch.sin(denoised_pha)),
+            dim=-1
+        )
+
+        return denoised_mag, denoised_pha, denoised_com

models/streaming_mamba_block2_SEMamba.py

@@ -0,0 +1,102 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn import init
+from torch.nn.parameter import Parameter
+from functools import partial
+from einops import rearrange
+from mamba_ssm import Mamba
+
+class Causal_MambaBlock(nn.Module):
+    def __init__(self, d_model, cfg):
+        super(Causal_MambaBlock, self).__init__()
+        
+        d_state = cfg['model_cfg']['d_state'] # 16
+        d_conv = cfg['model_cfg']['d_conv'] # 4
+        expand = cfg['model_cfg']['expand'] # 4
+
+        self.forward_blocks  = Mamba(d_model=d_model, d_state=d_state, d_conv=d_conv, expand=expand)
+        self.output_proj = nn.Linear(d_model, d_model)
+        self.norm = nn.LayerNorm(d_model)
+
+    def forward(self, x):
+        # x: [B, T, D]
+        out_fw = self.forward_blocks(x) + x
+        out = self.output_proj(out_fw)
+
+        # LayerNorm
+        return self.norm(out)
+
+class MambaBlock(nn.Module):
+    def __init__(self, d_model, cfg):
+        super(MambaBlock, self).__init__()
+        
+        d_state = cfg['model_cfg']['d_state'] # 16
+        d_conv = cfg['model_cfg']['d_conv'] # 4
+        expand = cfg['model_cfg']['expand'] # 4
+
+        self.forward_blocks  = Mamba(d_model=d_model, d_state=d_state, d_conv=d_conv, expand=expand)
+        self.backward_blocks = Mamba(d_model=d_model, d_state=d_state, d_conv=d_conv, expand=expand) 
+        self.output_proj = nn.Linear(2 * d_model, d_model)
+        self.norm = nn.LayerNorm(d_model)
+
+    def forward(self, x):
+        # x: [B, T, D]
+        out_fw = self.forward_blocks(x) + x
+
+        out_bw = self.backward_blocks(torch.flip(x, dims=[1])) + torch.flip(x, dims=[1])
+        out_bw = torch.flip(out_bw, dims=[1])
+        
+        out = torch.cat([out_fw, out_bw], dim=-1)
+        out = self.output_proj(out)
+
+        # LayerNorm
+        return self.norm(out)
+
+
+class TFMambaBlock(nn.Module):
+    """
+    Temporal-Frequency Mamba block for sequence modeling.
+    
+    Attributes:
+    cfg (Config): Configuration for the block.
+    time_mamba (MambaBlock): Mamba block for temporal dimension.
+    freq_mamba (MambaBlock): Mamba block for frequency dimension.
+    tlinear (ConvTranspose1d): ConvTranspose1d layer for temporal dimension.
+    flinear (ConvTranspose1d): ConvTranspose1d layer for frequency dimension.
+    """
+    def __init__(self, cfg):
+        super(TFMambaBlock, self).__init__()
+        self.cfg = cfg
+        self.hid_feature = cfg['model_cfg']['hid_feature']
+        
+        # Initialize Mamba blocks
+        self.time_mamba = Causal_MambaBlock(d_model=self.hid_feature, cfg=cfg)
+        self.freq_mamba = MambaBlock(d_model=self.hid_feature, cfg=cfg)
+    
+    def forward(self, x):
+        """
+        Forward pass of the TFMamba block.
+        
+        Parameters:
+        x (Tensor): Input tensor with shape (batch, channels, time, freq).
+        
+        Returns:
+        Tensor: Output tensor after applying temporal and frequency Mamba blocks.
+        """
+        b, c, t, f = x.size()
+        x = x.permute(0, 3, 2, 1).contiguous().view(b*f, t, c)
+        x = self.time_mamba(x) + x
+        x = x.view(b, f, t, c).permute(0, 2, 1, 3).contiguous().view(b*t, f, c)
+        x = self.freq_mamba(x) + x
+        x = x.view(b, t, f, c).permute(0, 3, 1, 2)
+        return x
+

noisy_audio/mic_test.wav

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

offline_enhanced_audio/mic_test.flac

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

offline_inference.py

@@ -0,0 +1,116 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+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.streaming_generator_SEMamba_time_d1_random_layer_ahead_sep_conv import SEMamba_decoder_list
+from utils.util import load_config, pad_or_trim_to_match
+from huggingface_hub import hf_hub_download
+RELU = nn.ReLU()
+
+config_path = hf_hub_download(repo_id="nvidia/Real-time_RE-USE", filename="config.json")
+
+def get_filepaths(directory, file_type=None):
+    file_paths = []  # List which will store all of the full filepaths.
+    # Walk the tree.
+    for root, directories, files in os.walk(directory):
+        for filename in files:
+            # Join the two strings in order to form the full filepath.
+            filepath = os.path.join(root, filename)
+            if file_type is not None:
+                if filepath.split('.')[-1] == file_type:
+                    file_paths.append(filepath)  # Add it to the list.
+            else:
+                file_paths.append(filepath)  # Add it to the list.
+    return file_paths  # Self-explanatory.
+
+def make_even(value):
+    value = int(round(value))
+    return value if value % 2 == 0 else value + 1
+
+def inference(args, device):
+    cfg = load_config(args.config)
+    n_fft, hop_size, win_size = cfg['stft_cfg']['n_fft'], cfg['stft_cfg']['hop_size'], cfg['stft_cfg']['win_size']
+    compress_factor = cfg['model_cfg']['compress_factor']
+    sampling_rate = cfg['stft_cfg']['sampling_rate']
+
+    SE_model = SEMamba_decoder_list.from_pretrained("nvidia/Real-time_RE-USE", cfg=cfg).to(device)
+    SE_model.eval()
+
+    os.makedirs(args.output_folder, exist_ok=True)
+    with torch.no_grad():
+        for i, fname in enumerate(get_filepaths(args.input_folder)):
+            print(fname)
+            try:
+                os.makedirs(args.output_folder + fname[0:fname.rfind('/')].replace(args.input_folder,''), exist_ok=True)
+                noisy_wav, noisy_sr = torchaudio.load(fname)
+            except Exception as e:
+                print(f"Warning: cannot read {fname}, skipping. ({e})")
+                continue
+            
+            if args.BWE is not None:
+                opts = {"res_type": "kaiser_best"}
+                noisy_wav = librosa.resample(noisy_wav.cpu().numpy(), orig_sr=noisy_sr, target_sr=int(args.BWE), **opts)
+                noisy_sr = int(args.BWE)
+
+            noisy_wav = torch.FloatTensor(noisy_wav).to(device)
+            n_fft_scaled = make_even(n_fft * noisy_sr // sampling_rate)
+            hop_size_scaled = make_even(hop_size * noisy_sr // sampling_rate)
+            win_size_scaled = make_even(win_size * noisy_sr // sampling_rate)
+
+            noisy_mag, noisy_pha, noisy_com = mag_phase_stft(
+                noisy_wav,
+                n_fft=n_fft_scaled,
+                hop_size=hop_size_scaled,
+                win_size=win_size_scaled,
+                compress_factor=compress_factor,
+                center=True,
+                addeps=False
+            )
+
+            # Offline inference!!
+            amp_g, pha_g, _ = SE_model(noisy_mag, noisy_pha, args.Exit_layer, args.look_ahead_frames)
+            # 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
+            assert audio_g.shape == noisy_wav.shape, audio_g.shape
+
+            output_file = os.path.join(args.output_folder + fname.replace(args.input_folder,'').split('.')[0]+'.flac') # save to .flac format
+            torchaudio.save(output_file, audio_g.cpu(), noisy_sr)
+            
+def main():
+    print('Initializing Inference Process...')
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--input_folder')
+    parser.add_argument('--output_folder')
+    parser.add_argument('--config')
+    parser.add_argument('--Exit_layer', type=int, required=True)
+    parser.add_argument('--look_ahead_frames', type=int, required=True)
+    parser.add_argument('--BWE', default=None)
+    args = parser.parse_args()
+
+    global device
+    if torch.cuda.is_available():
+        device = torch.device('cuda')
+    else:
+        raise RuntimeError("Currently, CPU mode is not supported.")
+        
+    inference(args, device)
+
+
+if __name__ == '__main__':
+    main()

offline_inference.sh

@@ -0,0 +1,8 @@
+CUDA_VISIBLE_DEVICES='0' python offline_inference.py \
+   --input_folder ./noisy_audio \
+   --output_folder ./offline_enhanced_audio \
+   --config recipes/USEMamba_12x1_lr_00002_norm_05_vq_067_nfft_320_hop_160_NRIR_012_pha_0005_com_04_early_005_release_random_layer_GAN_longer_1k.yaml \
+   --Exit_layer 8 \
+   --look_ahead_frames 0 \
+   #--BWE 16000 \
+# Exit_layer can be between 3~12 , look_ahead_frames can be between 0~2

online_inference.py

@@ -0,0 +1,470 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import torch
+import torchaudio
+import torchaudio.transforms as T
+import torch.nn.functional as F
+import torch.nn as nn
+import argparse
+import librosa
+import numpy as np
+import time
+from models.stfts import mag_phase_stft, mag_phase_istft
+from models.streaming_generator_SEMamba_time_d1_random_layer_ahead_sep_conv import SEMamba_decoder_list
+from utils.util import load_config, pad_or_trim_to_match
+from huggingface_hub import hf_hub_download
+RELU = nn.ReLU()
+
+config_path = hf_hub_download(repo_id="nvidia/Real-time_RE-USE", filename="config.json")
+################### Streaming inference modification start ######################################
+def get_causal_padding_2d(kernel_size, dilation=(1,1)):
+    """
+    Causal padding only along time axis.
+    Frequency axis uses symmetric padding.
+    """
+    pad_t =  kernel_size[0] * dilation[0] - dilation[0]  # all padding on left side
+    pad_f = (kernel_size[1] * dilation[1] - dilation[1]) // 2
+    return (pad_f, pad_f, pad_t, 0)  # (left, right, top, bottom)
+
+class StreamingCausalConv2d(nn.Module):
+    def __init__(self, conv: nn.Conv2d):
+        super().__init__()
+        self.conv = conv
+        kt, kf = conv.kernel_size
+        dt, df = conv.dilation
+        self.cache_len = (kt - 1) * dt
+        self.register_buffer("cache", None, persistent=False)
+        self.freq_pad = get_causal_padding_2d((1, kf), (dt, df))
+
+    def reset(self):
+        self.cache = None
+
+    def forward(self, x, t_pad=None):
+        # x: [B, C, 1, f]
+        B, C, _, f = x.shape
+
+        if self.cache is None:
+            # Assign the instance variable if no specific pad is provided
+            if t_pad is None:
+                t_pad = self.cache_len
+            self.cache = torch.zeros(
+                B, C, t_pad, f,
+                device=x.device,
+                dtype=x.dtype
+            )
+
+        x_cat = torch.cat([self.cache, x], dim=2)
+        y = self.conv(F.pad(x_cat, self.freq_pad, "constant", 0))
+        self.cache = x_cat[:, :, -self.cache_len:, :]
+        return y
+
+class StreamingCausalConv2d_FT(nn.Module):
+    def __init__(self, conv: nn.Conv2d):
+        super().__init__()
+        self.conv = conv
+        kf, kt = conv.kernel_size
+        df, dt = conv.dilation
+        self.cache_len = (kt - 1) * dt
+        self.register_buffer("cache", None, persistent=False)
+        #self.freq_pad = get_causal_padding_2d_FT((kf, 1), (df, dt))
+
+    def reset(self):
+        self.cache = None
+
+    def forward(self, x):
+        # x: [B, C, f, 1]
+        B, C, f, _ = x.shape
+
+        if self.cache is None:
+            self.cache = torch.zeros(
+                B, C, f, self.cache_len,
+                device=x.device,
+                dtype=x.dtype
+            )
+        x_cat = torch.cat([self.cache, x], dim=3)
+        y = self.conv(x_cat)
+        self.cache = x_cat[:, :, :, -self.cache_len:]
+        return y
+
+class StreamingDenseBlock(nn.Module):
+    def __init__(self, dense_block):
+        super().__init__()
+        self.blocks = nn.ModuleList()
+
+        for seq in dense_block.dense_block:
+            conv = seq[1]
+            self.blocks.append(nn.Sequential(
+                StreamingCausalConv2d(conv),
+                seq[2],  # ChannelLayerNorm
+                seq[3],  # PReLU
+            ))
+
+    def reset(self):
+        for b in self.blocks:
+            b[0].reset()
+
+    def forward(self, x):
+        skip = x
+        for b in self.blocks:
+            y = b(skip)
+            skip = torch.cat([y, skip], dim=1)
+        return y
+
+class StreamingDenseEncoder(nn.Module):
+    def __init__(self, encoder):
+        super().__init__()
+
+        # dense_conv_1_1
+        conv1_1 = encoder.dense_conv_1_1[0]
+        self.conv1_1 = nn.Sequential(
+            StreamingCausalConv2d(conv1_1),
+            encoder.dense_conv_1_1[1],
+            encoder.dense_conv_1_1[2],
+        )
+
+        # dense_conv_1_2
+        conv1_2 = encoder.dense_conv_1_2[0]
+        self.conv1_2 = nn.Sequential(
+            StreamingCausalConv2d(conv1_2),
+            encoder.dense_conv_1_2[1],
+            encoder.dense_conv_1_2[2],
+        )
+
+        # dense_conv_1_3
+        conv1_3 = encoder.dense_conv_1_3[0]
+        self.conv1_3 = nn.Sequential(
+            StreamingCausalConv2d(conv1_3),
+            encoder.dense_conv_1_3[1],
+            encoder.dense_conv_1_3[2],
+        )
+
+        self.dense_block = StreamingDenseBlock(encoder.dense_block)
+
+        # time kernel = 1 → no cache needed
+        self.conv2 = encoder.dense_conv_2
+
+    def reset(self):
+        self.conv1_1[0].reset()
+        self.conv1_2[0].reset()
+        self.conv1_3[0].reset()
+        self.dense_block.reset()
+
+    def forward(self, x, look_ahead_frames):
+        # x: [B, C, 1, F]
+        #x = F.pad(x, (1, 1, 2-number_ahead, number_ahead), "constant", 0)     
+        if look_ahead_frames == 0:
+            seq, pad = self.conv1_1, 2
+        elif look_ahead_frames == 1:
+            seq, pad = self.conv1_2, 1
+        elif look_ahead_frames == 2:
+            seq, pad = self.conv1_3, 0
+        else:
+            print('look_ahead_frames not support!')
+        
+        # Manually pipe through the first layer with the extra argument
+        # seq[0] is the StreamingCausalConv2d
+        x = seq[0](x, t_pad=pad)
+
+        # Pass through the remaining layers in the Sequential
+        for i in range(1, len(seq)):
+            x = seq[i](x)
+
+        x = self.dense_block(x)
+        x = self.conv2(x)
+        return x
+
+class StreamingMagDecoder(nn.Module):
+    def __init__(self, decoder):
+        super().__init__()
+
+        self.up_conv1 = decoder.up_conv1
+        self.up_conv2 = nn.Sequential(
+            StreamingCausalConv2d_FT(decoder.up_conv2[0].conv),
+            decoder.up_conv2[1],
+            decoder.up_conv2[2],
+        )
+        self.final_conv = decoder.final_conv
+        self.dense_block = StreamingDenseBlock(decoder.dense_block)
+
+    def reset(self):
+        self.dense_block.reset()
+        self.up_conv2[0].reset()
+
+    def forward(self, x):
+        # x: [B, C, 1, F]
+        x = self.dense_block(x)
+        x = self.up_conv1(x)
+        x = self.up_conv2(x.permute(0,1,3,2)).permute(0,1,3,2)
+        x = self.final_conv(x)
+        return x
+
+class StreamingPhaseDecoder(nn.Module):
+    def __init__(self, decoder):
+        super().__init__()
+
+        self.up_conv1 = decoder.up_conv1
+        self.up_conv2 = nn.Sequential(
+            StreamingCausalConv2d_FT(decoder.up_conv2[0].conv),
+            decoder.up_conv2[1],
+            decoder.up_conv2[2],
+        )
+        self.phase_conv_r = decoder.phase_conv_r
+        self.phase_conv_i = decoder.phase_conv_i
+        self.dense_block = StreamingDenseBlock(decoder.dense_block)
+
+    def reset(self):
+        self.dense_block.reset()
+        self.up_conv2[0].reset()
+
+    def forward(self, x):
+        # x: [B, C, 1, F]
+        x = self.dense_block(x)
+        x = self.up_conv1(x)
+        x = self.up_conv2(x.permute(0,1,3,2)).permute(0,1,3,2)
+        x_r = self.phase_conv_r(x)
+        x_i = self.phase_conv_i(x)
+        x = torch.atan2(x_i, x_r)
+        return x
+
+class StreamingCausalMambaBlock(nn.Module):
+    def __init__(self, block):
+        super().__init__()
+        self.mamba = block.forward_blocks
+        self.proj = block.output_proj
+        self.norm = block.norm
+        self.register_buffer("conv_state", None, persistent=False)
+        self.register_buffer("ssm_state", None, persistent=False)
+
+    def reset(self, batch_size):
+        self.conv_state, self.ssm_state = self.mamba.allocate_inference_cache(
+            batch_size=batch_size, # Frequecy_dim
+            max_seqlen=1,
+        )
+
+    def forward(self, x):
+        """
+        x: [B, 1, D]
+        """
+        y, self.conv_state, self.ssm_state = self.mamba.step(x, self.conv_state, self.ssm_state)
+        y = y + x
+        y = self.proj(y)
+        return self.norm(y)
+
+class StreamingTFMambaBlock(nn.Module):
+    def __init__(self, block):
+        super().__init__()
+        self.time_mamba = StreamingCausalMambaBlock(block.time_mamba)
+        self.freq_mamba = block.freq_mamba  # freq is full, no streaming needed
+
+    def reset(self, batch_size):
+        self.time_mamba.reset(batch_size)
+
+    def forward(self, x):
+        # x: [B, C, 1, F]
+        B, C, _, F = x.shape
+
+        # ---- Time Mamba (true streaming) ----
+        x_t = x.permute(0, 3, 2, 1).reshape(B * F, 1, C)
+        x_t = self.time_mamba(x_t) + x_t
+        x_t = x_t.view(B, F, 1, C).permute(0, 3, 2, 1)
+
+        # ---- Frequency Mamba (non-causal, full freq) ----
+        x_f = x_t.permute(0, 2, 3, 1).reshape(B, F, C)
+        x_f = self.freq_mamba(x_f) + x_f
+        x_f = x_f.view(B, 1, F, C).permute(0, 3, 1, 2)
+
+        return x_f
+
+class StreamingSEMamba(nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.encoder = StreamingDenseEncoder(model.dense_encoder)
+        self.mamba = nn.ModuleList(
+            [StreamingTFMambaBlock(b) for b in model.TSMamba]
+        )
+        self.mag_decoder_list = nn.ModuleList([StreamingMagDecoder(decoder) for decoder in model.mask_decoder_list])
+        self.pha_decoder_list = nn.ModuleList([StreamingPhaseDecoder(decoder) for decoder in model.phase_decoder_list])
+
+    def reset(self, freq_dim):
+        self.encoder.reset()
+        for l in self.mamba:
+            l.reset(freq_dim)
+        for dec in self.mag_decoder_list:
+            dec.reset()
+        for dec in self.pha_decoder_list:
+            dec.reset()
+
+    def forward(self, noisy_mag_t, noisy_pha_t, layer_use, look_ahead_frames):
+        """
+        noisy_mag_t: [B, F, T]
+        noisy_pha_t: [B, F, T]
+        """
+        x_mag = noisy_mag_t.permute(0,2,1).unsqueeze(1)  # [B,1,T,F]
+        x_pha = noisy_pha_t.permute(0,2,1).unsqueeze(1)  # [B,1,T,F]
+        x = torch.cat([x_mag, x_pha], dim=1)           # [B,2,T,F]
+
+        # match original zero-padding logic
+        zeros = torch.zeros(x.size(0), x.size(1), x.size(2), 2, device=x.device)
+        x = torch.cat([x, zeros], dim=-1)
+
+        # ---- Encoder ----
+        x = self.encoder(x, look_ahead_frames)
+
+        # ---- Mamba ----
+        #import pdb; pdb.set_trace()
+        for b in self.mamba[0:layer_use]:
+            x = b(x)
+
+        # ---- Decoders ----
+        mag = self.mag_decoder_list[layer_use-1](x)
+        pha = self.pha_decoder_list[layer_use-1](x)
+
+        mag = mag.squeeze(2).squeeze(1)  # [B,F]
+        pha = pha.squeeze(2).squeeze(1)
+
+         # Prevent unpredictable errors
+        mag = mag[:, 0:x_mag.shape[-1]]
+        pha = pha[:, 0:x_mag.shape[-1]]
+
+        return mag, pha
+
+################### Streaming inference modification end ######################################
+
+def make_even(value):
+    value = int(round(value))
+    return value if value % 2 == 0 else value + 1
+
+def inference(args, device):
+    cfg = load_config(args.config)
+    n_fft, hop_size, win_size = cfg['stft_cfg']['n_fft'], cfg['stft_cfg']['hop_size'], cfg['stft_cfg']['win_size']
+    compress_factor = cfg['model_cfg']['compress_factor']
+    sampling_rate = cfg['stft_cfg']['sampling_rate']
+
+    model = SEMamba_decoder_list.from_pretrained("nvidia/Real-time_RE-USE", cfg=cfg).to(device)
+    model.eval()
+    streaming_model = StreamingSEMamba(model).eval()
+
+    with torch.no_grad():
+        noisy_wav, noisy_sr = torchaudio.load('./noisy_audio/mic_test.wav')
+
+        # Leave online BWE as future work:   
+        if args.BWE is not None:
+            opts = {"res_type": "kaiser_fast"}
+            noisy_wav = librosa.resample(noisy_wav.cpu().numpy(), orig_sr=noisy_sr, target_sr=int(args.BWE), **opts)
+            noisy_sr = int(args.BWE)
+        noisy_wav = torch.FloatTensor(noisy_wav).to(device)
+        
+        n_fft_scaled = make_even(n_fft * noisy_sr // sampling_rate)
+        hop_size_scaled = make_even(hop_size * noisy_sr // sampling_rate)
+        win_size_scaled = make_even(win_size * noisy_sr // sampling_rate)
+
+        # Leave online STFT as future work:
+        noisy_mag, noisy_pha, noisy_com = mag_phase_stft(  # (B, F, T)
+            noisy_wav,
+            n_fft=n_fft_scaled,
+            hop_size=hop_size_scaled,
+            win_size=win_size_scaled,
+            compress_factor=compress_factor,
+            center=True,
+            addeps=False
+        )
+
+        ## Offline inference!!
+        mag_out2, pha_out2, _ = model(noisy_mag, noisy_pha, args.Exit_layer, args.look_ahead_frames)
+        # To remove "strange sweep artifact"
+        mag2 = torch.expm1(RELU(mag_out2)) # [1, F, T]
+        zero_portion = torch.sum(mag2==0, 1)/mag2.shape[1]
+        mag_out2[:,:,(zero_portion>0.5)[0]] = 0
+        audio_g2 = mag_phase_istft(mag_out2, pha_out2, n_fft_scaled, hop_size_scaled, win_size_scaled, compress_factor)
+        audio_g2 = pad_or_trim_to_match(noisy_wav.detach(), audio_g2, pad_value=1e-8)  # Align lengths using epsilon padding
+        torchaudio.save('./enhanced_audio/offline_enhanced_mic_test.flac', audio_g2.cpu(), noisy_sr)
+        
+        ## Online inference!! (one frame in, one frame out)
+        streaming_model.reset(noisy_mag.shape[1]//2+1)
+        noisy_mag = torch.cat([noisy_mag, torch.zeros(noisy_mag.shape[0],noisy_mag.shape[1], args.look_ahead_frames, device=noisy_mag.device)], dim=-1)
+        noisy_pha = torch.cat([noisy_pha, torch.zeros(noisy_pha.shape[0],noisy_pha.shape[1], args.look_ahead_frames, device=noisy_pha.device)], dim=-1)
+        
+        times, mag_out, pha_out = [], [], []
+        print('Start Online inferencing...')
+        # 1. wait for the enough look_ahead_frames
+        start = time.perf_counter()
+        mag_t, pha_t = streaming_model(
+            noisy_mag[:, :, 0:args.look_ahead_frames+1], 
+            noisy_pha[:, :, 0:args.look_ahead_frames+1], 
+            args.Exit_layer, 
+            args.look_ahead_frames
+            )
+        torch.cuda.synchronize()
+        end = time.perf_counter()
+        times.append(end - start)
+        mag_out.append(mag_t)
+        pha_out.append(pha_t)
+        
+        # 2. frame by frame inference
+        for t in range(args.look_ahead_frames+1, noisy_mag.shape[-1]): 
+            start = time.perf_counter()      
+            mag_t, pha_t = streaming_model(
+                noisy_mag[:, :, t:t+1],
+                noisy_pha[:, :, t:t+1],
+                args.Exit_layer, 
+                args.look_ahead_frames
+            )
+            torch.cuda.synchronize()
+            end = time.perf_counter()
+
+            # To remove "strange sweep artifact"
+            mag = torch.expm1(RELU(mag_t)) # [1, F, 1]
+            zero_portion = torch.sum(mag==0, 1)/mag.shape[1]
+            if zero_portion.item()>0.5:
+                mag_t = 0*mag_t
+            times.append(end - start)
+            mag_out.append(mag_t)
+            pha_out.append(pha_t)
+
+        mag_out = torch.stack(mag_out, dim=-1)
+        pha_out = torch.stack(pha_out, dim=-1)
+
+        total_audio_time = noisy_wav.shape[-1] / noisy_sr
+        
+        audio_g = mag_phase_istft(mag_out, pha_out, 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
+        torchaudio.save('./enhanced_audio/online_enhanced_mic_test.flac', audio_g.cpu(), noisy_sr)
+        print(f"Max waveform difference of Online and Offline inference: {(audio_g2-audio_g).max():.7f}")
+        #import pdb; pdb.set_trace()
+    return noisy_sr, sum(times)/len(times)*1000, sum(times)/total_audio_time
+            
+
+            
+def main():
+    print('Initializing Inference Process..')
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--config', default='results')
+    parser.add_argument('--Exit_layer', type=int, required=True)
+    parser.add_argument('--look_ahead_frames', type=int, required=True)
+    parser.add_argument('--BWE', default=None)
+    args = parser.parse_args()
+
+    global device
+    if torch.cuda.is_available():
+        device = torch.device('cuda')
+    else:
+        raise RuntimeError("Currently, CPU mode is not supported.")
+        
+    sr, latency_ms, rtf = inference(args, device)
+    #print(args.checkpoint_file)
+    print(f"Layer use: {args.Exit_layer:.0f}")
+    print(f"Look ahead frames: {args.look_ahead_frames:.0f}")
+    print(f"Sampling rate: {sr:.3f}")
+    print(f"Online Latency per chunk: {latency_ms:.3f} ms")
+    print(f"Online Real-Time Factor (RTF): {rtf:.3f}")
+
+
+if __name__ == '__main__':
+    main()
+

online_inference.sh

@@ -0,0 +1,7 @@
+pip install triton==2.3.0
+CUDA_VISIBLE_DEVICES='0' python online_inference.py \
+   --config recipes/USEMamba_12x1_lr_00002_norm_05_vq_067_nfft_320_hop_160_NRIR_012_pha_0005_com_04_early_005_release_random_layer_GAN_longer_1k.yaml \
+   --Exit_layer 8 \
+   --look_ahead_frames 0 \
+   #--BWE 16000 \
+# Exit_layer can be between 3~12 , look_ahead_frames can be between 0~2

recipes/USEMamba_12x1_lr_00002_norm_05_vq_067_nfft_320_hop_160_NRIR_012_pha_0005_com_04_early_005_release_random_layer_GAN_longer_1k.yaml

@@ -0,0 +1,44 @@
+# Environment Settings
+# These settings specify the hardware and distributed setup for the model training.
+# Adjust `num_gpus` and `dist_config` according to your distributed training environment.
+env_setting:
+  num_gpus: 8  # Number of GPUs. Now we don't support CPU mode. 
+  num_workers: 20 # 0  Number of worker threads for data loading.
+  persistent_workers: True # False  If you have large RAM, turn this to be True
+  prefetch_factor: 8 # null 
+  seed: 1234  # Seed for random number generators to ensure reproducibility.
+  stdout_interval: 1000
+  checkpoint_interval: 1000  # save model to ckpt every N steps
+  validation_interval: 1000 
+  dist_cfg:
+    dist_backend: nccl  # Distributed training backend, 'nccl' for NVIDIA GPUs.
+    dist_url: tcp://localhost:19478  # URL for initializing distributed training.
+    world_size: 1  # Total number of processes in the distributed training.
+  pin_memory: True # If you have large RAM, turn this to be True
+
+# STFT Configuration
+# Configuration for Short-Time Fourier Transform (STFT), crucial for audio processing models.
+stft_cfg:
+  sampling_rate: 8000  # Audio sampling rate in Hz.
+  n_fft: 320  # FFT components for transforming audio signals.
+  hop_size: 160  # Samples between successive frames.
+  win_size: 320  # Window size used in FFT.
+  sfi: True # Sampline Frequency Independent
+
+# Model Configuration
+# Defines the architecture specifics of the model, including layer configurations and feature compression.
+model_cfg:
+  hid_feature: 64  # Channels in dense layers.
+  compress_factor: relu_log1p  # Compression factor applied to extracted features.
+  num_tfmamba: 12  # Number of Time-Frequency Mamba (TFMamba) blocks in the model.
+  d_state: 16  # Dimensionality of the state vector in Mamba blocks.
+  d_conv: 4  # Convolutional layer dimensionality within Mamba blocks.
+  expand: 4  # Expansion factor for the layers within the Mamba blocks.
+  norm_epsilon: 0.00001  # Numerical stability in normalization layers within the Mamba blocks.
+  beta: 2.0  # Hyperparameter for the Learnable Sigmoid function.
+  input_channel: 2 # Magnitude and Phase
+  output_channel: 1  # Single Channel Speech Enhancement
+  inner_mamba_nlayer: 1 # Number of layer of Mamba in Bidirectional Mamba
+  nonlinear: None # last activation function for the mag encoder. 'softplus' or 'relu'
+  mapping: True # Otherwise, this should be masking model
+  maximum_layer: 12

utils/util.py

@@ -0,0 +1,37 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import yaml
+import torch
+import os
+import shutil
+import torch.nn.functional as F
+
+def load_config(config_path):
+    """Load configuration from a YAML file."""
+    with open(config_path, 'r') as file:
+        return yaml.safe_load(file)
+
+def pad_or_trim_to_match(reference: torch.Tensor, target: torch.Tensor, pad_value: float = 1e-6) -> torch.Tensor:
+    """
+    Extends the target tensor to match the reference tensor along dim=1
+    without breaking autograd, by creating a new tensor and copying data in.
+    """
+    B, ref_len = reference.shape
+    _, tgt_len = target.shape
+
+    if tgt_len == ref_len:
+        return target
+    elif tgt_len > ref_len:
+        return target[:, :ref_len]
+    
+    # Allocate padded tensor with grad support
+    padded = torch.full((B, ref_len), pad_value, dtype=target.dtype, device=target.device)
+    padded[:, :tgt_len] = target  # This preserves gradient tracking
+
+    return padded