# 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()