Clicko777/RVC_HFv2

1import torch2from torch.nn.functional import conv1d, conv2d3from typing import Union, Optional4from .utils import linspace, temperature_sigmoid, amp_to_db5 6 7class TorchGate(torch.nn.Module):8    """9    A PyTorch module that applies a spectral gate to an input signal.10 11    Arguments:12        sr {int} -- Sample rate of the input signal.13        nonstationary {bool} -- Whether to use non-stationary or stationary masking (default: {False}).14        n_std_thresh_stationary {float} -- Number of standard deviations above mean to threshold noise for15                                           stationary masking (default: {1.5}).16        n_thresh_nonstationary {float} -- Number of multiplies above smoothed magnitude spectrogram. for17                                        non-stationary masking (default: {1.3}).18        temp_coeff_nonstationary {float} -- Temperature coefficient for non-stationary masking (default: {0.1}).19        n_movemean_nonstationary {int} -- Number of samples for moving average smoothing in non-stationary masking20                                          (default: {20}).21        prop_decrease {float} -- Proportion to decrease signal by where the mask is zero (default: {1.0}).22        n_fft {int} -- Size of FFT for STFT (default: {1024}).23        win_length {[int]} -- Window length for STFT. If None, defaults to `n_fft` (default: {None}).24        hop_length {[int]} -- Hop length for STFT. If None, defaults to `win_length` // 4 (default: {None}).25        freq_mask_smooth_hz {float} -- Frequency smoothing width for mask (in Hz). If None, no smoothing is applied26                                     (default: {500}).27        time_mask_smooth_ms {float} -- Time smoothing width for mask (in ms). If None, no smoothing is applied28                                     (default: {50}).29    """30 31    @torch.no_grad()32    def __init__(33        self,34        sr: int,35        nonstationary: bool = False,36        n_std_thresh_stationary: float = 1.5,37        n_thresh_nonstationary: float = 1.3,38        temp_coeff_nonstationary: float = 0.1,39        n_movemean_nonstationary: int = 20,40        prop_decrease: float = 1.0,41        n_fft: int = 1024,42        win_length: bool = None,43        hop_length: int = None,44        freq_mask_smooth_hz: float = 500,45        time_mask_smooth_ms: float = 50,46    ):47        super().__init__()48 49        # General Params50        self.sr = sr51        self.nonstationary = nonstationary52        assert 0.0 <= prop_decrease <= 1.053        self.prop_decrease = prop_decrease54 55        # STFT Params56        self.n_fft = n_fft57        self.win_length = self.n_fft if win_length is None else win_length58        self.hop_length = self.win_length // 4 if hop_length is None else hop_length59 60        # Stationary Params61        self.n_std_thresh_stationary = n_std_thresh_stationary62 63        # Non-Stationary Params64        self.temp_coeff_nonstationary = temp_coeff_nonstationary65        self.n_movemean_nonstationary = n_movemean_nonstationary66        self.n_thresh_nonstationary = n_thresh_nonstationary67 68        # Smooth Mask Params69        self.freq_mask_smooth_hz = freq_mask_smooth_hz70        self.time_mask_smooth_ms = time_mask_smooth_ms71        self.register_buffer("smoothing_filter", self._generate_mask_smoothing_filter())72 73    @torch.no_grad()74    def _generate_mask_smoothing_filter(self) -> Union[torch.Tensor, None]:75        """76        A PyTorch module that applies a spectral gate to an input signal using the STFT.77 78        Returns:79            smoothing_filter (torch.Tensor): a 2D tensor representing the smoothing filter,80            with shape (n_grad_freq, n_grad_time), where n_grad_freq is the number of frequency81            bins to smooth and n_grad_time is the number of time frames to smooth.82            If both self.freq_mask_smooth_hz and self.time_mask_smooth_ms are None, returns None.83        """84        if self.freq_mask_smooth_hz is None and self.time_mask_smooth_ms is None:85            return None86 87        n_grad_freq = (88            189            if self.freq_mask_smooth_hz is None90            else int(self.freq_mask_smooth_hz / (self.sr / (self.n_fft / 2)))91        )92        if n_grad_freq < 1:93            raise ValueError(94                f"freq_mask_smooth_hz needs to be at least {int((self.sr / (self._n_fft / 2)))} Hz"95            )96 97        n_grad_time = (98            199            if self.time_mask_smooth_ms is None100            else int(self.time_mask_smooth_ms / ((self.hop_length / self.sr) * 1000))101        )102        if n_grad_time < 1:103            raise ValueError(104                f"time_mask_smooth_ms needs to be at least {int((self.hop_length / self.sr) * 1000)} ms"105            )106 107        if n_grad_time == 1 and n_grad_freq == 1:108            return None109 110        v_f = torch.cat(111            [112                linspace(0, 1, n_grad_freq + 1, endpoint=False),113                linspace(1, 0, n_grad_freq + 2),114            ]115        )[1:-1]116        v_t = torch.cat(117            [118                linspace(0, 1, n_grad_time + 1, endpoint=False),119                linspace(1, 0, n_grad_time + 2),120            ]121        )[1:-1]122        smoothing_filter = torch.outer(v_f, v_t).unsqueeze(0).unsqueeze(0)123 124        return smoothing_filter / smoothing_filter.sum()125 126    @torch.no_grad()127    def _stationary_mask(128        self, X_db: torch.Tensor, xn: Optional[torch.Tensor] = None129    ) -> torch.Tensor:130        """131        Computes a stationary binary mask to filter out noise in a log-magnitude spectrogram.132 133        Arguments:134            X_db (torch.Tensor): 2D tensor of shape (frames, freq_bins) containing the log-magnitude spectrogram.135            xn (torch.Tensor): 1D tensor containing the audio signal corresponding to X_db.136 137        Returns:138            sig_mask (torch.Tensor): Binary mask of the same shape as X_db, where values greater than the threshold139            are set to 1, and the rest are set to 0.140        """141        if xn is not None:142            XN = torch.stft(143                xn,144                n_fft=self.n_fft,145                hop_length=self.hop_length,146                win_length=self.win_length,147                return_complex=True,148                pad_mode="constant",149                center=True,150                window=torch.hann_window(self.win_length).to(xn.device),151            )152 153            XN_db = amp_to_db(XN).to(dtype=X_db.dtype)154        else:155            XN_db = X_db156 157        # calculate mean and standard deviation along the frequency axis158        std_freq_noise, mean_freq_noise = torch.std_mean(XN_db, dim=-1)159 160        # compute noise threshold161        noise_thresh = mean_freq_noise + std_freq_noise * self.n_std_thresh_stationary162 163        # create binary mask by thresholding the spectrogram164        sig_mask = X_db > noise_thresh.unsqueeze(2)165        return sig_mask166 167    @torch.no_grad()168    def _nonstationary_mask(self, X_abs: torch.Tensor) -> torch.Tensor:169        """170        Computes a non-stationary binary mask to filter out noise in a log-magnitude spectrogram.171 172        Arguments:173            X_abs (torch.Tensor): 2D tensor of shape (frames, freq_bins) containing the magnitude spectrogram.174 175        Returns:176            sig_mask (torch.Tensor): Binary mask of the same shape as X_abs, where values greater than the threshold177            are set to 1, and the rest are set to 0.178        """179        X_smoothed = (180            conv1d(181                X_abs.reshape(-1, 1, X_abs.shape[-1]),182                torch.ones(183                    self.n_movemean_nonstationary,184                    dtype=X_abs.dtype,185                    device=X_abs.device,186                ).view(1, 1, -1),187                padding="same",188            ).view(X_abs.shape)189            / self.n_movemean_nonstationary190        )191 192        # Compute slowness ratio and apply temperature sigmoid193        slowness_ratio = (X_abs - X_smoothed) / (X_smoothed + 1e-6)194        sig_mask = temperature_sigmoid(195            slowness_ratio, self.n_thresh_nonstationary, self.temp_coeff_nonstationary196        )197 198        return sig_mask199 200    def forward(201        self, x: torch.Tensor, xn: Optional[torch.Tensor] = None202    ) -> torch.Tensor:203        """204        Apply the proposed algorithm to the input signal.205 206        Arguments:207            x (torch.Tensor): The input audio signal, with shape (batch_size, signal_length).208            xn (Optional[torch.Tensor]): The noise signal used for stationary noise reduction. If `None`, the input209                                         signal is used as the noise signal. Default: `None`.210 211        Returns:212            torch.Tensor: The denoised audio signal, with the same shape as the input signal.213        """214        assert x.ndim == 2215        if x.shape[-1] < self.win_length * 2:216            raise Exception(f"x must be bigger than {self.win_length * 2}")217 218        assert xn is None or xn.ndim == 1 or xn.ndim == 2219        if xn is not None and xn.shape[-1] < self.win_length * 2:220            raise Exception(f"xn must be bigger than {self.win_length * 2}")221 222        # Compute short-time Fourier transform (STFT)223        X = torch.stft(224            x,225            n_fft=self.n_fft,226            hop_length=self.hop_length,227            win_length=self.win_length,228            return_complex=True,229            pad_mode="constant",230            center=True,231            window=torch.hann_window(self.win_length).to(x.device),232        )233 234        # Compute signal mask based on stationary or nonstationary assumptions235        if self.nonstationary:236            sig_mask = self._nonstationary_mask(X.abs())237        else:238            sig_mask = self._stationary_mask(amp_to_db(X), xn)239 240        # Propagate decrease in signal power241        sig_mask = self.prop_decrease * (sig_mask * 1.0 - 1.0) + 1.0242 243        # Smooth signal mask with 2D convolution244        if self.smoothing_filter is not None:245            sig_mask = conv2d(246                sig_mask.unsqueeze(1),247                self.smoothing_filter.to(sig_mask.dtype),248                padding="same",249            )250 251        # Apply signal mask to STFT magnitude and phase components252        Y = X * sig_mask.squeeze(1)253 254        # Inverse STFT to obtain time-domain signal255        y = torch.istft(256            Y,257            n_fft=self.n_fft,258            hop_length=self.hop_length,259            win_length=self.win_length,260            center=True,261            window=torch.hann_window(self.win_length).to(Y.device),262        )263 264        return y.to(dtype=x.dtype)265