Merge pull request #18 from Stardust-minus/dev

add japanese support and other fix

.github/workflows/pull_format.yml

@@ -0,0 +1,43 @@
+name: pull format
+
+on: [pull_request]
+
+permissions:
+  contents: write
+
+jobs:
+  pull_format:
+    runs-on: ${{ matrix.os }}
+
+    strategy:
+      matrix:
+        python-version: ["3.10"]
+        os: [ubuntu-latest]
+      fail-fast: false
+
+    continue-on-error: true
+
+    steps:
+      - name: checkout
+        continue-on-error: true
+        uses: actions/checkout@v3
+        with:
+          ref: ${{ github.head_ref }}
+          fetch-depth: 0
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v4
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install Black
+        run: pip install "black[jupyter]"
+
+      - name: Run Black
+        # run: black $(git ls-files '*.py')
+        run: black .
+
+      - name: Commit Back
+        uses: stefanzweifel/git-auto-commit-action@v4
+        with:
+          commit_message: Apply Code Formatter Change

.github/workflows/push_format.yml

@@ -0,0 +1,57 @@
+name: push format
+
+on:
+  push:
+    branches:
+      - master
+      - dev
+
+permissions:
+  contents: write
+  pull-requests: write
+
+jobs:
+  push_format:
+    runs-on: ${{ matrix.os }}
+
+    strategy:
+      matrix:
+        python-version: ["3.10"]
+        os: [ubuntu-latest]
+      fail-fast: false
+
+    steps:
+      - uses: actions/checkout@v3
+        with:
+          ref: ${{github.ref_name}}
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v4
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install Black
+        run: pip install "black[jupyter]"
+
+      - name: Run Black
+        # run: black $(git ls-files '*.py')
+        run: black .
+
+      - name: Commit Back
+        continue-on-error: true
+        id: commitback
+        run: |
+          git config --local user.email "github-actions[bot]@users.noreply.github.com"
+          git config --local user.name "github-actions[bot]"
+          git add --all
+          git commit -m "Format code"
+
+      - name: Create Pull Request
+        if: steps.commitback.outcome == 'success'
+        continue-on-error: true
+        uses: peter-evans/create-pull-request@v5
+        with:
+          delete-branch: true
+          body: Apply Code Formatter Change
+          title: Apply Code Formatter Change
+          commit-message: Automatic code format

.gitignore

@@ -162,3 +162,7 @@ cython_debug/
 .DS_Store
 /models
 /logs
