fix

README.md

@@ -1,5 +1,47 @@
 # Style-Bert-VITS2
 
+## 分離型 mel acoustic model / Vocoder 学習
+
+`train_ms_mel.py` は、従来の一体型 `SynthesizerTrn` とは別に、次の2段階の
+学習を提供します。
+
+1. 音素・BERT・style・speaker 条件から log-mel spectrogram を生成する
+   acoustic model の単体学習
+2. 学習済み acoustic model と学習済み Generator (Vocoder) を接続した
+   end-to-end fine-tuning
+
+Acoustic model の単体学習:
+
+```bash
+python train_ms_mel.py \
+  -c Data/your_model/config.json \
+  --stage acoustic
+```
+
+Vocoder の単体学習（ground-truth mel を入力）:
+
+```bash
+python train_ms_mel.py \
+  -c Data/your_model/config.json \
+  --stage vocoder
+```
+
+一貫学習（両方のチェックポイント指定が必須）:
+
+```bash
+python train_ms_mel.py \
+  -c Data/your_model/config.json \
+  --stage joint \
+  --acoustic-checkpoint Data/your_model/models/ACOUSTIC_10000.pth \
+  --vocoder-checkpoint Data/your_model/models/VOCODER_10000.pth
+```
+
+`--vocoder-checkpoint` は standalone Vocoder のほか、従来の `G_*.pth` に
+含まれる形状互換な `dec.*` も初期値として読み込めます。joint 学習時の CFM サンプリングは
+acoustic model まで勾配を通し、メモリ使用量を抑えるため
+`train.segment_size / data.hop_length` フレームだけを Vocoder に渡します。
+品質とメモリの調整には `--joint-timesteps` を使用します。
+
 **利用の際は必ず[お願いとデフォルトモデルの利用規約](/docs/TERMS_OF_USE.md)をお読みください。**
 
 Bert-VITS2 with more controllable voice styles.

style_bert_vits2/models/infer.py

