diff --git a/.gitignore b/.gitignore index 8625c49..cf33b74 100644 --- a/.gitignore +++ b/.gitignore @@ -1,6 +1,7 @@ __pycache__/ venv/ .venv/ +.uv-python/ dist/ .coverage* .ipynb_checkpoints/ diff --git a/Matcha-TTS b/Matcha-TTS new file mode 160000 index 0000000..bd4d90d --- /dev/null +++ b/Matcha-TTS @@ -0,0 +1 @@ +Subproject commit bd4d90d93214b37f7a159cf205ae85762c2c10aa diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md new file mode 100644 index 0000000..f57d7e8 --- /dev/null +++ b/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`. diff --git a/configs/config.json b/configs/config.json index fb8b41e..755f857 100644 --- a/configs/config.json +++ b/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, diff --git a/configs/config_jp_extra.json b/configs/config_jp_extra.json index d28ecc5..fd2d4fe 100644 --- a/configs/config_jp_extra.json +++ b/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, diff --git a/docs/MATCHA_FLOW.md b/docs/MATCHA_FLOW.md new file mode 100644 index 0000000..e917986 --- /dev/null +++ b/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. diff --git a/licenses/Matcha-TTS.LICENSE b/licenses/Matcha-TTS.LICENSE new file mode 100644 index 0000000..858018e --- /dev/null +++ b/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. diff --git a/pyproject.toml b/pyproject.toml index 3b20d05..861ab85 100644 --- a/pyproject.toml +++ b/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 \ No newline at end of file +lines-after-imports = 2 diff --git a/style_bert_vits2/models/hyper_parameters.py b/style_bert_vits2/models/hyper_parameters.py index e27eb42..b2081f8 100644 --- a/style_bert_vits2/models/hyper_parameters.py +++ b/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 diff --git a/style_bert_vits2/models/infer.py b/style_bert_vits2/models/infer.py index 0921c5f..c718fff 100644 --- a/style_bert_vits2/models/infer.py +++ b/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, diff --git a/style_bert_vits2/models/matcha_flow.py b/style_bert_vits2/models/matcha_flow.py new file mode 100644 index 0000000..7729ac0 --- /dev/null +++ b/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 diff --git a/style_bert_vits2/models/models.py b/style_bert_vits2/models/models.py index eaff2fa..22ccd23 100644 --- a/style_bert_vits2/models/models.py +++ b/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) diff --git a/style_bert_vits2/models/models_jp_extra.py b/style_bert_vits2/models/models_jp_extra.py index e8df7d9..ed90d07 100644 --- a/style_bert_vits2/models/models_jp_extra.py +++ b/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) diff --git a/tests/test_matcha_flow.py b/tests/test_matcha_flow.py new file mode 100644 index 0000000..d4e735f --- /dev/null +++ b/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 diff --git a/train_ms.py b/train_ms.py index 7df71d1..8fc1dee 100644 --- a/train_ms.py +++ b/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, diff --git a/train_ms_jp_extra.py b/train_ms_jp_extra.py index 1d4893f..c73083d 100644 --- a/train_ms_jp_extra.py +++ b/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,