ok

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