From 8125666e22d7f83171c432fa6ac58413256a2499 Mon Sep 17 00:00:00 2001 From: tuna2134 Date: Fri, 24 Jul 2026 16:24:50 +0900 Subject: [PATCH 1/3] split --- README.md | 42 +++ style_bert_vits2/models/matcha_flow.py | 15 +- style_bert_vits2/models/mel_synthesizer.py | 284 ++++++++++++++ style_bert_vits2/models/models.py | 86 +++++ style_bert_vits2/models/models_jp_extra.py | 85 +++++ tests/test_mel_synthesizer.py | 136 +++++++ train_ms_mel.py | 407 +++++++++++++++++++++ 7 files changed, 1053 insertions(+), 2 deletions(-) create mode 100644 style_bert_vits2/models/mel_synthesizer.py create mode 100644 tests/test_mel_synthesizer.py create mode 100644 train_ms_mel.py diff --git a/README.md b/README.md index ba4918c..72d0dbb 100644 --- a/README.md +++ b/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. diff --git a/style_bert_vits2/models/matcha_flow.py b/style_bert_vits2/models/matcha_flow.py index 0a4a50e..49b1079 100644 --- a/style_bert_vits2/models/matcha_flow.py +++ b/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 forward( + def sample( 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) diff --git a/style_bert_vits2/models/mel_synthesizer.py b/style_bert_vits2/models/mel_synthesizer.py new file mode 100644 index 0000000..b78f134 --- /dev/null +++ b/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] + 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) diff --git a/style_bert_vits2/models/models.py b/style_bert_vits2/models/models.py index d270ff3..994c748 100644 --- a/style_bert_vits2/models/models.py +++ b/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, + ) diff --git a/style_bert_vits2/models/models_jp_extra.py b/style_bert_vits2/models/models_jp_extra.py index 221e11e..31fbd46 100644 --- a/style_bert_vits2/models/models_jp_extra.py +++ b/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, + ) diff --git a/tests/test_mel_synthesizer.py b/tests/test_mel_synthesizer.py new file mode 100644 index 0000000..da7bc33 --- /dev/null +++ b/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 diff --git a/train_ms_mel.py b/train_ms_mel.py new file mode 100644 index 0000000..c472c4e --- /dev/null +++ b/train_ms_mel.py @@ -0,0 +1,407 @@ +"""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() -- 2.49.1 From 0d6eebc3a849b7b7bf87fdefaeed1656ee7abe77 Mon Sep 17 00:00:00 2001 From: tuna2134 Date: Fri, 24 Jul 2026 16:29:09 +0900 Subject: [PATCH 2/3] fix --- style_bert_vits2/models/mel_synthesizer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/style_bert_vits2/models/mel_synthesizer.py b/style_bert_vits2/models/mel_synthesizer.py index b78f134..b4b7e20 100644 --- a/style_bert_vits2/models/mel_synthesizer.py +++ b/style_bert_vits2/models/mel_synthesizer.py @@ -260,7 +260,7 @@ class JointMelSynthesizer(nn.Module): @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] + sid = args[2] if len(args) > 2 else kwargs["sid"] return self.vocoder(mel, sid), mel, attn, mel_mask -- 2.49.1 From db37175b87a93227aeb45ad25c91eb52016db2ae Mon Sep 17 00:00:00 2001 From: tuna2134 Date: Sun, 26 Jul 2026 00:49:25 +0900 Subject: [PATCH 3/3] fix --- style_bert_vits2/models/infer.py | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/style_bert_vits2/models/infer.py b/style_bert_vits2/models/infer.py index 835a042..c07b8c0 100644 --- a/style_bert_vits2/models/infer.py +++ b/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) - ja_bert = ja_bert.to(device).unsqueeze(0) - en_bert = en_bert.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) 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) -- 2.49.1