+
+filelists/*
+!/filelists/esd.list
+data/*

.pre-commit-config.yaml

@@ -0,0 +1,25 @@
+repos:
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v2.3.0
+    hooks:
+      - id: check-yaml
+      - id: end-of-file-fixer
+      - id: trailing-whitespace
+
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.0.280
+    hooks:
+      - id: ruff
+        args: [ --fix ]
+
+  - repo: https://github.com/psf/black
+    rev: 22.12.0
+    hooks:
+      - id: black
+
+  - repo: https://github.com/codespell-project/codespell
+    rev: v2.2.4
+    hooks:
+      - id: codespell
+        files: ^.*\.(py|md|rst|yml)$
+        args: [-L=fro]

README.md

@@ -3,5 +3,5 @@
 VITS2 Backbone with bert
 ## 成熟的旅行者/开拓者/舰长/博士/sensei/猎魔人/喵喵露/V应该参阅代码自己学习如何训练。
 ### 严禁将此项目用于一切违反《中华人民共和国宪法》，《中华人民共和国刑法》，《中华人民共和国治安管理处罚法》和《中华人民共和国民法典》之用途。
-#### Video:https://www.bilibili.com/video/BV1hp4y1K78E   
+#### Video:https://www.bilibili.com/video/BV1hp4y1K78E
 #### Demo:https://www.bilibili.com/video/BV1TF411k78w

attentions.py

@@ -1,4 +1,3 @@
-import copy
 import math
 import torch
 from torch import nn
@@ -9,336 +8,457 @@ import logging
 
 logger = logging.getLogger(__name__)
 
+
 class LayerNorm(nn.Module):
-  def __init__(self, channels, eps=1e-5):
-    super().__init__()
-    self.channels = channels
-    self.eps = eps
+    def __init__(self, channels, eps=1e-5):
+        super().__init__()
+        self.channels = channels
+        self.eps = eps
 
-    self.gamma = nn.Parameter(torch.ones(channels))
-    self.beta = nn.Parameter(torch.zeros(channels))
+        self.gamma = nn.Parameter(torch.ones(channels))
+        self.beta = nn.Parameter(torch.zeros(channels))
+
+    def forward(self, x):
+        x = x.transpose(1, -1)
+        x = F.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps)
+        return x.transpose(1, -1)
 
-  def forward(self, x):
-    x = x.transpose(1, -1)
-    x = F.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps)
-    return x.transpose(1, -1)
 
- 
-   
 @torch.jit.script
 def fused_add_tanh_sigmoid_multiply(input_a, input_b, n_channels):
-  n_channels_int = n_channels[0]
-  in_act = input_a + input_b
-  t_act = torch.tanh(in_act[:, :n_channels_int, :])
-  s_act = torch.sigmoid(in_act[:, n_channels_int:, :])
-  acts = t_act * s_act
-  return acts
+    n_channels_int = n_channels[0]
+    in_act = input_a + input_b
+    t_act = torch.tanh(in_act[:, :n_channels_int, :])
+    s_act = torch.sigmoid(in_act[:, n_channels_int:, :])
+    acts = t_act * s_act
+    return acts
+
 
 class Encoder(nn.Module):
-   def __init__(self, hidden_channels, filter_channels, n_heads, n_layers, kernel_size=1, p_dropout=0., window_size=4, isflow = True, **kwargs):
-    super().__init__()
-    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.window_size = window_size
-    #if isflow:
-    #  cond_layer = torch.nn.Conv1d(256, 2*hidden_channels*n_layers, 1)
-    #  self.cond_pre = torch.nn.Conv1d(hidden_channels, 2*hidden_channels, 1)
-    #  self.cond_layer = weight_norm(cond_layer, name='weight')
-    #  self.gin_channels = 256
-    self.cond_layer_idx = self.n_layers
-    if 'gin_channels' in kwargs:
-      self.gin_channels = kwargs['gin_channels']
-      if self.gin_channels != 0:
-        self.spk_emb_linear = nn.Linear(self.gin_channels, self.hidden_channels)
-        # vits2 says 3rd block, so idx is 2 by default
-        self.cond_layer_idx = kwargs['cond_layer_idx'] if 'cond_layer_idx' in kwargs else 2
-        logging.debug(self.gin_channels, self.cond_layer_idx)
-        assert self.cond_layer_idx < self.n_layers, 'cond_layer_idx should be less than n_layers'
-    self.drop = nn.Dropout(p_dropout)
-    self.attn_layers = nn.ModuleList()
-    self.norm_layers_1 = nn.ModuleList()
-    self.ffn_layers = nn.ModuleList()
-    self.norm_layers_2 = nn.ModuleList()
-    for i in range(self.n_layers):
-      self.attn_layers.append(MultiHeadAttention(hidden_channels, hidden_channels, n_heads, p_dropout=p_dropout, window_size=window_size))
-      self.norm_layers_1.append(LayerNorm(hidden_channels))
-      self.ffn_layers.append(FFN(hidden_channels, hidden_channels, filter_channels, kernel_size, p_dropout=p_dropout))
-      self.norm_layers_2.append(LayerNorm(hidden_channels))
-   def forward(self, x, x_mask, g=None):
-    attn_mask = x_mask.unsqueeze(2) * x_mask.unsqueeze(-1)
-    x = x * x_mask
-    for i in range(self.n_layers):
-      if i == self.cond_layer_idx and g is not None:
-        g = self.spk_emb_linear(g.transpose(1, 2))
-        g = g.transpose(1, 2)
-        x = x + g
-        x = x * x_mask
-      y = self.attn_layers[i](x, x, attn_mask)
-      y = self.drop(y)
-      x = self.norm_layers_1[i](x + y)
+    def __init__(
+        self,
+        hidden_channels,
+        filter_channels,
+        n_heads,
+        n_layers,
+        kernel_size=1,
+        p_dropout=0.0,
+        window_size=4,
+        isflow=True,
+        **kwargs
+    ):
+        super().__init__()
+        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.window_size = window_size
+        # if isflow:
+        #  cond_layer = torch.nn.Conv1d(256, 2*hidden_channels*n_layers, 1)
+        #  self.cond_pre = torch.nn.Conv1d(hidden_channels, 2*hidden_channels, 1)
+        #  self.cond_layer = weight_norm(cond_layer, name='weight')
+        #  self.gin_channels = 256
+        self.cond_layer_idx = self.n_layers
+        if "gin_channels" in kwargs:
+            self.gin_channels = kwargs["gin_channels"]
+            if self.gin_channels != 0:
+                self.spk_emb_linear = nn.Linear(self.gin_channels, self.hidden_channels)
+                # vits2 says 3rd block, so idx is 2 by default
+                self.cond_layer_idx = (
+                    kwargs["cond_layer_idx"] if "cond_layer_idx" in kwargs else 2
+                )
+                logging.debug(self.gin_channels, self.cond_layer_idx)
+                assert (
+                    self.cond_layer_idx < self.n_layers
+                ), "cond_layer_idx should be less than n_layers"
+        self.drop = nn.Dropout(p_dropout)
+        self.attn_layers = nn.ModuleList()
+        self.norm_layers_1 = nn.ModuleList()
+        self.ffn_layers = nn.ModuleList()
+        self.norm_layers_2 = nn.ModuleList()
+        for i in range(self.n_layers):
+            self.attn_layers.append(
+                MultiHeadAttention(
+                    hidden_channels,
+                    hidden_channels,
+                    n_heads,
+                    p_dropout=p_dropout,
+                    window_size=window_size,
+                )
+            )
+            self.norm_layers_1.append(LayerNorm(hidden_channels))
+            self.ffn_layers.append(
+                FFN(
+                    hidden_channels,
+                    hidden_channels,
+                    filter_channels,
+                    kernel_size,
+                    p_dropout=p_dropout,
+                )
+            )
+            self.norm_layers_2.append(LayerNorm(hidden_channels))
 
-      y = self.ffn_layers[i](x, x_mask)
-      y = self.drop(y)
-      x = self.norm_layers_2[i](x + y)
-    x = x * x_mask
-    return x
+    def forward(self, x, x_mask, g=None):
+        attn_mask = x_mask.unsqueeze(2) * x_mask.unsqueeze(-1)
+        x = x * x_mask
+        for i in range(self.n_layers):
+            if i == self.cond_layer_idx and g is not None:
+                g = self.spk_emb_linear(g.transpose(1, 2))
+                g = g.transpose(1, 2)
+                x = x + g
+                x = x * x_mask
+            y = self.attn_layers[i](x, x, attn_mask)
+            y = self.drop(y)
+            x = self.norm_layers_1[i](x + y)
+
+            y = self.ffn_layers[i](x, x_mask)
+            y = self.drop(y)
+            x = self.norm_layers_2[i](x + y)
+        x = x * x_mask
+        return x
 
 
 class Decoder(nn.Module):
-  def __init__(self, hidden_channels, filter_channels, n_heads, n_layers, kernel_size=1, p_dropout=0., proximal_bias=False, proximal_init=True, **kwargs):
-    super().__init__()
-    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.proximal_bias = proximal_bias
-    self.proximal_init = proximal_init
+    def __init__(
+        self,
+        hidden_channels,
+        filter_channels,
+        n_heads,
+        n_layers,
+        kernel_size=1,
+        p_dropout=0.0,
+        proximal_bias=False,
+        proximal_init=True,
+        **kwargs
+    ):
+        super().__init__()
+        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.proximal_bias = proximal_bias
+        self.proximal_init = proximal_init
 
-    self.drop = nn.Dropout(p_dropout)
-    self.self_attn_layers = nn.ModuleList()
-    self.norm_layers_0 = nn.ModuleList()
-    self.encdec_attn_layers = nn.ModuleList()
-    self.norm_layers_1 = nn.ModuleList()
-    self.ffn_layers = nn.ModuleList()
-    self.norm_layers_2 = nn.ModuleList()
-    for i in range(self.n_layers):
-      self.self_attn_layers.append(MultiHeadAttention(hidden_channels, hidden_channels, n_heads, p_dropout=p_dropout, proximal_bias=proximal_bias, proximal_init=proximal_init))
-      self.norm_layers_0.append(LayerNorm(hidden_channels))
-      self.encdec_attn_layers.append(MultiHeadAttention(hidden_channels, hidden_channels, n_heads, p_dropout=p_dropout))
-      self.norm_layers_1.append(LayerNorm(hidden_channels))
-      self.ffn_layers.append(FFN(hidden_channels, hidden_channels, filter_channels, kernel_size, p_dropout=p_dropout, causal=True))
-      self.norm_layers_2.append(LayerNorm(hidden_channels))
+        self.drop = nn.Dropout(p_dropout)
+        self.self_attn_layers = nn.ModuleList()
+        self.norm_layers_0 = nn.ModuleList()
+        self.encdec_attn_layers = nn.ModuleList()
+        self.norm_layers_1 = nn.ModuleList()
+        self.ffn_layers = nn.ModuleList()
+        self.norm_layers_2 = nn.ModuleList()
+        for i in range(self.n_layers):
+            self.self_attn_layers.append(
+                MultiHeadAttention(
+                    hidden_channels,
+                    hidden_channels,
+                    n_heads,
+                    p_dropout=p_dropout,
+                    proximal_bias=proximal_bias,
+                    proximal_init=proximal_init,
+                )
+            )
+            self.norm_layers_0.append(LayerNorm(hidden_channels))
+            self.encdec_attn_layers.append(
+                MultiHeadAttention(
+                    hidden_channels, hidden_channels, n_heads, p_dropout=p_dropout
+                )
+            )
+            self.norm_layers_1.append(LayerNorm(hidden_channels))
+            self.ffn_layers.append(
+                FFN(
+                    hidden_channels,
+                    hidden_channels,
+                    filter_channels,
+                    kernel_size,
+                    p_dropout=p_dropout,
+                    causal=True,
+                )
+            )
+            self.norm_layers_2.append(LayerNorm(hidden_channels))
 
-  def forward(self, x, x_mask, h, h_mask):
-    """
-    x: decoder input
-    h: encoder output
-    """
-    self_attn_mask = commons.subsequent_mask(x_mask.size(2)).to(device=x.device, dtype=x.dtype)
-    encdec_attn_mask = h_mask.unsqueeze(2) * x_mask.unsqueeze(-1)
-    x = x * x_mask
-    for i in range(self.n_layers):
-      y = self.self_attn_layers[i](x, x, self_attn_mask)
-      y = self.drop(y)
-      x = self.norm_layers_0[i](x + y)
+    def forward(self, x, x_mask, h, h_mask):
+        """
+        x: decoder input
+        h: encoder output
+        """
+        self_attn_mask = commons.subsequent_mask(x_mask.size(2)).to(
+            device=x.device, dtype=x.dtype
+        )
+        encdec_attn_mask = h_mask.unsqueeze(2) * x_mask.unsqueeze(-1)
+        x = x * x_mask
+        for i in range(self.n_layers):
+            y = self.self_attn_layers[i](x, x, self_attn_mask)
+            y = self.drop(y)
+            x = self.norm_layers_0[i](x + y)
 
-      y = self.encdec_attn_layers[i](x, h, encdec_attn_mask)
-      y = self.drop(y)
-      x = self.norm_layers_1[i](x + y)
-      
-      y = self.ffn_layers[i](x, x_mask)
-      y = self.drop(y)
-      x = self.norm_layers_2[i](x + y)
-    x = x * x_mask
-    return x
+            y = self.encdec_attn_layers[i](x, h, encdec_attn_mask)
+            y = self.drop(y)
+            x = self.norm_layers_1[i](x + y)
+
+            y = self.ffn_layers[i](x, x_mask)
+            y = self.drop(y)
+            x = self.norm_layers_2[i](x + y)
+        x = x * x_mask
+        return x
 
 
 class MultiHeadAttention(nn.Module):
-  def __init__(self, channels, out_channels, n_heads, p_dropout=0., window_size=None, heads_share=True, block_length=None, proximal_bias=False, proximal_init=False):
-    super().__init__()
-    assert channels % n_heads == 0
+    def __init__(
+        self,
+        channels,
+        out_channels,
+        n_heads,
+        p_dropout=0.0,
+        window_size=None,
+        heads_share=True,
+        block_length=None,
+        proximal_bias=False,
+        proximal_init=False,
+    ):
+        super().__init__()
+        assert channels % n_heads == 0
 
-    self.channels = channels
-    self.out_channels = out_channels
-    self.n_heads = n_heads
-    self.p_dropout = p_dropout
-    self.window_size = window_size
-    self.heads_share = heads_share
-    self.block_length = block_length
-    self.proximal_bias = proximal_bias
-    self.proximal_init = proximal_init
-    self.attn = None
+        self.channels = channels
+        self.out_channels = out_channels
+        self.n_heads = n_heads
+        self.p_dropout = p_dropout
+        self.window_size = window_size
+        self.heads_share = heads_share
+        self.block_length = block_length
+        self.proximal_bias = proximal_bias
+        self.proximal_init = proximal_init
+        self.attn = None
 
-    self.k_channels = channels // n_heads
-    self.conv_q = nn.Conv1d(channels, channels, 1)
-    self.conv_k = nn.Conv1d(channels, channels, 1)
-    self.conv_v = nn.Conv1d(channels, channels, 1)
-    self.conv_o = nn.Conv1d(channels, out_channels, 1)
-    self.drop = nn.Dropout(p_dropout)
+        self.k_channels = channels // n_heads
+        self.conv_q = nn.Conv1d(channels, channels, 1)
+        self.conv_k = nn.Conv1d(channels, channels, 1)
+        self.conv_v = nn.Conv1d(channels, channels, 1)
+        self.conv_o = nn.Conv1d(channels, out_channels, 1)
+        self.drop = nn.Dropout(p_dropout)
 
-    if window_size is not None:
-      n_heads_rel = 1 if heads_share else n_heads
-      rel_stddev = self.k_channels**-0.5
-      self.emb_rel_k = nn.Parameter(torch.randn(n_heads_rel, window_size * 2 + 1, self.k_channels) * rel_stddev)
-      self.emb_rel_v = nn.Parameter(torch.randn(n_heads_rel, window_size * 2 + 1, self.k_channels) * rel_stddev)
+        if window_size is not None:
+            n_heads_rel = 1 if heads_share else n_heads
+            rel_stddev = self.k_channels**-0.5
+            self.emb_rel_k = nn.Parameter(
+                torch.randn(n_heads_rel, window_size * 2 + 1, self.k_channels)
+                * rel_stddev
+            )
+            self.emb_rel_v = nn.Parameter(
+                torch.randn(n_heads_rel, window_size * 2 + 1, self.k_channels)
+                * rel_stddev
+            )
 
-    nn.init.xavier_uniform_(self.conv_q.weight)
-    nn.init.xavier_uniform_(self.conv_k.weight)
-    nn.init.xavier_uniform_(self.conv_v.weight)
-    if proximal_init:
-      with torch.no_grad():
-        self.conv_k.weight.copy_(self.conv_q.weight)
-        self.conv_k.bias.copy_(self.conv_q.bias)
-      
-  def forward(self, x, c, attn_mask=None):
-    q = self.conv_q(x)
-    k = self.conv_k(c)
-    v = self.conv_v(c)
-    
-    x, self.attn = self.attention(q, k, v, mask=attn_mask)
+        nn.init.xavier_uniform_(self.conv_q.weight)
+        nn.init.xavier_uniform_(self.conv_k.weight)
+        nn.init.xavier_uniform_(self.conv_v.weight)
+        if proximal_init:
+            with torch.no_grad():
+                self.conv_k.weight.copy_(self.conv_q.weight)
+                self.conv_k.bias.copy_(self.conv_q.bias)
 
-    x = self.conv_o(x)
-    return x
+    def forward(self, x, c, attn_mask=None):
+        q = self.conv_q(x)
+        k = self.conv_k(c)
+        v = self.conv_v(c)
 
-  def attention(self, query, key, value, mask=None):
-    # reshape [b, d, t] -> [b, n_h, t, d_k]
-    b, d, t_s, t_t = (*key.size(), query.size(2))
-    query = query.view(b, self.n_heads, self.k_channels, t_t).transpose(2, 3)
-    key = key.view(b, self.n_heads, self.k_channels, t_s).transpose(2, 3)
-    value = value.view(b, self.n_heads, self.k_channels, t_s).transpose(2, 3)
+        x, self.attn = self.attention(q, k, v, mask=attn_mask)
 
-    scores = torch.matmul(query / math.sqrt(self.k_channels), key.transpose(-2, -1))
-    if self.window_size is not None:
-      assert t_s == t_t, "Relative attention is only available for self-attention."
-      key_relative_embeddings = self._get_relative_embeddings(self.emb_rel_k, t_s)
-      rel_logits = self._matmul_with_relative_keys(query /math.sqrt(self.k_channels), key_relative_embeddings)
-      scores_local = self._relative_position_to_absolute_position(rel_logits)
-      scores = scores + scores_local
-    if self.proximal_bias:
-      assert t_s == t_t, "Proximal bias is only available for self-attention."
-      scores = scores + self._attention_bias_proximal(t_s).to(device=scores.device, dtype=scores.dtype)
-    if mask is not None:
-      scores = scores.masked_fill(mask == 0, -1e4)
-      if self.block_length is not None:
-        assert t_s == t_t, "Local attention is only available for self-attention."
-        block_mask = torch.ones_like(scores).triu(-self.block_length).tril(self.block_length)
-        scores = scores.masked_fill(block_mask == 0, -1e4)
-    p_attn = F.softmax(scores, dim=-1) # [b, n_h, t_t, t_s]
-    p_attn = self.drop(p_attn)
-    output = torch.matmul(p_attn, value)
-    if self.window_size is not None:
-      relative_weights = self._absolute_position_to_relative_position(p_attn)
-      value_relative_embeddings = self._get_relative_embeddings(self.emb_rel_v, t_s)
-      output = output + self._matmul_with_relative_values(relative_weights, value_relative_embeddings)
-    output = output.transpose(2, 3).contiguous().view(b, d, t_t) # [b, n_h, t_t, d_k] -> [b, d, t_t]
-    return output, p_attn
+        x = self.conv_o(x)
+        return x
 
-  def _matmul_with_relative_values(self, x, y):
-    """
-    x: [b, h, l, m]
-    y: [h or 1, m, d]
-    ret: [b, h, l, d]
-    """
-    ret = torch.matmul(x, y.unsqueeze(0))
-    return ret
+    def attention(self, query, key, value, mask=None):
+        # reshape [b, d, t] -> [b, n_h, t, d_k]
+        b, d, t_s, t_t = (*key.size(), query.size(2))
+        query = query.view(b, self.n_heads, self.k_channels, t_t).transpose(2, 3)
+        key = key.view(b, self.n_heads, self.k_channels, t_s).transpose(2, 3)
+        value = value.view(b, self.n_heads, self.k_channels, t_s).transpose(2, 3)
 
-  def _matmul_with_relative_keys(self, x, y):
-    """
-    x: [b, h, l, d]
-    y: [h or 1, m, d]
-    ret: [b, h, l, m]
-    """
-    ret = torch.matmul(x, y.unsqueeze(0).transpose(-2, -1))
-    return ret
+        scores = torch.matmul(query / math.sqrt(self.k_channels), key.transpose(-2, -1))
+        if self.window_size is not None:
+            assert (
+                t_s == t_t
+            ), "Relative attention is only available for self-attention."
+            key_relative_embeddings = self._get_relative_embeddings(self.emb_rel_k, t_s)
+            rel_logits = self._matmul_with_relative_keys(
+                query / math.sqrt(self.k_channels), key_relative_embeddings
+            )
+            scores_local = self._relative_position_to_absolute_position(rel_logits)
+            scores = scores + scores_local
+        if self.proximal_bias:
+            assert t_s == t_t, "Proximal bias is only available for self-attention."
+            scores = scores + self._attention_bias_proximal(t_s).to(
+                device=scores.device, dtype=scores.dtype
+            )
+        if mask is not None:
+            scores = scores.masked_fill(mask == 0, -1e4)
+            if self.block_length is not None:
+                assert (
+                    t_s == t_t
+                ), "Local attention is only available for self-attention."
+                block_mask = (
+                    torch.ones_like(scores)
+                    .triu(-self.block_length)
+                    .tril(self.block_length)
+                )
+                scores = scores.masked_fill(block_mask == 0, -1e4)
+        p_attn = F.softmax(scores, dim=-1)  # [b, n_h, t_t, t_s]
+        p_attn = self.drop(p_attn)
+        output = torch.matmul(p_attn, value)
+        if self.window_size is not None:
+            relative_weights = self._absolute_position_to_relative_position(p_attn)
+            value_relative_embeddings = self._get_relative_embeddings(
+                self.emb_rel_v, t_s
+            )
+            output = output + self._matmul_with_relative_values(
+                relative_weights, value_relative_embeddings
+            )
+        output = (
+            output.transpose(2, 3).contiguous().view(b, d, t_t)
+        )  # [b, n_h, t_t, d_k] -> [b, d, t_t]
+        return output, p_attn
 
-  def _get_relative_embeddings(self, relative_embeddings, length):
-    max_relative_position = 2 * self.window_size + 1
-    # Pad first before slice to avoid using cond ops.
-    pad_length = max(length - (self.window_size + 1), 0)
-    slice_start_position = max((self.window_size + 1) - length, 0)
-    slice_end_position = slice_start_position + 2 * length - 1
-    if pad_length > 0:
-      padded_relative_embeddings = F.pad(
-          relative_embeddings,
-          commons.convert_pad_shape([[0, 0], [pad_length, pad_length], [0, 0]]))
-    else:
-      padded_relative_embeddings = relative_embeddings
-    used_relative_embeddings = padded_relative_embeddings[:,slice_start_position:slice_end_position]
-    return used_relative_embeddings
+    def _matmul_with_relative_values(self, x, y):
+        """
+        x: [b, h, l, m]
+        y: [h or 1, m, d]
+        ret: [b, h, l, d]
+        """
+        ret = torch.matmul(x, y.unsqueeze(0))
+        return ret
 
-  def _relative_position_to_absolute_position(self, x):
-    """
-    x: [b, h, l, 2*l-1]
-    ret: [b, h, l, l]
-    """
-    batch, heads, length, _ = x.size()
-    # Concat columns of pad to shift from relative to absolute indexing.
-    x = F.pad(x, commons.convert_pad_shape([[0,0],[0,0],[0,0],[0,1]]))
+    def _matmul_with_relative_keys(self, x, y):
+        """
+        x: [b, h, l, d]
+        y: [h or 1, m, d]
+        ret: [b, h, l, m]
+        """
+        ret = torch.matmul(x, y.unsqueeze(0).transpose(-2, -1))
+        return ret
 
-    # Concat extra elements so to add up to shape (len+1, 2*len-1).
-    x_flat = x.view([batch, heads, length * 2 * length])
-    x_flat = F.pad(x_flat, commons.convert_pad_shape([[0,0],[0,0],[0,length-1]]))
+    def _get_relative_embeddings(self, relative_embeddings, length):
+        2 * self.window_size + 1
+        # Pad first before slice to avoid using cond ops.
+        pad_length = max(length - (self.window_size + 1), 0)
+        slice_start_position = max((self.window_size + 1) - length, 0)
+        slice_end_position = slice_start_position + 2 * length - 1
+        if pad_length > 0:
+            padded_relative_embeddings = F.pad(
+                relative_embeddings,
+                commons.convert_pad_shape([[0, 0], [pad_length, pad_length], [0, 0]]),
+            )
+        else:
+            padded_relative_embeddings = relative_embeddings
+        used_relative_embeddings = padded_relative_embeddings[
+            :, slice_start_position:slice_end_position
+        ]
+        return used_relative_embeddings
 
-    # Reshape and slice out the padded elements.
-    x_final = x_flat.view([batch, heads, length+1, 2*length-1])[:, :, :length, length-1:]
-    return x_final
+    def _relative_position_to_absolute_position(self, x):
+        """
+        x: [b, h, l, 2*l-1]
+        ret: [b, h, l, l]
+        """
+        batch, heads, length, _ = x.size()
+        # Concat columns of pad to shift from relative to absolute indexing.
+        x = F.pad(x, commons.convert_pad_shape([[0, 0], [0, 0], [0, 0], [0, 1]]))
 
-  def _absolute_position_to_relative_position(self, x):
-    """
-    x: [b, h, l, l]
-    ret: [b, h, l, 2*l-1]
-    """
-    batch, heads, length, _ = x.size()
-    # padd along column
-    x = F.pad(x, commons.convert_pad_shape([[0, 0], [0, 0], [0, 0], [0, length-1]]))
-    x_flat = x.view([batch, heads, length**2 + length*(length -1)])
-    # add 0's in the beginning that will skew the elements after reshape
-    x_flat = F.pad(x_flat, commons.convert_pad_shape([[0, 0], [0, 0], [length, 0]]))
-    x_final = x_flat.view([batch, heads, length, 2*length])[:,:,:,1:]
-    return x_final
+        # Concat extra elements so to add up to shape (len+1, 2*len-1).
+        x_flat = x.view([batch, heads, length * 2 * length])
+        x_flat = F.pad(
+            x_flat, commons.convert_pad_shape([[0, 0], [0, 0], [0, length - 1]])
+        )
 
-  def _attention_bias_proximal(self, length):
-    """Bias for self-attention to encourage attention to close positions.
-    Args:
-      length: an integer scalar.
-    Returns:
-      a Tensor with shape [1, 1, length, length]
-    """
-    r = torch.arange(length, dtype=torch.float32)
-    diff = torch.unsqueeze(r, 0) - torch.unsqueeze(r, 1)
-    return torch.unsqueeze(torch.unsqueeze(-torch.log1p(torch.abs(diff)), 0), 0)
+        # Reshape and slice out the padded elements.
+        x_final = x_flat.view([batch, heads, length + 1, 2 * length - 1])[
+            :, :, :length, length - 1 :
+        ]
+        return x_final
+
+    def _absolute_position_to_relative_position(self, x):
+        """
+        x: [b, h, l, l]
+        ret: [b, h, l, 2*l-1]
+        """
+        batch, heads, length, _ = x.size()
+        # pad along column
+        x = F.pad(
+            x, commons.convert_pad_shape([[0, 0], [0, 0], [0, 0], [0, length - 1]])
+        )
+        x_flat = x.view([batch, heads, length**2 + length * (length - 1)])
+        # add 0's in the beginning that will skew the elements after reshape
+        x_flat = F.pad(x_flat, commons.convert_pad_shape([[0, 0], [0, 0], [length, 0]]))
+        x_final = x_flat.view([batch, heads, length, 2 * length])[:, :, :, 1:]
+        return x_final
+
+    def _attention_bias_proximal(self, length):
+        """Bias for self-attention to encourage attention to close positions.
+        Args:
+          length: an integer scalar.
+        Returns:
+          a Tensor with shape [1, 1, length, length]
+        """
+        r = torch.arange(length, dtype=torch.float32)
+        diff = torch.unsqueeze(r, 0) - torch.unsqueeze(r, 1)
+        return torch.unsqueeze(torch.unsqueeze(-torch.log1p(torch.abs(diff)), 0), 0)
 
 
 class FFN(nn.Module):
-  def __init__(self, in_channels, out_channels, filter_channels, kernel_size, p_dropout=0., activation=None, causal=False):
-    super().__init__()
-    self.in_channels = in_channels
-    self.out_channels = out_channels
-    self.filter_channels = filter_channels
-    self.kernel_size = kernel_size
-    self.p_dropout = p_dropout
-    self.activation = activation
-    self.causal = causal
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        filter_channels,
+        kernel_size,
+        p_dropout=0.0,
+        activation=None,
+        causal=False,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.filter_channels = filter_channels
+        self.kernel_size = kernel_size
+        self.p_dropout = p_dropout
+        self.activation = activation
+        self.causal = causal
 
-    if causal:
-      self.padding = self._causal_padding
-    else:
-      self.padding = self._same_padding
+        if causal:
+            self.padding = self._causal_padding
+        else:
+            self.padding = self._same_padding
 
-    self.conv_1 = nn.Conv1d(in_channels, filter_channels, kernel_size)
-    self.conv_2 = nn.Conv1d(filter_channels, out_channels, kernel_size)
-    self.drop = nn.Dropout(p_dropout)
+        self.conv_1 = nn.Conv1d(in_channels, filter_channels, kernel_size)
+        self.conv_2 = nn.Conv1d(filter_channels, out_channels, kernel_size)
+        self.drop = nn.Dropout(p_dropout)
 
-  def forward(self, x, x_mask):
-    x = self.conv_1(self.padding(x * x_mask))
-    if self.activation == "gelu":
-      x = x * torch.sigmoid(1.702 * x)
-    else:
-      x = torch.relu(x)
-    x = self.drop(x)
-    x = self.conv_2(self.padding(x * x_mask))
-    return x * x_mask
-  
-  def _causal_padding(self, x):
-    if self.kernel_size == 1:
-      return x
-    pad_l = self.kernel_size - 1
-    pad_r = 0
-    padding = [[0, 0], [0, 0], [pad_l, pad_r]]
-    x = F.pad(x, commons.convert_pad_shape(padding))
-    return x
+    def forward(self, x, x_mask):
+        x = self.conv_1(self.padding(x * x_mask))
+        if self.activation == "gelu":
+            x = x * torch.sigmoid(1.702 * x)
+        else:
+            x = torch.relu(x)
+        x = self.drop(x)
+        x = self.conv_2(self.padding(x * x_mask))
+        return x * x_mask
 
-  def _same_padding(self, x):
-    if self.kernel_size == 1:
-      return x
-    pad_l = (self.kernel_size - 1) // 2
-    pad_r = self.kernel_size // 2
-    padding = [[0, 0], [0, 0], [pad_l, pad_r]]
-    x = F.pad(x, commons.convert_pad_shape(padding))
-    return x
+    def _causal_padding(self, x):
+        if self.kernel_size == 1:
+            return x
+        pad_l = self.kernel_size - 1
+        pad_r = 0
+        padding = [[0, 0], [0, 0], [pad_l, pad_r]]
+        x = F.pad(x, commons.convert_pad_shape(padding))
+        return x
+
+    def _same_padding(self, x):
+        if self.kernel_size == 1:
+            return x
+        pad_l = (self.kernel_size - 1) // 2
+        pad_r = self.kernel_size // 2
+        padding = [[0, 0], [0, 0], [pad_l, pad_r]]
+        x = F.pad(x, commons.convert_pad_shape(padding))
+        return x

bert/bert-base-japanese-v3/README.md

@@ -0,0 +1,53 @@
+---
+license: apache-2.0
+datasets:
+- cc100
+- wikipedia
+language:
+- ja
+widget:
+- text: 東北大学で[MASK]の研究をしています。
+---
+
+# BERT base Japanese (unidic-lite with whole word masking, CC-100 and jawiki-20230102)
+
+This is a [BERT](https://github.com/google-research/bert) model pretrained on texts in the Japanese language.
+
+This version of the model processes input texts with word-level tokenization based on the Unidic 2.1.2 dictionary (available in [unidic-lite](https://pypi.org/project/unidic-lite/) package), followed by the WordPiece subword tokenization.
+Additionally, the model is trained with the whole word masking enabled for the masked language modeling (MLM) objective.
+
+The codes for the pretraining are available at [cl-tohoku/bert-japanese](https://github.com/cl-tohoku/bert-japanese/).
+
+## Model architecture
+
+The model architecture is the same as the original BERT base model; 12 layers, 768 dimensions of hidden states, and 12 attention heads.
+
+## Training Data
+
+The model is trained on the Japanese portion of [CC-100 dataset](https://data.statmt.org/cc-100/) and the Japanese version of Wikipedia.
+For Wikipedia, we generated a text corpus from the [Wikipedia Cirrussearch dump file](https://dumps.wikimedia.org/other/cirrussearch/) as of January 2, 2023.
+The corpus files generated from CC-100 and Wikipedia are 74.3GB and 4.9GB in size and consist of approximately 392M and 34M sentences, respectively.
+
+For the purpose of splitting texts into sentences, we used [fugashi](https://github.com/polm/fugashi) with [mecab-ipadic-NEologd](https://github.com/neologd/mecab-ipadic-neologd) dictionary (v0.0.7).
+
+## Tokenization
+
+The texts are first tokenized by MeCab with the Unidic 2.1.2 dictionary and then split into subwords by the WordPiece algorithm.
+The vocabulary size is 32768.
+
+We used [fugashi](https://github.com/polm/fugashi) and [unidic-lite](https://github.com/polm/unidic-lite) packages for the tokenization.
+
+## Training
+
+We trained the model first on the CC-100 corpus for 1M steps and then on the Wikipedia corpus for another 1M steps.
+For training of the MLM (masked language modeling) objective, we introduced whole word masking in which all of the subword tokens corresponding to a single word (tokenized by MeCab) are masked at once.
+
+For training of each model, we used a v3-8 instance of Cloud TPUs provided by [TPU Research Cloud](https://sites.research.google/trc/about/).
+
+## Licenses
+
+The pretrained models are distributed under the Apache License 2.0.
+
+## Acknowledgments
+
+This model is trained with Cloud TPUs provided by [TPU Research Cloud](https://sites.research.google/trc/about/) program.

bert/bert-base-japanese-v3/config.json

@@ -0,0 +1,19 @@
+{
+    "architectures": [
+        "BertForPreTraining"
+    ],
+    "attention_probs_dropout_prob": 0.1,
+    "hidden_act": "gelu",
+    "hidden_dropout_prob": 0.1,
+    "hidden_size": 768,
+    "initializer_range": 0.02,
+    "intermediate_size": 3072,
+    "layer_norm_eps": 1e-12,
+    "max_position_embeddings": 512,
+    "model_type": "bert",
+    "num_attention_heads": 12,
+    "num_hidden_layers": 12,
+    "pad_token_id": 0,
+    "type_vocab_size": 2,
+    "vocab_size": 32768
+}

bert_gen.py

@@ -1,23 +1,23 @@
 import torch
-from torch.utils.data import DataLoader
 from multiprocessing import Pool
 import commons
 import utils
-from data_utils import TextAudioSpeakerLoader, TextAudioSpeakerCollate
 from tqdm import tqdm
-import warnings
-
 from text import cleaned_text_to_sequence, get_bert
+import argparse
+import torch.multiprocessing as mp
 
-config_path = 'configs/config.json'
-hps = utils.get_hparams_from_file(config_path)
 
 def process_line(line):
-    _id, spk, language_str, text, phones, tone, word2ph = line.strip().split("|")
+    rank = mp.current_process()._identity
+    rank = rank[0] if len(rank) > 0 else 0
+    if torch.cuda.is_available():
+        gpu_id = rank % torch.cuda.device_count()
+        device = torch.device(f"cuda:{gpu_id}")
+    wav_path, _, language_str, text, phones, tone, word2ph = line.strip().split("|")
     phone = phones.split(" ")
     tone = [int(i) for i in tone.split(" ")]
     word2ph = [int(i) for i in word2ph.split(" ")]
-    w2pho = [i for i in word2ph]
     word2ph = [i for i in word2ph]
     phone, tone, language = cleaned_text_to_sequence(phone, tone, language_str)
 
@@ -28,26 +28,33 @@ def process_line(line):
         for i in range(len(word2ph)):
             word2ph[i] = word2ph[i] * 2
         word2ph[0] += 1
-    wav_path = f'{_id}'
 
     bert_path = wav_path.replace(".wav", ".bert.pt")
+
     try:
         bert = torch.load(bert_path)
         assert bert.shape[-1] == len(phone)
-    except:
-        bert = get_bert(text, word2ph, language_str)
+    except Exception:
+        bert = get_bert(text, word2ph, language_str, device)
         assert bert.shape[-1] == len(phone)
         torch.save(bert, bert_path)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-c", "--config", type=str, default="configs/config.json")
+    parser.add_argument("--num_processes", type=int, default=2)
+    args = parser.parse_args()
+    config_path = args.config
+    hps = utils.get_hparams_from_file(config_path)
     lines = []
-    with open(hps.data.training_files, encoding='utf-8' ) as f:
+    with open(hps.data.training_files, encoding="utf-8") as f:
         lines.extend(f.readlines())
 
-    with open(hps.data.validation_files, encoding='utf-8' ) as f:
+    with open(hps.data.validation_files, encoding="utf-8") as f:
         lines.extend(f.readlines())
 
-    with Pool(processes=12) as pool: #A100 40GB suitable config,if coom,please decrease the processess number.
-        for _ in tqdm(pool.imap_unordered(process_line, lines)):
+    num_processes = args.num_processes
+    with Pool(processes=num_processes) as pool:
+        for _ in tqdm(pool.imap_unordered(process_line, lines), total=len(lines)):
             pass

commons.py

@@ -1,161 +1,160 @@
 import math
-import numpy as np
 import torch
-from torch import nn
 from torch.nn import functional as F
 
 
 def init_weights(m, mean=0.0, std=0.01):
-  classname = m.__class__.__name__
-  if classname.find("Conv") != -1:
-    m.weight.data.normal_(mean, std)
+    classname = m.__class__.__name__
+    if classname.find("Conv") != -1:
+        m.weight.data.normal_(mean, std)
 
 
 def get_padding(kernel_size, dilation=1):
-  return int((kernel_size*dilation - dilation)/2)
+    return int((kernel_size * dilation - dilation) / 2)
 
 
 def convert_pad_shape(pad_shape):
-  l = pad_shape[::-1]
-  pad_shape = [item for sublist in l for item in sublist]
-  return pad_shape
+    layer = pad_shape[::-1]
+    pad_shape = [item for sublist in layer for item in sublist]
+    return pad_shape
 
 
 def intersperse(lst, item):
-  result = [item] * (len(lst) * 2 + 1)
-  result[1::2] = lst
-  return result
+    result = [item] * (len(lst) * 2 + 1)
+    result[1::2] = lst
+    return result
 
 
 def kl_divergence(m_p, logs_p, m_q, logs_q):
-  """KL(P||Q)"""
-  kl = (logs_q - logs_p) - 0.5
-  kl += 0.5 * (torch.exp(2. * logs_p) + ((m_p - m_q)**2)) * torch.exp(-2. * logs_q)
-  return kl
+    """KL(P||Q)"""
+    kl = (logs_q - logs_p) - 0.5
+    kl += (
+        0.5 * (torch.exp(2.0 * logs_p) + ((m_p - m_q) ** 2)) * torch.exp(-2.0 * logs_q)
+    )
+    return kl
 
 
 def rand_gumbel(shape):
-  """Sample from the Gumbel distribution, protect from overflows."""
-  uniform_samples = torch.rand(shape) * 0.99998 + 0.00001
-  return -torch.log(-torch.log(uniform_samples))
+    """Sample from the Gumbel distribution, protect from overflows."""
+    uniform_samples = torch.rand(shape) * 0.99998 + 0.00001
+    return -torch.log(-torch.log(uniform_samples))
 
 
 def rand_gumbel_like(x):
-  g = rand_gumbel(x.size()).to(dtype=x.dtype, device=x.device)
-  return g
+    g = rand_gumbel(x.size()).to(dtype=x.dtype, device=x.device)
+    return g
 
 
 def slice_segments(x, ids_str, segment_size=4):
-  ret = torch.zeros_like(x[:, :, :segment_size])
-  for i in range(x.size(0)):
-    idx_str = ids_str[i]
-    idx_end = idx_str + segment_size
-    ret[i] = x[i, :, idx_str:idx_end]
-  return ret
+    ret = torch.zeros_like(x[:, :, :segment_size])
+    for i in range(x.size(0)):
+        idx_str = ids_str[i]
+        idx_end = idx_str + segment_size
+        ret[i] = x[i, :, idx_str:idx_end]
+    return ret
 
 
 def rand_slice_segments(x, x_lengths=None, segment_size=4):
-  b, d, t = x.size()
-  if x_lengths is None:
-    x_lengths = t
-  ids_str_max = x_lengths - segment_size + 1
-  ids_str = (torch.rand([b]).to(device=x.device) * ids_str_max).to(dtype=torch.long)
-  ret = slice_segments(x, ids_str, segment_size)
-  return ret, ids_str
+    b, d, t = x.size()
+    if x_lengths is None:
+        x_lengths = t
+    ids_str_max = x_lengths - segment_size + 1
+    ids_str = (torch.rand([b]).to(device=x.device) * ids_str_max).to(dtype=torch.long)
+    ret = slice_segments(x, ids_str, segment_size)
+    return ret, ids_str
 
 
-def get_timing_signal_1d(
-    length, channels, min_timescale=1.0, max_timescale=1.0e4):
-  position = torch.arange(length, dtype=torch.float)
-  num_timescales = channels // 2
-  log_timescale_increment = (
-      math.log(float(max_timescale) / float(min_timescale)) /
-      (num_timescales - 1))
-  inv_timescales = min_timescale * torch.exp(
-      torch.arange(num_timescales, dtype=torch.float) * -log_timescale_increment)
-  scaled_time = position.unsqueeze(0) * inv_timescales.unsqueeze(1)
-  signal = torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], 0)
-  signal = F.pad(signal, [0, 0, 0, channels % 2])
-  signal = signal.view(1, channels, length)
-  return signal
+def get_timing_signal_1d(length, channels, min_timescale=1.0, max_timescale=1.0e4):
+    position = torch.arange(length, dtype=torch.float)
+    num_timescales = channels // 2
+    log_timescale_increment = math.log(float(max_timescale) / float(min_timescale)) / (
+        num_timescales - 1
+    )
+    inv_timescales = min_timescale * torch.exp(
+        torch.arange(num_timescales, dtype=torch.float) * -log_timescale_increment
+    )
+    scaled_time = position.unsqueeze(0) * inv_timescales.unsqueeze(1)
+    signal = torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], 0)
+    signal = F.pad(signal, [0, 0, 0, channels % 2])
+    signal = signal.view(1, channels, length)
+    return signal
 
 
 def add_timing_signal_1d(x, min_timescale=1.0, max_timescale=1.0e4):
-  b, channels, length = x.size()
-  signal = get_timing_signal_1d(length, channels, min_timescale, max_timescale)
-  return x + signal.to(dtype=x.dtype, device=x.device)
+    b, channels, length = x.size()
+    signal = get_timing_signal_1d(length, channels, min_timescale, max_timescale)
+    return x + signal.to(dtype=x.dtype, device=x.device)
 
 
 def cat_timing_signal_1d(x, min_timescale=1.0, max_timescale=1.0e4, axis=1):
-  b, channels, length = x.size()
-  signal = get_timing_signal_1d(length, channels, min_timescale, max_timescale)
-  return torch.cat([x, signal.to(dtype=x.dtype, device=x.device)], axis)
+    b, channels, length = x.size()
+    signal = get_timing_signal_1d(length, channels, min_timescale, max_timescale)
+    return torch.cat([x, signal.to(dtype=x.dtype, device=x.device)], axis)
 
 
 def subsequent_mask(length):
-  mask = torch.tril(torch.ones(length, length)).unsqueeze(0).unsqueeze(0)
-  return mask
+    mask = torch.tril(torch.ones(length, length)).unsqueeze(0).unsqueeze(0)
+    return mask
 
 
 @torch.jit.script
 def fused_add_tanh_sigmoid_multiply(input_a, input_b, n_channels):
-  n_channels_int = n_channels[0]
-  in_act = input_a + input_b
-  t_act = torch.tanh(in_act[:, :n_channels_int, :])
-  s_act = torch.sigmoid(in_act[:, n_channels_int:, :])
-  acts = t_act * s_act
-  return acts
+    n_channels_int = n_channels[0]
+    in_act = input_a + input_b
+    t_act = torch.tanh(in_act[:, :n_channels_int, :])
+    s_act = torch.sigmoid(in_act[:, n_channels_int:, :])
+    acts = t_act * s_act
+    return acts
 
 
 def convert_pad_shape(pad_shape):
-  l = pad_shape[::-1]
-  pad_shape = [item for sublist in l for item in sublist]
-  return pad_shape
+    layer = pad_shape[::-1]
+    pad_shape = [item for sublist in layer for item in sublist]
+    return pad_shape
 
 
 def shift_1d(x):
-  x = F.pad(x, convert_pad_shape([[0, 0], [0, 0], [1, 0]]))[:, :, :-1]
-  return x
+    x = F.pad(x, convert_pad_shape([[0, 0], [0, 0], [1, 0]]))[:, :, :-1]
+    return x
 
 
 def sequence_mask(length, max_length=None):
-  if max_length is None:
-    max_length = length.max()
-  x = torch.arange(max_length, dtype=length.dtype, device=length.device)
-  return x.unsqueeze(0) < length.unsqueeze(1)
+    if max_length is None:
+        max_length = length.max()
+    x = torch.arange(max_length, dtype=length.dtype, device=length.device)
+    return x.unsqueeze(0) < length.unsqueeze(1)
 
 
 def generate_path(duration, mask):
-  """
-  duration: [b, 1, t_x]
-  mask: [b, 1, t_y, t_x]
-  """
-  device = duration.device
-  
-  b, _, t_y, t_x = mask.shape
-  cum_duration = torch.cumsum(duration, -1)
-  
-  cum_duration_flat = cum_duration.view(b * t_x)
-  path = sequence_mask(cum_duration_flat, t_y).to(mask.dtype)
-  path = path.view(b, t_x, t_y)
-  path = path - F.pad(path, convert_pad_shape([[0, 0], [1, 0], [0, 0]]))[:, :-1]
-  path = path.unsqueeze(1).transpose(2,3) * mask
-  return path
+    """
+    duration: [b, 1, t_x]
+    mask: [b, 1, t_y, t_x]
+    """
+
+    b, _, t_y, t_x = mask.shape
+    cum_duration = torch.cumsum(duration, -1)
+
+    cum_duration_flat = cum_duration.view(b * t_x)
+    path = sequence_mask(cum_duration_flat, t_y).to(mask.dtype)
+    path = path.view(b, t_x, t_y)
+    path = path - F.pad(path, convert_pad_shape([[0, 0], [1, 0], [0, 0]]))[:, :-1]
+    path = path.unsqueeze(1).transpose(2, 3) * mask
+    return path
 
 
 def clip_grad_value_(parameters, clip_value, norm_type=2):
-  if isinstance(parameters, torch.Tensor):
-    parameters = [parameters]
-  parameters = list(filter(lambda p: p.grad is not None, parameters))
-  norm_type = float(norm_type)
-  if clip_value is not None:
-    clip_value = float(clip_value)
-
-  total_norm = 0
-  for p in parameters:
-    param_norm = p.grad.data.norm(norm_type)
-    total_norm += param_norm.item() ** norm_type
+    if isinstance(parameters, torch.Tensor):
+        parameters = [parameters]
+    parameters = list(filter(lambda p: p.grad is not None, parameters))
+    norm_type = float(norm_type)
     if clip_value is not None:
-      p.grad.data.clamp_(min=-clip_value, max=clip_value)
-  total_norm = total_norm ** (1. / norm_type)
-  return total_norm
+        clip_value = float(clip_value)
+
+    total_norm = 0
+    for p in parameters:
+        param_norm = p.grad.data.norm(norm_type)
+        total_norm += param_norm.item() ** norm_type
+        if clip_value is not None:
+            p.grad.data.clamp_(min=-clip_value, max=clip_value)
+    total_norm = total_norm ** (1.0 / norm_type)
+    return total_norm

configs/config.json

@@ -17,10 +17,10 @@
     "init_lr_ratio": 1,
     "warmup_epochs": 0,
     "c_mel": 45,
-    "c_kl": 1.0
+    "c_kl": 1.0,
+    "skip_optimizer": true
   },
   "data": {
-    "use_mel_posterior_encoder": false,
     "training_files": "filelists/train.list",
     "validation_files": "filelists/val.list",
     "max_wav_value": 32768.0,

data_utils.py

@@ -1,11 +1,11 @@
-import time
 import os
 import random
-import numpy as np
 import torch
 import torch.utils.data
+from tqdm import tqdm
+from loguru import logger
 import commons
-from mel_processing import spectrogram_torch, mel_spectrogram_torch, spec_to_mel_torch
+from mel_processing import spectrogram_torch, mel_spectrogram_torch
 from utils import load_wav_to_torch, load_filepaths_and_text
 from text import cleaned_text_to_sequence, get_bert
 
@@ -14,9 +14,9 @@ from text import cleaned_text_to_sequence, get_bert
 
 class TextAudioSpeakerLoader(torch.utils.data.Dataset):
     """
-        1) loads audio, speaker_id, text pairs
-        2) normalizes text and converts them to sequences of integers
-        3) computes spectrograms from audio files.
+    1) loads audio, speaker_id, text pairs
+    2) normalizes text and converts them to sequences of integers
+    3) computes spectrograms from audio files.
     """
 
     def __init__(self, audiopaths_sid_text, hparams):
@@ -30,7 +30,9 @@ class TextAudioSpeakerLoader(torch.utils.data.Dataset):
         self.spk_map = hparams.spk2id
         self.hparams = hparams
 
-        self.use_mel_spec_posterior = getattr(hparams, "use_mel_posterior_encoder", False)
+        self.use_mel_spec_posterior = getattr(
+            hparams, "use_mel_posterior_encoder", False
+        )
         if self.use_mel_spec_posterior:
             self.n_mel_channels = getattr(hparams, "n_mel_channels", 80)
 
@@ -55,17 +57,27 @@ class TextAudioSpeakerLoader(torch.utils.data.Dataset):
         audiopaths_sid_text_new = []
         lengths = []
         skipped = 0
-        for _id, spk, language, text, phones, tone, word2ph in self.audiopaths_sid_text:
-            audiopath = f'{_id}'
+        logger.info("Init dataset...")
+        for _id, spk, language, text, phones, tone, word2ph in tqdm(
+            self.audiopaths_sid_text
+        ):
+            audiopath = f"{_id}"
             if self.min_text_len <= len(phones) and len(phones) <= self.max_text_len:
                 phones = phones.split(" ")
                 tone = [int(i) for i in tone.split(" ")]
                 word2ph = [int(i) for i in word2ph.split(" ")]
-                audiopaths_sid_text_new.append([audiopath, spk, language, text, phones, tone, word2ph])
+                audiopaths_sid_text_new.append(
+                    [audiopath, spk, language, text, phones, tone, word2ph]
+                )
                 lengths.append(os.path.getsize(audiopath) // (2 * self.hop_length))
             else:
                 skipped += 1
-        print("skipped: ", skipped, ", total: ", len(self.audiopaths_sid_text))
+        logger.info(
+            "skipped: "
+            + str(skipped)
+            + ", total: "
+            + str(len(self.audiopaths_sid_text))
+        )
         self.audiopaths_sid_text = audiopaths_sid_text_new
         self.lengths = lengths
 
@@ -73,17 +85,22 @@ class TextAudioSpeakerLoader(torch.utils.data.Dataset):
         # separate filename, speaker_id and text
         audiopath, sid, language, text, phones, tone, word2ph = audiopath_sid_text
 
-        bert, phones, tone, language = self.get_text(text, word2ph, phones, tone, language, audiopath)
+        bert, ja_bert, phones, tone, language = self.get_text(
+            text, word2ph, phones, tone, language, audiopath
+        )
 
         spec, wav = self.get_audio(audiopath)
         sid = torch.LongTensor([int(self.spk_map[sid])])
-        return (phones, spec, wav, sid, tone, language, bert)
+        return (phones, spec, wav, sid, tone, language, bert, ja_bert)
 
     def get_audio(self, filename):
         audio, sampling_rate = load_wav_to_torch(filename)
         if sampling_rate != self.sampling_rate:
-            raise ValueError("{} {} SR doesn't match target {} SR".format(
-                sampling_rate, self.sampling_rate))
+            raise ValueError(
+                "{} {} SR doesn't match target {} SR".format(
+                    filename, sampling_rate, self.sampling_rate
+                )
+            )
         audio_norm = audio / self.max_wav_value
         audio_norm = audio_norm.unsqueeze(0)
         spec_filename = filename.replace(".wav", ".spec.pt")
@@ -93,31 +110,35 @@ class TextAudioSpeakerLoader(torch.utils.data.Dataset):
             spec = torch.load(spec_filename)
         except:
             if self.use_mel_spec_posterior:
-                spec = mel_spectrogram_torch(audio_norm, self.filter_length,
-                    self.n_mel_channels, self.sampling_rate, self.hop_length,
-                    self.win_length, self.hparams.mel_fmin, self.hparams.mel_fmax, center=False)
+                spec = mel_spectrogram_torch(
+                    audio_norm,
+                    self.filter_length,
+                    self.n_mel_channels,
+                    self.sampling_rate,
+                    self.hop_length,
+                    self.win_length,
+                    self.hparams.mel_fmin,
+                    self.hparams.mel_fmax,
+                    center=False,
+                )
             else:
-                spec = spectrogram_torch(audio_norm, self.filter_length,
-                    self.sampling_rate, self.hop_length, self.win_length,
-                    center=False)
+                spec = spectrogram_torch(
+                    audio_norm,
+                    self.filter_length,
+                    self.sampling_rate,
+                    self.hop_length,
+                    self.win_length,
+                    center=False,
+                )
             spec = torch.squeeze(spec, 0)
             torch.save(spec, spec_filename)
         return spec, audio_norm
 
     def get_text(self, text, word2ph, phone, tone, language_str, wav_path):
-        pold = phone
-        w2pho = [i for i in word2ph]
-        word2ph = [i for i in word2ph]
         phone, tone, language = cleaned_text_to_sequence(phone, tone, language_str)
-        pold2 = phone
-
         if self.add_blank:
-            p1 = len(phone)
             phone = commons.intersperse(phone, 0)
-            p2 = len(phone)
-            t1 = len(tone)
             tone = commons.intersperse(tone, 0)
-            t2 = len(tone)
             language = commons.intersperse(language, 0)
             for i in range(len(word2ph)):
                 word2ph[i] = word2ph[i] * 2
@@ -129,15 +150,35 @@ class TextAudioSpeakerLoader(torch.utils.data.Dataset):
         except:
             bert = get_bert(text, word2ph, language_str)
             torch.save(bert, bert_path)
-            #print(bert.shape[-1], bert_path, text, pold)
-            assert bert.shape[-1] == len(phone)
+            assert bert.shape[-1] == len(phone), phone
 
+        if language_str == "ZH":
+            bert = bert
+            ja_bert = torch.zeros(768, len(phone))
+        elif language_str == "JA":
+            ja_bert = bert
+            bert = torch.zeros(1024, len(phone))
+        else:
+            bert = torch.zeros(1024, len(phone))
+            ja_bert = torch.zeros(768, len(phone))
         assert bert.shape[-1] == len(phone), (
-        bert.shape, len(phone), sum(word2ph), p1, p2, t1, t2, pold, pold2, word2ph, text, w2pho)
+            bert.shape,
+            len(phone),
+            sum(word2ph),
+            p1,
+            p2,
+            t1,
+            t2,
+            pold,
+            pold2,
+            word2ph,
+            text,
+            w2pho,
+        )
         phone = torch.LongTensor(phone)
         tone = torch.LongTensor(tone)
         language = torch.LongTensor(language)
-        return bert, phone, tone, language
+        return bert, ja_bert, phone, tone, language
 
     def get_sid(self, sid):
         sid = torch.LongTensor([int(sid)])
@@ -150,9 +191,8 @@ class TextAudioSpeakerLoader(torch.utils.data.Dataset):
         return len(self.audiopaths_sid_text)
 
 
-class TextAudioSpeakerCollate():
-    """ Zero-pads model inputs and targets
-    """
+class TextAudioSpeakerCollate:
+    """Zero-pads model inputs and targets"""
 
     def __init__(self, return_ids=False):
         self.return_ids = return_ids
@@ -165,8 +205,8 @@ class TextAudioSpeakerCollate():
         """
         # Right zero-pad all one-hot text sequences to max input length
         _, ids_sorted_decreasing = torch.sort(
-            torch.LongTensor([x[1].size(1) for x in batch]),
-            dim=0, descending=True)
+            torch.LongTensor([x[1].size(1) for x in batch]), dim=0, descending=True
+        )
 
         max_text_len = max([len(x[0]) for x in batch])
         max_spec_len = max([x[1].size(1) for x in batch])
@@ -181,6 +221,7 @@ class TextAudioSpeakerCollate():
         tone_padded = torch.LongTensor(len(batch), max_text_len)
         language_padded = torch.LongTensor(len(batch), max_text_len)
         bert_padded = torch.FloatTensor(len(batch), 1024, max_text_len)
+        ja_bert_padded = torch.FloatTensor(len(batch), 768, max_text_len)
 
         spec_padded = torch.FloatTensor(len(batch), batch[0][1].size(0), max_spec_len)
         wav_padded = torch.FloatTensor(len(batch), 1, max_wav_len)
@@ -190,33 +231,49 @@ class TextAudioSpeakerCollate():
         spec_padded.zero_()
         wav_padded.zero_()
         bert_padded.zero_()
+        ja_bert_padded.zero_()
         for i in range(len(ids_sorted_decreasing)):
             row = batch[ids_sorted_decreasing[i]]
 
             text = row[0]
-            text_padded[i, :text.size(0)] = text
+            text_padded[i, : text.size(0)] = text
             text_lengths[i] = text.size(0)
 
             spec = row[1]
-            spec_padded[i, :, :spec.size(1)] = spec
+            spec_padded[i, :, : spec.size(1)] = spec
             spec_lengths[i] = spec.size(1)
 
             wav = row[2]
-            wav_padded[i, :, :wav.size(1)] = wav
+            wav_padded[i, :, : wav.size(1)] = wav
             wav_lengths[i] = wav.size(1)
 
             sid[i] = row[3]
 
             tone = row[4]
-            tone_padded[i, :tone.size(0)] = tone
+            tone_padded[i, : tone.size(0)] = tone
 
             language = row[5]
-            language_padded[i, :language.size(0)] = language
+            language_padded[i, : language.size(0)] = language
 
             bert = row[6]
-            bert_padded[i, :, :bert.size(1)] = bert
+            bert_padded[i, :, : bert.size(1)] = bert
 
-        return text_padded, text_lengths, spec_padded, spec_lengths, wav_padded, wav_lengths, sid, tone_padded, language_padded, bert_padded
+            ja_bert = row[7]
+            ja_bert_padded[i, :, : ja_bert.size(1)] = ja_bert
+
+        return (
+            text_padded,
+            text_lengths,
+            spec_padded,
+            spec_lengths,
+            wav_padded,
+            wav_lengths,
+            sid,
+            tone_padded,
+            language_padded,
+            bert_padded,
+            ja_bert_padded,
+        )
 
 
 class DistributedBucketSampler(torch.utils.data.distributed.DistributedSampler):
@@ -229,7 +286,15 @@ class DistributedBucketSampler(torch.utils.data.distributed.DistributedSampler):
     Ex) boundaries = [b1, b2, b3] -> any x s.t. length(x) <= b1 or length(x) > b3 are discarded.
     """
 
-    def __init__(self, dataset, batch_size, boundaries, num_replicas=None, rank=None, shuffle=True):
+    def __init__(
+        self,
+        dataset,
+        batch_size,
+        boundaries,
+        num_replicas=None,
+        rank=None,
+        shuffle=True,
+    ):
         super().__init__(dataset, num_replicas=num_replicas, rank=rank, shuffle=shuffle)
         self.lengths = dataset.lengths
         self.batch_size = batch_size
@@ -247,16 +312,27 @@ class DistributedBucketSampler(torch.utils.data.distributed.DistributedSampler):
             if idx_bucket != -1:
                 buckets[idx_bucket].append(i)
 
-        for i in range(len(buckets) - 1, 0, -1):
-            if len(buckets[i]) == 0:
-                buckets.pop(i)
-                self.boundaries.pop(i + 1)
+        try:
+            for i in range(len(buckets) - 1, 0, -1):
+                if len(buckets[i]) == 0:
+                    buckets.pop(i)
+                    self.boundaries.pop(i + 1)
+            assert all(len(bucket) > 0 for bucket in buckets)
+        # When one bucket is not traversed
+        except Exception as e:
+            print("Bucket warning ", e)
+            for i in range(len(buckets) - 1, -1, -1):
+                if len(buckets[i]) == 0:
+                    buckets.pop(i)
+                    self.boundaries.pop(i + 1)
 
         num_samples_per_bucket = []
         for i in range(len(buckets)):
             len_bucket = len(buckets[i])
             total_batch_size = self.num_replicas * self.batch_size
-            rem = (total_batch_size - (len_bucket % total_batch_size)) % total_batch_size
+            rem = (
+                total_batch_size - (len_bucket % total_batch_size)
+            ) % total_batch_size
             num_samples_per_bucket.append(len_bucket + rem)
         return buckets, num_samples_per_bucket
 
@@ -277,21 +353,30 @@ class DistributedBucketSampler(torch.utils.data.distributed.DistributedSampler):
         for i in range(len(self.buckets)):
             bucket = self.buckets[i]
             len_bucket = len(bucket)
-            if (len_bucket == 0):
+            if len_bucket == 0:
                 continue
             ids_bucket = indices[i]
             num_samples_bucket = self.num_samples_per_bucket[i]
 
             # add extra samples to make it evenly divisible
             rem = num_samples_bucket - len_bucket
-            ids_bucket = ids_bucket + ids_bucket * (rem // len_bucket) + ids_bucket[:(rem % len_bucket)]
+            ids_bucket = (
+                ids_bucket
+                + ids_bucket * (rem // len_bucket)
+                + ids_bucket[: (rem % len_bucket)]
+            )
 
             # subsample
-            ids_bucket = ids_bucket[self.rank::self.num_replicas]
+            ids_bucket = ids_bucket[self.rank :: self.num_replicas]
 
             # batching
             for j in range(len(ids_bucket) // self.batch_size):
-                batch = [bucket[idx] for idx in ids_bucket[j * self.batch_size:(j + 1) * self.batch_size]]
+                batch = [
+                    bucket[idx]
+                    for idx in ids_bucket[
+                        j * self.batch_size : (j + 1) * self.batch_size
+                    ]
+                ]
                 batches.append(batch)
 
         if self.shuffle:

filelists/esd.list

@@ -1,2 +1,3 @@
 Example:
 {wav_path}|{speaker_name}|{language}|{text}
+派蒙_1.wav|派蒙|ZH|前面的区域，以后再来探索吧！

losses.py

@@ -1,61 +1,58 @@
-import torch 
-from torch.nn import functional as F
-
-import commons
+import torch
 
 
 def feature_loss(fmap_r, fmap_g):
-  loss = 0
-  for dr, dg in zip(fmap_r, fmap_g):
-    for rl, gl in zip(dr, dg):
-      rl = rl.float().detach()
-      gl = gl.float()
-      loss += torch.mean(torch.abs(rl - gl))
+    loss = 0
+    for dr, dg in zip(fmap_r, fmap_g):
+        for rl, gl in zip(dr, dg):
+            rl = rl.float().detach()
+            gl = gl.float()
+            loss += torch.mean(torch.abs(rl - gl))
 
-  return loss * 2 
+    return loss * 2
 
 
 def discriminator_loss(disc_real_outputs, disc_generated_outputs):
-  loss = 0
-  r_losses = []
-  g_losses = []
-  for dr, dg in zip(disc_real_outputs, disc_generated_outputs):
-    dr = dr.float()
-    dg = dg.float()
-    r_loss = torch.mean((1-dr)**2)
-    g_loss = torch.mean(dg**2)
-    loss += (r_loss + g_loss)
-    r_losses.append(r_loss.item())
-    g_losses.append(g_loss.item())
+    loss = 0
+    r_losses = []
+    g_losses = []
+    for dr, dg in zip(disc_real_outputs, disc_generated_outputs):
+        dr = dr.float()
+        dg = dg.float()
+        r_loss = torch.mean((1 - dr) ** 2)
+        g_loss = torch.mean(dg**2)
+        loss += r_loss + g_loss
+        r_losses.append(r_loss.item())
+        g_losses.append(g_loss.item())
 
-  return loss, r_losses, g_losses
+    return loss, r_losses, g_losses
 
 
 def generator_loss(disc_outputs):
-  loss = 0
-  gen_losses = []
-  for dg in disc_outputs:
-    dg = dg.float()
-    l = torch.mean((1-dg)**2)
-    gen_losses.append(l)
-    loss += l
+    loss = 0
+    gen_losses = []
+    for dg in disc_outputs:
+        dg = dg.float()
+        l = torch.mean((1 - dg) ** 2)
+        gen_losses.append(l)
+        loss += l
 
-  return loss, gen_losses
+    return loss, gen_losses
 
 
 def kl_loss(z_p, logs_q, m_p, logs_p, z_mask):
-  """
-  z_p, logs_q: [b, h, t_t]
-  m_p, logs_p: [b, h, t_t]
-  """
-  z_p = z_p.float()
-  logs_q = logs_q.float()
-  m_p = m_p.float()
-  logs_p = logs_p.float()
-  z_mask = z_mask.float()
+    """
+    z_p, logs_q: [b, h, t_t]
+    m_p, logs_p: [b, h, t_t]
+    """
+    z_p = z_p.float()
+    logs_q = logs_q.float()
+    m_p = m_p.float()
+    logs_p = logs_p.float()
+    z_mask = z_mask.float()
 
-  kl = logs_p - logs_q - 0.5
-  kl += 0.5 * ((z_p - m_p)**2) * torch.exp(-2. * logs_p)
-  kl = torch.sum(kl * z_mask)
-  l = kl / torch.sum(z_mask)
-  return l
+    kl = logs_p - logs_q - 0.5
+    kl += 0.5 * ((z_p - m_p) ** 2) * torch.exp(-2.0 * logs_p)
+    kl = torch.sum(kl * z_mask)
+    l = kl / torch.sum(z_mask)
+    return l

mel_processing.py

@@ -1,16 +1,5 @@
-import math
-import os
-import random
 import torch
-from torch import nn
-import torch.nn.functional as F
 import torch.utils.data
-import numpy as np
-import librosa
-import librosa.util as librosa_util
-from librosa.util import normalize, pad_center, tiny
-from scipy.signal import get_window
-from scipy.io.wavfile import read
 from librosa.filters import mel as librosa_mel_fn
 
 MAX_WAV_VALUE = 32768.0
@@ -49,22 +38,38 @@ hann_window = {}
 
 
 def spectrogram_torch(y, n_fft, sampling_rate, hop_size, win_size, center=False):
-    if torch.min(y) < -1.:
-        print('min value is ', torch.min(y))
-    if torch.max(y) > 1.:
-        print('max value is ', torch.max(y))
+    if torch.min(y) < -1.0:
+        print("min value is ", torch.min(y))
+    if torch.max(y) > 1.0:
+        print("max value is ", torch.max(y))
 
     global hann_window
-    dtype_device = str(y.dtype) + '_' + str(y.device)
-    wnsize_dtype_device = str(win_size) + '_' + dtype_device
+    dtype_device = str(y.dtype) + "_" + str(y.device)
+    wnsize_dtype_device = str(win_size) + "_" + dtype_device
     if wnsize_dtype_device not in hann_window:
-        hann_window[wnsize_dtype_device] = torch.hann_window(win_size).to(dtype=y.dtype, device=y.device)
+        hann_window[wnsize_dtype_device] = torch.hann_window(win_size).to(
+            dtype=y.dtype, device=y.device
+        )
 
-    y = torch.nn.functional.pad(y.unsqueeze(1), (int((n_fft-hop_size)/2), int((n_fft-hop_size)/2)), mode='reflect')
+    y = torch.nn.functional.pad(
+        y.unsqueeze(1),
+        (int((n_fft - hop_size) / 2), int((n_fft - hop_size) / 2)),
+        mode="reflect",
+    )
     y = y.squeeze(1)
 
-    spec = torch.stft(y, n_fft, hop_length=hop_size, win_length=win_size, window=hann_window[wnsize_dtype_device],
-                      center=center, pad_mode='reflect', normalized=False, onesided=True, return_complex=False)
+    spec = torch.stft(
+        y,
+        n_fft,
+        hop_length=hop_size,
+        win_length=win_size,
+        window=hann_window[wnsize_dtype_device],
+        center=center,
+        pad_mode="reflect",
+        normalized=False,
+        onesided=True,
+        return_complex=False,
+    )
 
     spec = torch.sqrt(spec.pow(2).sum(-1) + 1e-6)
     return spec
@@ -72,37 +77,59 @@ def spectrogram_torch(y, n_fft, sampling_rate, hop_size, win_size, center=False)
 
 def spec_to_mel_torch(spec, n_fft, num_mels, sampling_rate, fmin, fmax):
     global mel_basis
-    dtype_device = str(spec.dtype) + '_' + str(spec.device)
-    fmax_dtype_device = str(fmax) + '_' + dtype_device
+    dtype_device = str(spec.dtype) + "_" + str(spec.device)
+    fmax_dtype_device = str(fmax) + "_" + dtype_device
     if fmax_dtype_device not in mel_basis:
         mel = librosa_mel_fn(sampling_rate, n_fft, num_mels, fmin, fmax)
-        mel_basis[fmax_dtype_device] = torch.from_numpy(mel).to(dtype=spec.dtype, device=spec.device)
+        mel_basis[fmax_dtype_device] = torch.from_numpy(mel).to(
+            dtype=spec.dtype, device=spec.device
+        )
     spec = torch.matmul(mel_basis[fmax_dtype_device], spec)
     spec = spectral_normalize_torch(spec)
     return spec
 
 
-def mel_spectrogram_torch(y, n_fft, num_mels, sampling_rate, hop_size, win_size, fmin, fmax, center=False):
-    if torch.min(y) < -1.:
-        print('min value is ', torch.min(y))
-    if torch.max(y) > 1.:
-        print('max value is ', torch.max(y))
+def mel_spectrogram_torch(
+    y, n_fft, num_mels, sampling_rate, hop_size, win_size, fmin, fmax, center=False
+):
+    if torch.min(y) < -1.0:
+        print("min value is ", torch.min(y))
+    if torch.max(y) > 1.0:
+        print("max value is ", torch.max(y))
 
     global mel_basis, hann_window
-    dtype_device = str(y.dtype) + '_' + str(y.device)
-    fmax_dtype_device = str(fmax) + '_' + dtype_device
-    wnsize_dtype_device = str(win_size) + '_' + dtype_device
+    dtype_device = str(y.dtype) + "_" + str(y.device)
+    fmax_dtype_device = str(fmax) + "_" + dtype_device
+    wnsize_dtype_device = str(win_size) + "_" + dtype_device
     if fmax_dtype_device not in mel_basis:
         mel = librosa_mel_fn(sampling_rate, n_fft, num_mels, fmin, fmax)
-        mel_basis[fmax_dtype_device] = torch.from_numpy(mel).to(dtype=y.dtype, device=y.device)
+        mel_basis[fmax_dtype_device] = torch.from_numpy(mel).to(
+            dtype=y.dtype, device=y.device
+        )
     if wnsize_dtype_device not in hann_window:
-        hann_window[wnsize_dtype_device] = torch.hann_window(win_size).to(dtype=y.dtype, device=y.device)
+        hann_window[wnsize_dtype_device] = torch.hann_window(win_size).to(
+            dtype=y.dtype, device=y.device
+        )
 
-    y = torch.nn.functional.pad(y.unsqueeze(1), (int((n_fft-hop_size)/2), int((n_fft-hop_size)/2)), mode='reflect')
+    y = torch.nn.functional.pad(
+        y.unsqueeze(1),
+        (int((n_fft - hop_size) / 2), int((n_fft - hop_size) / 2)),
+        mode="reflect",
+    )
     y = y.squeeze(1)
 
-    spec = torch.stft(y, n_fft, hop_length=hop_size, win_length=win_size, window=hann_window[wnsize_dtype_device],
-                          center=center, pad_mode='reflect', normalized=False, onesided=True, return_complex=False)
+    spec = torch.stft(
+        y,
+        n_fft,
+        hop_length=hop_size,
+        win_length=win_size,
+        window=hann_window[wnsize_dtype_device],
+        center=center,
+        pad_mode="reflect",
+        normalized=False,
+        onesided=True,
+        return_complex=False,
+    )
 
     spec = torch.sqrt(spec.pow(2).sum(-1) + 1e-6)
 

modules.py

@@ -1,12 +1,9 @@
-import copy
 import math
-import numpy as np
-import scipy
 import torch
 from torch import nn
 from torch.nn import functional as F
 
-from torch.nn import Conv1d, ConvTranspose1d, AvgPool1d, Conv2d
+from torch.nn import Conv1d
 from torch.nn.utils import weight_norm, remove_weight_norm
 
 import commons
@@ -16,193 +13,284 @@ from attentions import Encoder
 
 LRELU_SLOPE = 0.1
 
+
 class LayerNorm(nn.Module):
-  def __init__(self, channels, eps=1e-5):
-    super().__init__()
-    self.channels = channels
-    self.eps = eps
+    def __init__(self, channels, eps=1e-5):
+        super().__init__()
+        self.channels = channels
+        self.eps = eps
+
+        self.gamma = nn.Parameter(torch.ones(channels))
+        self.beta = nn.Parameter(torch.zeros(channels))
+
+    def forward(self, x):
+        x = x.transpose(1, -1)
+        x = F.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps)
+        return x.transpose(1, -1)
 
-    self.gamma = nn.Parameter(torch.ones(channels))
-    self.beta = nn.Parameter(torch.zeros(channels))
 
-  def forward(self, x):
-    x = x.transpose(1, -1)
-    x = F.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps)
-    return x.transpose(1, -1)
-    
 class ConvReluNorm(nn.Module):
-  def __init__(self, in_channels, hidden_channels, out_channels, kernel_size, n_layers, p_dropout):
-    super().__init__()
-    self.in_channels = in_channels
-    self.hidden_channels = hidden_channels
-    self.out_channels = out_channels
-    self.kernel_size = kernel_size
-    self.n_layers = n_layers
-    self.p_dropout = p_dropout
-    assert n_layers > 1, "Number of layers should be larger than 0."
+    def __init__(
+        self,
+        in_channels,
+        hidden_channels,
+        out_channels,
+        kernel_size,
+        n_layers,
+        p_dropout,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.hidden_channels = hidden_channels
+        self.out_channels = out_channels
+        self.kernel_size = kernel_size
+        self.n_layers = n_layers
+        self.p_dropout = p_dropout
+        assert n_layers > 1, "Number of layers should be larger than 0."
 
-    self.conv_layers = nn.ModuleList()
-    self.norm_layers = nn.ModuleList()
-    self.conv_layers.append(nn.Conv1d(in_channels, hidden_channels, kernel_size, padding=kernel_size//2))
-    self.norm_layers.append(LayerNorm(hidden_channels))
-    self.relu_drop = nn.Sequential(
-        nn.ReLU(),
-        nn.Dropout(p_dropout))
-    for _ in range(n_layers-1):
-      self.conv_layers.append(nn.Conv1d(hidden_channels, hidden_channels, kernel_size, padding=kernel_size//2))
-      self.norm_layers.append(LayerNorm(hidden_channels))
-    self.proj = nn.Conv1d(hidden_channels, out_channels, 1)
-    self.proj.weight.data.zero_()
-    self.proj.bias.data.zero_()
+        self.conv_layers = nn.ModuleList()
+        self.norm_layers = nn.ModuleList()
+        self.conv_layers.append(
+            nn.Conv1d(
+                in_channels, hidden_channels, kernel_size, padding=kernel_size // 2
+            )
+        )
+        self.norm_layers.append(LayerNorm(hidden_channels))
+        self.relu_drop = nn.Sequential(nn.ReLU(), nn.Dropout(p_dropout))
+        for _ in range(n_layers - 1):
+            self.conv_layers.append(
+                nn.Conv1d(
+                    hidden_channels,
+                    hidden_channels,
+                    kernel_size,
+                    padding=kernel_size // 2,
+                )
+            )
+            self.norm_layers.append(LayerNorm(hidden_channels))
+        self.proj = nn.Conv1d(hidden_channels, out_channels, 1)
+        self.proj.weight.data.zero_()
+        self.proj.bias.data.zero_()
 
-  def forward(self, x, x_mask):
-    x_org = x
-    for i in range(self.n_layers):
-      x = self.conv_layers[i](x * x_mask)
-      x = self.norm_layers[i](x)
-      x = self.relu_drop(x)
-    x = x_org + self.proj(x)
-    return x * x_mask
+    def forward(self, x, x_mask):
+        x_org = x
+        for i in range(self.n_layers):
+            x = self.conv_layers[i](x * x_mask)
+            x = self.norm_layers[i](x)
+            x = self.relu_drop(x)
+        x = x_org + self.proj(x)
+        return x * x_mask
 
 
 class DDSConv(nn.Module):
-  """
-  Dialted and Depth-Separable Convolution
-  """
-  def __init__(self, channels, kernel_size, n_layers, p_dropout=0.):
-    super().__init__()
-    self.channels = channels
-    self.kernel_size = kernel_size
-    self.n_layers = n_layers
-    self.p_dropout = p_dropout
+    """
+    Dialted and Depth-Separable Convolution
+    """
 
-    self.drop = nn.Dropout(p_dropout)
-    self.convs_sep = nn.ModuleList()
-    self.convs_1x1 = nn.ModuleList()
-    self.norms_1 = nn.ModuleList()
-    self.norms_2 = nn.ModuleList()
-    for i in range(n_layers):
-      dilation = kernel_size ** i
-      padding = (kernel_size * dilation - dilation) // 2
-      self.convs_sep.append(nn.Conv1d(channels, channels, kernel_size, 
-          groups=channels, dilation=dilation, padding=padding
-      ))
-      self.convs_1x1.append(nn.Conv1d(channels, channels, 1))
-      self.norms_1.append(LayerNorm(channels))
-      self.norms_2.append(LayerNorm(channels))
+    def __init__(self, channels, kernel_size, n_layers, p_dropout=0.0):
+        super().__init__()
+        self.channels = channels
+        self.kernel_size = kernel_size
+        self.n_layers = n_layers
+        self.p_dropout = p_dropout
 
-  def forward(self, x, x_mask, g=None):
-    if g is not None:
-      x = x + g
-    for i in range(self.n_layers):
-      y = self.convs_sep[i](x * x_mask)
-      y = self.norms_1[i](y)
-      y = F.gelu(y)
-      y = self.convs_1x1[i](y)
-      y = self.norms_2[i](y)
-      y = F.gelu(y)
-      y = self.drop(y)
-      x = x + y
-    return x * x_mask
+        self.drop = nn.Dropout(p_dropout)
+        self.convs_sep = nn.ModuleList()
+        self.convs_1x1 = nn.ModuleList()
+        self.norms_1 = nn.ModuleList()
+        self.norms_2 = nn.ModuleList()
+        for i in range(n_layers):
+            dilation = kernel_size**i
+            padding = (kernel_size * dilation - dilation) // 2
+            self.convs_sep.append(
+                nn.Conv1d(
+                    channels,
+                    channels,
+                    kernel_size,
+                    groups=channels,
+                    dilation=dilation,
+                    padding=padding,
+                )
+            )
+            self.convs_1x1.append(nn.Conv1d(channels, channels, 1))
+            self.norms_1.append(LayerNorm(channels))
+            self.norms_2.append(LayerNorm(channels))
+
+    def forward(self, x, x_mask, g=None):
+        if g is not None:
+            x = x + g
+        for i in range(self.n_layers):
+            y = self.convs_sep[i](x * x_mask)
+            y = self.norms_1[i](y)
+            y = F.gelu(y)
+            y = self.convs_1x1[i](y)
+            y = self.norms_2[i](y)
+            y = F.gelu(y)
+            y = self.drop(y)
+            x = x + y
+        return x * x_mask
 
 
 class WN(torch.nn.Module):
-  def __init__(self, hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=0, p_dropout=0):
-    super(WN, self).__init__()
-    assert(kernel_size % 2 == 1)
-    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.p_dropout = p_dropout
+    def __init__(
+        self,
+        hidden_channels,
+        kernel_size,
+        dilation_rate,
+        n_layers,
+        gin_channels=0,
+        p_dropout=0,
+    ):
+        super(WN, self).__init__()
+        assert kernel_size % 2 == 1
+        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.p_dropout = p_dropout
 
-    self.in_layers = torch.nn.ModuleList()
-    self.res_skip_layers = torch.nn.ModuleList()
-    self.drop = nn.Dropout(p_dropout)
+        self.in_layers = torch.nn.ModuleList()
+        self.res_skip_layers = torch.nn.ModuleList()
+        self.drop = nn.Dropout(p_dropout)
 
-    if gin_channels != 0:
-      cond_layer = torch.nn.Conv1d(gin_channels, 2*hidden_channels*n_layers, 1)
-      self.cond_layer = torch.nn.utils.weight_norm(cond_layer, name='weight')
+        if gin_channels != 0:
+            cond_layer = torch.nn.Conv1d(
+                gin_channels, 2 * hidden_channels * n_layers, 1
+            )
+            self.cond_layer = torch.nn.utils.weight_norm(cond_layer, name="weight")
 
-    for i in range(n_layers):
-      dilation = dilation_rate ** i
-      padding = int((kernel_size * dilation - dilation) / 2)
-      in_layer = torch.nn.Conv1d(hidden_channels, 2*hidden_channels, kernel_size,
-                                 dilation=dilation, padding=padding)
-      in_layer = torch.nn.utils.weight_norm(in_layer, name='weight')
-      self.in_layers.append(in_layer)
+        for i in range(n_layers):
+            dilation = dilation_rate**i
+            padding = int((kernel_size * dilation - dilation) / 2)
+            in_layer = torch.nn.Conv1d(
+                hidden_channels,
+                2 * hidden_channels,
+                kernel_size,
+                dilation=dilation,
+                padding=padding,
+            )
+            in_layer = torch.nn.utils.weight_norm(in_layer, name="weight")
+            self.in_layers.append(in_layer)
 
-      # last one is not necessary
-      if i < n_layers - 1:
-        res_skip_channels = 2 * hidden_channels
-      else:
-        res_skip_channels = hidden_channels
+            # last one is not necessary
+            if i < n_layers - 1:
+                res_skip_channels = 2 * hidden_channels
+            else:
+                res_skip_channels = hidden_channels
 
-      res_skip_layer = torch.nn.Conv1d(hidden_channels, res_skip_channels, 1)
-      res_skip_layer = torch.nn.utils.weight_norm(res_skip_layer, name='weight')
-      self.res_skip_layers.append(res_skip_layer)
+            res_skip_layer = torch.nn.Conv1d(hidden_channels, res_skip_channels, 1)
+            res_skip_layer = torch.nn.utils.weight_norm(res_skip_layer, name="weight")
+            self.res_skip_layers.append(res_skip_layer)
 
-  def forward(self, x, x_mask, g=None, **kwargs):
-    output = torch.zeros_like(x)
-    n_channels_tensor = torch.IntTensor([self.hidden_channels])
+    def forward(self, x, x_mask, g=None, **kwargs):
+        output = torch.zeros_like(x)
+        n_channels_tensor = torch.IntTensor([self.hidden_channels])
 
-    if g is not None:
-      g = self.cond_layer(g)
+        if g is not None:
+            g = self.cond_layer(g)
 
-    for i in range(self.n_layers):
-      x_in = self.in_layers[i](x)
-      if g is not None:
-        cond_offset = i * 2 * self.hidden_channels
-        g_l = g[:,cond_offset:cond_offset+2*self.hidden_channels,:]
-      else:
-        g_l = torch.zeros_like(x_in)
+        for i in range(self.n_layers):
+            x_in = self.in_layers[i](x)
+            if g is not None:
+                cond_offset = i * 2 * self.hidden_channels
+                g_l = g[:, cond_offset : cond_offset + 2 * self.hidden_channels, :]
+            else:
+                g_l = torch.zeros_like(x_in)
 
-      acts = commons.fused_add_tanh_sigmoid_multiply(
-          x_in,
-          g_l,
-          n_channels_tensor)
-      acts = self.drop(acts)
+            acts = commons.fused_add_tanh_sigmoid_multiply(x_in, g_l, n_channels_tensor)
+            acts = self.drop(acts)
 
-      res_skip_acts = self.res_skip_layers[i](acts)
-      if i < self.n_layers - 1:
-        res_acts = res_skip_acts[:,:self.hidden_channels,:]
-        x = (x + res_acts) * x_mask
-        output = output + res_skip_acts[:,self.hidden_channels:,:]
-      else:
-        output = output + res_skip_acts
-    return output * x_mask
+            res_skip_acts = self.res_skip_layers[i](acts)
+            if i < self.n_layers - 1:
+                res_acts = res_skip_acts[:, : self.hidden_channels, :]
+                x = (x + res_acts) * x_mask
+                output = output + res_skip_acts[:, self.hidden_channels :, :]
+            else:
+                output = output + res_skip_acts
+        return output * x_mask
 
-  def remove_weight_norm(self):
-    if self.gin_channels != 0:
-      torch.nn.utils.remove_weight_norm(self.cond_layer)
-    for l in self.in_layers:
-      torch.nn.utils.remove_weight_norm(l)
-    for l in self.res_skip_layers:
-     torch.nn.utils.remove_weight_norm(l)
+    def remove_weight_norm(self):
+        if self.gin_channels != 0:
+            torch.nn.utils.remove_weight_norm(self.cond_layer)
+        for l in self.in_layers:
+            torch.nn.utils.remove_weight_norm(l)
+        for l in self.res_skip_layers:
+            torch.nn.utils.remove_weight_norm(l)
 
 
 class ResBlock1(torch.nn.Module):
     def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5)):
         super(ResBlock1, self).__init__()
-        self.convs1 = nn.ModuleList([
-            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0],
-                               padding=get_padding(kernel_size, dilation[0]))),
-            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1],
-                               padding=get_padding(kernel_size, dilation[1]))),
-            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[2],
-                               padding=get_padding(kernel_size, dilation[2])))
-        ])
+        self.convs1 = nn.ModuleList(
+            [
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[0],
+                        padding=get_padding(kernel_size, dilation[0]),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[1],
+                        padding=get_padding(kernel_size, dilation[1]),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[2],
+                        padding=get_padding(kernel_size, dilation[2]),
+                    )
+                ),
+            ]
+        )
         self.convs1.apply(init_weights)
 
-        self.convs2 = nn.ModuleList([
-            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
-                               padding=get_padding(kernel_size, 1))),
-            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
-                               padding=get_padding(kernel_size, 1))),
-            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
-                               padding=get_padding(kernel_size, 1)))
-        ])
+        self.convs2 = nn.ModuleList(
+            [
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=1,
+                        padding=get_padding(kernel_size, 1),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=1,
+                        padding=get_padding(kernel_size, 1),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=1,
+                        padding=get_padding(kernel_size, 1),
+                    )
+                ),
+            ]
+        )
         self.convs2.apply(init_weights)
 
     def forward(self, x, x_mask=None):
@@ -230,12 +318,30 @@ class ResBlock1(torch.nn.Module):
 class ResBlock2(torch.nn.Module):
     def __init__(self, channels, kernel_size=3, dilation=(1, 3)):
         super(ResBlock2, self).__init__()
-        self.convs = nn.ModuleList([
-            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0],
-                               padding=get_padding(kernel_size, dilation[0]))),
-            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1],
-                               padding=get_padding(kernel_size, dilation[1])))
-        ])
+        self.convs = nn.ModuleList(
+            [
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[0],
+                        padding=get_padding(kernel_size, dilation[0]),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[1],
+                        padding=get_padding(kernel_size, dilation[1]),
+                    )
+                ),
+            ]
+        )
         self.convs.apply(init_weights)
 
     def forward(self, x, x_mask=None):
@@ -255,198 +361,237 @@ class ResBlock2(torch.nn.Module):
 
 
 class Log(nn.Module):
-  def forward(self, x, x_mask, reverse=False, **kwargs):
-    if not reverse:
-      y = torch.log(torch.clamp_min(x, 1e-5)) * x_mask
-      logdet = torch.sum(-y, [1, 2])
-      return y, logdet
-    else:
-      x = torch.exp(x) * x_mask
-      return x
-    
+    def forward(self, x, x_mask, reverse=False, **kwargs):
+        if not reverse:
+            y = torch.log(torch.clamp_min(x, 1e-5)) * x_mask
+            logdet = torch.sum(-y, [1, 2])
+            return y, logdet
+        else:
+            x = torch.exp(x) * x_mask
+            return x
+
 
 class Flip(nn.Module):
-  def forward(self, x, *args, reverse=False, **kwargs):
-    x = torch.flip(x, [1])
-    if not reverse:
-      logdet = torch.zeros(x.size(0)).to(dtype=x.dtype, device=x.device)
-      return x, logdet
-    else:
-      return x
+    def forward(self, x, *args, reverse=False, **kwargs):
+        x = torch.flip(x, [1])
+        if not reverse:
+            logdet = torch.zeros(x.size(0)).to(dtype=x.dtype, device=x.device)
+            return x, logdet
+        else:
+            return x
 
 
 class ElementwiseAffine(nn.Module):
-  def __init__(self, channels):
-    super().__init__()
-    self.channels = channels
-    self.m = nn.Parameter(torch.zeros(channels,1))
-    self.logs = nn.Parameter(torch.zeros(channels,1))
+    def __init__(self, channels):
+        super().__init__()
+        self.channels = channels
+        self.m = nn.Parameter(torch.zeros(channels, 1))
+        self.logs = nn.Parameter(torch.zeros(channels, 1))
 
-  def forward(self, x, x_mask, reverse=False, **kwargs):
-    if not reverse:
-      y = self.m + torch.exp(self.logs) * x
-      y = y * x_mask
-      logdet = torch.sum(self.logs * x_mask, [1,2])
-      return y, logdet
-    else:
-      x = (x - self.m) * torch.exp(-self.logs) * x_mask
-      return x
+    def forward(self, x, x_mask, reverse=False, **kwargs):
+        if not reverse:
+            y = self.m + torch.exp(self.logs) * x
+            y = y * x_mask
+            logdet = torch.sum(self.logs * x_mask, [1, 2])
+            return y, logdet
+        else:
+            x = (x - self.m) * torch.exp(-self.logs) * x_mask
+            return x
 
 
 class ResidualCouplingLayer(nn.Module):
-  def __init__(self,
-      channels,
-      hidden_channels,
-      kernel_size,
-      dilation_rate,
-      n_layers,
-      p_dropout=0,
-      gin_channels=0,
-      mean_only=False):
-    assert channels % 2 == 0, "channels should be divisible by 2"
-    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.half_channels = channels // 2
-    self.mean_only = mean_only
+    def __init__(
+        self,
+        channels,
+        hidden_channels,
+        kernel_size,
+        dilation_rate,
+        n_layers,
+        p_dropout=0,
+        gin_channels=0,
+        mean_only=False,
+    ):
+        assert channels % 2 == 0, "channels should be divisible by 2"
+        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.half_channels = channels // 2
+        self.mean_only = mean_only
 
-    self.pre = nn.Conv1d(self.half_channels, hidden_channels, 1)
-    self.enc = WN(hidden_channels, kernel_size, dilation_rate, n_layers, p_dropout=p_dropout, gin_channels=gin_channels)
-    self.post = nn.Conv1d(hidden_channels, self.half_channels * (2 - mean_only), 1)
-    self.post.weight.data.zero_()
-    self.post.bias.data.zero_()
+        self.pre = nn.Conv1d(self.half_channels, hidden_channels, 1)
+        self.enc = WN(
+            hidden_channels,
+            kernel_size,
+            dilation_rate,
+            n_layers,
+            p_dropout=p_dropout,
+            gin_channels=gin_channels,
+        )
+        self.post = nn.Conv1d(hidden_channels, self.half_channels * (2 - mean_only), 1)
+        self.post.weight.data.zero_()
+        self.post.bias.data.zero_()
 
-  def forward(self, x, x_mask, g=None, reverse=False):
-    x0, x1 = torch.split(x, [self.half_channels]*2, 1)
-    h = self.pre(x0) * x_mask
-    h = self.enc(h, x_mask, g=g)
-    stats = self.post(h) * x_mask
-    if not self.mean_only:
-      m, logs = torch.split(stats, [self.half_channels]*2, 1)
-    else:
-      m = stats
-      logs = torch.zeros_like(m)
+    def forward(self, x, x_mask, g=None, reverse=False):
+        x0, x1 = torch.split(x, [self.half_channels] * 2, 1)
+        h = self.pre(x0) * x_mask
+        h = self.enc(h, x_mask, g=g)
+        stats = self.post(h) * x_mask
+        if not self.mean_only:
+            m, logs = torch.split(stats, [self.half_channels] * 2, 1)
+        else:
+            m = stats
+            logs = torch.zeros_like(m)
 
-    if not reverse:
-      x1 = m + x1 * torch.exp(logs) * x_mask
-      x = torch.cat([x0, x1], 1)
-      logdet = torch.sum(logs, [1,2])
-      return x, logdet
-    else:
-      x1 = (x1 - m) * torch.exp(-logs) * x_mask
-      x = torch.cat([x0, x1], 1)
-      return x
+        if not reverse:
+            x1 = m + x1 * torch.exp(logs) * x_mask
+            x = torch.cat([x0, x1], 1)
+            logdet = torch.sum(logs, [1, 2])
+            return x, logdet
+        else:
+            x1 = (x1 - m) * torch.exp(-logs) * x_mask
+            x = torch.cat([x0, x1], 1)
+            return x
 
 
 class ConvFlow(nn.Module):
-  def __init__(self, in_channels, filter_channels, kernel_size, n_layers, num_bins=10, tail_bound=5.0):
-    super().__init__()
-    self.in_channels = in_channels
-    self.filter_channels = filter_channels
-    self.kernel_size = kernel_size
-    self.n_layers = n_layers
-    self.num_bins = num_bins
-    self.tail_bound = tail_bound
-    self.half_channels = in_channels // 2
+    def __init__(
+        self,
+        in_channels,
+        filter_channels,
+        kernel_size,
+        n_layers,
+        num_bins=10,
+        tail_bound=5.0,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.filter_channels = filter_channels
+        self.kernel_size = kernel_size
+        self.n_layers = n_layers
+        self.num_bins = num_bins
+        self.tail_bound = tail_bound
+        self.half_channels = in_channels // 2
 
-    self.pre = nn.Conv1d(self.half_channels, filter_channels, 1)
-    self.convs = DDSConv(filter_channels, kernel_size, n_layers, p_dropout=0.)
-    self.proj = nn.Conv1d(filter_channels, self.half_channels * (num_bins * 3 - 1), 1)
-    self.proj.weight.data.zero_()
-    self.proj.bias.data.zero_()
+        self.pre = nn.Conv1d(self.half_channels, filter_channels, 1)
+        self.convs = DDSConv(filter_channels, kernel_size, n_layers, p_dropout=0.0)
+        self.proj = nn.Conv1d(
+            filter_channels, self.half_channels * (num_bins * 3 - 1), 1
+        )
+        self.proj.weight.data.zero_()
+        self.proj.bias.data.zero_()
 
-  def forward(self, x, x_mask, g=None, reverse=False):
-    x0, x1 = torch.split(x, [self.half_channels]*2, 1)
-    h = self.pre(x0)
-    h = self.convs(h, x_mask, g=g)
-    h = self.proj(h) * x_mask
+    def forward(self, x, x_mask, g=None, reverse=False):
+        x0, x1 = torch.split(x, [self.half_channels] * 2, 1)
+        h = self.pre(x0)
+        h = self.convs(h, x_mask, g=g)
+        h = self.proj(h) * x_mask
 
-    b, c, t = x0.shape
-    h = h.reshape(b, c, -1, t).permute(0, 1, 3, 2) # [b, cx?, t] -> [b, c, t, ?]
+        b, c, t = x0.shape
+        h = h.reshape(b, c, -1, t).permute(0, 1, 3, 2)  # [b, cx?, t] -> [b, c, t, ?]
 
-    unnormalized_widths = h[..., :self.num_bins] / math.sqrt(self.filter_channels)
-    unnormalized_heights = h[..., self.num_bins:2*self.num_bins] / math.sqrt(self.filter_channels)
-    unnormalized_derivatives = h[..., 2 * self.num_bins:]
+        unnormalized_widths = h[..., : self.num_bins] / math.sqrt(self.filter_channels)
+        unnormalized_heights = h[..., self.num_bins : 2 * self.num_bins] / math.sqrt(
+            self.filter_channels
+        )
+        unnormalized_derivatives = h[..., 2 * self.num_bins :]
+
+        x1, logabsdet = piecewise_rational_quadratic_transform(
+            x1,
+            unnormalized_widths,
+            unnormalized_heights,
+            unnormalized_derivatives,
+            inverse=reverse,
+            tails="linear",
+            tail_bound=self.tail_bound,
+        )
+
+        x = torch.cat([x0, x1], 1) * x_mask
+        logdet = torch.sum(logabsdet * x_mask, [1, 2])
+        if not reverse:
+            return x, logdet
+        else:
+            return x
 
-    x1, logabsdet = piecewise_rational_quadratic_transform(x1,
-        unnormalized_widths,
-        unnormalized_heights,
-        unnormalized_derivatives,
-        inverse=reverse,
-        tails='linear',
-        tail_bound=self.tail_bound
-    )
 
-    x = torch.cat([x0, x1], 1) * x_mask
-    logdet = torch.sum(logabsdet * x_mask, [1,2])
-    if not reverse:
-        return x, logdet
-    else:
-        return x
 class TransformerCouplingLayer(nn.Module):
-  def __init__(self,
-      channels,
-      hidden_channels,
-      kernel_size,
-      n_layers,
-      n_heads,
-      p_dropout=0,
-      filter_channels=0,
-      mean_only=False,
-      wn_sharing_parameter=None,
-      gin_channels = 0
-      ):
-    assert channels % 2 == 0, "channels should be divisible by 2"
-    super().__init__()
-    self.channels = channels
-    self.hidden_channels = hidden_channels
-    self.kernel_size = kernel_size
-    self.n_layers = n_layers
-    self.half_channels = channels // 2
-    self.mean_only = mean_only
+    def __init__(
+        self,
+        channels,
+        hidden_channels,
+        kernel_size,
+        n_layers,
+        n_heads,
+        p_dropout=0,
+        filter_channels=0,
+        mean_only=False,
+        wn_sharing_parameter=None,
+        gin_channels=0,
+    ):
+        assert channels % 2 == 0, "channels should be divisible by 2"
+        super().__init__()
+        self.channels = channels
+        self.hidden_channels = hidden_channels
+        self.kernel_size = kernel_size
+        self.n_layers = n_layers
+        self.half_channels = channels // 2
+        self.mean_only = mean_only
 
-    self.pre = nn.Conv1d(self.half_channels, hidden_channels, 1)
-    self.enc = Encoder(hidden_channels, filter_channels, n_heads, n_layers, kernel_size, p_dropout, isflow = True, gin_channels = gin_channels) if wn_sharing_parameter is None else wn_sharing_parameter
-    self.post = nn.Conv1d(hidden_channels, self.half_channels * (2 - mean_only), 1)
-    self.post.weight.data.zero_()
-    self.post.bias.data.zero_()
+        self.pre = nn.Conv1d(self.half_channels, hidden_channels, 1)
+        self.enc = (
+            Encoder(
+                hidden_channels,
+                filter_channels,
+                n_heads,
+                n_layers,
+                kernel_size,
+                p_dropout,
+                isflow=True,
+                gin_channels=gin_channels,
+            )
+            if wn_sharing_parameter is None
+            else wn_sharing_parameter
+        )
+        self.post = nn.Conv1d(hidden_channels, self.half_channels * (2 - mean_only), 1)
+        self.post.weight.data.zero_()
+        self.post.bias.data.zero_()
 
-  def forward(self, x, x_mask, g=None, reverse=False):
-    x0, x1 = torch.split(x, [self.half_channels]*2, 1)
-    h = self.pre(x0) * x_mask
-    h = self.enc(h, x_mask, g=g)
-    stats = self.post(h) * x_mask
-    if not self.mean_only:
-      m, logs = torch.split(stats, [self.half_channels]*2, 1)
-    else:
-      m = stats
-      logs = torch.zeros_like(m)
+    def forward(self, x, x_mask, g=None, reverse=False):
+        x0, x1 = torch.split(x, [self.half_channels] * 2, 1)
+        h = self.pre(x0) * x_mask
+        h = self.enc(h, x_mask, g=g)
+        stats = self.post(h) * x_mask
+        if not self.mean_only:
+            m, logs = torch.split(stats, [self.half_channels] * 2, 1)
+        else:
+            m = stats
+            logs = torch.zeros_like(m)
 
-    if not reverse:
-      x1 = m + x1 * torch.exp(logs) * x_mask
-      x = torch.cat([x0, x1], 1)
-      logdet = torch.sum(logs, [1,2])
-      return x, logdet
-    else:
-      x1 = (x1 - m) * torch.exp(-logs) * x_mask
-      x = torch.cat([x0, x1], 1)
-      return x
+        if not reverse:
+            x1 = m + x1 * torch.exp(logs) * x_mask
+            x = torch.cat([x0, x1], 1)
+            logdet = torch.sum(logs, [1, 2])
+            return x, logdet
+        else:
+            x1 = (x1 - m) * torch.exp(-logs) * x_mask
+            x = torch.cat([x0, x1], 1)
+            return x
 
-    x1, logabsdet = piecewise_rational_quadratic_transform(x1,
-        unnormalized_widths,
-        unnormalized_heights,
-        unnormalized_derivatives,
-        inverse=reverse,
-        tails='linear',
-        tail_bound=self.tail_bound
-    )
+        x1, logabsdet = piecewise_rational_quadratic_transform(
+            x1,
+            unnormalized_widths,
+            unnormalized_heights,
+            unnormalized_derivatives,
+            inverse=reverse,
+            tails="linear",
+            tail_bound=self.tail_bound,
+        )
 
-    x = torch.cat([x0, x1], 1) * x_mask
-    logdet = torch.sum(logabsdet * x_mask, [1,2])
-    if not reverse:
-        return x, logdet
-    else:
-        return x
+        x = torch.cat([x0, x1], 1) * x_mask
+        logdet = torch.sum(logabsdet * x_mask, [1, 2])
+        if not reverse:
+            return x, logdet
+        else:
+            return x

monotonic_align/__init__.py

@@ -3,13 +3,14 @@ from torch import from_numpy
 
 from .core import maximum_path_jit
 
-def maximum_path(neg_cent, mask):
-  device = neg_cent.device
-  dtype = neg_cent.dtype
-  neg_cent = neg_cent.data.cpu().numpy().astype(float32)
-  path = zeros(neg_cent.shape, dtype=int32)
 
-  t_t_max = mask.sum(1)[:, 0].data.cpu().numpy().astype(int32)
-  t_s_max = mask.sum(2)[:, 0].data.cpu().numpy().astype(int32)
-  maximum_path_jit(path, neg_cent, t_t_max, t_s_max)
-  return from_numpy(path).to(device=device, dtype=dtype)
+def maximum_path(neg_cent, mask):
+    device = neg_cent.device
+    dtype = neg_cent.dtype
+    neg_cent = neg_cent.data.cpu().numpy().astype(float32)
+    path = zeros(neg_cent.shape, dtype=int32)
+
+    t_t_max = mask.sum(1)[:, 0].data.cpu().numpy().astype(int32)
+    t_s_max = mask.sum(2)[:, 0].data.cpu().numpy().astype(int32)
+    maximum_path_jit(path, neg_cent, t_t_max, t_s_max)
+    return from_numpy(path).to(device=device, dtype=dtype)

monotonic_align/core.py

@@ -1,35 +1,46 @@
 import numba
 
 
-@numba.jit(numba.void(numba.int32[:,:,::1], numba.float32[:,:,::1], numba.int32[::1], numba.int32[::1]), nopython=True, nogil=True)
+@numba.jit(
+    numba.void(
+        numba.int32[:, :, ::1],
+        numba.float32[:, :, ::1],
+        numba.int32[::1],
+        numba.int32[::1],
+    ),
+    nopython=True,
+    nogil=True,
+)
 def maximum_path_jit(paths, values, t_ys, t_xs):
-  b = paths.shape[0]
-  max_neg_val=-1e9
-  for i in range(int(b)):
-    path = paths[i]
-    value = values[i]
-    t_y = t_ys[i]
-    t_x = t_xs[i]
+    b = paths.shape[0]
+    max_neg_val = -1e9
+    for i in range(int(b)):
+        path = paths[i]
+        value = values[i]
+        t_y = t_ys[i]
+        t_x = t_xs[i]
 
-    v_prev = v_cur = 0.0
-    index = t_x - 1
+        v_prev = v_cur = 0.0
+        index = t_x - 1
 
-    for y in range(t_y):
-      for x in range(max(0, t_x + y - t_y), min(t_x, y + 1)):
-        if x == y:
-          v_cur = max_neg_val
-        else:
-          v_cur = value[y-1, x]
-        if x == 0:
-          if y == 0:
-            v_prev = 0.
-          else:
-            v_prev = max_neg_val
-        else:
-          v_prev = value[y-1, x-1]
-        value[y, x] += max(v_prev, v_cur)
+        for y in range(t_y):
+            for x in range(max(0, t_x + y - t_y), min(t_x, y + 1)):
+                if x == y:
+                    v_cur = max_neg_val
+                else:
+                    v_cur = value[y - 1, x]
+                if x == 0:
+                    if y == 0:
+                        v_prev = 0.0
+                    else:
+                        v_prev = max_neg_val
+                else:
+                    v_prev = value[y - 1, x - 1]
+                value[y, x] += max(v_prev, v_cur)
 
-    for y in range(t_y - 1, -1, -1):
-      path[y, index] = 1
-      if index != 0 and (index == y or value[y-1, index] < value[y-1, index-1]):
-        index = index - 1
+        for y in range(t_y - 1, -1, -1):
+            path[y, index] = 1
+            if index != 0 and (
+                index == y or value[y - 1, index] < value[y - 1, index - 1]
+            ):
+                index = index - 1

preprocess_text.py

@@ -1,64 +1,105 @@
 import json
-from random import shuffle
-
-import tqdm
-from text.cleaner import clean_text
 from collections import defaultdict
-stage = [1,2,3]
+from random import shuffle
+from typing import Optional
 
-transcription_path = 'filelists/genshin.list'
-train_path = 'filelists/train.list'
-val_path = 'filelists/val.list'
-config_path = "configs/config.json"
-val_per_spk = 4
-max_val_total = 8
+from tqdm import tqdm
+import click
+from text.cleaner import clean_text
 
-if 1 in stage:
-    with open( transcription_path+'.cleaned', 'w', encoding='utf-8') as f:
-        for line in tqdm.tqdm(open(transcription_path, encoding='utf-8').readlines()):
+
+@click.command()
+@click.option(
+    "--transcription-path",
+    default="filelists/genshin.list",
+    type=click.Path(exists=True, file_okay=True, dir_okay=False),
+)
+@click.option("--cleaned-path", default=None)
+@click.option("--train-path", default="filelists/train.list")
+@click.option("--val-path", default="filelists/val.list")
+@click.option(
+    "--config-path",
+    default="configs/config.json",
+    type=click.Path(exists=True, file_okay=True, dir_okay=False),
+)
+@click.option("--val-per-spk", default=4)
+@click.option("--max-val-total", default=8)
+@click.option("--clean/--no-clean", default=True)
+def main(
+    transcription_path: str,
+    cleaned_path: Optional[str],
+    train_path: str,
+    val_path: str,
+    config_path: str,
+    val_per_spk: int,
+    max_val_total: int,
+    clean: bool,
+):
+    if cleaned_path is None:
+        cleaned_path = transcription_path + ".cleaned"
+
+    if clean:
+        out_file = open(cleaned_path, "w", encoding="utf-8")
+        for line in tqdm(open(transcription_path, encoding="utf-8").readlines()):
             try:
-                utt, spk, language, text = line.strip().split('|')
+                utt, spk, language, text = line.strip().split("|")
                 norm_text, phones, tones, word2ph = clean_text(text, language)
-                f.write('{}|{}|{}|{}|{}|{}|{}\n'.format(utt, spk, language, norm_text, ' '.join(phones),
-                                                     " ".join([str(i) for i in tones]),
-                                                     " ".join([str(i) for i in word2ph])))
-            except Exception as error :
-                print("err!", utt, error)
+                out_file.write(
+                    "{}|{}|{}|{}|{}|{}|{}\n".format(
+                        utt,
+                        spk,
+                        language,
+                        norm_text,
+                        " ".join(phones),
+                        " ".join([str(i) for i in tones]),
+                        " ".join([str(i) for i in word2ph]),
+                    )
+                )
+            except Exception as error:
+                print("err!", line, error)
+
+        out_file.close()
+
+        transcription_path = cleaned_path
 
-if 2 in stage:
     spk_utt_map = defaultdict(list)
     spk_id_map = {}
     current_sid = 0
 
-    with open( transcription_path+'.cleaned', encoding='utf-8') as f:
+    with open(transcription_path, encoding="utf-8") as f:
         for line in f.readlines():
-            utt, spk, language, text, phones, tones, word2ph = line.strip().split('|')
+            utt, spk, language, text, phones, tones, word2ph = line.strip().split("|")
             spk_utt_map[spk].append(line)
+
             if spk not in spk_id_map.keys():
                 spk_id_map[spk] = current_sid
                 current_sid += 1
+
     train_list = []
     val_list = []
-    
+
     for spk, utts in spk_utt_map.items():
         shuffle(utts)
-        val_list+=utts[:val_per_spk]
-        train_list+=utts[val_per_spk:]
+        val_list += utts[:val_per_spk]
+        train_list += utts[val_per_spk:]
+
     if len(val_list) > max_val_total:
-        train_list+=val_list[max_val_total:]
+        train_list += val_list[max_val_total:]
         val_list = val_list[:max_val_total]
-    
-    with open( train_path,"w", encoding='utf-8') as f:
+
+    with open(train_path, "w", encoding="utf-8") as f:
         for line in train_list:
             f.write(line)
-    
-    with open(val_path, "w", encoding='utf-8') as f:
+
+    with open(val_path, "w", encoding="utf-8") as f:
         for line in val_list:
             f.write(line)
 
-if 3 in stage:
-    assert 2 in stage
-    config = json.load(open(config_path, encoding='utf-8'))
-    config["data"]['spk2id'] = spk_id_map
-    with open(config_path, 'w', encoding='utf-8') as f:
+    config = json.load(open(config_path, encoding="utf-8"))
+    config["data"]["spk2id"] = spk_id_map
+    with open(config_path, "w", encoding="utf-8") as f:
         json.dump(config, f, indent=2, ensure_ascii=False)
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -15,3 +15,6 @@ pypinyin
 cn2an
 gradio
 av
+mecab-python3
+loguru
+unidic-lite

resample.py

@@ -1,11 +1,9 @@
 import os
 import argparse
 import librosa
-import numpy as np
 from multiprocessing import Pool, cpu_count
 
 import soundfile
-from scipy.io import wavfile
 from tqdm import tqdm
 
 
@@ -13,30 +11,38 @@ def process(item):
     spkdir, wav_name, args = item
     speaker = spkdir.replace("\\", "/").split("/")[-1]
     wav_path = os.path.join(args.in_dir, speaker, wav_name)
-    if os.path.exists(wav_path) and '.wav' in wav_path:
+    if os.path.exists(wav_path) and ".wav" in wav_path:
         os.makedirs(os.path.join(args.out_dir, speaker), exist_ok=True)
         wav, sr = librosa.load(wav_path, sr=args.sr)
-        soundfile.write(
-            os.path.join(args.out_dir, speaker, wav_name),
-            wav,
-            sr
-        )
-
+        soundfile.write(os.path.join(args.out_dir, speaker, wav_name), wav, sr)
 
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument("--sr", type=int, default=44100, help="sampling rate")
-    parser.add_argument("--in_dir", type=str, default="./raw", help="path to source dir")
-    parser.add_argument("--out_dir", type=str, default="./dataset", help="path to target dir")
+    parser.add_argument(
+        "--in_dir", type=str, default="./raw", help="path to source dir"
+    )
+    parser.add_argument(
+        "--out_dir", type=str, default="./dataset", help="path to target dir"
+    )
     args = parser.parse_args()
-    # processs = 8
-    processs = cpu_count()-2 if cpu_count() >4 else 1
-    pool = Pool(processes=processs)
+    # processes = 8
+    processes = cpu_count() - 2 if cpu_count() > 4 else 1
+    pool = Pool(processes=processes)
 
     for speaker in os.listdir(args.in_dir):
         spk_dir = os.path.join(args.in_dir, speaker)
         if os.path.isdir(spk_dir):
             print(spk_dir)
-            for _ in tqdm(pool.imap_unordered(process, [(spk_dir, i, args) for i in os.listdir(spk_dir) if i.endswith("wav")])):
+            for _ in tqdm(
+                pool.imap_unordered(
+                    process,
+                    [
+                        (spk_dir, i, args)
+                        for i in os.listdir(spk_dir)
+                        if i.endswith("wav")
+                    ],
+                )
+            ):
                 pass

server.py

@@ -13,10 +13,11 @@ from scipy.io import wavfile
 
 # Flask Init
 app = Flask(__name__)
-app.config['JSON_AS_ASCII'] = False
+app.config["JSON_AS_ASCII"] = False
+
+
 def get_text(text, language_str, hps):
     norm_text, phone, tone, word2ph = clean_text(text, language_str)
-    print([f"{p}{t}" for p, t in zip(phone, tone)])
     phone, tone, language = cleaned_text_to_sequence(phone, tone, language_str)
 
     if hps.data.add_blank:
@@ -27,97 +28,143 @@ def get_text(text, language_str, hps):
             word2ph[i] = word2ph[i] * 2
         word2ph[0] += 1
     bert = get_bert(norm_text, word2ph, language_str)
+    del word2ph
+    assert bert.shape[-1] == len(phone), phone
 
-    assert bert.shape[-1] == len(phone)
-
+    if language_str == "ZH":
+        bert = bert
+        ja_bert = torch.zeros(768, len(phone))
+    elif language_str == "JA":
+        ja_bert = bert
+        bert = torch.zeros(1024, len(phone))
+    else:
+        bert = torch.zeros(1024, len(phone))
+        ja_bert = torch.zeros(768, len(phone))
+    assert bert.shape[-1] == len(
+        phone
+    ), f"Bert seq len {bert.shape[-1]} != {len(phone)}"
     phone = torch.LongTensor(phone)
     tone = torch.LongTensor(tone)
     language = torch.LongTensor(language)
+    return bert, ja_bert, phone, tone, language
 
-    return bert, phone, tone, language
 
-def infer(text, sdp_ratio, noise_scale, noise_scale_w,length_scale,sid):
-    bert, phones, tones, lang_ids = get_text(text,"ZH", hps,)
+def infer(text, sdp_ratio, noise_scale, noise_scale_w, length_scale, sid, language):
+    bert, ja_bert, phones, tones, lang_ids = get_text(text, language, hps)
     with torch.no_grad():
-        x_tst=phones.to(dev).unsqueeze(0)
-        tones=tones.to(dev).unsqueeze(0)
-        lang_ids=lang_ids.to(dev).unsqueeze(0)
+        x_tst = phones.to(dev).unsqueeze(0)
+        tones = tones.to(dev).unsqueeze(0)
+        lang_ids = lang_ids.to(dev).unsqueeze(0)
         bert = bert.to(dev).unsqueeze(0)
+        ja_bert = ja_bert.to(device).unsqueeze(0)
         x_tst_lengths = torch.LongTensor([phones.size(0)]).to(dev)
         speakers = torch.LongTensor([hps.data.spk2id[sid]]).to(dev)
-        audio = net_g.infer(x_tst, x_tst_lengths, speakers, tones, lang_ids,bert, sdp_ratio=sdp_ratio
-                           , noise_scale=noise_scale, noise_scale_w=noise_scale_w, length_scale=length_scale)[0][0,0].data.cpu().float().numpy()
+        audio = (
+            net_g.infer(
+                x_tst,
+                x_tst_lengths,
+                speakers,
+                tones,
+                lang_ids,
+                bert,
+                ja_bert,
+                sdp_ratio=sdp_ratio,
+                noise_scale=noise_scale,
+                noise_scale_w=noise_scale_w,
+                length_scale=length_scale,
+            )[0][0, 0]
+            .data.cpu()
+            .float()
+            .numpy()
+        )
         return audio
 
+
 def replace_punctuation(text, i=2):
     punctuation = "，。？！"
     for char in punctuation:
         text = text.replace(char, char * i)
     return text
 
+
 def wav2(i, o, format):
-    inp = avopen(i, 'rb')
-    out = avopen(o, 'wb', format=format)
-    if format == "ogg": format = "libvorbis"
+    inp = avopen(i, "rb")
+    out = avopen(o, "wb", format=format)
+    if format == "ogg":
+        format = "libvorbis"
 
     ostream = out.add_stream(format)
 
     for frame in inp.decode(audio=0):
-        for p in ostream.encode(frame): out.mux(p)
+        for p in ostream.encode(frame):
+            out.mux(p)
 
-    for p in ostream.encode(None): out.mux(p)
+    for p in ostream.encode(None):
+        out.mux(p)
 
     out.close()
     inp.close()
 
+
 # Load Generator
 hps = utils.get_hparams_from_file("./configs/config.json")
 
-dev='cuda'
+dev = "cuda"
 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,
-    **hps.model).to(dev)
+    **hps.model,
+).to(dev)
 _ = net_g.eval()
 
-_ = utils.load_checkpoint("logs/G_649000.pth", net_g, None,skip_optimizer=True)
+_ = utils.load_checkpoint("logs/G_649000.pth", net_g, None, skip_optimizer=True)
 
-@app.route("/",methods=['GET','POST'])
+
+@app.route("/")
 def main():
-    if request.method == 'GET':
-        try:
-            speaker = request.args.get('speaker')
-            text = request.args.get('text').replace("/n","")
-            sdp_ratio = float(request.args.get("sdp_ratio", 0.2))
-            noise = float(request.args.get("noise", 0.5))
-            noisew = float(request.args.get("noisew", 0.6))
-            length = float(request.args.get("length", 1.2))
-            if length >= 2:
-                return "Too big length"
-            if len(text) >=200:
-                return "Too long text"
-            fmt = request.args.get("format", "wav")
-            if None in (speaker, text):
-                return "Missing Parameter"
-            if fmt not in ("mp3", "wav", "ogg"):
-                return "Invalid Format"
-        except:
-            return "Invalid Parameter"
+    try:
+        speaker = request.args.get("speaker")
+        text = request.args.get("text").replace("/n", "")
+        sdp_ratio = float(request.args.get("sdp_ratio", 0.2))
+        noise = float(request.args.get("noise", 0.5))
+        noisew = float(request.args.get("noisew", 0.6))
+        length = float(request.args.get("length", 1.2))
+        language = request.args.get("language")
+        if length >= 2:
+            return "Too big length"
+        if len(text) >= 250:
+            return "Too long text"
+        fmt = request.args.get("format", "wav")
+        if None in (speaker, text):
+            return "Missing Parameter"
+        if fmt not in ("mp3", "wav", "ogg"):
+            return "Invalid Format"
+        if language not in ("JA", "ZH"):
+            return "Invalid language"
+    except:
+        return "Invalid Parameter"
 
-        with torch.no_grad():
-            audio = infer(text, sdp_ratio=sdp_ratio, noise_scale=noise, noise_scale_w=noisew, length_scale=length, sid=speaker)
+    with torch.no_grad():
+        audio = infer(
+            text,
+            sdp_ratio=sdp_ratio,
+            noise_scale=noise,
+            noise_scale_w=noisew,
+            length_scale=length,
+            sid=speaker,
+            language=language,
+        )
 
-        with BytesIO() as wav:
-            wavfile.write(wav, hps.data.sampling_rate, audio)
-            torch.cuda.empty_cache()
-            if fmt == "wav":
-                return Response(wav.getvalue(), mimetype="audio/wav")
-            wav.seek(0, 0)
-            with BytesIO() as ofp:
-                wav2(wav, ofp, fmt)
-                return Response(
-                    ofp.getvalue(),
-                    mimetype="audio/mpeg" if fmt == "mp3" else "audio/ogg"
-                )
+    with BytesIO() as wav:
+        wavfile.write(wav, hps.data.sampling_rate, audio)
+        torch.cuda.empty_cache()
+        if fmt == "wav":
+            return Response(wav.getvalue(), mimetype="audio/wav")
+        wav.seek(0, 0)
+        with BytesIO() as ofp:
+            wav2(wav, ofp, fmt)
+            return Response(
+                ofp.getvalue(), mimetype="audio/mpeg" if fmt == "mp3" else "audio/ogg"
+            )

text/__init__.py

@@ -3,26 +3,27 @@ from text.symbols import *
 
 _symbol_to_id = {s: i for i, s in enumerate(symbols)}
 
+
 def cleaned_text_to_sequence(cleaned_text, tones, language):
-  '''Converts a string of text to a sequence of IDs corresponding to the symbols in the text.
+    """Converts a string of text to a sequence of IDs corresponding to the symbols in the text.
     Args:
       text: string to convert to a sequence
     Returns:
       List of integers corresponding to the symbols in the text
-  '''
-  phones = [_symbol_to_id[symbol] for symbol in cleaned_text]
-  tone_start = language_tone_start_map[language]
-  tones = [i + tone_start for i in tones]
-  lang_id = language_id_map[language]
-  lang_ids = [lang_id for i in phones]
-  return phones, tones, lang_ids
+    """
+    phones = [_symbol_to_id[symbol] for symbol in cleaned_text]
+    tone_start = language_tone_start_map[language]
+    tones = [i + tone_start for i in tones]
+    lang_id = language_id_map[language]
+    lang_ids = [lang_id for i in phones]
+    return phones, tones, lang_ids
 
-def get_bert(norm_text, word2ph, language):
-  from .chinese_bert import get_bert_feature as zh_bert
-  from .english_bert_mock import get_bert_feature as en_bert
-  lang_bert_func_map = {
-    'ZH': zh_bert,
-    'EN': en_bert
-  }
-  bert = lang_bert_func_map[language](norm_text, word2ph)
-  return bert
+
+def get_bert(norm_text, word2ph, language, device):
+    from .chinese_bert import get_bert_feature as zh_bert
+    from .english_bert_mock import get_bert_feature as en_bert
+    from .japanese_bert import get_bert_feature as jp_bert
+
+    lang_bert_func_map = {"ZH": zh_bert, "EN": en_bert, "JP": jp_bert}
+    bert = lang_bert_func_map[language](norm_text, word2ph, device)
+    return bert

text/chinese.py

@@ -4,70 +4,74 @@ import re
 import cn2an
 from pypinyin import lazy_pinyin, Style
 
-from text import symbols
 from text.symbols import punctuation
 from text.tone_sandhi import ToneSandhi
 
 current_file_path = os.path.dirname(__file__)
-pinyin_to_symbol_map = {line.split("\t")[0]: line.strip().split("\t")[1] for line in
-                        open(os.path.join(current_file_path, 'opencpop-strict.txt')).readlines()}
+pinyin_to_symbol_map = {
+    line.split("\t")[0]: line.strip().split("\t")[1]
+    for line in open(os.path.join(current_file_path, "opencpop-strict.txt")).readlines()
+}
 
 import jieba.posseg as psg
 
 
 rep_map = {
-    '：': ',',
-    '；': ',',
-    '，': ',',
-    '。': '.',
-    '！': '!',
-    '？': '?',
-    '\n': '.',
+    "：": ",",
+    "；": ",",
+    "，": ",",
+    "。": ".",
+    "！": "!",
+    "？": "?",
+    "\n": ".",
     "·": ",",
-    '、': ",",
-    '...': '…',
-    '$': '.',
-    '“': "'",
-    '”': "'",
-    '‘': "'",
-    '’': "'",
-    '（': "'",
-    '）': "'",
-    '(': "'",
-    ')': "'",
-    '《': "'",
-    '》': "'",
-    '【': "'",
-    '】': "'",
-    '[': "'",
-    ']': "'",
-    '—': "-",
-    '～': "-",
-    '~': "-",
-    '「': "'",
-    '」': "'",
-
+    "、": ",",
+    "...": "…",
+    "$": ".",
+    "“": "'",
+    "”": "'",
+    "‘": "'",
+    "’": "'",
+    "（": "'",
+    "）": "'",
+    "(": "'",
+    ")": "'",
+    "《": "'",
+    "》": "'",
+    "【": "'",
+    "】": "'",
+    "[": "'",
+    "]": "'",
+    "—": "-",
+    "～": "-",
+    "~": "-",
+    "「": "'",
+    "」": "'",
 }
 
 tone_modifier = ToneSandhi()
 
+
 def replace_punctuation(text):
-    text = text.replace("嗯", "恩").replace("呣","母")
-    pattern = re.compile('|'.join(re.escape(p) for p in rep_map.keys()))
+    text = text.replace("嗯", "恩").replace("呣", "母")
+    pattern = re.compile("|".join(re.escape(p) for p in rep_map.keys()))
 
     replaced_text = pattern.sub(lambda x: rep_map[x.group()], text)
 
-    replaced_text = re.sub(r'[^\u4e00-\u9fa5'+"".join(punctuation)+r']+', '', replaced_text)
+    replaced_text = re.sub(
+        r"[^\u4e00-\u9fa5" + "".join(punctuation) + r"]+", "", replaced_text
+    )
 
     return replaced_text
 
+
 def g2p(text):
-    pattern = r'(?<=[{0}])\s*'.format(''.join(punctuation))
-    sentences = [i for i in re.split(pattern, text) if i.strip()!='']
+    pattern = r"(?<=[{0}])\s*".format("".join(punctuation))
+    sentences = [i for i in re.split(pattern, text) if i.strip() != ""]
     phones, tones, word2ph = _g2p(sentences)
     assert sum(word2ph) == len(phones)
-    assert len(word2ph) == len(text) #Sometimes it will crash,you can add a try-catch.
-    phones = ['_'] + phones + ["_"]
+    assert len(word2ph) == len(text)  # Sometimes it will crash,you can add a try-catch.
+    phones = ["_"] + phones + ["_"]
     tones = [0] + tones + [0]
     word2ph = [1] + word2ph + [1]
     return phones, tones, word2ph
@@ -76,10 +80,10 @@ def g2p(text):
 def _get_initials_finals(word):
     initials = []
     finals = []
-    orig_initials = lazy_pinyin(
-        word, neutral_tone_with_five=True, style=Style.INITIALS)
+    orig_initials = lazy_pinyin(word, neutral_tone_with_five=True, style=Style.INITIALS)
     orig_finals = lazy_pinyin(
-        word, neutral_tone_with_five=True, style=Style.FINALS_TONE3)
+        word, neutral_tone_with_five=True, style=Style.FINALS_TONE3
+    )
     for c, v in zip(orig_initials, orig_finals):
         initials.append(c)
         finals.append(v)
@@ -91,19 +95,17 @@ def _g2p(segments):
     tones_list = []
     word2ph = []
     for seg in segments:
-        pinyins = []
         # Replace all English words in the sentence
-        seg = re.sub('[a-zA-Z]+', '', seg)
+        seg = re.sub("[a-zA-Z]+", "", seg)
         seg_cut = psg.lcut(seg)
         initials = []
         finals = []
         seg_cut = tone_modifier.pre_merge_for_modify(seg_cut)
         for word, pos in seg_cut:
-            if pos == 'eng':
+            if pos == "eng":
                 continue
             sub_initials, sub_finals = _get_initials_finals(word)
-            sub_finals = tone_modifier.modified_tone(word, pos,
-                                                          sub_finals)
+            sub_finals = tone_modifier.modified_tone(word, pos, sub_finals)
             initials.append(sub_initials)
             finals.append(sub_finals)
 
@@ -112,52 +114,52 @@ def _g2p(segments):
         finals = sum(finals, [])
         #
         for c, v in zip(initials, finals):
-            raw_pinyin = c+v
+            raw_pinyin = c + v
             # NOTE: post process for pypinyin outputs
             # we discriminate i, ii and iii
             if c == v:
                 assert c in punctuation
                 phone = [c]
-                tone = '0'
+                tone = "0"
                 word2ph.append(1)
             else:
                 v_without_tone = v[:-1]
                 tone = v[-1]
 
-                pinyin = c+v_without_tone
-                assert tone in '12345'
+                pinyin = c + v_without_tone
+                assert tone in "12345"
 
                 if c:
                     # 多音节
                     v_rep_map = {
-                        "uei": 'ui',
-                        'iou': 'iu',
-                        'uen': 'un',
+                        "uei": "ui",
+                        "iou": "iu",
+                        "uen": "un",
                     }
                     if v_without_tone in v_rep_map.keys():
-                        pinyin = c+v_rep_map[v_without_tone]
+                        pinyin = c + v_rep_map[v_without_tone]
                 else:
                     # 单音节
                     pinyin_rep_map = {
-                        'ing': 'ying',
-                        'i': 'yi',
-                        'in': 'yin',
-                        'u': 'wu',
+                        "ing": "ying",
+                        "i": "yi",
+                        "in": "yin",
+                        "u": "wu",
                     }
                     if pinyin in pinyin_rep_map.keys():
                         pinyin = pinyin_rep_map[pinyin]
                     else:
                         single_rep_map = {
-                            'v': 'yu',
-                            'e': 'e',
-                            'i': 'y',
-                            'u': 'w',
+                            "v": "yu",
+                            "e": "e",
+                            "i": "y",
+                            "u": "w",
                         }
                         if pinyin[0] in single_rep_map.keys():
-                            pinyin = single_rep_map[pinyin[0]]+pinyin[1:]
+                            pinyin = single_rep_map[pinyin[0]] + pinyin[1:]
 
                 assert pinyin in pinyin_to_symbol_map.keys(), (pinyin, seg, raw_pinyin)
-                phone = pinyin_to_symbol_map[pinyin].split(' ')
+                phone = pinyin_to_symbol_map[pinyin].split(" ")
                 word2ph.append(len(phone))
 
             phones_list += phone
@@ -165,20 +167,23 @@ def _g2p(segments):
     return phones_list, tones_list, word2ph
 
 
-
 def text_normalize(text):
-    numbers = re.findall(r'\d+(?:\.?\d+)?', text)
+    numbers = re.findall(r"\d+(?:\.?\d+)?", text)
     for number in numbers:
         text = text.replace(number, cn2an.an2cn(number), 1)
     text = replace_punctuation(text)
     return text
 
+
 def get_bert_feature(text, word2ph):
-    from  text import chinese_bert
+    from text import chinese_bert
+
     return chinese_bert.get_bert_feature(text, word2ph)
 
-if __name__ == '__main__':
+
+if __name__ == "__main__":
     from text.chinese_bert import get_bert_feature
+
     text = "啊！但是《原神》是由,米哈\游自主，  [研发]的一款全.新开放世界.冒险游戏"
     text = text_normalize(text)
     print(text)

text/chinese_bert.py

@@ -2,28 +2,29 @@ import torch
 import sys
 from transformers import AutoTokenizer, AutoModelForMaskedLM
 
-device = torch.device(
-        "cuda"
-        if torch.cuda.is_available()
-        else (
-            "mps"
-            if sys.platform == "darwin" and torch.backends.mps.is_available()
-            else "cpu"
-        )
-    )
-
 tokenizer = AutoTokenizer.from_pretrained("./bert/chinese-roberta-wwm-ext-large")
-model = AutoModelForMaskedLM.from_pretrained("./bert/chinese-roberta-wwm-ext-large").to(device)
 
-def get_bert_feature(text, word2ph):
+
+def get_bert_feature(text, word2ph, device=None):
+    if (
+        sys.platform == "darwin"
+        and torch.backends.mps.is_available()
+        and device == "cpu"
+    ):
+        device = "mps"
+    if not device:
+        device = "cuda"
+    model = AutoModelForMaskedLM.from_pretrained(
+        "./bert/chinese-roberta-wwm-ext-large"
+    ).to(device)
     with torch.no_grad():
-        inputs = tokenizer(text, return_tensors='pt')
+        inputs = tokenizer(text, return_tensors="pt")
         for i in inputs:
             inputs[i] = inputs[i].to(device)
         res = model(**inputs, output_hidden_states=True)
-        res = torch.cat(res['hidden_states'][-3:-2], -1)[0].cpu()
+        res = torch.cat(res["hidden_states"][-3:-2], -1)[0].cpu()
 
-    assert len(word2ph) == len(text)+2
+    assert len(word2ph) == len(text) + 2
     word2phone = word2ph
     phone_level_feature = []
     for i in range(len(word2phone)):
@@ -32,15 +33,53 @@ def get_bert_feature(text, word2ph):
 
     phone_level_feature = torch.cat(phone_level_feature, dim=0)
 
-
     return phone_level_feature.T
 
-if __name__ == '__main__':
-    # feature = get_bert_feature('你好,我是说的道理。')
+
+if __name__ == "__main__":
     import torch
 
     word_level_feature = torch.rand(38, 1024)  # 12个词,每个词1024维特征
-    word2phone = [1, 2, 1, 2, 2, 1, 2, 2, 1, 2, 2, 1, 2, 2, 2, 2, 2, 1, 1, 2, 2, 1, 2, 2, 2, 2, 1, 2, 2, 2, 2, 2, 1, 2, 2, 2, 2, 1]
+    word2phone = [
+        1,
+        2,
+        1,
+        2,
+        2,
+        1,
+        2,
+        2,
+        1,
+        2,
+        2,
+        1,
+        2,
+        2,
+        2,
+        2,
+        2,
+        1,
+        1,
+        2,
+        2,
+        1,
+        2,
+        2,
+        2,
+        2,
+        1,
+        2,
+        2,
+        2,
+        2,
+        2,
+        1,
+        2,
+        2,
+        2,
+        2,
+        1,
+    ]
 
     # 计算总帧数
     total_frames = sum(word2phone)
@@ -56,4 +95,3 @@ if __name__ == '__main__':
 
     phone_level_feature = torch.cat(phone_level_feature, dim=0)
     print(phone_level_feature.shape)  # torch.Size([36, 1024])
-

text/cleaner.py

@@ -1,9 +1,7 @@
-from text import chinese, cleaned_text_to_sequence
+from text import chinese, japanese, cleaned_text_to_sequence
 
 
-language_module_map = {
-    'ZH': chinese
-}
+language_module_map = {"ZH": chinese, "JP": japanese}
 
 
 def clean_text(text, language):
@@ -12,6 +10,7 @@ def clean_text(text, language):
     phones, tones, word2ph = language_module.g2p(norm_text)
     return norm_text, phones, tones, word2ph
 
+
 def clean_text_bert(text, language):
     language_module = language_module_map[language]
     norm_text = language_module.text_normalize(text)
@@ -19,9 +18,11 @@ def clean_text_bert(text, language):
     bert = language_module.get_bert_feature(norm_text, word2ph)
     return phones, tones, bert
 
+
 def text_to_sequence(text, language):
     norm_text, phones, tones, word2ph = clean_text(text, language)
     return cleaned_text_to_sequence(phones, tones, language)
 
-if __name__ == '__main__':
+
+if __name__ == "__main__":
     pass

text/cmudict.rep

@@ -7,7 +7,7 @@
 ## origin.
 ##
 ## cmudict.0.6 is the fifth release of cmudict, first released as cmudict.0.1
-## in September of 1993.  There was no generally available public release 
+## in September of 1993.  There was no generally available public release
 ## of version 0.5.
 ##
 ## See the README in this directory before you use this dictionary.
@@ -16,11 +16,11 @@
 ## Alex Rudnicky, Jack Mostow, Roni Rosenfeld, Richard Stern,
 ## Matthew Siegler, Kevin Lenzo, Maxine Eskenazi, Mosur Ravishankar,
 ## Eric Thayer, Kristie Seymore, and Raj Reddy at CMU; Lin Chase at
-## LIMSI; Doug Paul at MIT Lincoln Labs; Ben Serridge at MIT SLS; Murray 
-## Spiegel at Bellcore; Tony Robinson at Cambridge UK; David Bowness of 
-## CAE Electronics Ltd. and CRIM; Stephen Hocking; Jerry Quinn at BNR 
-## Canada, and Marshal Midden for bringing to our attention problems and 
-## inadequacies with the first releases. Most special thanks to Bob Weide 
+## LIMSI; Doug Paul at MIT Lincoln Labs; Ben Serridge at MIT SLS; Murray
+## Spiegel at Bellcore; Tony Robinson at Cambridge UK; David Bowness of
+## CAE Electronics Ltd. and CRIM; Stephen Hocking; Jerry Quinn at BNR
+## Canada, and Marshal Midden for bringing to our attention problems and
+## inadequacies with the first releases. Most special thanks to Bob Weide
 ## for all his work on prior versions of the dictionary.
 ##
 ## We welcome input from users and will continue to acknowledge such input
@@ -37,12 +37,12 @@
 ## so keep your eyes open for problems and mail them to me.
 ##
 ## We hope this dictionary is an improvement over cmudict.0.4.
-## 
+##
 ## email: cmudict@cs.cmu.edu
 ## web:   http://www.speech.cs.cmu.edu/cgi-bin/cmudict
 ## ftp:   ftp://ftp.cs.cmu.edu/project/speech/dict/
 ##
-## Thank you for your continued interest in the CMU Pronouncing 
+## Thank you for your continued interest in the CMU Pronouncing
 ## Dictionary.  Further addictions and improvements are planned
 ## for forthcoming releases.
 ##

text/english.py

@@ -2,40 +2,112 @@ import pickle
 import os
 import re
 from g2p_en import G2p
-from string import punctuation
 
 from text import symbols
 
 current_file_path = os.path.dirname(__file__)
-CMU_DICT_PATH = os.path.join(current_file_path, 'cmudict.rep')
-CACHE_PATH = os.path.join(current_file_path, 'cmudict_cache.pickle')
+CMU_DICT_PATH = os.path.join(current_file_path, "cmudict.rep")
+CACHE_PATH = os.path.join(current_file_path, "cmudict_cache.pickle")
 _g2p = G2p()
 
-arpa = {'AH0', 'S', 'AH1', 'EY2', 'AE2', 'EH0', 'OW2', 'UH0', 'NG', 'B', 'G', 'AY0', 'M', 'AA0', 'F', 'AO0', 'ER2', 'UH1', 'IY1', 'AH2', 'DH', 'IY0', 'EY1', 'IH0', 'K', 'N', 'W', 'IY2', 'T', 'AA1', 'ER1', 'EH2', 'OY0', 'UH2', 'UW1', 'Z', 'AW2', 'AW1', 'V', 'UW2', 'AA2', 'ER', 'AW0', 'UW0', 'R', 'OW1', 'EH1', 'ZH', 'AE0', 'IH2', 'IH', 'Y', 'JH', 'P', 'AY1', 'EY0', 'OY2', 'TH', 'HH', 'D', 'ER0', 'CH', 'AO1', 'AE1', 'AO2', 'OY1', 'AY2', 'IH1', 'OW0', 'L', 'SH'}
+arpa = {
+    "AH0",
+    "S",
+    "AH1",
+    "EY2",
+    "AE2",
+    "EH0",
+    "OW2",
+    "UH0",
+    "NG",
+    "B",
+    "G",
+    "AY0",
+    "M",
+    "AA0",
+    "F",
+    "AO0",
+    "ER2",
+    "UH1",
+    "IY1",
+    "AH2",
+    "DH",
+    "IY0",
+    "EY1",
+    "IH0",
+    "K",
+    "N",
+    "W",
+    "IY2",
+    "T",
+    "AA1",
+    "ER1",
+    "EH2",
+    "OY0",
+    "UH2",
+    "UW1",
+    "Z",
+    "AW2",
+    "AW1",
+    "V",
+    "UW2",
+    "AA2",
+    "ER",
+    "AW0",
+    "UW0",
+    "R",
+    "OW1",
+    "EH1",
+    "ZH",
+    "AE0",
+    "IH2",
+    "IH",
+    "Y",
+    "JH",
+    "P",
+    "AY1",
+    "EY0",
+    "OY2",
+    "TH",
+    "HH",
+    "D",
+    "ER0",
+    "CH",
+    "AO1",
+    "AE1",
+    "AO2",
+    "OY1",
+    "AY2",
+    "IH1",
+    "OW0",
+    "L",
+    "SH",
+}
 
 
 def post_replace_ph(ph):
     rep_map = {
-        '：': ',',
-        '；': ',',
-        '，': ',',
-        '。': '.',
-        '！': '!',
-        '？': '?',
-        '\n': '.',
+        "：": ",",
+        "；": ",",
+        "，": ",",
+        "。": ".",
+        "！": "!",
+        "？": "?",
+        "\n": ".",
         "·": ",",
-        '、': ",",
-        '...': '…',
-        'v': "V"
+        "、": ",",
+        "...": "…",
+        "v": "V",
     }
     if ph in rep_map.keys():
         ph = rep_map[ph]
     if ph in symbols:
         return ph
     if ph not in symbols:
-        ph = 'UNK'
+        ph = "UNK"
     return ph
 
+
 def read_dict():
     g2p_dict = {}
     start_line = 49
@@ -45,13 +117,13 @@ def read_dict():
         while line:
             if line_index >= start_line:
                 line = line.strip()
-                word_split = line.split('  ')
+                word_split = line.split("  ")
                 word = word_split[0]
 
-                syllable_split = word_split[1].split(' - ')
+                syllable_split = word_split[1].split(" - ")
                 g2p_dict[word] = []
                 for syllable in syllable_split:
-                    phone_split = syllable.split(' ')
+                    phone_split = syllable.split(" ")
                     g2p_dict[word].append(phone_split)
 
             line_index = line_index + 1
@@ -61,13 +133,13 @@ def read_dict():
 
 
 def cache_dict(g2p_dict, file_path):
-    with open(file_path, 'wb') as pickle_file:
+    with open(file_path, "wb") as pickle_file:
         pickle.dump(g2p_dict, pickle_file)
 
 
 def get_dict():
     if os.path.exists(CACHE_PATH):
-        with open(CACHE_PATH, 'rb') as pickle_file:
+        with open(CACHE_PATH, "rb") as pickle_file:
             g2p_dict = pickle.load(pickle_file)
     else:
         g2p_dict = read_dict()
@@ -75,15 +147,18 @@ def get_dict():
 
     return g2p_dict
 
+
 eng_dict = get_dict()
 
+
 def refine_ph(phn):
     tone = 0
-    if re.search(r'\d$', phn):
+    if re.search(r"\d$", phn):
         tone = int(phn[-1]) + 1
         phn = phn[:-1]
     return phn.lower(), tone
 
+
 def refine_syllables(syllables):
     tones = []
     phonemes = []
@@ -100,8 +175,8 @@ def text_normalize(text):
     # todo: eng text normalize
     return text
 
-def g2p(text):
 
+def g2p(text):
     phones = []
     tones = []
     words = re.split(r"([,;.\-\?\!\s+])", text)
@@ -126,6 +201,7 @@ def g2p(text):
     phones = [post_replace_ph(i) for i in phones]
     return phones, tones, word2ph
 
+
 if __name__ == "__main__":
     # print(get_dict())
     # print(eng_word_to_phoneme("hello"))
@@ -135,4 +211,4 @@ if __name__ == "__main__":
     #     for group in syllables:
     #         for ph in group:
     #             all_phones.add(ph)
-    # print(all_phones)
+    # print(all_phones)

text/japanese.py

@@ -1,104 +1,586 @@
-# modified from https://github.com/CjangCjengh/vits/blob/main/text/japanese.py
+# Convert Japanese text to phonemes which is
+# compatible with Julius https://github.com/julius-speech/segmentation-kit
 import re
-import sys
+import unicodedata
 
-import pyopenjtalk
+from transformers import AutoTokenizer
 
-from text import symbols
+from text import punctuation, symbols
 
-# Regular expression matching Japanese without punctuation marks:
-_japanese_characters = re.compile(
-    r'[A-Za-z\d\u3005\u3040-\u30ff\u4e00-\u9fff\uff11-\uff19\uff21-\uff3a\uff41-\uff5a\uff66-\uff9d]')
+try:
+    import MeCab
+except ImportError as e:
+    raise ImportError("Japanese requires mecab-python3 and unidic-lite.") from e
+from num2words import num2words
 
-# Regular expression matching non-Japanese characters or punctuation marks:
-_japanese_marks = re.compile(
-    r'[^A-Za-z\d\u3005\u3040-\u30ff\u4e00-\u9fff\uff11-\uff19\uff21-\uff3a\uff41-\uff5a\uff66-\uff9d]')
+_CONVRULES = [
+    # Conversion of 2 letters
+    "アァ/ a a",
+    "イィ/ i i",
+    "イェ/ i e",
+    "イャ/ y a",
+    "ウゥ/ u:",
+    "エェ/ e e",
+    "オォ/ o:",
+    "カァ/ k a:",
+    "キィ/ k i:",
+    "クゥ/ k u:",
+    "クャ/ ky a",
+    "クュ/ ky u",
+    "クョ/ ky o",
+    "ケェ/ k e:",
+    "コォ/ k o:",
+    "ガァ/ g a:",
+    "ギィ/ g i:",
+    "グゥ/ g u:",
+    "グャ/ gy a",
+    "グュ/ gy u",
+    "グョ/ gy o",
+    "ゲェ/ g e:",
+    "ゴォ/ g o:",
+    "サァ/ s a:",
+    "シィ/ sh i:",
+    "スゥ/ s u:",
+    "スャ/ sh a",
+    "スュ/ sh u",
+    "スョ/ sh o",
+    "セェ/ s e:",
+    "ソォ/ s o:",
+    "ザァ/ z a:",
+    "ジィ/ j i:",
+    "ズゥ/ z u:",
+    "ズャ/ zy a",
+    "ズュ/ zy u",
+    "ズョ/ zy o",
+    "ゼェ/ z e:",
+    "ゾォ/ z o:",
+    "タァ/ t a:",
+    "チィ/ ch i:",
+    "ツァ/ ts a",
+    "ツィ/ ts i",
+    "ツゥ/ ts u:",
+    "ツャ/ ch a",
+    "ツュ/ ch u",
+    "ツョ/ ch o",
+    "ツェ/ ts e",
+    "ツォ/ ts o",
+    "テェ/ t e:",
+    "トォ/ t o:",
+    "ダァ/ d a:",
+    "ヂィ/ j i:",
+    "ヅゥ/ d u:",
+    "ヅャ/ zy a",
+    "ヅュ/ zy u",
+    "ヅョ/ zy o",
+    "デェ/ d e:",
+    "ドォ/ d o:",
+    "ナァ/ n a:",
+    "ニィ/ n i:",
+    "ヌゥ/ n u:",
+    "ヌャ/ ny a",
+    "ヌュ/ ny u",
+    "ヌョ/ ny o",
+    "ネェ/ n e:",
+    "ノォ/ n o:",
+    "ハァ/ h a:",
+    "ヒィ/ h i:",
+    "フゥ/ f u:",
+    "フャ/ hy a",
+    "フュ/ hy u",
+    "フョ/ hy o",
+    "ヘェ/ h e:",
+    "ホォ/ h o:",
+    "バァ/ b a:",
+    "ビィ/ b i:",
+    "ブゥ/ b u:",
+    "フャ/ hy a",
+    "ブュ/ by u",
+    "フョ/ hy o",
+    "ベェ/ b e:",
+    "ボォ/ b o:",
+    "パァ/ p a:",
+    "ピィ/ p i:",
+    "プゥ/ p u:",
+    "プャ/ py a",
+    "プュ/ py u",
+    "プョ/ py o",
+    "ペェ/ p e:",
+    "ポォ/ p o:",
+    "マァ/ m a:",
+    "ミィ/ m i:",
+    "ムゥ/ m u:",
+    "ムャ/ my a",
+    "ムュ/ my u",
+    "ムョ/ my o",
+    "メェ/ m e:",
+    "モォ/ m o:",
+    "ヤァ/ y a:",
+    "ユゥ/ y u:",
+    "ユャ/ y a:",
+    "ユュ/ y u:",
+    "ユョ/ y o:",
+    "ヨォ/ y o:",
+    "ラァ/ r a:",
+    "リィ/ r i:",
+    "ルゥ/ r u:",
+    "ルャ/ ry a",
+    "ルュ/ ry u",
+    "ルョ/ ry o",
+    "レェ/ r e:",
+    "ロォ/ r o:",
+    "ワァ/ w a:",
+    "ヲォ/ o:",
+    "ディ/ d i",
+    "デェ/ d e:",
+    "デャ/ dy a",
+    "デュ/ dy u",
+    "デョ/ dy o",
+    "ティ/ t i",
+    "テェ/ t e:",
+    "テャ/ ty a",
+    "テュ/ ty u",
+    "テョ/ ty o",
+    "スィ/ s i",
+    "ズァ/ z u a",
+    "ズィ/ z i",
+    "ズゥ/ z u",
+    "ズャ/ zy a",
+    "ズュ/ zy u",
+    "ズョ/ zy o",
+    "ズェ/ z e",
+    "ズォ/ z o",
+    "キャ/ ky a",
+    "キュ/ ky u",
+    "キョ/ ky o",
+    "シャ/ sh a",
+    "シュ/ sh u",
+    "シェ/ sh e",
+    "ショ/ sh o",
+    "チャ/ ch a",
+    "チュ/ ch u",
+    "チェ/ ch e",
+    "チョ/ ch o",
+    "トゥ/ t u",
+    "トャ/ ty a",
+    "トュ/ ty u",
+    "トョ/ ty o",
+    "ドァ/ d o a",
+    "ドゥ/ d u",
+    "ドャ/ dy a",
+    "ドュ/ dy u",
+    "ドョ/ dy o",
+    "ドォ/ d o:",
+    "ニャ/ ny a",
+    "ニュ/ ny u",
+    "ニョ/ ny o",
+    "ヒャ/ hy a",
+    "ヒュ/ hy u",
+    "ヒョ/ hy o",
+    "ミャ/ my a",
+    "ミュ/ my u",
+    "ミョ/ my o",
+    "リャ/ ry a",
+    "リュ/ ry u",
+    "リョ/ ry o",
+    "ギャ/ gy a",
+    "ギュ/ gy u",
+    "ギョ/ gy o",
+    "ヂェ/ j e",
+    "ヂャ/ j a",
+    "ヂュ/ j u",
+    "ヂョ/ j o",
+    "ジェ/ j e",
+    "ジャ/ j a",
+    "ジュ/ j u",
+    "ジョ/ j o",
+    "ビャ/ by a",
+    "ビュ/ by u",
+    "ビョ/ by o",
+    "ピャ/ py a",
+    "ピュ/ py u",
+    "ピョ/ py o",
+    "ウァ/ u a",
+    "ウィ/ w i",
+    "ウェ/ w e",
+    "ウォ/ w o",
+    "ファ/ f a",
+    "フィ/ f i",
+    "フゥ/ f u",
+    "フャ/ hy a",
+    "フュ/ hy u",
+    "フョ/ hy o",
+    "フェ/ f e",
+    "フォ/ f o",
+    "ヴァ/ b a",
+    "ヴィ/ b i",
+    "ヴェ/ b e",
+    "ヴォ/ b o",
+    "ヴュ/ by u",
+    # Conversion of 1 letter
+    "ア/ a",
+    "イ/ i",
+    "ウ/ u",
+    "エ/ e",
+    "オ/ o",
+    "カ/ k a",
+    "キ/ k i",
+    "ク/ k u",
+    "ケ/ k e",
+    "コ/ k o",
+    "サ/ s a",
+    "シ/ sh i",
+    "ス/ s u",
+    "セ/ s e",
+    "ソ/ s o",
+    "タ/ t a",
+    "チ/ ch i",
+    "ツ/ ts u",
+    "テ/ t e",
+    "ト/ t o",
+    "ナ/ n a",
+    "ニ/ n i",
+    "ヌ/ n u",
+    "ネ/ n e",
+    "ノ/ n o",
+    "ハ/ h a",
+    "ヒ/ h i",
+    "フ/ f u",
+    "ヘ/ h e",
+    "ホ/ h o",
+    "マ/ m a",
+    "ミ/ m i",
+    "ム/ m u",
+    "メ/ m e",
+    "モ/ m o",
+    "ラ/ r a",
+    "リ/ r i",
+    "ル/ r u",
+    "レ/ r e",
+    "ロ/ r o",
+    "ガ/ g a",
+    "ギ/ g i",
+    "グ/ g u",
+    "ゲ/ g e",
+    "ゴ/ g o",
+    "ザ/ z a",
+    "ジ/ j i",
+    "ズ/ z u",
+    "ゼ/ z e",
+    "ゾ/ z o",
+    "ダ/ d a",
+    "ヂ/ j i",
+    "ヅ/ z u",
+    "デ/ d e",
+    "ド/ d o",
+    "バ/ b a",
+    "ビ/ b i",
+    "ブ/ b u",
+    "ベ/ b e",
+    "ボ/ b o",
+    "パ/ p a",
+    "ピ/ p i",
+    "プ/ p u",
+    "ペ/ p e",
+    "ポ/ p o",
+    "ヤ/ y a",
+    "ユ/ y u",
+    "ヨ/ y o",
+    "ワ/ w a",
+    "ヰ/ i",
+    "ヱ/ e",
+    "ヲ/ o",
+    "ン/ N",
+    "ッ/ q",
+    "ヴ/ b u",
+    "ー/:",
+    # Try converting broken text
+    "ァ/ a",
+    "ィ/ i",
+    "ゥ/ u",
+    "ェ/ e",
+    "ォ/ o",
+    "ヮ/ w a",
+    "ォ/ o",
+    # Symbols
+    "、/ ,",
+    "。/ .",
+    "！/ !",
+    "？/ ?",
+    "・/ ,",
+]
 
-# List of (symbol, Japanese) pairs for marks:
-_symbols_to_japanese = [(re.compile('%s' % x[0]), x[1]) for x in [
-    ('％', 'パーセント')
-]]
+_COLON_RX = re.compile(":+")
+_REJECT_RX = re.compile("[^ a-zA-Z:,.?]")
 
 
-# List of (consonant, sokuon) pairs:
-_real_sokuon = [(re.compile('%s' % x[0]), x[1]) for x in [
-    (r'Q([↑↓]*[kg])', r'k#\1'),
-    (r'Q([↑↓]*[tdjʧ])', r't#\1'),
-    (r'Q([↑↓]*[sʃ])', r's\1'),
-    (r'Q([↑↓]*[pb])', r'p#\1')
-]]
-
-# List of (consonant, hatsuon) pairs:
-_real_hatsuon = [(re.compile('%s' % x[0]), x[1]) for x in [
-    (r'N([↑↓]*[pbm])', r'm\1'),
-    (r'N([↑↓]*[ʧʥj])', r'n^\1'),
-    (r'N([↑↓]*[tdn])', r'n\1'),
-    (r'N([↑↓]*[kg])', r'ŋ\1')
-]]
+def _makerulemap():
+    l = [tuple(x.split("/")) for x in _CONVRULES]
+    return tuple({k: v for k, v in l if len(k) == i} for i in (1, 2))
 
 
-
-def post_replace_ph(ph):
-    rep_map = {
-        '：': ',',
-        '；': ',',
-        '，': ',',
-        '。': '.',
-        '！': '!',
-        '？': '?',
-        '\n': '.',
-        "·": ",",
-        '、': ",",
-        '...': '…',
-        'v': "V"
-    }
-    if ph in rep_map.keys():
-        ph = rep_map[ph]
-    if ph in symbols:
-        return ph
-    if ph not in symbols:
-        ph = 'UNK'
-    return ph
-
-def symbols_to_japanese(text):
-    for regex, replacement in _symbols_to_japanese:
-        text = re.sub(regex, replacement, text)
-    return text
+_RULEMAP1, _RULEMAP2 = _makerulemap()
 
 
-def preprocess_jap(text):
-    '''Reference https://r9y9.github.io/ttslearn/latest/notebooks/ch10_Recipe-Tacotron.html'''
-    text = symbols_to_japanese(text)
-    sentences = re.split(_japanese_marks, text)
-    marks = re.findall(_japanese_marks, text)
-    text = []
-    for i, sentence in enumerate(sentences):
-        if re.match(_japanese_characters, sentence):
-            p = pyopenjtalk.g2p(sentence)
-            text += p.split(" ")
+def kata2phoneme(text: str) -> str:
+    """Convert katakana text to phonemes."""
+    text = text.strip()
+    res = []
+    while text:
+        if len(text) >= 2:
+            x = _RULEMAP2.get(text[:2])
+            if x is not None:
+                text = text[2:]
+                res += x.split(" ")[1:]
+                continue
+        x = _RULEMAP1.get(text[0])
+        if x is not None:
+            text = text[1:]
+            res += x.split(" ")[1:]
+            continue
+        res.append(text[0])
+        text = text[1:]
+    # res = _COLON_RX.sub(":", res)
+    return res
+
+
+_KATAKANA = "".join(chr(ch) for ch in range(ord("ァ"), ord("ン") + 1))
+_HIRAGANA = "".join(chr(ch) for ch in range(ord("ぁ"), ord("ん") + 1))
+_HIRA2KATATRANS = str.maketrans(_HIRAGANA, _KATAKANA)
+
+
+def hira2kata(text: str) -> str:
+    text = text.translate(_HIRA2KATATRANS)
+    return text.replace("う゛", "ヴ")
+
+
+_SYMBOL_TOKENS = set(list("・、。？！"))
+_NO_YOMI_TOKENS = set(list("「」『』―（）［］[]"))
+_TAGGER = MeCab.Tagger()
+
+
+def text2kata(text: str) -> str:
+    parsed = _TAGGER.parse(text)
+    res = []
+    for line in parsed.split("\n"):
+        if line == "EOS":
+            break
+        parts = line.split("\t")
+
+        word, yomi = parts[0], parts[1]
+        if yomi:
+            res.append(yomi)
+        else:
+            if word in _SYMBOL_TOKENS:
+                res.append(word)
+            elif word in ("っ", "ッ"):
+                res.append("ッ")
+            elif word in _NO_YOMI_TOKENS:
+                pass
+            else:
+                res.append(word)
+    return hira2kata("".join(res))
+
+
+_ALPHASYMBOL_YOMI = {
+    "#": "シャープ",
+    "%": "パーセント",
+    "&": "アンド",
+    "+": "プラス",
+    "-": "マイナス",
+    ":": "コロン",
+    ";": "セミコロン",
+    "<": "小なり",
+    "=": "イコール",
+    ">": "大なり",
+    "@": "アット",
+    "a": "エー",
+    "b": "ビー",
+    "c": "シー",
+    "d": "ディー",
+    "e": "イー",
+    "f": "エフ",
+    "g": "ジー",
+    "h": "エイチ",
+    "i": "アイ",
+    "j": "ジェー",
+    "k": "ケー",
+    "l": "エル",
+    "m": "エム",
+    "n": "エヌ",
+    "o": "オー",
+    "p": "ピー",
+    "q": "キュー",
+    "r": "アール",
+    "s": "エス",
+    "t": "ティー",
+    "u": "ユー",
+    "v": "ブイ",
+    "w": "ダブリュー",
+    "x": "エックス",
+    "y": "ワイ",
+    "z": "ゼット",
+    "α": "アルファ",
+    "β": "ベータ",
+    "γ": "ガンマ",
+    "δ": "デルタ",
+    "ε": "イプシロン",
+    "ζ": "ゼータ",
+    "η": "イータ",
+    "θ": "シータ",
+    "ι": "イオタ",
+    "κ": "カッパ",
+    "λ": "ラムダ",
+    "μ": "ミュー",
+    "ν": "ニュー",
+    "ξ": "クサイ",
+    "ο": "オミクロン",
+    "π": "パイ",
+    "ρ": "ロー",
+    "σ": "シグマ",
+    "τ": "タウ",
+    "υ": "ウプシロン",
+    "φ": "ファイ",
+    "χ": "カイ",
+    "ψ": "プサイ",
+    "ω": "オメガ",
+}
+
+
+_NUMBER_WITH_SEPARATOR_RX = re.compile("[0-9]{1,3}(,[0-9]{3})+")
+_CURRENCY_MAP = {"$": "ドル", "¥": "円", "£": "ポンド", "€": "ユーロ"}
+_CURRENCY_RX = re.compile(r"([$¥£€])([0-9.]*[0-9])")
+_NUMBER_RX = re.compile(r"[0-9]+(\.[0-9]+)?")
+
+
+def japanese_convert_numbers_to_words(text: str) -> str:
+    res = _NUMBER_WITH_SEPARATOR_RX.sub(lambda m: m[0].replace(",", ""), text)
+    res = _CURRENCY_RX.sub(lambda m: m[2] + _CURRENCY_MAP.get(m[1], m[1]), res)
+    res = _NUMBER_RX.sub(lambda m: num2words(m[0], lang="ja"), res)
+    return res
+
+
+def japanese_convert_alpha_symbols_to_words(text: str) -> str:
+    return "".join([_ALPHASYMBOL_YOMI.get(ch, ch) for ch in text.lower()])
+
+
+def japanese_text_to_phonemes(text: str) -> str:
+    """Convert Japanese text to phonemes."""
+    res = unicodedata.normalize("NFKC", text)
+    res = japanese_convert_numbers_to_words(res)
+    # res = japanese_convert_alpha_symbols_to_words(res)
+    res = text2kata(res)
+    res = kata2phoneme(res)
+    return res
+
+
+def is_japanese_character(char):
+    # 定义日语文字系统的 Unicode 范围
+    japanese_ranges = [
+        (0x3040, 0x309F),  # 平假名
+        (0x30A0, 0x30FF),  # 片假名
+        (0x4E00, 0x9FFF),  # 汉字 (CJK Unified Ideographs)
+        (0x3400, 0x4DBF),  # 汉字扩展 A
+        (0x20000, 0x2A6DF),  # 汉字扩展 B
+        # 可以根据需要添加其他汉字扩展范围
+    ]
+
+    # 将字符的 Unicode 编码转换为整数
+    char_code = ord(char)
+
+    # 检查字符是否在任何一个日语范围内
+    for start, end in japanese_ranges:
+        if start <= char_code <= end:
+            return True
+
+    return False
+
+
+rep_map = {
+    "：": ",",
+    "；": ",",
+    "，": ",",
+    "。": ".",
+    "！": "!",
+    "？": "?",
+    "\n": ".",
+    "·": ",",
+    "、": ",",
+    "...": "…",
+}
+
+
+def replace_punctuation(text):
+    pattern = re.compile("|".join(re.escape(p) for p in rep_map.keys()))
+
+    replaced_text = pattern.sub(lambda x: rep_map[x.group()], text)
+
+    replaced_text = re.sub(
+        r"[^\u3040-\u309F\u30A0-\u30FF\u4E00-\u9FFF\u3400-\u4DBF"
+        + "".join(punctuation)
+        + r"]+",
+        "",
+        replaced_text,
+    )
+
+    return replaced_text
 
-        if i < len(marks):
-            text += [marks[i].replace(' ', '')]
-    return text
 
 def text_normalize(text):
-    # todo: jap text normalize
-    return text
+    res = unicodedata.normalize("NFKC", text)
+    res = japanese_convert_numbers_to_words(res)
+    # res = "".join([i for i in res if is_japanese_character(i)])
+    res = replace_punctuation(res)
+    return res
+
+
+def distribute_phone(n_phone, n_word):
+    phones_per_word = [0] * n_word
+    for task in range(n_phone):
+        min_tasks = min(phones_per_word)
+        min_index = phones_per_word.index(min_tasks)
+        phones_per_word[min_index] += 1
+    return phones_per_word
+
+
+tokenizer = AutoTokenizer.from_pretrained("./bert/bert-base-japanese-v3")
+
 
 def g2p(norm_text):
-    phones = preprocess_jap(norm_text)
-    phones = [post_replace_ph(i) for i in phones]
-    # todo: implement tones and word2ph
+    tokenized = tokenizer.tokenize(norm_text)
+    phs = []
+    ph_groups = []
+    for t in tokenized:
+        if not t.startswith("#"):
+            ph_groups.append([t])
+        else:
+            ph_groups[-1].append(t.replace("#", ""))
+    word2ph = []
+    for group in ph_groups:
+        phonemes = kata2phoneme(text2kata("".join(group)))
+        # phonemes = [i for i in phonemes if i in symbols]
+        for i in phonemes:
+            assert i in symbols, (group, norm_text, tokenized)
+        phone_len = len(phonemes)
+        word_len = len(group)
+
+        aaa = distribute_phone(phone_len, word_len)
+        word2ph += aaa
+
+        phs += phonemes
+    phones = ["_"] + phs + ["_"]
     tones = [0 for i in phones]
-    word2ph = [1 for i in phones]
+    word2ph = [1] + word2ph + [1]
     return phones, tones, word2ph
 
 
-if __name__ == '__main__':
-    for line in open("../../../Downloads/transcript_utf8.txt").readlines():
-        text = line.split(":")[1]
-        phones, tones, word2ph = g2p(text)
-        for p in phones:
-            if p == "z":
-                print(text, phones)
-                sys.exit(0)
+if __name__ == "__main__":
+    tokenizer = AutoTokenizer.from_pretrained("./bert/bert-base-japanese-v3")
+    text = "hello,こんにちは、世界！……"
+    from text.japanese_bert import get_bert_feature
+
+    text = text_normalize(text)
+    print(text)
+    phones, tones, word2ph = g2p(text)
+    bert = get_bert_feature(text, word2ph)
+
+    print(phones, tones, word2ph, bert.shape)

text/japanese_bert.py

@@ -0,0 +1,35 @@
+import torch
+from transformers import AutoTokenizer, AutoModelForMaskedLM
+import sys
+
+tokenizer = AutoTokenizer.from_pretrained("./bert/bert-base-japanese-v3")
+
+
+def get_bert_feature(text, word2ph, device=None):
+    if (
+        sys.platform == "darwin"
+        and torch.backends.mps.is_available()
+        and device == "cpu"
+    ):
+        device = "mps"
+    if not device:
+        device = "cuda"
+    model = AutoModelForMaskedLM.from_pretrained("./bert/bert-base-japanese-v3").to(
+        device
+    )
+    with torch.no_grad():
+        inputs = tokenizer(text, return_tensors="pt")
+        for i in inputs:
+            inputs[i] = inputs[i].to(device)
+        res = model(**inputs, output_hidden_states=True)
+        res = torch.cat(res["hidden_states"][-3:-2], -1)[0].cpu()
+    assert inputs["input_ids"].shape[-1] == len(word2ph)
+    word2phone = word2ph
+    phone_level_feature = []
+    for i in range(len(word2phone)):
+        repeat_feature = res[i].repeat(word2phone[i], 1)
+        phone_level_feature.append(repeat_feature)
+
+    phone_level_feature = torch.cat(phone_level_feature, dim=0)
+
+    return phone_level_feature.T

text/symbols.py

@@ -1,25 +1,166 @@
-punctuation = ['!', '?', '…', ",", ".", "'", '-']
+punctuation = ["!", "?", "…", ",", ".", "'", "-"]
 pu_symbols = punctuation + ["SP", "UNK"]
-pad = '_'
+pad = "_"
 
 # chinese
-zh_symbols = ['E', 'En', 'a', 'ai', 'an', 'ang', 'ao', 'b', 'c', 'ch', 'd', 'e', 'ei', 'en', 'eng', 'er', 'f', 'g', 'h',
-       'i', 'i0', 'ia', 'ian', 'iang', 'iao', 'ie', 'in', 'ing', 'iong', 'ir', 'iu', 'j', 'k', 'l', 'm', 'n', 'o',
-       'ong',
-       'ou', 'p', 'q', 'r', 's', 'sh', 't', 'u', 'ua', 'uai', 'uan', 'uang', 'ui', 'un', 'uo', 'v', 'van', 've', 'vn',
-       'w', 'x', 'y', 'z', 'zh',
-        "AA", "EE", "OO"]
+zh_symbols = [
+    "E",
+    "En",
+    "a",
+    "ai",
+    "an",
+    "ang",
+    "ao",
+    "b",
+    "c",
+    "ch",
+    "d",
+    "e",
+    "ei",
+    "en",
+    "eng",
+    "er",
+    "f",
+    "g",
+    "h",
+    "i",
+    "i0",
+    "ia",
+    "ian",
+    "iang",
+    "iao",
+    "ie",
+    "in",
+    "ing",
+    "iong",
+    "ir",
+    "iu",
+    "j",
+    "k",
+    "l",
+    "m",
+    "n",
+    "o",
+    "ong",
+    "ou",
+    "p",
+    "q",
+    "r",
+    "s",
+    "sh",
+    "t",
+    "u",
+    "ua",
+    "uai",
+    "uan",
+    "uang",
+    "ui",
+    "un",
+    "uo",
+    "v",
+    "van",
+    "ve",
+    "vn",
+    "w",
+    "x",
+    "y",
+    "z",
+    "zh",
+    "AA",
+    "EE",
+    "OO",
+]
 num_zh_tones = 6
 
 # japanese
-ja_symbols = ['I', 'N', 'U', 'a', 'b', 'by', 'ch', 'cl', 'd', 'dy', 'e', 'f', 'g', 'gy', 'h', 'hy', 'i', 'j', 'k', 'ky',
-              'm', 'my', 'n', 'ny', 'o', 'p', 'py', 'r', 'ry', 's', 'sh', 't', 'ts', 'u', 'V', 'w', 'y', 'z']
+ja_symbols = [
+    "N",
+    "a",
+    "a:",
+    "b",
+    "by",
+    "ch",
+    "d",
+    "dy",
+    "e",
+    "e:",
+    "f",
+    "g",
+    "gy",
+    "h",
+    "hy",
+    "i",
+    "i:",
+    "j",
+    "k",
+    "ky",
+    "m",
+    "my",
+    "n",
+    "ny",
+    "o",
+    "o:",
+    "p",
+    "py",
+    "q",
+    "r",
+    "ry",
+    "s",
+    "sh",
+    "t",
+    "ts",
+    "ty",
+    "u",
+    "u:",
+    "w",
+    "y",
+    "z",
+    "zy",
+]
 num_ja_tones = 1
 
 # English
-en_symbols = ['aa', 'ae', 'ah', 'ao', 'aw', 'ay', 'b', 'ch', 'd', 'dh', 'eh', 'er', 'ey', 'f', 'g', 'hh', 'ih', 'iy',
-              'jh', 'k', 'l', 'm', 'n', 'ng', 'ow', 'oy', 'p', 'r', 's',
-              'sh', 't', 'th', 'uh', 'uw', 'V', 'w', 'y', 'z', 'zh']
+en_symbols = [
+    "aa",
+    "ae",
+    "ah",
+    "ao",
+    "aw",
+    "ay",
+    "b",
+    "ch",
+    "d",
+    "dh",
+    "eh",
+    "er",
+    "ey",
+    "f",
+    "g",
+    "hh",
+    "ih",
+    "iy",
+    "jh",
+    "k",
+    "l",
+    "m",
+    "n",
+    "ng",
+    "ow",
+    "oy",
+    "p",
+    "r",
+    "s",
+    "sh",
+    "t",
+    "th",
+    "uh",
+    "uw",
+    "V",
+    "w",
+    "y",
+    "z",
+    "zh",
+]
 num_en_tones = 4
 
 # combine all symbols
@@ -31,21 +172,16 @@ sil_phonemes_ids = [symbols.index(i) for i in pu_symbols]
 num_tones = num_zh_tones + num_ja_tones + num_en_tones
 
 # language maps
-language_id_map = {
-    'ZH': 0,
-    "JA": 1,
-    "EN": 2
-}
+language_id_map = {"ZH": 0, "JP": 1, "EN": 2}
 num_languages = len(language_id_map.keys())
 
 language_tone_start_map = {
-    'ZH': 0,
-    "JA": num_zh_tones,
-    "EN": num_zh_tones + num_ja_tones
+    "ZH": 0,
+    "JP": num_zh_tones,
+    "EN": num_zh_tones + num_ja_tones,
 }
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     a = set(zh_symbols)
     b = set(en_symbols)
-    print(sorted(a&b))
-
+    print(sorted(a & b))

text/tone_sandhi.py

@@ -19,51 +19,442 @@ from pypinyin import lazy_pinyin
 from pypinyin import Style
 
 
-class ToneSandhi():
+class ToneSandhi:
     def __init__(self):
         self.must_neural_tone_words = {
-            '麻烦', '麻利', '鸳鸯', '高粱', '骨头', '骆驼', '马虎', '首饰', '馒头', '馄饨', '风筝',
-            '难为', '队伍', '阔气', '闺女', '门道', '锄头', '铺盖', '铃铛', '铁匠', '钥匙', '里脊',
-            '里头', '部分', '那么', '道士', '造化', '迷糊', '连累', '这么', '这个', '运气', '过去',
-            '软和', '转悠', '踏实', '跳蚤', '跟头', '趔趄', '财主', '豆腐', '讲究', '记性', '记号',
-            '认识', '规矩', '见识', '裁缝', '补丁', '衣裳', '衣服', '衙门', '街坊', '行李', '行当',
-            '蛤蟆', '蘑菇', '薄荷', '葫芦', '葡萄', '萝卜', '荸荠', '苗条', '苗头', '苍蝇', '芝麻',
-            '舒服', '舒坦', '舌头', '自在', '膏药', '脾气', '脑袋', '脊梁', '能耐', '胳膊', '胭脂',
-            '胡萝', '胡琴', '胡同', '聪明', '耽误', '耽搁', '耷拉', '耳朵', '老爷', '老实', '老婆',
-            '老头', '老太', '翻腾', '罗嗦', '罐头', '编辑', '结实', '红火', '累赘', '糨糊', '糊涂',
-            '精神', '粮食', '簸箕', '篱笆', '算计', '算盘', '答应', '笤帚', '笑语', '笑话', '窟窿',
-            '窝囊', '窗户', '稳当', '稀罕', '称呼', '秧歌', '秀气', '秀才', '福气', '祖宗', '砚台',
-            '码头', '石榴', '石头', '石匠', '知识', '眼睛', '眯缝', '眨巴', '眉毛', '相声', '盘算',
-            '白净', '痢疾', '痛快', '疟疾', '疙瘩', '疏忽', '畜生', '生意', '甘蔗', '琵琶', '琢磨',
-            '琉璃', '玻璃', '玫瑰', '玄乎', '狐狸', '状元', '特务', '牲口', '牙碜', '牌楼', '爽快',
-            '爱人', '热闹', '烧饼', '烟筒', '烂糊', '点心', '炊帚', '灯笼', '火候', '漂亮', '滑溜',
-            '溜达', '温和', '清楚', '消息', '浪头', '活泼', '比方', '正经', '欺负', '模糊', '槟榔',
-            '棺材', '棒槌', '棉花', '核桃', '栅栏', '柴火', '架势', '枕头', '枇杷', '机灵', '本事',
-            '木头', '木匠', '朋友', '月饼', '月亮', '暖和', '明白', '时候', '新鲜', '故事', '收拾',
-            '收成', '提防', '挖苦', '挑剔', '指甲', '指头', '拾掇', '拳头', '拨弄', '招牌', '招呼',
-            '抬举', '护士', '折腾', '扫帚', '打量', '打算', '打点', '打扮', '打听', '打发', '扎实',
-            '扁担', '戒指', '懒得', '意识', '意思', '情形', '悟性', '怪物', '思量', '怎么', '念头',
-            '念叨', '快活', '忙活', '志气', '心思', '得罪', '张罗', '弟兄', '开通', '应酬', '庄稼',
-            '干事', '帮手', '帐篷', '希罕', '师父', '师傅', '巴结', '巴掌', '差事', '工夫', '岁数',
-            '屁股', '尾巴', '少爷', '小气', '小伙', '将就', '对头', '对付', '寡妇', '家伙', '客气',
-            '实在', '官司', '学问', '学生', '字号', '嫁妆', '媳妇', '媒人', '婆家', '娘家', '委屈',
-            '姑娘', '姐夫', '妯娌', '妥当', '妖精', '奴才', '女婿', '头发', '太阳', '大爷', '大方',
-            '大意', '大夫', '多少', '多么', '外甥', '壮实', '地道', '地方', '在乎', '困难', '嘴巴',
-            '嘱咐', '嘟囔', '嘀咕', '喜欢', '喇嘛', '喇叭', '商量', '唾沫', '哑巴', '哈欠', '哆嗦',
-            '咳嗽', '和尚', '告诉', '告示', '含糊', '吓唬', '后头', '名字', '名堂', '合同', '吆喝',
-            '叫唤', '口袋', '厚道', '厉害', '千斤', '包袱', '包涵', '匀称', '勤快', '动静', '动弹',
-            '功夫', '力气', '前头', '刺猬', '刺激', '别扭', '利落', '利索', '利害', '分析', '出息',
-            '凑合', '凉快', '冷战', '冤枉', '冒失', '养活', '关系', '先生', '兄弟', '便宜', '使唤',
-            '佩服', '作坊', '体面', '位置', '似的', '伙计', '休息', '什么', '人家', '亲戚', '亲家',
-            '交情', '云彩', '事情', '买卖', '主意', '丫头', '丧气', '两口', '东西', '东家', '世故',
-            '不由', '不在', '下水', '下巴', '上头', '上司', '丈夫', '丈人', '一辈', '那个', '菩萨',
-            '父亲', '母亲', '咕噜', '邋遢', '费用', '冤家', '甜头', '介绍', '荒唐', '大人', '泥鳅',
-            '幸福', '熟悉', '计划', '扑腾', '蜡烛', '姥爷', '照顾', '喉咙', '吉他', '弄堂', '蚂蚱',
-            '凤凰', '拖沓', '寒碜', '糟蹋', '倒腾', '报复', '逻辑', '盘缠', '喽啰', '牢骚', '咖喱',
-            '扫把', '惦记'
+            "麻烦",
+            "麻利",
+            "鸳鸯",
+            "高粱",
+            "骨头",
+            "骆驼",
+            "马虎",
+            "首饰",
+            "馒头",
+            "馄饨",
+            "风筝",
+            "难为",
+            "队伍",
+            "阔气",
+            "闺女",
+            "门道",
+            "锄头",
+            "铺盖",
+            "铃铛",
+            "铁匠",
+            "钥匙",
+            "里脊",
+            "里头",
+            "部分",
+            "那么",
+            "道士",
+            "造化",
+            "迷糊",
+            "连累",
+            "这么",
+            "这个",
+            "运气",
+            "过去",
+            "软和",
+            "转悠",
+            "踏实",
+            "跳蚤",
+            "跟头",
+            "趔趄",
+            "财主",
+            "豆腐",
+            "讲究",
+            "记性",
+            "记号",
+            "认识",
+            "规矩",
+            "见识",
+            "裁缝",
+            "补丁",
+            "衣裳",
+            "衣服",
+            "衙门",
+            "街坊",
+            "行李",
+            "行当",
+            "蛤蟆",
+            "蘑菇",
+            "薄荷",
+            "葫芦",
+            "葡萄",
+            "萝卜",
+            "荸荠",
+            "苗条",
+            "苗头",
+            "苍蝇",
+            "芝麻",
+            "舒服",
+            "舒坦",
+            "舌头",
+            "自在",
+            "膏药",
+            "脾气",
+            "脑袋",
+            "脊梁",
+            "能耐",
+            "胳膊",
+            "胭脂",
+            "胡萝",
+            "胡琴",
+            "胡同",
+            "聪明",
+            "耽误",
+            "耽搁",
+            "耷拉",
+            "耳朵",
+            "老爷",
+            "老实",
+            "老婆",
+            "老头",
+            "老太",
+            "翻腾",
+            "罗嗦",
+            "罐头",
+            "编辑",
+            "结实",
+            "红火",
+            "累赘",
+            "糨糊",
+            "糊涂",
+            "精神",
+            "粮食",
+            "簸箕",
+            "篱笆",
+            "算计",
+            "算盘",
+            "答应",
+            "笤帚",
+            "笑语",
+            "笑话",
+            "窟窿",
+            "窝囊",
+            "窗户",
+            "稳当",
+            "稀罕",
+            "称呼",
+            "秧歌",
+            "秀气",
+            "秀才",
+            "福气",
+            "祖宗",
+            "砚台",
+            "码头",
+            "石榴",
+            "石头",
+            "石匠",
+            "知识",
+            "眼睛",
+            "眯缝",
+            "眨巴",
+            "眉毛",
+            "相声",
+            "盘算",
+            "白净",
+            "痢疾",
+            "痛快",
+            "疟疾",
+            "疙瘩",
+            "疏忽",
+            "畜生",
+            "生意",
+            "甘蔗",
+            "琵琶",
+            "琢磨",
+            "琉璃",
+            "玻璃",
+            "玫瑰",
+            "玄乎",
+            "狐狸",
+            "状元",
+            "特务",
+            "牲口",
+            "牙碜",
+            "牌楼",
+            "爽快",
+            "爱人",
+            "热闹",
+            "烧饼",
+            "烟筒",
+            "烂糊",
+            "点心",
+            "炊帚",
+            "灯笼",
+            "火候",
+            "漂亮",
+            "滑溜",
+            "溜达",
+            "温和",
+            "清楚",
+            "消息",
+            "浪头",
+            "活泼",
+            "比方",
+            "正经",
+            "欺负",
+            "模糊",
+            "槟榔",
+            "棺材",
+            "棒槌",
+            "棉花",
+            "核桃",
+            "栅栏",
+            "柴火",
+            "架势",
+            "枕头",
+            "枇杷",
+            "机灵",
+            "本事",
+            "木头",
+            "木匠",
+            "朋友",
+            "月饼",
+            "月亮",
+            "暖和",
+            "明白",
+            "时候",
+            "新鲜",
+            "故事",
+            "收拾",
+            "收成",
+            "提防",
+            "挖苦",
+            "挑剔",
+            "指甲",
+            "指头",
+            "拾掇",
+            "拳头",
+            "拨弄",
+            "招牌",
+            "招呼",
+            "抬举",
+            "护士",
+            "折腾",
+            "扫帚",
+            "打量",
+            "打算",
+            "打点",
+            "打扮",
+            "打听",
+            "打发",
+            "扎实",
+            "扁担",
+            "戒指",
+            "懒得",
+            "意识",
+            "意思",
+            "情形",
+            "悟性",
+            "怪物",
+            "思量",
+            "怎么",
+            "念头",
+            "念叨",
+            "快活",
+            "忙活",
+            "志气",
+            "心思",
+            "得罪",
+            "张罗",
+            "弟兄",
+            "开通",
+            "应酬",
+            "庄稼",
+            "干事",
+            "帮手",
+            "帐篷",
+            "希罕",
+            "师父",
+            "师傅",
+            "巴结",
+            "巴掌",
+            "差事",
+            "工夫",
+            "岁数",
+            "屁股",
+            "尾巴",
+            "少爷",
+            "小气",
+            "小伙",
+            "将就",
+            "对头",
+            "对付",
+            "寡妇",
+            "家伙",
+            "客气",
+            "实在",
+            "官司",
+            "学问",
+            "学生",
+            "字号",
+            "嫁妆",
+            "媳妇",
+            "媒人",
+            "婆家",
+            "娘家",
+            "委屈",
+            "姑娘",
+            "姐夫",
+            "妯娌",
+            "妥当",
+            "妖精",
+            "奴才",
+            "女婿",
+            "头发",
+            "太阳",
+            "大爷",
+            "大方",
+            "大意",
+            "大夫",
+            "多少",
+            "多么",
+            "外甥",
+            "壮实",
+            "地道",
+            "地方",
+            "在乎",
+            "困难",
+            "嘴巴",
+            "嘱咐",
+            "嘟囔",
+            "嘀咕",
+            "喜欢",
+            "喇嘛",
+            "喇叭",
+            "商量",
+            "唾沫",
+            "哑巴",
+            "哈欠",
+            "哆嗦",
+            "咳嗽",
+            "和尚",
+            "告诉",
+            "告示",
+            "含糊",
+            "吓唬",
+            "后头",
+            "名字",
+            "名堂",
+            "合同",
+            "吆喝",
+            "叫唤",
+            "口袋",
+            "厚道",
+            "厉害",
+            "千斤",
+            "包袱",
+            "包涵",
+            "匀称",
+            "勤快",
+            "动静",
+            "动弹",
+            "功夫",
+            "力气",
+            "前头",
+            "刺猬",
+            "刺激",
+            "别扭",
+            "利落",
+            "利索",
+            "利害",
+            "分析",
+            "出息",
+            "凑合",
+            "凉快",
+            "冷战",
+            "冤枉",
+            "冒失",
+            "养活",
+            "关系",
+            "先生",
+            "兄弟",
+            "便宜",
+            "使唤",
+            "佩服",
+            "作坊",
+            "体面",
+            "位置",
+            "似的",
+            "伙计",
+            "休息",
+            "什么",
+            "人家",
+            "亲戚",
+            "亲家",
+            "交情",
+            "云彩",
+            "事情",
+            "买卖",
+            "主意",
+            "丫头",
+            "丧气",
+            "两口",
+            "东西",
+            "东家",
+            "世故",
+            "不由",
+            "不在",
+            "下水",
+            "下巴",
+            "上头",
+            "上司",
+            "丈夫",
+            "丈人",
+            "一辈",
+            "那个",
+            "菩萨",
+            "父亲",
+            "母亲",
+            "咕噜",
+            "邋遢",
+            "费用",
+            "冤家",
+            "甜头",
+            "介绍",
+            "荒唐",
+            "大人",
+            "泥鳅",
+            "幸福",
+            "熟悉",
+            "计划",
+            "扑腾",
+            "蜡烛",
+            "姥爷",
+            "照顾",
+            "喉咙",
+            "吉他",
+            "弄堂",
+            "蚂蚱",
+            "凤凰",
+            "拖沓",
+            "寒碜",
+            "糟蹋",
+            "倒腾",
+            "报复",
+            "逻辑",
+            "盘缠",
+            "喽啰",
+            "牢骚",
+            "咖喱",
+            "扫把",
+            "惦记",
         }
         self.must_not_neural_tone_words = {
-            "男子", "女子", "分子", "原子", "量子", "莲子", "石子", "瓜子", "电子", "人人", "虎虎"
+            "男子",
+            "女子",
+            "分子",
+            "原子",
+            "量子",
+            "莲子",
+            "石子",
+            "瓜子",
+            "电子",
+            "人人",
+            "虎虎",
         }
         self.punc = "：，；。？！“”‘’':,;.?!"
 
@@ -72,14 +463,15 @@ class ToneSandhi():
     # word: "家里"
     # pos: "s"
     # finals: ['ia1', 'i3']
-    def _neural_sandhi(self, word: str, pos: str,
-                       finals: List[str]) -> List[str]:
-
+    def _neural_sandhi(self, word: str, pos: str, finals: List[str]) -> List[str]:
         # reduplication words for n. and v. e.g. 奶奶, 试试, 旺旺
         for j, item in enumerate(word):
-            if j - 1 >= 0 and item == word[j - 1] and pos[0] in {
-                    "n", "v", "a"
-            } and word not in self.must_not_neural_tone_words:
+            if (
+                j - 1 >= 0
+                and item == word[j - 1]
+                and pos[0] in {"n", "v", "a"}
+                and word not in self.must_not_neural_tone_words
+            ):
                 finals[j] = finals[j][:-1] + "5"
         ge_idx = word.find("个")
         if len(word) >= 1 and word[-1] in "吧呢啊呐噻嘛吖嗨呐哦哒额滴哩哟喽啰耶喔诶":
@@ -89,9 +481,12 @@ class ToneSandhi():
         # e.g. 走了, 看着, 去过
         # elif len(word) == 1 and word in "了着过" and pos in {"ul", "uz", "ug"}:
         #     finals[-1] = finals[-1][:-1] + "5"
-        elif len(word) > 1 and word[-1] in "们子" and pos in {
-                "r", "n"
-        } and word not in self.must_not_neural_tone_words:
+        elif (
+            len(word) > 1
+            and word[-1] in "们子"
+            and pos in {"r", "n"}
+            and word not in self.must_not_neural_tone_words
+        ):
             finals[-1] = finals[-1][:-1] + "5"
         # e.g. 桌上, 地下, 家里
         elif len(word) > 1 and word[-1] in "上下里" and pos in {"s", "l", "f"}:
@@ -100,21 +495,26 @@ class ToneSandhi():
         elif len(word) > 1 and word[-1] in "来去" and word[-2] in "上下进出回过起开":
             finals[-1] = finals[-1][:-1] + "5"
         # 个做量词
-        elif (ge_idx >= 1 and
-              (word[ge_idx - 1].isnumeric() or
-               word[ge_idx - 1] in "几有两半多各整每做是")) or word == '个':
+        elif (
+            ge_idx >= 1
+            and (word[ge_idx - 1].isnumeric() or word[ge_idx - 1] in "几有两半多各整每做是")
+        ) or word == "个":
             finals[ge_idx] = finals[ge_idx][:-1] + "5"
         else:
-            if word in self.must_neural_tone_words or word[
-                    -2:] in self.must_neural_tone_words:
+            if (
+                word in self.must_neural_tone_words
+                or word[-2:] in self.must_neural_tone_words
+            ):
                 finals[-1] = finals[-1][:-1] + "5"
 
         word_list = self._split_word(word)
-        finals_list = [finals[:len(word_list[0])], finals[len(word_list[0]):]]
+        finals_list = [finals[: len(word_list[0])], finals[len(word_list[0]) :]]
         for i, word in enumerate(word_list):
             # conventional neural in Chinese
-            if word in self.must_neural_tone_words or word[
-                    -2:] in self.must_neural_tone_words:
+            if (
+                word in self.must_neural_tone_words
+                or word[-2:] in self.must_neural_tone_words
+            ):
                 finals_list[i][-1] = finals_list[i][-1][:-1] + "5"
         finals = sum(finals_list, [])
         return finals
@@ -126,17 +526,17 @@ class ToneSandhi():
         else:
             for i, char in enumerate(word):
                 # "不" before tone4 should be bu2, e.g. 不怕
-                if char == "不" and i + 1 < len(word) and finals[i +
-                                                                1][-1] == "4":
+                if char == "不" and i + 1 < len(word) and finals[i + 1][-1] == "4":
                     finals[i] = finals[i][:-1] + "2"
         return finals
 
     def _yi_sandhi(self, word: str, finals: List[str]) -> List[str]:
         # "一" in number sequences, e.g. 一零零, 二一零
         if word.find("一") != -1 and all(
-            [item.isnumeric() for item in word if item != "一"]):
+            [item.isnumeric() for item in word if item != "一"]
+        ):
             return finals
-        # "一" between reduplication words shold be yi5, e.g. 看一看
+        # "一" between reduplication words should be yi5, e.g. 看一看
         elif len(word) == 3 and word[1] == "一" and word[0] == word[-1]:
             finals[1] = finals[1][:-1] + "5"
         # when "一" is ordinal word, it should be yi1
@@ -161,10 +561,10 @@ class ToneSandhi():
         first_subword = word_list[0]
         first_begin_idx = word.find(first_subword)
         if first_begin_idx == 0:
-            second_subword = word[len(first_subword):]
+            second_subword = word[len(first_subword) :]
             new_word_list = [first_subword, second_subword]
         else:
-            second_subword = word[:-len(first_subword)]
+            second_subword = word[: -len(first_subword)]
             new_word_list = [second_subword, first_subword]
         return new_word_list
 
@@ -182,18 +582,19 @@ class ToneSandhi():
                 elif len(word_list[0]) == 1:
                     finals[1] = finals[1][:-1] + "2"
             else:
-                finals_list = [
-                    finals[:len(word_list[0])], finals[len(word_list[0]):]
-                ]
+                finals_list = [finals[: len(word_list[0])], finals[len(word_list[0]) :]]
                 if len(finals_list) == 2:
                     for i, sub in enumerate(finals_list):
                         # e.g. 所有/人
                         if self._all_tone_three(sub) and len(sub) == 2:
                             finals_list[i][0] = finals_list[i][0][:-1] + "2"
                         # e.g. 好/喜欢
-                        elif i == 1 and not self._all_tone_three(sub) and finals_list[i][0][-1] == "3" and \
-                                finals_list[0][-1][-1] == "3":
-
+                        elif (
+                            i == 1
+                            and not self._all_tone_three(sub)
+                            and finals_list[i][0][-1] == "3"
+                            and finals_list[0][-1][-1] == "3"
+                        ):
                             finals_list[0][-1] = finals_list[0][-1][:-1] + "2"
                         finals = sum(finals_list, [])
         # split idiom into two words who's length is 2
@@ -222,7 +623,7 @@ class ToneSandhi():
                 new_seg.append((word, pos))
             last_word = word[:]
         if last_word == "不":
-            new_seg.append((last_word, 'd'))
+            new_seg.append((last_word, "d"))
             last_word = ""
         return new_seg
 
@@ -236,12 +637,21 @@ class ToneSandhi():
         new_seg = []
         # function 1
         for i, (word, pos) in enumerate(seg):
-            if i - 1 >= 0 and word == "一" and i + 1 < len(seg) and seg[i - 1][
-                    0] == seg[i + 1][0] and seg[i - 1][1] == "v":
+            if (
+                i - 1 >= 0
+                and word == "一"
+                and i + 1 < len(seg)
+                and seg[i - 1][0] == seg[i + 1][0]
+                and seg[i - 1][1] == "v"
+            ):
                 new_seg[i - 1][0] = new_seg[i - 1][0] + "一" + new_seg[i - 1][0]
             else:
-                if i - 2 >= 0 and seg[i - 1][0] == "一" and seg[i - 2][
-                        0] == word and pos == "v":
+                if (
+                    i - 2 >= 0
+                    and seg[i - 1][0] == "一"
+                    and seg[i - 2][0] == word
+                    and pos == "v"
+                ):
                     continue
                 else:
                     new_seg.append([word, pos])
@@ -257,22 +667,27 @@ class ToneSandhi():
 
     # the first and the second words are all_tone_three
     def _merge_continuous_three_tones(
-            self, seg: List[Tuple[str, str]]) -> List[Tuple[str, str]]:
+        self, seg: List[Tuple[str, str]]
+    ) -> List[Tuple[str, str]]:
         new_seg = []
         sub_finals_list = [
-            lazy_pinyin(
-                word, neutral_tone_with_five=True, style=Style.FINALS_TONE3)
+            lazy_pinyin(word, neutral_tone_with_five=True, style=Style.FINALS_TONE3)
             for (word, pos) in seg
         ]
         assert len(sub_finals_list) == len(seg)
         merge_last = [False] * len(seg)
         for i, (word, pos) in enumerate(seg):
-            if i - 1 >= 0 and self._all_tone_three(
-                    sub_finals_list[i - 1]) and self._all_tone_three(
-                        sub_finals_list[i]) and not merge_last[i - 1]:
+            if (
+                i - 1 >= 0
+                and self._all_tone_three(sub_finals_list[i - 1])
+                and self._all_tone_three(sub_finals_list[i])
+                and not merge_last[i - 1]
+            ):
                 # if the last word is reduplication, not merge, because reduplication need to be _neural_sandhi
-                if not self._is_reduplication(seg[i - 1][0]) and len(
-                        seg[i - 1][0]) + len(seg[i][0]) <= 3:
+                if (
+                    not self._is_reduplication(seg[i - 1][0])
+                    and len(seg[i - 1][0]) + len(seg[i][0]) <= 3
+                ):
                     new_seg[-1][0] = new_seg[-1][0] + seg[i][0]
                     merge_last[i] = True
                 else:
@@ -287,21 +702,27 @@ class ToneSandhi():
 
     # the last char of first word and the first char of second word is tone_three
     def _merge_continuous_three_tones_2(
-            self, seg: List[Tuple[str, str]]) -> List[Tuple[str, str]]:
+        self, seg: List[Tuple[str, str]]
+    ) -> List[Tuple[str, str]]:
         new_seg = []
         sub_finals_list = [
-            lazy_pinyin(
-                word, neutral_tone_with_five=True, style=Style.FINALS_TONE3)
+            lazy_pinyin(word, neutral_tone_with_five=True, style=Style.FINALS_TONE3)
             for (word, pos) in seg
         ]
         assert len(sub_finals_list) == len(seg)
         merge_last = [False] * len(seg)
         for i, (word, pos) in enumerate(seg):
-            if i - 1 >= 0 and sub_finals_list[i - 1][-1][-1] == "3" and sub_finals_list[i][0][-1] == "3" and not \
-                    merge_last[i - 1]:
+            if (
+                i - 1 >= 0
+                and sub_finals_list[i - 1][-1][-1] == "3"
+                and sub_finals_list[i][0][-1] == "3"
+                and not merge_last[i - 1]
+            ):
                 # if the last word is reduplication, not merge, because reduplication need to be _neural_sandhi
-                if not self._is_reduplication(seg[i - 1][0]) and len(
-                        seg[i - 1][0]) + len(seg[i][0]) <= 3:
+                if (
+                    not self._is_reduplication(seg[i - 1][0])
+                    and len(seg[i - 1][0]) + len(seg[i][0]) <= 3
+                ):
                     new_seg[-1][0] = new_seg[-1][0] + seg[i][0]
                     merge_last[i] = True
                 else:
@@ -313,14 +734,13 @@ class ToneSandhi():
     def _merge_er(self, seg: List[Tuple[str, str]]) -> List[Tuple[str, str]]:
         new_seg = []
         for i, (word, pos) in enumerate(seg):
-            if i - 1 >= 0 and word == "儿" and seg[i-1][0] != "#":
+            if i - 1 >= 0 and word == "儿" and seg[i - 1][0] != "#":
                 new_seg[-1][0] = new_seg[-1][0] + seg[i][0]
             else:
                 new_seg.append([word, pos])
         return new_seg
 
-    def _merge_reduplication(
-            self, seg: List[Tuple[str, str]]) -> List[Tuple[str, str]]:
+    def _merge_reduplication(self, seg: List[Tuple[str, str]]) -> List[Tuple[str, str]]:
         new_seg = []
         for i, (word, pos) in enumerate(seg):
             if new_seg and word == new_seg[-1][0]:
@@ -329,8 +749,7 @@ class ToneSandhi():
                 new_seg.append([word, pos])
         return new_seg
 
-    def pre_merge_for_modify(
-            self, seg: List[Tuple[str, str]]) -> List[Tuple[str, str]]:
+    def pre_merge_for_modify(self, seg: List[Tuple[str, str]]) -> List[Tuple[str, str]]:
         seg = self._merge_bu(seg)
         try:
             seg = self._merge_yi(seg)
@@ -342,8 +761,7 @@ class ToneSandhi():
         seg = self._merge_er(seg)
         return seg
 
-    def modified_tone(self, word: str, pos: str,
-                      finals: List[str]) -> List[str]:
+    def modified_tone(self, word: str, pos: str, finals: List[str]) -> List[str]:
         finals = self._bu_sandhi(word, finals)
         finals = self._yi_sandhi(word, finals)
         finals = self._neural_sandhi(word, pos, finals)

train_ms.py

@@ -1,65 +1,58 @@
+# flake8: noqa: E402
+
 import os
-import json
-import argparse
-import itertools
-import math
 import torch
-from torch import nn, optim
 from torch.nn import functional as F
 from torch.utils.data import DataLoader
 from torch.utils.tensorboard import SummaryWriter
-import torch.multiprocessing as mp
 import torch.distributed as dist
 from torch.nn.parallel import DistributedDataParallel as DDP
 from torch.cuda.amp import autocast, GradScaler
 from tqdm import tqdm
 import logging
-logging.getLogger('numba').setLevel(logging.WARNING)
+
+logging.getLogger("numba").setLevel(logging.WARNING)
 import commons
 import utils
 from data_utils import (
     TextAudioSpeakerLoader,
     TextAudioSpeakerCollate,
-    DistributedBucketSampler
+    DistributedBucketSampler,
 )
 from models import (
     SynthesizerTrn,
     MultiPeriodDiscriminator,
     DurationDiscriminator,
 )
-from losses import (
-    generator_loss,
-    discriminator_loss,
-    feature_loss,
-    kl_loss
-)
+from losses import generator_loss, discriminator_loss, feature_loss, kl_loss
 from mel_processing import mel_spectrogram_torch, spec_to_mel_torch
 from text.symbols import symbols
 
-torch.backends.cudnn.benchmark = True
 torch.backends.cuda.matmul.allow_tf32 = True
-torch.backends.cudnn.allow_tf32 = True  # If encontered training problem,please try to disable TF32.
-torch.set_float32_matmul_precision('medium')
+torch.backends.cudnn.allow_tf32 = (
+    True  # If encontered training problem,please try to disable TF32.
+)
+torch.set_float32_matmul_precision("medium")
+torch.backends.cudnn.benchmark = True
 torch.backends.cuda.sdp_kernel("flash")
 torch.backends.cuda.enable_flash_sdp(True)
-torch.backends.cuda.enable_mem_efficient_sdp(True)  # Not avaliable if torch version is lower than 2.0
+torch.backends.cuda.enable_mem_efficient_sdp(
+    True
+)  # Not available if torch version is lower than 2.0
 torch.backends.cuda.enable_math_sdp(True)
 global_step = 0
 
 
-def main():
-    """Assume Single Node Multi GPUs Training Only"""
-    assert torch.cuda.is_available(), "CPU training is not allowed."
-
-    n_gpus = torch.cuda.device_count()
-    os.environ['MASTER_ADDR'] = 'localhost'
-    os.environ['MASTER_PORT'] = '65280'
-
+def run():
+    dist.init_process_group(
+        backend="gloo",
+        init_method="env://",  # Due to some training problem,we proposed to use gloo instead of nccl.
+    )  # Use torchrun instead of mp.spawn
+    rank = dist.get_rank()
+    n_gpus = dist.get_world_size()
     hps = utils.get_hparams()
-    mp.spawn(run, nprocs=n_gpus, args=(n_gpus, hps,))
-
-
-def run(rank, n_gpus, hps):
+    torch.manual_seed(hps.train.seed)
+    torch.cuda.set_device(rank)
     global global_step
     if rank == 0:
         logger = utils.get_logger(hps.model_dir)
@@ -67,11 +60,6 @@ def run(rank, n_gpus, hps):
         utils.check_git_hash(hps.model_dir)
         writer = SummaryWriter(log_dir=hps.model_dir)
         writer_eval = SummaryWriter(log_dir=os.path.join(hps.model_dir, "eval"))
-
-    dist.init_process_group(backend='nccl', init_method='env://', world_size=n_gpus, rank=rank)
-    torch.manual_seed(hps.train.seed)
-    torch.cuda.set_device(rank)
-
     train_dataset = TextAudioSpeakerLoader(hps.data.training_files, hps.data)
     train_sampler = DistributedBucketSampler(
         train_dataset,
@@ -79,76 +67,94 @@ def run(rank, n_gpus, hps):
         [32, 300, 400, 500, 600, 700, 800, 900, 1000],
         num_replicas=n_gpus,
         rank=rank,
-        shuffle=True)
+        shuffle=True,
+    )
     collate_fn = TextAudioSpeakerCollate()
-    train_loader = DataLoader(train_dataset, num_workers=24, shuffle=False, pin_memory=True,
-                              collate_fn=collate_fn, batch_sampler=train_sampler,
-                              persistent_workers=True,prefetch_factor=4)  #256G Memory suitable loader.
+    train_loader = DataLoader(
+        train_dataset,
+        num_workers=16,
+        shuffle=False,
+        pin_memory=True,
+        collate_fn=collate_fn,
+        batch_sampler=train_sampler,
+        persistent_workers=True,
+        prefetch_factor=4,
+    )  # 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 == True:
+        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
+    ):
         print("Using noise scaled MAS for VITS2")
-        use_noise_scaled_mas = True
         mas_noise_scale_initial = 0.01
         noise_scale_delta = 2e-6
     else:
         print("Using normal MAS for VITS1")
-        use_noise_scaled_mas = False
         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 == True:
+    if (
+        "use_duration_discriminator" in hps.model.keys()
+        and hps.model.use_duration_discriminator is True
+    ):
         print("Using duration discriminator for VITS2")
-        use_duration_discriminator = True
         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(rank)
-    if "use_spk_conditioned_encoder" in hps.model.keys() and hps.model.use_spk_conditioned_encoder == True:
+            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(rank)
+    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")
-        use_spk_conditioned_encoder = True
+            raise ValueError(
+                "n_speakers must be > 0 when using spk conditioned encoder to train multi-speaker model"
+            )
     else:
         print("Using normal encoder for VITS1")
-        use_spk_conditioned_encoder = False
 
     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(rank)
-
-    freeze_enc = getattr(hps.model, "freeze_enc", False)
-    if freeze_enc:
-        print("freeze encoder !!!")
-        for param in net_g.enc_p.parameters():
-            param.requires_grad = False
+        mas_noise_scale_initial=mas_noise_scale_initial,
+        noise_scale_delta=noise_scale_delta,
+        **hps.model,
+    ).cuda(rank)
 
     net_d = MultiPeriodDiscriminator(hps.model.use_spectral_norm).cuda(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)
+        eps=hps.train.eps,
+    )
     optim_d = torch.optim.AdamW(
         net_d.parameters(),
         hps.train.learning_rate,
         betas=hps.train.betas,
-        eps=hps.train.eps)
+        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)
+            net_dur_disc.parameters(),
+            hps.train.learning_rate,
+            betas=hps.train.betas,
+            eps=hps.train.eps,
+        )
     else:
         optim_dur_disc = None
     net_g = DDP(net_g, device_ids=[rank], find_unused_parameters=True)
@@ -157,40 +163,94 @@ def run(rank, n_gpus, hps):
         net_dur_disc = DDP(net_dur_disc, device_ids=[rank], find_unused_parameters=True)
     try:
         if net_dur_disc is not None:
-            _, _, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(hps.model_dir, "DUR_*.pth"), net_dur_disc, optim_dur_disc, skip_optimizer=True)
-        _, optim_g, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(hps.model_dir, "G_*.pth"), net_g,
-                                                   optim_g, skip_optimizer=True)
-        _, optim_d, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(hps.model_dir, "D_*.pth"), net_d,
-                                                   optim_d, skip_optimizer=True)
-            
+            _, _, dur_resume_lr, epoch_str = utils.load_checkpoint(
+                utils.latest_checkpoint_path(hps.model_dir, "DUR_*.pth"),
+                net_dur_disc,
+                optim_dur_disc,
+                skip_optimizer=hps.train.skip_optimizer
+                if "skip_optimizer" in hps.train
+                else True,
+            )
+            _, optim_g, g_resume_lr, epoch_str = utils.load_checkpoint(
+                utils.latest_checkpoint_path(hps.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(hps.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)
     except Exception as e:
-            print(e)
-            epoch_str = 1
-            global_step = 0
+        print(e)
+        epoch_str = 1
+        global_step = 0
 
-
-    scheduler_g = torch.optim.lr_scheduler.ExponentialLR(optim_g, gamma=hps.train.lr_decay, last_epoch=epoch_str - 2)
-    scheduler_d = torch.optim.lr_scheduler.ExponentialLR(optim_d, gamma=hps.train.lr_decay, last_epoch=epoch_str - 2)
+    scheduler_g = torch.optim.lr_scheduler.ExponentialLR(
+        optim_g, gamma=hps.train.lr_decay, last_epoch=epoch_str - 2
+    )
+    scheduler_d = torch.optim.lr_scheduler.ExponentialLR(
+        optim_d, gamma=hps.train.lr_decay, last_epoch=epoch_str - 2
+    )
     if net_dur_disc is not None:
-        scheduler_dur_disc = torch.optim.lr_scheduler.ExponentialLR(optim_dur_disc, gamma=hps.train.lr_decay, last_epoch=epoch_str-2)
+        if not optim_dur_disc.param_groups[0].get("initial_lr"):
+            optim_dur_disc.param_groups[0]["initial_lr"] = dur_resume_lr
+        scheduler_dur_disc = torch.optim.lr_scheduler.ExponentialLR(
+            optim_dur_disc, gamma=hps.train.lr_decay, last_epoch=epoch_str - 2
+        )
     else:
         scheduler_dur_disc = None
     scaler = GradScaler(enabled=hps.train.fp16_run)
 
     for epoch in range(epoch_str, hps.train.epochs + 1):
         if rank == 0:
-            train_and_evaluate(rank, epoch, hps, [net_g, net_d, net_dur_disc], [optim_g, optim_d, optim_dur_disc], [scheduler_g, scheduler_d, scheduler_dur_disc], scaler, [train_loader, eval_loader], logger, [writer, writer_eval])
+            train_and_evaluate(
+                rank,
+                epoch,
+                hps,
+                [net_g, net_d, net_dur_disc],
+                [optim_g, optim_d, optim_dur_disc],
+                [scheduler_g, scheduler_d, scheduler_dur_disc],
+                scaler,
+                [train_loader, eval_loader],
+                logger,
+                [writer, writer_eval],
+            )
         else:
-            train_and_evaluate(rank, epoch, hps, [net_g, net_d, net_dur_disc], [optim_g, optim_d, optim_dur_disc], [scheduler_g, scheduler_d, scheduler_dur_disc], scaler, [train_loader, None], None, None)
+            train_and_evaluate(
+                rank,
+                epoch,
+                hps,
+                [net_g, net_d, net_dur_disc],
+                [optim_g, optim_d, optim_dur_disc],
+                [scheduler_g, scheduler_d, scheduler_dur_disc],
+                scaler,
+                [train_loader, None],
+                None,
+                None,
+            )
         scheduler_g.step()
         scheduler_d.step()
         if net_dur_disc is not None:
             scheduler_dur_disc.step()
 
 
-def train_and_evaluate(rank, epoch, hps, nets, optims, schedulers, scaler, loaders, logger, writers):
+def train_and_evaluate(
+    rank, epoch, hps, nets, optims, schedulers, scaler, loaders, logger, writers
+):
     net_g, net_d, net_dur_disc = nets
     optim_g, optim_d, optim_dur_disc = optims
     scheduler_g, scheduler_d, scheduler_dur_disc = schedulers
@@ -205,29 +265,72 @@ def train_and_evaluate(rank, epoch, hps, nets, optims, schedulers, scaler, loade
     net_d.train()
     if net_dur_disc is not None:
         net_dur_disc.train()
-    for batch_idx, (x, x_lengths, spec, spec_lengths, y, y_lengths, speakers, tone, language, bert) in tqdm(enumerate(train_loader)):
+    for batch_idx, (
+        x,
+        x_lengths,
+        spec,
+        spec_lengths,
+        y,
+        y_lengths,
+        speakers,
+        tone,
+        language,
+        bert,
+        ja_bert,
+    ) in tqdm(enumerate(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
+            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(rank, non_blocking=True), x_lengths.cuda(rank, non_blocking=True)
-        spec, spec_lengths = spec.cuda(rank, non_blocking=True), spec_lengths.cuda(rank, non_blocking=True)
-        y, y_lengths = y.cuda(rank, non_blocking=True), y_lengths.cuda(rank, non_blocking=True)
+        x, x_lengths = x.cuda(rank, non_blocking=True), x_lengths.cuda(
+            rank, non_blocking=True
+        )
+        spec, spec_lengths = spec.cuda(rank, non_blocking=True), spec_lengths.cuda(
+            rank, non_blocking=True
+        )
+        y, y_lengths = y.cuda(rank, non_blocking=True), y_lengths.cuda(
+            rank, non_blocking=True
+        )
         speakers = speakers.cuda(rank, non_blocking=True)
         tone = tone.cuda(rank, non_blocking=True)
         language = language.cuda(rank, non_blocking=True)
         bert = bert.cuda(rank, non_blocking=True)
+        ja_bert = ja_bert.cuda(rank, non_blocking=True)
 
         with autocast(enabled=hps.train.fp16_run):
-            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_) = net_g(x, x_lengths, spec, spec_lengths, speakers, tone, language, bert)
+            (
+                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_),
+            ) = net_g(
+                x,
+                x_lengths,
+                spec,
+                spec_lengths,
+                speakers,
+                tone,
+                language,
+                bert,
+                ja_bert,
+            )
             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)
+                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),
                 hps.data.filter_length,
@@ -236,26 +339,36 @@ def train_and_evaluate(rank, epoch, hps, nets, optims, schedulers, scaler, loade
                 hps.data.hop_length,
                 hps.data.win_length,
                 hps.data.mel_fmin,
-                hps.data.mel_fmax
+                hps.data.mel_fmax,
             )
 
-            y = commons.slice_segments(y, ids_slice * hps.data.hop_length, hps.train.segment_size)  # slice
+            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=False):
-                loss_disc, losses_disc_r, losses_disc_g = discriminator_loss(y_d_hat_r, y_d_hat_g)
+                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())
+                y_dur_hat_r, y_dur_hat_g = net_dur_disc(
+                    hidden_x.detach(), x_mask.detach(), logw.detach(), logw_.detach()
+                )
                 with autocast(enabled=False):
-                 # 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)
+                    # 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)
-                grad_norm_dur_disc = commons.clip_grad_value_(net_dur_disc.parameters(), None)
+                commons.clip_grad_value_(net_dur_disc.parameters(), None)
                 scaler.step(optim_dur_disc)
 
         optim_d.zero_grad()
@@ -289,51 +402,97 @@ def train_and_evaluate(rank, epoch, hps, nets, optims, schedulers, scaler, loade
 
         if rank == 0:
             if global_step % hps.train.log_interval == 0:
-                lr = optim_g.param_groups[0]['lr']
+                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. * batch_idx / len(train_loader)))
+                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 = {
+                    "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)})
-          
+                    {
+                        "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)}
+                )
+
                 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())
+                    "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)
+                    scalars=scalar_dict,
+                )
 
             if global_step % hps.train.eval_interval == 0:
                 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)))
+                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)))    
-                keep_ckpts = getattr(hps.train, 'keep_ckpts', 5)
+                    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)),
+                    )
+                keep_ckpts = getattr(hps.train, "keep_ckpts", 5)
                 if keep_ckpts > 0:
-                    utils.clean_checkpoints(path_to_models=hps.model_dir, n_ckpts_to_keep=keep_ckpts, sort_by_time=True)
-
+                    utils.clean_checkpoints(
+                        path_to_models=hps.model_dir,
+                        n_ckpts_to_keep=keep_ckpts,
+                        sort_by_time=True,
+                    )
 
         global_step += 1
 
     if rank == 0:
-        logger.info('====> Epoch: {}'.format(epoch))
-
+        logger.info("====> Epoch: {}".format(epoch))
 
 
 def evaluate(hps, generator, eval_loader, writer_eval):
@@ -342,16 +501,40 @@ def evaluate(hps, generator, eval_loader, writer_eval):
     audio_dict = {}
     print("Evaluating ...")
     with torch.no_grad():
-        for batch_idx, (x, x_lengths, spec, spec_lengths, y, y_lengths, speakers, tone, language, bert) in enumerate(eval_loader):
+        for batch_idx, (
+            x,
+            x_lengths,
+            spec,
+            spec_lengths,
+            y,
+            y_lengths,
+            speakers,
+            tone,
+            language,
+            bert,
+            ja_bert,
+        ) 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()
+            ja_bert = ja_bert.cuda()
             tone = tone.cuda()
             language = language.cuda()
             for use_sdp in [True, False]:
-                y_hat, attn, mask, *_ = generator.module.infer(x, x_lengths, speakers, tone, language, bert, y=spec, max_len=1000, sdp_ratio=0.0 if not use_sdp else 1.0)
+                y_hat, attn, mask, *_ = generator.module.infer(
+                    x,
+                    x_lengths,
+                    speakers,
+                    tone,
+                    language,
+                    bert,
+                    ja_bert,
+                    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(
@@ -360,7 +543,8 @@ def evaluate(hps, generator, eval_loader, writer_eval):
                     hps.data.n_mel_channels,
                     hps.data.sampling_rate,
                     hps.data.mel_fmin,
-                    hps.data.mel_fmax)
+                    hps.data.mel_fmax,
+                )
                 y_hat_mel = mel_spectrogram_torch(
                     y_hat.squeeze(1).float(),
                     hps.data.filter_length,
@@ -369,25 +553,40 @@ def evaluate(hps, generator, eval_loader, writer_eval):
                     hps.data.hop_length,
                     hps.data.win_length,
                     hps.data.mel_fmin,
-                    hps.data.mel_fmax
+                    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]]})
+                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
+        audio_sampling_rate=hps.data.sampling_rate,
     )
     generator.train()
 
+
 if __name__ == "__main__":
-    main()
+    run()

transforms.py

@@ -9,66 +9,63 @@ DEFAULT_MIN_BIN_HEIGHT = 1e-3
 DEFAULT_MIN_DERIVATIVE = 1e-3
 
 
-def piecewise_rational_quadratic_transform(inputs, 
-                                           unnormalized_widths,
-                                           unnormalized_heights,
-                                           unnormalized_derivatives,
-                                           inverse=False,
-                                           tails=None, 
-                                           tail_bound=1.,
-                                           min_bin_width=DEFAULT_MIN_BIN_WIDTH,
-                                           min_bin_height=DEFAULT_MIN_BIN_HEIGHT,
-                                           min_derivative=DEFAULT_MIN_DERIVATIVE):
-
+def piecewise_rational_quadratic_transform(
+    inputs,
+    unnormalized_widths,
+    unnormalized_heights,
+    unnormalized_derivatives,
+    inverse=False,
+    tails=None,
+    tail_bound=1.0,
+    min_bin_width=DEFAULT_MIN_BIN_WIDTH,
+    min_bin_height=DEFAULT_MIN_BIN_HEIGHT,
+    min_derivative=DEFAULT_MIN_DERIVATIVE,
+):
     if tails is None:
         spline_fn = rational_quadratic_spline
         spline_kwargs = {}
     else:
         spline_fn = unconstrained_rational_quadratic_spline
-        spline_kwargs = {
-            'tails': tails,
-            'tail_bound': tail_bound
-        }
+        spline_kwargs = {"tails": tails, "tail_bound": tail_bound}
 
     outputs, logabsdet = spline_fn(
-            inputs=inputs,
-            unnormalized_widths=unnormalized_widths,
-            unnormalized_heights=unnormalized_heights,
-            unnormalized_derivatives=unnormalized_derivatives,
-            inverse=inverse,
-            min_bin_width=min_bin_width,
-            min_bin_height=min_bin_height,
-            min_derivative=min_derivative,
-            **spline_kwargs
+        inputs=inputs,
+        unnormalized_widths=unnormalized_widths,
+        unnormalized_heights=unnormalized_heights,
+        unnormalized_derivatives=unnormalized_derivatives,
+        inverse=inverse,
+        min_bin_width=min_bin_width,
+        min_bin_height=min_bin_height,
+        min_derivative=min_derivative,
+        **spline_kwargs
     )
     return outputs, logabsdet
 
 
 def searchsorted(bin_locations, inputs, eps=1e-6):
     bin_locations[..., -1] += eps
-    return torch.sum(
-        inputs[..., None] >= bin_locations,
-        dim=-1
-    ) - 1
+    return torch.sum(inputs[..., None] >= bin_locations, dim=-1) - 1
 
 
-def unconstrained_rational_quadratic_spline(inputs,
-                                            unnormalized_widths,
-                                            unnormalized_heights,
-                                            unnormalized_derivatives,
-                                            inverse=False,
-                                            tails='linear',
-                                            tail_bound=1.,
-                                            min_bin_width=DEFAULT_MIN_BIN_WIDTH,
-                                            min_bin_height=DEFAULT_MIN_BIN_HEIGHT,
-                                            min_derivative=DEFAULT_MIN_DERIVATIVE):
+def unconstrained_rational_quadratic_spline(
+    inputs,
+    unnormalized_widths,
+    unnormalized_heights,
+    unnormalized_derivatives,
+    inverse=False,
+    tails="linear",
+    tail_bound=1.0,
+    min_bin_width=DEFAULT_MIN_BIN_WIDTH,
+    min_bin_height=DEFAULT_MIN_BIN_HEIGHT,
+    min_derivative=DEFAULT_MIN_DERIVATIVE,
+):
     inside_interval_mask = (inputs >= -tail_bound) & (inputs <= tail_bound)
     outside_interval_mask = ~inside_interval_mask
 
     outputs = torch.zeros_like(inputs)
     logabsdet = torch.zeros_like(inputs)
 
-    if tails == 'linear':
+    if tails == "linear":
         unnormalized_derivatives = F.pad(unnormalized_derivatives, pad=(1, 1))
         constant = np.log(np.exp(1 - min_derivative) - 1)
         unnormalized_derivatives[..., 0] = constant
@@ -77,45 +74,57 @@ def unconstrained_rational_quadratic_spline(inputs,
         outputs[outside_interval_mask] = inputs[outside_interval_mask]
         logabsdet[outside_interval_mask] = 0
     else:
-        raise RuntimeError('{} tails are not implemented.'.format(tails))
+        raise RuntimeError("{} tails are not implemented.".format(tails))
 
-    outputs[inside_interval_mask], logabsdet[inside_interval_mask] = rational_quadratic_spline(
+    (
+        outputs[inside_interval_mask],
+        logabsdet[inside_interval_mask],
+    ) = rational_quadratic_spline(
         inputs=inputs[inside_interval_mask],
         unnormalized_widths=unnormalized_widths[inside_interval_mask, :],
         unnormalized_heights=unnormalized_heights[inside_interval_mask, :],
         unnormalized_derivatives=unnormalized_derivatives[inside_interval_mask, :],
         inverse=inverse,
-        left=-tail_bound, right=tail_bound, bottom=-tail_bound, top=tail_bound,
+        left=-tail_bound,
+        right=tail_bound,
+        bottom=-tail_bound,
+        top=tail_bound,
         min_bin_width=min_bin_width,
         min_bin_height=min_bin_height,
-        min_derivative=min_derivative
+        min_derivative=min_derivative,
     )
 
     return outputs, logabsdet
 
-def rational_quadratic_spline(inputs,
-                              unnormalized_widths,
-                              unnormalized_heights,
-                              unnormalized_derivatives,
-                              inverse=False,
-                              left=0., right=1., bottom=0., top=1.,
-                              min_bin_width=DEFAULT_MIN_BIN_WIDTH,
-                              min_bin_height=DEFAULT_MIN_BIN_HEIGHT,
-                              min_derivative=DEFAULT_MIN_DERIVATIVE):
+
+def rational_quadratic_spline(
+    inputs,
+    unnormalized_widths,
+    unnormalized_heights,
+    unnormalized_derivatives,
+    inverse=False,
+    left=0.0,
+    right=1.0,
+    bottom=0.0,
+    top=1.0,
+    min_bin_width=DEFAULT_MIN_BIN_WIDTH,
+    min_bin_height=DEFAULT_MIN_BIN_HEIGHT,
+    min_derivative=DEFAULT_MIN_DERIVATIVE,
+):
     if torch.min(inputs) < left or torch.max(inputs) > right:
-        raise ValueError('Input to a transform is not within its domain')
+        raise ValueError("Input to a transform is not within its domain")
 
     num_bins = unnormalized_widths.shape[-1]
 
     if min_bin_width * num_bins > 1.0:
-        raise ValueError('Minimal bin width too large for the number of bins')
+        raise ValueError("Minimal bin width too large for the number of bins")
     if min_bin_height * num_bins > 1.0:
-        raise ValueError('Minimal bin height too large for the number of bins')
+        raise ValueError("Minimal bin height too large for the number of bins")
 
     widths = F.softmax(unnormalized_widths, dim=-1)
     widths = min_bin_width + (1 - min_bin_width * num_bins) * widths
     cumwidths = torch.cumsum(widths, dim=-1)
-    cumwidths = F.pad(cumwidths, pad=(1, 0), mode='constant', value=0.0)
+    cumwidths = F.pad(cumwidths, pad=(1, 0), mode="constant", value=0.0)
     cumwidths = (right - left) * cumwidths + left
     cumwidths[..., 0] = left
     cumwidths[..., -1] = right
@@ -126,7 +135,7 @@ def rational_quadratic_spline(inputs,
     heights = F.softmax(unnormalized_heights, dim=-1)
     heights = min_bin_height + (1 - min_bin_height * num_bins) * heights
     cumheights = torch.cumsum(heights, dim=-1)
-    cumheights = F.pad(cumheights, pad=(1, 0), mode='constant', value=0.0)
+    cumheights = F.pad(cumheights, pad=(1, 0), mode="constant", value=0.0)
     cumheights = (top - bottom) * cumheights + bottom
     cumheights[..., 0] = bottom
     cumheights[..., -1] = top
@@ -150,15 +159,13 @@ def rational_quadratic_spline(inputs,
     input_heights = heights.gather(-1, bin_idx)[..., 0]
 
     if inverse:
-        a = (((inputs - input_cumheights) * (input_derivatives
-                                             + input_derivatives_plus_one
-                                             - 2 * input_delta)
-              + input_heights * (input_delta - input_derivatives)))
-        b = (input_heights * input_derivatives
-             - (inputs - input_cumheights) * (input_derivatives
-                                              + input_derivatives_plus_one
-                                              - 2 * input_delta))
-        c = - input_delta * (inputs - input_cumheights)
+        a = (inputs - input_cumheights) * (
+            input_derivatives + input_derivatives_plus_one - 2 * input_delta
+        ) + input_heights * (input_delta - input_derivatives)
+        b = input_heights * input_derivatives - (inputs - input_cumheights) * (
+            input_derivatives + input_derivatives_plus_one - 2 * input_delta
+        )
+        c = -input_delta * (inputs - input_cumheights)
 
         discriminant = b.pow(2) - 4 * a * c
         assert (discriminant >= 0).all()
@@ -167,11 +174,15 @@ def rational_quadratic_spline(inputs,
         outputs = root * input_bin_widths + input_cumwidths
 
         theta_one_minus_theta = root * (1 - root)
-        denominator = input_delta + ((input_derivatives + input_derivatives_plus_one - 2 * input_delta)
-                                     * theta_one_minus_theta)
-        derivative_numerator = input_delta.pow(2) * (input_derivatives_plus_one * root.pow(2)
-                                                     + 2 * input_delta * theta_one_minus_theta
-                                                     + input_derivatives * (1 - root).pow(2))
+        denominator = input_delta + (
+            (input_derivatives + input_derivatives_plus_one - 2 * input_delta)
+            * theta_one_minus_theta
+        )
+        derivative_numerator = input_delta.pow(2) * (
+            input_derivatives_plus_one * root.pow(2)
+            + 2 * input_delta * theta_one_minus_theta
+            + input_derivatives * (1 - root).pow(2)
+        )
         logabsdet = torch.log(derivative_numerator) - 2 * torch.log(denominator)
 
         return outputs, -logabsdet
@@ -179,15 +190,20 @@ def rational_quadratic_spline(inputs,
         theta = (inputs - input_cumwidths) / input_bin_widths
         theta_one_minus_theta = theta * (1 - theta)
 
-        numerator = input_heights * (input_delta * theta.pow(2)
-                                     + input_derivatives * theta_one_minus_theta)
-        denominator = input_delta + ((input_derivatives + input_derivatives_plus_one - 2 * input_delta)
-                                     * theta_one_minus_theta)
+        numerator = input_heights * (
+            input_delta * theta.pow(2) + input_derivatives * theta_one_minus_theta
+        )
+        denominator = input_delta + (
+            (input_derivatives + input_derivatives_plus_one - 2 * input_delta)
+            * theta_one_minus_theta
+        )
         outputs = input_cumheights + numerator / denominator
 
-        derivative_numerator = input_delta.pow(2) * (input_derivatives_plus_one * theta.pow(2)
-                                                     + 2 * input_delta * theta_one_minus_theta
-                                                     + input_derivatives * (1 - theta).pow(2))
+        derivative_numerator = input_delta.pow(2) * (
+            input_derivatives_plus_one * theta.pow(2)
+            + 2 * input_delta * theta_one_minus_theta
+            + input_derivatives * (1 - theta).pow(2)
+        )
         logabsdet = torch.log(derivative_numerator) - 2 * torch.log(denominator)
 
         return outputs, logabsdet

utils.py

@@ -1,6 +1,5 @@
 import os
 import glob
-import sys
 import argparse
 import logging
 import json
@@ -16,62 +15,98 @@ logger = logging.getLogger(__name__)
 
 def load_checkpoint(checkpoint_path, model, optimizer=None, skip_optimizer=False):
     assert os.path.isfile(checkpoint_path)
-    checkpoint_dict = torch.load(checkpoint_path, map_location='cpu')
-    iteration = checkpoint_dict['iteration']
-    learning_rate = checkpoint_dict['learning_rate']
-    if optimizer is not None and not skip_optimizer and checkpoint_dict['optimizer'] is not None:
-        optimizer.load_state_dict(checkpoint_dict['optimizer'])
-    elif optimizer is None and not skip_optimizer:  
-    #else:      Disable this line if Infer and resume checkpoint,then enable the line upper
+    checkpoint_dict = torch.load(checkpoint_path, map_location="cpu")
+    iteration = checkpoint_dict["iteration"]
+    learning_rate = checkpoint_dict["learning_rate"]
+    if (
+        optimizer is not None
+        and not skip_optimizer
+        and checkpoint_dict["optimizer"] is not None
+    ):
+        optimizer.load_state_dict(checkpoint_dict["optimizer"])
+    elif optimizer is None and not skip_optimizer:
+        # else:      Disable this line if Infer and resume checkpoint,then enable the line upper
         new_opt_dict = optimizer.state_dict()
-        new_opt_dict_params = new_opt_dict['param_groups'][0]['params']
-        new_opt_dict['param_groups'] = checkpoint_dict['optimizer']['param_groups']
-        new_opt_dict['param_groups'][0]['params'] = new_opt_dict_params
+        new_opt_dict_params = new_opt_dict["param_groups"][0]["params"]
+        new_opt_dict["param_groups"] = checkpoint_dict["optimizer"]["param_groups"]
+        new_opt_dict["param_groups"][0]["params"] = new_opt_dict_params
         optimizer.load_state_dict(new_opt_dict)
-    saved_state_dict = checkpoint_dict['model']
-    if hasattr(model, 'module'):
+
+    saved_state_dict = checkpoint_dict["model"]
+    if hasattr(model, "module"):
         state_dict = model.module.state_dict()
     else:
         state_dict = model.state_dict()
+
     new_state_dict = {}
     for k, v in state_dict.items():
         try:
-            #assert "emb_g" not in k
-            # print("load", k)
+            # assert "emb_g" not in k
             new_state_dict[k] = saved_state_dict[k]
-            assert saved_state_dict[k].shape == v.shape, (saved_state_dict[k].shape, v.shape)
+            assert saved_state_dict[k].shape == v.shape, (
+                saved_state_dict[k].shape,
+                v.shape,
+            )
         except:
-            logger.error("%s is not in the checkpoint" % k)
+            # For upgrading from the old version
+            if "ja_bert_proj" in k:
+                v = torch.zeros_like(v)
+                logger.warn(
+                    f"Seems you are using the old version of the model, the {k} is automatically set to zero for backward compatibility"
+                )
+            else:
+                logger.error(f"{k} is not in the checkpoint")
+
             new_state_dict[k] = v
-    if hasattr(model, 'module'):
+
+    if hasattr(model, "module"):
         model.module.load_state_dict(new_state_dict, strict=False)
     else:
         model.load_state_dict(new_state_dict, strict=False)
-    logger.info("Loaded checkpoint '{}' (iteration {})".format(
-        checkpoint_path, iteration))
+
+    logger.info(
+        "Loaded checkpoint '{}' (iteration {})".format(checkpoint_path, iteration)
+    )
+
     return model, optimizer, learning_rate, iteration
 
 
 def save_checkpoint(model, optimizer, learning_rate, iteration, checkpoint_path):
-    logger.info("Saving model and optimizer state at iteration {} to {}".format(
-        iteration, checkpoint_path))
-    if hasattr(model, 'module'):
+    logger.info(
+        "Saving model and optimizer state at iteration {} to {}".format(
+            iteration, checkpoint_path
+        )
+    )
+    if hasattr(model, "module"):
         state_dict = model.module.state_dict()
     else:
         state_dict = model.state_dict()
-    torch.save({'model': state_dict,
-                'iteration': iteration,
-                'optimizer': optimizer.state_dict(),
-                'learning_rate': learning_rate}, checkpoint_path)
+    torch.save(
+        {
+            "model": state_dict,
+            "iteration": iteration,
+            "optimizer": optimizer.state_dict(),
+            "learning_rate": learning_rate,
+        },
+        checkpoint_path,
+    )
 
 
-def summarize(writer, global_step, scalars={}, histograms={}, images={}, audios={}, audio_sampling_rate=22050):
+def summarize(
+    writer,
+    global_step,
+    scalars={},
+    histograms={},
+    images={},
+    audios={},
+    audio_sampling_rate=22050,
+):
     for k, v in scalars.items():
         writer.add_scalar(k, v, global_step)
     for k, v in histograms.items():
         writer.add_histogram(k, v, global_step)
     for k, v in images.items():
-        writer.add_image(k, v, global_step, dataformats='HWC')
+        writer.add_image(k, v, global_step, dataformats="HWC")
     for k, v in audios.items():
         writer.add_audio(k, v, global_step, audio_sampling_rate)
 
@@ -80,7 +115,6 @@ def latest_checkpoint_path(dir_path, regex="G_*.pth"):
     f_list = glob.glob(os.path.join(dir_path, regex))
     f_list.sort(key=lambda f: int("".join(filter(str.isdigit, f))))
     x = f_list[-1]
-    print(x)
     return x
 
 
@@ -88,23 +122,23 @@ def plot_spectrogram_to_numpy(spectrogram):
     global MATPLOTLIB_FLAG
     if not MATPLOTLIB_FLAG:
         import matplotlib
+
         matplotlib.use("Agg")
         MATPLOTLIB_FLAG = True
-        mpl_logger = logging.getLogger('matplotlib')
+        mpl_logger = logging.getLogger("matplotlib")
         mpl_logger.setLevel(logging.WARNING)
     import matplotlib.pylab as plt
     import numpy as np
 
     fig, ax = plt.subplots(figsize=(10, 2))
-    im = ax.imshow(spectrogram, aspect="auto", origin="lower",
-                   interpolation='none')
+    im = ax.imshow(spectrogram, aspect="auto", origin="lower", interpolation="none")
     plt.colorbar(im, ax=ax)
     plt.xlabel("Frames")
     plt.ylabel("Channels")
     plt.tight_layout()
 
     fig.canvas.draw()
-    data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep='')
+    data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep="")
     data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
     plt.close()
     return data
@@ -114,26 +148,28 @@ def plot_alignment_to_numpy(alignment, info=None):
     global MATPLOTLIB_FLAG
     if not MATPLOTLIB_FLAG:
         import matplotlib
+
         matplotlib.use("Agg")
         MATPLOTLIB_FLAG = True
-        mpl_logger = logging.getLogger('matplotlib')
+        mpl_logger = logging.getLogger("matplotlib")
         mpl_logger.setLevel(logging.WARNING)
     import matplotlib.pylab as plt
     import numpy as np
 
     fig, ax = plt.subplots(figsize=(6, 4))
-    im = ax.imshow(alignment.transpose(), aspect='auto', origin='lower',
-                   interpolation='none')
+    im = ax.imshow(
+        alignment.transpose(), aspect="auto", origin="lower", interpolation="none"
+    )
     fig.colorbar(im, ax=ax)
-    xlabel = 'Decoder timestep'
+    xlabel = "Decoder timestep"
     if info is not None:
-        xlabel += '\n\n' + info
+        xlabel += "\n\n" + info
     plt.xlabel(xlabel)
-    plt.ylabel('Encoder timestep')
+    plt.ylabel("Encoder timestep")
     plt.tight_layout()
 
     fig.canvas.draw()
-    data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep='')
+    data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep="")
     data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
     plt.close()
     return data
@@ -145,17 +181,21 @@ def load_wav_to_torch(full_path):
 
 
 def load_filepaths_and_text(filename, split="|"):
-    with open(filename, encoding='utf-8') as f:
+    with open(filename, encoding="utf-8") as f:
         filepaths_and_text = [line.strip().split(split) for line in f]
     return filepaths_and_text
 
 
 def get_hparams(init=True):
     parser = argparse.ArgumentParser()
-    parser.add_argument('-c', '--config', type=str, default="./configs/base.json",
-                        help='JSON file for configuration')
-    parser.add_argument('-m', '--model', type=str, required=True,
-                        help='Model name')
+    parser.add_argument(
+        "-c",
+        "--config",
+        type=str,
+        default="./configs/base.json",
+        help="JSON file for configuration",
+    )
+    parser.add_argument("-m", "--model", type=str, required=True, help="Model name")
 
     args = parser.parse_args()
     model_dir = os.path.join("./logs", args.model)
@@ -180,31 +220,54 @@ def get_hparams(init=True):
     return hparams
 
 
-def clean_checkpoints(path_to_models='logs/44k/', n_ckpts_to_keep=2, sort_by_time=True):
+def clean_checkpoints(path_to_models="logs/44k/", n_ckpts_to_keep=2, sort_by_time=True):
     """Freeing up space by deleting saved ckpts
 
-  Arguments:
-  path_to_models    --  Path to the model directory
-  n_ckpts_to_keep   --  Number of ckpts to keep, excluding G_0.pth and D_0.pth
-  sort_by_time      --  True -> chronologically delete ckpts
-                        False -> lexicographically delete ckpts
-  """
+    Arguments:
+    path_to_models    --  Path to the model directory
+    n_ckpts_to_keep   --  Number of ckpts to keep, excluding G_0.pth and D_0.pth
+    sort_by_time      --  True -> chronologically delete ckpts
+                          False -> lexicographically delete ckpts
+    """
     import re
-    ckpts_files = [f for f in os.listdir(path_to_models) if os.path.isfile(os.path.join(path_to_models, f))]
-    name_key = (lambda _f: int(re.compile('._(\d+)\.pth').match(_f).group(1)))
-    time_key = (lambda _f: os.path.getmtime(os.path.join(path_to_models, _f)))
+
+    ckpts_files = [
+        f
+        for f in os.listdir(path_to_models)
+        if os.path.isfile(os.path.join(path_to_models, f))
+    ]
+
+    def name_key(_f):
+        return int(re.compile("._(\\d+)\\.pth").match(_f).group(1))
+
+    def time_key(_f):
+        return os.path.getmtime(os.path.join(path_to_models, _f))
+
     sort_key = time_key if sort_by_time else name_key
-    x_sorted = lambda _x: sorted([f for f in ckpts_files if f.startswith(_x) and not f.endswith('_0.pth')],
-                                 key=sort_key)
-    to_del = [os.path.join(path_to_models, fn) for fn in
-              (x_sorted('G')[:-n_ckpts_to_keep] + x_sorted('D')[:-n_ckpts_to_keep])]
-    del_info = lambda fn: logger.info(f".. Free up space by deleting ckpt {fn}")
-    del_routine = lambda x: [os.remove(x), del_info(x)]
-    rs = [del_routine(fn) for fn in to_del]
+
+    def x_sorted(_x):
+        return sorted(
+            [f for f in ckpts_files if f.startswith(_x) and not f.endswith("_0.pth")],
+            key=sort_key,
+        )
+
+    to_del = [
+        os.path.join(path_to_models, fn)
+        for fn in (x_sorted("G")[:-n_ckpts_to_keep] + x_sorted("D")[:-n_ckpts_to_keep])
+    ]
+
+    def del_info(fn):
+        return logger.info(f".. Free up space by deleting ckpt {fn}")
+
+    def del_routine(x):
+        return [os.remove(x), del_info(x)]
+
+    [del_routine(fn) for fn in to_del]
+
 
 def get_hparams_from_dir(model_dir):
     config_save_path = os.path.join(model_dir, "config.json")
-    with open(config_save_path, "r", encoding='utf-8') as f:
+    with open(config_save_path, "r", encoding="utf-8") as f:
         data = f.read()
     config = json.loads(data)
 
@@ -214,7 +277,7 @@ def get_hparams_from_dir(model_dir):
 
 
 def get_hparams_from_file(config_path):
-    with open(config_path, "r", encoding='utf-8') as f:
+    with open(config_path, "r", encoding="utf-8") as f:
         data = f.read()
     config = json.loads(data)
 
@@ -225,9 +288,11 @@ def get_hparams_from_file(config_path):
 def check_git_hash(model_dir):
     source_dir = os.path.dirname(os.path.realpath(__file__))
     if not os.path.exists(os.path.join(source_dir, ".git")):
-        logger.warn("{} is not a git repository, therefore hash value comparison will be ignored.".format(
-            source_dir
-        ))
+        logger.warn(
+            "{} is not a git repository, therefore hash value comparison will be ignored.".format(
+                source_dir
+            )
+        )
         return
 
     cur_hash = subprocess.getoutput("git rev-parse HEAD")
@@ -236,8 +301,11 @@ def check_git_hash(model_dir):
     if os.path.exists(path):
         saved_hash = open(path).read()
         if saved_hash != cur_hash:
-            logger.warn("git hash values are different. {}(saved) != {}(current)".format(
-                saved_hash[:8], cur_hash[:8]))
+            logger.warn(
+                "git hash values are different. {}(saved) != {}(current)".format(
+                    saved_hash[:8], cur_hash[:8]
+                )
+            )
     else:
         open(path, "w").write(cur_hash)
 
@@ -257,7 +325,7 @@ def get_logger(model_dir, filename="train.log"):
     return logger
 
 
-class HParams():
+class HParams:
     def __init__(self, **kwargs):
         for k, v in kwargs.items():
             if type(v) == dict:

webui.py

@@ -1,8 +1,6 @@
+# flake8: noqa: E402
+
 import sys, os
-
-if sys.platform == "darwin":
-    os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
-
 import logging
 
 logging.getLogger("numba").setLevel(logging.WARNING)
@@ -10,7 +8,9 @@ logging.getLogger("markdown_it").setLevel(logging.WARNING)
 logging.getLogger("urllib3").setLevel(logging.WARNING)
 logging.getLogger("matplotlib").setLevel(logging.WARNING)
 
-logging.basicConfig(level=logging.INFO, format="| %(name)s | %(levelname)s | %(message)s")
+logging.basicConfig(
+    level=logging.INFO, format="| %(name)s | %(levelname)s | %(message)s"
+)
 
 logger = logging.getLogger(__name__)
 
@@ -25,9 +25,14 @@ from text.cleaner import clean_text
 import gradio as gr
 import webbrowser
 
-
 net_g = None
 
+if sys.platform == "darwin" and torch.backends.mps.is_available():
+    device = "mps"
+    os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
+else:
+    device = "cuda"
+
 
 def get_text(text, language_str, hps):
     norm_text, phone, tone, word2ph = clean_text(text, language_str)
@@ -40,45 +45,98 @@ def get_text(text, language_str, hps):
         for i in range(len(word2ph)):
             word2ph[i] = word2ph[i] * 2
         word2ph[0] += 1
-    bert = get_bert(norm_text, word2ph, language_str)
+    bert = get_bert(norm_text, word2ph, language_str, device)
     del word2ph
+    assert bert.shape[-1] == len(phone), phone
 
-    assert bert.shape[-1] == len(phone)
+    if language_str == "ZH":
+        bert = bert
+        ja_bert = torch.zeros(768, len(phone))
+    elif language_str == "JA":
+        ja_bert = bert
+        bert = torch.zeros(1024, len(phone))
+    else:
+        bert = torch.zeros(1024, len(phone))
+        ja_bert = torch.zeros(768, len(phone))
+
+    assert bert.shape[-1] == len(
+        phone
+    ), f"Bert seq len {bert.shape[-1]} != {len(phone)}"
 
     phone = torch.LongTensor(phone)
     tone = torch.LongTensor(tone)
     language = torch.LongTensor(language)
+    return bert, ja_bert, phone, tone, language
 
-    return bert, phone, tone, language
 
-def infer(text, sdp_ratio, noise_scale, noise_scale_w, length_scale, sid):
+def infer(text, sdp_ratio, noise_scale, noise_scale_w, length_scale, sid, language):
     global net_g
-    bert, phones, tones, lang_ids = get_text(text, "ZH", hps)
+    bert, ja_bert, phones, tones, lang_ids = get_text(text, language, hps)
     with torch.no_grad():
-        x_tst=phones.to(device).unsqueeze(0)
-        tones=tones.to(device).unsqueeze(0)
-        lang_ids=lang_ids.to(device).unsqueeze(0)
+        x_tst = phones.to(device).unsqueeze(0)
+        tones = tones.to(device).unsqueeze(0)
+        lang_ids = lang_ids.to(device).unsqueeze(0)
         bert = bert.to(device).unsqueeze(0)
+        ja_bert = ja_bert.to(device).unsqueeze(0)
         x_tst_lengths = torch.LongTensor([phones.size(0)]).to(device)
         del phones
         speakers = torch.LongTensor([hps.data.spk2id[sid]]).to(device)
-        audio = net_g.infer(x_tst, x_tst_lengths, speakers, tones, lang_ids, bert, sdp_ratio=sdp_ratio
-                           , noise_scale=noise_scale, noise_scale_w=noise_scale_w, length_scale=length_scale)[0][0,0].data.cpu().float().numpy()
+        audio = (
+            net_g.infer(
+                x_tst,
+                x_tst_lengths,
+                speakers,
+                tones,
+                lang_ids,
+                bert,
+                ja_bert,
+                sdp_ratio=sdp_ratio,
+                noise_scale=noise_scale,
+                noise_scale_w=noise_scale_w,
+                length_scale=length_scale,
+            )[0][0, 0]
+            .data.cpu()
+            .float()
+            .numpy()
+        )
         del x_tst, tones, lang_ids, bert, x_tst_lengths, speakers
         return audio
 
-def tts_fn(text, speaker, sdp_ratio, noise_scale, noise_scale_w, length_scale):
+
+def tts_fn(
+    text, speaker, sdp_ratio, noise_scale, noise_scale_w, length_scale, language
+):
     with torch.no_grad():
-        audio = infer(text, sdp_ratio=sdp_ratio, noise_scale=noise_scale, noise_scale_w=noise_scale_w, length_scale=length_scale, sid=speaker)
+        audio = infer(
+            text,
+            sdp_ratio=sdp_ratio,
+            noise_scale=noise_scale,
+            noise_scale_w=noise_scale_w,
+            length_scale=length_scale,
+            sid=speaker,
+            language=language,
+        )
+        torch.cuda.empty_cache()
     return "Success", (hps.data.sampling_rate, audio)
 
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
-    parser.add_argument("-m", "--model", default="./logs/as/G_8000.pth", help="path of your model")
-    parser.add_argument("-c", "--config", default="./configs/config.json", help="path of your config file")
-    parser.add_argument("--share", default=False, help="make link public")
-    parser.add_argument("-d", "--debug", action="store_true", help="enable DEBUG-LEVEL log")
+    parser.add_argument(
+        "-m", "--model", default="./logs/as/G_8000.pth", help="path of your model"
+    )
+    parser.add_argument(
+        "-c",
+        "--config",
+        default="./configs/config.json",
+        help="path of your config file",
+    )
+    parser.add_argument(
+        "--share", default=False, help="make link public", action="store_true"
+    )
+    parser.add_argument(
+        "-d", "--debug", action="store_true", help="enable DEBUG-LEVEL log"
+    )
 
     args = parser.parse_args()
     if args.debug:
@@ -100,31 +158,59 @@ if __name__ == "__main__":
         hps.data.filter_length // 2 + 1,
         hps.train.segment_size // hps.data.hop_length,
         n_speakers=hps.data.n_speakers,
-        **hps.model).to(device)
+        **hps.model,
+    ).to(device)
     _ = net_g.eval()
 
     _ = utils.load_checkpoint(args.model, net_g, None, skip_optimizer=True)
 
     speaker_ids = hps.data.spk2id
     speakers = list(speaker_ids.keys())
+    languages = ["ZH", "JP"]
     with gr.Blocks() as app:
         with gr.Row():
             with gr.Column():
-                text = gr.TextArea(label="Text", placeholder="Input Text Here",
-                                      value="吃葡萄不吐葡萄皮，不吃葡萄倒吐葡萄皮。")
-                speaker = gr.Dropdown(choices=speakers, value=speakers[0], label='Speaker')
-                sdp_ratio = gr.Slider(minimum=0, maximum=1, value=0.2, step=0.1, label='SDP Ratio')
-                noise_scale = gr.Slider(minimum=0.1, maximum=2, value=0.6, step=0.1, label='Noise Scale')
-                noise_scale_w = gr.Slider(minimum=0.1, maximum=2, value=0.8, step=0.1, label='Noise Scale W')
-                length_scale = gr.Slider(minimum=0.1, maximum=2, value=1, step=0.1, label='Length Scale')
+                text = gr.TextArea(
+                    label="Text",
+                    placeholder="Input Text Here",
+                    value="吃葡萄不吐葡萄皮，不吃葡萄倒吐葡萄皮。",
+                )
+                speaker = gr.Dropdown(
+                    choices=speakers, value=speakers[0], label="Speaker"
+                )
+                sdp_ratio = gr.Slider(
+                    minimum=0, maximum=1, value=0.2, step=0.1, label="SDP Ratio"
+                )
+                noise_scale = gr.Slider(
+                    minimum=0.1, maximum=2, value=0.6, step=0.1, label="Noise Scale"
+                )
+                noise_scale_w = gr.Slider(
+                    minimum=0.1, maximum=2, value=0.8, step=0.1, label="Noise Scale W"
+                )
+                length_scale = gr.Slider(
+                    minimum=0.1, maximum=2, value=1, step=0.1, label="Length Scale"
+                )
+                language = gr.Dropdown(
+                    choices=languages, value=languages[0], label="Language"
+                )
                 btn = gr.Button("Generate!", variant="primary")
             with gr.Column():
                 text_output = gr.Textbox(label="Message")
                 audio_output = gr.Audio(label="Output Audio")
 
-        btn.click(tts_fn,
-                inputs=[text, speaker, sdp_ratio, noise_scale, noise_scale_w, length_scale],
-                outputs=[text_output, audio_output])
-    
+        btn.click(
+            tts_fn,
+            inputs=[
+                text,
+                speaker,
+                sdp_ratio,
+                noise_scale,
+                noise_scale_w,
+                length_scale,
+                language,
+            ],
+            outputs=[text_output, audio_output],
+        )
+
     webbrowser.open("http://127.0.0.1:7860")
     app.launch(share=args.share)