From e566ac62f373557106dd48af15b4c3deb942c355 Mon Sep 17 00:00:00 2001
From: litagin02 <litagin02@gmail.com>
Date: Fri, 2 Feb 2024 21:11:23 +0900
Subject: [PATCH] Update

---
 .gitignore                                   |    5 +
 configs/config.json                          |    2 +-
 configs/configs_jp_extra.json                |   77 ++
 models_JP_Extra.py                           |   58 +-
 models_V210.py                               | 1043 ------------------
 preprocess_text.py                           |    6 +-
 safetensors.ipynb                            |  115 ++
 slm/wavlm-base-plus/.gitattributes           |   27 +
 slm/wavlm-base-plus/README.md                |   65 ++
 slm/wavlm-base-plus/config.json              |   99 ++
 slm/wavlm-base-plus/preprocessor_config.json |    9 +
 text/cleaner.py                              |    7 +-
 text/japanese.py                             |   47 +-
 text/japanese_mora_list.py                   |    7 +-
 train_ms.py                                  |   19 +-
 train_ms_jp_extra.py                         |  998 +++++++++++++++++
 webui_train.py                               |  113 +-
 17 files changed, 1535 insertions(+), 1162 deletions(-)
 create mode 100644 configs/configs_jp_extra.json
 delete mode 100644 models_V210.py
 create mode 100644 safetensors.ipynb
 create mode 100644 slm/wavlm-base-plus/.gitattributes
 create mode 100644 slm/wavlm-base-plus/README.md
 create mode 100644 slm/wavlm-base-plus/config.json
 create mode 100644 slm/wavlm-base-plus/preprocessor_config.json
 create mode 100644 train_ms_jp_extra.py

diff --git a/.gitignore b/.gitignore
index 9e3cd02..85c0cf2 100644
--- a/.gitignore
+++ b/.gitignore
@@ -15,6 +15,11 @@ venv/
 /pretrained/*.safetensors
 /pretrained/*.pth
 
+/pretrained_jp_extra/*.safetensors
+/pretrained_jp_extra/*.pth
+
+/slm/*/*.bin
+
 /scripts/test/
 *.zip
 *.csv
diff --git a/configs/config.json b/configs/config.json
index 7728a75..e4dd39a 100644
--- a/configs/config.json
+++ b/configs/config.json
@@ -67,5 +67,5 @@
     "use_spectral_norm": false,
     "gin_channels": 256
   },
-  "version": "1.3"
+  "version": "2.0"
 }
diff --git a/configs/configs_jp_extra.json b/configs/configs_jp_extra.json
new file mode 100644
index 0000000..addb364
--- /dev/null
+++ b/configs/configs_jp_extra.json
@@ -0,0 +1,77 @@
+{
+  "train": {
+    "log_interval": 200,
+    "eval_interval": 1000,
+    "seed": 42,
+    "epochs": 1000,
+    "learning_rate": 0.0001,
+    "betas": [0.8, 0.99],
+    "eps": 1e-9,
+    "batch_size": 24,
+    "bf16_run": false,
+    "fp16_run": false,
+    "lr_decay": 0.99996,
+    "segment_size": 16384,
+    "init_lr_ratio": 1,
+    "warmup_epochs": 0,
+    "c_mel": 45,
+    "c_kl": 1.0,
+    "c_commit": 100,
+    "skip_optimizer": true,
+    "freeze_ZH_bert": false,
+    "freeze_JP_bert": false,
+    "freeze_EN_bert": false,
+    "freeze_emo": false,
+    "freeze_style": false
+  },
+  "data": {
+    "training_files": "filelists/train.list",
+    "validation_files": "filelists/val.list",
+    "max_wav_value": 32768.0,
+    "sampling_rate": 44100,
+    "filter_length": 2048,
+    "hop_length": 512,
+    "win_length": 2048,
+    "n_mel_channels": 128,
+    "mel_fmin": 0.0,
+    "mel_fmax": null,
+    "add_blank": true,
+    "n_speakers": 512,
+    "cleaned_text": true
+  },
+  "model": {
+    "use_spk_conditioned_encoder": true,
+    "use_noise_scaled_mas": true,
+    "use_mel_posterior_encoder": false,
+    "use_duration_discriminator": false,
+    "use_wavlm_discriminator": true,
+    "inter_channels": 192,
+    "hidden_channels": 192,
+    "filter_channels": 768,
+    "n_heads": 2,
+    "n_layers": 6,
+    "kernel_size": 3,
+    "p_dropout": 0.1,
+    "resblock": "1",
+    "resblock_kernel_sizes": [3, 7, 11],
+    "resblock_dilation_sizes": [
+      [1, 3, 5],
+      [1, 3, 5],
+      [1, 3, 5]
+    ],
+    "upsample_rates": [8, 8, 2, 2, 2],
+    "upsample_initial_channel": 512,
+    "upsample_kernel_sizes": [16, 16, 8, 2, 2],
+    "n_layers_q": 3,
+    "use_spectral_norm": false,
+    "gin_channels": 512,
+    "slm": {
+      "model": "./slm/wavlm-base-plus",
+      "sr": 16000,
+      "hidden": 768,
+      "nlayers": 13,
+      "initial_channel": 64
+    }
+  },
+  "version": "2.0-JP-Extra"
+}
diff --git a/models_JP_Extra.py b/models_JP_Extra.py
index ed3a2fe..a5d828d 100644
--- a/models_JP_Extra.py
+++ b/models_JP_Extra.py
@@ -362,35 +362,9 @@ class TextEncoder(nn.Module):
         self.language_emb = nn.Embedding(num_languages, hidden_channels)
         nn.init.normal_(self.language_emb.weight, 0.0, hidden_channels**-0.5)
         self.bert_proj = nn.Conv1d(1024, hidden_channels, 1)
-        #self.bert_pre_proj = nn.Conv1d(2048, 1024, 1)
-        # self.en_bert_proj = nn.Conv1d(1024, hidden_channels, 1)
-        self.in_feature_net = nn.Sequential(
-            # input is assumed to an already normalized embedding
-            nn.Linear(512, 1028, bias=False),
-            nn.GELU(),
-            nn.LayerNorm(1028),
-            *[Block(1028, 512) for _ in range(1)],
-            nn.Linear(1028, 512, bias=False),
-            # normalize before passing to VQ?
-            # nn.GELU(),
-            # nn.LayerNorm(512),
-        )
-        self.emo_vq = VectorQuantize(
-            dim=512,
-            # codebook_size=128,
-            codebook_size=256,
-            codebook_dim=16,
-            # codebook_dim=32,
-            commitment_weight=0.1,
-            decay=0.99,
-            heads=32,
-            kmeans_iters=20,
-            separate_codebook_per_head=True,
-            stochastic_sample_codes=True,
-            threshold_ema_dead_code=2,
-            use_cosine_sim = True,
-        )
-        self.out_feature_net = nn.Linear(512, hidden_channels)
+
+        # Remove emo_vq since it's not working well.
+        self.style_proj = nn.Linear(256, hidden_channels)
 
         self.encoder = attentions.Encoder(
             hidden_channels,
@@ -403,20 +377,15 @@ class TextEncoder(nn.Module):
         )
         self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
 
-    def forward(self, x, x_lengths, tone, language, bert, emo, g=None):
+    def forward(self, x, x_lengths, tone, language, bert, style_vec, g=None):
         bert_emb = self.bert_proj(bert).transpose(1, 2)
-        # en_bert_emb = self.en_bert_proj(en_bert).transpose(1, 2)
-        emo_emb = self.in_feature_net(emo)
-        emo_emb, _, loss_commit = self.emo_vq(emo_emb.unsqueeze(1))
-        loss_commit = loss_commit.mean()
-        emo_emb = self.out_feature_net(emo_emb)
+        style_emb = self.style_proj(style_vec.unsqueeze(1))
         x = (
             self.emb(x)
             + self.tone_emb(tone)
             + self.language_emb(language)
             + bert_emb
-            # + en_bert_emb
-            + emo_emb
+            + style_emb
         ) * math.sqrt(
             self.hidden_channels
         )  # [b, t, h]
@@ -429,7 +398,7 @@ class TextEncoder(nn.Module):
         stats = self.proj(x) * x_mask
 
         m, logs = torch.split(stats, self.out_channels, dim=1)
-        return x, m, logs, x_mask, loss_commit
+        return x, m, logs, x_mask
 
 
 class ResidualCouplingBlock(nn.Module):
@@ -976,14 +945,14 @@ class SynthesizerTrn(nn.Module):
         tone,
         language,
         bert,
-        emo,
+        style_vec,
     ):
         if self.n_speakers > 0:
             g = self.emb_g(sid).unsqueeze(-1)  # [b, h, 1]
         else:
             g = self.ref_enc(y.transpose(1, 2)).unsqueeze(-1)
