ok

style_bert_vits2/models/hyper_parameters.py

@@ -27,6 +27,7 @@ class HyperParametersTrain(BaseModel):
     warmup_epochs: int = 0
     c_mel: int = 45
     c_kl: float = 1.0
+    c_matcha: float = 1.0
     c_commit: int = 100
     skip_optimizer: bool = False
     freeze_ZH_bert: bool = False
@@ -76,6 +77,12 @@ class HyperParametersModel(BaseModel):
     use_mel_posterior_encoder: bool = False
     use_duration_discriminator: bool = False
     use_wavlm_discriminator: bool = True
+    use_matcha: bool = False
+    matcha_channels: int = 192
+    matcha_num_heads: int = 2
+    matcha_dropout: float = 0.05
+    matcha_sigma_min: float = 0.0001
+    matcha_n_timesteps: int = 8
     inter_channels: int = 192
     hidden_channels: int = 192
     filter_channels: int = 768

style_bert_vits2/models/infer.py

@@ -35,6 +35,12 @@ def get_net_g(
             use_mel_posterior_encoder=hps.model.use_mel_posterior_encoder,
             use_duration_discriminator=hps.model.use_duration_discriminator,
             use_wavlm_discriminator=hps.model.use_wavlm_discriminator,
+            use_matcha=hps.model.use_matcha,
+            matcha_channels=hps.model.matcha_channels,
+            matcha_num_heads=hps.model.matcha_num_heads,
+            matcha_dropout=hps.model.matcha_dropout,
+            matcha_sigma_min=hps.model.matcha_sigma_min,
+            matcha_n_timesteps=hps.model.matcha_n_timesteps,
             inter_channels=hps.model.inter_channels,
             hidden_channels=hps.model.hidden_channels,
             filter_channels=hps.model.filter_channels,
@@ -66,6 +72,12 @@ def get_net_g(
             use_mel_posterior_encoder=hps.model.use_mel_posterior_encoder,
             use_duration_discriminator=hps.model.use_duration_discriminator,
             use_wavlm_discriminator=hps.model.use_wavlm_discriminator,
+            use_matcha=hps.model.use_matcha,
+            matcha_channels=hps.model.matcha_channels,
+            matcha_num_heads=hps.model.matcha_num_heads,
+            matcha_dropout=hps.model.matcha_dropout,
+            matcha_sigma_min=hps.model.matcha_sigma_min,
+            matcha_n_timesteps=hps.model.matcha_n_timesteps,
             inter_channels=hps.model.inter_channels,
             hidden_channels=hps.model.hidden_channels,
             filter_channels=hps.model.filter_channels,

style_bert_vits2/models/matcha_flow.py

@@ -0,0 +1,198 @@
+"""Matcha-TTS inspired conditional flow matching for VITS latent features.
+
+The flow-matching formulation and the masked 1-D U-Net layout are adapted from
+Matcha-TTS (MIT License).  This implementation only depends on PyTorch and is
+conditioned on Style-Bert-VITS2's aligned text prior and speaker embedding.
+"""
+
+import math
+from typing import Optional
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+class SinusoidalTimeEmbedding(nn.Module):
+    def __init__(self, channels: int) -> None:
+        super().__init__()
+        if channels % 2:
+            raise ValueError("Time embedding channels must be even")
+        self.channels = channels
+
+    def forward(self, time: torch.Tensor) -> torch.Tensor:
+        if time.ndim == 0:
+            time = time.unsqueeze(0)
+        half = self.channels // 2
+        scale = -math.log(10000) / max(half - 1, 1)
+        frequencies = torch.exp(
+            torch.arange(half, device=time.device, dtype=torch.float32) * scale
+        ).to(time.dtype)
+        angles = 1000 * time[:, None] * frequencies[None]
+        return torch.cat((angles.sin(), angles.cos()), dim=-1)
+
+
+class MaskedResBlock(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int, time_channels: int) -> None:
+        super().__init__()
+        groups = math.gcd(out_channels, 8)
+        self.conv1 = nn.Conv1d(in_channels, out_channels, 3, padding=1)
+        self.norm1 = nn.GroupNorm(groups, out_channels)
+        self.conv2 = nn.Conv1d(out_channels, out_channels, 3, padding=1)
+        self.norm2 = nn.GroupNorm(groups, out_channels)
+        self.time_proj = nn.Linear(time_channels, out_channels)
+        self.residual = (
+            nn.Identity()
+            if in_channels == out_channels
+            else nn.Conv1d(in_channels, out_channels, 1)
+        )
+
+    def forward(
+        self, x: torch.Tensor, mask: torch.Tensor, time: torch.Tensor
+    ) -> torch.Tensor:
+        residual = self.residual(x * mask)
+        x = F.mish(self.norm1(self.conv1(x * mask)))
+        x = x + self.time_proj(F.mish(time)).unsqueeze(-1)
+        x = self.norm2(self.conv2(x * mask))
+        return F.mish(x + residual) * mask
+
+
+class MaskedTransformerBlock(nn.Module):
+    def __init__(self, channels: int, num_heads: int, dropout: float) -> None:
+        super().__init__()
+        if channels % num_heads:
+            raise ValueError("Matcha channels must be divisible by attention heads")
+        self.norm1 = nn.LayerNorm(channels)
+        self.attention = nn.MultiheadAttention(
+            channels, num_heads, dropout=dropout, batch_first=True
+        )
+        self.norm2 = nn.LayerNorm(channels)
+        self.feed_forward = nn.Sequential(
+            nn.Linear(channels, channels * 4),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(channels * 4, channels),
+        )
+
+    def forward(self, x: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+        x = x.transpose(1, 2)
+        valid = mask[:, 0].bool()
+        hidden = self.norm1(x)
+        hidden, _ = self.attention(
+            hidden, hidden, hidden, key_padding_mask=~valid, need_weights=False
+        )
+        x = x + hidden
+        x = x + self.feed_forward(self.norm2(x))
+        return x.transpose(1, 2) * mask
+
+
+class MatchaEstimator(nn.Module):
+    """A compact masked 1-D U-Net velocity estimator."""
+
+    def __init__(
+        self,
+        latent_channels: int,
+        speaker_channels: int,
+        channels: int,
+        num_heads: int,
+        dropout: float,
+    ) -> None:
+        super().__init__()
+        time_channels = channels * 4
+        input_channels = latent_channels * 2 + speaker_channels
+        self.time_embedding = SinusoidalTimeEmbedding(latent_channels)
+        self.time_mlp = nn.Sequential(
+            nn.Linear(latent_channels, time_channels),
+            nn.SiLU(),
+            nn.Linear(time_channels, time_channels),
+        )
+        self.down_block = MaskedResBlock(input_channels, channels, time_channels)
+        self.down_attention = MaskedTransformerBlock(channels, num_heads, dropout)
+        self.downsample = nn.Conv1d(channels, channels, 3, stride=2, padding=1)
+        self.mid_block = MaskedResBlock(channels, channels, time_channels)
+        self.mid_attention = MaskedTransformerBlock(channels, num_heads, dropout)
+        self.up_block = MaskedResBlock(channels * 2, channels, time_channels)
+        self.up_attention = MaskedTransformerBlock(channels, num_heads, dropout)
+        self.output = nn.Conv1d(channels, latent_channels, 1)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        mask: torch.Tensor,
+        mu: torch.Tensor,
+        time: torch.Tensor,
+        speaker: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if time.ndim == 0:
+            time = time.expand(x.shape[0])
+        time_embedding = self.time_mlp(self.time_embedding(time))
+        inputs = [x, mu]
+        if speaker is not None:
+            inputs.append(speaker.expand(-1, -1, x.shape[-1]))
+        x = torch.cat(inputs, dim=1)
+
+        skip = self.down_attention(self.down_block(x, mask, time_embedding), mask)
+        down_mask = F.interpolate(mask, size=(skip.shape[-1] + 1) // 2, mode="nearest")
+        x = self.downsample(skip * mask)
+        x = self.mid_attention(self.mid_block(x, down_mask, time_embedding), down_mask)
+        x = F.interpolate(x, size=skip.shape[-1], mode="nearest")
+        x = torch.cat((x, skip), dim=1)
+        x = self.up_attention(self.up_block(x, mask, time_embedding), mask)
+        return self.output(x * mask) * mask
+
+
+class MatchaFlow(nn.Module):
+    """Conditional flow matching decoder operating in VITS prior space."""
+
+    def __init__(
+        self,
+        latent_channels: int,
+        speaker_channels: int,
+        channels: int = 192,
+        num_heads: int = 2,
+        dropout: float = 0.05,
+        sigma_min: float = 1e-4,
+    ) -> None:
+        super().__init__()
+        self.sigma_min = sigma_min
+        self.estimator = MatchaEstimator(
+            latent_channels, speaker_channels, channels, num_heads, dropout
+        )
+
+    def compute_loss(
+        self,
+        target: torch.Tensor,
+        mask: torch.Tensor,
+        mu: torch.Tensor,
+        speaker: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        batch = target.shape[0]
+        time = torch.rand(batch, 1, 1, device=target.device, dtype=target.dtype)
+        noise = torch.randn_like(target)
+        sample = (1 - (1 - self.sigma_min) * time) * noise + time * target
+        velocity = target - (1 - self.sigma_min) * noise
+        prediction = self.estimator(
+            sample, mask, mu, time.flatten(), speaker=speaker
+        )
+        error = (prediction - velocity).square() * mask
+        return error.sum() / (mask.sum().clamp_min(1) * target.shape[1])
+
+    @torch.inference_mode()
+    def forward(
+        self,
+        mu: torch.Tensor,
+        mask: torch.Tensor,
+        n_timesteps: int,
+        temperature: float = 1.0,
+        speaker: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if n_timesteps < 1:
+            raise ValueError("Matcha inference requires at least one timestep")
+        x = torch.randn_like(mu) * temperature
+        dt = 1.0 / n_timesteps
+        for step in range(n_timesteps):
+            time = torch.full(
+                (mu.shape[0],), step * dt, device=mu.device, dtype=mu.dtype
+            )
+            x = x + dt * self.estimator(x, mask, mu, time, speaker=speaker)
+        return x * mask

style_bert_vits2/models/models.py

@@ -8,6 +8,7 @@ from torch.nn import functional as F
 from torch.nn.utils import remove_weight_norm, spectral_norm, weight_norm
 
 from style_bert_vits2.models import attentions, commons, modules, monotonic_alignment
+from style_bert_vits2.models.matcha_flow import MatchaFlow
 from style_bert_vits2.nlp.symbols import NUM_LANGUAGES, NUM_TONES, SYMBOLS
 
 
@@ -880,6 +881,8 @@ class SynthesizerTrn(nn.Module):
             "use_spk_conditioned_encoder", True
         )
         self.use_sdp = use_sdp
+        self.use_matcha = kwargs.get("use_matcha", False)
+        self.matcha_n_timesteps = kwargs.get("matcha_n_timesteps", 8)
         self.use_noise_scaled_mas = kwargs.get("use_noise_scaled_mas", False)
         self.mas_noise_scale_initial = kwargs.get("mas_noise_scale_initial", 0.01)
         self.noise_scale_delta = kwargs.get("noise_scale_delta", 2e-6)
@@ -950,6 +953,15 @@ class SynthesizerTrn(nn.Module):
             self.emb_g = nn.Embedding(n_speakers, gin_channels)
         else:
             self.ref_enc = ReferenceEncoder(spec_channels, gin_channels)
+        if self.use_matcha:
+            self.matcha = MatchaFlow(
+                latent_channels=inter_channels,
+                speaker_channels=gin_channels,
+                channels=kwargs.get("matcha_channels", hidden_channels),
+                num_heads=kwargs.get("matcha_num_heads", n_heads),
+                dropout=kwargs.get("matcha_dropout", 0.05),
+                sigma_min=kwargs.get("matcha_sigma_min", 1e-4),
+            )
 
     def forward(
         self,
@@ -971,6 +983,7 @@ class SynthesizerTrn(nn.Module):
         torch.Tensor,
         torch.Tensor,
         torch.Tensor,
+        torch.Tensor,
         tuple[torch.Tensor, ...],
         tuple[torch.Tensor, ...],
     ]:
@@ -1037,10 +1050,20 @@ class SynthesizerTrn(nn.Module):
         z_slice, ids_slice = commons.rand_slice_segments(
             z, y_lengths, self.segment_size
         )
+        if self.use_matcha:
+            matcha_target = commons.slice_segments(z_p, ids_slice, self.segment_size)
+            matcha_mu = commons.slice_segments(m_p, ids_slice, self.segment_size)
+            matcha_mask = commons.slice_segments(y_mask, ids_slice, self.segment_size)
+            matcha_loss = self.matcha.compute_loss(
+                matcha_target.detach(), matcha_mask, matcha_mu, speaker=g
+            )
+        else:
+            matcha_loss = z.new_zeros(())
         o = self.dec(z_slice, g=g)
         return (
             o,
             l_length,
+            matcha_loss,
             attn,
             ids_slice,
             x_mask,
@@ -1096,7 +1119,16 @@ class SynthesizerTrn(nn.Module):
             1, 2
         )  # [b, t', t], [b, t, d] -> [b, d, t']
 
-        z_p = m_p + torch.randn_like(m_p) * torch.exp(logs_p) * noise_scale
+        if self.use_matcha:
+            z_p = self.matcha(
+                m_p,
+                y_mask,
+                n_timesteps=self.matcha_n_timesteps,
+                temperature=noise_scale,
+                speaker=g,
+            )
+        else:
+            z_p = m_p + torch.randn_like(m_p) * torch.exp(logs_p) * noise_scale
         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)

style_bert_vits2/models/models_jp_extra.py

@@ -8,6 +8,7 @@ from torch.nn import functional as F
 from torch.nn.utils import remove_weight_norm, spectral_norm, weight_norm
 
 from style_bert_vits2.models import attentions, commons, modules, monotonic_alignment
+from style_bert_vits2.models.matcha_flow import MatchaFlow
 from style_bert_vits2.nlp.symbols import NUM_LANGUAGES, NUM_TONES, SYMBOLS
 
 
@@ -938,6 +939,8 @@ class SynthesizerTrn(nn.Module):
             "use_spk_conditioned_encoder", True
         )
         self.use_sdp = use_sdp
+        self.use_matcha = kwargs.get("use_matcha", False)
+        self.matcha_n_timesteps = kwargs.get("matcha_n_timesteps", 8)
         self.use_noise_scaled_mas = kwargs.get("use_noise_scaled_mas", False)
         self.mas_noise_scale_initial = kwargs.get("mas_noise_scale_initial", 0.01)
         self.noise_scale_delta = kwargs.get("noise_scale_delta", 2e-6)
@@ -1007,6 +1010,15 @@ class SynthesizerTrn(nn.Module):
             self.emb_g = nn.Embedding(n_speakers, gin_channels)
         else:
             self.ref_enc = ReferenceEncoder(spec_channels, gin_channels)
+        if self.use_matcha:
+            self.matcha = MatchaFlow(
+                latent_channels=inter_channels,
+                speaker_channels=gin_channels,
+                channels=kwargs.get("matcha_channels", hidden_channels),
+                num_heads=kwargs.get("matcha_num_heads", n_heads),
+                dropout=kwargs.get("matcha_dropout", 0.05),
+                sigma_min=kwargs.get("matcha_sigma_min", 1e-4),
+            )
 
     def forward(
         self,
@@ -1029,6 +1041,7 @@ class SynthesizerTrn(nn.Module):
         torch.Tensor,
         tuple[torch.Tensor, ...],
         tuple[torch.Tensor, ...],
+        torch.Tensor,
     ]:
         if self.n_speakers > 0:
             g = self.emb_g(sid).unsqueeze(-1)  # [b, h, 1]
@@ -1093,10 +1106,20 @@ class SynthesizerTrn(nn.Module):
         z_slice, ids_slice = commons.rand_slice_segments(
             z, y_lengths, self.segment_size
         )
+        if self.use_matcha:
+            matcha_target = commons.slice_segments(z_p, ids_slice, self.segment_size)
+            matcha_mu = commons.slice_segments(m_p, ids_slice, self.segment_size)
+            matcha_mask = commons.slice_segments(y_mask, ids_slice, self.segment_size)
+            matcha_loss = self.matcha.compute_loss(
+                matcha_target.detach(), matcha_mask, matcha_mu, speaker=g
+            )
+        else:
+            matcha_loss = z.new_zeros(())
         o = self.dec(z_slice, g=g)
         return (
             o,
             l_length,
+            matcha_loss,
             attn,
             ids_slice,
             x_mask,
@@ -1151,7 +1174,16 @@ class SynthesizerTrn(nn.Module):
             1, 2
         )  # [b, t', t], [b, t, d] -> [b, d, t']
 
-        z_p = m_p + torch.randn_like(m_p) * torch.exp(logs_p) * noise_scale
+        if self.use_matcha:
+            z_p = self.matcha(
+                m_p,
+                y_mask,
+                n_timesteps=self.matcha_n_timesteps,
+                temperature=noise_scale,
+                speaker=g,
+            )
+        else:
+            z_p = m_p + torch.randn_like(m_p) * torch.exp(logs_p) * noise_scale
         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)