fix

Dockerfile.deploy

@@ -2,7 +2,7 @@
 
 # See https://huggingface.co/docs/hub/spaces-sdks-docker-first-demo
 
-FROM python:3.10
+FROM python:3.12
 
 RUN useradd -m -u 1000 user
 

Dockerfile.train

@@ -5,8 +5,8 @@
 
 # 主なバージョン等
 # Ubuntu 22.04
-# Python 3.10
-# PyTorch 2.1.2 (CUDA 11.8)
+# Python 3.12
+# PyTorch 2.7.x (CUDA 11.8)
 # CUDA Toolkit 12.0, CUDNN 8.9.7
 
 
@@ -24,7 +24,7 @@ ENV SHELL=/bin/bash
 ENV DEBIAN_FRONTEND=noninteractive
 
 ENV APT_INSTALL="apt-get install -y --no-install-recommends"
-ENV PIP_INSTALL="python3 -m pip --no-cache-dir install --upgrade"
+ENV PIP_INSTALL="python3.12 -m pip --no-cache-dir install --upgrade"
 ENV GIT_CLONE="git clone --depth 10"
 
 # ==================================================================
@@ -44,10 +44,14 @@ RUN apt-get update && \
     nano \
     ffmpeg \
     software-properties-common \
-    gnupg \
-    python3 \
-    python3-pip \
-    python3-dev
+    gnupg
+
+RUN add-apt-repository ppa:deadsnakes/ppa && \
+    apt-get update && \
+    $APT_INSTALL python3.12 python3.12-dev && \
+    curl -sS https://bootstrap.pypa.io/get-pip.py -o /tmp/get-pip.py && \
+    python3.12 /tmp/get-pip.py && \
+    rm /tmp/get-pip.py
 
 # ==================================================================
 # Git-lfs
@@ -57,8 +61,9 @@ RUN curl -s https://packagecloud.io/install/repositories/github/git-lfs/script.d
     $APT_INSTALL git-lfs
 
 
-# Add symlink so python and python3 commands use same python3.9 executable
-RUN ln -s /usr/bin/python3 /usr/local/bin/python    
+# Add symlinks so python and python3 commands use Python 3.12.
+RUN ln -s /usr/bin/python3.12 /usr/local/bin/python && \
+    ln -s /usr/bin/python3.12 /usr/local/bin/python3
 
 # ==================================================================
 # Installing CUDA packages (CUDA Toolkit 12.0 and CUDNN 8.9.7)
@@ -106,4 +111,4 @@ RUN rm /tmp/requirements.txt
 
 EXPOSE 8888 6006
 
