Refactor: moved monotonic_align/ to style_bert_vits2/models/monotonic_alignment.py

style_bert_vits2/models/infer.py

@@ -10,10 +10,6 @@ from style_bert_vits2.text_processing import clean_text, cleaned_text_to_sequenc
 from style_bert_vits2.text_processing.symbols import SYMBOLS
 
 
-class InvalidToneError(ValueError):
-    pass
-
-
 def get_net_g(model_path: str, version: str, device: str, hps):
     if version.endswith("JP-Extra"):
         logger.info("Using JP-Extra model")
@@ -315,3 +311,7 @@ def infer_multilang(
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
         return audio
+
+
+class InvalidToneError(ValueError):
+    pass

style_bert_vits2/models/models.py

@@ -6,10 +6,10 @@ from torch.nn import Conv1d, Conv2d, ConvTranspose1d
 from torch.nn import functional as F
 from torch.nn.utils import remove_weight_norm, spectral_norm, weight_norm
 
-import monotonic_align
 from style_bert_vits2.models import attentions
 from style_bert_vits2.models import commons
 from style_bert_vits2.models import modules
+from style_bert_vits2.models import monotonic_alignment
 from style_bert_vits2.models.commons import get_padding, init_weights
 from style_bert_vits2.text_processing.symbols import NUM_LANGUAGES, NUM_TONES, SYMBOLS
 
@@ -932,7 +932,7 @@ class SynthesizerTrn(nn.Module):
 
             attn_mask = torch.unsqueeze(x_mask, 2) * torch.unsqueeze(y_mask, -1)
             attn = (
-                monotonic_align.maximum_path(neg_cent, attn_mask.squeeze(1))
+                monotonic_alignment.maximum_path(neg_cent, attn_mask.squeeze(1))
                 .unsqueeze(1)
                 .detach()
             )

style_bert_vits2/models/models_jp_extra.py

@@ -6,10 +6,10 @@ from torch.nn import Conv1d, Conv2d, ConvTranspose1d
 from torch.nn import functional as F
 from torch.nn.utils import remove_weight_norm, spectral_norm, weight_norm
 
-import monotonic_align
 from style_bert_vits2.models import attentions
 from style_bert_vits2.models import commons
 from style_bert_vits2.models import modules
+from style_bert_vits2.models import monotonic_alignment
 from style_bert_vits2.text_processing.symbols import SYMBOLS, NUM_TONES, NUM_LANGUAGES
 
 
@@ -979,7 +979,7 @@ class SynthesizerTrn(nn.Module):
 
             attn_mask = torch.unsqueeze(x_mask, 2) * torch.unsqueeze(y_mask, -1)
             attn = (
-                monotonic_align.maximum_path(neg_cent, attn_mask.squeeze(1))
+                monotonic_alignment.maximum_path(neg_cent, attn_mask.squeeze(1))
                 .unsqueeze(1)
                 .detach()
             )

style_bert_vits2/models/modules.py

@@ -10,6 +10,7 @@ from transforms import piecewise_rational_quadratic_transform
 from style_bert_vits2.models import commons
 from style_bert_vits2.models.attentions import Encoder
 
+
 LRELU_SLOPE = 0.1
 
 

style_bert_vits2/models/monotonic_alignment.py

@@ -0,0 +1,88 @@
+"""
+以下に記述されている関数のコメントはリファクタリング時に GPT-4 に生成させたもので、
+コードと完全に一致している保証はない。あくまで参考程度とすること。
+"""
+
+import numba
+import torch
+from numpy import int32, float32, zeros
+from typing import Any
+
+
+def maximum_path(neg_cent: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+    """
+    与えられた負の中心とマスクを使用して最大パスを計算する
+
+    Args:
+        neg_cent (torch.Tensor): 負の中心を表すテンソル
+        mask (torch.Tensor): マスクを表すテンソル
+
+    Returns:
+        Tensor: 計算された最大パスを表すテンソル
+    """
+
+    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 torch.from_numpy(path).to(device=device, dtype=dtype)
+
+
+@numba.jit(
+    numba.void(
+        numba.int32[:, :, ::1],
+        numba.float32[:, :, ::1],
+        numba.int32[::1],
+        numba.int32[::1],
+    ),
+    nopython = True,
+    nogil = True,
+)  # type: ignore
+def maximum_path_jit(paths: Any, values: Any, t_ys: Any, t_xs: Any) -> None:
+    """
+    与えられたパス、値、およびターゲットの y と x 座標を使用して JIT で最大パスを計算する
+
+    Args:
+        paths: 計算されたパスを格納するための整数型の 3 次元配列
+        values: 値を格納するための浮動小数点型の 3 次元配列
+        t_ys: ターゲットの y 座標を格納するための整数型の 1 次元配列
+        t_xs: ターゲットの x 座標を格納するための整数型の 1 次元配列
+    """
+
+    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
+
+        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