split

style_bert_vits2/models/models.py

@@ -1133,3 +1133,89 @@ class SynthesizerTrn(nn.Module):
         z = self.flow(z_p, y_mask, g=g, reverse=True)
         o = self.dec((z * y_mask)[:, :, :max_len], g=g)
         return o, attn, y_mask, (z, z_p, m_p, logs_p)
+
+
+def build_mel_synthesizer(
+    n_vocab: int,
+    n_mel_channels: int,
+    hidden_channels: int,
+    filter_channels: int,
+    n_heads: int,
+    n_layers: int,
+    kernel_size: int,
+    p_dropout: float,
+    n_speakers: int,
+    gin_channels: int,
+    matcha_channels: int = 192,
+    matcha_num_heads: int = 2,
+    matcha_dropout: float = 0.05,
+    matcha_sigma_min: float = 1e-4,
+    matcha_n_timesteps: int = 8,
+    matcha_use_diff_attention: bool = False,
+) -> "MelAcousticModel":
+    """Build the multilingual text-to-mel model without a waveform generator."""
+    from style_bert_vits2.models.mel_synthesizer import MelAcousticModel
+
+    encoder = TextEncoder(
+        n_vocab,
+        n_mel_channels,
+        hidden_channels,
+        filter_channels,
+        n_heads,
+        n_layers,
+        kernel_size,
+        p_dropout,
+        n_speakers,
+        gin_channels=gin_channels,
+    )
+    sdp = StochasticDurationPredictor(
+        hidden_channels, 192, 3, 0.5, 4, gin_channels=gin_channels
+    )
+    dp = DurationPredictor(
+        hidden_channels, 256, 3, 0.5, gin_channels=gin_channels
+    )
+    return MelAcousticModel(
+        encoder,
+        sdp,
+        dp,
+        n_mel_channels,
+        n_speakers,
+        gin_channels,
+        hidden_channels,
+        matcha_channels,
+        matcha_num_heads,
+        matcha_dropout,
+        matcha_sigma_min,
+        matcha_n_timesteps,
+        matcha_use_diff_attention,
+    )
+
+
+def build_mel_vocoder(
+    n_mel_channels: int,
+    resblock: str,
+    resblock_kernel_sizes: list[int],
+    resblock_dilation_sizes: list[list[int]],
+    upsample_rates: list[int],
+    upsample_initial_channel: int,
+    upsample_kernel_sizes: list[int],
+    n_speakers: int,
+    gin_channels: int,
+) -> "MelVocoder":
+    """Build a standalone HiFi-GAN-style mel-to-wave Generator."""
+    from style_bert_vits2.models.mel_synthesizer import MelVocoder
+
+    return MelVocoder(
+        Generator(
+            n_mel_channels,
+            resblock,
+            resblock_kernel_sizes,
+            resblock_dilation_sizes,
+            upsample_rates,
+            upsample_initial_channel,
+            upsample_kernel_sizes,
+            gin_channels=gin_channels,
+        ),
+        n_speakers=n_speakers,
+        gin_channels=gin_channels,
+    )