@@ -224,14 +224,23 @@ def infer(
         en_bert = en_bert[:, :-2]
 
     with torch.no_grad():
+        # BERT may run in fp16 (for example on Colab) while a loaded TTS model
+        # remains fp32. Conv1d/Linear require inputs and parameters to use the
+        # same dtype outside autocast, so normalize all floating conditioning
+        # tensors to the TTS model dtype at this boundary.
+        model_dtype = next(net_g.parameters()).dtype
         x_tst = phones.to(device).unsqueeze(0)
         tones = tones.to(device).unsqueeze(0)
         lang_ids = lang_ids.to(device).unsqueeze(0)
-        bert = bert.to(device).unsqueeze(0)
+        bert = bert.to(device=device, dtype=model_dtype).unsqueeze(0)
-        ja_bert = ja_bert.to(device).unsqueeze(0)
+        ja_bert = ja_bert.to(device=device, dtype=model_dtype).unsqueeze(0)
-        en_bert = en_bert.to(device).unsqueeze(0)
+        en_bert = en_bert.to(device=device, dtype=model_dtype).unsqueeze(0)
         x_tst_lengths = torch.LongTensor([phones.size(0)]).to(device)
-        style_vec_tensor = torch.from_numpy(style_vec).to(device).unsqueeze(0)
+        style_vec_tensor = (
+            torch.from_numpy(style_vec)
+            .to(device=device, dtype=model_dtype)
+            .unsqueeze(0)
+        )
         del phones
         sid_tensor = torch.LongTensor([sid]).to(device)
 

style_bert_vits2/models/matcha_flow.py

@@ -270,8 +270,7 @@ class MatchaFlow(nn.Module):
         error = (prediction - velocity).square() * mask
         return error.sum() / (mask.sum().clamp_min(1) * target.shape[1])
 
-    @torch.inference_mode()
+    def sample(
-    def forward(
         self,
         mu: torch.Tensor,
         mask: torch.Tensor,
@@ -289,3 +288,15 @@ class MatchaFlow(nn.Module):
             )
             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)

style_bert_vits2/models/mel_synthesizer.py

@@ -0,0 +1,284 @@
+"""Separated mel acoustic model and mel-to-wave vocoder composition.
+
+Unlike :class:`SynthesizerTrn`, the acoustic model in this module ends at a
+log-mel spectrogram.  The vocoder is therefore independently pretrainable and
+replaceable, while :class:`JointMelSynthesizer` keeps the boundary
+differentiable for end-to-end fine-tuning.
+"""
+
+import math
+from typing import Any, Optional
+
+import torch
+from torch import nn
+
+from style_bert_vits2.models import commons
+from style_bert_vits2.models.matcha_flow import MatchaFlow
+
+
+class MelAcousticModel(nn.Module):
+    """Text-to-mel model shared by the standard and JP-Extra front ends."""
+
+    def __init__(
+        self,
+        text_encoder: nn.Module,
+        stochastic_duration_predictor: nn.Module,
+        duration_predictor: nn.Module,
+        n_mel_channels: int,
+        n_speakers: int,
+        gin_channels: int,
+        hidden_channels: int,
+        matcha_channels: int,
+        matcha_num_heads: int,
+        matcha_dropout: float,
+        matcha_sigma_min: float,
+        matcha_n_timesteps: int,
+        matcha_use_diff_attention: bool,
+    ) -> None:
+        super().__init__()
+        self.enc_p = text_encoder
+        self.sdp = stochastic_duration_predictor
+        self.dp = duration_predictor
+        self.n_mel_channels = n_mel_channels
+        self.n_speakers = n_speakers
+        self.gin_channels = gin_channels
+        self.matcha_n_timesteps = matcha_n_timesteps
+        if n_speakers < 1:
+            raise ValueError("Separated mel models require at least one speaker")
+        self.emb_g = nn.Embedding(n_speakers, gin_channels)
+        self.matcha = MatchaFlow(
+            latent_channels=n_mel_channels,
+            speaker_channels=gin_channels,
+            channels=matcha_channels or hidden_channels,
+            num_heads=matcha_num_heads,
+            dropout=matcha_dropout,
+            sigma_min=matcha_sigma_min,
+            use_diff_attention=matcha_use_diff_attention,
+        )
+
+    def _encode(
+        self,
+        x: torch.Tensor,
+        x_lengths: torch.Tensor,
+        sid: torch.Tensor,
+        encoder_args: tuple[torch.Tensor, ...],
+    ) -> tuple[torch.Tensor, ...]:
+        g = self.emb_g(sid).unsqueeze(-1)
+        hidden, mu, logs, x_mask = self.enc_p(
+            x, x_lengths, *encoder_args, g=g
+        )
+        return hidden, mu, logs, x_mask, g
+
+    @staticmethod
+    def _align(
+        mel: torch.Tensor,
+        mu: torch.Tensor,
+        logs: torch.Tensor,
+        x_mask: torch.Tensor,
+        mel_mask: torch.Tensor,
+    ) -> torch.Tensor:
+        """Run MAS using the encoder's diagonal Gaussian mel prior."""
+        from style_bert_vits2.models import monotonic_alignment
+
+        with torch.no_grad():
+            inv_variance = torch.exp(-2 * logs)
+            neg_cent1 = torch.sum(
+                -0.5 * math.log(2 * math.pi) - logs, 1, keepdim=True
+            )
+            neg_cent2 = torch.matmul(
+                -0.5 * (mel**2).transpose(1, 2), inv_variance
+            )
+            neg_cent3 = torch.matmul(
+                mel.transpose(1, 2), mu * inv_variance
+            )
+            neg_cent4 = torch.sum(
+                -0.5 * (mu**2) * inv_variance, 1, keepdim=True
+            )
+            score = neg_cent1 + neg_cent2 + neg_cent3 + neg_cent4
+            attn_mask = x_mask.unsqueeze(2) * mel_mask.unsqueeze(-1)
+            return (
+                monotonic_alignment.maximum_path(score, attn_mask.squeeze(1))
+                .unsqueeze(1)
+                .detach()
+            )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_lengths: torch.Tensor,
+        mel: torch.Tensor,
+        mel_lengths: torch.Tensor,
+        sid: torch.Tensor,
+        *encoder_args: torch.Tensor,
+        out_size: Optional[int] = None,
+        generate: bool = False,
+        n_timesteps: Optional[int] = None,
+        temperature: float = 1.0,
+    ) -> dict[str, torch.Tensor]:
+        """Compute acoustic losses and optionally a differentiable mel sample."""
+        hidden, mu, logs, x_mask, g = self._encode(
+            x, x_lengths, sid, encoder_args
+        )
+        mel_mask = commons.sequence_mask(mel_lengths, mel.shape[-1]).unsqueeze(1)
+        mel_mask = mel_mask.to(dtype=mu.dtype, device=mu.device)
+        attn = self._align(mel, mu, logs, x_mask, mel_mask)
+        durations = attn.sum(2)
+        log_durations = torch.log(durations + 1e-6) * x_mask
+        predicted_log_durations = self.dp(hidden, x_mask, g=g)
+        duration_loss = torch.sum(
+            (predicted_log_durations - log_durations) ** 2
+        ) / x_mask.sum().clamp_min(1)
+        duration_loss = duration_loss + (
+            self.sdp(hidden, x_mask, durations, g=g).sum()
+            / x_mask.sum().clamp_min(1)
+        )
+
+        aligned_mu = torch.matmul(
+            attn.squeeze(1), mu.transpose(1, 2)
+        ).transpose(1, 2)
+        aligned_logs = torch.matmul(
+            attn.squeeze(1), logs.transpose(1, 2)
+        ).transpose(1, 2)
+
+        ids_slice: Optional[torch.Tensor] = None
+        if out_size is not None:
+            mel, ids_slice = commons.rand_slice_segments(
+                mel, mel_lengths, out_size
+            )
+            aligned_mu = commons.slice_segments(aligned_mu, ids_slice, out_size)
+            aligned_logs = commons.slice_segments(
+                aligned_logs, ids_slice, out_size
+            )
+            mel_mask = commons.slice_segments(mel_mask, ids_slice, out_size)
+
+        flow_loss = self.matcha.compute_loss(mel, mel_mask, aligned_mu, g)
+        prior = (
+            aligned_logs
+            + 0.5 * ((mel - aligned_mu) ** 2) * torch.exp(-2 * aligned_logs)
+            + 0.5 * math.log(2 * math.pi)
+        )
+        prior_loss = (prior * mel_mask).sum() / (
+            mel_mask.sum().clamp_min(1) * self.n_mel_channels
+        )
+        result = {
+            "duration_loss": duration_loss,
+            "prior_loss": prior_loss,
+            "flow_loss": flow_loss,
+            "attn": attn,
+            "mel_mask": mel_mask,
+            "target_mel": mel,
+        }
+        if ids_slice is not None:
+            result["ids_slice"] = ids_slice
+        if generate:
+            result["generated_mel"] = self.matcha.sample(
+                aligned_mu,
+                mel_mask,
+                n_timesteps or self.matcha_n_timesteps,
+                temperature,
+                g,
+            )
+        return result
+
+    @torch.inference_mode()
+    def infer_mel(
+        self,
+        x: torch.Tensor,
+        x_lengths: torch.Tensor,
+        sid: torch.Tensor,
+        *encoder_args: torch.Tensor,
+        noise_scale: float = 1.0,
+        length_scale: float = 1.0,
+        noise_scale_w: float = 0.8,
+        sdp_ratio: float = 0.0,
+        n_timesteps: Optional[int] = None,
+        max_len: Optional[int] = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        hidden, mu, _, x_mask, g = self._encode(
+            x, x_lengths, sid, encoder_args
+        )
+        logw = self.sdp(
+            hidden, x_mask, g=g, reverse=True, noise_scale=noise_scale_w
+        ) * sdp_ratio + self.dp(hidden, x_mask, g=g) * (1 - sdp_ratio)
+        durations = torch.ceil(torch.exp(logw) * x_mask * length_scale)
+        mel_lengths = torch.clamp_min(durations.sum((1, 2)), 1).long()
+        mel_mask = commons.sequence_mask(mel_lengths, None).unsqueeze(1).to(x_mask)
+        attn_mask = x_mask.unsqueeze(2) * mel_mask.unsqueeze(-1)
+        attn = commons.generate_path(durations, attn_mask)
+        aligned_mu = torch.matmul(
+            attn.squeeze(1), mu.transpose(1, 2)
+        ).transpose(1, 2)
+        if max_len is not None:
+            aligned_mu = aligned_mu[:, :, :max_len]
+            mel_mask = mel_mask[:, :, :max_len]
+            attn = attn[:, :, :max_len]
+        generated_mel = self.matcha(
+            aligned_mu,
+            mel_mask,
+            n_timesteps or self.matcha_n_timesteps,
+            noise_scale,
+            g,
+        )
+        return generated_mel, attn, mel_mask
+
+
+class JointMelSynthesizer(nn.Module):
+    """Differentiable composition of a pretrained acoustic model and vocoder."""
+
+    def __init__(self, acoustic_model: MelAcousticModel, vocoder: nn.Module) -> None:
+        super().__init__()
+        self.acoustic_model = acoustic_model
+        self.vocoder = vocoder
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_lengths: torch.Tensor,
+        mel: torch.Tensor,
+        mel_lengths: torch.Tensor,
+        sid: torch.Tensor,
+        *encoder_args: torch.Tensor,
+        out_size: int,
+        n_timesteps: Optional[int] = None,
+        temperature: float = 1.0,
+    ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
+        acoustic = self.acoustic_model(
+            x,
+            x_lengths,
+            mel,
+            mel_lengths,
+            sid,
+            *encoder_args,
+            out_size=out_size,
+            generate=True,
+            n_timesteps=n_timesteps,
+            temperature=temperature,
+        )
+        waveform = self.vocoder(acoustic["generated_mel"], sid)
+        return waveform, acoustic
+
+    @torch.inference_mode()
+    def infer(self, *args: Any, **kwargs: Any) -> tuple[torch.Tensor, ...]:
+        mel, attn, mel_mask = self.acoustic_model.infer_mel(*args, **kwargs)
+        sid = args[2] if len(args) > 2 else kwargs["sid"]
+        return self.vocoder(mel, sid), mel, attn, mel_mask
+
+
+class MelVocoder(nn.Module):
+    """Standalone mel-to-wave Generator with its own speaker embedding."""
+
+    def __init__(
+        self,
+        generator: nn.Module,
+        n_speakers: int,
+        gin_channels: int,
+    ) -> None:
+        super().__init__()
+        if n_speakers < 1:
+            raise ValueError("Separated mel vocoders require at least one speaker")
+        self.generator = generator
+        self.emb_g = nn.Embedding(n_speakers, gin_channels)
+
+    def forward(self, mel: torch.Tensor, sid: torch.Tensor) -> torch.Tensor:
+        speaker = self.emb_g(sid).unsqueeze(-1)
+        return self.generator(mel, g=speaker)

style_bert_vits2/models/models.py

@@ -1133,3 +1133,89 @@ class SynthesizerTrn(nn.Module):
         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)
