merge

README.md

@@ -1,47 +1,5 @@
 # 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,23 +224,14 @@ 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=device, dtype=model_dtype).unsqueeze(0)
-        ja_bert = ja_bert.to(device=device, dtype=model_dtype).unsqueeze(0)
-        en_bert = en_bert.to(device=device, dtype=model_dtype).unsqueeze(0)
+        bert = bert.to(device).unsqueeze(0)
+        ja_bert = ja_bert.to(device).unsqueeze(0)
+        en_bert = en_bert.to(device).unsqueeze(0)
         x_tst_lengths = torch.LongTensor([phones.size(0)]).to(device)
-        style_vec_tensor = (
-            torch.from_numpy(style_vec)
-            .to(device=device, dtype=model_dtype)
-            .unsqueeze(0)
-        )
+        style_vec_tensor = torch.from_numpy(style_vec).to(device).unsqueeze(0)
         del phones
         sid_tensor = torch.LongTensor([sid]).to(device)
 

style_bert_vits2/models/matcha_flow.py

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

@@ -1,284 +0,0 @@
-"""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,89 +1133,3 @@ 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,88 +1188,3 @@ 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

@@ -1,136 +0,0 @@
-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

@@ -1,407 +0,0 @@
-"""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 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
-
-
-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]:
-    (
-        x,
-        x_lengths,
-        spec,
-        spec_lengths,
-        waveform,
-        _,
-        speakers,
-        tone,
-        language,
-        bert,
-        ja_bert,
-        en_bert,
-        style_vec,
-    ) = (item.to(device, non_blocking=True) for item in batch)
-    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,
-        )
-    )
-    encoder_args = (
-        (tone, language, bert, style_vec)
-        if hps.data.use_jp_extra
-        else (tone, language, bert, ja_bert, en_bert, style_vec, speakers)
-    )
-    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)
-    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(),
-        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
-
-    for epoch in range(1, hps.train.epochs + 1):
-        model.train()
-        for batch in 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 global_step % hps.train.log_interval == 0:
-                print(
-                    f"epoch={epoch} step={global_step} "
-                    f"stage={args.stage} loss={loss.item():.5f}"
-                )
-            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 __name__ == "__main__":
-    run()