tiktokman/DiffRhythm2

1# Implementation adapted from https://github.com/EdwardDixon/snake under the MIT license.2#   LICENSE is in incl_licenses directory.3 4import torch5from torch import nn, sin, pow6from torch.nn import Parameter7 8 9class Snake(nn.Module):10    """11    Implementation of a sine-based periodic activation function12    Shape:13        - Input: (B, C, T)14        - Output: (B, C, T), same shape as the input15    Parameters:16        - alpha - trainable parameter17    References:18        - This activation function is from this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda:19        https://arxiv.org/abs/2006.0819520    Examples:21        >>> a1 = snake(256)22        >>> x = torch.randn(256)23        >>> x = a1(x)24    """25 26    def __init__(27        self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False28    ):29        """30        Initialization.31        INPUT:32            - in_features: shape of the input33            - alpha: trainable parameter34            alpha is initialized to 1 by default, higher values = higher-frequency.35            alpha will be trained along with the rest of your model.36        """37        super(Snake, self).__init__()38        self.in_features = in_features39 40        # Initialize alpha41        self.alpha_logscale = alpha_logscale42        if self.alpha_logscale:  # Log scale alphas initialized to zeros43            self.alpha = Parameter(torch.zeros(in_features) * alpha)44        else:  # Linear scale alphas initialized to ones45            self.alpha = Parameter(torch.ones(in_features) * alpha)46 47        self.alpha.requires_grad = alpha_trainable48 49        self.no_div_by_zero = 0.00000000150 51    def forward(self, x):52        """53        Forward pass of the function.54        Applies the function to the input elementwise.55        Snake ∶= x + 1/a * sin^2 (xa)56        """57        alpha = self.alpha.unsqueeze(0).unsqueeze(-1)  # Line up with x to [B, C, T]58        if self.alpha_logscale:59            alpha = torch.exp(alpha)60        x = x + (1.0 / (alpha + self.no_div_by_zero)) * pow(sin(x * alpha), 2)61 62        return x63 64 65class SnakeBeta(nn.Module):66    """67    A modified Snake function which uses separate parameters for the magnitude of the periodic components68    Shape:69        - Input: (B, C, T)70        - Output: (B, C, T), same shape as the input71    Parameters:72        - alpha - trainable parameter that controls frequency73        - beta - trainable parameter that controls magnitude74    References:75        - This activation function is a modified version based on this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda:76        https://arxiv.org/abs/2006.0819577    Examples:78        >>> a1 = snakebeta(256)79        >>> x = torch.randn(256)80        >>> x = a1(x)81    """82 83    def __init__(84        self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False85    ):86        """87        Initialization.88        INPUT:89            - in_features: shape of the input90            - alpha - trainable parameter that controls frequency91            - beta - trainable parameter that controls magnitude92            alpha is initialized to 1 by default, higher values = higher-frequency.93            beta is initialized to 1 by default, higher values = higher-magnitude.94            alpha will be trained along with the rest of your model.95        """96        super(SnakeBeta, self).__init__()97        self.in_features = in_features98 99        # Initialize alpha100        self.alpha_logscale = alpha_logscale101        if self.alpha_logscale:  # Log scale alphas initialized to zeros102            self.alpha = Parameter(torch.zeros(in_features) * alpha)103            self.beta = Parameter(torch.zeros(in_features) * alpha)104        else:  # Linear scale alphas initialized to ones105            self.alpha = Parameter(torch.ones(in_features) * alpha)106            self.beta = Parameter(torch.ones(in_features) * alpha)107 108        self.alpha.requires_grad = alpha_trainable109        self.beta.requires_grad = alpha_trainable110 111        self.no_div_by_zero = 0.000000001112 113    def forward(self, x):114        """115        Forward pass of the function.116        Applies the function to the input elementwise.117        SnakeBeta ∶= x + 1/b * sin^2 (xa)118        """119        alpha = self.alpha.unsqueeze(0).unsqueeze(-1)  # Line up with x to [B, C, T]120        beta = self.beta.unsqueeze(0).unsqueeze(-1)121        if self.alpha_logscale:122            alpha = torch.exp(alpha)123            beta = torch.exp(beta)124        x = x + (1.0 / (beta + self.no_div_by_zero)) * pow(sin(x * alpha), 2)125 126        return x127