+
+
+def build_mel_synthesizer(
+    n_vocab: int,
+    n_mel_channels: int,
+    hidden_channels: int,
+    filter_channels: int,
+    n_heads: int,
+    n_layers: int,
+    kernel_size: int,
+    p_dropout: float,
+    n_speakers: int,
+    gin_channels: int,
+    matcha_channels: int = 192,
+    matcha_num_heads: int = 2,
+    matcha_dropout: float = 0.05,
+    matcha_sigma_min: float = 1e-4,
+    matcha_n_timesteps: int = 8,
+    matcha_use_diff_attention: bool = False,
+) -> "MelAcousticModel":
+    """Build the multilingual text-to-mel model without a waveform generator."""
+    from style_bert_vits2.models.mel_synthesizer import MelAcousticModel
+
+    encoder = TextEncoder(
+        n_vocab,
+        n_mel_channels,
+        hidden_channels,
+        filter_channels,
+        n_heads,
+        n_layers,
+        kernel_size,
+        p_dropout,
+        n_speakers,
+        gin_channels=gin_channels,
+    )
+    sdp = StochasticDurationPredictor(
+        hidden_channels, 192, 3, 0.5, 4, gin_channels=gin_channels
+    )
+    dp = DurationPredictor(
+        hidden_channels, 256, 3, 0.5, gin_channels=gin_channels
+    )
+    return MelAcousticModel(
+        encoder,
+        sdp,
+        dp,
+        n_mel_channels,
+        n_speakers,
+        gin_channels,
+        hidden_channels,
+        matcha_channels,
+        matcha_num_heads,
+        matcha_dropout,
+        matcha_sigma_min,
+        matcha_n_timesteps,
+        matcha_use_diff_attention,
+    )
+
+
+def build_mel_vocoder(
+    n_mel_channels: int,
+    resblock: str,
+    resblock_kernel_sizes: list[int],
+    resblock_dilation_sizes: list[list[int]],
+    upsample_rates: list[int],
+    upsample_initial_channel: int,
+    upsample_kernel_sizes: list[int],
+    n_speakers: int,
+    gin_channels: int,
+) -> "MelVocoder":
+    """Build a standalone HiFi-GAN-style mel-to-wave Generator."""
+    from style_bert_vits2.models.mel_synthesizer import MelVocoder
+
+    return MelVocoder(
+        Generator(
+            n_mel_channels,
+            resblock,
+            resblock_kernel_sizes,
+            resblock_dilation_sizes,
+            upsample_rates,
+            upsample_initial_channel,
+            upsample_kernel_sizes,
+            gin_channels=gin_channels,
+        ),
+        n_speakers=n_speakers,
+        gin_channels=gin_channels,
+    )

