Spaces:
Running
Running
| # Refer from https://github.com/NVIDIA/BigVGAN | |
| import math | |
| import torch | |
| import torch.nn as nn | |
| from torch import nn | |
| from torch.nn.utils.parametrizations import weight_norm | |
| from .alias_free_torch import DownSample1d, UpSample1d | |
| class SnakeBeta(nn.Module): | |
| """ | |
| A modified Snake function which uses separate parameters for the magnitude of the periodic components | |
| Shape: | |
| - Input: (B, C, T) | |
| - Output: (B, C, T), same shape as the input | |
| Parameters: | |
| - alpha - trainable parameter that controls frequency | |
| - beta - trainable parameter that controls magnitude | |
| References: | |
| - This activation function is a modified version based on this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda: | |
| https://arxiv.org/abs/2006.08195 | |
| Examples: | |
| >>> a1 = snakebeta(256) | |
| >>> x = torch.randn(256) | |
| >>> x = a1(x) | |
| """ | |
| def __init__(self, in_features, alpha=1.0, clamp=(1e-2, 50)): | |
| """ | |
| Initialization. | |
| INPUT: | |
| - in_features: shape of the input | |
| - alpha - trainable parameter that controls frequency | |
| - beta - trainable parameter that controls magnitude | |
| alpha is initialized to 1 by default, higher values = higher-frequency. | |
| beta is initialized to 1 by default, higher values = higher-magnitude. | |
| alpha will be trained along with the rest of your model. | |
| """ | |
| super().__init__() | |
| self.in_features = in_features | |
| self.log_alpha = nn.Parameter(torch.zeros(in_features) + math.log(alpha)) | |
| self.log_beta = nn.Parameter(torch.zeros(in_features) + math.log(alpha)) | |
| self.clamp = clamp | |
| def forward(self, x): | |
| """ | |
| Forward pass of the function. | |
| Applies the function to the input elementwise. | |
| SnakeBeta ∶= x + 1/b * sin^2 (xa) | |
| """ | |
| alpha = self.log_alpha.exp().clamp(*self.clamp) | |
| alpha = alpha[None, :, None] | |
| beta = self.log_beta.exp().clamp(*self.clamp) | |
| beta = beta[None, :, None] | |
| x = x + (1.0 / beta) * (x * alpha).sin().pow(2) | |
| return x | |
| class UpActDown(nn.Module): | |
| def __init__( | |
| self, | |
| act, | |
| up_ratio: int = 2, | |
| down_ratio: int = 2, | |
| up_kernel_size: int = 12, | |
| down_kernel_size: int = 12, | |
| ): | |
| super().__init__() | |
| self.up_ratio = up_ratio | |
| self.down_ratio = down_ratio | |
| self.act = act | |
| self.upsample = UpSample1d(up_ratio, up_kernel_size) | |
| self.downsample = DownSample1d(down_ratio, down_kernel_size) | |
| def forward(self, x): | |
| # x: [B,C,T] | |
| x = self.upsample(x) | |
| x = self.act(x) | |
| x = self.downsample(x) | |
| return x | |
| class AMPBlock(nn.Sequential): | |
| def __init__(self, channels, *, kernel_size=3, dilations=(1, 3, 5)): | |
| super().__init__(*(self._make_layer(channels, kernel_size, d) for d in dilations)) | |
| def _make_layer(self, channels, kernel_size, dilation): | |
| return nn.Sequential( | |
| weight_norm(nn.Conv1d(channels, channels, kernel_size, dilation=dilation, padding="same")), | |
| UpActDown(act=SnakeBeta(channels)), | |
| weight_norm(nn.Conv1d(channels, channels, kernel_size, padding="same")), | |
| ) | |
| def forward(self, x): | |
| return x + super().forward(x) | |