fix

style_bert_vits2/models/matcha_flow.py

@@ -13,6 +13,21 @@ 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__()
@@ -58,13 +73,24 @@ class MaskedResBlock(nn.Module):
 
 
 class MaskedTransformerBlock(nn.Module):
-    def __init__(self, channels: int, num_heads: int, dropout: float) -> None:
+    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.attention = nn.MultiheadAttention(
-            channels, num_heads, dropout=dropout, batch_first=True
+        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(
@@ -78,14 +104,68 @@ class MaskedTransformerBlock(nn.Module):
         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
-        )
+        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."""
 
@@ -96,6 +176,7 @@ class MatchaEstimator(nn.Module):
         channels: int,
         num_heads: int,
         dropout: float,
+        use_diff_attention: bool,
     ) -> None:
         super().__init__()
         time_channels = channels * 4
@@ -107,12 +188,18 @@ class MatchaEstimator(nn.Module):
             nn.Linear(time_channels, time_channels),
         )
         self.down_block = MaskedResBlock(input_channels, channels, time_channels)
-        self.down_attention = MaskedTransformerBlock(channels, num_heads, dropout)
+        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)
+        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)
+        self.up_attention = MaskedTransformerBlock(
+            channels, num_heads, dropout, use_diff_attention
+        )
         self.output = nn.Conv1d(channels, latent_channels, 1)
 
     def forward(
@@ -152,11 +239,17 @@ class MatchaFlow(nn.Module):
         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
+            latent_channels,
+            speaker_channels,
+            channels,
+            num_heads,
+            dropout,
+            use_diff_attention,
         )
 
     def compute_loss(