style_bert_vits2/models/models_jp_extra.py

@@ -1188,3 +1188,88 @@ class SynthesizerTrn(nn.Module):
         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)
+
+
+def build_mel_synthesizer(
+    n_vocab: int,
+    n_mel_channels: int,
+    hidden_channels: int,
+    filter_channels: int,
+    n_heads: int,
+    n_layers: int,
+    kernel_size: int,
+    p_dropout: float,
+    n_speakers: int,
+    gin_channels: int,
+    matcha_channels: int = 192,
+    matcha_num_heads: int = 2,
+    matcha_dropout: float = 0.05,
+    matcha_sigma_min: float = 1e-4,
+    matcha_n_timesteps: int = 8,
+    matcha_use_diff_attention: bool = False,
+) -> "MelAcousticModel":
+    """Build the JP-Extra text-to-mel model without a waveform generator."""
+    from style_bert_vits2.models.mel_synthesizer import MelAcousticModel
+
+    encoder = TextEncoder(
+        n_vocab,
+        n_mel_channels,
+        hidden_channels,
+        filter_channels,
+        n_heads,
+        n_layers,
+        kernel_size,
+        p_dropout,
+        gin_channels=gin_channels,
+    )
+    sdp = StochasticDurationPredictor(
+        hidden_channels, 192, 3, 0.5, 4, gin_channels=gin_channels
+    )
+    dp = DurationPredictor(
+        hidden_channels, 256, 3, 0.5, gin_channels=gin_channels
+    )
+    return MelAcousticModel(
+        encoder,
+        sdp,
+        dp,
+        n_mel_channels,
+        n_speakers,
+        gin_channels,
+        hidden_channels,
+        matcha_channels,
+        matcha_num_heads,
+        matcha_dropout,
+        matcha_sigma_min,
+        matcha_n_timesteps,
+        matcha_use_diff_attention,
+    )
+
+
+def build_mel_vocoder(
+    n_mel_channels: int,
+    resblock: str,
+    resblock_kernel_sizes: list[int],
+    resblock_dilation_sizes: list[list[int]],
+    upsample_rates: list[int],
+    upsample_initial_channel: int,
+    upsample_kernel_sizes: list[int],
+    n_speakers: int,
+    gin_channels: int,
+) -> "MelVocoder":
+    """Build a standalone HiFi-GAN-style mel-to-wave Generator."""
+    from style_bert_vits2.models.mel_synthesizer import MelVocoder
+
+    return MelVocoder(
+        Generator(
+            n_mel_channels,
+            resblock,
+            resblock_kernel_sizes,
+            resblock_dilation_sizes,
+            upsample_rates,
+            upsample_initial_channel,
+            upsample_kernel_sizes,
+            gin_channels=gin_channels,
+        ),
+        n_speakers=n_speakers,
+        gin_channels=gin_channels,
+    )

tests/test_mel_synthesizer.py