-CMD jupyter lab --allow-root --ip=0.0.0.0 --no-browser --ServerApp.trust_xheaders=True --ServerApp.disable_check_xsrf=False --ServerApp.allow_remote_access=True --ServerApp.allow_origin='*' --ServerApp.allow_credentials=True
+CMD jupyter lab --allow-root --ip=0.0.0.0 --no-browser --ServerApp.trust_xheaders=True --ServerApp.disable_check_xsrf=False --ServerApp.allow_remote_access=True --ServerApp.allow_origin='*' --ServerApp.allow_credentials=True

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.
@@ -79,9 +121,9 @@ Pythonの仮想環境・パッケージ管理ツールである[uv](https://gith
 powershell -c "irm https://astral.sh/uv/install.ps1 | iex"
 git clone https://github.com/litagin02/Style-Bert-VITS2.git
 cd Style-Bert-VITS2
-uv venv venv
+uv venv --python 3.12 venv
 venv\Scripts\activate
-uv pip install "torch<2.4" "torchaudio<2.4" --index-url https://download.pytorch.org/whl/cu118
+uv pip install "torch>=2.8,<2.9" "torchaudio>=2.8,<2.9" --index-url https://download.pytorch.org/whl/cu126
 uv pip install -r requirements.txt
 python initialize.py  # 必要なモデルとデフォルトTTSモデルをダウンロード
 ```

app.py

@@ -49,7 +49,7 @@ model_holder = TTSModelHolder(
     ignore_onnx=True,
 )
 
-with gr.Blocks(theme=GRADIO_THEME) as app:
+with gr.Blocks() as app:
     gr.Markdown(f"# Style-Bert-VITS2 WebUI (version {VERSION})")
     with gr.Tabs():
         with gr.Tab("音声合成"):
@@ -66,6 +66,7 @@ with gr.Blocks(theme=GRADIO_THEME) as app:
             create_onnx_app(model_holder=model_holder)
 
 app.launch(
+    theme=GRADIO_THEME,
     server_name=args.host,
     server_port=args.port,
     inbrowser=not args.no_autolaunch,

bert_gen.py

@@ -55,7 +55,7 @@ def process_line(x: tuple[str, bool]):
     bert_path = wav_path.replace(".WAV", ".wav").replace(".wav", ".bert.pt")
 
     try:
-        bert = torch.load(bert_path)
+        bert = torch.load(bert_path, weights_only=True)
         assert bert.shape[-1] == len(phone)
     except Exception:
         bert = extract_bert_feature(text, word2ph, Languages(language_str), device)

data_utils.py

@@ -129,7 +129,7 @@ class TextAudioSpeakerLoader(torch.utils.data.Dataset):
         if self.use_mel_spec_posterior:
             spec_filename = spec_filename.replace(".spec.pt", ".mel.pt")
         try:
-            spec = torch.load(spec_filename)
+            spec = torch.load(spec_filename, weights_only=True)
         except:
             if self.use_mel_spec_posterior:
                 spec = mel_spectrogram_torch(
@@ -168,7 +168,7 @@ class TextAudioSpeakerLoader(torch.utils.data.Dataset):
             word2ph[0] += 1
         bert_path = wav_path.replace(".wav", ".bert.pt")
         try:
-            bert_ori = torch.load(bert_path)
+            bert_ori = torch.load(bert_path, weights_only=True)
             assert bert_ori.shape[-1] == len(phone)
         except Exception as e:
             logger.warning("Bert load Failed")

gradio_tabs/convert_onnx.py

@@ -52,7 +52,7 @@ def create_onnx_app(model_holder: TTSModelHolder) -> gr.Blocks:
     initial_id = 0
     initial_pth_files = get_model_files(model_names[initial_id])
 
-    with gr.Blocks(theme=GRADIO_THEME) as app:
+    with gr.Blocks() as app:
         gr.Markdown(initial_md)
         with gr.Row():
             with gr.Column():
@@ -100,4 +100,4 @@ if __name__ == "__main__":
     assets_root = path_config.assets_root
     model_holder = TTSModelHolder(assets_root, "cpu", "", ignore_onnx=True)
     app = create_onnx_app(model_holder)
-    app.launch(inbrowser=True)
+    app.launch(inbrowser=True, theme=GRADIO_THEME)

gradio_tabs/dataset.py

@@ -115,7 +115,7 @@ Style-Bert-VITS2の学習用データセットを作成するためのツール
 
 
 def create_dataset_app() -> gr.Blocks:
-    with gr.Blocks(theme=GRADIO_THEME) as app:
+    with gr.Blocks() as app:
         gr.Markdown(
             "**既に1ファイル2-12秒程度の音声ファイル集とその書き起こしデータがある場合は、このタブは使用せずに学習できます。**"
         )
@@ -273,4 +273,4 @@ def create_dataset_app() -> gr.Blocks:
 
 if __name__ == "__main__":
     app = create_dataset_app()
-    app.launch(inbrowser=True)
+    app.launch(inbrowser=True, theme=GRADIO_THEME)

gradio_tabs/inference.py

@@ -361,7 +361,7 @@ def create_inference_app(model_holder: TTSModelHolder) -> gr.Blocks:
     initial_id = 0
     initial_pth_files = get_model_files(model_names[initial_id])
 
-    with gr.Blocks(theme=GRADIO_THEME) as app:
+    with gr.Blocks() as app:
         gr.Markdown(initial_md)
         gr.Markdown(terms_of_use_md)
         null_models = gr.State({})
@@ -742,4 +742,4 @@ if __name__ == "__main__":
         assets_root, device, torch_device_to_onnx_providers(device)
     )
     app = create_inference_app(model_holder)
-    app.launch(inbrowser=True)
+    app.launch(inbrowser=True, theme=GRADIO_THEME)

gradio_tabs/merge.py

@@ -1042,7 +1042,7 @@ def create_merge_app(model_holder: TTSModelHolder) -> gr.Blocks:
         str(f) for f in model_holder.model_files_dict[model_names[initial_id]]
     ]
 
-    with gr.Blocks(theme=GRADIO_THEME) as app:
+    with gr.Blocks() as app:
         gr.Markdown(
             "複数のStyle-Bert-VITS2モデルから、声質・話し方・話す速さを取り替えたり混ぜたり引いたりして新しいモデルを作成できます。"
         )
@@ -1554,4 +1554,4 @@ if __name__ == "__main__":
         assets_root, device, torch_device_to_onnx_providers(device), ignore_onnx=True
     )
     app = create_merge_app(model_holder)
-    app.launch(inbrowser=True)
+    app.launch(inbrowser=True, theme=GRADIO_THEME)

gradio_tabs/style_vectors.py

@@ -416,7 +416,7 @@ https://ja.wikipedia.org/wiki/DBSCAN
 
 
 def create_style_vectors_app():
-    with gr.Blocks(theme=GRADIO_THEME) as app:
+    with gr.Blocks() as app:
         with gr.Accordion("使い方", open=False):
             gr.Markdown(how_to_md)
         model_name = gr.Textbox(placeholder="your_model_name", label="モデル名")
@@ -583,4 +583,4 @@ def create_style_vectors_app():
 
 if __name__ == "__main__":
     app = create_style_vectors_app()
-    app.launch(inbrowser=True)
+    app.launch(inbrowser=True, theme=GRADIO_THEME)

gradio_tabs/train.py

@@ -492,7 +492,7 @@ english_teacher.wav|Mary|EN|How are you? I'm fine, thank you, and you?
 
 
 def create_train_app():
-    with gr.Blocks(theme=GRADIO_THEME).queue() as app:
+    with gr.Blocks().queue() as app:
         gr.Markdown(change_log_md)
         with gr.Accordion("使い方", open=False):
             gr.Markdown(how_to_md)
@@ -852,4 +852,4 @@ def create_train_app():
 
 if __name__ == "__main__":
     app = create_train_app()
-    app.launch(inbrowser=True)
+    app.launch(inbrowser=True, theme=GRADIO_THEME)

pyproject.toml

@@ -7,7 +7,7 @@ name = "style-bert-vits2"
 dynamic = ["version"]
 description = "Style-Bert-VITS2: Bert-VITS2 with more controllable voice styles."
 readme = "README.md"
-requires-python = ">=3.9"
+requires-python = ">=3.12"
 license = "AGPL-3.0"
 keywords = []
 authors = [
@@ -16,9 +16,7 @@ authors = [
 classifiers = [
     "Development Status :: 4 - Beta",
     "Programming Language :: Python",
-    "Programming Language :: Python :: 3.9",
-    "Programming Language :: Python :: 3.10",
-    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
     "Programming Language :: Python :: Implementation :: CPython",
 ]
 dependencies = [
@@ -27,13 +25,13 @@ dependencies = [
     "g2p_en",
     "jieba",
     "loguru",
-    "nltk<=3.8.1",
+    "nltk>=3.9,<4",
     "num2words",
     "numba",
-    "numpy<2",
+    "numpy>=2,<2.5",
     "onnxruntime",
     "pydantic>=2.0",
-    "pyopenjtalk-dict",
+    "pyopenjtalk-plus",
     "pypinyin",
     "pyworld-prebuilt",
     "safetensors",
@@ -43,7 +41,7 @@ dependencies = [
 [project.optional-dependencies]
 torch = [
     "accelerate",
-    "torch>=2.1",
+    "torch>=2.8,<2.9",
 ]
 
 [project.urls]
@@ -93,7 +91,7 @@ cov-report = ["- coverage combine", "coverage report"]
 # Usage: `hatch run test:cov`
 cov = ["test-cov", "cov-report"]
 [[tool.hatch.envs.test.matrix]]
-python = ["3.9", "3.10", "3.11"]
+python = ["3.12"]
 
 # for ONNX inference (without PyTorch dependency)
 [tool.hatch.envs.test-onnx]
@@ -120,7 +118,7 @@ cov-report = ["- coverage combine", "coverage report"]
 # Usage: `hatch run test-onnx:cov`
 cov = ["test-cov", "cov-report"]
 [[tool.hatch.envs.test-onnx.matrix]]
-python = ["3.9", "3.10", "3.11"]
+python = ["3.12"]
 
 [tool.hatch.envs.style]
 detached = true

requirements-colab.txt

@@ -4,19 +4,19 @@ cn2an
 g2p_en
 gradio>=4.32
 jieba
-librosa==0.9.2
+librosa>=0.11,<1
 loguru
-nltk<=3.8.1
+nltk>=3.9,<4
 num2words
-numpy<2
+numpy>=2,<2.5
 onnx
 onnxconverter-common
 onnxruntime
 onnxruntime-gpu
 onnxsim-prebuilt
-pyannote.audio>=3.1.0
+pyannote.audio>=4,<5
 pyloudnorm
-pyopenjtalk-dict
+pyopenjtalk-plus
 pypinyin
 pyworld-prebuilt
 torch

requirements-infer.txt

@@ -1,33 +1,33 @@
 accelerate
 cmudict
 cn2an
-# faster-whisper==0.10.1
+# faster-whisper>=1.2,<2
 g2p_en
 GPUtil
 gradio>=4.32
 jieba
-# librosa==0.9.2
+# librosa>=0.11,<1
 loguru
-nltk<=3.8.1
+nltk>=3.9,<4
 num2words
-numpy<2
+numpy>=2,<2.5
 onnx
 onnxconverter-common
 onnxruntime
 onnxruntime-directml; sys_platform == 'win32'
 onnxruntime-gpu; sys_platform != 'darwin'
 onnxsim-prebuilt
-# protobuf==4.25
+# protobuf>=5,<7
 psutil
 # punctuators
-pyannote.audio>=3.1.0
+pyannote.audio>=4,<5
 # pyloudnorm
-pyopenjtalk-dict
+pyopenjtalk-plus
 pypinyin
 pyworld-prebuilt
 # stable_ts
 # tensorboard
-torch<2.4
-torchaudio<2.4
+torch>=2.8,<2.9
+torchaudio>=2.8,<2.9
 transformers
 umap-learn

requirements.txt

@@ -1,33 +1,33 @@
 accelerate
 cmudict
 cn2an
-faster-whisper==0.10.1
+faster-whisper>=1.2,<2
 g2p_en
 GPUtil
 gradio>=4.32
 jieba
-librosa==0.9.2
+librosa>=0.11,<1
 loguru
-nltk<=3.8.1
+nltk>=3.9,<4
 num2words
-numpy<2
+numpy>=2,<2.5
 onnx
 onnxconverter-common
 onnxruntime
 onnxruntime-directml; sys_platform == 'win32'
 onnxruntime-gpu; sys_platform != 'darwin'
 onnxsim-prebuilt
-protobuf==4.25
+protobuf>=5,<7
 psutil
 punctuators
-pyannote.audio>=3.1.0
+pyannote.audio>=4,<5
 pyloudnorm
-pyopenjtalk-dict
+pyopenjtalk-plus
 pypinyin
 pyworld-prebuilt
 stable_ts
 tensorboard
-torch<2.4
-torchaudio<2.4
+torch>=2.8,<2.9
+torchaudio>=2.8,<2.9
 transformers
 umap-learn

scripts/Install-Style-Bert-VITS2.bat

@@ -110,8 +110,8 @@ if !errorlevel! neq 0 ( pause & popd & exit /b !errorlevel! )
 echo --------------------------------------------------
 echo Installing PyTorch...
 echo --------------------------------------------------
-echo Executing: uv pip install "torch<2.4" "torchaudio<2.4" --index-url https://download.pytorch.org/whl/cu118
-uv pip install "torch<2.4" "torchaudio<2.4" --index-url https://download.pytorch.org/whl/cu118
+echo Executing: uv pip install "torch>=2.8,<2.9" "torchaudio>=2.8,<2.9" --index-url https://download.pytorch.org/whl/cu126
+uv pip install "torch>=2.8,<2.9" "torchaudio>=2.8,<2.9" --index-url https://download.pytorch.org/whl/cu126
 if !errorlevel! neq 0 ( pause & popd & exit /b !errorlevel! )
 
 echo --------------------------------------------------

scripts/Setup-Python.bat

@@ -40,8 +40,8 @@ if not exist "%PYTHON_DIR%"\ (
 	echo --------------------------------------------------
 	echo Downloading Python...
 	echo --------------------------------------------------
-	echo Executing: %CURL_CMD% -o python.zip https://www.python.org/ftp/python/3.10.11/python-3.10.11-embed-amd64.zip
-	%CURL_CMD% -o python.zip https://www.python.org/ftp/python/3.10.11/python-3.10.11-embed-amd64.zip
+	echo Executing: %CURL_CMD% -o python.zip https://www.python.org/ftp/python/3.12.10/python-3.12.10-embed-amd64.zip
+	%CURL_CMD% -o python.zip https://www.python.org/ftp/python/3.12.10/python-3.12.10-embed-amd64.zip
 	if !errorlevel! neq 0 ( pause & exit /b !errorlevel! )
 
 	echo --------------------------------------------------
@@ -61,8 +61,8 @@ if not exist "%PYTHON_DIR%"\ (
 	echo --------------------------------------------------
 	echo Enabling 'site' module in the embedded Python environment...
 	echo --------------------------------------------------
-	echo Executing: %PS_CMD% "&{(Get-Content '%PYTHON_DIR%/python310._pth') -creplace '#import site', 'import site' | Set-Content '%PYTHON_DIR%/python310._pth' }"
-	%PS_CMD% "&{(Get-Content '%PYTHON_DIR%/python310._pth') -creplace '#import site', 'import site' | Set-Content '%PYTHON_DIR%/python310._pth' }"
+	echo Executing: %PS_CMD% "&{(Get-Content '%PYTHON_DIR%/python312._pth') -creplace '#import site', 'import site' | Set-Content '%PYTHON_DIR%/python312._pth' }"
+	%PS_CMD% "&{(Get-Content '%PYTHON_DIR%/python312._pth') -creplace '#import site', 'import site' | Set-Content '%PYTHON_DIR%/python312._pth' }"
 	if !errorlevel! neq 0 ( pause & exit /b !errorlevel! )
 
 	echo --------------------------------------------------

slice.py

@@ -171,7 +171,7 @@ if __name__ == "__main__":
 
     # モデルをダウンロードしておく
     _ = torch.hub.load(
-        repo_or_dir="litagin02/silero-vad",
+        repo_or_dir="snakers4/silero-vad",
         model="silero_vad",
         onnx=True,
         trust_repo=True,
@@ -186,7 +186,7 @@ if __name__ == "__main__":
     ):
         # logger.debug("Worker started.")
         vad_model, utils = torch.hub.load(
-            repo_or_dir="litagin02/silero-vad",
+            repo_or_dir="snakers4/silero-vad",
             model="silero_vad",
             onnx=True,
             trust_repo=True,

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)
 

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)

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,
+    )

style_bert_vits2/models/utils/checkpoints.py

@@ -32,7 +32,9 @@ def load_checkpoint(
     """
 
     assert os.path.isfile(checkpoint_path)
-    checkpoint_dict = torch.load(checkpoint_path, map_location=device)
+    checkpoint_dict = torch.load(
+        checkpoint_path, map_location=device, weights_only=True
+    )
     iteration = checkpoint_dict["iteration"]
     learning_rate = checkpoint_dict["learning_rate"]
     logger.info(

style_bert_vits2/nlp/japanese/user_dict/word_model.py

@@ -9,7 +9,13 @@ from enum import Enum
 from re import findall, fullmatch
 from typing import List, Optional
 
-from pydantic import BaseModel, Field, validator
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+    ValidationInfo,
+    field_validator,
+)
 
 
 USER_DICT_MIN_PRIORITY = 0
@@ -39,11 +45,11 @@ class UserDictWord(BaseModel):
     mora_count: Optional[int] = Field(title="モーラ数", default=None)
     accent_associative_rule: str = Field(title="アクセント結合規則")
 
-    class Config:
-        validate_assignment = True
+    model_config = ConfigDict(validate_assignment=True, validate_default=True)
 
-    @validator("surface")
-    def convert_to_zenkaku(cls, surface):
+    @field_validator("surface")
+    @classmethod
+    def convert_to_zenkaku(cls, surface: str) -> str:
         return surface.translate(
             str.maketrans(
                 "".join(chr(0x21 + i) for i in range(94)),
@@ -51,8 +57,9 @@ class UserDictWord(BaseModel):
             )
         )
 
-    @validator("pronunciation", pre=True)
-    def check_is_katakana(cls, pronunciation):
+    @field_validator("pronunciation", mode="before")
+    @classmethod
+    def check_is_katakana(cls, pronunciation: str) -> str:
         if not fullmatch(r"[ァ-ヴー]+", pronunciation):
             raise ValueError("発音は有効なカタカナでなくてはいけません。")
         sutegana = ["ァ", "ィ", "ゥ", "ェ", "ォ", "ャ", "ュ", "ョ", "ヮ", "ッ"]
@@ -75,9 +82,12 @@ class UserDictWord(BaseModel):
                     )
         return pronunciation
 
-    @validator("mora_count", pre=True, always=True)
-    def check_mora_count_and_accent_type(cls, mora_count, values):
-        if "pronunciation" not in values or "accent_type" not in values:
+    @field_validator("mora_count", mode="before")
+    @classmethod
+    def check_mora_count_and_accent_type(
+        cls, mora_count: Optional[int], info: ValidationInfo
+    ) -> Optional[int]:
+        if "pronunciation" not in info.data or "accent_type" not in info.data:
             # 適切な場所でエラーを出すようにする
             return mora_count
 
@@ -91,14 +101,14 @@ class UserDictWord(BaseModel):
             mora_count = len(
                 findall(
                     f"(?:{rule_others}|{rule_line_i}|{rule_line_u}|{rule_one_mora})",
-                    values["pronunciation"],
+                    info.data["pronunciation"],
                 )
             )
 
-        if not 0 <= values["accent_type"] <= mora_count:
+        if not 0 <= info.data["accent_type"] <= mora_count:
             raise ValueError(
                 "誤ったアクセント型です({})。 expect: 0 <= accent_type <= {}".format(
-                    values["accent_type"], mora_count
+                    info.data["accent_type"], mora_count
                 )
             )
         return mora_count

tests/test_library_compatibility.py

@@ -0,0 +1,58 @@
+from pathlib import Path
+
+import pytest
+import torch
+from pydantic import ValidationError
+
+from style_bert_vits2.models.utils.checkpoints import load_checkpoint, save_checkpoint
+from style_bert_vits2.nlp.japanese.user_dict.word_model import UserDictWord
+
+
+def _user_dict_word(**overrides: object) -> UserDictWord:
+    values = {
+        "surface": "test",
+        "priority": 5,
+        "part_of_speech": "名詞",
+        "part_of_speech_detail_1": "一般",
+        "part_of_speech_detail_2": "*",
+        "part_of_speech_detail_3": "*",
+        "inflectional_type": "*",
+        "inflectional_form": "*",
+        "stem": "*",
+        "yomi": "テスト",
+        "pronunciation": "テスト",
+        "accent_type": 1,
+        "accent_associative_rule": "*",
+    }
+    values.update(overrides)
+    return UserDictWord.model_validate(values)
+
+
+def test_user_dict_word_uses_pydantic_v2_validation() -> None:
+    word = _user_dict_word()
+
+    assert word.surface == "ｔｅｓｔ"
+    assert word.mora_count == 3
+
+    with pytest.raises(ValidationError):
+        _user_dict_word(pronunciation="てすと")
+
+
+def test_checkpoint_round_trip_with_safe_torch_load(tmp_path: Path) -> None:
+    checkpoint_path = tmp_path / "checkpoint.pth"
+    source = torch.nn.Linear(2, 1)
+    optimizer = torch.optim.AdamW(source.parameters())
+    save_checkpoint(source, optimizer, 1e-4, 12, checkpoint_path)
+
+    destination = torch.nn.Linear(2, 1)
+    destination_optimizer = torch.optim.AdamW(destination.parameters())
+    _, _, learning_rate, iteration = load_checkpoint(
+        checkpoint_path, destination, destination_optimizer
+    )
+
+    assert learning_rate == 1e-4
+    assert iteration == 12
+    for source_parameter, destination_parameter in zip(
+        source.parameters(), destination.parameters()
+    ):
+        assert torch.equal(source_parameter, destination_parameter)

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,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()