Upload synthesizers.py

synthesizers.py

@@ -0,0 +1,250 @@
+import torch
+from typing import Optional
+from rvc.lib.algorithm.generators.hifigan_mrf import HiFiGANMRFGenerator
+from rvc.lib.algorithm.generators.hifigan_nsf import HiFiGANNSFGenerator
+from rvc.lib.algorithm.generators.hifigan import HiFiGANGenerator
+from rvc.lib.algorithm.generators.refinegan import RefineGANGenerator
+from rvc.lib.algorithm.commons import slice_segments, rand_slice_segments
+from rvc.lib.algorithm.residuals import ResidualCouplingBlock
+from rvc.lib.algorithm.encoders import TextEncoder, PosteriorEncoder
+
+
+class Synthesizer(torch.nn.Module):
+    """
+    Base Synthesizer model.
+
+    Args:
+        spec_channels (int): Number of channels in the spectrogram.
+        segment_size (int): Size of the audio segment.
+        inter_channels (int): Number of channels in the intermediate layers.
+        hidden_channels (int): Number of channels in the hidden layers.
+        filter_channels (int): Number of channels in the filter layers.
+        n_heads (int): Number of attention heads.
+        n_layers (int): Number of layers in the encoder.
+        kernel_size (int): Size of the convolution kernel.
+        p_dropout (float): Dropout probability.
+        resblock (str): Type of residual block.
+        resblock_kernel_sizes (list): Kernel sizes for the residual blocks.
+        resblock_dilation_sizes (list): Dilation sizes for the residual blocks.
+        upsample_rates (list): Upsampling rates for the decoder.
+        upsample_initial_channel (int): Number of channels in the initial upsampling layer.
+        upsample_kernel_sizes (list): Kernel sizes for the upsampling layers.
+        spk_embed_dim (int): Dimension of the speaker embedding.
+        gin_channels (int): Number of channels in the global conditioning vector.
+        sr (int): Sampling rate of the audio.
+        use_f0 (bool): Whether to use F0 information.
+        text_enc_hidden_dim (int): Hidden dimension for the text encoder.
+        kwargs: Additional keyword arguments.
+    """
+
+    def __init__(
+        self,
+        spec_channels: int,
+        segment_size: int,
+        inter_channels: int,
+        hidden_channels: int,
+        filter_channels: int,
+        n_heads: int,
+        n_layers: int,
+        kernel_size: int,
+        p_dropout: float,
+        resblock: str,
+        resblock_kernel_sizes: list,
+        resblock_dilation_sizes: list,
+        upsample_rates: list,
+        upsample_initial_channel: int,
+        upsample_kernel_sizes: list,
+        spk_embed_dim: int,
+        gin_channels: int,
+        sr: int,
+        use_f0: bool,
+        text_enc_hidden_dim: int = 768,
+        vocoder: str = "RefineGAN",
+        randomized: bool = True,
+        checkpointing: bool = False,
+        **kwargs,
+    ):
+        super().__init__()
+        self.segment_size = segment_size
+        self.use_f0 = use_f0
+        self.randomized = randomized
+
+        self.enc_p = TextEncoder(
+            inter_channels,
+            hidden_channels,
+            filter_channels,
+            n_heads,
+            n_layers,
+            kernel_size,
+            p_dropout,
+            text_enc_hidden_dim,
+            f0=use_f0,
+        )
+        print(f"Using {vocoder} vocoder")
+        if use_f0:
+            if vocoder == "MRF HiFi-GAN":
+                self.dec = HiFiGANMRFGenerator(
+                    in_channel=inter_channels,
+                    upsample_initial_channel=upsample_initial_channel,
+                    upsample_rates=upsample_rates,
+                    upsample_kernel_sizes=upsample_kernel_sizes,
+                    resblock_kernel_sizes=resblock_kernel_sizes,
+                    resblock_dilations=resblock_dilation_sizes,
+                    gin_channels=gin_channels,
+                    sample_rate=sr,
+                    harmonic_num=8,
+                    checkpointing=checkpointing,
+                )
+            elif vocoder == "RefineGAN":
+                self.dec = RefineGANGenerator(
+                    sample_rate=sr,
+                    downsample_rates=upsample_rates[::-1],
+                    upsample_rates=upsample_rates,
+                    start_channels=16,
+                    num_mels=inter_channels,
+                    checkpointing=checkpointing,
+                )
+            else:
+                self.dec = HiFiGANNSFGenerator(
+                    inter_channels,
+                    resblock_kernel_sizes,
+                    resblock_dilation_sizes,
+                    upsample_rates,
+                    upsample_initial_channel,
+                    upsample_kernel_sizes,
+                    gin_channels=gin_channels,
+                    sr=sr,
+                    checkpointing=checkpointing,
+                )
+        else:
+            if vocoder == "MRF HiFi-GAN":
+                print("Using RefineGAN without pitch guidance (experimental).")
+                self.dec = None
+            elif vocoder == "RefineGAN":
+                print("RefineGAN does not support training without pitch guidance.")
+                self.dec = RefineGANGenerator(
+                   sample_rate=sr,
+                   downsample_rates=upsample_rates[::-1],
+                   upsample_rates=upsample_rates,
+                   start_channels=16,
+                   num_mels=inter_channels,
+                   checkpointing=checkpointing,
+                )
+            else:
+                self.dec = HiFiGANGenerator(
+                    inter_channels,
+                    resblock_kernel_sizes,
+                    resblock_dilation_sizes,
+                    upsample_rates,
+                    upsample_initial_channel,
+                    upsample_kernel_sizes,
+                    gin_channels=gin_channels,
+                )
+        self.enc_q = PosteriorEncoder(
+            spec_channels,
+            inter_channels,
+            hidden_channels,
+            5,
+            1,
+            16,
+            gin_channels=gin_channels,
+        )
+        self.flow = ResidualCouplingBlock(
+            inter_channels,
+            hidden_channels,
+            5,
+            1,
+            3,
+            gin_channels=gin_channels,
+        )
+        self.emb_g = torch.nn.Embedding(spk_embed_dim, gin_channels)
+
+    def _remove_weight_norm_from(self, module):
+        for hook in module._forward_pre_hooks.values():
+            if getattr(hook, "__class__", None).__name__ == "WeightNorm":
+                torch.nn.utils.remove_weight_norm(module)
+
+    def remove_weight_norm(self):
+        for module in [self.dec, self.flow, self.enc_q]:
+            self._remove_weight_norm_from(module)
+
+    def __prepare_scriptable__(self):
+        self.remove_weight_norm()
+        return self
+
+    def forward(
+        self,
+        phone: torch.Tensor,
+        phone_lengths: torch.Tensor,
+        pitch: Optional[torch.Tensor] = None,
+        pitchf: Optional[torch.Tensor] = None,
+        y: Optional[torch.Tensor] = None,
+        y_lengths: Optional[torch.Tensor] = None,
+        ds: Optional[torch.Tensor] = None,
+    ):
+        g = self.emb_g(ds).unsqueeze(-1)
+        m_p, logs_p, x_mask = self.enc_p(phone, pitch, phone_lengths)
+
+        if y is not None:
+            z, m_q, logs_q, y_mask = self.enc_q(y, y_lengths, g=g)
+            z_p = self.flow(z, y_mask, g=g)
+            # regular old training method using random slices
+            if self.randomized:
+                z_slice, ids_slice = rand_slice_segments(
+                    z, y_lengths, self.segment_size
+                )
+                if self.use_f0:
+                    pitchf = slice_segments(pitchf, ids_slice, self.segment_size, 2)
+                    o = self.dec(z_slice, pitchf, g=g)
+                else:
+                    o = self.dec(z_slice, g=g)
+                return o, ids_slice, x_mask, y_mask, (z, z_p, m_p, logs_p, m_q, logs_q)
+            # future use for finetuning using the entire dataset each pass
+            else:
+                if self.use_f0:
+                    o = self.dec(z, pitchf, g=g)
+                else:
+                    o = self.dec(z, g=g)
+                return o, None, x_mask, y_mask, (z, z_p, m_p, logs_p, m_q, logs_q)
+        else:
+            return None, None, x_mask, None, (None, None, m_p, logs_p, None, None)
+
+    @torch.jit.export
+    def infer(
+        self,
+        phone: torch.Tensor,
+        phone_lengths: torch.Tensor,
+        pitch: Optional[torch.Tensor] = None,
+        nsff0: Optional[torch.Tensor] = None,
+        sid: torch.Tensor = None,
+        rate: Optional[torch.Tensor] = None,
+    ):
+        """
+        Inference of the model.
+
+        Args:
+            phone (torch.Tensor): Phoneme sequence.
+            phone_lengths (torch.Tensor): Lengths of the phoneme sequences.
+            pitch (torch.Tensor, optional): Pitch sequence.
+            nsff0 (torch.Tensor, optional): Fine-grained pitch sequence.
+            sid (torch.Tensor): Speaker embedding.
+            rate (torch.Tensor, optional): Rate for time-stretching.
+        """
+        g = self.emb_g(sid).unsqueeze(-1)
+        m_p, logs_p, x_mask = self.enc_p(phone, pitch, phone_lengths)
+        z_p = (m_p + torch.exp(logs_p) * torch.randn_like(m_p) * 0.66666) * x_mask
+
+        if rate is not None:
+            head = int(z_p.shape[2] * (1.0 - rate.item()))
+            z_p, x_mask = z_p[:, :, head:], x_mask[:, :, head:]
+            if self.use_f0 and nsff0 is not None:
+                nsff0 = nsff0[:, head:]
+
+        z = self.flow(z_p, x_mask, g=g, reverse=True)
+        o = (
+            self.dec(z * x_mask, nsff0, g=g)
+            if self.use_f0
+            else self.dec(z * x_mask, g=g)
+        )
+
+        return o, x_mask, (z, z_p, m_p, logs_p)