@@ -0,0 +1,136 @@
+import sys
+import types
+
+import torch
+
+
+def maximum_path(score: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+    path = torch.zeros_like(score)
+    for batch in range(score.shape[0]):
+        mel_length = int(mask[batch].any(-1).sum())
+        text_length = int(mask[batch].any(0).sum())
+        for frame in range(mel_length):
+            token = min(frame * text_length // mel_length, text_length - 1)
+            path[batch, frame, token] = 1
+    return path
+
+
+alignment_stub = types.ModuleType(
+    "style_bert_vits2.models.monotonic_alignment"
+)
+alignment_stub.maximum_path = maximum_path
+sys.modules["style_bert_vits2.models.monotonic_alignment"] = alignment_stub
+
+from style_bert_vits2.models.mel_synthesizer import (
+    JointMelSynthesizer,
+    MelAcousticModel,
+    MelVocoder,
+)
+
+
+class DummyEncoder(torch.nn.Module):
+    def __init__(self, hidden: int, n_mels: int) -> None:
+        super().__init__()
+        self.embedding = torch.nn.Embedding(16, hidden)
+        self.projection = torch.nn.Conv1d(hidden, n_mels * 2, 1)
+
+    def forward(self, x, x_lengths, *_args, g=None):
+        hidden = self.embedding(x).transpose(1, 2)
+        mask = (
+            torch.arange(x.shape[1], device=x.device)[None, :] < x_lengths[:, None]
+        ).unsqueeze(1).to(hidden)
+        mu, logs = self.projection(hidden * mask).chunk(2, dim=1)
+        return hidden, mu * mask, logs.tanh() * mask, mask
+
+
+class DummyDurationPredictor(torch.nn.Module):
+    def __init__(self, hidden: int) -> None:
+        super().__init__()
+        self.projection = torch.nn.Conv1d(hidden, 1, 1)
+
+    def forward(self, x, mask, g=None):
+        return self.projection(x) * mask
+
+
+class DummyStochasticDurationPredictor(DummyDurationPredictor):
+    def forward(self, x, mask, w=None, g=None, reverse=False, noise_scale=1.0):
+        prediction = super().forward(x, mask, g)
+        if reverse:
+            return prediction
+        return ((prediction - torch.log(w + 1e-6)) ** 2 * mask).sum()
+
+
+class DummyVocoder(torch.nn.Module):
+    def __init__(self, n_mels: int) -> None:
+        super().__init__()
+        self.projection = torch.nn.Conv1d(n_mels, 1, 1)
+
+    def forward(self, mel, g=None):
+        return self.projection(mel)
+
+
+def make_acoustic() -> MelAcousticModel:
+    hidden, n_mels, speaker_channels = 8, 4, 2
+    return MelAcousticModel(
+        DummyEncoder(hidden, n_mels),
+        DummyStochasticDurationPredictor(hidden),
+        DummyDurationPredictor(hidden),
+        n_mel_channels=n_mels,
+        n_speakers=2,
+        gin_channels=speaker_channels,
+        hidden_channels=hidden,
+        matcha_channels=8,
+        matcha_num_heads=2,
+        matcha_dropout=0.0,
+        matcha_sigma_min=1e-4,
+        matcha_n_timesteps=2,
+        matcha_use_diff_attention=False,
+    )
+
+
+def test_acoustic_model_outputs_mel_and_losses() -> None:
+    model = make_acoustic()
+    x = torch.tensor([[1, 2, 3], [1, 2, 0]])
+    x_lengths = torch.tensor([3, 2])
+    mel = torch.randn(2, 4, 7)
+    mel_lengths = torch.tensor([7, 5])
+    sid = torch.tensor([0, 1])
+
+    output = model(
+        x, x_lengths, mel, mel_lengths, sid, out_size=4, generate=True
+    )
+    loss = (
+        output["duration_loss"] + output["prior_loss"] + output["flow_loss"]
+        + output["generated_mel"].abs().mean()
+    )
+    loss.backward()
+
+    assert output["generated_mel"].shape == (2, 4, 4)
+    assert output["target_mel"].shape == (2, 4, 4)
+    assert all(torch.isfinite(output[name]) for name in (
+        "duration_loss", "prior_loss", "flow_loss"
+    ))
+    assert model.enc_p.embedding.weight.grad is not None
+    assert any(p.grad is not None for p in model.matcha.parameters())
+
+
+def test_joint_model_backpropagates_through_vocoder_to_acoustic() -> None:
+    acoustic = make_acoustic()
+    vocoder = MelVocoder(DummyVocoder(4), n_speakers=2, gin_channels=2)
+    model = JointMelSynthesizer(acoustic, vocoder)
+    waveform, output = model(
+        torch.tensor([[1, 2, 3]]),
+        torch.tensor([3]),
+        torch.randn(1, 4, 6),
+        torch.tensor([6]),
+        torch.tensor([0]),
+        out_size=4,
+        n_timesteps=2,
+    )
+
+    waveform.square().mean().backward()
+
+    assert waveform.shape == (1, 1, 4)
+    assert model.vocoder.generator.projection.weight.grad is not None
+    assert any(p.grad is not None for p in acoustic.matcha.parameters())
+    assert output["generated_mel"].grad_fn is not None

train_ms_mel.py

@@ -0,0 +1,466 @@
+"""Train the separated text-to-mel model, then fine-tune it with a vocoder.
+
+Examples:
+  python train_ms_mel.py -c Data/model/config.json --stage acoustic
+  python train_ms_mel.py -c Data/model/config.json --stage vocoder
+  python train_ms_mel.py -c Data/model/config.json --stage joint \
+      --acoustic-checkpoint Data/model/models/ACOUSTIC_10000.pth \
+      --vocoder-checkpoint Data/model/models/VOCODER_10000.pth
+"""
+
+import argparse
+import os
+from pathlib import Path
+from typing import Any
+
+import torch
+from torch.nn import functional as F
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+from data_utils import TextAudioSpeakerCollate, TextAudioSpeakerLoader
+from losses import discriminator_loss, feature_loss, generator_loss
+from mel_processing import mel_spectrogram_torch, spec_to_mel_torch
+from style_bert_vits2.models import commons
+from style_bert_vits2.models.hyper_parameters import HyperParameters
+from style_bert_vits2.models.mel_synthesizer import JointMelSynthesizer
+from style_bert_vits2.models.models import MultiPeriodDiscriminator
+from style_bert_vits2.nlp.symbols import SYMBOLS
+from style_bert_vits2.utils.stdout_wrapper import SAFE_STDOUT
+
+
+def _checkpoint_state(path: str) -> dict[str, torch.Tensor]:
+    if path.endswith(".safetensors"):
+        from safetensors.torch import load_file
+
+        return load_file(path)
+    checkpoint = torch.load(path, map_location="cpu", weights_only=True)
+    return checkpoint.get("model", checkpoint)
+
+
+def load_component(
+    module: torch.nn.Module, path: str, prefixes: tuple[str, ...] = ()
+) -> None:
+    """Load either a standalone or a prefixed joint/legacy component."""
+    saved = _checkpoint_state(path)
+    target = module.state_dict()
+    selected: dict[str, torch.Tensor] = {}
+    for key, value in saved.items():
+        candidates = [key]
+        if key.startswith("dec."):
+            candidates.append("generator." + key[len("dec.") :])
+        if key.startswith("module.dec."):
+            candidates.append("generator." + key[len("module.dec.") :])
+        for prefix in prefixes:
+            if key.startswith(prefix):
+                candidates.append(key[len(prefix) :])
+        for candidate in candidates:
+            if candidate in target and target[candidate].shape == value.shape:
+                selected[candidate] = value
+                break
+    if not selected:
+        raise ValueError(f"No compatible parameters found in {path}")
+    result = module.load_state_dict(selected, strict=False)
+    print(
+        f"Loaded {len(selected)}/{len(target)} tensors from {path} "
+        f"({len(result.missing_keys)} parameters kept at initialization)"
+    )
+
+
+def save_training_checkpoint(
+    path: Path,
+    model: torch.nn.Module,
+    optimizer: torch.optim.Optimizer,
+    epoch: int,
+    step: int,
+) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    torch.save(
+        {
+            "model": model.state_dict(),
+            "optimizer": optimizer.state_dict(),
+            "iteration": epoch,
+            "global_step": step,
+            "learning_rate": optimizer.param_groups[0]["lr"],
+        },
+        path,
+    )
+
+
+def build_models(hps: HyperParameters) -> tuple[torch.nn.Module, torch.nn.Module]:
+    module = (
+        __import__(
+            "style_bert_vits2.models.models_jp_extra",
+            fromlist=["build_mel_synthesizer"],
+        )
+        if hps.data.use_jp_extra
+        else __import__(
+            "style_bert_vits2.models.models",
+            fromlist=["build_mel_synthesizer"],
+        )
+    )
+    acoustic = module.build_mel_synthesizer(
+        len(SYMBOLS),
+        hps.data.n_mel_channels,
+        hps.model.hidden_channels,
+        hps.model.filter_channels,
+        hps.model.n_heads,
+        hps.model.n_layers,
+        hps.model.kernel_size,
+        hps.model.p_dropout,
+        hps.data.n_speakers,
+        hps.model.gin_channels,
+        hps.model.matcha_channels,
+        hps.model.matcha_num_heads,
+        hps.model.matcha_dropout,
+        hps.model.matcha_sigma_min,
+        hps.model.matcha_n_timesteps,
+        hps.model.matcha_use_diff_attention,
+    )
+    vocoder = module.build_mel_vocoder(
+        hps.data.n_mel_channels,
+        hps.model.resblock,
+        hps.model.resblock_kernel_sizes,
+        hps.model.resblock_dilation_sizes,
+        hps.model.upsample_rates,
+        hps.model.upsample_initial_channel,
+        hps.model.upsample_kernel_sizes,
+        hps.data.n_speakers,
+        hps.model.gin_channels,
+    )
+    return acoustic, vocoder
+
+
+def unpack_batch(
+    batch: tuple[torch.Tensor, ...],
+    hps: HyperParameters,
+    device: torch.device,
+) -> tuple[dict[str, Any], torch.Tensor]:
+    batch = tuple(item.to(device, non_blocking=True) for item in batch)
+    if hps.data.use_jp_extra:
+        (
+            x,
+            x_lengths,
+            spec,
+            spec_lengths,
+            waveform,
+            _,
+            speakers,
+            tone,
+            language,
+            bert,
+            style_vec,
+        ) = batch
+        encoder_args = (tone, language, bert, style_vec)
+    else:
+        (
+            x,
+            x_lengths,
+            spec,
+            spec_lengths,
+            waveform,
+            _,
+            speakers,
+            tone,
+            language,
+            bert,
+            ja_bert,
+            en_bert,
+            style_vec,
+        ) = batch
+        encoder_args = (
+            tone,
+            language,
+            bert,
+            ja_bert,
+            en_bert,
+            style_vec,
+            speakers,
+        )
+    mel = (
+        spec
+        if hps.model.use_mel_posterior_encoder
+        else spec_to_mel_torch(
+            spec,
+            hps.data.filter_length,
+            hps.data.n_mel_channels,
+            hps.data.sampling_rate,
+            hps.data.mel_fmin,
+            hps.data.mel_fmax,
+        )
+    )
+    return {
+        "x": x,
+        "x_lengths": x_lengths,
+        "mel": mel,
+        "mel_lengths": spec_lengths,
+        "sid": speakers,
+        "encoder_args": encoder_args,
+    }, waveform
+
+
+def acoustic_loss(outputs: dict[str, torch.Tensor], hps: HyperParameters) -> torch.Tensor:
+    return (
+        outputs["duration_loss"]
+        + outputs["prior_loss"]
+        + outputs["flow_loss"] * hps.train.c_matcha
+    )
+
+
+def run() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-c", "--config", required=True)
+    parser.add_argument(
+        "--stage", choices=("acoustic", "vocoder", "joint"), required=True
+    )
+    parser.add_argument("--acoustic-checkpoint")
+    parser.add_argument("--vocoder-checkpoint")
+    parser.add_argument("--output-dir")
+    parser.add_argument("--save-every", type=int, default=1000)
+    parser.add_argument("--num-workers", type=int, default=1)
+    parser.add_argument("--joint-timesteps", type=int, default=2)
+    parser.add_argument(
+        "--no_progress_bar",
+        action="store_true",
+        help="Disable the tqdm training progress bar.",
+    )
+    args = parser.parse_args()
+
+    if args.stage == "joint" and (
+        not args.acoustic_checkpoint or not args.vocoder_checkpoint
+    ):
+        parser.error(
+            "joint stage requires --acoustic-checkpoint and --vocoder-checkpoint"
+        )
+
+    hps = HyperParameters.load_from_json(args.config)
+    if hps.data.n_speakers < 1:
+        raise ValueError("Separated mel training requires data.n_speakers >= 1")
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    acoustic, vocoder = build_models(hps)
+    if args.acoustic_checkpoint:
+        load_component(
+            acoustic,
+            args.acoustic_checkpoint,
+            ("acoustic_model.", "module.acoustic_model.", "module."),
+        )
+    if args.vocoder_checkpoint:
+        load_component(
+            vocoder,
+            args.vocoder_checkpoint,
+            ("vocoder.", "module.vocoder.", "dec.", "module.dec.", "module."),
+        )
+
+    discriminator = None
+    if args.stage == "acoustic":
+        model: torch.nn.Module = acoustic.to(device)
+    elif args.stage == "vocoder":
+        model = vocoder.to(device)
+        discriminator = MultiPeriodDiscriminator(
+            hps.model.use_spectral_norm
+        ).to(device)
+    else:
+        model = JointMelSynthesizer(acoustic, vocoder).to(device)
+        discriminator = MultiPeriodDiscriminator(
+            hps.model.use_spectral_norm
+        ).to(device)
+
+    optimizer = torch.optim.AdamW(
+        model.parameters(),
+        hps.train.learning_rate,
+        betas=hps.train.betas,
+        eps=hps.train.eps,
+    )
+    optimizer_d = (
+        torch.optim.AdamW(
+            discriminator.parameters(),
+            hps.train.learning_rate,
+            betas=hps.train.betas,
+            eps=hps.train.eps,
+        )
+        if discriminator is not None
+        else None
+    )
+    dataset = TextAudioSpeakerLoader(hps.data.training_files, hps.data)
+    loader = DataLoader(
+        dataset,
+        batch_size=hps.train.batch_size,
+        shuffle=True,
+        num_workers=args.num_workers,
+        pin_memory=device.type == "cuda",
+        collate_fn=TextAudioSpeakerCollate(use_jp_extra=hps.data.use_jp_extra),
+        drop_last=True,
+    )
+    output_dir = Path(args.output_dir or Path(args.config).parent / "models")
+    segment_frames = hps.train.segment_size // hps.data.hop_length
+    global_step = 0
+    pbar = None
+    if not args.no_progress_bar:
+        pbar = tqdm(
+            total=hps.train.epochs * len(loader),
+            initial=global_step,
+            smoothing=0.05,
+            file=SAFE_STDOUT,
+            dynamic_ncols=True,
+        )
+
+    for epoch in range(1, hps.train.epochs + 1):
+        model.train()
+        for batch_idx, batch in enumerate(loader):
+            inputs, waveform = unpack_batch(batch, hps, device)
+            if args.stage == "acoustic":
+                outputs = acoustic(
+                    inputs["x"],
+                    inputs["x_lengths"],
+                    inputs["mel"],
+                    inputs["mel_lengths"],
+                    inputs["sid"],
+                    *inputs["encoder_args"],
+                    out_size=segment_frames,
+                )
+                loss = acoustic_loss(outputs, hps)
+            elif args.stage == "vocoder":
+                assert discriminator is not None and optimizer_d is not None
+                target_mel, ids = commons.rand_slice_segments(
+                    inputs["mel"], inputs["mel_lengths"], segment_frames
+                )
+                real = commons.slice_segments(
+                    waveform,
+                    ids * hps.data.hop_length,
+                    hps.train.segment_size,
+                )
+                generated = vocoder(target_mel, inputs["sid"])
+                real_scores, fake_scores, _, _ = discriminator(
+                    real, generated.detach()
+                )
+                loss_d, _, _ = discriminator_loss(real_scores, fake_scores)
+                optimizer_d.zero_grad(set_to_none=True)
+                loss_d.backward()
+                optimizer_d.step()
+                _, fake_scores, fmap_r, fmap_g = discriminator(real, generated)
+                generated_mel = mel_spectrogram_torch(
+                    generated.squeeze(1).float(),
+                    hps.data.filter_length,
+                    hps.data.n_mel_channels,
+                    hps.data.sampling_rate,
+                    hps.data.hop_length,
+                    hps.data.win_length,
+                    hps.data.mel_fmin,
+                    hps.data.mel_fmax,
+                )
+                loss_g, _ = generator_loss(fake_scores)
+                loss = (
+                    loss_g
+                    + feature_loss(fmap_r, fmap_g)
+                    + F.l1_loss(target_mel, generated_mel) * hps.train.c_mel
+                )
+            else:
+                assert discriminator is not None and optimizer_d is not None
+                generated, outputs = model(
+                    inputs["x"],
+                    inputs["x_lengths"],
+                    inputs["mel"],
+                    inputs["mel_lengths"],
+                    inputs["sid"],
+                    *inputs["encoder_args"],
+                    out_size=segment_frames,
+                    n_timesteps=args.joint_timesteps,
+                )
+                ids = outputs["ids_slice"]
+                real = commons.slice_segments(
+                    waveform,
+                    ids * hps.data.hop_length,
+                    hps.train.segment_size,
+                )
+                real_scores, fake_scores, _, _ = discriminator(
+                    real, generated.detach()
+                )
+                loss_d, _, _ = discriminator_loss(real_scores, fake_scores)
+                optimizer_d.zero_grad(set_to_none=True)
+                loss_d.backward()
+                optimizer_d.step()
+
+                real_scores, fake_scores, fmap_r, fmap_g = discriminator(
+                    real, generated
+                )
+                generated_mel = mel_spectrogram_torch(
+                    generated.squeeze(1).float(),
+                    hps.data.filter_length,
+                    hps.data.n_mel_channels,
+                    hps.data.sampling_rate,
+                    hps.data.hop_length,
+                    hps.data.win_length,
+                    hps.data.mel_fmin,
+                    hps.data.mel_fmax,
+                )
+                loss_g, _ = generator_loss(fake_scores)
+                loss = (
+                    acoustic_loss(outputs, hps)
+                    + loss_g
+                    + feature_loss(fmap_r, fmap_g)
+                    + F.l1_loss(outputs["target_mel"], generated_mel)
+                    * hps.train.c_mel
+                )
+
+            optimizer.zero_grad(set_to_none=True)
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 500)
+            optimizer.step()
+            global_step += 1
+            if pbar is not None:
+                epoch_progress = 100.0 * (batch_idx + 1) / len(loader)
+                pbar.set_description(
+                    f"{args.stage.capitalize()} "
+                    f"{epoch}({epoch_progress:.0f}%)/{hps.train.epochs}"
+                )
+                if args.stage == "acoustic":
+                    pbar.set_postfix(
+                        loss=f"{loss.item():.4f}",
+                        dur=f"{outputs['duration_loss'].item():.4f}",
+                        prior=f"{outputs['prior_loss'].item():.4f}",
+                        flow=f"{outputs['flow_loss'].item():.4f}",
+                    )
+                else:
+                    pbar.set_postfix(loss=f"{loss.item():.4f}")
+                pbar.update()
+            if global_step % hps.train.log_interval == 0:
+                message = (
+                    f"epoch={epoch} step={global_step} "
+                    f"stage={args.stage} loss={loss.item():.5f}"
+                )
+                if pbar is not None:
+                    pbar.write(message)
+                else:
+                    print(message)
+            if global_step % args.save_every == 0:
+                name = args.stage.upper()
+                save_training_checkpoint(
+                    output_dir / f"{name}_{global_step}.pth",
+                    model,
+                    optimizer,
+                    epoch,
+                    global_step,
+                )
+                if optimizer_d is not None and discriminator is not None:
+                    save_training_checkpoint(
+                        output_dir / f"D_{global_step}.pth",
+                        discriminator,
+                        optimizer_d,
+                        epoch,
+                        global_step,
+                    )
+
+    name = args.stage.upper()
+    save_training_checkpoint(
+        output_dir / f"{name}_{global_step}.pth",
+        model,
+        optimizer,
+        hps.train.epochs,
+        global_step,
+    )
+    if pbar is not None:
+        pbar.close()
+
+
+if __name__ == "__main__":
+    run()