ASesYusuf1/SESA_Audio_Separation

1import torch2import torch.nn as nn3from torch.nn.modules.rnn import LSTM4 5 6class FeatureConversion(nn.Module):7    """8    Integrates into the adjacent Dual-Path layer.9 10    Args:11        channels (int): Number of input channels.12        inverse (bool): If True, uses ifft; otherwise, uses rfft.13    """14 15    def __init__(self, channels, inverse):16        super().__init__()17        self.inverse = inverse18        self.channels = channels19 20    def forward(self, x):21        # B, C, F, T = x.shape22        if self.inverse:23            x = x.float()24            x_r = x[:, :self.channels // 2, :, :]25            x_i = x[:, self.channels // 2:, :, :]26            x = torch.complex(x_r, x_i)27            x = torch.fft.irfft(x, dim=3, norm="ortho")28        else:29            x = x.float()30            x = torch.fft.rfft(x, dim=3, norm="ortho")31            x_real = x.real32            x_imag = x.imag33            x = torch.cat([x_real, x_imag], dim=1)34        return x35 36 37class DualPathRNN(nn.Module):38    """39    Dual-Path RNN in Separation Network.40 41    Args:42        d_model (int): The number of expected features in the input (input_size).43        expand (int): Expansion factor used to calculate the hidden_size of LSTM.44        bidirectional (bool): If True, becomes a bidirectional LSTM.45    """46 47    def __init__(self, d_model, expand, bidirectional=True):48        super(DualPathRNN, self).__init__()49 50        self.d_model = d_model51        self.hidden_size = d_model * expand52        self.bidirectional = bidirectional53        # Initialize LSTM layers and normalization layers54        self.lstm_layers = nn.ModuleList([self._init_lstm_layer(self.d_model, self.hidden_size) for _ in range(2)])55        self.linear_layers = nn.ModuleList([nn.Linear(self.hidden_size * 2, self.d_model) for _ in range(2)])56        self.norm_layers = nn.ModuleList([nn.GroupNorm(1, d_model) for _ in range(2)])57 58    def _init_lstm_layer(self, d_model, hidden_size):59        return LSTM(d_model, hidden_size, num_layers=1, bidirectional=self.bidirectional, batch_first=True)60 61    def forward(self, x):62        B, C, F, T = x.shape63 64        # Process dual-path rnn65        original_x = x66        # Frequency-path67        x = self.norm_layers[0](x)68        x = x.transpose(1, 3).contiguous().view(B * T, F, C)69        x, _ = self.lstm_layers[0](x)70        x = self.linear_layers[0](x)71        x = x.view(B, T, F, C).transpose(1, 3)72        x = x + original_x73 74        original_x = x75        # Time-path76        x = self.norm_layers[1](x)77        x = x.transpose(1, 2).contiguous().view(B * F, C, T).transpose(1, 2)78        x, _ = self.lstm_layers[1](x)79        x = self.linear_layers[1](x)80        x = x.transpose(1, 2).contiguous().view(B, F, C, T).transpose(1, 2)81        x = x + original_x82 83        return x84 85 86class SeparationNet(nn.Module):87    """88    Implements a simplified Sparse Down-sample block in an encoder architecture.89 90    Args:91    - channels (int): Number input channels.92    - expand (int): Expansion factor used to calculate the hidden_size of LSTM.93    - num_layers (int): Number of dual-path layers.94    """95 96    def __init__(self, channels, expand=1, num_layers=6):97        super(SeparationNet, self).__init__()98 99        self.num_layers = num_layers100 101        self.dp_modules = nn.ModuleList([102            DualPathRNN(channels * (2 if i % 2 == 1 else 1), expand) for i in range(num_layers)103        ])104 105        self.feature_conversion = nn.ModuleList([106            FeatureConversion(channels * 2, inverse=False if i % 2 == 0 else True) for i in range(num_layers)107        ])108 109    def forward(self, x):110        for i in range(self.num_layers):111            x = self.dp_modules[i](x)112            x = self.feature_conversion[i](x)113        return x114