303 lines
11 KiB
Python
303 lines
11 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
|
|
|
|
|
|
def configure_matcha_only_training(model: nn.Module) -> int:
|
|
"""Freeze a synthesizer except for its Matcha branch.
|
|
|
|
Returns the number of trainable parameters after configuration.
|
|
"""
|
|
matcha = getattr(model, "matcha", None)
|
|
if not getattr(model, "use_matcha", False) or matcha is None:
|
|
raise ValueError("Matcha-only training requires an enabled Matcha branch")
|
|
for parameter in model.parameters():
|
|
parameter.requires_grad = False
|
|
for parameter in matcha.parameters():
|
|
parameter.requires_grad = True
|
|
return sum(parameter.numel() for parameter in matcha.parameters())
|
|
|
|
|
|
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,
|
|
use_diff_attention: bool,
|
|
) -> None:
|
|
super().__init__()
|
|
if channels % num_heads:
|
|
raise ValueError("Matcha channels must be divisible by attention heads")
|
|
self.norm1 = nn.LayerNorm(channels)
|
|
self.use_diff_attention = use_diff_attention
|
|
self.attention = (
|
|
DifferentialAttentionV2(channels, num_heads, dropout)
|
|
if use_diff_attention
|
|
else 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)
|
|
if self.use_diff_attention:
|
|
hidden = self.attention(hidden, valid)
|
|
else:
|
|
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 DifferentialAttentionV2(nn.Module):
|
|
"""Differential Attention V2 without FlashAttention or custom kernels.
|
|
|
|
Each logical head has two query heads that share one key/value head. Their
|
|
contexts are combined as ``context_1 - sigmoid(lambda) * context_2``, where
|
|
lambda is projected per token and logical head.
|
|
"""
|
|
|
|
def __init__(self, channels: int, num_heads: int, dropout: float) -> None:
|
|
super().__init__()
|
|
if channels % num_heads:
|
|
raise ValueError("Channels must be divisible by attention heads")
|
|
self.num_heads = num_heads
|
|
self.head_channels = channels // num_heads
|
|
self.scale = self.head_channels**-0.5
|
|
self.query_projection = nn.Linear(channels, channels * 2)
|
|
self.key_projection = nn.Linear(channels, channels)
|
|
self.value_projection = nn.Linear(channels, channels)
|
|
self.lambda_projection = nn.Linear(channels, num_heads)
|
|
self.output_projection = nn.Linear(channels, channels)
|
|
self.attention_dropout = nn.Dropout(dropout)
|
|
|
|
def forward(self, x: torch.Tensor, valid: torch.Tensor) -> torch.Tensor:
|
|
batch, length, channels = x.shape
|
|
queries = self.query_projection(x).view(
|
|
batch, length, self.num_heads, 2, self.head_channels
|
|
)
|
|
keys = self.key_projection(x).view(
|
|
batch, length, self.num_heads, self.head_channels
|
|
)
|
|
values = self.value_projection(x).view(
|
|
batch, length, self.num_heads, self.head_channels
|
|
)
|
|
|
|
queries = queries.permute(0, 2, 3, 1, 4)
|
|
keys = keys.permute(0, 2, 1, 3)
|
|
values = values.permute(0, 2, 1, 3)
|
|
scores = torch.matmul(queries, keys.unsqueeze(2).transpose(-1, -2))
|
|
scores = scores * self.scale
|
|
scores = scores.masked_fill(~valid[:, None, None, None, :], float("-inf"))
|
|
attention = self.attention_dropout(torch.softmax(scores, dim=-1))
|
|
contexts = torch.matmul(attention, values.unsqueeze(2))
|
|
|
|
lambda_value = torch.sigmoid(self.lambda_projection(x))
|
|
lambda_value = lambda_value.permute(0, 2, 1).unsqueeze(-1)
|
|
contexts = contexts[:, :, 0] - lambda_value * contexts[:, :, 1]
|
|
contexts = contexts.permute(0, 2, 1, 3).reshape(batch, length, channels)
|
|
contexts = self.output_projection(contexts)
|
|
return contexts * valid.unsqueeze(-1).to(contexts.dtype)
|
|
|
|
|
|
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,
|
|
use_diff_attention: bool,
|
|
) -> 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, use_diff_attention
|
|
)
|
|
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, use_diff_attention
|
|
)
|
|
self.up_block = MaskedResBlock(channels * 2, channels, time_channels)
|
|
self.up_attention = MaskedTransformerBlock(
|
|
channels, num_heads, dropout, use_diff_attention
|
|
)
|
|
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,
|
|
use_diff_attention: bool = False,
|
|
) -> None:
|
|
super().__init__()
|
|
self.sigma_min = sigma_min
|
|
self.estimator = MatchaEstimator(
|
|
latent_channels,
|
|
speaker_channels,
|
|
channels,
|
|
num_heads,
|
|
dropout,
|
|
use_diff_attention,
|
|
)
|
|
|
|
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])
|
|
|
|
def sample(
|
|
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
|
|
|
|
@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:
|
|
"""Sample without building a graph (normal inference path)."""
|
|
return self.sample(mu, mask, n_timesteps, temperature, speaker)
|