199 lines
7.3 KiB
Python
199 lines
7.3 KiB
Python
"""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
|