Update modules/quantize.py

modules/quantize.py

@@ -1,229 +1,229 @@
-from dac.nn.quantize import ResidualVectorQuantize
-from torch import nn
-from modules.wavenet import WN
-import torch
-import torchaudio
-import torchaudio.functional as audio_F
-import numpy as np
-from .alias_free_torch import *
-from torch.nn.utils import weight_norm
-from torch import nn, sin, pow
-from einops.layers.torch import Rearrange
-from dac.model.encodec import SConv1d
-
-def init_weights(m):
-    if isinstance(m, nn.Conv1d):
-        nn.init.trunc_normal_(m.weight, std=0.02)
-        nn.init.constant_(m.bias, 0)
-
-
-def WNConv1d(*args, **kwargs):
-    return weight_norm(nn.Conv1d(*args, **kwargs))
-
-
-def WNConvTranspose1d(*args, **kwargs):
-    return weight_norm(nn.ConvTranspose1d(*args, **kwargs))
-
-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, alpha_trainable=True, alpha_logscale=False
-    ):
-        """
-        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(SnakeBeta, self).__init__()
-        self.in_features = in_features
-
-        # initialize alpha
-        self.alpha_logscale = alpha_logscale
-        if self.alpha_logscale:  # log scale alphas initialized to zeros
-            self.alpha = nn.Parameter(torch.zeros(in_features) * alpha)
-            self.beta = nn.Parameter(torch.zeros(in_features) * alpha)
-        else:  # linear scale alphas initialized to ones
-            self.alpha = nn.Parameter(torch.ones(in_features) * alpha)
-            self.beta = nn.Parameter(torch.ones(in_features) * alpha)
-
-        self.alpha.requires_grad = alpha_trainable
-        self.beta.requires_grad = alpha_trainable
-
-        self.no_div_by_zero = 0.000000001
-
-    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.alpha.unsqueeze(0).unsqueeze(-1)  # line up with x to [B, C, T]
-        beta = self.beta.unsqueeze(0).unsqueeze(-1)
-        if self.alpha_logscale:
-            alpha = torch.exp(alpha)
-            beta = torch.exp(beta)
-        x = x + (1.0 / (beta + self.no_div_by_zero)) * pow(sin(x * alpha), 2)
-
-        return x
-
-class ResidualUnit(nn.Module):
-    def __init__(self, dim: int = 16, dilation: int = 1):
-        super().__init__()
-        pad = ((7 - 1) * dilation) // 2
-        self.block = nn.Sequential(
-            Activation1d(activation=SnakeBeta(dim, alpha_logscale=True)),
-            WNConv1d(dim, dim, kernel_size=7, dilation=dilation, padding=pad),
-            Activation1d(activation=SnakeBeta(dim, alpha_logscale=True)),
-            WNConv1d(dim, dim, kernel_size=1),
-        )
-
-    def forward(self, x):
-        return x + self.block(x)
-
-class CNNLSTM(nn.Module):
-    def __init__(self, indim, outdim, head, global_pred=False):
-        super().__init__()
-        self.global_pred = global_pred
-        self.model = nn.Sequential(
-            ResidualUnit(indim, dilation=1),
-            ResidualUnit(indim, dilation=2),
-            ResidualUnit(indim, dilation=3),
-            Activation1d(activation=SnakeBeta(indim, alpha_logscale=True)),
-            Rearrange("b c t -> b t c"),
-        )
-        self.heads = nn.ModuleList([nn.Linear(indim, outdim) for i in range(head)])
-
-    def forward(self, x):
-        # x: [B, C, T]
-        x = self.model(x)
-        if self.global_pred:
-            x = torch.mean(x, dim=1, keepdim=False)
-        outs = [head(x) for head in self.heads]
-        return outs
-
-def sequence_mask(length, max_length=None):
-  if max_length is None:
-    max_length = length.max()
-  x = torch.arange(max_length, dtype=length.dtype, device=length.device)
-  return x.unsqueeze(0) < length.unsqueeze(1)
-class FAquantizer(nn.Module):
-    def __init__(self, in_dim=1024,
-                 n_p_codebooks=1,
-                 n_c_codebooks=2,
-                 n_t_codebooks=2,
-                 n_r_codebooks=3,
-                 codebook_size=1024,
-                 codebook_dim=8,
-                 quantizer_dropout=0.5,
-                 causal=False,
-                 separate_prosody_encoder=False,
-                 timbre_norm=False,):
-        super(FAquantizer, self).__init__()
-        conv1d_type = SConv1d# if causal else nn.Conv1d
-        self.prosody_quantizer = ResidualVectorQuantize(
-            input_dim=in_dim,
-            n_codebooks=n_p_codebooks,
-            codebook_size=codebook_size,
-            codebook_dim=codebook_dim,
-            quantizer_dropout=quantizer_dropout,
-        )
-
-        self.content_quantizer = ResidualVectorQuantize(
-            input_dim=in_dim,
-            n_codebooks=n_c_codebooks,
-            codebook_size=codebook_size,
-            codebook_dim=codebook_dim,
-            quantizer_dropout=quantizer_dropout,
-        )
-
-        self.residual_quantizer = ResidualVectorQuantize(
-            input_dim=in_dim,
-            n_codebooks=n_r_codebooks,
-            codebook_size=codebook_size,
-            codebook_dim=codebook_dim,
-            quantizer_dropout=quantizer_dropout,
-        )
-
-        self.melspec_linear = conv1d_type(in_channels=20, out_channels=256, kernel_size=1, causal=causal)
-        self.melspec_encoder = WN(hidden_channels=256, kernel_size=5, dilation_rate=1, n_layers=8, gin_channels=0, p_dropout=0.2, causal=causal)
-        self.melspec_linear2 = conv1d_type(in_channels=256, out_channels=1024, kernel_size=1, causal=causal)
-
-        self.prob_random_mask_residual = 0.75
-
-        SPECT_PARAMS = {
-            "n_fft": 2048,
-            "win_length": 1200,
-            "hop_length": 300,
-        }
-        MEL_PARAMS = {
-            "n_mels": 80,
-        }
-
-        self.to_mel = torchaudio.transforms.MelSpectrogram(
-            n_mels=MEL_PARAMS["n_mels"], sample_rate=24000, **SPECT_PARAMS
-        )
-        self.mel_mean, self.mel_std = -4, 4
-        self.frame_rate = 24000 / 300
-        self.hop_length = 300
-
-    def preprocess(self, wave_tensor, n_bins=20):
-        mel_tensor = self.to_mel(wave_tensor.squeeze(1))
-        mel_tensor = (torch.log(1e-5 + mel_tensor) - self.mel_mean) / self.mel_std
-        return mel_tensor[:, :n_bins, :int(wave_tensor.size(-1) / self.hop_length)]
-
-    def forward(self, x, wave_segments):
-        outs = 0
-        prosody_feature = self.preprocess(wave_segments)
-
-        f0_input = prosody_feature  # (B, T, 20)
-        f0_input = self.melspec_linear(f0_input)
-        f0_input = self.melspec_encoder(f0_input, torch.ones(f0_input.shape[0], 1, f0_input.shape[2]).to(
-            f0_input.device).bool())
-        f0_input = self.melspec_linear2(f0_input)
-
-        common_min_size = min(f0_input.size(2), x.size(2))
-        f0_input = f0_input[:, :, :common_min_size]
-
-        x = x[:, :, :common_min_size]
-
-        z_p, codes_p, latents_p, commitment_loss_p, codebook_loss_p = self.prosody_quantizer(
-            f0_input, 1
-        )
-        outs += z_p.detach()
-
-        z_c, codes_c, latents_c, commitment_loss_c, codebook_loss_c = self.content_quantizer(
-            x, 2
-        )
-        outs += z_c.detach()
-
-        residual_feature = x - z_p.detach() - z_c.detach()
-
-        z_r, codes_r, latents_r, commitment_loss_r, codebook_loss_r = self.residual_quantizer(
-            residual_feature, 3
-        )
-
-        quantized = [z_p, z_c, z_r]
-        codes = [codes_p, codes_c, codes_r]
-
+from dac.nn.quantize import ResidualVectorQuantize
+from torch import nn
+from modules.wavenet import WN
+import torch
+import torchaudio
+import torchaudio.functional as audio_F
+import numpy as np
+from .bigvgan import *
+from torch.nn.utils import weight_norm
+from torch import nn, sin, pow
+from einops.layers.torch import Rearrange
+from dac.model.encodec import SConv1d
+
+def init_weights(m):
+    if isinstance(m, nn.Conv1d):
+        nn.init.trunc_normal_(m.weight, std=0.02)
+        nn.init.constant_(m.bias, 0)
+
+
+def WNConv1d(*args, **kwargs):
+    return weight_norm(nn.Conv1d(*args, **kwargs))
+
+
+def WNConvTranspose1d(*args, **kwargs):
+    return weight_norm(nn.ConvTranspose1d(*args, **kwargs))
+
+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, alpha_trainable=True, alpha_logscale=False
+    ):
+        """
+        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(SnakeBeta, self).__init__()
+        self.in_features = in_features
+
+        # initialize alpha
+        self.alpha_logscale = alpha_logscale
+        if self.alpha_logscale:  # log scale alphas initialized to zeros
+            self.alpha = nn.Parameter(torch.zeros(in_features) * alpha)
+            self.beta = nn.Parameter(torch.zeros(in_features) * alpha)
+        else:  # linear scale alphas initialized to ones
+            self.alpha = nn.Parameter(torch.ones(in_features) * alpha)
+            self.beta = nn.Parameter(torch.ones(in_features) * alpha)
+
+        self.alpha.requires_grad = alpha_trainable
+        self.beta.requires_grad = alpha_trainable
+
+        self.no_div_by_zero = 0.000000001
+
+    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.alpha.unsqueeze(0).unsqueeze(-1)  # line up with x to [B, C, T]
+        beta = self.beta.unsqueeze(0).unsqueeze(-1)
+        if self.alpha_logscale:
+            alpha = torch.exp(alpha)
+            beta = torch.exp(beta)
+        x = x + (1.0 / (beta + self.no_div_by_zero)) * pow(sin(x * alpha), 2)
+
+        return x
+
+class ResidualUnit(nn.Module):
+    def __init__(self, dim: int = 16, dilation: int = 1):
+        super().__init__()
+        pad = ((7 - 1) * dilation) // 2
+        self.block = nn.Sequential(
+            Activation1d(activation=SnakeBeta(dim, alpha_logscale=True)),
+            WNConv1d(dim, dim, kernel_size=7, dilation=dilation, padding=pad),
+            Activation1d(activation=SnakeBeta(dim, alpha_logscale=True)),
+            WNConv1d(dim, dim, kernel_size=1),
+        )
+
+    def forward(self, x):
+        return x + self.block(x)
+
+class CNNLSTM(nn.Module):
+    def __init__(self, indim, outdim, head, global_pred=False):
+        super().__init__()
+        self.global_pred = global_pred
+        self.model = nn.Sequential(
+            ResidualUnit(indim, dilation=1),
+            ResidualUnit(indim, dilation=2),
+            ResidualUnit(indim, dilation=3),
+            Activation1d(activation=SnakeBeta(indim, alpha_logscale=True)),
+            Rearrange("b c t -> b t c"),
+        )
+        self.heads = nn.ModuleList([nn.Linear(indim, outdim) for i in range(head)])
+
+    def forward(self, x):
+        # x: [B, C, T]
+        x = self.model(x)
+        if self.global_pred:
+            x = torch.mean(x, dim=1, keepdim=False)
+        outs = [head(x) for head in self.heads]
+        return outs
+
+def sequence_mask(length, max_length=None):
+  if max_length is None:
+    max_length = length.max()
+  x = torch.arange(max_length, dtype=length.dtype, device=length.device)
+  return x.unsqueeze(0) < length.unsqueeze(1)
+class FAquantizer(nn.Module):
+    def __init__(self, in_dim=1024,
+                 n_p_codebooks=1,
+                 n_c_codebooks=2,
+                 n_t_codebooks=2,
+                 n_r_codebooks=3,
+                 codebook_size=1024,
+                 codebook_dim=8,
+                 quantizer_dropout=0.5,
+                 causal=False,
+                 separate_prosody_encoder=False,
+                 timbre_norm=False,):
+        super(FAquantizer, self).__init__()
+        conv1d_type = SConv1d# if causal else nn.Conv1d
+        self.prosody_quantizer = ResidualVectorQuantize(
+            input_dim=in_dim,
+            n_codebooks=n_p_codebooks,
+            codebook_size=codebook_size,
+            codebook_dim=codebook_dim,
+            quantizer_dropout=quantizer_dropout,
+        )
+
+        self.content_quantizer = ResidualVectorQuantize(
+            input_dim=in_dim,
+            n_codebooks=n_c_codebooks,
+            codebook_size=codebook_size,
+            codebook_dim=codebook_dim,
+            quantizer_dropout=quantizer_dropout,
+        )
+
+        self.residual_quantizer = ResidualVectorQuantize(
+            input_dim=in_dim,
+            n_codebooks=n_r_codebooks,
+            codebook_size=codebook_size,
+            codebook_dim=codebook_dim,
+            quantizer_dropout=quantizer_dropout,
+        )
+
+        self.melspec_linear = conv1d_type(in_channels=20, out_channels=256, kernel_size=1, causal=causal)
+        self.melspec_encoder = WN(hidden_channels=256, kernel_size=5, dilation_rate=1, n_layers=8, gin_channels=0, p_dropout=0.2, causal=causal)
+        self.melspec_linear2 = conv1d_type(in_channels=256, out_channels=1024, kernel_size=1, causal=causal)
+
+        self.prob_random_mask_residual = 0.75
+
+        SPECT_PARAMS = {
+            "n_fft": 2048,
+            "win_length": 1200,
+            "hop_length": 300,
+        }
+        MEL_PARAMS = {
+            "n_mels": 80,
+        }
+
+        self.to_mel = torchaudio.transforms.MelSpectrogram(
+            n_mels=MEL_PARAMS["n_mels"], sample_rate=24000, **SPECT_PARAMS
+        )
+        self.mel_mean, self.mel_std = -4, 4
+        self.frame_rate = 24000 / 300
+        self.hop_length = 300
+
+    def preprocess(self, wave_tensor, n_bins=20):
+        mel_tensor = self.to_mel(wave_tensor.squeeze(1))
+        mel_tensor = (torch.log(1e-5 + mel_tensor) - self.mel_mean) / self.mel_std
+        return mel_tensor[:, :n_bins, :int(wave_tensor.size(-1) / self.hop_length)]
+
+    def forward(self, x, wave_segments):
+        outs = 0
+        prosody_feature = self.preprocess(wave_segments)
+
+        f0_input = prosody_feature  # (B, T, 20)
+        f0_input = self.melspec_linear(f0_input)
+        f0_input = self.melspec_encoder(f0_input, torch.ones(f0_input.shape[0], 1, f0_input.shape[2]).to(
+            f0_input.device).bool())
+        f0_input = self.melspec_linear2(f0_input)
+
+        common_min_size = min(f0_input.size(2), x.size(2))
+        f0_input = f0_input[:, :, :common_min_size]
+
+        x = x[:, :, :common_min_size]
+
+        z_p, codes_p, latents_p, commitment_loss_p, codebook_loss_p = self.prosody_quantizer(
+            f0_input, 1
+        )
+        outs += z_p.detach()
+
+        z_c, codes_c, latents_c, commitment_loss_c, codebook_loss_c = self.content_quantizer(
+            x, 2
+        )
+        outs += z_c.detach()
+
+        residual_feature = x - z_p.detach() - z_c.detach()
+
+        z_r, codes_r, latents_r, commitment_loss_r, codebook_loss_r = self.residual_quantizer(
+            residual_feature, 3
+        )
+
+        quantized = [z_p, z_c, z_r]
+        codes = [codes_p, codes_c, codes_r]
+
         return quantized, codes