ok

.gitignore

@@ -1,6 +1,7 @@
 __pycache__/
 venv/
 .venv/
+.uv-python/
 dist/
 .coverage*
 .ipynb_checkpoints/

THIRD_PARTY_NOTICES.md

@@ -0,0 +1,9 @@
+# Third-party notices
+
+## Matcha-TTS
+
+`style_bert_vits2/models/matcha_flow.py` adapts the conditional flow matching
+formulation and masked 1-D U-Net design from Matcha-TTS.
+
+Copyright (c) 2023 Shivam Mehta. Matcha-TTS is distributed under the MIT
+License; see `licenses/Matcha-TTS.LICENSE`.

configs/config.json

@@ -16,6 +16,7 @@
     "warmup_epochs": 0,
     "c_mel": 45,
     "c_kl": 1.0,
+    "c_matcha": 1.0,
     "skip_optimizer": false,
     "freeze_ZH_bert": false,
     "freeze_JP_bert": false,
@@ -48,6 +49,12 @@
     "use_noise_scaled_mas": true,
     "use_mel_posterior_encoder": false,
     "use_duration_discriminator": true,
+    "use_matcha": true,
+    "matcha_channels": 192,
+    "matcha_num_heads": 2,
+    "matcha_dropout": 0.05,
+    "matcha_sigma_min": 0.0001,
+    "matcha_n_timesteps": 8,
     "inter_channels": 192,
     "hidden_channels": 192,
     "filter_channels": 768,

configs/config_jp_extra.json

@@ -17,6 +17,7 @@
     "warmup_epochs": 0,
     "c_mel": 45,
     "c_kl": 1.0,
+    "c_matcha": 1.0,
     "c_commit": 100,
     "skip_optimizer": false,
     "freeze_ZH_bert": false,
@@ -48,6 +49,12 @@
     "use_mel_posterior_encoder": false,
     "use_duration_discriminator": false,
     "use_wavlm_discriminator": true,
+    "use_matcha": true,
+    "matcha_channels": 192,
+    "matcha_num_heads": 2,
+    "matcha_dropout": 0.05,
+    "matcha_sigma_min": 0.0001,
+    "matcha_n_timesteps": 8,
     "inter_channels": 192,
     "hidden_channels": 192,
     "filter_channels": 768,

docs/MATCHA_FLOW.md

@@ -0,0 +1,24 @@
+# Matcha flow integration
+
+Style-Bert-VITS2 can optionally generate its aligned prior latent with a
+Matcha-TTS-inspired conditional flow matching (CFM) decoder. The existing
+BERT/style encoder, duration predictors, VITS normalizing flow, and waveform
+generator are retained.
+
+Set `model.use_matcha` to `true` for a newly trained or fine-tuned model. The
+template configurations enable it. Models whose configuration does not contain
+this key continue to use the original Gaussian prior sampling path.
+
+Relevant settings:
+
+- `train.c_matcha`: weight of the CFM velocity loss.
+- `model.matcha_channels`: U-Net hidden width.
+- `model.matcha_num_heads`: attention head count.
+- `model.matcha_dropout`: transformer dropout.
+- `model.matcha_sigma_min`: minimum flow path noise.
+- `model.matcha_n_timesteps`: Euler steps at inference; more steps trade speed
+  for refinement quality.
+
+The Matcha branch has new parameters, so enabling it on an existing checkpoint
+requires fine-tuning before inference. It is trained on the same random latent
+segments used by the waveform generator to keep memory use bounded.

licenses/Matcha-TTS.LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 Shivam Mehta
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

pyproject.toml

@@ -63,6 +63,8 @@ only-include = [
     "tests",
     "LGPL_LICENSE",
     "LICENSE",
+    "THIRD_PARTY_NOTICES.md",
+    "licenses",
     "pyproject.toml",
     "README.md",
 ]
@@ -153,4 +155,4 @@ exclude_lines = ["no cov", "if __name__ == .__main__.:", "if TYPE_CHECKING:"]
 extend-select = ["I"]
 
 [tool.ruff.lint.isort]
-lines-after-imports = 2
+lines-after-imports = 2

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)

tests/test_matcha_flow.py

@@ -0,0 +1,31 @@
+import torch
+
+from style_bert_vits2.models.matcha_flow import MatchaFlow
+
+
+def test_matcha_flow_loss_and_sampling_with_odd_length() -> None:
+    model = MatchaFlow(
+        latent_channels=8,
+        speaker_channels=4,
+        channels=16,
+        num_heads=2,
+        dropout=0.0,
+    )
+    target = torch.randn(2, 8, 7)
+    mu = torch.randn_like(target)
+    mask = torch.tensor(
+        [[[1, 1, 1, 1, 1, 1, 1]], [[1, 1, 1, 1, 1, 0, 0]]],
+        dtype=target.dtype,
+    )
+    speaker = torch.randn(2, 4, 1)
+
+    loss = model.compute_loss(target, mask, mu, speaker)
+    loss.backward()
+
+    assert loss.ndim == 0
+    assert torch.isfinite(loss)
+    assert any(parameter.grad is not None for parameter in model.parameters())
+
+    sample = model(mu, mask, n_timesteps=2, speaker=speaker)
+    assert sample.shape == target.shape
+    assert torch.count_nonzero(sample[1, :, 5:]) == 0