-        x, m_p, logs_p, x_mask, loss_commit = self.enc_p(
-            x, x_lengths, tone, language, bert, emo, g=g
+        x, m_p, logs_p, x_mask = self.enc_p(
+            x, x_lengths, tone, language, bert, style_vec, g=g
         )
         z, m_q, logs_q, y_mask = self.enc_q(y, y_lengths, g=g)
         z_p = self.flow(z, y_mask, g=g)
@@ -1052,7 +1021,6 @@ class SynthesizerTrn(nn.Module):
             (z, z_p, m_p, logs_p, m_q, logs_q),
             (x, logw, logw_),  # , logw_sdp),
             g,
-            loss_commit,
         )
 
     def infer(
@@ -1063,7 +1031,7 @@ class SynthesizerTrn(nn.Module):
         tone,
         language,
         bert,
-        emo,
+        style_vec,
         noise_scale=0.667,
         length_scale=1,
         noise_scale_w=0.8,
@@ -1077,8 +1045,8 @@ class SynthesizerTrn(nn.Module):
             g = self.emb_g(sid).unsqueeze(-1)  # [b, h, 1]
         else:
             g = self.ref_enc(y.transpose(1, 2)).unsqueeze(-1)
-        x, m_p, logs_p, x_mask, _ = self.enc_p(
-            x, x_lengths, tone, language, bert, emo, g=g
+        x, m_p, logs_p, x_mask = self.enc_p(
+            x, x_lengths, tone, language, bert, style_vec, g=g
         )
         logw = self.sdp(x, x_mask, g=g, reverse=True, noise_scale=noise_scale_w) * (
             sdp_ratio
diff --git a/models_V210.py b/models_V210.py
deleted file mode 100644
index 7525a76..0000000
--- a/models_V210.py
+++ /dev/null
@@ -1,1043 +0,0 @@
-"""
-Original models in Bert-VITS2 ver 2.1.
-"""
-import math
-import torch
-from torch import nn
-from torch.nn import functional as F
-
-import commons
-import modules
-import attentions
-import monotonic_align
-
-from torch.nn import Conv1d, ConvTranspose1d, Conv2d
-from torch.nn.utils import weight_norm, remove_weight_norm, spectral_norm
-from vector_quantize_pytorch import VectorQuantize
-
-from commons import init_weights, get_padding
-from .text import symbols, num_tones, num_languages
-
-
-class DurationDiscriminator(nn.Module):  # vits2
-    def __init__(
-        self, in_channels, filter_channels, kernel_size, p_dropout, gin_channels=0
-    ):
-        super().__init__()
-
-        self.in_channels = in_channels
-        self.filter_channels = filter_channels
-        self.kernel_size = kernel_size
-        self.p_dropout = p_dropout
-        self.gin_channels = gin_channels
-
-        self.drop = nn.Dropout(p_dropout)
-        self.conv_1 = nn.Conv1d(
-            in_channels, filter_channels, kernel_size, padding=kernel_size // 2
-        )
-        self.norm_1 = modules.LayerNorm(filter_channels)
-        self.conv_2 = nn.Conv1d(
-            filter_channels, filter_channels, kernel_size, padding=kernel_size // 2
-        )
-        self.norm_2 = modules.LayerNorm(filter_channels)
-        self.dur_proj = nn.Conv1d(1, filter_channels, 1)
-
-        self.pre_out_conv_1 = nn.Conv1d(
-            2 * filter_channels, filter_channels, kernel_size, padding=kernel_size // 2
-        )
-        self.pre_out_norm_1 = modules.LayerNorm(filter_channels)
-        self.pre_out_conv_2 = nn.Conv1d(
-            filter_channels, filter_channels, kernel_size, padding=kernel_size // 2
-        )
-        self.pre_out_norm_2 = modules.LayerNorm(filter_channels)
-
-        if gin_channels != 0:
-            self.cond = nn.Conv1d(gin_channels, in_channels, 1)
-
-        self.output_layer = nn.Sequential(nn.Linear(filter_channels, 1), nn.Sigmoid())
-
-    def forward_probability(self, x, x_mask, dur, g=None):
-        dur = self.dur_proj(dur)
-        x = torch.cat([x, dur], dim=1)
-        x = self.pre_out_conv_1(x * x_mask)
-        x = torch.relu(x)
-        x = self.pre_out_norm_1(x)
-        x = self.drop(x)
-        x = self.pre_out_conv_2(x * x_mask)
-        x = torch.relu(x)
-        x = self.pre_out_norm_2(x)
-        x = self.drop(x)
-        x = x * x_mask
-        x = x.transpose(1, 2)
-        output_prob = self.output_layer(x)
-        return output_prob
-
-    def forward(self, x, x_mask, dur_r, dur_hat, g=None):
-        x = torch.detach(x)
-        if g is not None:
-            g = torch.detach(g)
-            x = x + self.cond(g)
-        x = self.conv_1(x * x_mask)
-        x = torch.relu(x)
-        x = self.norm_1(x)
-        x = self.drop(x)
-        x = self.conv_2(x * x_mask)
-        x = torch.relu(x)
-        x = self.norm_2(x)
-        x = self.drop(x)
-
-        output_probs = []
-        for dur in [dur_r, dur_hat]:
-            output_prob = self.forward_probability(x, x_mask, dur, g)
-            output_probs.append(output_prob)
-
-        return output_probs
-
-
-class TransformerCouplingBlock(nn.Module):
-    def __init__(
-        self,
-        channels,
-        hidden_channels,
-        filter_channels,
-        n_heads,
-        n_layers,
-        kernel_size,
-        p_dropout,
-        n_flows=4,
-        gin_channels=0,
-        share_parameter=False,
-    ):
-        super().__init__()
-        self.channels = channels
-        self.hidden_channels = hidden_channels
-        self.kernel_size = kernel_size
-        self.n_layers = n_layers
-        self.n_flows = n_flows
-        self.gin_channels = gin_channels
-
-        self.flows = nn.ModuleList()
-
-        self.wn = (
-            attentions.FFT(
-                hidden_channels,
-                filter_channels,
-                n_heads,
-                n_layers,
-                kernel_size,
-                p_dropout,
-                isflow=True,
-                gin_channels=self.gin_channels,
-            )
-            if share_parameter
-            else None
-        )
-
-        for i in range(n_flows):
-            self.flows.append(
-                modules.TransformerCouplingLayer(
-                    channels,
-                    hidden_channels,
-                    kernel_size,
-                    n_layers,
-                    n_heads,
-                    p_dropout,
-                    filter_channels,
-                    mean_only=True,
-                    wn_sharing_parameter=self.wn,
-                    gin_channels=self.gin_channels,
-                )
-            )
-            self.flows.append(modules.Flip())
-
-    def forward(self, x, x_mask, g=None, reverse=False):
-        if not reverse:
-            for flow in self.flows:
-                x, _ = flow(x, x_mask, g=g, reverse=reverse)
-        else:
-            for flow in reversed(self.flows):
-                x = flow(x, x_mask, g=g, reverse=reverse)
-        return x
-
-
-class StochasticDurationPredictor(nn.Module):
-    def __init__(
-        self,
-        in_channels,
-        filter_channels,
-        kernel_size,
-        p_dropout,
-        n_flows=4,
-        gin_channels=0,
-    ):
-        super().__init__()
-        filter_channels = in_channels  # it needs to be removed from future version.
-        self.in_channels = in_channels
-        self.filter_channels = filter_channels
-        self.kernel_size = kernel_size
-        self.p_dropout = p_dropout
-        self.n_flows = n_flows
-        self.gin_channels = gin_channels
-
-        self.log_flow = modules.Log()
-        self.flows = nn.ModuleList()
-        self.flows.append(modules.ElementwiseAffine(2))
-        for i in range(n_flows):
-            self.flows.append(
-                modules.ConvFlow(2, filter_channels, kernel_size, n_layers=3)
-            )
-            self.flows.append(modules.Flip())
-
-        self.post_pre = nn.Conv1d(1, filter_channels, 1)
-        self.post_proj = nn.Conv1d(filter_channels, filter_channels, 1)
-        self.post_convs = modules.DDSConv(
-            filter_channels, kernel_size, n_layers=3, p_dropout=p_dropout
-        )
-        self.post_flows = nn.ModuleList()
-        self.post_flows.append(modules.ElementwiseAffine(2))
-        for i in range(4):
-            self.post_flows.append(
-                modules.ConvFlow(2, filter_channels, kernel_size, n_layers=3)
-            )
-            self.post_flows.append(modules.Flip())
-
-        self.pre = nn.Conv1d(in_channels, filter_channels, 1)
-        self.proj = nn.Conv1d(filter_channels, filter_channels, 1)
-        self.convs = modules.DDSConv(
-            filter_channels, kernel_size, n_layers=3, p_dropout=p_dropout
-        )
-        if gin_channels != 0:
-            self.cond = nn.Conv1d(gin_channels, filter_channels, 1)
-
-    def forward(self, x, x_mask, w=None, g=None, reverse=False, noise_scale=1.0):
-        x = torch.detach(x)
-        x = self.pre(x)
-        if g is not None:
-            g = torch.detach(g)
-            x = x + self.cond(g)
-        x = self.convs(x, x_mask)
-        x = self.proj(x) * x_mask
-
-        if not reverse:
-            flows = self.flows
-            assert w is not None
-
-            logdet_tot_q = 0
-            h_w = self.post_pre(w)
-            h_w = self.post_convs(h_w, x_mask)
-            h_w = self.post_proj(h_w) * x_mask
-            e_q = (
-                torch.randn(w.size(0), 2, w.size(2)).to(device=x.device, dtype=x.dtype)
-                * x_mask
-            )
-            z_q = e_q
-            for flow in self.post_flows:
-                z_q, logdet_q = flow(z_q, x_mask, g=(x + h_w))
-                logdet_tot_q += logdet_q
-            z_u, z1 = torch.split(z_q, [1, 1], 1)
-            u = torch.sigmoid(z_u) * x_mask
-            z0 = (w - u) * x_mask
-            logdet_tot_q += torch.sum(
-                (F.logsigmoid(z_u) + F.logsigmoid(-z_u)) * x_mask, [1, 2]
-            )
-            logq = (
-                torch.sum(-0.5 * (math.log(2 * math.pi) + (e_q**2)) * x_mask, [1, 2])
-                - logdet_tot_q
-            )
-
-            logdet_tot = 0
-            z0, logdet = self.log_flow(z0, x_mask)
-            logdet_tot += logdet
-            z = torch.cat([z0, z1], 1)
-            for flow in flows:
-                z, logdet = flow(z, x_mask, g=x, reverse=reverse)
-                logdet_tot = logdet_tot + logdet
-            nll = (
-                torch.sum(0.5 * (math.log(2 * math.pi) + (z**2)) * x_mask, [1, 2])
-                - logdet_tot
-            )
-            return nll + logq  # [b]
-        else:
-            flows = list(reversed(self.flows))
-            flows = flows[:-2] + [flows[-1]]  # remove a useless vflow
-            z = (
-                torch.randn(x.size(0), 2, x.size(2)).to(device=x.device, dtype=x.dtype)
-                * noise_scale
-            )
-            for flow in flows:
-                z = flow(z, x_mask, g=x, reverse=reverse)
-            z0, z1 = torch.split(z, [1, 1], 1)
-            logw = z0
-            return logw
-
-
-class DurationPredictor(nn.Module):
-    def __init__(
-        self, in_channels, filter_channels, kernel_size, p_dropout, gin_channels=0
-    ):
-        super().__init__()
-
-        self.in_channels = in_channels
-        self.filter_channels = filter_channels
-        self.kernel_size = kernel_size
-        self.p_dropout = p_dropout
-        self.gin_channels = gin_channels
-
-        self.drop = nn.Dropout(p_dropout)
-        self.conv_1 = nn.Conv1d(
-            in_channels, filter_channels, kernel_size, padding=kernel_size // 2
-        )
-        self.norm_1 = modules.LayerNorm(filter_channels)
-        self.conv_2 = nn.Conv1d(
-            filter_channels, filter_channels, kernel_size, padding=kernel_size // 2
-        )
-        self.norm_2 = modules.LayerNorm(filter_channels)
-        self.proj = nn.Conv1d(filter_channels, 1, 1)
-
-        if gin_channels != 0:
-            self.cond = nn.Conv1d(gin_channels, in_channels, 1)
-
-    def forward(self, x, x_mask, g=None):
-        x = torch.detach(x)
-        if g is not None:
-            g = torch.detach(g)
-            x = x + self.cond(g)
-        x = self.conv_1(x * x_mask)
-        x = torch.relu(x)
-        x = self.norm_1(x)
-        x = self.drop(x)
-        x = self.conv_2(x * x_mask)
-        x = torch.relu(x)
-        x = self.norm_2(x)
-        x = self.drop(x)
-        x = self.proj(x * x_mask)
-        return x * x_mask
-
-
-class TextEncoder(nn.Module):
-    def __init__(
-        self,
-        n_vocab,
-        out_channels,
-        hidden_channels,
-        filter_channels,
-        n_heads,
-        n_layers,
-        kernel_size,
-        p_dropout,
-        n_speakers,
-        gin_channels=0,
-    ):
-        super().__init__()
-        self.n_vocab = n_vocab
-        self.out_channels = out_channels
-        self.hidden_channels = hidden_channels
-        self.filter_channels = filter_channels
-        self.n_heads = n_heads
-        self.n_layers = n_layers
-        self.kernel_size = kernel_size
-        self.p_dropout = p_dropout
-        self.gin_channels = gin_channels
-        self.emb = nn.Embedding(len(symbols), hidden_channels)
-        nn.init.normal_(self.emb.weight, 0.0, hidden_channels**-0.5)
-        self.tone_emb = nn.Embedding(num_tones, hidden_channels)
-        nn.init.normal_(self.tone_emb.weight, 0.0, hidden_channels**-0.5)
-        self.language_emb = nn.Embedding(num_languages, hidden_channels)
-        nn.init.normal_(self.language_emb.weight, 0.0, hidden_channels**-0.5)
-        self.bert_proj = nn.Conv1d(1024, hidden_channels, 1)
-        self.ja_bert_proj = nn.Conv1d(1024, hidden_channels, 1)
-        self.en_bert_proj = nn.Conv1d(1024, hidden_channels, 1)
-        self.emo_proj = nn.Linear(1024, 1024)
-        self.emo_quantizer = VectorQuantize(
-            dim=1024,
-            codebook_size=10,
-            decay=0.8,
-            commitment_weight=1.0,
-            learnable_codebook=True,
-            ema_update=False,
-        )
-        self.emo_q_proj = nn.Linear(1024, hidden_channels)
-
-        self.encoder = attentions.Encoder(
-            hidden_channels,
-            filter_channels,
-            n_heads,
-            n_layers,
-            kernel_size,
-            p_dropout,
-            gin_channels=self.gin_channels,
-        )
-        self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
-
-    def forward(
-        self, x, x_lengths, tone, language, bert, ja_bert, en_bert, emo, sid, g=None
-    ):
-        bert_emb = self.bert_proj(bert).transpose(1, 2)
-        ja_bert_emb = self.ja_bert_proj(ja_bert).transpose(1, 2)
-        en_bert_emb = self.en_bert_proj(en_bert).transpose(1, 2)
-        if emo.size(-1) == 1024:
-            emo_emb = self.emo_proj(emo.unsqueeze(1))
-            emo_commit_loss = torch.zeros(1).to(emo_emb.device)
-            emo_emb_ = []
-            for i in range(emo_emb.size(0)):
-                temp_emo_emb, _, temp_emo_commit_loss = self.emo_quantizer(
-                    emo_emb[i].unsqueeze(0)
-                )
-                emo_commit_loss += temp_emo_commit_loss
-                emo_emb_.append(temp_emo_emb)
-            emo_emb = torch.cat(emo_emb_, dim=0).to(emo_emb.device)
-            emo_commit_loss = emo_commit_loss.to(emo_emb.device)
-        else:
-            emo_emb = (
-                self.emo_quantizer.get_output_from_indices(emo.to(torch.int))
-                .unsqueeze(0)
-                .to(emo.device)
-            )
-            emo_commit_loss = torch.zeros(1)
-        x = (
-            self.emb(x)
-            + self.tone_emb(tone)
-            + self.language_emb(language)
-            + bert_emb
-            + ja_bert_emb
-            + en_bert_emb
-            + self.emo_q_proj(emo_emb)
-        ) * math.sqrt(
-            self.hidden_channels
-        )  # [b, t, h]
-        x = torch.transpose(x, 1, -1)  # [b, h, t]
-        x_mask = torch.unsqueeze(commons.sequence_mask(x_lengths, x.size(2)), 1).to(
-            x.dtype
-        )
-
-        x = self.encoder(x * x_mask, x_mask, g=g)
-        stats = self.proj(x) * x_mask
-
-        m, logs = torch.split(stats, self.out_channels, dim=1)
-        return x, m, logs, x_mask, emo_commit_loss
-
-
-class ResidualCouplingBlock(nn.Module):
-    def __init__(
-        self,
-        channels,
-        hidden_channels,
-        kernel_size,
-        dilation_rate,
-        n_layers,
-        n_flows=4,
-        gin_channels=0,
-    ):
-        super().__init__()
-        self.channels = channels
-        self.hidden_channels = hidden_channels
-        self.kernel_size = kernel_size
-        self.dilation_rate = dilation_rate
-        self.n_layers = n_layers
-        self.n_flows = n_flows
-        self.gin_channels = gin_channels
-
-        self.flows = nn.ModuleList()
-        for i in range(n_flows):
-            self.flows.append(
-                modules.ResidualCouplingLayer(
-                    channels,
-                    hidden_channels,
-                    kernel_size,
-                    dilation_rate,
-                    n_layers,
-                    gin_channels=gin_channels,
-                    mean_only=True,
-                )
-            )
-            self.flows.append(modules.Flip())
-
-    def forward(self, x, x_mask, g=None, reverse=False):
-        if not reverse:
-            for flow in self.flows:
-                x, _ = flow(x, x_mask, g=g, reverse=reverse)
-        else:
-            for flow in reversed(self.flows):
-                x = flow(x, x_mask, g=g, reverse=reverse)
-        return x
-
-
-class PosteriorEncoder(nn.Module):
-    def __init__(
-        self,
-        in_channels,
-        out_channels,
-        hidden_channels,
-        kernel_size,
-        dilation_rate,
-        n_layers,
-        gin_channels=0,
-    ):
-        super().__init__()
-        self.in_channels = in_channels
-        self.out_channels = out_channels
-        self.hidden_channels = hidden_channels
-        self.kernel_size = kernel_size
-        self.dilation_rate = dilation_rate
-        self.n_layers = n_layers
-        self.gin_channels = gin_channels
-
-        self.pre = nn.Conv1d(in_channels, hidden_channels, 1)
-        self.enc = modules.WN(
-            hidden_channels,
-            kernel_size,
-            dilation_rate,
-            n_layers,
-            gin_channels=gin_channels,
-        )
-        self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
-
-    def forward(self, x, x_lengths, g=None):
-        x_mask = torch.unsqueeze(commons.sequence_mask(x_lengths, x.size(2)), 1).to(
-            x.dtype
-        )
-        x = self.pre(x) * x_mask
-        x = self.enc(x, x_mask, g=g)
-        stats = self.proj(x) * x_mask
-        m, logs = torch.split(stats, self.out_channels, dim=1)
-        z = (m + torch.randn_like(m) * torch.exp(logs)) * x_mask
-        return z, m, logs, x_mask
-
-
-class Generator(torch.nn.Module):
-    def __init__(
-        self,
-        initial_channel,
-        resblock,
-        resblock_kernel_sizes,
-        resblock_dilation_sizes,
-        upsample_rates,
-        upsample_initial_channel,
-        upsample_kernel_sizes,
-        gin_channels=0,
-    ):
-        super(Generator, self).__init__()
-        self.num_kernels = len(resblock_kernel_sizes)
-        self.num_upsamples = len(upsample_rates)
-        self.conv_pre = Conv1d(
-            initial_channel, upsample_initial_channel, 7, 1, padding=3
-        )
-        resblock = modules.ResBlock1 if resblock == "1" else modules.ResBlock2
-
-        self.ups = nn.ModuleList()
-        for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
-            self.ups.append(
-                weight_norm(
-                    ConvTranspose1d(
-                        upsample_initial_channel // (2**i),
-                        upsample_initial_channel // (2 ** (i + 1)),
-                        k,
-                        u,
-                        padding=(k - u) // 2,
-                    )
-                )
-            )
-
-        self.resblocks = nn.ModuleList()
-        for i in range(len(self.ups)):
-            ch = upsample_initial_channel // (2 ** (i + 1))
-            for j, (k, d) in enumerate(
-                zip(resblock_kernel_sizes, resblock_dilation_sizes)
-            ):
-                self.resblocks.append(resblock(ch, k, d))
-
-        self.conv_post = Conv1d(ch, 1, 7, 1, padding=3, bias=False)
-        self.ups.apply(init_weights)
-
-        if gin_channels != 0:
-            self.cond = nn.Conv1d(gin_channels, upsample_initial_channel, 1)
-
-    def forward(self, x, g=None):
-        x = self.conv_pre(x)
-        if g is not None:
-            x = x + self.cond(g)
-
-        for i in range(self.num_upsamples):
-            x = F.leaky_relu(x, modules.LRELU_SLOPE)
-            x = self.ups[i](x)
-            xs = None
-            for j in range(self.num_kernels):
-                if xs is None:
-                    xs = self.resblocks[i * self.num_kernels + j](x)
-                else:
-                    xs += self.resblocks[i * self.num_kernels + j](x)
-            x = xs / self.num_kernels
-        x = F.leaky_relu(x)
-        x = self.conv_post(x)
-        x = torch.tanh(x)
-
-        return x
-
-    def remove_weight_norm(self):
-        print("Removing weight norm...")
-        for layer in self.ups:
-            remove_weight_norm(layer)
-        for layer in self.resblocks:
-            layer.remove_weight_norm()
-
-
-class DiscriminatorP(torch.nn.Module):
-    def __init__(self, period, kernel_size=5, stride=3, use_spectral_norm=False):
-        super(DiscriminatorP, self).__init__()
-        self.period = period
-        self.use_spectral_norm = use_spectral_norm
-        norm_f = weight_norm if use_spectral_norm is False else spectral_norm
-        self.convs = nn.ModuleList(
-            [
-                norm_f(
-                    Conv2d(
-                        1,
-                        32,
-                        (kernel_size, 1),
-                        (stride, 1),
-                        padding=(get_padding(kernel_size, 1), 0),
-                    )
-                ),
-                norm_f(
-                    Conv2d(
-                        32,
-                        128,
-                        (kernel_size, 1),
-                        (stride, 1),
-                        padding=(get_padding(kernel_size, 1), 0),
-                    )
-                ),
-                norm_f(
-                    Conv2d(
-                        128,
-                        512,
-                        (kernel_size, 1),
-                        (stride, 1),
-                        padding=(get_padding(kernel_size, 1), 0),
-                    )
-                ),
-                norm_f(
-                    Conv2d(
-                        512,
-                        1024,
-                        (kernel_size, 1),
-                        (stride, 1),
-                        padding=(get_padding(kernel_size, 1), 0),
-                    )
-                ),
-                norm_f(
-                    Conv2d(
-                        1024,
-                        1024,
-                        (kernel_size, 1),
-                        1,
-                        padding=(get_padding(kernel_size, 1), 0),
-                    )
-                ),
-            ]
-        )
-        self.conv_post = norm_f(Conv2d(1024, 1, (3, 1), 1, padding=(1, 0)))
-
-    def forward(self, x):
-        fmap = []
-
-        # 1d to 2d
-        b, c, t = x.shape
-        if t % self.period != 0:  # pad first
-            n_pad = self.period - (t % self.period)
-            x = F.pad(x, (0, n_pad), "reflect")
-            t = t + n_pad
-        x = x.view(b, c, t // self.period, self.period)
-
-        for layer in self.convs:
-            x = layer(x)
-            x = F.leaky_relu(x, modules.LRELU_SLOPE)
-            fmap.append(x)
-        x = self.conv_post(x)
-        fmap.append(x)
-        x = torch.flatten(x, 1, -1)
-
-        return x, fmap
-
-
-class DiscriminatorS(torch.nn.Module):
-    def __init__(self, use_spectral_norm=False):
-        super(DiscriminatorS, self).__init__()
-        norm_f = weight_norm if use_spectral_norm is False else spectral_norm
-        self.convs = nn.ModuleList(
-            [
-                norm_f(Conv1d(1, 16, 15, 1, padding=7)),
-                norm_f(Conv1d(16, 64, 41, 4, groups=4, padding=20)),
-                norm_f(Conv1d(64, 256, 41, 4, groups=16, padding=20)),
-                norm_f(Conv1d(256, 1024, 41, 4, groups=64, padding=20)),
-                norm_f(Conv1d(1024, 1024, 41, 4, groups=256, padding=20)),
-                norm_f(Conv1d(1024, 1024, 5, 1, padding=2)),
-            ]
-        )
-        self.conv_post = norm_f(Conv1d(1024, 1, 3, 1, padding=1))
-
-    def forward(self, x):
-        fmap = []
-
-        for layer in self.convs:
-            x = layer(x)
-            x = F.leaky_relu(x, modules.LRELU_SLOPE)
-            fmap.append(x)
-        x = self.conv_post(x)
-        fmap.append(x)
-        x = torch.flatten(x, 1, -1)
-
-        return x, fmap
-
-
-class MultiPeriodDiscriminator(torch.nn.Module):
-    def __init__(self, use_spectral_norm=False):
-        super(MultiPeriodDiscriminator, self).__init__()
-        periods = [2, 3, 5, 7, 11]
-
-        discs = [DiscriminatorS(use_spectral_norm=use_spectral_norm)]
-        discs = discs + [
-            DiscriminatorP(i, use_spectral_norm=use_spectral_norm) for i in periods
-        ]
-        self.discriminators = nn.ModuleList(discs)
-
-    def forward(self, y, y_hat):
-        y_d_rs = []
-        y_d_gs = []
-        fmap_rs = []
-        fmap_gs = []
-        for i, d in enumerate(self.discriminators):
-            y_d_r, fmap_r = d(y)
-            y_d_g, fmap_g = d(y_hat)
-            y_d_rs.append(y_d_r)
-            y_d_gs.append(y_d_g)
-            fmap_rs.append(fmap_r)
-            fmap_gs.append(fmap_g)
-
-        return y_d_rs, y_d_gs, fmap_rs, fmap_gs
-
-
-class ReferenceEncoder(nn.Module):
-    """
-    inputs --- [N, Ty/r, n_mels*r]  mels
-    outputs --- [N, ref_enc_gru_size]
-    """
-
-    def __init__(self, spec_channels, gin_channels=0):
-        super().__init__()
-        self.spec_channels = spec_channels
-        ref_enc_filters = [32, 32, 64, 64, 128, 128]
-        K = len(ref_enc_filters)
-        filters = [1] + ref_enc_filters
-        convs = [
-            weight_norm(
-                nn.Conv2d(
-                    in_channels=filters[i],
-                    out_channels=filters[i + 1],
-                    kernel_size=(3, 3),
-                    stride=(2, 2),
-                    padding=(1, 1),
-                )
-            )
-            for i in range(K)
-        ]
-        self.convs = nn.ModuleList(convs)
-        # self.wns = nn.ModuleList([weight_norm(num_features=ref_enc_filters[i]) for i in range(K)]) # noqa: E501
-
-        out_channels = self.calculate_channels(spec_channels, 3, 2, 1, K)
-        self.gru = nn.GRU(
-            input_size=ref_enc_filters[-1] * out_channels,
-            hidden_size=256 // 2,
-            batch_first=True,
-        )
-        self.proj = nn.Linear(128, gin_channels)
-
-    def forward(self, inputs, mask=None):
-        N = inputs.size(0)
-        out = inputs.view(N, 1, -1, self.spec_channels)  # [N, 1, Ty, n_freqs]
-        for conv in self.convs:
-            out = conv(out)
-            # out = wn(out)
-            out = F.relu(out)  # [N, 128, Ty//2^K, n_mels//2^K]
-
-        out = out.transpose(1, 2)  # [N, Ty//2^K, 128, n_mels//2^K]
-        T = out.size(1)
-        N = out.size(0)
-        out = out.contiguous().view(N, T, -1)  # [N, Ty//2^K, 128*n_mels//2^K]
-
-        self.gru.flatten_parameters()
-        memory, out = self.gru(out)  # out --- [1, N, 128]
-
-        return self.proj(out.squeeze(0))
-
-    def calculate_channels(self, L, kernel_size, stride, pad, n_convs):
-        for i in range(n_convs):
-            L = (L - kernel_size + 2 * pad) // stride + 1
-        return L
-
-
-class SynthesizerTrn(nn.Module):
-    """
-    Synthesizer for Training
-    """
-
-    def __init__(
-        self,
-        n_vocab,
-        spec_channels,
-        segment_size,
-        inter_channels,
-        hidden_channels,
-        filter_channels,
-        n_heads,
-        n_layers,
-        kernel_size,
-        p_dropout,
-        resblock,
-        resblock_kernel_sizes,
-        resblock_dilation_sizes,
-        upsample_rates,
-        upsample_initial_channel,
-        upsample_kernel_sizes,
-        n_speakers=256,
-        gin_channels=256,
-        use_sdp=True,
-        n_flow_layer=4,
-        n_layers_trans_flow=4,
-        flow_share_parameter=False,
-        use_transformer_flow=True,
-        **kwargs
-    ):
-        super().__init__()
-        self.n_vocab = n_vocab
-        self.spec_channels = spec_channels
-        self.inter_channels = inter_channels
-        self.hidden_channels = hidden_channels
-        self.filter_channels = filter_channels
-        self.n_heads = n_heads
-        self.n_layers = n_layers
-        self.kernel_size = kernel_size
-        self.p_dropout = p_dropout
-        self.resblock = resblock
-        self.resblock_kernel_sizes = resblock_kernel_sizes
-        self.resblock_dilation_sizes = resblock_dilation_sizes
-        self.upsample_rates = upsample_rates
-        self.upsample_initial_channel = upsample_initial_channel
-        self.upsample_kernel_sizes = upsample_kernel_sizes
-        self.segment_size = segment_size
-        self.n_speakers = n_speakers
-        self.gin_channels = gin_channels
-        self.n_layers_trans_flow = n_layers_trans_flow
-        self.use_spk_conditioned_encoder = kwargs.get(
-            "use_spk_conditioned_encoder", True
-        )
-        self.use_sdp = use_sdp
-        self.use_noise_scaled_mas = kwargs.get("use_noise_scaled_mas", False)
-        self.mas_noise_scale_initial = kwargs.get("mas_noise_scale_initial", 0.01)
-        self.noise_scale_delta = kwargs.get("noise_scale_delta", 2e-6)
-        self.current_mas_noise_scale = self.mas_noise_scale_initial
-        if self.use_spk_conditioned_encoder and gin_channels > 0:
-            self.enc_gin_channels = gin_channels
-        self.enc_p = TextEncoder(
-            n_vocab,
-            inter_channels,
-            hidden_channels,
-            filter_channels,
-            n_heads,
-            n_layers,
-            kernel_size,
-            p_dropout,
-            self.n_speakers,
-            gin_channels=self.enc_gin_channels,
-        )
-        self.dec = Generator(
-            inter_channels,
-            resblock,
-            resblock_kernel_sizes,
-            resblock_dilation_sizes,
-            upsample_rates,
-            upsample_initial_channel,
-            upsample_kernel_sizes,
-            gin_channels=gin_channels,
-        )
-        self.enc_q = PosteriorEncoder(
-            spec_channels,
-            inter_channels,
-            hidden_channels,
-            5,
-            1,
-            16,
-            gin_channels=gin_channels,
-        )
-        if use_transformer_flow:
-            self.flow = TransformerCouplingBlock(
-                inter_channels,
-                hidden_channels,
-                filter_channels,
-                n_heads,
-                n_layers_trans_flow,
-                5,
-                p_dropout,
-                n_flow_layer,
-                gin_channels=gin_channels,
-                share_parameter=flow_share_parameter,
-            )
-        else:
-            self.flow = ResidualCouplingBlock(
-                inter_channels,
-                hidden_channels,
-                5,
-                1,
-                n_flow_layer,
-                gin_channels=gin_channels,
-            )
-        self.sdp = StochasticDurationPredictor(
-            hidden_channels, 192, 3, 0.5, 4, gin_channels=gin_channels
-        )
-        self.dp = DurationPredictor(
-            hidden_channels, 256, 3, 0.5, gin_channels=gin_channels
-        )
-
-        if n_speakers >= 1:
-            self.emb_g = nn.Embedding(n_speakers, gin_channels)
-        else:
-            self.ref_enc = ReferenceEncoder(spec_channels, gin_channels)
-
-    def forward(
-        self,
-        x,
-        x_lengths,
-        y,
-        y_lengths,
-        sid,
-        tone,
-        language,
-        bert,
-        ja_bert,
-        en_bert,
-        emo=None,
-    ):
-        if self.n_speakers > 0:
-            g = self.emb_g(sid).unsqueeze(-1)  # [b, h, 1]
-        else:
-            g = self.ref_enc(y.transpose(1, 2)).unsqueeze(-1)
-        x, m_p, logs_p, x_mask, loss_commit = self.enc_p(
-            x, x_lengths, tone, language, bert, ja_bert, en_bert, emo, sid, g=g
-        )
-        z, m_q, logs_q, y_mask = self.enc_q(y, y_lengths, g=g)
-        z_p = self.flow(z, y_mask, g=g)
-
-        with torch.no_grad():
-            # negative cross-entropy
-            s_p_sq_r = torch.exp(-2 * logs_p)  # [b, d, t]
-            neg_cent1 = torch.sum(
-                -0.5 * math.log(2 * math.pi) - logs_p, [1], keepdim=True
-            )  # [b, 1, t_s]
-            neg_cent2 = torch.matmul(
-                -0.5 * (z_p**2).transpose(1, 2), s_p_sq_r
-            )  # [b, t_t, d] x [b, d, t_s] = [b, t_t, t_s]
-            neg_cent3 = torch.matmul(
-                z_p.transpose(1, 2), (m_p * s_p_sq_r)
-            )  # [b, t_t, d] x [b, d, t_s] = [b, t_t, t_s]
-            neg_cent4 = torch.sum(
-                -0.5 * (m_p**2) * s_p_sq_r, [1], keepdim=True
-            )  # [b, 1, t_s]
-            neg_cent = neg_cent1 + neg_cent2 + neg_cent3 + neg_cent4
-            if self.use_noise_scaled_mas:
-                epsilon = (
-                    torch.std(neg_cent)
-                    * torch.randn_like(neg_cent)
-                    * self.current_mas_noise_scale
-                )
-                neg_cent = neg_cent + epsilon
-
-            attn_mask = torch.unsqueeze(x_mask, 2) * torch.unsqueeze(y_mask, -1)
-            attn = (
-                monotonic_align.maximum_path(neg_cent, attn_mask.squeeze(1))
-                .unsqueeze(1)
-                .detach()
-            )
-
-        w = attn.sum(2)
-
-        l_length_sdp = self.sdp(x, x_mask, w, g=g)
-        l_length_sdp = l_length_sdp / torch.sum(x_mask)
-
-        logw_ = torch.log(w + 1e-6) * x_mask
-        logw = self.dp(x, x_mask, g=g)
-        l_length_dp = torch.sum((logw - logw_) ** 2, [1, 2]) / torch.sum(
-            x_mask
-        )  # for averaging
-
-        l_length = l_length_dp + l_length_sdp
-
-        # expand prior
-        m_p = torch.matmul(attn.squeeze(1), m_p.transpose(1, 2)).transpose(1, 2)
-        logs_p = torch.matmul(attn.squeeze(1), logs_p.transpose(1, 2)).transpose(1, 2)
-
-        z_slice, ids_slice = commons.rand_slice_segments(
-            z, y_lengths, self.segment_size
-        )
-        o = self.dec(z_slice, g=g)
-        return (
-            o,
-            l_length,
-            attn,
-            ids_slice,
-            x_mask,
-            y_mask,
-            (z, z_p, m_p, logs_p, m_q, logs_q),
-            (x, logw, logw_),
-            loss_commit,
-        )
-
-    def infer(
-        self,
-        x,
-        x_lengths,
-        sid,
-        tone,
-        language,
-        bert,
-        ja_bert,
-        en_bert,
-        emo=None,
-        noise_scale=0.667,
-        length_scale=1,
-        noise_scale_w=0.8,
-        max_len=None,
-        sdp_ratio=0,
-        y=None,
-    ):
-        # x, m_p, logs_p, x_mask = self.enc_p(x, x_lengths, tone, language, bert)
-        # g = self.gst(y)
-        if self.n_speakers > 0:
-            g = self.emb_g(sid).unsqueeze(-1)  # [b, h, 1]
-        else:
-            g = self.ref_enc(y.transpose(1, 2)).unsqueeze(-1)
-        x, m_p, logs_p, x_mask, _ = self.enc_p(
-            x, x_lengths, tone, language, bert, ja_bert, en_bert, emo, sid, g=g
-        )
-        logw = self.sdp(x, x_mask, g=g, reverse=True, noise_scale=noise_scale_w) * (
-            sdp_ratio
-        ) + self.dp(x, x_mask, g=g) * (1 - sdp_ratio)
-        w = torch.exp(logw) * x_mask * length_scale
-        w_ceil = torch.ceil(w)
-        y_lengths = torch.clamp_min(torch.sum(w_ceil, [1, 2]), 1).long()
-        y_mask = torch.unsqueeze(commons.sequence_mask(y_lengths, None), 1).to(
-            x_mask.dtype
-        )
-        attn_mask = torch.unsqueeze(x_mask, 2) * torch.unsqueeze(y_mask, -1)
-        attn = commons.generate_path(w_ceil, attn_mask)
-
-        m_p = torch.matmul(attn.squeeze(1), m_p.transpose(1, 2)).transpose(
-            1, 2
-        )  # [b, t', t], [b, t, d] -> [b, d, t']
-        logs_p = torch.matmul(attn.squeeze(1), logs_p.transpose(1, 2)).transpose(
-            1, 2
-        )  # [b, t', t], [b, t, d] -> [b, d, t']
-
-        z_p = m_p + torch.randn_like(m_p) * torch.exp(logs_p) * noise_scale
-        z = self.flow(z_p, y_mask, g=g, reverse=True)
-        o = self.dec((z * y_mask)[:, :, :max_len], g=g)
-        return o, attn, y_mask, (z, z_p, m_p, logs_p)
diff --git a/preprocess_text.py b/preprocess_text.py
index d205f85..9fba5eb 100644
--- a/preprocess_text.py
+++ b/preprocess_text.py
@@ -33,6 +33,7 @@ preprocess_text_config = config.preprocess_text_config
 @click.option("--max-val-total", default=preprocess_text_config.max_val_total)
 @click.option("--clean/--no-clean", default=preprocess_text_config.clean)
 @click.option("-y", "--yml_config")
+@click.option("--use_jp_extra", is_flag=True)
 def preprocess(
     transcription_path: str,
     cleaned_path: Optional[str],
@@ -43,6 +44,7 @@ def preprocess(
     max_val_total: int,
     clean: bool,
     yml_config: str,  # 这个不要删
+    use_jp_extra: bool,
 ):
     if cleaned_path == "" or cleaned_path is None:
         cleaned_path = transcription_path + ".cleaned"
@@ -53,7 +55,9 @@ def preprocess(
                 for line in tqdm(trans_file, file=SAFE_STDOUT):
                     try:
                         utt, spk, language, text = line.strip().split("|")
-                        norm_text, phones, tones, word2ph = clean_text(text, language)
+                        norm_text, phones, tones, word2ph = clean_text(
+                            text, language, use_jp_extra
+                        )
                         out_file.write(
                             "{}|{}|{}|{}|{}|{}|{}\n".format(
                                 utt,
diff --git a/safetensors.ipynb b/safetensors.ipynb
new file mode 100644
index 0000000..6722654
--- /dev/null
+++ b/safetensors.ipynb
@@ -0,0 +1,115 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "\n",
+    "model_dir = \"pretrained-Japanese-Extra\"\n",
+    "step = 0\n",
+    "\n",
+    "g_model = torch.load(f\"{model_dir}/G_{step}.pth\", map_location=\"cpu\")\n",
+    "d_model = torch.load(f\"{model_dir}/D_{step}.pth\", map_location=\"cpu\")\n",
+    "wd_model = torch.load(f\"{model_dir}/WD_{step}.pth\", map_location=\"cpu\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "skip enc_p.in_feature_net.0.weight\n",
+      "skip enc_p.in_feature_net.2.weight\n",
+      "skip enc_p.in_feature_net.2.bias\n",
+      "skip enc_p.in_feature_net.3.norm.weight\n",
+      "skip enc_p.in_feature_net.3.norm.bias\n",
+      "skip enc_p.in_feature_net.3.mlp.c_fc1.weight\n",
+      "skip enc_p.in_feature_net.3.mlp.c_fc2.weight\n",
+      "skip enc_p.in_feature_net.3.mlp.c_proj.weight\n",
+      "skip enc_p.in_feature_net.4.weight\n",
+      "skip enc_p.emo_vq._codebook.initted\n",
+      "skip enc_p.emo_vq._codebook.cluster_size\n",
+      "skip enc_p.emo_vq._codebook.embed_avg\n",
+      "skip enc_p.emo_vq._codebook.embed\n",
+      "skip enc_p.out_feature_net.weight\n",
+      "skip enc_p.out_feature_net.bias\n"
+     ]
+    }
+   ],
+   "source": [
+    "g_dict = {}\n",
+    "skip_list = [\"enc_p.in_feature_net.\", \"enc_p.emo_vq.\", \"enc_p.out_feature_net.\"]\n",
+    "for key in g_model[\"model\"].keys():\n",
+    "    if any([key.startswith(s) for s in skip_list]):\n",
+    "        print(f\"skip {key}\")\n",
+    "        continue\n",
+    "    else:\n",
+    "        g_dict[key] = g_model[\"model\"][key]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "d_dict = {}\n",
+    "for key in d_model[\"model\"].keys():\n",
+    "    d_dict[key] = d_model[\"model\"][key]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "wd_dict = {}\n",
+    "for key in wd_model[\"model\"].keys():\n",
+    "    wd_dict[key] = wd_model[\"model\"][key]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from safetensors.torch import save_file\n",
+    "\n",
+    "\n",
+    "save_file(g_dict, f\"G_0.safetensors\")\n",
+    "save_file (d_dict, f\"D_0.safetensors\")\n",
+    "save_file (wd_dict, f\"WD_0.safetensors\")"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "venv",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/slm/wavlm-base-plus/.gitattributes b/slm/wavlm-base-plus/.gitattributes
new file mode 100644
index 0000000..957b257
--- /dev/null
+++ b/slm/wavlm-base-plus/.gitattributes
@@ -0,0 +1,27 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bin.* filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zstandard filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
diff --git a/slm/wavlm-base-plus/README.md b/slm/wavlm-base-plus/README.md
new file mode 100644
index 0000000..541b2e0
--- /dev/null
+++ b/slm/wavlm-base-plus/README.md
@@ -0,0 +1,65 @@
+---
+language:
+- en
+datasets:
+tags:
+- speech
+inference: false
+---
+
+# WavLM-Base-Plus
+
+[Microsoft's WavLM](https://github.com/microsoft/unilm/tree/master/wavlm)
+
+The base model pretrained on 16kHz sampled speech audio. When using the model, make sure that your speech input is also sampled at 16kHz.
+
+**Note**: This model does not have a tokenizer as it was pretrained on audio alone. In order to use this model **speech recognition**, a tokenizer should be created and the model should be fine-tuned on labeled text data. Check out [this blog](https://huggingface.co/blog/fine-tune-wav2vec2-english) for more in-detail explanation of how to fine-tune the model.
+
+The model was pre-trained on:
+
+- 60,000 hours of [Libri-Light](https://arxiv.org/abs/1912.07875)
+- 10,000 hours of [GigaSpeech](https://arxiv.org/abs/2106.06909)
+- 24,000 hours of [VoxPopuli](https://arxiv.org/abs/2101.00390)
+
+[Paper: WavLM: Large-Scale Self-Supervised Pre-Training for Full Stack Speech Processing](https://arxiv.org/abs/2110.13900)
+
+Authors: Sanyuan Chen, Chengyi Wang, Zhengyang Chen, Yu Wu, Shujie Liu, Zhuo Chen, Jinyu Li, Naoyuki Kanda, Takuya Yoshioka, Xiong Xiao, Jian Wu, Long Zhou, Shuo Ren, Yanmin Qian, Yao Qian, Jian Wu, Michael Zeng, Furu Wei
+
+**Abstract**
+*Self-supervised learning (SSL) achieves great success in speech recognition, while limited exploration has been attempted for other speech processing tasks. As speech signal contains multi-faceted information including speaker identity, paralinguistics, spoken content, etc., learning universal representations for all speech tasks is challenging. In this paper, we propose a new pre-trained model, WavLM, to solve full-stack downstream speech tasks. WavLM is built based on the HuBERT framework, with an emphasis on both spoken content modeling and speaker identity preservation. We first equip the Transformer structure with gated relative position bias to improve its capability on recognition tasks. For better speaker discrimination, we propose an utterance mixing training strategy, where additional overlapped utterances are created unsupervisely and incorporated during model training. Lastly, we scale up the training dataset from 60k hours to 94k hours. WavLM Large achieves state-of-the-art performance on the SUPERB benchmark, and brings significant improvements for various speech processing tasks on their representative benchmarks.*
+
+The original model can be found under https://github.com/microsoft/unilm/tree/master/wavlm.
+
+# Usage
+
+This is an English pre-trained speech model that has to be fine-tuned on a downstream task like speech recognition or audio classification before it can be
+used in inference. The model was pre-trained in English and should therefore perform well only in English. The model has been shown to work well on the [SUPERB benchmark](https://superbbenchmark.org/).
+
+**Note**: The model was pre-trained on phonemes rather than characters. This means that one should make sure that the input text is converted to a sequence
+of phonemes before fine-tuning.
+
+## Speech Recognition
+
+To fine-tune the model for speech recognition, see [the official speech recognition example](https://github.com/huggingface/transformers/tree/master/examples/pytorch/speech-recognition).
+
+## Speech Classification
+
+To fine-tune the model for speech classification, see [the official audio classification example](https://github.com/huggingface/transformers/tree/master/examples/pytorch/audio-classification).
+
+## Speaker Verification
+
+TODO
+
+## Speaker Diarization
+
+TODO
+
+# Contribution
+
+The model was contributed by [cywang](https://huggingface.co/cywang) and [patrickvonplaten](https://huggingface.co/patrickvonplaten).
+
+# License
+
+The official license can be found [here](https://github.com/microsoft/UniSpeech/blob/main/LICENSE)
+
+![design](https://raw.githubusercontent.com/patrickvonplaten/scientific_images/master/wavlm.png)
diff --git a/slm/wavlm-base-plus/config.json b/slm/wavlm-base-plus/config.json
new file mode 100644
index 0000000..b7b4e5f
--- /dev/null
+++ b/slm/wavlm-base-plus/config.json
@@ -0,0 +1,99 @@
+{
+  "_name_or_path": "wavlm-base-plus",
+  "activation_dropout": 0.0,
+  "adapter_kernel_size": 3,
+  "adapter_stride": 2,
+  "add_adapter": false,
+  "apply_spec_augment": true,
+  "architectures": [
+    "WavLMModel"
+  ],
+  "attention_dropout": 0.1,
+  "bos_token_id": 1,
+  "classifier_proj_size": 256,
+  "codevector_dim": 256,
+  "contrastive_logits_temperature": 0.1,
+  "conv_bias": false,
+  "conv_dim": [
+    512,
+    512,
+    512,
+    512,
+    512,
+    512,
+    512
+  ],
+  "conv_kernel": [
+    10,
+    3,
+    3,
+    3,
+    3,
+    2,
+    2
+  ],
+  "conv_stride": [
+    5,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2
+  ],
+  "ctc_loss_reduction": "sum",
+  "ctc_zero_infinity": false,
+  "diversity_loss_weight": 0.1,
+  "do_stable_layer_norm": false,
+  "eos_token_id": 2,
+  "feat_extract_activation": "gelu",
+  "feat_extract_norm": "group",
+  "feat_proj_dropout": 0.1,
+  "feat_quantizer_dropout": 0.0,
+  "final_dropout": 0.0,
+  "freeze_feat_extract_train": true,
+  "hidden_act": "gelu",
+  "hidden_dropout": 0.1,
+  "hidden_size": 768,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "layer_norm_eps": 1e-05,
+  "layerdrop": 0.05,
+  "mask_channel_length": 10,
+  "mask_channel_min_space": 1,
+  "mask_channel_other": 0.0,
+  "mask_channel_prob": 0.0,
+  "mask_channel_selection": "static",
+  "mask_feature_length": 10,
+  "mask_feature_min_masks": 0,
+  "mask_feature_prob": 0.0,
+  "mask_time_length": 10,
+  "mask_time_min_masks": 2,
+  "mask_time_min_space": 1,
+  "mask_time_other": 0.0,
+  "mask_time_prob": 0.05,
+  "mask_time_selection": "static",
+  "model_type": "wavlm",
+  "no_mask_channel_overlap": false,
+  "no_mask_time_overlap": false,
+  "num_adapter_layers": 3,
+  "num_attention_heads": 12,
+  "num_buckets": 320,
+  "num_codevector_groups": 2,
+  "num_codevectors_per_group": 320,
+  "num_conv_pos_embedding_groups": 16,
+  "num_conv_pos_embeddings": 128,
+  "num_ctc_classes": 80,
+  "num_feat_extract_layers": 7,
+  "num_hidden_layers": 12,
+  "num_negatives": 100,
+  "output_hidden_size": 768,
+  "pad_token_id": 0,
+  "proj_codevector_dim": 256,
+  "replace_prob": 0.5,
+  "torch_dtype": "float32",
+  "transformers_version": "4.13.0.dev0",
+  "use_weighted_layer_sum": false,
+  "vocab_size": 32,
+  "tokenizer_class": "Wav2Vec2CTCTokenizer"
+}
diff --git a/slm/wavlm-base-plus/preprocessor_config.json b/slm/wavlm-base-plus/preprocessor_config.json
new file mode 100644
index 0000000..10f6def
--- /dev/null
+++ b/slm/wavlm-base-plus/preprocessor_config.json
@@ -0,0 +1,9 @@
+{
+  "do_normalize": false,
+  "feature_extractor_type": "Wav2Vec2FeatureExtractor",
+  "feature_size": 1,
+  "padding_side": "right",
+  "padding_value": 0.0,
+  "return_attention_mask": true,
+  "sampling_rate": 16000
+}
diff --git a/text/cleaner.py b/text/cleaner.py
index f1efda2..ab59b4e 100644
--- a/text/cleaner.py
+++ b/text/cleaner.py
@@ -4,10 +4,13 @@ from text import chinese, japanese, english, cleaned_text_to_sequence
 language_module_map = {"ZH": chinese, "JP": japanese, "EN": english}
 
 
-def clean_text(text, language):
+def clean_text(text, language, use_jp_extra=True):
     language_module = language_module_map[language]
     norm_text = language_module.text_normalize(text)
-    phones, tones, word2ph = language_module.g2p(norm_text)
+    if language == "JP":
+        phones, tones, word2ph = language_module.g2p(norm_text, use_jp_extra)
+    else:
+        phones, tones, word2ph = language_module.g2p(norm_text)
     return norm_text, phones, tones, word2ph
 
 
diff --git a/text/japanese.py b/text/japanese.py
index 2c15841..970578b 100644
--- a/text/japanese.py
+++ b/text/japanese.py
@@ -23,8 +23,8 @@ COSONANTS = set(
     ]
 )
 
-# 母音の集合
-VOWELS = {"a", "i", "u", "e", "o"}
+# 母音の集合、便宜上「ん」を含める
+VOWELS = {"a", "i", "u", "e", "o", "N"}
 
 
 # 正規化で記号を変換するための辞書
@@ -62,7 +62,7 @@ rep_map = {
     "—": "-",
     "−": "-",
     # "～": "-",  # これは長音記号「ー」として扱うよう変更
-    "~": "-",
+    # "~": "-",  # これも長音記号「ー」として扱うよう変更
     "「": "'",
     "」": "'",
 }
@@ -145,7 +145,9 @@ def japanese_convert_numbers_to_words(text: str) -> str:
     return res
 
 
-def g2p(norm_text: str) -> tuple[list[str], list[int], list[int]]:
+def g2p(
+    norm_text: str, use_jp_extra: bool = True
+) -> tuple[list[str], list[int], list[int]]:
     """
     他で使われるメインの関数。`text_normalize()`で正規化された`norm_text`を受け取り、
     - phones: 音素のリスト（ただし`!`や`,`や`.`等punctuationが含まれうる）
@@ -153,6 +155,7 @@ def g2p(norm_text: str) -> tuple[list[str], list[int], list[int]]:
     - word2ph: 元のテキストの各文字に音素が何個割り当てられるかを表すリスト
     のタプルを返す。
     ただし`phones`と`tones`の最初と終わりに`_`が入り、応じて`word2ph`の最初と最後に1が追加される。
+    use_jp_extra: Falseの場合、「ん」の音素を「N」ではなく「n」とする。
     """
     # pyopenjtalkのフルコンテキストラベルを使ってアクセントを取り出すと、punctuationの位置が消えてしまい情報が失われてしまう：
     # 「こんにちは、世界。」と「こんにちは！世界。」と「こんにちは！！！？？？世界……。」は全て同じになる。
@@ -160,7 +163,7 @@ def g2p(norm_text: str) -> tuple[list[str], list[int], list[int]]:
     # それとは別にpyopenjtalk.run_frontend()で得られる音素リスト（こちらはpunctuationが保持される）を使い、
     # アクセント割当をしなおすことによってpunctuationを含めた音素とアクセントのリストを作る。
 
-    # punctuationがすべて消えた、音素とアクセントのタプルのリスト
+    # punctuationがすべて消えた、音素とアクセントのタプルのリスト（「ん」は「N」）
     phone_tone_list_wo_punct = g2phone_tone_wo_punct(norm_text)
 
     # sep_text: 単語単位の単語のリスト
@@ -185,7 +188,9 @@ def g2p(norm_text: str) -> tuple[list[str], list[int], list[int]]:
     sep_tokenized: list[list[str]] = []
     for i in sep_text:
         if i not in punctuation:
-            sep_tokenized.append(tokenizer.tokenize(i))  # ここでおそらく`i`が文字単位に分割される
+            sep_tokenized.append(
+                tokenizer.tokenize(i)
+            )  # ここでおそらく`i`が文字単位に分割される
         else:
             sep_tokenized.append([i])
 
@@ -205,11 +210,15 @@ def g2p(norm_text: str) -> tuple[list[str], list[int], list[int]]:
 
     assert len(phones) == sum(word2ph), f"{len(phones)} != {sum(word2ph)}"
 
+    # 最後にuse_jp_extraでない場合は「N」を「n」に変換
+    if not use_jp_extra:
+        phones = [phone if phone != "N" else "n" for phone in phones]
+
     return phones, tones, word2ph
 
 
 def g2kata_tone(norm_text: str) -> list[tuple[str, int]]:
-    phones, tones, _ = g2p(norm_text)
+    phones, tones, _ = g2p(norm_text, use_jp_extra=True)
     return phone_tone2kata_tone(list(zip(phones, tones)))
 
 
@@ -225,22 +234,12 @@ def phone_tone2kata_tone(phone_tone: list[tuple[str, int]]) -> list[tuple[str, i
         if phone in punctuation:
             result.append((phone, tone))
             continue
-        if phone in COSONANTS and phone != "n":  # n以外の子音の場合
+        if phone in COSONANTS:  # n以外の子音の場合
             assert current_mora == "", f"Unexpected {phone} after {current_mora}"
             assert tone == next_tone, f"Unexpected {phone} tone {tone} != {next_tone}"
             current_mora = phone
-        elif phone == "n":
-            assert current_mora == "", f"Unexpected {phone} after {current_mora}"
-            if next_phone not in VOWELS:  # 次の音素が母音でない場合
-                result.append(("ン", tone))
-            else:
-                # 子音のnの場合
-                assert (
-                    tone == next_tone
-                ), f"Unexpected {phone} tone {tone} != {next_tone}"
-                current_mora = "n"
         else:
-            # phoneが母音
+            # phoneが母音もしくは「N」
             current_mora += phone
             result.append((mora_phonemes_to_mora_kata[current_mora], tone))
             current_mora = ""
@@ -269,9 +268,9 @@ def g2phone_tone_wo_punct(text: str) -> list[tuple[str, int]]:
     テキストに対して、音素とアクセント（0か1）のペアのリストを返す。
     ただし「!」「.」「?」等の非音素記号(punctuation)は全て消える（ポーズ記号も残さない）。
     非音素記号を含める処理は`align_tones()`で行われる。
-    また「っ」は「q」に、「ん」は「n」に変換される。
+    また「っ」は「q」に、「ん」は「N」に変換される。
     例: "こんにちは、世界ー。。元気？！" →
-    [('k', 0), ('o', 0), ('n', 1), ('n', 1), ('i', 1), ('ch', 1), ('i', 1), ('w', 1), ('a', 1), ('s', 1), ('e', 1), ('k', 0), ('a', 0), ('i', 0), ('i', 0), ('g', 1), ('e', 1), ('n', 0), ('k', 0), ('i', 0)]
+    [('k', 0), ('o', 0), ('N', 1), ('n', 1), ('i', 1), ('ch', 1), ('i', 1), ('w', 1), ('a', 1), ('s', 1), ('e', 1), ('k', 0), ('a', 0), ('i', 0), ('i', 0), ('g', 1), ('e', 1), ('N', 0), ('k', 0), ('i', 0)]
     """
     prosodies = pyopenjtalk_g2p_prosody(text, drop_unvoiced_vowels=True)
     # logger.debug(f"prosodies: {prosodies}")
@@ -306,8 +305,8 @@ def g2phone_tone_wo_punct(text: str) -> list[tuple[str, int]]:
         else:
             if letter == "cl":  # 「っ」の処理
                 letter = "q"
-            elif letter == "N":  # 「ん」の処理
-                letter = "n"
+            # elif letter == "N":  # 「ん」の処理
+            #     letter = "n"
             current_phrase.append((letter, current_tone))
     return result
 
@@ -364,7 +363,7 @@ def text2sep_kata(norm_text: str) -> tuple[list[str], list[str]]:
     return sep_text, sep_kata
 
 
-# ESPnetの実装から引用、変更点無し
+# ESPnetの実装から引用、変更点無し。「ん」は「N」なことに注意。
 # https://github.com/espnet/espnet/blob/master/espnet2/text/phoneme_tokenizer.py
 def pyopenjtalk_g2p_prosody(text: str, drop_unvoiced_vowels: bool = True) -> list[str]:
     """Extract phoneme + prosoody symbol sequence from input full-context labels.
diff --git a/text/japanese_mora_list.py b/text/japanese_mora_list.py
index 1abe5ba..b43e54d 100644
--- a/text/japanese_mora_list.py
+++ b/text/japanese_mora_list.py
@@ -2,6 +2,7 @@
 VOICEVOXのソースコードからお借りして最低限に改造したコード。
 https://github.com/VOICEVOX/voicevox_engine/blob/master/voicevox_engine/tts_pipeline/mora_list.py
 """
+
 """
 以下のモーラ対応表はOpenJTalkのソースコードから取得し、
 カタカナ表記とモーラが一対一対応するように改造した。
@@ -48,8 +49,8 @@ POSSIBILITY OF SUCH DAMAGE.
 from typing import Optional
 
 # (カタカナ, 子音, 母音)の順。子音がない場合はNoneを入れる。
-# 但し「ン」と「ッ」は母音のみという扱いで、「ン」は「n」、「ッ」は「q」に変更
-# （元々はそれぞれ「N」「cl」）
+# 但し「ン」と「ッ」は母音のみという扱いで、「ン」は「N」、「ッ」は「q」とする。
+# （元々「ッ」は「cl」）
 # また「デェ = dy e」はpyopenjtalkの出力（de e）と合わないため削除
 _mora_list_minimum: list[tuple[str, Optional[str], str]] = [
     ("ヴォ", "v", "o"),
@@ -57,7 +58,7 @@ _mora_list_minimum: list[tuple[str, Optional[str], str]] = [
     ("ヴィ", "v", "i"),
     ("ヴァ", "v", "a"),
     ("ヴ", "v", "u"),
-    ("ン", None, "n"),  # 「N」から「n」に変更
+    ("ン", None, "N"),
     ("ワ", "w", "a"),
     ("ロ", "r", "o"),
     ("レ", "r", "e"),
diff --git a/train_ms.py b/train_ms.py
index 4b75991..dcf7eca 100644
--- a/train_ms.py
+++ b/train_ms.py
@@ -1,6 +1,5 @@
 import argparse
 import datetime
-import gc
 import os
 import platform
 
@@ -299,9 +298,9 @@ def run():
                 utils.latest_checkpoint_path(model_dir, "DUR_*.pth"),
                 net_dur_disc,
                 optim_dur_disc,
-                skip_optimizer=hps.train.skip_optimizer
-                if "skip_optimizer" in hps.train
-                else True,
+                skip_optimizer=(
+                    hps.train.skip_optimizer if "skip_optimizer" in hps.train else True
+                ),
             )
             if not optim_dur_disc.param_groups[0].get("initial_lr"):
                 optim_dur_disc.param_groups[0]["initial_lr"] = dur_resume_lr
@@ -309,17 +308,17 @@ def run():
             utils.latest_checkpoint_path(model_dir, "G_*.pth"),
             net_g,
             optim_g,
-            skip_optimizer=hps.train.skip_optimizer
-            if "skip_optimizer" in hps.train
-            else True,
+            skip_optimizer=(
+                hps.train.skip_optimizer if "skip_optimizer" in hps.train else True
+            ),
         )
         _, optim_d, d_resume_lr, epoch_str = utils.load_checkpoint(
             utils.latest_checkpoint_path(model_dir, "D_*.pth"),
             net_d,
             optim_d,
-            skip_optimizer=hps.train.skip_optimizer
-            if "skip_optimizer" in hps.train
-            else True,
+            skip_optimizer=(
+                hps.train.skip_optimizer if "skip_optimizer" in hps.train else True
+            ),
         )
         if not optim_g.param_groups[0].get("initial_lr"):
             optim_g.param_groups[0]["initial_lr"] = g_resume_lr
diff --git a/train_ms_jp_extra.py b/train_ms_jp_extra.py
new file mode 100644
index 0000000..8531d64
--- /dev/null
+++ b/train_ms_jp_extra.py
@@ -0,0 +1,998 @@
+import argparse
+import datetime
+import os
+import platform
+
+import torch
+import torch.distributed as dist
+from torch.cuda.amp import GradScaler, autocast
+from torch.nn import functional as F
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.utils.data import DataLoader
+from torch.utils.tensorboard import SummaryWriter
+from tqdm import tqdm
+
+# logging.getLogger("numba").setLevel(logging.WARNING)
+import commons
+import default_style
+import utils
+from common.log import logger
+from common.stdout_wrapper import SAFE_STDOUT
+from config import config
+from data_utils import (
+    DistributedBucketSampler,
+    TextAudioSpeakerCollate,
+    TextAudioSpeakerLoader,
+)
+from losses import WavLMLoss, discriminator_loss, feature_loss, generator_loss, kl_loss
+from mel_processing import mel_spectrogram_torch, spec_to_mel_torch
+from models_jp_extra import (
+    DurationDiscriminator,
+    MultiPeriodDiscriminator,
+    SynthesizerTrn,
+    WavLMDiscriminator,
+)
+from text.symbols import symbols
+
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = (
+    True  # If encontered training problem,please try to disable TF32.
+)
+torch.set_num_threads(1)
+torch.set_float32_matmul_precision("medium")
+torch.backends.cuda.sdp_kernel("flash")
+torch.backends.cuda.enable_flash_sdp(True)
+torch.backends.cuda.enable_mem_efficient_sdp(
+    True
+)  # Not available if torch version is lower than 2.0
+global_step = 0
+
+
+def run():
+    # Command line configuration is not recommended unless necessary, use config.yml
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "-c",
+        "--config",
+        type=str,
+        default=config.train_ms_config.config_path,
+        help="JSON file for configuration",
+    )
+    parser.add_argument(
+        "-m",
+        "--model",
+        type=str,
+        help="数据集文件夹路径，请注意，数据不再默认放在/logs文件夹下。如果需要用命令行配置，请声明相对于根目录的路径",
+        default=config.dataset_path,
+    )
+    parser.add_argument(
+        "--assets_root",
+        type=str,
+        help="Root directory of model assets needed for inference.",
+        default=config.assets_root,
+    )
+    parser.add_argument(
+        "--skip_default_style",
+        action="store_true",
+        help="Skip saving default style config and mean vector.",
+    )
+    parser.add_argument(
+        "--no_progress_bar",
+        action="store_true",
+        help="Do not show the progress bar while training.",
+    )
+    args = parser.parse_args()
+
+    # Set log file
+    model_dir = os.path.join(args.model, config.train_ms_config.model_dir)
+    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+    logger.add(os.path.join(args.model, f"train_{timestamp}.log"))
+
+    # Parsing environment variables
+    envs = config.train_ms_config.env
+    for env_name, env_value in envs.items():
+        if env_name not in os.environ.keys():
+            logger.info("Loading configuration from config {}".format(str(env_value)))
+            os.environ[env_name] = str(env_value)
+    logger.info(
+        "Loading environment variables \nMASTER_ADDR: {},\nMASTER_PORT: {},\nWORLD_SIZE: {},\nRANK: {},\nLOCAL_RANK: {}".format(
+            os.environ["MASTER_ADDR"],
+            os.environ["MASTER_PORT"],
+            os.environ["WORLD_SIZE"],
+            os.environ["RANK"],
+            os.environ["LOCAL_RANK"],
+        )
+    )
+
+    backend = "nccl"
+    if platform.system() == "Windows":
+        backend = "gloo"  # If Windows,switch to gloo backend.
+    dist.init_process_group(
+        backend=backend,
+        init_method="env://",
+        timeout=datetime.timedelta(seconds=300),
+    )  # Use torchrun instead of mp.spawn
+    rank = dist.get_rank()
+    local_rank = int(os.environ["LOCAL_RANK"])
+    n_gpus = dist.get_world_size()
+
+    hps = utils.get_hparams_from_file(args.config)
+    # This is needed because we have to pass `model_dir` to `train_and_evaluate()`
+    hps.model_dir = model_dir
+
+    # 比较路径是否相同
+    if os.path.realpath(args.config) != os.path.realpath(
+        config.train_ms_config.config_path
+    ):
+        with open(args.config, "r", encoding="utf-8") as f:
+            data = f.read()
+        os.makedirs(os.path.dirname(config.train_ms_config.config_path), exist_ok=True)
+        with open(config.train_ms_config.config_path, "w", encoding="utf-8") as f:
+            f.write(data)
+
+    """
+    Path constants are a bit complicated...
+    TODO: Refactor or rename these?
+    (Both `config.yml` and `config.json` are used, which is confusing I think.)
+
+    args.model: For saving all info needed for training.
+        default: `Data/{model_name}`.
+    hps.model_dir := model_dir: For saving checkpoints (for resuming training).
+        default: `Data/{model_name}/models`.
+        (Use `hps` since we have to pass `model_dir` to `train_and_evaluate()`.
+
+    args.assets_root: The root directory of model assets needed for inference.
+        default: config.assets_root == `model_assets`.
+
+    config.out_dir: The directory for model assets of this model (for inference).
+        default: `model_assets/{model_name}`.
+    """
+    os.makedirs(config.out_dir, exist_ok=True)
+
+    if not args.skip_default_style:
+        # Save default style to out_dir
+        default_style.set_style_config(
+            args.config, os.path.join(config.out_dir, "config.json")
+        )
+        default_style.save_mean_vector(
+            os.path.join(args.model, "wavs"),
+            os.path.join(config.out_dir, "style_vectors.npy"),
+        )
+
+    torch.manual_seed(hps.train.seed)
+    torch.cuda.set_device(local_rank)
+
+    global global_step
+    if rank == 0:
+        # logger = utils.get_logger(hps.model_dir)
+        # logger.info(hps)
+        utils.check_git_hash(model_dir)
+        writer = SummaryWriter(log_dir=model_dir)
+        writer_eval = SummaryWriter(log_dir=os.path.join(model_dir, "eval"))
+    train_dataset = TextAudioSpeakerLoader(hps.data.training_files, hps.data)
+    train_sampler = DistributedBucketSampler(
+        train_dataset,
+        hps.train.batch_size,
+        [32, 300, 400, 500, 600, 700, 800, 900, 1000],
+        num_replicas=n_gpus,
+        rank=rank,
+        shuffle=True,
+    )
+    collate_fn = TextAudioSpeakerCollate()
+    train_loader = DataLoader(
+        train_dataset,
+        # num_workers=min(config.train_ms_config.num_workers, os.cpu_count() - 1),
+        # Slow and often freezes, so use only half of the cores.
+        num_workers=min(config.train_ms_config.num_workers, os.cpu_count() // 2),
+        shuffle=False,
+        pin_memory=True,
+        collate_fn=collate_fn,
+        batch_sampler=train_sampler,
+        persistent_workers=True,
+        prefetch_factor=6,
+    )  # DataLoader config could be adjusted.
+    if rank == 0:
+        eval_dataset = TextAudioSpeakerLoader(hps.data.validation_files, hps.data)
+        eval_loader = DataLoader(
+            eval_dataset,
+            num_workers=0,
+            shuffle=False,
+            batch_size=1,
+            pin_memory=True,
+            drop_last=False,
+            collate_fn=collate_fn,
+        )
+    if (
+        "use_noise_scaled_mas" in hps.model.keys()
+        and hps.model.use_noise_scaled_mas is True
+    ):
+        logger.info("Using noise scaled MAS for VITS2")
+        mas_noise_scale_initial = 0.01
+        noise_scale_delta = 2e-6
+    else:
+        logger.info("Using normal MAS for VITS1")
+        mas_noise_scale_initial = 0.0
+        noise_scale_delta = 0.0
+    if (
+        "use_duration_discriminator" in hps.model.keys()
+        and hps.model.use_duration_discriminator is True
+    ):
+        logger.info("Using duration discriminator for VITS2")
+        net_dur_disc = DurationDiscriminator(
+            hps.model.hidden_channels,
+            hps.model.hidden_channels,
+            3,
+            0.1,
+            gin_channels=hps.model.gin_channels if hps.data.n_speakers != 0 else 0,
+        ).cuda(local_rank)
+    else:
+        net_dur_disc = None
+    if (
+        "use_wavlm_discriminator" in hps.model.keys()
+        and hps.model.use_wavlm_discriminator is True
+    ):
+        net_wd = WavLMDiscriminator(
+            hps.model.slm.hidden, hps.model.slm.nlayers, hps.model.slm.initial_channel
+        ).cuda(local_rank)
+    else:
+        net_wd = None
+    if (
+        "use_spk_conditioned_encoder" in hps.model.keys()
+        and hps.model.use_spk_conditioned_encoder is True
+    ):
+        if hps.data.n_speakers == 0:
+            raise ValueError(
+                "n_speakers must be > 0 when using spk conditioned encoder to train multi-speaker model"
+            )
+    else:
+        logger.info("Using normal encoder for VITS1")
+
+    net_g = SynthesizerTrn(
+        len(symbols),
+        hps.data.filter_length // 2 + 1,
+        hps.train.segment_size // hps.data.hop_length,
+        n_speakers=hps.data.n_speakers,
+        mas_noise_scale_initial=mas_noise_scale_initial,
+        noise_scale_delta=noise_scale_delta,
+        **hps.model,
+    ).cuda(local_rank)
+    if getattr(hps.train, "freeze_JP_bert", False):
+        logger.info("Freezing (JP) bert encoder !!!")
+        for param in net_g.enc_p.bert_proj.parameters():
+            param.requires_grad = False
+    if getattr(hps.train, "freeze_style", False):
+        logger.info("Freezing style encoder !!!")
+        for param in net_g.enc_p.style_proj.parameters():
+            param.requires_grad = False
+
+    net_d = MultiPeriodDiscriminator(hps.model.use_spectral_norm).cuda(local_rank)
+    optim_g = torch.optim.AdamW(
+        filter(lambda p: p.requires_grad, net_g.parameters()),
+        hps.train.learning_rate,
+        betas=hps.train.betas,
+        eps=hps.train.eps,
+    )
+    optim_d = torch.optim.AdamW(
+        net_d.parameters(),
+        hps.train.learning_rate,
+        betas=hps.train.betas,
+        eps=hps.train.eps,
+    )
+    if net_dur_disc is not None:
+        optim_dur_disc = torch.optim.AdamW(
+            net_dur_disc.parameters(),
+            hps.train.learning_rate,
+            betas=hps.train.betas,
+            eps=hps.train.eps,
+        )
+    else:
+        optim_dur_disc = None
+    if net_wd is not None:
+        optim_wd = torch.optim.AdamW(
+            net_wd.parameters(),
+            hps.train.learning_rate,
+            betas=hps.train.betas,
+            eps=hps.train.eps,
+        )
+    else:
+        optim_wd = None
+    net_g = DDP(net_g, device_ids=[local_rank], bucket_cap_mb=512)
+    net_d = DDP(net_d, device_ids=[local_rank], bucket_cap_mb=512)
+    if net_dur_disc is not None:
+        net_dur_disc = DDP(
+            net_dur_disc,
+            device_ids=[local_rank],
+            bucket_cap_mb=512,
+        )
+    if net_wd is not None:
+        net_wd = DDP(net_wd, device_ids=[local_rank], bucket_cap_mb=512)
+
+    if utils.is_resuming(model_dir):
+        if net_dur_disc is not None:
+            try:
+                _, _, dur_resume_lr, epoch_str = utils.load_checkpoint(
+                    utils.latest_checkpoint_path(model_dir, "DUR_*.pth"),
+                    net_dur_disc,
+                    optim_dur_disc,
+                    skip_optimizer=(
+                        hps.train.skip_optimizer
+                        if "skip_optimizer" in hps.train
+                        else True
+                    ),
+                )
+                if not optim_dur_disc.param_groups[0].get("initial_lr"):
+                    optim_dur_disc.param_groups[0]["initial_lr"] = dur_resume_lr
+            except:
+                if not optim_dur_disc.param_groups[0].get("initial_lr"):
+                    optim_dur_disc.param_groups[0]["initial_lr"] = dur_resume_lr
+                print("Initialize dur_disc")
+        if net_wd is not None:
+            try:
+                _, optim_wd, wd_resume_lr, epoch_str = utils.load_checkpoint(
+                    utils.latest_checkpoint_path(model_dir, "WD_*.pth"),
+                    net_wd,
+                    optim_wd,
+                    skip_optimizer=(
+                        hps.train.skip_optimizer
+                        if "skip_optimizer" in hps.train
+                        else True
+                    ),
+                )
+                if not optim_wd.param_groups[0].get("initial_lr"):
+                    optim_wd.param_groups[0]["initial_lr"] = wd_resume_lr
+            except:
+                if not optim_wd.param_groups[0].get("initial_lr"):
+                    optim_wd.param_groups[0]["initial_lr"] = wd_resume_lr
+                logger.info("Initialize wavlm")
+
+        try:
+            _, optim_g, g_resume_lr, epoch_str = utils.load_checkpoint(
+                utils.latest_checkpoint_path(model_dir, "G_*.pth"),
+                net_g,
+                optim_g,
+                skip_optimizer=(
+                    hps.train.skip_optimizer if "skip_optimizer" in hps.train else True
+                ),
+            )
+            _, optim_d, d_resume_lr, epoch_str = utils.load_checkpoint(
+                utils.latest_checkpoint_path(model_dir, "D_*.pth"),
+                net_d,
+                optim_d,
+                skip_optimizer=(
+                    hps.train.skip_optimizer if "skip_optimizer" in hps.train else True
+                ),
+            )
+            if not optim_g.param_groups[0].get("initial_lr"):
+                optim_g.param_groups[0]["initial_lr"] = g_resume_lr
+            if not optim_d.param_groups[0].get("initial_lr"):
+                optim_d.param_groups[0]["initial_lr"] = d_resume_lr
+
+            epoch_str = max(epoch_str, 1)
+            # global_step = (epoch_str - 1) * len(train_loader)
+            global_step = int(
+                utils.get_steps(utils.latest_checkpoint_path(model_dir, "G_*.pth"))
+            )
+            logger.info(
+                f"******************Found the model. Current epoch is {epoch_str}, gloabl step is {global_step}*********************"
+            )
+        except Exception as e:
+            logger.warning(e)
+            logger.warning(
+                "It seems that you are not using the pretrained models, so we will train from scratch."
+            )
+            epoch_str = 1
+            global_step = 0
+    else:
+        try:
+            _ = utils.load_safetensors(
+                os.path.join(model_dir, "G_0.safetensors"), net_g
+            )
+            _ = utils.load_safetensors(
+                os.path.join(model_dir, "D_0.safetensors"), net_d
+            )
+            if net_dur_disc is not None:
+                _ = utils.load_safetensors(
+                    os.path.join(model_dir, "DUR_0.safetensors"), net_dur_disc
+                )
+            logger.info("Loaded the pretrained models.")
+        except Exception as e:
+            logger.warning(e)
+            logger.warning(
+                "It seems that you are not using the pretrained models, so we will train from scratch."
+            )
+        finally:
+            epoch_str = 1
+            global_step = 0
+
+    def lr_lambda(epoch):
+        """
+        Learning rate scheduler for warmup and exponential decay.
+        - During the warmup period, the learning rate increases linearly.
+        - After the warmup period, the learning rate decreases exponentially.
+        """
+        if epoch < hps.train.warmup_epochs:
+            return float(epoch) / float(max(1, hps.train.warmup_epochs))
+        else:
+            return hps.train.lr_decay ** (epoch - hps.train.warmup_epochs)
+
+    scheduler_last_epoch = epoch_str - 2
+    scheduler_g = torch.optim.lr_scheduler.LambdaLR(
+        optim_g, lr_lambda=lr_lambda, last_epoch=scheduler_last_epoch
+    )
+    scheduler_d = torch.optim.lr_scheduler.LambdaLR(
+        optim_d, lr_lambda=lr_lambda, last_epoch=scheduler_last_epoch
+    )
+    if net_dur_disc is not None:
+        scheduler_dur_disc = torch.optim.lr_scheduler.LambdaLR(
+            optim_dur_disc, lr_lambda=lr_lambda, last_epoch=scheduler_last_epoch
+        )
+    else:
+        scheduler_dur_disc = None
+    if net_wd is not None:
+        scheduler_wd = torch.optim.lr_scheduler.LambdaLR(
+            optim_wd, lr_lambda=lr_lambda, last_epoch=scheduler_last_epoch
+        )
+        wl = WavLMLoss(
+            hps.model.slm.model,
+            net_wd,
+            hps.data.sampling_rate,
+            hps.model.slm.sr,
+        ).to(local_rank)
+    else:
+        scheduler_wd = None
+        wl = None
+    scaler = GradScaler(enabled=hps.train.bf16_run)
+    logger.info("Start training.")
+
+    diff = abs(
+        epoch_str * len(train_loader) - (hps.train.epochs + 1) * len(train_loader)
+    )
+    pbar = None
+    if not args.no_progress_bar:
+        pbar = tqdm(
+            total=global_step + diff,
+            initial=global_step,
+            smoothing=0.05,
+            file=SAFE_STDOUT,
+        )
+    initial_step = global_step
+
+    for epoch in range(epoch_str, hps.train.epochs + 1):
+        if rank == 0:
+            train_and_evaluate(
+                rank,
+                local_rank,
+                epoch,
+                hps,
+                [net_g, net_d, net_dur_disc, net_wd, wl],
+                [optim_g, optim_d, optim_dur_disc, optim_wd],
+                [scheduler_g, scheduler_d, scheduler_dur_disc, scheduler_wd],
+                scaler,
+                [train_loader, eval_loader],
+                logger,
+                [writer, writer_eval],
+                pbar,
+                initial_step,
+            )
+        else:
+            train_and_evaluate(
+                rank,
+                local_rank,
+                epoch,
+                hps,
+                [net_g, net_d, net_dur_disc, net_wd, wl],
+                [optim_g, optim_d, optim_dur_disc, optim_wd],
+                [scheduler_g, scheduler_d, scheduler_dur_disc, scheduler_wd],
+                scaler,
+                [train_loader, None],
+                None,
+                None,
+                pbar,
+                initial_step,
+            )
+        scheduler_g.step()
+        scheduler_d.step()
+        if net_dur_disc is not None:
+            scheduler_dur_disc.step()
+        if net_wd is not None:
+            scheduler_wd.step()
+        if epoch == hps.train.epochs:
+            # Save the final models
+            utils.save_checkpoint(
+                net_g,
+                optim_g,
+                hps.train.learning_rate,
+                epoch,
+                os.path.join(model_dir, "G_{}.pth".format(global_step)),
+            )
+            utils.save_checkpoint(
+                net_d,
+                optim_d,
+                hps.train.learning_rate,
+                epoch,
+                os.path.join(model_dir, "D_{}.pth".format(global_step)),
+            )
+            if net_dur_disc is not None:
+                utils.save_checkpoint(
+                    net_dur_disc,
+                    optim_dur_disc,
+                    hps.train.learning_rate,
+                    epoch,
+                    os.path.join(model_dir, "DUR_{}.pth".format(global_step)),
+                )
+            if net_wd is not None:
+                utils.save_checkpoint(
+                    net_wd,
+                    optim_wd,
+                    hps.train.learning_rate,
+                    epoch,
+                    os.path.join(model_dir, "WD_{}.pth".format(global_step)),
+                )
+            utils.save_safetensors(
+                net_g,
+                epoch,
+                os.path.join(
+                    config.out_dir,
+                    f"{config.model_name}_e{epoch}_s{global_step}.safetensors",
+                ),
+                for_infer=True,
+            )
+
+    if pbar is not None:
+        pbar.close()
+
+
+def train_and_evaluate(
+    rank,
+    local_rank,
+    epoch,
+    hps,
+    nets,
+    optims,
+    schedulers,
+    scaler,
+    loaders,
+    logger,
+    writers,
+    pbar: tqdm,
+    initial_step: int,
+):
+    net_g, net_d, net_dur_disc, net_wd, wl = nets
+    optim_g, optim_d, optim_dur_disc, optim_wd = optims
+    scheduler_g, scheduler_d, scheduler_dur_disc, scheduler_wd = schedulers
+    train_loader, eval_loader = loaders
+    if writers is not None:
+        writer, writer_eval = writers
+
+    train_loader.batch_sampler.set_epoch(epoch)
+    global global_step
+
+    net_g.train()
+    net_d.train()
+    if net_dur_disc is not None:
+        net_dur_disc.train()
+    if net_wd is not None:
+        net_wd.train()
+    for batch_idx, (
+        x,
+        x_lengths,
+        spec,
+        spec_lengths,
+        y,
+        y_lengths,
+        speakers,
+        tone,
+        language,
+        bert,
+        style_vec,
+    ) in enumerate(tqdm(train_loader)):
+        if net_g.module.use_noise_scaled_mas:
+            current_mas_noise_scale = (
+                net_g.module.mas_noise_scale_initial
+                - net_g.module.noise_scale_delta * global_step
+            )
+            net_g.module.current_mas_noise_scale = max(current_mas_noise_scale, 0.0)
+        x, x_lengths = x.cuda(local_rank, non_blocking=True), x_lengths.cuda(
+            local_rank, non_blocking=True
+        )
+        spec, spec_lengths = spec.cuda(
+            local_rank, non_blocking=True
+        ), spec_lengths.cuda(local_rank, non_blocking=True)
+        y, y_lengths = y.cuda(local_rank, non_blocking=True), y_lengths.cuda(
+            local_rank, non_blocking=True
+        )
+        speakers = speakers.cuda(local_rank, non_blocking=True)
+        tone = tone.cuda(local_rank, non_blocking=True)
+        language = language.cuda(local_rank, non_blocking=True)
+        bert = bert.cuda(local_rank, non_blocking=True)
+        style_vec = style_vec.cuda(local_rank, non_blocking=True)
+
+        with autocast(enabled=hps.train.bf16_run, dtype=torch.bfloat16):
+            (
+                y_hat,
+                l_length,
+                attn,
+                ids_slice,
+                x_mask,
+                z_mask,
+                (z, z_p, m_p, logs_p, m_q, logs_q),
+                (hidden_x, logw, logw_),  # , logw_sdp),
+                g,
+            ) = net_g(
+                x,
+                x_lengths,
+                spec,
+                spec_lengths,
+                speakers,
+                tone,
+                language,
+                bert,
+                style_vec,
+            )
+            mel = 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,
+            )
+            y_mel = commons.slice_segments(
+                mel, ids_slice, hps.train.segment_size // hps.data.hop_length
+            )
+            y_hat_mel = mel_spectrogram_torch(
+                y_hat.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,
+            )
+
+            y = commons.slice_segments(
+                y, ids_slice * hps.data.hop_length, hps.train.segment_size
+            )  # slice
+
+            # Discriminator
+            y_d_hat_r, y_d_hat_g, _, _ = net_d(y, y_hat.detach())
+            with autocast(enabled=hps.train.bf16_run, dtype=torch.bfloat16):
+                loss_disc, losses_disc_r, losses_disc_g = discriminator_loss(
+                    y_d_hat_r, y_d_hat_g
+                )
+                loss_disc_all = loss_disc
+            if net_dur_disc is not None:
+                y_dur_hat_r, y_dur_hat_g = net_dur_disc(
+                    hidden_x.detach(),
+                    x_mask.detach(),
+                    logw_.detach(),
+                    logw.detach(),
+                    g.detach(),
+                )
+                with autocast(enabled=hps.train.bf16_run, dtype=torch.bfloat16):
+                    # TODO: I think need to mean using the mask, but for now, just mean all
+                    (
+                        loss_dur_disc,
+                        losses_dur_disc_r,
+                        losses_dur_disc_g,
+                    ) = discriminator_loss(y_dur_hat_r, y_dur_hat_g)
+                    loss_dur_disc_all = loss_dur_disc
+                optim_dur_disc.zero_grad()
+                scaler.scale(loss_dur_disc_all).backward()
+                scaler.unscale_(optim_dur_disc)
+                # torch.nn.utils.clip_grad_norm_(
+                # parameters=net_dur_disc.parameters(), max_norm=5
+                # )
+                grad_norm_dur = commons.clip_grad_value_(
+                    net_dur_disc.parameters(), None
+                )
+                scaler.step(optim_dur_disc)
+            if net_wd is not None:
+                with autocast(enabled=hps.train.bf16_run, dtype=torch.bfloat16):
+                    loss_slm = wl.discriminator(
+                        y.detach().squeeze(), y_hat.detach().squeeze()
+                    ).mean()
+
+                optim_wd.zero_grad()
+                scaler.scale(loss_slm).backward()
+                scaler.unscale_(optim_wd)
+                # torch.nn.utils.clip_grad_norm_(parameters=net_wd.parameters(), max_norm=200)
+                grad_norm_wd = commons.clip_grad_value_(net_wd.parameters(), None)
+                scaler.step(optim_wd)
+
+        optim_d.zero_grad()
+        scaler.scale(loss_disc_all).backward()
+        scaler.unscale_(optim_d)
+        if getattr(hps.train, "bf16_run", False):
+            torch.nn.utils.clip_grad_norm_(parameters=net_d.parameters(), max_norm=200)
+        grad_norm_d = commons.clip_grad_value_(net_d.parameters(), None)
+        scaler.step(optim_d)
+
+        with autocast(enabled=hps.train.bf16_run, dtype=torch.bfloat16):
+            # Generator
+            y_d_hat_r, y_d_hat_g, fmap_r, fmap_g = net_d(y, y_hat)
+            if net_dur_disc is not None:
+                _, y_dur_hat_g = net_dur_disc(hidden_x, x_mask, logw_, logw, g)
+            if net_wd is not None:
+                loss_lm = wl(y.detach().squeeze(), y_hat.squeeze()).mean()
+                loss_lm_gen = wl.generator(y_hat.squeeze())
+            with autocast(enabled=hps.train.bf16_run, dtype=torch.bfloat16):
+                loss_dur = torch.sum(l_length.float())
+                loss_mel = F.l1_loss(y_mel, y_hat_mel) * hps.train.c_mel
+                loss_kl = kl_loss(z_p, logs_q, m_p, logs_p, z_mask) * hps.train.c_kl
+
+                loss_fm = feature_loss(fmap_r, fmap_g)
+                loss_gen, losses_gen = generator_loss(y_d_hat_g)
+                # loss_commit = loss_commit * hps.train.c_commit
+
+                loss_gen_all = loss_gen + loss_fm + loss_mel + loss_dur + loss_kl
+                if net_dur_disc is not None:
+                    loss_dur_gen, losses_dur_gen = generator_loss(y_dur_hat_g)
+                    if net_wd is not None:
+                        loss_gen_all += loss_dur_gen + loss_lm + loss_lm_gen
+                    else:
+                        loss_gen_all += loss_dur_gen
+        optim_g.zero_grad()
+        scaler.scale(loss_gen_all).backward()
+        scaler.unscale_(optim_g)
+        # if getattr(hps.train, "bf16_run", False):
+        torch.nn.utils.clip_grad_norm_(parameters=net_g.parameters(), max_norm=500)
+        grad_norm_g = commons.clip_grad_value_(net_g.parameters(), None)
+        scaler.step(optim_g)
+        scaler.update()
+
+        if rank == 0:
+            if global_step % hps.train.log_interval == 0:
+                lr = optim_g.param_groups[0]["lr"]
+                losses = [loss_disc, loss_gen, loss_fm, loss_mel, loss_dur, loss_kl]
+                # logger.info(
+                #     "Train Epoch: {} [{:.0f}%]".format(
+                #         epoch, 100.0 * batch_idx / len(train_loader)
+                #     )
+                # )
+                # logger.info([x.item() for x in losses] + [global_step, lr])
+
+                scalar_dict = {
+                    "loss/g/total": loss_gen_all,
+                    "loss/d/total": loss_disc_all,
+                    "learning_rate": lr,
+                    "grad_norm_d": grad_norm_d,
+                    "grad_norm_g": grad_norm_g,
+                }
+                scalar_dict.update(
+                    {
+                        "loss/g/fm": loss_fm,
+                        "loss/g/mel": loss_mel,
+                        "loss/g/dur": loss_dur,
+                        "loss/g/kl": loss_kl,
+                    }
+                )
+                scalar_dict.update(
+                    {"loss/g/{}".format(i): v for i, v in enumerate(losses_gen)}
+                )
+                scalar_dict.update(
+                    {"loss/d_r/{}".format(i): v for i, v in enumerate(losses_disc_r)}
+                )
+                scalar_dict.update(
+                    {"loss/d_g/{}".format(i): v for i, v in enumerate(losses_disc_g)}
+                )
+
+                if net_dur_disc is not None:
+                    scalar_dict.update({"loss/dur_disc/total": loss_dur_disc_all})
+
+                    scalar_dict.update(
+                        {
+                            "loss/dur_disc_g/{}".format(i): v
+                            for i, v in enumerate(losses_dur_disc_g)
+                        }
+                    )
+                    scalar_dict.update(
+                        {
+                            "loss/dur_disc_r/{}".format(i): v
+                            for i, v in enumerate(losses_dur_disc_r)
+                        }
+                    )
+
+                    scalar_dict.update({"loss/g/dur_gen": loss_dur_gen})
+                    scalar_dict.update(
+                        {
+                            "loss/g/dur_gen_{}".format(i): v
+                            for i, v in enumerate(losses_dur_gen)
+                        }
+                    )
+
+                if net_wd is not None:
+                    scalar_dict.update(
+                        {
+                            "loss/wd/total": loss_slm,
+                            "grad_norm_wd": grad_norm_wd,
+                            "loss/g/lm": loss_lm,
+                            "loss/g/lm_gen": loss_lm_gen,
+                        }
+                    )
+                image_dict = {
+                    "slice/mel_org": utils.plot_spectrogram_to_numpy(
+                        y_mel[0].data.cpu().numpy()
+                    ),
+                    "slice/mel_gen": utils.plot_spectrogram_to_numpy(
+                        y_hat_mel[0].data.cpu().numpy()
+                    ),
+                    "all/mel": utils.plot_spectrogram_to_numpy(
+                        mel[0].data.cpu().numpy()
+                    ),
+                    "all/attn": utils.plot_alignment_to_numpy(
+                        attn[0, 0].data.cpu().numpy()
+                    ),
+                }
+                utils.summarize(
+                    writer=writer,
+                    global_step=global_step,
+                    images=image_dict,
+                    scalars=scalar_dict,
+                )
+
+            if (
+                global_step % hps.train.eval_interval == 0
+                and global_step != 0
+                and initial_step != global_step
+            ):
+                evaluate(hps, net_g, eval_loader, writer_eval)
+                utils.save_checkpoint(
+                    net_g,
+                    optim_g,
+                    hps.train.learning_rate,
+                    epoch,
+                    os.path.join(hps.model_dir, "G_{}.pth".format(global_step)),
+                )
+                utils.save_checkpoint(
+                    net_d,
+                    optim_d,
+                    hps.train.learning_rate,
+                    epoch,
+                    os.path.join(hps.model_dir, "D_{}.pth".format(global_step)),
+                )
+                if net_dur_disc is not None:
+                    utils.save_checkpoint(
+                        net_dur_disc,
+                        optim_dur_disc,
+                        hps.train.learning_rate,
+                        epoch,
+                        os.path.join(hps.model_dir, "DUR_{}.pth".format(global_step)),
+                    )
+                if net_wd is not None:
+                    utils.save_checkpoint(
+                        net_wd,
+                        optim_wd,
+                        hps.train.learning_rate,
+                        epoch,
+                        os.path.join(hps.model_dir, "WD_{}.pth".format(global_step)),
+                    )
+                keep_ckpts = config.train_ms_config.keep_ckpts
+                if keep_ckpts > 0:
+                    utils.clean_checkpoints(
+                        path_to_models=hps.model_dir,
+                        n_ckpts_to_keep=keep_ckpts,
+                        sort_by_time=True,
+                    )
+                # Save safetensors (for inference) to `model_assets/{model_name}`
+                utils.save_safetensors(
+                    net_g,
+                    epoch,
+                    os.path.join(
+                        config.out_dir,
+                        f"{config.model_name}_e{epoch}_s{global_step}.safetensors",
+                    ),
+                    for_infer=True,
+                )
+
+        global_step += 1
+        if pbar is not None:
+            pbar.set_description(
+                "Epoch {}({:.0f}%)/{}".format(
+                    epoch, 100.0 * batch_idx / len(train_loader), hps.train.epochs
+                )
+            )
+            pbar.update()
+    # 本家ではこれをスピードアップのために消すと書かれていたので、一応消してみる
+    # gc.collect()
+    # torch.cuda.empty_cache()
+    if pbar is None and rank == 0:
+        logger.info(f"====> Epoch: {epoch}, step: {global_step}")
+
+
+def evaluate(hps, generator, eval_loader, writer_eval):
+    generator.eval()
+    image_dict = {}
+    audio_dict = {}
+    logger.info("Evaluating ...")
+    with torch.no_grad():
+        for batch_idx, (
+            x,
+            x_lengths,
+            spec,
+            spec_lengths,
+            y,
+            y_lengths,
+            speakers,
+            tone,
+            language,
+            bert,
+            style_vec,
+        ) in enumerate(eval_loader):
+            x, x_lengths = x.cuda(), x_lengths.cuda()
+            spec, spec_lengths = spec.cuda(), spec_lengths.cuda()
+            y, y_lengths = y.cuda(), y_lengths.cuda()
+            speakers = speakers.cuda()
+            bert = bert.cuda()
+            tone = tone.cuda()
+            language = language.cuda()
+            style_vec = style_vec.cuda()
+            for use_sdp in [True, False]:
+                y_hat, attn, mask, *_ = generator.module.infer(
+                    x,
+                    x_lengths,
+                    speakers,
+                    tone,
+                    language,
+                    bert,
+                    style_vec,
+                    y=spec,
+                    max_len=1000,
+                    sdp_ratio=0.0 if not use_sdp else 1.0,
+                )
+                y_hat_lengths = mask.sum([1, 2]).long() * hps.data.hop_length
+
+                mel = 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,
+                )
+                y_hat_mel = mel_spectrogram_torch(
+                    y_hat.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,
+                )
+                image_dict.update(
+                    {
+                        f"gen/mel_{batch_idx}": utils.plot_spectrogram_to_numpy(
+                            y_hat_mel[0].cpu().numpy()
+                        )
+                    }
+                )
+                audio_dict.update(
+                    {
+                        f"gen/audio_{batch_idx}_{use_sdp}": y_hat[
+                            0, :, : y_hat_lengths[0]
+                        ]
+                    }
+                )
+                image_dict.update(
+                    {
+                        f"gt/mel_{batch_idx}": utils.plot_spectrogram_to_numpy(
+                            mel[0].cpu().numpy()
+                        )
+                    }
+                )
+                audio_dict.update({f"gt/audio_{batch_idx}": y[0, :, : y_lengths[0]]})
+
+    utils.summarize(
+        writer=writer_eval,
+        global_step=global_step,
+        images=image_dict,
+        audios=audio_dict,
+        audio_sampling_rate=hps.data.sampling_rate,
+    )
+    generator.train()
+
+
+if __name__ == "__main__":
+    run()
diff --git a/webui_train.py b/webui_train.py
index 94b1bc3..db55738 100644
--- a/webui_train.py
+++ b/webui_train.py
@@ -40,6 +40,7 @@ def initialize(
     freeze_JP_bert,
     freeze_ZH_bert,
     freeze_style,
+    use_jp_extra,
 ):
     global logger_handler
     dataset_path, _, train_path, val_path, config_path = get_path(model_name)
@@ -52,14 +53,20 @@ def initialize(
     logger_handler = logger.add(os.path.join(dataset_path, file_name))
 
     logger.info(
-        f"Step 1: start initialization...\nmodel_name: {model_name}, batch_size: {batch_size}, epochs: {epochs}, save_every_steps: {save_every_steps}, bf16_run: {bf16_run}, freeze_ZH_bert: {freeze_ZH_bert}, freeze_JP_bert: {freeze_JP_bert}, freeze_EN_bert: {freeze_EN_bert}, freeze_style: {freeze_style}"
+        f"Step 1: start initialization...\nmodel_name: {model_name}, batch_size: {batch_size}, epochs: {epochs}, save_every_steps: {save_every_steps}, bf16_run: {bf16_run}, freeze_ZH_bert: {freeze_ZH_bert}, freeze_JP_bert: {freeze_JP_bert}, freeze_EN_bert: {freeze_EN_bert}, freeze_style: {freeze_style}, use_jp_extra: {use_jp_extra}"
     )
     if os.path.isfile(config_path):
         with open(config_path, "r", encoding="utf-8") as f:
             config = json.load(f)
     else:
         # Use default config
-        with open("configs/config.json", "r", encoding="utf-8") as f:
+        default_config_path = (
+            "configs/config.json"
+            if not use_jp_extra
+            else "configs/configs_jp_extra.json"
+        )
+
+        with open(default_config_path, "r", encoding="utf-8") as f:
             config = json.load(f)
     config["model_name"] = model_name
     config["data"]["training_files"] = train_path
@@ -83,14 +90,15 @@ def initialize(
             dirs_exist_ok=True,
         )
         shutil.rmtree(model_path)
+    pretrained_dir = "pretrained" if not use_jp_extra else "pretrained_jp_extra"
     try:
         shutil.copytree(
-            src="pretrained",
+            src=pretrained_dir,
             dst=model_path,
         )
     except FileNotFoundError:
-        logger.error("Step 1: `pretrained` folder not found.")
-        return False, "Step 1, Error: pretrainedフォルダが見つかりません。"
+        logger.error(f"Step 1: {pretrained_dir} folder not found.")
+        return False, f"Step 1, Error: {pretrained_dir}フォルダが見つかりません。"
 
     with open(config_path, "w", encoding="utf-8") as f:
         json.dump(config, f, indent=2, ensure_ascii=False)
@@ -138,7 +146,7 @@ def resample(model_name, normalize, trim, num_processes):
     return True, "Step 2, Success: 音声ファイルの前処理が完了しました"
 
 
-def preprocess_text(model_name):
+def preprocess_text(model_name, use_jp_extra):
     logger.info("Step 3: start preprocessing text...")
     dataset_path, lbl_path, train_path, val_path, config_path = get_path(model_name)
     try:
@@ -153,25 +161,32 @@ def preprocess_text(model_name):
                 "\\", "/"
             )
             f.writelines(f"{path}|{spk}|{language}|{text}\n")
-    success, message = run_script_with_log(
-        [
-            "preprocess_text.py",
-            "--config-path",
-            config_path,
-            "--transcription-path",
-            lbl_path,
-            "--train-path",
-            train_path,
-            "--val-path",
-            val_path,
-        ]
-    )
+    cmd = [
+        "preprocess_text.py",
+        "--config-path",
+        config_path,
+        "--transcription-path",
+        lbl_path,
+        "--train-path",
+        train_path,
+        "--val-path",
+        val_path,
+    ]
+    if use_jp_extra:
+        cmd.append("--use_jp_extra")
+    success, message = run_script_with_log(cmd)
     if not success:
         logger.error(f"Step 3: preprocessing text failed.")
-        return False, f"Step 3, Error: 書き起こしファイルの前処理に失敗しました:\n{message}"
+        return (
+            False,
+            f"Step 3, Error: 書き起こしファイルの前処理に失敗しました:\n{message}",
+        )
     elif message:
         logger.warning(f"Step 3: preprocessing text finished with stderr.")
-        return True, f"Step 3, Success: 書き起こしファイルの前処理が完了しました:\n{message}"
+        return (
+            True,
+            f"Step 3, Success: 書き起こしファイルの前処理が完了しました:\n{message}",
+        )
     logger.success("Step 3: preprocessing text finished.")
     return True, "Step 3, Success: 書き起こしファイルの前処理が完了しました"
 
@@ -193,7 +208,10 @@ def bert_gen(model_name):
         return False, f"Step 4, Error: BERT特徴ファイルの生成に失敗しました:\n{message}"
     elif message:
         logger.warning(f"Step 4: bert_gen finished with stderr.")
-        return True, f"Step 4, Success: BERT特徴ファイルの生成が完了しました:\n{message}"
+        return (
+            True,
+            f"Step 4, Success: BERT特徴ファイルの生成が完了しました:\n{message}",
+        )
     logger.success("Step 4: bert_gen finished.")
     return True, "Step 4, Success: BERT特徴ファイルの生成が完了しました"
 
@@ -212,10 +230,16 @@ def style_gen(model_name, num_processes):
     )
     if not success:
         logger.error(f"Step 5: style_gen failed.")
-        return False, f"Step 5, Error: スタイル特徴ファイルの生成に失敗しました:\n{message}"
+        return (
+            False,
+            f"Step 5, Error: スタイル特徴ファイルの生成に失敗しました:\n{message}",
+        )
     elif message:
         logger.warning(f"Step 5: style_gen finished with stderr.")
-        return True, f"Step 5, Success: スタイル特徴ファイルの生成が完了しました:\n{message}"
+        return (
+            True,
+            f"Step 5, Success: スタイル特徴ファイルの生成が完了しました:\n{message}",
+        )
     logger.success("Step 5: style_gen finished.")
     return True, "Step 5, Success: スタイル特徴ファイルの生成が完了しました"
 
@@ -233,6 +257,7 @@ def preprocess_all(
     freeze_JP_bert,
     freeze_ZH_bert,
     freeze_style,
+    use_jp_extra,
 ):
     if model_name == "":
         return False, "Error: モデル名を入力してください"
@@ -246,13 +271,14 @@ def preprocess_all(
         freeze_JP_bert,
         freeze_ZH_bert,
         freeze_style,
+        use_jp_extra,
     )
     if not success:
         return False, message
     success, message = resample(model_name, normalize, trim, num_processes)
     if not success:
         return False, message
-    success, message = preprocess_text(model_name)
+    success, message = preprocess_text(model_name, use_jp_extra)
     if not success:
         return False, message
     success, message = bert_gen(model_name)  # bert_genは重いのでプロセス数いじらない
@@ -262,10 +288,13 @@ def preprocess_all(
     if not success:
         return False, message
     logger.success("Success: All preprocess finished!")
-    return True, "Success: 全ての前処理が完了しました。ターミナルを確認しておかしいところがないか確認するのをおすすめします。"
+    return (
+        True,
+        "Success: 全ての前処理が完了しました。ターミナルを確認しておかしいところがないか確認するのをおすすめします。",
+    )
 
 
-def train(model_name, skip_style=False):
+def train(model_name, skip_style=False, use_jp_extra=True):
     dataset_path, _, _, _, config_path = get_path(model_name)
     # 学習再開の場合は念のためconfig.ymlの名前等を更新
     with open("config.yml", "r", encoding="utf-8") as f:
@@ -274,12 +303,12 @@ def train(model_name, skip_style=False):
     yml_data["dataset_path"] = dataset_path
     with open("config.yml", "w", encoding="utf-8") as f:
         yaml.dump(yml_data, f, allow_unicode=True)
-    cmd = ["train_ms.py", "--config", config_path, "--model", dataset_path]
+
+    train_py = "train_ms.py" if not use_jp_extra else "train_ms_jp_extra.py"
+    cmd = [train_py, "--config", config_path, "--model", dataset_path]
     if skip_style:
         cmd.append("--skip_default_style")
-    success, message = run_script_with_log(
-        ["train_ms.py", "--config", config_path, "--model", dataset_path]
-    )
+    success, message = run_script_with_log(cmd)
     if not success:
         logger.error(f"Train failed.")
         return False, f"Error: 学習に失敗しました:\n{message}"
@@ -348,6 +377,10 @@ if __name__ == "__main__":
         gr.Markdown("### 自動前処理")
         with gr.Row(variant="panel"):
             with gr.Column():
+                use_jp_extra = gr.Checkbox(
+                    label="日本語特化版を使う（日本語の性能が上がるが英語と中国語は話せなくなる）",
+                    value=False,
+                )
                 batch_size = gr.Slider(
                     label="バッチサイズ",
                     value=4,
@@ -412,12 +445,18 @@ if __name__ == "__main__":
                     )
 
             with gr.Column():
-                preprocess_button = gr.Button(value="自動前処理を実行", variant="primary")
+                preprocess_button = gr.Button(
+                    value="自動前処理を実行", variant="primary"
+                )
                 info_all = gr.Textbox(label="状況")
         with gr.Accordion(open=False, label="手動前処理"):
             with gr.Row(variant="panel"):
                 with gr.Column():
                     gr.Markdown(value="#### Step 1: 設定ファイルの生成")
+                    use_jp_extra_manual = gr.Checkbox(
+                        label="日本語特化版を使う",
+                        value=False,
+                    )
                     batch_size_manual = gr.Slider(
                         label="バッチサイズ",
                         value=4,
@@ -515,6 +554,10 @@ if __name__ == "__main__":
                 info="学習再開の場合の場合はチェックしてください",
                 value=False,
             )
+            use_jp_extra_train = gr.Checkbox(
+                label="日本語特化版を使う",
+                value=True,
+            )
             train_btn = gr.Button(value="学習を開始する", variant="primary")
             info_train = gr.Textbox(label="状況")
 
@@ -533,6 +576,7 @@ if __name__ == "__main__":
                 freeze_JP_bert,
                 freeze_ZH_bert,
                 freeze_style,
+                use_jp_extra,
             ],
             outputs=[info_all],
         )
@@ -548,6 +592,7 @@ if __name__ == "__main__":
                 freeze_JP_bert_manual,
                 freeze_ZH_bert_manual,
                 freeze_style_manual,
+                use_jp_extra_manual,
             ],
             outputs=[info_init],
         )
@@ -577,7 +622,9 @@ if __name__ == "__main__":
             outputs=[info_style],
         )
         train_btn.click(
-            second_elem_of(train), inputs=[model_name, skip_style], outputs=[info_train]
+            second_elem_of(train),
+            inputs=[model_name, skip_style, use_jp_extra_train],
+            outputs=[info_train],
         )
 
     parser = argparse.ArgumentParser()