train_ms.py

@@ -316,6 +316,12 @@ def run():
         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,
@@ -658,6 +664,7 @@ def train_and_evaluate(
             (
                 y_hat,
                 l_length,
+                matcha_loss,
                 attn,
                 ids_slice,
                 x_mask,
@@ -745,10 +752,18 @@ def train_and_evaluate(
                 loss_dur = torch.sum(l_length.float())
                 loss_mel = F.l1_loss(y_mel, y_hat_mel) * hps.train.c_mel
                 loss_kl = kl_loss(z_p, logs_q, m_p, logs_p, z_mask) * hps.train.c_kl
+                loss_matcha = matcha_loss.float() * hps.train.c_matcha
 
                 loss_fm = feature_loss(fmap_r, fmap_g)
                 loss_gen, losses_gen = generator_loss(y_d_hat_g)
-                loss_gen_all = loss_gen + loss_fm + loss_mel + loss_dur + loss_kl
+                loss_gen_all = (
+                    loss_gen
+                    + loss_fm
+                    + loss_mel
+                    + loss_dur
+                    + loss_kl
+                    + loss_matcha
+                )
                 if net_dur_disc is not None:
                     loss_dur_gen, losses_dur_gen = generator_loss(y_dur_hat_g)
                     loss_gen_all += loss_dur_gen
@@ -764,7 +779,15 @@ def train_and_evaluate(
         if rank == 0:
             if global_step % hps.train.log_interval == 0 and not hps.speedup:
                 lr = optim_g.param_groups[0]["lr"]
-                losses = [loss_disc, loss_gen, loss_fm, loss_mel, loss_dur, loss_kl]
+                losses = [
+                    loss_disc,
+                    loss_gen,
+                    loss_fm,
+                    loss_mel,
+                    loss_dur,
+                    loss_kl,
+                    loss_matcha,
+                ]
                 # logger.info(
                 #     "Train Epoch: {} [{:.0f}%]".format(
                 #         epoch, 100.0 * batch_idx / len(train_loader)
@@ -774,6 +797,7 @@ def train_and_evaluate(
 
                 scalar_dict = {
                     "loss/g/total": loss_gen_all,
+                    "loss/g/matcha": loss_matcha,
                     "loss/d/total": loss_disc_all,
                     "learning_rate": lr,
                     "grad_norm_d": grad_norm_d,

train_ms_jp_extra.py

@@ -329,6 +329,12 @@ def run():
         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,
@@ -744,6 +750,7 @@ def train_and_evaluate(
             (
                 y_hat,
                 l_length,
+                matcha_loss,
                 attn,
                 ids_slice,
                 x_mask,
@@ -856,12 +863,20 @@ def train_and_evaluate(
                 loss_dur = torch.sum(l_length.float())
                 loss_mel = F.l1_loss(y_mel, y_hat_mel) * hps.train.c_mel
                 loss_kl = kl_loss(z_p, logs_q, m_p, logs_p, z_mask) * hps.train.c_kl
+                loss_matcha = matcha_loss.float() * hps.train.c_matcha
 
                 loss_fm = feature_loss(fmap_r, fmap_g)
                 loss_gen, losses_gen = generator_loss(y_d_hat_g)
                 # loss_commit = loss_commit * hps.train.c_commit
 
-                loss_gen_all = loss_gen + loss_fm + loss_mel + loss_dur + loss_kl
+                loss_gen_all = (
+                    loss_gen
+                    + loss_fm
+                    + loss_mel
+                    + loss_dur
+                    + loss_kl
+                    + loss_matcha
+                )
                 if net_dur_disc is not None:
                     loss_dur_gen, losses_dur_gen = generator_loss(y_dur_hat_g)
                     if net_wd is not None:
@@ -880,7 +895,15 @@ def train_and_evaluate(
         if rank == 0:
             if global_step % hps.train.log_interval == 0 and not hps.speedup:
                 lr = optim_g.param_groups[0]["lr"]
-                losses = [loss_disc, loss_gen, loss_fm, loss_mel, loss_dur, loss_kl]
+                losses = [
+                    loss_disc,
+                    loss_gen,
+                    loss_fm,
+                    loss_mel,
+                    loss_dur,
+                    loss_kl,
+                    loss_matcha,
+                ]
                 # logger.info(
                 #     "Train Epoch: {} [{:.0f}%]".format(
                 #         epoch, 100.0 * batch_idx / len(train_loader)
@@ -890,6 +913,7 @@ def train_and_evaluate(
 
                 scalar_dict = {
                     "loss/g/total": loss_gen_all,
+                    "loss/g/matcha": loss_matcha,
                     "loss/d/total": loss_disc_all,
                     "learning_rate": lr,
                     "grad_norm_d": grad_norm_d,