Merge pull request #92 from zunan-islands/dev

コードベースを大規模にリファクタリングし、ライブラリ (Python パッケージ) としてほかの Python コードから利用できるように改善

.gitignore

@@ -1,7 +1,8 @@
-.vscode/
-
 __pycache__/
 venv/
+.venv/
+dist/
+.coverage
 .ipynb_checkpoints/
 
 /*.yml

.vscode/extensions.json

@@ -0,0 +1,6 @@
+{
+    "recommendations": [
+        "ms-python.python",
+        "ms-python.vscode-pylance"
+    ]
+}

.vscode/settings.json

@@ -0,0 +1,22 @@
+{
+    // Pylance の Type Checking を有効化
+    "python.languageServer": "Pylance",
+    "python.analysis.typeCheckingMode": "strict",
+    // Pylance の Type Checking のうち、いくつかのエラー報告を抑制する
+    "python.analysis.diagnosticSeverityOverrides": {
+        "reportConstantRedefinition": "none",
+        "reportGeneralTypeIssues": "warning",
+        "reportMissingParameterType": "warning",
+        "reportMissingTypeStubs": "none",
+        "reportPrivateImportUsage": "none",
+        "reportPrivateUsage": "warning",
+        "reportShadowedImports": "none",
+        "reportUnnecessaryComparison": "none",
+        "reportUnknownArgumentType": "none",
+        "reportUnknownMemberType": "none",
+        "reportUnknownParameterType": "warning",
+        "reportUnknownVariableType": "none",
+        "reportUnusedFunction": "none",
+        "reportUnusedVariable": "information",
+    },
+}

app.py

@@ -5,8 +5,8 @@ import gradio as gr
 import torch
 import yaml
 
-from common.constants import GRADIO_THEME, LATEST_VERSION
+from style_bert_vits2.constants import GRADIO_THEME, VERSION
-from common.tts_model import ModelHolder
+from style_bert_vits2.tts_model import TTSModelHolder
 from webui import (
     create_dataset_app,
     create_inference_app,
@@ -15,6 +15,7 @@ from webui import (
     create_train_app,
 )
 
+
 # Get path settings
 with Path("configs/paths.yml").open("r", encoding="utf-8") as f:
     path_config: dict[str, str] = yaml.safe_load(f.read())
@@ -23,6 +24,8 @@ with Path("configs/paths.yml").open("r", encoding="utf-8") as f:
 
 parser = argparse.ArgumentParser()
 parser.add_argument("--device", type=str, default="cuda")
+parser.add_argument("--host", type=str, default="127.0.0.1")
+parser.add_argument("--port", type=int, default=7860)
 parser.add_argument("--no_autolaunch", action="store_true")
 parser.add_argument("--share", action="store_true")
 
@@ -31,10 +34,10 @@ device = args.device
 if device == "cuda" and not torch.cuda.is_available():
     device = "cpu"
 
-model_holder = ModelHolder(Path(assets_root), device)
+model_holder = TTSModelHolder(Path(assets_root), device)
 
 with gr.Blocks(theme=GRADIO_THEME) as app:
-    gr.Markdown(f"# Style-Bert-VITS2 WebUI (version {LATEST_VERSION})")
+    gr.Markdown(f"# Style-Bert-VITS2 WebUI (version {VERSION})")
     with gr.Tabs():
         with gr.Tab("音声合成"):
             create_inference_app(model_holder=model_holder)
@@ -48,4 +51,4 @@ with gr.Blocks(theme=GRADIO_THEME) as app:
             create_merge_app(model_holder=model_holder)
 
 
-app.launch(inbrowser=not args.no_autolaunch, share=args.share)
+app.launch(server_name=args.host, server_port=args.port, inbrowser=not args.no_autolaunch, share=args.share)

bert_gen.py

@@ -5,14 +5,12 @@ import torch
 import torch.multiprocessing as mp
 from tqdm import tqdm
 
-import commons
-from text.get_bert import get_bert
-import text.pyopenjtalk_worker as pyopenjtalk
-import utils
-from common.log import logger
-from common.stdout_wrapper import SAFE_STDOUT
 from config import config
-from text import cleaned_text_to_sequence
+from style_bert_vits2.logging import logger
+from style_bert_vits2.models import commons
+from style_bert_vits2.models.hyper_parameters import HyperParameters
+from style_bert_vits2.nlp import cleaned_text_to_sequence, extract_bert_feature
+from style_bert_vits2.utils.stdout_wrapper import SAFE_STDOUT
 
 
 def process_line(x):
@@ -47,7 +45,7 @@ def process_line(x):
         bert = torch.load(bert_path)
         assert bert.shape[-1] == len(phone)
     except Exception:
-        bert = get_bert(text, word2ph, language_str, device)
+        bert = extract_bert_feature(text, word2ph, language_str, device)
         assert bert.shape[-1] == len(phone)
         torch.save(bert, bert_path)
 
@@ -64,7 +62,7 @@ if __name__ == "__main__":
     )
     args, _ = parser.parse_known_args()
     config_path = args.config
-    hps = utils.get_hparams_from_file(config_path)
+    hps = HyperParameters.load_from_json(config_path)
     lines = []
     with open(hps.data.training_files, encoding="utf-8") as f:
         lines.extend(f.readlines())

common/constants.py

@@ -1,28 +0,0 @@
-import enum
-
-# Built-in theme: "default", "base", "monochrome", "soft", "glass"
-# See https://huggingface.co/spaces/gradio/theme-gallery for more themes
-GRADIO_THEME: str = "NoCrypt/miku"
-
-LATEST_VERSION: str = "2.4"
-
-USER_DICT_DIR = "dict_data"
-
-DEFAULT_STYLE: str = "Neutral"
-DEFAULT_STYLE_WEIGHT: float = 5.0
-
-
-class Languages(str, enum.Enum):
-    JP = "JP"
-    EN = "EN"
-    ZH = "ZH"
-
-
-DEFAULT_SDP_RATIO: float = 0.2
-DEFAULT_NOISE: float = 0.6
-DEFAULT_NOISEW: float = 0.8
-DEFAULT_LENGTH: float = 1.0
-DEFAULT_LINE_SPLIT: bool = True
-DEFAULT_SPLIT_INTERVAL: float = 0.5
-DEFAULT_ASSIST_TEXT_WEIGHT: float = 0.7
-DEFAULT_ASSIST_TEXT_WEIGHT: float = 1.0

common/log.py

@@ -1,18 +0,0 @@
-"""
-logger封装
-"""
-
-from loguru import logger
-
-from .stdout_wrapper import SAFE_STDOUT
-
-# 移除所有默认的处理器
-logger.remove()
-
-# 自定义格式并添加到标准输出
-log_format = (
-    "<g>{time:MM-DD HH:mm:ss}</g> |<lvl>{level:^8}</lvl>| {file}:{line} | {message}"
-)
-
-# logger.add(SAFE_STDOUT, format=log_format, backtrace=True, diagnose=True, level="TRACE")
-logger.add(SAFE_STDOUT, format=log_format, backtrace=True, diagnose=True)

common/subprocess_utils.py

@@ -1,33 +0,0 @@
-import subprocess
-import sys
-
-from .log import logger
-from .stdout_wrapper import SAFE_STDOUT
-
-python = sys.executable
-
-
-def run_script_with_log(cmd: list[str], ignore_warning=False) -> tuple[bool, str]:
-    logger.info(f"Running: {' '.join(cmd)}")
-    result = subprocess.run(
-        [python] + cmd,
-        stdout=SAFE_STDOUT,  # type: ignore
-        stderr=subprocess.PIPE,
-        text=True,
-        encoding="utf-8",
-    )
-    if result.returncode != 0:
-        logger.error(f"Error: {' '.join(cmd)}\n{result.stderr}")
-        return False, result.stderr
-    elif result.stderr and not ignore_warning:
-        logger.warning(f"Warning: {' '.join(cmd)}\n{result.stderr}")
-        return True, result.stderr
-    logger.success(f"Success: {' '.join(cmd)}")
-    return True, ""
-
-
-def second_elem_of(original_function):
-    def inner_function(*args, **kwargs):
-        return original_function(*args, **kwargs)[1]
-
-    return inner_function

common/tts_model.py

@@ -1,327 +0,0 @@
-import os
-import warnings
-from pathlib import Path
-from typing import Optional, Union
-
-import gradio as gr
-import numpy as np
-
-import torch
-from gradio.processing_utils import convert_to_16_bit_wav
-
-import utils
-from infer import get_net_g, infer
-from models import SynthesizerTrn
-from models_jp_extra import SynthesizerTrn as SynthesizerTrnJPExtra
-
-from .constants import (
-    DEFAULT_ASSIST_TEXT_WEIGHT,
-    DEFAULT_LENGTH,
-    DEFAULT_LINE_SPLIT,
-    DEFAULT_NOISE,
-    DEFAULT_NOISEW,
-    DEFAULT_SDP_RATIO,
-    DEFAULT_SPLIT_INTERVAL,
-    DEFAULT_STYLE,
-    DEFAULT_STYLE_WEIGHT,
-)
-from .log import logger
-
-
-def adjust_voice(fs, wave, pitch_scale, intonation_scale):
-    if pitch_scale == 1.0 and intonation_scale == 1.0:
-        # 初期値の場合は、音質劣化を避けるためにそのまま返す
-        return fs, wave
-
-    try:
-        import pyworld
-    except ImportError:
-        raise ImportError(
-            "pyworld is not installed. Please install it by `pip install pyworld`"
-        )
-
-    # pyworldでf0を加工して合成
-    # pyworldよりもよいのがあるかもしれないが……
-
-    wave = wave.astype(np.double)
-    f0, t = pyworld.harvest(wave, fs)
-    # 質が高そうだしとりあえずharvestにしておく
-
-    sp = pyworld.cheaptrick(wave, f0, t, fs)
-    ap = pyworld.d4c(wave, f0, t, fs)
-
-    non_zero_f0 = [f for f in f0 if f != 0]
-    f0_mean = sum(non_zero_f0) / len(non_zero_f0)
-
-    for i, f in enumerate(f0):
-        if f == 0:
-            continue
-        f0[i] = pitch_scale * f0_mean + intonation_scale * (f - f0_mean)
-
-    wave = pyworld.synthesize(f0, sp, ap, fs)
-    return fs, wave
-
-
-class Model:
-    def __init__(
-        self, model_path: Path, config_path: Path, style_vec_path: Path, device: str
-    ):
-        self.model_path: Path = model_path
-        self.config_path: Path = config_path
-        self.style_vec_path: Path = style_vec_path
-        self.device: str = device
-        self.hps: utils.HParams = utils.get_hparams_from_file(self.config_path)
-        self.spk2id: dict[str, int] = self.hps.data.spk2id
-        self.id2spk: dict[int, str] = {v: k for k, v in self.spk2id.items()}
-
-        self.num_styles: int = self.hps.data.num_styles
-        if hasattr(self.hps.data, "style2id"):
-            self.style2id: dict[str, int] = self.hps.data.style2id
-        else:
-            self.style2id: dict[str, int] = {str(i): i for i in range(self.num_styles)}
-        if len(self.style2id) != self.num_styles:
-            raise ValueError(
-                f"Number of styles ({self.num_styles}) does not match the number of style2id ({len(self.style2id)})"
-            )
-
-        self.style_vectors: np.ndarray = np.load(self.style_vec_path)
-        if self.style_vectors.shape[0] != self.num_styles:
-            raise ValueError(
-                f"The number of styles ({self.num_styles}) does not match the number of style vectors ({self.style_vectors.shape[0]})"
-            )
-
-        self.net_g: Union[SynthesizerTrn, SynthesizerTrnJPExtra, None] = None
-
-    def load_net_g(self):
-        self.net_g = get_net_g(
-            model_path=str(self.model_path),
-            version=self.hps.version,
-            device=self.device,
-            hps=self.hps,
-        )
-
-    def get_style_vector(self, style_id: int, weight: float = 1.0) -> np.ndarray:
-        mean = self.style_vectors[0]
-        style_vec = self.style_vectors[style_id]
-        style_vec = mean + (style_vec - mean) * weight
-        return style_vec
-
-    def get_style_vector_from_audio(
-        self, audio_path: str, weight: float = 1.0
-    ) -> np.ndarray:
-        from style_gen import get_style_vector
-
-        xvec = get_style_vector(audio_path)
-        mean = self.style_vectors[0]
-        xvec = mean + (xvec - mean) * weight
-        return xvec
-
-    def infer(
-        self,
-        text: str,
-        language: str = "JP",
-        sid: int = 0,
-        reference_audio_path: Optional[str] = None,
-        sdp_ratio: float = DEFAULT_SDP_RATIO,
-        noise: float = DEFAULT_NOISE,
-        noisew: float = DEFAULT_NOISEW,
-        length: float = DEFAULT_LENGTH,
-        line_split: bool = DEFAULT_LINE_SPLIT,
-        split_interval: float = DEFAULT_SPLIT_INTERVAL,
-        assist_text: Optional[str] = None,
-        assist_text_weight: float = DEFAULT_ASSIST_TEXT_WEIGHT,
-        use_assist_text: bool = False,
-        style: str = DEFAULT_STYLE,
-        style_weight: float = DEFAULT_STYLE_WEIGHT,
-        given_tone: Optional[list[int]] = None,
-        pitch_scale: float = 1.0,
-        intonation_scale: float = 1.0,
-    ) -> tuple[int, np.ndarray]:
-        logger.info(f"Start generating audio data from text:\n{text}")
-        if language != "JP" and self.hps.version.endswith("JP-Extra"):
-            raise ValueError(
-                "The model is trained with JP-Extra, but the language is not JP"
-            )
-        if reference_audio_path == "":
-            reference_audio_path = None
-        if assist_text == "" or not use_assist_text:
-            assist_text = None
-
-        if self.net_g is None:
-            self.load_net_g()
-        if reference_audio_path is None:
-            style_id = self.style2id[style]
-            style_vector = self.get_style_vector(style_id, style_weight)
-        else:
-            style_vector = self.get_style_vector_from_audio(
-                reference_audio_path, style_weight
-            )
-        if not line_split:
-            with torch.no_grad():
-                audio = infer(
-                    text=text,
-                    sdp_ratio=sdp_ratio,
-                    noise_scale=noise,
-                    noise_scale_w=noisew,
-                    length_scale=length,
-                    sid=sid,
-                    language=language,
-                    hps=self.hps,
-                    net_g=self.net_g,
-                    device=self.device,
-                    assist_text=assist_text,
-                    assist_text_weight=assist_text_weight,
-                    style_vec=style_vector,
-                    given_tone=given_tone,
-                )
-        else:
-            texts = text.split("\n")
-            texts = [t for t in texts if t != ""]
-            audios = []
-            with torch.no_grad():
-                for i, t in enumerate(texts):
-                    audios.append(
-                        infer(
-                            text=t,
-                            sdp_ratio=sdp_ratio,
-                            noise_scale=noise,
-                            noise_scale_w=noisew,
-                            length_scale=length,
-                            sid=sid,
-                            language=language,
-                            hps=self.hps,
-                            net_g=self.net_g,
-                            device=self.device,
-                            assist_text=assist_text,
-                            assist_text_weight=assist_text_weight,
-                            style_vec=style_vector,
-                        )
-                    )
-                    if i != len(texts) - 1:
-                        audios.append(np.zeros(int(44100 * split_interval)))
-                audio = np.concatenate(audios)
-        logger.info("Audio data generated successfully")
-        if not (pitch_scale == 1.0 and intonation_scale == 1.0):
-            _, audio = adjust_voice(
-                fs=self.hps.data.sampling_rate,
-                wave=audio,
-                pitch_scale=pitch_scale,
-                intonation_scale=intonation_scale,
-            )
-        with warnings.catch_warnings():
-            warnings.simplefilter("ignore")
-            audio = convert_to_16_bit_wav(audio)
-        return (self.hps.data.sampling_rate, audio)
-
-
-class ModelHolder:
-    def __init__(self, root_dir: Path, device: str):
-        self.root_dir: Path = root_dir
-        self.device: str = device
-        self.model_files_dict: dict[str, list[Path]] = {}
-        self.current_model: Optional[Model] = None
-        self.model_names: list[str] = []
-        self.models: list[Model] = []
-        self.models_info: list[dict[str, Union[str, list[str]]]] = []
-        self.refresh()
-
-    def refresh(self):
-        self.model_files_dict = {}
-        self.model_names = []
-        self.current_model = None
-        self.models_info = []
-
-        model_dirs = [d for d in self.root_dir.iterdir() if d.is_dir()]
-        for model_dir in model_dirs:
-            model_files = [
-                f
-                for f in model_dir.iterdir()
-                if f.suffix in [".pth", ".pt", ".safetensors"]
-            ]
-            if len(model_files) == 0:
-                logger.warning(f"No model files found in {model_dir}, so skip it")
-                continue
-            config_path = model_dir / "config.json"
-            if not config_path.exists():
-                logger.warning(
-                    f"Config file {config_path} not found, so skip {model_dir}"
-                )
-                continue
-            self.model_files_dict[model_dir.name] = model_files
-            self.model_names.append(model_dir.name)
-            hps = utils.get_hparams_from_file(config_path)
-            style2id: dict[str, int] = hps.data.style2id
-            styles = list(style2id.keys())
-            spk2id: dict[str, int] = hps.data.spk2id
-            speakers = list(spk2id.keys())
-            self.models_info.append(
-                {
-                    "name": model_dir.name,
-                    "files": [str(f) for f in model_files],
-                    "styles": styles,
-                    "speakers": speakers,
-                }
-            )
-
-    def load_model(self, model_name: str, model_path_str: str):
-        model_path = Path(model_path_str)
-        if model_name not in self.model_files_dict:
-            raise ValueError(f"Model `{model_name}` is not found")
-        if model_path not in self.model_files_dict[model_name]:
-            raise ValueError(f"Model file `{model_path}` is not found")
-        if self.current_model is None or self.current_model.model_path != model_path:
-            self.current_model = Model(
-                model_path=model_path,
-                config_path=self.root_dir / model_name / "config.json",
-                style_vec_path=self.root_dir / model_name / "style_vectors.npy",
-                device=self.device,
-            )
-        return self.current_model
-
-    def load_model_gr(
-        self, model_name: str, model_path_str: str
-    ) -> tuple[gr.Dropdown, gr.Button, gr.Dropdown]:
-        model_path = Path(model_path_str)
-        if model_name not in self.model_files_dict:
-            raise ValueError(f"Model `{model_name}` is not found")
-        if model_path not in self.model_files_dict[model_name]:
-            raise ValueError(f"Model file `{model_path}` is not found")
-        if (
-            self.current_model is not None
-            and self.current_model.model_path == model_path
-        ):
-            # Already loaded
-            speakers = list(self.current_model.spk2id.keys())
-            styles = list(self.current_model.style2id.keys())
-            return (
-                gr.Dropdown(choices=styles, value=styles[0]),
-                gr.Button(interactive=True, value="音声合成"),
-                gr.Dropdown(choices=speakers, value=speakers[0]),
-            )
-        self.current_model = Model(
-            model_path=model_path,
-            config_path=self.root_dir / model_name / "config.json",
-            style_vec_path=self.root_dir / model_name / "style_vectors.npy",
-            device=self.device,
-        )
-        speakers = list(self.current_model.spk2id.keys())
-        styles = list(self.current_model.style2id.keys())
-        return (
-            gr.Dropdown(choices=styles, value=styles[0]),
-            gr.Button(interactive=True, value="音声合成"),
-            gr.Dropdown(choices=speakers, value=speakers[0]),
-        )
-
-    def update_model_files_gr(self, model_name: str) -> gr.Dropdown:
-        model_files = self.model_files_dict[model_name]
-        return gr.Dropdown(choices=model_files, value=model_files[0])
-
-    def update_model_names_gr(self) -> tuple[gr.Dropdown, gr.Dropdown, gr.Button]:
-        self.refresh()
-        initial_model_name = self.model_names[0]
-        initial_model_files = self.model_files_dict[initial_model_name]
-        return (
-            gr.Dropdown(choices=self.model_names, value=initial_model_name),
-            gr.Dropdown(choices=initial_model_files, value=initial_model_files[0]),
-            gr.Button(interactive=False),  # For tts_button
-        )

commons.py

@@ -1,152 +0,0 @@
-import math
-import torch
-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)
-
-
-def get_padding(kernel_size, dilation=1):
-    return int((kernel_size * dilation - dilation) / 2)
-
-
-def convert_pad_shape(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
-
-
-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.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))
-
-
-def rand_gumbel_like(x):
-    g = rand_gumbel(x.size()).to(dtype=x.dtype, device=x.device)
-    return g
-
-
-def slice_segments(x, ids_str, segment_size=4):
-    gather_indices = ids_str.view(x.size(0), 1, 1).repeat(
-        1, x.size(1), 1
-    ) + torch.arange(segment_size, device=x.device)
-    return torch.gather(x, 2, gather_indices)
-
-
-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 = torch.clamp(x_lengths - segment_size + 1, min=0)
-    ids_str = (torch.rand([b], 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 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)
-
-
-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)
-
-
-def subsequent_mask(length):
-    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
-
-
-def shift_1d(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)
-
-
-def generate_path(duration, mask):
-    """
-    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 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

config.py

@@ -9,7 +9,7 @@ from typing import Dict, List
 import torch
 import yaml
 
-from common.log import logger
+from style_bert_vits2.logging import logger
 
 # If not cuda available, set possible devices to cpu
 cuda_available = torch.cuda.is_available()

data_utils.py

@@ -7,12 +7,13 @@ import torch
 import torch.utils.data
 from tqdm import tqdm
 
-import commons
 from config import config
 from mel_processing import mel_spectrogram_torch, spectrogram_torch
-from text import cleaned_text_to_sequence
+from style_bert_vits2.logging import logger
-from common.log import logger
+from style_bert_vits2.models import commons
-from utils import load_filepaths_and_text, load_wav_to_torch
+from style_bert_vits2.models.hyper_parameters import HyperParametersData
+from style_bert_vits2.models.utils import load_filepaths_and_text, load_wav_to_torch
+from style_bert_vits2.nlp import cleaned_text_to_sequence
 
 """Multi speaker version"""
 
@@ -24,7 +25,7 @@ class TextAudioSpeakerLoader(torch.utils.data.Dataset):
     3) computes spectrograms from audio files.
     """
 
-    def __init__(self, audiopaths_sid_text, hparams):
+    def __init__(self, audiopaths_sid_text: str, hparams: HyperParametersData):
         self.audiopaths_sid_text = load_filepaths_and_text(audiopaths_sid_text)
         self.max_wav_value = hparams.max_wav_value
         self.sampling_rate = hparams.sampling_rate

default_style.py

@@ -1,6 +1,6 @@
 import os
-from common.log import logger
+from style_bert_vits2.constants import DEFAULT_STYLE
-from common.constants import DEFAULT_STYLE
+from style_bert_vits2.logging import logger
 
 import numpy as np
 import json

gen_yaml.py

@@ -1,7 +1,8 @@
+import argparse
 import os
 import shutil
 import yaml
-import argparse
+
 
 parser = argparse.ArgumentParser(
     description="config.ymlの生成。あらかじめ前準備をしたデータをバッチファイルなどで連続で学習する時にtrain_ms.pyより前に使用する。"

infer.py

@@ -1,315 +0,0 @@
-import torch
-
-import commons
-from text.get_bert import get_bert
-import utils
-from models import SynthesizerTrn
-from models_jp_extra import SynthesizerTrn as SynthesizerTrnJPExtra
-from text import cleaned_text_to_sequence
-from text.cleaner import clean_text
-from text.symbols import symbols
-from common.log import logger
-
-
-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")
-        net_g = SynthesizerTrnJPExtra(
-            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(device)
-    else:
-        logger.info("Using normal model")
-        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(device)
-    net_g.state_dict()
-    _ = net_g.eval()
-    if model_path.endswith(".pth") or model_path.endswith(".pt"):
-        _ = utils.load_checkpoint(model_path, net_g, None, skip_optimizer=True)
-    elif model_path.endswith(".safetensors"):
-        _ = utils.load_safetensors(model_path, net_g, True)
-    else:
-        raise ValueError(f"Unknown model format: {model_path}")
-    return net_g
-
-
-def get_text(
-    text,
-    language_str,
-    hps,
-    device,
-    assist_text=None,
-    assist_text_weight=0.7,
-    given_tone=None,
-):
-    use_jp_extra = hps.version.endswith("JP-Extra")
-    # 推論のときにのみ呼び出されるので、raise_yomi_errorはFalseに設定
-    norm_text, phone, tone, word2ph = clean_text(
-        text, language_str, use_jp_extra, raise_yomi_error=False
-    )
-    if given_tone is not None:
-        if len(given_tone) != len(phone):
-            raise InvalidToneError(
-                f"Length of given_tone ({len(given_tone)}) != length of phone ({len(phone)})"
-            )
-        tone = given_tone
-    phone, tone, language = cleaned_text_to_sequence(phone, tone, language_str)
-
-    if hps.data.add_blank:
-        phone = commons.intersperse(phone, 0)
-        tone = commons.intersperse(tone, 0)
-        language = commons.intersperse(language, 0)
-        for i in range(len(word2ph)):
-            word2ph[i] = word2ph[i] * 2
-        word2ph[0] += 1
-    bert_ori = get_bert(
-        norm_text,
-        word2ph,
-        language_str,
-        device,
-        assist_text,
-        assist_text_weight,
-    )
-    del word2ph
-    assert bert_ori.shape[-1] == len(phone), phone
-
-    if language_str == "ZH":
-        bert = bert_ori
-        ja_bert = torch.zeros(1024, len(phone))
-        en_bert = torch.zeros(1024, len(phone))
-    elif language_str == "JP":
-        bert = torch.zeros(1024, len(phone))
-        ja_bert = bert_ori
-        en_bert = torch.zeros(1024, len(phone))
-    elif language_str == "EN":
-        bert = torch.zeros(1024, len(phone))
-        ja_bert = torch.zeros(1024, len(phone))
-        en_bert = bert_ori
-    else:
-        raise ValueError("language_str should be ZH, JP or EN")
-
-    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, en_bert, phone, tone, language
-
-
-def infer(
-    text,
-    style_vec,
-    sdp_ratio,
-    noise_scale,
-    noise_scale_w,
-    length_scale,
-    sid: int,  # In the original Bert-VITS2, its speaker_name: str, but here it's id
-    language,
-    hps,
-    net_g,
-    device,
-    skip_start=False,
-    skip_end=False,
-    assist_text=None,
-    assist_text_weight=0.7,
-    given_tone=None,
-):
-    is_jp_extra = hps.version.endswith("JP-Extra")
-    bert, ja_bert, en_bert, phones, tones, lang_ids = get_text(
-        text,
-        language,
-        hps,
-        device,
-        assist_text=assist_text,
-        assist_text_weight=assist_text_weight,
-        given_tone=given_tone,
-    )
-    if skip_start:
-        phones = phones[3:]
-        tones = tones[3:]
-        lang_ids = lang_ids[3:]
-        bert = bert[:, 3:]
-        ja_bert = ja_bert[:, 3:]
-        en_bert = en_bert[:, 3:]
-    if skip_end:
-        phones = phones[:-2]
-        tones = tones[:-2]
-        lang_ids = lang_ids[:-2]
-        bert = bert[:, :-2]
-        ja_bert = ja_bert[:, :-2]
-        en_bert = en_bert[:, :-2]
-    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)
-        bert = bert.to(device).unsqueeze(0)
-        ja_bert = ja_bert.to(device).unsqueeze(0)
-        en_bert = en_bert.to(device).unsqueeze(0)
-        x_tst_lengths = torch.LongTensor([phones.size(0)]).to(device)
-        style_vec = torch.from_numpy(style_vec).to(device).unsqueeze(0)
-        del phones
-        sid_tensor = torch.LongTensor([sid]).to(device)
-        if is_jp_extra:
-            output = net_g.infer(
-                x_tst,
-                x_tst_lengths,
-                sid_tensor,
-                tones,
-                lang_ids,
-                ja_bert,
-                style_vec=style_vec,
-                sdp_ratio=sdp_ratio,
-                noise_scale=noise_scale,
-                noise_scale_w=noise_scale_w,
-                length_scale=length_scale,
-            )
-        else:
-            output = net_g.infer(
-                x_tst,
-                x_tst_lengths,
-                sid_tensor,
-                tones,
-                lang_ids,
-                bert,
-                ja_bert,
-                en_bert,
-                style_vec=style_vec,
-                sdp_ratio=sdp_ratio,
-                noise_scale=noise_scale,
-                noise_scale_w=noise_scale_w,
-                length_scale=length_scale,
-            )
-        audio = output[0][0, 0].data.cpu().float().numpy()
-        del (
-            x_tst,
-            tones,
-            lang_ids,
-            bert,
-            x_tst_lengths,
-            sid_tensor,
-            ja_bert,
-            en_bert,
-            style_vec,
-        )  # , emo
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-        return audio
-
-
-def infer_multilang(
-    text,
-    style_vec,
-    sdp_ratio,
-    noise_scale,
-    noise_scale_w,
-    length_scale,
-    sid,
-    language,
-    hps,
-    net_g,
-    device,
-    skip_start=False,
-    skip_end=False,
-):
-    bert, ja_bert, en_bert, phones, tones, lang_ids = [], [], [], [], [], []
-    # emo = get_emo_(reference_audio, emotion, sid)
-    # if isinstance(reference_audio, np.ndarray):
-    #     emo = get_clap_audio_feature(reference_audio, device)
-    # else:
-    #     emo = get_clap_text_feature(emotion, device)
-    # emo = torch.squeeze(emo, dim=1)
-    for idx, (txt, lang) in enumerate(zip(text, language)):
-        _skip_start = (idx != 0) or (skip_start and idx == 0)
-        _skip_end = (idx != len(language) - 1) or skip_end
-        (
-            temp_bert,
-            temp_ja_bert,
-            temp_en_bert,
-            temp_phones,
-            temp_tones,
-            temp_lang_ids,
-        ) = get_text(txt, lang, hps, device)
-        if _skip_start:
-            temp_bert = temp_bert[:, 3:]
-            temp_ja_bert = temp_ja_bert[:, 3:]
-            temp_en_bert = temp_en_bert[:, 3:]
-            temp_phones = temp_phones[3:]
-            temp_tones = temp_tones[3:]
-            temp_lang_ids = temp_lang_ids[3:]
-        if _skip_end:
-            temp_bert = temp_bert[:, :-2]
-            temp_ja_bert = temp_ja_bert[:, :-2]
-            temp_en_bert = temp_en_bert[:, :-2]
-            temp_phones = temp_phones[:-2]
-            temp_tones = temp_tones[:-2]
-            temp_lang_ids = temp_lang_ids[:-2]
-        bert.append(temp_bert)
-        ja_bert.append(temp_ja_bert)
-        en_bert.append(temp_en_bert)
-        phones.append(temp_phones)
-        tones.append(temp_tones)
-        lang_ids.append(temp_lang_ids)
-    bert = torch.concatenate(bert, dim=1)
-    ja_bert = torch.concatenate(ja_bert, dim=1)
-    en_bert = torch.concatenate(en_bert, dim=1)
-    phones = torch.concatenate(phones, dim=0)
-    tones = torch.concatenate(tones, dim=0)
-    lang_ids = torch.concatenate(lang_ids, dim=0)
-    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)
-        bert = bert.to(device).unsqueeze(0)
-        ja_bert = ja_bert.to(device).unsqueeze(0)
-        en_bert = en_bert.to(device).unsqueeze(0)
-        # emo = emo.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,
-                ja_bert,
-                en_bert,
-                style_vec=style_vec,
-                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,
-            ja_bert,
-            en_bert,
-        )  # , emo
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-        return audio

initialize.py

@@ -5,7 +5,7 @@ from pathlib import Path
 import yaml
 from huggingface_hub import hf_hub_download
 
-from common.log import logger
+from style_bert_vits2.logging import logger
 
 
 def download_bert_models():

losses.py

@@ -2,8 +2,6 @@ import torch
 import torchaudio
 from transformers import AutoModel
 
-from common.log import logger
-
 
 def feature_loss(fmap_r, fmap_g):
     loss = 0

monotonic_align/__init__.py

@@ -1,16 +0,0 @@
-from numpy import zeros, int32, float32
-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)

monotonic_align/core.py

@@ -1,46 +0,0 @@
-import numba
-
-
-@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]
-
-        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

preprocess_text.py

@@ -7,10 +7,10 @@ from typing import Optional
 import click
 from tqdm import tqdm
 
-from common.log import logger
-from common.stdout_wrapper import SAFE_STDOUT
 from config import config
-from text.cleaner import clean_text
+from style_bert_vits2.logging import logger
+from style_bert_vits2.nlp import clean_text
+from style_bert_vits2.utils.stdout_wrapper import SAFE_STDOUT
 
 preprocess_text_config = config.preprocess_text_config
 
@@ -72,7 +72,7 @@ def preprocess(
                         utt, spk, language, text = line.strip().split("|")
                         norm_text, phones, tones, word2ph = clean_text(
                             text=text,
-                            language=language,
+                            language=language,  # type: ignore
                             use_jp_extra=use_jp_extra,
                             raise_yomi_error=(yomi_error != "use"),
                         )

pyproject.toml

@@ -0,0 +1,95 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "style-bert-vits2"
+dynamic = ["version"]
+description = 'Style-Bert-VITS2: Bert-VITS2 with more controllable voice styles.'
+readme = "README.md"
+requires-python = ">=3.9"
+license = "AGPL-3.0"
+keywords = []
+authors = [
+    { name = "litagin02", email = "139731664+litagin02@users.noreply.github.com" },
+]
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "Programming Language :: Python",
+    "Programming Language :: Python :: 3.9",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: Implementation :: CPython",
+]
+dependencies = [
+    'cmudict',
+    'cn2an',
+    'g2p_en',
+    'gradio',
+    'jieba',
+    'librosa==0.9.2',
+    'loguru',
+    'num2words',
+    'numba',
+    'numpy',
+    'pyannote.audio>=3.1.0',
+    'pydantic>=2.0',
+    'pyopenjtalk-dict',
+    'pypinyin',
+    'pyworld-prebuilt',
+    'safetensors',
+    'scipy',
+    'torch>=2.1,<2.2',
+    'transformers',
+]
+
+[project.urls]
+Documentation = "https://github.com/litagin02/Style-Bert-VITS2#readme"
+Issues = "https://github.com/litagin02/Style-Bert-VITS2/issues"
+Source = "https://github.com/litagin02/Style-Bert-VITS2"
+
+[tool.hatch.version]
+path = "style_bert_vits2/constants.py"
+
+[tool.hatch.envs.test]
+dependencies = [
+    "coverage[toml]>=6.5",
+    "pytest",
+]
+[tool.hatch.envs.test.scripts]
+# Usage: `hatch run test:test`
+test = "pytest {args:tests}"
+# Usage: `hatch run test:coverage`
+test-cov = "coverage run -m pytest {args:tests}"
+# Usage: `hatch run test:cov-report`
+cov-report = [
+    "- coverage combine",
+    "coverage report",
+]
+# Usage: `hatch run test:cov`
+cov = [
+    "test-cov",
+    "cov-report",
+]
+
+[[tool.hatch.envs.test.matrix]]
+python = ["3.9", "3.10", "3.11"]
+
+[tool.coverage.run]
+source_pkgs = ["style_bert_vits2", "tests"]
+branch = true
+parallel = true
+omit = [
+    "style_bert_vits2/constants.py",
+]
+
+[tool.coverage.paths]
+style_bert_vits2 = ["style_bert_vits2", "*/style-bert-vits2/style_bert_vits2"]
+tests = ["tests", "*/style-bert-vits2/tests"]
+
+[tool.coverage.report]
+exclude_lines = [
+    "no cov",
+    "if __name__ == .__main__.:",
+    "if TYPE_CHECKING:",
+]

requirements.txt

@@ -14,11 +14,12 @@ numba
 numpy
 psutil
 pyannote.audio>=3.1.0
+pydantic>=2.0
 pyloudnorm
 # pyopenjtalk-prebuilt  # Should be manually uninstalled
 pyopenjtalk-dict
 pypinyin
-# pyworld  # Not supported on Windows without Cython...
+pyworld-prebuilt
 PyYAML
 requests
 safetensors

resample.py

@@ -7,8 +7,8 @@ import pyloudnorm as pyln
 import soundfile
 from tqdm import tqdm
 
-from common.log import logger
+from style_bert_vits2.logging import logger
-from common.stdout_wrapper import SAFE_STDOUT
+from style_bert_vits2.utils.stdout_wrapper import SAFE_STDOUT
 from config import config
 
 DEFAULT_BLOCK_SIZE: float = 0.400  # seconds

server_editor.py

@@ -30,20 +30,30 @@ from fastapi.staticfiles import StaticFiles
 from pydantic import BaseModel
 from scipy.io import wavfile
 
-from common.constants import (
+from style_bert_vits2.constants import (
     DEFAULT_ASSIST_TEXT_WEIGHT,
     DEFAULT_NOISE,
     DEFAULT_NOISEW,
     DEFAULT_SDP_RATIO,
     DEFAULT_STYLE,
     DEFAULT_STYLE_WEIGHT,
-    LATEST_VERSION,
+    VERSION,
     Languages,
 )
-from common.log import logger
+from style_bert_vits2.logging import logger
-from common.tts_model import ModelHolder
+from style_bert_vits2.nlp import bert_models
-from text.japanese import g2kata_tone, kata_tone2phone_tone, text_normalize
+from style_bert_vits2.nlp.japanese import pyopenjtalk_worker as pyopenjtalk
-from text.user_dict import apply_word, delete_word, read_dict, rewrite_word, update_dict
+from style_bert_vits2.nlp.japanese.g2p_utils import g2kata_tone, kata_tone2phone_tone
+from style_bert_vits2.nlp.japanese.normalizer import normalize_text
+from style_bert_vits2.nlp.japanese.user_dict import (
+    apply_word,
+    delete_word,
+    read_dict,
+    rewrite_word,
+    update_dict,
+)
+from style_bert_vits2.tts_model import TTSModelHolder, TTSModelInfo
+
 
 # ---フロントエンド部分に関する処理---
 
@@ -139,6 +149,19 @@ def save_last_download(latest_release):
 # ---フロントエンド部分に関する処理ここまで---
 # 以降はAPIの設定
 
+# pyopenjtalk_worker を起動
+## pyopenjtalk_worker は TCP ソケットサーバーのため、ここで起動する
+pyopenjtalk.initialize_worker()
+
+# pyopenjtalk の辞書を更新
+update_dict()
+
+# 事前に BERT モデル/トークナイザーをロードしておく
+## ここでロードしなくても必要になった際に自動ロードされるが、時間がかかるため事前にロードしておいた方が体験が良い
+## server_editor.py は日本語にしか対応していないため、日本語の BERT モデル/トークナイザーのみロードする
+bert_models.load_model(Languages.JP)
+bert_models.load_tokenizer(Languages.JP)
+
 
 class AudioResponse(Response):
     media_type = "audio/wav"
@@ -175,7 +198,7 @@ if device == "cuda" and not torch.cuda.is_available():
 model_dir = Path(args.model_dir)
 port = int(args.port)
 
-model_holder = ModelHolder(model_dir, device)
+model_holder = TTSModelHolder(model_dir, device)
 if len(model_holder.model_names) == 0:
     logger.error(f"Models not found in {model_dir}.")
     sys.exit(1)
@@ -196,7 +219,7 @@ router = APIRouter()
 
 @router.get("/version")
 def version() -> str:
-    return LATEST_VERSION
+    return VERSION
 
 
 class MoraTone(BaseModel):
@@ -212,7 +235,7 @@ class TextRequest(BaseModel):
 async def read_item(item: TextRequest):
     try:
         # 最初に正規化しないと整合性がとれない
-        text = text_normalize(item.text)
+        text = normalize_text(item.text)
         kata_tone_list = g2kata_tone(text)
     except Exception as e:
         raise HTTPException(
@@ -223,11 +246,11 @@ async def read_item(item: TextRequest):
 
 
 @router.post("/normalize")
-async def normalize_text(item: TextRequest):
+async def normalize(item: TextRequest):
-    return text_normalize(item.text)
+    return normalize_text(item.text)
 
 
-@router.get("/models_info")
+@router.get("/models_info", response_model=list[TTSModelInfo])
 def models_info():
     return model_holder.models_info
 
@@ -260,7 +283,7 @@ def synthesis(request: SynthesisRequest):
             detail=f"1行の文字数は{args.line_length}文字以下にしてください。",
         )
     try:
-        model = model_holder.load_model(
+        model = model_holder.get_model(
             model_name=request.model, model_path_str=request.modelFile
         )
     except Exception as e:
@@ -284,10 +307,10 @@ def synthesis(request: SynthesisRequest):
         )
     sr, audio = model.infer(
         text=text,
-        language=request.language.value,
+        language=request.language,
         sdp_ratio=request.sdpRatio,
         noise=request.noise,
-        noisew=request.noisew,
+        noise_w=request.noisew,
         length=1 / request.speed,
         given_tone=tone,
         style=request.style,
@@ -298,7 +321,7 @@ def synthesis(request: SynthesisRequest):
         line_split=False,
         pitch_scale=request.pitchScale,
         intonation_scale=request.intonationScale,
-        sid=sid,
+        speaker_id=sid,
     )
 
     with BytesIO() as wavContent:
@@ -319,6 +342,7 @@ def multi_synthesis(request: MultiSynthesisRequest):
             detail=f"行数は{args.line_count}行以下にしてください。",
         )
     audios = []
+    sr = None
     for i, req in enumerate(lines):
         if args.line_length is not None and len(req.text) > args.line_length:
             raise HTTPException(
@@ -326,7 +350,7 @@ def multi_synthesis(request: MultiSynthesisRequest):
                 detail=f"1行の文字数は{args.line_length}文字以下にしてください。",
             )
         try:
-            model = model_holder.load_model(
+            model = model_holder.get_model(
                 model_name=req.model, model_path_str=req.modelFile
             )
         except Exception as e:
@@ -343,10 +367,10 @@ def multi_synthesis(request: MultiSynthesisRequest):
         tone = [t for _, t in phone_tone]
         sr, audio = model.infer(
             text=text,
-            language=req.language.value,
+            language=req.language,
             sdp_ratio=req.sdpRatio,
             noise=req.noise,
-            noisew=req.noisew,
+            noise_w=req.noisew,
             length=1 / req.speed,
             given_tone=tone,
             style=req.style,

server_fastapi.py

@@ -20,7 +20,8 @@ from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import FileResponse, Response
 from scipy.io import wavfile
 
-from common.constants import (
+from config import config
+from style_bert_vits2.constants import (
     DEFAULT_ASSIST_TEXT_WEIGHT,
     DEFAULT_LENGTH,
     DEFAULT_LINE_SPLIT,
@@ -32,13 +33,28 @@ from common.constants import (
     DEFAULT_STYLE_WEIGHT,
     Languages,
 )
-from common.log import logger
+from style_bert_vits2.logging import logger
-from common.tts_model import Model, ModelHolder
+from style_bert_vits2.nlp import bert_models
-from config import config
+from style_bert_vits2.nlp.japanese import pyopenjtalk_worker as pyopenjtalk
+from style_bert_vits2.tts_model import TTSModel, TTSModelHolder
 
 ln = config.server_config.language
 
 
+# pyopenjtalk_worker を起動
+## pyopenjtalk_worker は TCP ソケットサーバーのため、ここで起動する
+pyopenjtalk.initialize_worker()
+
+# 事前に BERT モデル/トークナイザーをロードしておく
+## ここでロードしなくても必要になった際に自動ロードされるが、時間がかかるため事前にロードしておいた方が体験が良い
+bert_models.load_model(Languages.JP)
+bert_models.load_tokenizer(Languages.JP)
+bert_models.load_model(Languages.EN)
+bert_models.load_tokenizer(Languages.EN)
+bert_models.load_model(Languages.ZH)
+bert_models.load_tokenizer(Languages.ZH)
+
+
 def raise_validation_error(msg: str, param: str):
     logger.warning(f"Validation error: {msg}")
     raise HTTPException(
@@ -51,16 +67,16 @@ class AudioResponse(Response):
     media_type = "audio/wav"
 
 
-def load_models(model_holder: ModelHolder):
+def load_models(model_holder: TTSModelHolder):
     model_holder.models = []
     for model_name, model_paths in model_holder.model_files_dict.items():
-        model = Model(
+        model = TTSModel(
             model_path=model_paths[0],
             config_path=model_holder.root_dir / model_name / "config.json",
             style_vec_path=model_holder.root_dir / model_name / "style_vectors.npy",
             device=model_holder.device,
         )
-        model.load_net_g()
+        model.load()
         model_holder.models.append(model)
 
 
@@ -78,7 +94,7 @@ if __name__ == "__main__":
         device = "cuda" if torch.cuda.is_available() else "cpu"
 
     model_dir = Path(args.dir)
-    model_holder = ModelHolder(model_dir, device)
+    model_holder = TTSModelHolder(model_dir, device)
     if len(model_holder.model_names) == 0:
         logger.error(f"Models not found in {model_dir}.")
         sys.exit(1)
@@ -104,7 +120,7 @@ if __name__ == "__main__":
     @app.get("/voice", response_class=AudioResponse)
     async def voice(
         request: Request,
-        text: str = Query(..., min_length=1, max_length=limit, description=f"セリフ"),
+        text: str = Query(..., min_length=1, max_length=limit, description="セリフ"),
         encoding: str = Query(None, description="textをURLデコードする(ex, `utf-8`)"),
         model_id: int = Query(
             0, description="モデルID。`GET /models/info`のkeyの値を指定ください"
@@ -132,7 +148,7 @@ if __name__ == "__main__":
             DEFAULT_LENGTH,
             description="話速。基準は1で大きくするほど音声は長くなり読み上げが遅まる",
         ),
-        language: Languages = Query(ln, description=f"textの言語"),
+        language: Languages = Query(ln, description="textの言語"),
         auto_split: bool = Query(DEFAULT_LINE_SPLIT, description="改行で分けて生成"),
         split_interval: float = Query(
             DEFAULT_SPLIT_INTERVAL, description="分けた場合に挟む無音の長さ（秒）"
@@ -178,11 +194,11 @@ if __name__ == "__main__":
         sr, audio = model.infer(
             text=text,
             language=language,
-            sid=speaker_id,
+            speaker_id=speaker_id,
             reference_audio_path=reference_audio_path,
             sdp_ratio=sdp_ratio,
             noise=noise,
-            noisew=noisew,
+            noise_w=noisew,
             length=length,
             line_split=auto_split,
             split_interval=split_interval,

slice.py

@@ -8,8 +8,8 @@ import torch
 import yaml
 from tqdm import tqdm
 
-from common.log import logger
+from style_bert_vits2.logging import logger
-from common.stdout_wrapper import SAFE_STDOUT
+from style_bert_vits2.utils.stdout_wrapper import SAFE_STDOUT
 
 vad_model, utils = torch.hub.load(
     repo_or_dir="snakers4/silero-vad",

spec_gen.py

@@ -1,87 +0,0 @@
-import torch
-from tqdm import tqdm
-from multiprocessing import Pool
-from mel_processing import spectrogram_torch, mel_spectrogram_torch
-from utils import load_wav_to_torch
-
-
-class AudioProcessor:
-    def __init__(
-        self,
-        max_wav_value,
-        use_mel_spec_posterior,
-        filter_length,
-        n_mel_channels,
-        sampling_rate,
-        hop_length,
-        win_length,
-        mel_fmin,
-        mel_fmax,
-    ):
-        self.max_wav_value = max_wav_value
-        self.use_mel_spec_posterior = use_mel_spec_posterior
-        self.filter_length = filter_length
-        self.n_mel_channels = n_mel_channels
-        self.sampling_rate = sampling_rate
-        self.hop_length = hop_length
-        self.win_length = win_length
-        self.mel_fmin = mel_fmin
-        self.mel_fmax = mel_fmax
-
-    def process_audio(self, filename):
-        audio, sampling_rate = load_wav_to_torch(filename)
-        audio_norm = audio / self.max_wav_value
-        audio_norm = audio_norm.unsqueeze(0)
-        spec_filename = filename.replace(".wav", ".spec.pt")
-        if self.use_mel_spec_posterior:
-            spec_filename = spec_filename.replace(".spec.pt", ".mel.pt")
-        try:
-            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.mel_fmin,
-                    self.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 = torch.squeeze(spec, 0)
-            torch.save(spec, spec_filename)
-        return spec, audio_norm
-
-
-# 使用示例
-processor = AudioProcessor(
-    max_wav_value=32768.0,
-    use_mel_spec_posterior=False,
-    filter_length=2048,
-    n_mel_channels=128,
-    sampling_rate=44100,
-    hop_length=512,
-    win_length=2048,
-    mel_fmin=0.0,
-    mel_fmax="null",
-)
-
-with open("filelists/train.list", "r") as f:
-    filepaths = [line.split("|")[0] for line in f]  # 取每一行的第一部分作为audiopath
-
-# 使用多进程处理
-with Pool(processes=32) as pool:  # 使用4个进程
-    with tqdm(total=len(filepaths)) as pbar:
-        for i, _ in enumerate(pool.imap_unordered(processor.process_audio, filepaths)):
-            pbar.update()

speech_mos.py

@@ -10,9 +10,9 @@ import pandas as pd
 import torch
 from tqdm import tqdm
 
-from common.log import logger
-from common.tts_model import Model
 from config import config
+from style_bert_vits2.logging import logger
+from style_bert_vits2.tts_model import TTSModel
 
 warnings.filterwarnings("ignore")
 
@@ -54,7 +54,7 @@ safetensors_files = model_path.glob("*.safetensors")
 
 
 def get_model(model_file: Path):
-    return Model(
+    return TTSModel(
         model_path=str(model_file),
         config_path=str(model_file.parent / "config.json"),
         style_vec_path=str(model_file.parent / "style_vectors.npy"),

style_bert_vits2/.editorconfig

@@ -0,0 +1,15 @@
+root = true
+
+[*]
+charset = utf-8
+end_of_line = lf
+insert_final_newline = true
+indent_size = 4
+indent_style = space
+trim_trailing_whitespace = true
+
+[*.md]
+trim_trailing_whitespace = false
+
+[*.yml]
+indent_size = 2

style_bert_vits2/constants.py

@@ -0,0 +1,46 @@
+from pathlib import Path
+
+from style_bert_vits2.utils.strenum import StrEnum
+
+
+# Style-Bert-VITS2 のバージョン
+VERSION = "2.4"
+
+# Style-Bert-VITS2 のベースディレクトリ
+BASE_DIR = Path(__file__).parent.parent
+
+# 利用可能な言語
+## JP-Extra モデル利用時は JP 以外の言語の音声合成はできない
+class Languages(StrEnum):
+    JP = "JP"
+    EN = "EN"
+    ZH = "ZH"
+
+# 言語ごとのデフォルトの BERT トークナイザーのパス
+DEFAULT_BERT_TOKENIZER_PATHS = {
+    Languages.JP: BASE_DIR / "bert" / "deberta-v2-large-japanese-char-wwm",
+    Languages.EN: BASE_DIR / "bert" / "deberta-v3-large",
+    Languages.ZH: BASE_DIR / "bert" / "chinese-roberta-wwm-ext-large",
+}
+
+# デフォルトのユーザー辞書ディレクトリ
+## style_bert_vits2.nlp.japanese.user_dict モジュールのデフォルト値として利用される
+## ライブラリとしての利用などで外部のユーザー辞書を指定したい場合は、user_dict 以下の各関数の実行時、引数に辞書データファイルのパスを指定する
+DEFAULT_USER_DICT_DIR = BASE_DIR / "dict_data"
+
+# デフォルトの推論パラメータ
+DEFAULT_STYLE = "Neutral"
+DEFAULT_STYLE_WEIGHT = 5.0
+DEFAULT_SDP_RATIO = 0.2
+DEFAULT_NOISE = 0.6
+DEFAULT_NOISEW = 0.8
+DEFAULT_LENGTH = 1.0
+DEFAULT_LINE_SPLIT = True
+DEFAULT_SPLIT_INTERVAL = 0.5
+DEFAULT_ASSIST_TEXT_WEIGHT = 0.7
+DEFAULT_ASSIST_TEXT_WEIGHT = 1.0
+
+# Gradio のテーマ
+## Built-in theme: "default", "base", "monochrome", "soft", "glass"
+## See https://huggingface.co/spaces/gradio/theme-gallery for more themes
+GRADIO_THEME = "NoCrypt/miku"

style_bert_vits2/logging.py

@@ -0,0 +1,15 @@
+from loguru import logger
+
+from style_bert_vits2.utils.stdout_wrapper import SAFE_STDOUT
+
+
+# Remove all default handlers
+logger.remove()
+
+# Add a new handler
+logger.add(
+    SAFE_STDOUT,
+    format = "<g>{time:MM-DD HH:mm:ss}</g> |<lvl>{level:^8}</lvl>| {file}:{line} | {message}",
+    backtrace = True,
+    diagnose = True,
+)

style_bert_vits2/models/attentions.py

@@ -1,14 +1,15 @@
+from typing import Any, Optional
+
 import math
 import torch
 from torch import nn
 from torch.nn import functional as F
 
-import commons
+from style_bert_vits2.models import commons
-from common.log import logger as logging
 
 
 class LayerNorm(nn.Module):
-    def __init__(self, channels, eps=1e-5):
+    def __init__(self, channels: int, eps: float = 1e-5) -> None:
         super().__init__()
         self.channels = channels
         self.eps = eps
@@ -16,14 +17,14 @@ class LayerNorm(nn.Module):
         self.gamma = nn.Parameter(torch.ones(channels))
         self.beta = nn.Parameter(torch.zeros(channels))
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         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
+@torch.jit.script  # type: ignore
-def fused_add_tanh_sigmoid_multiply(input_a, input_b, n_channels):
+def fused_add_tanh_sigmoid_multiply(input_a: torch.Tensor, input_b: torch.Tensor, n_channels: list[int]) -> torch.Tensor:
     n_channels_int = n_channels[0]
     in_act = input_a + input_b
     t_act = torch.tanh(in_act[:, :n_channels_int, :])
@@ -35,16 +36,16 @@ def fused_add_tanh_sigmoid_multiply(input_a, input_b, n_channels):
 class Encoder(nn.Module):
     def __init__(
         self,
-        hidden_channels,
+        hidden_channels: int,
-        filter_channels,
+        filter_channels: int,
-        n_heads,
+        n_heads: int,
-        n_layers,
+        n_layers: int,
-        kernel_size=1,
+        kernel_size: int = 1,
-        p_dropout=0.0,
+        p_dropout: float = 0.0,
-        window_size=4,
+        window_size: int = 4,
-        isflow=True,
+        isflow: bool = True,
-        **kwargs
+        **kwargs: Any
-    ):
+    ) -> None:
         super().__init__()
         self.hidden_channels = hidden_channels
         self.filter_channels = filter_channels
@@ -67,7 +68,7 @@ class Encoder(nn.Module):
                 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)
+                # logger.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"
@@ -98,12 +99,13 @@ class Encoder(nn.Module):
             )
             self.norm_layers_2.append(LayerNorm(hidden_channels))
 
-    def forward(self, x, x_mask, g=None):
+    def forward(self, x: torch.Tensor, x_mask: torch.Tensor, g: Optional[torch.Tensor] = None) -> torch.Tensor:
         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))
+                assert g is not None
                 g = g.transpose(1, 2)
                 x = x + g
                 x = x * x_mask
@@ -121,16 +123,16 @@ class Encoder(nn.Module):
 class Decoder(nn.Module):
     def __init__(
         self,
-        hidden_channels,
+        hidden_channels: int,
-        filter_channels,
+        filter_channels: int,
-        n_heads,
+        n_heads: int,
-        n_layers,
+        n_layers: int,
-        kernel_size=1,
+        kernel_size: int = 1,
-        p_dropout=0.0,
+        p_dropout: float = 0.0,
-        proximal_bias=False,
+        proximal_bias: bool = False,
-        proximal_init=True,
+        proximal_init: bool = True,
-        **kwargs
+        **kwargs: Any
-    ):
+    ) -> None:
         super().__init__()
         self.hidden_channels = hidden_channels
         self.filter_channels = filter_channels
@@ -178,7 +180,7 @@ class Decoder(nn.Module):
             )
             self.norm_layers_2.append(LayerNorm(hidden_channels))
 
-    def forward(self, x, x_mask, h, h_mask):
+    def forward(self, x: torch.Tensor, x_mask: torch.Tensor, h: torch.Tensor, h_mask: torch.Tensor) -> torch.Tensor:
         """
         x: decoder input
         h: encoder output
@@ -207,16 +209,16 @@ class Decoder(nn.Module):
 class MultiHeadAttention(nn.Module):
     def __init__(
         self,
-        channels,
+        channels: int,
-        out_channels,
+        out_channels: int,
-        n_heads,
+        n_heads: int,
-        p_dropout=0.0,
+        p_dropout: float = 0.0,
-        window_size=None,
+        window_size: Optional[int] = None,
-        heads_share=True,
+        heads_share: bool = True,
-        block_length=None,
+        block_length: Optional[int] = None,
-        proximal_bias=False,
+        proximal_bias: bool = False,
-        proximal_init=False,
+        proximal_init: bool = False,
-    ):
+    ) -> None:
         super().__init__()
         assert channels % n_heads == 0
 
@@ -256,9 +258,11 @@ class MultiHeadAttention(nn.Module):
         if proximal_init:
             with torch.no_grad():
                 self.conv_k.weight.copy_(self.conv_q.weight)
+                assert self.conv_k.bias is not None
+                assert self.conv_q.bias is not None
                 self.conv_k.bias.copy_(self.conv_q.bias)
 
-    def forward(self, x, c, attn_mask=None):
+    def forward(self, x: torch.Tensor, c: torch.Tensor, attn_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
         q = self.conv_q(x)
         k = self.conv_k(c)
         v = self.conv_v(c)
@@ -268,7 +272,13 @@ class MultiHeadAttention(nn.Module):
         x = self.conv_o(x)
         return x
 
-    def attention(self, query, key, value, mask=None):
+    def attention(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        mask: Optional[torch.Tensor] = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         # 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)
@@ -319,7 +329,7 @@ class MultiHeadAttention(nn.Module):
         )  # [b, n_h, t_t, d_k] -> [b, d, t_t]
         return output, p_attn
 
-    def _matmul_with_relative_values(self, x, y):
+    def _matmul_with_relative_values(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
         """
         x: [b, h, l, m]
         y: [h or 1, m, d]
@@ -328,7 +338,7 @@ class MultiHeadAttention(nn.Module):
         ret = torch.matmul(x, y.unsqueeze(0))
         return ret
 
-    def _matmul_with_relative_keys(self, x, y):
+    def _matmul_with_relative_keys(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
         """
         x: [b, h, l, d]
         y: [h or 1, m, d]
@@ -337,8 +347,9 @@ class MultiHeadAttention(nn.Module):
         ret = torch.matmul(x, y.unsqueeze(0).transpose(-2, -1))
         return ret
 
-    def _get_relative_embeddings(self, relative_embeddings, length):
+    def _get_relative_embeddings(self, relative_embeddings: torch.Tensor, length: int) -> torch.Tensor:
-        2 * self.window_size + 1
+        assert self.window_size is not None
+        2 * self.window_size + 1  # type: ignore
         # 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)
@@ -355,7 +366,7 @@ class MultiHeadAttention(nn.Module):
         ]
         return used_relative_embeddings
 
-    def _relative_position_to_absolute_position(self, x):
+    def _relative_position_to_absolute_position(self, x: torch.Tensor) -> torch.Tensor:
         """
         x: [b, h, l, 2*l-1]
         ret: [b, h, l, l]
@@ -376,7 +387,7 @@ class MultiHeadAttention(nn.Module):
         ]
         return x_final
 
-    def _absolute_position_to_relative_position(self, x):
+    def _absolute_position_to_relative_position(self, x: torch.Tensor) -> torch.Tensor:
         """
         x: [b, h, l, l]
         ret: [b, h, l, 2*l-1]
@@ -392,7 +403,7 @@ class MultiHeadAttention(nn.Module):
         x_final = x_flat.view([batch, heads, length, 2 * length])[:, :, :, 1:]
         return x_final
 
-    def _attention_bias_proximal(self, length):
+    def _attention_bias_proximal(self, length: int) -> torch.Tensor:
         """Bias for self-attention to encourage attention to close positions.
         Args:
           length: an integer scalar.
@@ -407,14 +418,14 @@ class MultiHeadAttention(nn.Module):
 class FFN(nn.Module):
     def __init__(
         self,
-        in_channels,
+        in_channels: int,
-        out_channels,
+        out_channels: int,
-        filter_channels,
+        filter_channels: int,
-        kernel_size,
+        kernel_size: int,
-        p_dropout=0.0,
+        p_dropout: float = 0.0,
-        activation=None,
+        activation: Optional[str] = None,
-        causal=False,
+        causal: bool = False,
-    ):
+    ) -> None:
         super().__init__()
         self.in_channels = in_channels
         self.out_channels = out_channels
@@ -433,7 +444,7 @@ class FFN(nn.Module):
         self.conv_2 = nn.Conv1d(filter_channels, out_channels, kernel_size)
         self.drop = nn.Dropout(p_dropout)
 
-    def forward(self, x, x_mask):
+    def forward(self, x: torch.Tensor, x_mask: torch.Tensor) -> torch.Tensor:
         x = self.conv_1(self.padding(x * x_mask))
         if self.activation == "gelu":
             x = x * torch.sigmoid(1.702 * x)
@@ -443,7 +454,7 @@ class FFN(nn.Module):
         x = self.conv_2(self.padding(x * x_mask))
         return x * x_mask
 
-    def _causal_padding(self, x):
+    def _causal_padding(self, x: torch.Tensor) -> torch.Tensor:
         if self.kernel_size == 1:
             return x
         pad_l = self.kernel_size - 1
@@ -452,7 +463,7 @@ class FFN(nn.Module):
         x = F.pad(x, commons.convert_pad_shape(padding))
         return x
 
-    def _same_padding(self, x):
+    def _same_padding(self, x: torch.Tensor) -> torch.Tensor:
         if self.kernel_size == 1:
             return x
         pad_l = (self.kernel_size - 1) // 2

style_bert_vits2/models/commons.py

@@ -0,0 +1,210 @@
+"""
+以下に記述されている関数のコメントはリファクタリング時に GPT-4 に生成させたもので、
+コードと完全に一致している保証はない。あくまで参考程度とすること。
+"""
+
+import torch
+from torch.nn import functional as F
+from typing import Any, Optional, Union
+
+
+def init_weights(m: torch.nn.Module, mean: float = 0.0, std: float = 0.01) -> None:
+    """
+    モジュールの重みを初期化する
+
+    Args:
+        m (torch.nn.Module): 重みを初期化する対象のモジュール
+        mean (float): 正規分布の平均
+        std (float): 正規分布の標準偏差
+    """
+    classname = m.__class__.__name__
+    if classname.find("Conv") != -1:
+        m.weight.data.normal_(mean, std)
+
+
+def get_padding(kernel_size: int, dilation: int = 1) -> int:
+    """
+    カーネルサイズと膨張率からパディングの大きさを計算する
+
+    Args:
+        kernel_size (int): カーネルのサイズ
+        dilation (int): 膨張率
+
+    Returns:
+        int: 計算されたパディングの大きさ
+    """
+    return int((kernel_size * dilation - dilation) / 2)
+
+
+def convert_pad_shape(pad_shape: list[list[Any]]) -> list[Any]:
+    """
+    パディングの形状を変換する
+
+    Args:
+        pad_shape (list[list[Any]]): 変換前のパディングの形状
+
+    Returns:
+        list[Any]: 変換後のパディングの形状
+    """
+    layer = pad_shape[::-1]
+    new_pad_shape = [item for sublist in layer for item in sublist]
+    return new_pad_shape
+
+
+def intersperse(lst: list[Any], item: Any) -> list[Any]:
+    """
+    リストの要素の間に特定のアイテムを挿入する
+
+    Args:
+        lst (list[Any]): 元のリスト
+        item (Any): 挿入するアイテム
+
+    Returns:
+        list[Any]: 新しいリスト
+    """
+    result = [item] * (len(lst) * 2 + 1)
+    result[1::2] = lst
+    return result
+
+
+def slice_segments(x: torch.Tensor, ids_str: torch.Tensor, segment_size: int = 4) -> torch.Tensor:
+    """
+    テンソルからセグメントをスライスする
+
+    Args:
+        x (torch.Tensor): 入力テンソル
+        ids_str (torch.Tensor): スライスを開始するインデックス
+        segment_size (int, optional): スライスのサイズ (デフォルト: 4)
+
+    Returns:
+        torch.Tensor: スライスされたセグメント
+    """
+    gather_indices = ids_str.view(x.size(0), 1, 1).repeat(
+        1, x.size(1), 1
+    ) + torch.arange(segment_size, device=x.device)
+    return torch.gather(x, 2, gather_indices)
+
+
+def rand_slice_segments(x: torch.Tensor, x_lengths: Optional[torch.Tensor] = None, segment_size: int = 4) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    ランダムなセグメントをスライスする
+
+    Args:
+        x (torch.Tensor): 入力テンソル
+        x_lengths (Optional[torch.Tensor], optional): 各バッチの長さ (デフォルト: None)
+        segment_size (int, optional): スライスのサイズ (デフォルト: 4)
+
+    Returns:
+        tuple[torch.Tensor, torch.Tensor]: スライスされたセグメントと開始インデックス
+    """
+    b, d, t = x.size()
+    if x_lengths is None:
+        x_lengths = t  # type: ignore
+    ids_str_max = torch.clamp(x_lengths - segment_size + 1, min=0)  # type: ignore
+    ids_str = (torch.rand([b], device=x.device) * ids_str_max).to(dtype=torch.long)
+    ret = slice_segments(x, ids_str, segment_size)
+    return ret, ids_str
+
+
+def subsequent_mask(length: int) -> torch.Tensor:
+    """
+    後続のマスクを生成する
+
+    Args:
+        length (int): マスクのサイズ
+
+    Returns:
+        torch.Tensor: 生成されたマスク
+    """
+    mask = torch.tril(torch.ones(length, length)).unsqueeze(0).unsqueeze(0)
+    return mask
+
+
+@torch.jit.script  # type: ignore
+def fused_add_tanh_sigmoid_multiply(input_a: torch.Tensor, input_b: torch.Tensor, n_channels: torch.Tensor) -> torch.Tensor:
+    """
+    加算、tanh、sigmoid の活性化関数を組み合わせた演算を行う
+
+    Args:
+        input_a (torch.Tensor): 入力テンソル A
+        input_b (torch.Tensor): 入力テンソル B
+        n_channels (torch.Tensor): チャネル数
+
+    Returns:
+        torch.Tensor: 演算結果
+    """
+    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 sequence_mask(length: torch.Tensor, max_length: Optional[int] = None) -> torch.Tensor:
+    """
+    シーケンスマスクを生成する
+
+    Args:
+        length (torch.Tensor): 各シーケンスの長さ
+        max_length (Optional[int]): 最大のシーケンス長さ。指定されていない場合は length の最大値を使用
+
+    Returns:
+        torch.Tensor: 生成されたシーケンスマスク
+    """
+    if max_length is None:
+        max_length = length.max()  # type: ignore
+    x = torch.arange(max_length, dtype=length.dtype, device=length.device)  # type: ignore
+    return x.unsqueeze(0) < length.unsqueeze(1)
+
+
+def generate_path(duration: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+    """
+    パスを生成する
+
+    Args:
+        duration (torch.Tensor): 各時間ステップの持続時間
+        mask (torch.Tensor): マスクテンソル
+
+    Returns:
+        torch.Tensor: 生成されたパス
+    """
+    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: Union[torch.Tensor, list[torch.Tensor]], clip_value: Optional[float], norm_type: float = 2.0) -> float:
+    """
+    勾配の値をクリップする
+
+    Args:
+        parameters (Union[torch.Tensor, list[torch.Tensor]]): クリップするパラメータ
+        clip_value (Optional[float]): クリップする値。None の場合はクリップしない
+        norm_type (float): ノルムの種類
+
+    Returns:
+        float: 総ノルム
+    """
+    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.0
+    for p in parameters:
+        assert p.grad is not None
+        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

style_bert_vits2/models/hyper_parameters.py

@@ -0,0 +1,129 @@
+"""
+Style-Bert-VITS2 モデルのハイパーパラメータを表す Pydantic モデル。
+デフォルト値は configs/configs_jp_extra.json 内の定義と概ね同一で、
+万が一ロードした config.json に存在しないキーがあった際のフェイルセーフとして適用される。
+"""
+
+from pathlib import Path
+from typing import Optional, Union
+
+from pydantic import BaseModel, ConfigDict
+
+
+class HyperParametersTrain(BaseModel):
+    log_interval: int = 200
+    eval_interval: int = 1000
+    seed: int = 42
+    epochs: int = 1000
+    learning_rate: float = 0.0001
+    betas: list[float] = [0.8, 0.99]
+    eps: float = 1e-9
+    batch_size: int = 2
+    bf16_run: bool = False
+    fp16_run: bool = False
+    lr_decay: float = 0.99996
+    segment_size: int = 16384
+    init_lr_ratio: int = 1
+    warmup_epochs: int = 0
+    c_mel: int = 45
+    c_kl: float = 1.0
+    c_commit: int = 100
+    skip_optimizer: bool = False
+    freeze_ZH_bert: bool = False
+    freeze_JP_bert: bool = False
+    freeze_EN_bert: bool = False
+    freeze_emo: bool = False
+    freeze_style: bool = False
+    freeze_decoder: bool = False
+
+
+class HyperParametersData(BaseModel):
+    use_jp_extra: bool = True
+    training_files: str = "Data/Dummy/train.list"
+    validation_files: str = "Data/Dummy/val.list"
+    max_wav_value: float = 32768.0
+    sampling_rate: int = 44100
+    filter_length: int = 2048
+    hop_length: int = 512
+    win_length: int = 2048
+    n_mel_channels: int = 128
+    mel_fmin: float = 0.0
+    mel_fmax: Optional[float] = None
+    add_blank: bool = True
+    n_speakers: int = 512
+    cleaned_text: bool = True
+    spk2id: dict[str, int] = {
+      "Dummy": 0,
+    }
+    num_styles: int = 1
+    style2id: dict[str, int] = {
+      "Neutral": 0,
+    }
+
+
+class HyperParametersModelSLM(BaseModel):
+    model: str = "./slm/wavlm-base-plus"
+    sr: int = 16000
+    hidden: int = 768
+    nlayers: int = 13
+    initial_channel: int = 64
+
+class HyperParametersModel(BaseModel):
+    use_spk_conditioned_encoder: bool = True
+    use_noise_scaled_mas: bool = True
+    use_mel_posterior_encoder: bool = False
+    use_duration_discriminator: bool = False
+    use_wavlm_discriminator: bool = True
+    inter_channels: int = 192
+    hidden_channels: int = 192
+    filter_channels: int = 768
+    n_heads: int = 2
+    n_layers: int = 6
+    kernel_size: int = 3
+    p_dropout: float = 0.1
+    resblock: str = "1"
+    resblock_kernel_sizes: list[int] = [3, 7, 11]
+    resblock_dilation_sizes: list[list[int]] = [
+        [1, 3, 5],
+        [1, 3, 5],
+        [1, 3, 5],
+    ]
+    upsample_rates: list[int] = [8, 8, 2, 2, 2]
+    upsample_initial_channel: int = 512
+    upsample_kernel_sizes: list[int] = [16, 16, 8, 2, 2]
+    n_layers_q: int = 3
+    use_spectral_norm: bool = False
+    gin_channels: int = 512
+    slm: HyperParametersModelSLM = HyperParametersModelSLM()
+
+
+class HyperParameters(BaseModel):
+    model_name: str = 'Dummy'
+    version: str = "2.0-JP-Extra"
+    train: HyperParametersTrain = HyperParametersTrain()
+    data: HyperParametersData = HyperParametersData()
+    model: HyperParametersModel = HyperParametersModel()
+
+    # 以下は学習時にのみ動的に設定されるパラメータ (通常 config.json には存在しない)
+    model_dir: Optional[str] = None
+    speedup: bool = False
+    repo_id: Optional[str] = None
+
+    # model_ 以下を Pydantic の保護対象から除外する
+    model_config = ConfigDict(protected_namespaces=())
+
+
+    @staticmethod
+    def load_from_json(json_path: Union[str, Path]) -> "HyperParameters":
+        """
+        与えられた JSON ファイルからハイパーパラメータを読み込む。
+
+        Args:
+            json_path (Union[str, Path]): JSON ファイルのパス
+
+        Returns:
+            HyperParameters: ハイパーパラメータ
+        """
+
+        with open(json_path, "r") as f:
+            return HyperParameters.model_validate_json(f.read())

style_bert_vits2/models/infer.py

@@ -0,0 +1,260 @@
+from typing import Any, cast, Optional, Union
+
+import torch
+from numpy.typing import NDArray
+
+from style_bert_vits2.constants import Languages
+from style_bert_vits2.logging import logger
+from style_bert_vits2.models import commons
+from style_bert_vits2.models import utils
+from style_bert_vits2.models.hyper_parameters import HyperParameters
+from style_bert_vits2.models.models import SynthesizerTrn
+from style_bert_vits2.models.models_jp_extra import SynthesizerTrn as SynthesizerTrnJPExtra
+from style_bert_vits2.nlp import clean_text, cleaned_text_to_sequence, extract_bert_feature
+from style_bert_vits2.nlp.symbols import SYMBOLS
+
+
+def get_net_g(model_path: str, version: str, device: str, hps: HyperParameters):
+    if version.endswith("JP-Extra"):
+        logger.info("Using JP-Extra model")
+        net_g = SynthesizerTrnJPExtra(
+            n_vocab = len(SYMBOLS),
+            spec_channels = hps.data.filter_length // 2 + 1,
+            segment_size = hps.train.segment_size // hps.data.hop_length,
+            n_speakers = hps.data.n_speakers,
+            # hps.model 以下のすべての値を引数に渡す
+            use_spk_conditioned_encoder = hps.model.use_spk_conditioned_encoder,
+            use_noise_scaled_mas = hps.model.use_noise_scaled_mas,
+            use_mel_posterior_encoder = hps.model.use_mel_posterior_encoder,
+            use_duration_discriminator = hps.model.use_duration_discriminator,
+            use_wavlm_discriminator = hps.model.use_wavlm_discriminator,
+            inter_channels = hps.model.inter_channels,
+            hidden_channels = hps.model.hidden_channels,
+            filter_channels = hps.model.filter_channels,
+            n_heads = hps.model.n_heads,
+            n_layers = hps.model.n_layers,
+            kernel_size = hps.model.kernel_size,
+            p_dropout = hps.model.p_dropout,
+            resblock = hps.model.resblock,
+            resblock_kernel_sizes = hps.model.resblock_kernel_sizes,
+            resblock_dilation_sizes = hps.model.resblock_dilation_sizes,
+            upsample_rates = hps.model.upsample_rates,
+            upsample_initial_channel = hps.model.upsample_initial_channel,
+            upsample_kernel_sizes = hps.model.upsample_kernel_sizes,
+            n_layers_q = hps.model.n_layers_q,
+            use_spectral_norm = hps.model.use_spectral_norm,
+            gin_channels = hps.model.gin_channels,
+            slm = hps.model.slm,
+        ).to(device)
+    else:
+        logger.info("Using normal model")
+        net_g = SynthesizerTrn(
+            n_vocab = len(SYMBOLS),
+            spec_channels = hps.data.filter_length // 2 + 1,
+            segment_size = hps.train.segment_size // hps.data.hop_length,
+            n_speakers=hps.data.n_speakers,
+            # hps.model 以下のすべての値を引数に渡す
+            use_spk_conditioned_encoder = hps.model.use_spk_conditioned_encoder,
+            use_noise_scaled_mas = hps.model.use_noise_scaled_mas,
+            use_mel_posterior_encoder = hps.model.use_mel_posterior_encoder,
+            use_duration_discriminator = hps.model.use_duration_discriminator,
+            use_wavlm_discriminator = hps.model.use_wavlm_discriminator,
+            inter_channels = hps.model.inter_channels,
+            hidden_channels = hps.model.hidden_channels,
+            filter_channels = hps.model.filter_channels,
+            n_heads = hps.model.n_heads,
+            n_layers = hps.model.n_layers,
+            kernel_size = hps.model.kernel_size,
+            p_dropout = hps.model.p_dropout,
+            resblock = hps.model.resblock,
+            resblock_kernel_sizes = hps.model.resblock_kernel_sizes,
+            resblock_dilation_sizes = hps.model.resblock_dilation_sizes,
+            upsample_rates = hps.model.upsample_rates,
+            upsample_initial_channel = hps.model.upsample_initial_channel,
+            upsample_kernel_sizes = hps.model.upsample_kernel_sizes,
+            n_layers_q = hps.model.n_layers_q,
+            use_spectral_norm = hps.model.use_spectral_norm,
+            gin_channels = hps.model.gin_channels,
+            slm = hps.model.slm,
+        ).to(device)
+    net_g.state_dict()
+    _ = net_g.eval()
+    if model_path.endswith(".pth") or model_path.endswith(".pt"):
+        _ = utils.checkpoints.load_checkpoint(model_path, net_g, None, skip_optimizer=True)
+    elif model_path.endswith(".safetensors"):
+        _ = utils.safetensors.load_safetensors(model_path, net_g, True)
+    else:
+        raise ValueError(f"Unknown model format: {model_path}")
+    return net_g
+
+
+def get_text(
+    text: str,
+    language_str: Languages,
+    hps: HyperParameters,
+    device: str,
+    assist_text: Optional[str] = None,
+    assist_text_weight: float = 0.7,
+    given_tone: Optional[list[int]] = None,
+):
+    use_jp_extra = hps.version.endswith("JP-Extra")
+    # 推論時のみ呼び出されるので、raise_yomi_error は False に設定
+    norm_text, phone, tone, word2ph = clean_text(
+        text,
+        language_str,
+        use_jp_extra = use_jp_extra,
+        raise_yomi_error = False,
+    )
+    if given_tone is not None:
+        if len(given_tone) != len(phone):
+            raise InvalidToneError(
+                f"Length of given_tone ({len(given_tone)}) != length of phone ({len(phone)})"
+            )
+        tone = given_tone
+    phone, tone, language = cleaned_text_to_sequence(phone, tone, language_str)
+
+    if hps.data.add_blank:
+        phone = commons.intersperse(phone, 0)
+        tone = commons.intersperse(tone, 0)
+        language = commons.intersperse(language, 0)
+        for i in range(len(word2ph)):
+            word2ph[i] = word2ph[i] * 2
+        word2ph[0] += 1
+    bert_ori = extract_bert_feature(
+        norm_text,
+        word2ph,
+        language_str,
+        device,
+        assist_text,
+        assist_text_weight,
+    )
+    del word2ph
+    assert bert_ori.shape[-1] == len(phone), phone
+
+    if language_str == Languages.ZH:
+        bert = bert_ori
+        ja_bert = torch.zeros(1024, len(phone))
+        en_bert = torch.zeros(1024, len(phone))
+    elif language_str == Languages.JP:
+        bert = torch.zeros(1024, len(phone))
+        ja_bert = bert_ori
+        en_bert = torch.zeros(1024, len(phone))
+    elif language_str == Languages.EN:
+        bert = torch.zeros(1024, len(phone))
+        ja_bert = torch.zeros(1024, len(phone))
+        en_bert = bert_ori
+    else:
+        raise ValueError("language_str should be ZH, JP or EN")
+
+    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, en_bert, phone, tone, language
+
+
+def infer(
+    text: str,
+    style_vec: NDArray[Any],
+    sdp_ratio: float,
+    noise_scale: float,
+    noise_scale_w: float,
+    length_scale: float,
+    sid: int,  # In the original Bert-VITS2, its speaker_name: str, but here it's id
+    language: Languages,
+    hps: HyperParameters,
+    net_g: Union[SynthesizerTrn, SynthesizerTrnJPExtra],
+    device: str,
+    skip_start: bool = False,
+    skip_end: bool = False,
+    assist_text: Optional[str] = None,
+    assist_text_weight: float = 0.7,
+    given_tone: Optional[list[int]] = None,
+):
+    is_jp_extra = hps.version.endswith("JP-Extra")
+    bert, ja_bert, en_bert, phones, tones, lang_ids = get_text(
+        text,
+        language,
+        hps,
+        device,
+        assist_text=assist_text,
+        assist_text_weight=assist_text_weight,
+        given_tone=given_tone,
+    )
+    if skip_start:
+        phones = phones[3:]
+        tones = tones[3:]
+        lang_ids = lang_ids[3:]
+        bert = bert[:, 3:]
+        ja_bert = ja_bert[:, 3:]
+        en_bert = en_bert[:, 3:]
+    if skip_end:
+        phones = phones[:-2]
+        tones = tones[:-2]
+        lang_ids = lang_ids[:-2]
+        bert = bert[:, :-2]
+        ja_bert = ja_bert[:, :-2]
+        en_bert = en_bert[:, :-2]
+    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)
+        bert = bert.to(device).unsqueeze(0)
+        ja_bert = ja_bert.to(device).unsqueeze(0)
+        en_bert = en_bert.to(device).unsqueeze(0)
+        x_tst_lengths = torch.LongTensor([phones.size(0)]).to(device)
+        style_vec_tensor = torch.from_numpy(style_vec).to(device).unsqueeze(0)
+        del phones
+        sid_tensor = torch.LongTensor([sid]).to(device)
+        if is_jp_extra:
+            output = cast(SynthesizerTrnJPExtra, net_g).infer(
+                x_tst,
+                x_tst_lengths,
+                sid_tensor,
+                tones,
+                lang_ids,
+                ja_bert,
+                style_vec=style_vec_tensor,
+                sdp_ratio=sdp_ratio,
+                noise_scale=noise_scale,
+                noise_scale_w=noise_scale_w,
+                length_scale=length_scale,
+            )
+        else:
+            output = cast(SynthesizerTrn, net_g).infer(
+                x_tst,
+                x_tst_lengths,
+                sid_tensor,
+                tones,
+                lang_ids,
+                bert,
+                ja_bert,
+                en_bert,
+                style_vec=style_vec_tensor,
+                sdp_ratio=sdp_ratio,
+                noise_scale=noise_scale,
+                noise_scale_w=noise_scale_w,
+                length_scale=length_scale,
+            )
+        audio = output[0][0, 0].data.cpu().float().numpy()
+        del (
+            x_tst,
+            tones,
+            lang_ids,
+            bert,
+            x_tst_lengths,
+            sid_tensor,
+            ja_bert,
+            en_bert,
+            style_vec,
+        )  # , emo
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        return audio
+
+
+class InvalidToneError(ValueError):
+    pass

style_bert_vits2/models/models.py

@@ -1,5 +1,5 @@
 import math
-import warnings
+from typing import Any, Optional
 
 import torch
 from torch import nn
@@ -7,18 +7,22 @@ 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 attentions
+from style_bert_vits2.models import attentions
-import commons
+from style_bert_vits2.models import commons
-import modules
+from style_bert_vits2.models import modules
-import monotonic_align
+from style_bert_vits2.models import monotonic_alignment
-from commons import get_padding, init_weights
+from style_bert_vits2.nlp.symbols import NUM_LANGUAGES, NUM_TONES, SYMBOLS
-from text import num_languages, num_tones, symbols
 
 
 class DurationDiscriminator(nn.Module):  # vits2
     def __init__(
-        self, in_channels, filter_channels, kernel_size, p_dropout, gin_channels=0
+        self,
-    ):
+        in_channels: int,
+        filter_channels: int,
+        kernel_size: int,
+        p_dropout: float,
+        gin_channels: int = 0
+    ) -> None:
         super().__init__()
 
         self.in_channels = in_channels
@@ -52,7 +56,13 @@ class DurationDiscriminator(nn.Module):  # vits2
 
         self.output_layer = nn.Sequential(nn.Linear(filter_channels, 1), nn.Sigmoid())
 
-    def forward_probability(self, x, x_mask, dur, g=None):
+    def forward_probability(
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        dur: torch.Tensor,
+        g: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
         dur = self.dur_proj(dur)
         x = torch.cat([x, dur], dim=1)
         x = self.pre_out_conv_1(x * x_mask)
@@ -68,7 +78,14 @@ class DurationDiscriminator(nn.Module):  # vits2
         output_prob = self.output_layer(x)
         return output_prob
 
-    def forward(self, x, x_mask, dur_r, dur_hat, g=None):
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        dur_r: torch.Tensor,
+        dur_hat: torch.Tensor,
+        g: Optional[torch.Tensor] = None,
+    ) -> list[torch.Tensor]:
         x = torch.detach(x)
         if g is not None:
             g = torch.detach(g)
@@ -93,17 +110,17 @@ class DurationDiscriminator(nn.Module):  # vits2
 class TransformerCouplingBlock(nn.Module):
     def __init__(
         self,
-        channels,
+        channels: int,
-        hidden_channels,
+        hidden_channels: int,
-        filter_channels,
+        filter_channels: int,
-        n_heads,
+        n_heads: int,
-        n_layers,
+        n_layers: int,
-        kernel_size,
+        kernel_size: int,
-        p_dropout,
+        p_dropout: float,
-        n_flows=4,
+        n_flows: int = 4,
-        gin_channels=0,
+        gin_channels: int = 0,
-        share_parameter=False,
+        share_parameter: bool = False,
-    ):
+    ) -> None:
         super().__init__()
         self.channels = channels
         self.hidden_channels = hidden_channels
@@ -115,16 +132,17 @@ class TransformerCouplingBlock(nn.Module):
         self.flows = nn.ModuleList()
 
         self.wn = (
-            attentions.FFT(
+            # attentions.FFT(
-                hidden_channels,
+            #     hidden_channels,
-                filter_channels,
+            #     filter_channels,
-                n_heads,
+            #     n_heads,
-                n_layers,
+            #     n_layers,
-                kernel_size,
+            #     kernel_size,
-                p_dropout,
+            #     p_dropout,
-                isflow=True,
+            #     isflow=True,
-                gin_channels=self.gin_channels,
+            #     gin_channels=self.gin_channels,
-            )
+            # )
+            None
             if share_parameter
             else None
         )
@@ -146,7 +164,13 @@ class TransformerCouplingBlock(nn.Module):
             )
             self.flows.append(modules.Flip())
 
-    def forward(self, x, x_mask, g=None, reverse=False):
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        g: Optional[torch.Tensor] = None,
+        reverse: bool = False,
+    ) -> torch.Tensor:
         if not reverse:
             for flow in self.flows:
                 x, _ = flow(x, x_mask, g=g, reverse=reverse)
@@ -159,13 +183,13 @@ class TransformerCouplingBlock(nn.Module):
 class StochasticDurationPredictor(nn.Module):
     def __init__(
         self,
-        in_channels,
+        in_channels: int,
-        filter_channels,
+        filter_channels: int,
-        kernel_size,
+        kernel_size: int,
-        p_dropout,
+        p_dropout: float,
-        n_flows=4,
+        n_flows: int = 4,
-        gin_channels=0,
+        gin_channels: int = 0,
-    ):
+    ) -> None:
         super().__init__()
         filter_channels = in_channels  # it needs to be removed from future version.
         self.in_channels = in_channels
@@ -205,7 +229,15 @@ class StochasticDurationPredictor(nn.Module):
         if gin_channels != 0:
             self.cond = nn.Conv1d(gin_channels, filter_channels, 1)
 
-    def forward(self, x, x_mask, w=None, g=None, reverse=False, noise_scale=1.0):
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        w: Optional[torch.Tensor] = None,
+        g: Optional[torch.Tensor] = None,
+        reverse: bool = False,
+        noise_scale: float = 1.0,
+    ) -> torch.Tensor:
         x = torch.detach(x)
         x = self.pre(x)
         if g is not None:
@@ -269,8 +301,13 @@ class StochasticDurationPredictor(nn.Module):
 
 class DurationPredictor(nn.Module):
     def __init__(
-        self, in_channels, filter_channels, kernel_size, p_dropout, gin_channels=0
+        self,
-    ):
+        in_channels: int,
+        filter_channels: int,
+        kernel_size: int,
+        p_dropout: float,
+        gin_channels: int = 0,
+    ) -> None:
         super().__init__()
 
         self.in_channels = in_channels
@@ -293,7 +330,7 @@ class DurationPredictor(nn.Module):
         if gin_channels != 0:
             self.cond = nn.Conv1d(gin_channels, in_channels, 1)
 
-    def forward(self, x, x_mask, g=None):
+    def forward(self, x: torch.Tensor, x_mask: torch.Tensor, g: Optional[torch.Tensor] = None) -> torch.Tensor:
         x = torch.detach(x)
         if g is not None:
             g = torch.detach(g)
@@ -313,17 +350,17 @@ class DurationPredictor(nn.Module):
 class TextEncoder(nn.Module):
     def __init__(
         self,
-        n_vocab,
+        n_vocab: int,
-        out_channels,
+        out_channels: int,
-        hidden_channels,
+        hidden_channels: int,
-        filter_channels,
+        filter_channels: int,
-        n_heads,
+        n_heads: int,
-        n_layers,
+        n_layers: int,
-        kernel_size,
+        kernel_size: int,
-        p_dropout,
+        p_dropout: float,
-        n_speakers,
+        n_speakers: int,
-        gin_channels=0,
+        gin_channels: int = 0,
-    ):
+    ) -> None:
         super().__init__()
         self.n_vocab = n_vocab
         self.out_channels = out_channels
@@ -334,11 +371,11 @@ class TextEncoder(nn.Module):
         self.kernel_size = kernel_size
         self.p_dropout = p_dropout
         self.gin_channels = gin_channels
-        self.emb = nn.Embedding(len(symbols), hidden_channels)
+        self.emb = nn.Embedding(len(SYMBOLS), hidden_channels)
         nn.init.normal_(self.emb.weight, 0.0, hidden_channels**-0.5)
-        self.tone_emb = nn.Embedding(num_tones, hidden_channels)
+        self.tone_emb = nn.Embedding(NUM_TONES, hidden_channels)
         nn.init.normal_(self.tone_emb.weight, 0.0, hidden_channels**-0.5)
-        self.language_emb = nn.Embedding(num_languages, hidden_channels)
+        self.language_emb = nn.Embedding(NUM_LANGUAGES, hidden_channels)
         nn.init.normal_(self.language_emb.weight, 0.0, hidden_channels**-0.5)
         self.bert_proj = nn.Conv1d(1024, hidden_channels, 1)
         self.ja_bert_proj = nn.Conv1d(1024, hidden_channels, 1)
@@ -358,17 +395,17 @@ class TextEncoder(nn.Module):
 
     def forward(
         self,
-        x,
+        x: torch.Tensor,
-        x_lengths,
+        x_lengths: torch.Tensor,
-        tone,
+        tone: torch.Tensor,
-        language,
+        language: torch.Tensor,
-        bert,
+        bert: torch.Tensor,
-        ja_bert,
+        ja_bert: torch.Tensor,
-        en_bert,
+        en_bert: torch.Tensor,
-        style_vec,
+        style_vec: torch.Tensor,
-        sid,
+        sid: torch.Tensor,
-        g=None,
+        g: Optional[torch.Tensor] = None,
-    ):
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
         bert_emb = self.bert_proj(bert).transpose(1, 2)
         ja_bert_emb = self.ja_bert_proj(ja_bert).transpose(1, 2)
         en_bert_emb = self.en_bert_proj(en_bert).transpose(1, 2)
@@ -400,14 +437,14 @@ class TextEncoder(nn.Module):
 class ResidualCouplingBlock(nn.Module):
     def __init__(
         self,
-        channels,
+        channels: int,
-        hidden_channels,
+        hidden_channels: int,
-        kernel_size,
+        kernel_size: int,
-        dilation_rate,
+        dilation_rate: int,
-        n_layers,
+        n_layers: int,
-        n_flows=4,
+        n_flows: int = 4,
-        gin_channels=0,
+        gin_channels: int = 0,
-    ):
+    ) -> None:
         super().__init__()
         self.channels = channels
         self.hidden_channels = hidden_channels
@@ -432,7 +469,13 @@ class ResidualCouplingBlock(nn.Module):
             )
             self.flows.append(modules.Flip())
 
-    def forward(self, x, x_mask, g=None, reverse=False):
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        g: Optional[torch.Tensor] = None,
+        reverse: bool = False,
+    ) -> torch.Tensor:
         if not reverse:
             for flow in self.flows:
                 x, _ = flow(x, x_mask, g=g, reverse=reverse)
@@ -445,14 +488,14 @@ class ResidualCouplingBlock(nn.Module):
 class PosteriorEncoder(nn.Module):
     def __init__(
         self,
-        in_channels,
+        in_channels: int,
-        out_channels,
+        out_channels: int,
-        hidden_channels,
+        hidden_channels: int,
-        kernel_size,
+        kernel_size: int,
-        dilation_rate,
+        dilation_rate: int,
-        n_layers,
+        n_layers: int,
-        gin_channels=0,
+        gin_channels: int = 0,
-    ):
+    ) -> None:
         super().__init__()
         self.in_channels = in_channels
         self.out_channels = out_channels
@@ -472,7 +515,12 @@ class PosteriorEncoder(nn.Module):
         )
         self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
 
-    def forward(self, x, x_lengths, g=None):
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_lengths: torch.Tensor,
+        g: Optional[torch.Tensor] = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
         x_mask = torch.unsqueeze(commons.sequence_mask(x_lengths, x.size(2)), 1).to(
             x.dtype
         )
@@ -487,22 +535,22 @@ class PosteriorEncoder(nn.Module):
 class Generator(torch.nn.Module):
     def __init__(
         self,
-        initial_channel,
+        initial_channel: int,
-        resblock,
+        resblock_str: str,
-        resblock_kernel_sizes,
+        resblock_kernel_sizes: list[int],
-        resblock_dilation_sizes,
+        resblock_dilation_sizes: list[list[int]],
-        upsample_rates,
+        upsample_rates: list[int],
-        upsample_initial_channel,
+        upsample_initial_channel: int,
-        upsample_kernel_sizes,
+        upsample_kernel_sizes: list[int],
-        gin_channels=0,
+        gin_channels: int = 0,
-    ):
+    ) -> None:
         super(Generator, self).__init__()
         self.num_kernels = len(resblock_kernel_sizes)
         self.num_upsamples = len(upsample_rates)
         self.conv_pre = Conv1d(
             initial_channel, upsample_initial_channel, 7, 1, padding=3
         )
-        resblock = modules.ResBlock1 if resblock == "1" else modules.ResBlock2
+        resblock = modules.ResBlock1 if resblock_str == "1" else modules.ResBlock2
 
         self.ups = nn.ModuleList()
         for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
@@ -519,20 +567,22 @@ class Generator(torch.nn.Module):
             )
 
         self.resblocks = nn.ModuleList()
+        ch = None
         for i in range(len(self.ups)):
             ch = upsample_initial_channel // (2 ** (i + 1))
             for j, (k, d) in enumerate(
                 zip(resblock_kernel_sizes, resblock_dilation_sizes)
             ):
-                self.resblocks.append(resblock(ch, k, d))
+                self.resblocks.append(resblock(ch, k, d))  # type: ignore
 
+        assert ch is not None
         self.conv_post = Conv1d(ch, 1, 7, 1, padding=3, bias=False)
-        self.ups.apply(init_weights)
+        self.ups.apply(commons.init_weights)
 
         if gin_channels != 0:
             self.cond = nn.Conv1d(gin_channels, upsample_initial_channel, 1)
 
-    def forward(self, x, g=None):
+    def forward(self, x: torch.Tensor, g: Optional[torch.Tensor] = None) -> torch.Tensor:
         x = self.conv_pre(x)
         if g is not None:
             x = x + self.cond(g)
@@ -546,6 +596,7 @@ class Generator(torch.nn.Module):
                     xs = self.resblocks[i * self.num_kernels + j](x)
                 else:
                     xs += self.resblocks[i * self.num_kernels + j](x)
+            assert xs is not None
             x = xs / self.num_kernels
         x = F.leaky_relu(x)
         x = self.conv_post(x)
@@ -553,7 +604,7 @@ class Generator(torch.nn.Module):
 
         return x
 
-    def remove_weight_norm(self):
+    def remove_weight_norm(self) -> None:
         print("Removing weight norm...")
         for layer in self.ups:
             remove_weight_norm(layer)
@@ -562,7 +613,7 @@ class Generator(torch.nn.Module):
 
 
 class DiscriminatorP(torch.nn.Module):
-    def __init__(self, period, kernel_size=5, stride=3, use_spectral_norm=False):
+    def __init__(self, period: int, kernel_size: int = 5, stride: int = 3, use_spectral_norm: bool = False) -> None:
         super(DiscriminatorP, self).__init__()
         self.period = period
         self.use_spectral_norm = use_spectral_norm
@@ -575,7 +626,7 @@ class DiscriminatorP(torch.nn.Module):
                         32,
                         (kernel_size, 1),
                         (stride, 1),
-                        padding=(get_padding(kernel_size, 1), 0),
+                        padding=(commons.get_padding(kernel_size, 1), 0),
                     )
                 ),
                 norm_f(
@@ -584,7 +635,7 @@ class DiscriminatorP(torch.nn.Module):
                         128,
                         (kernel_size, 1),
                         (stride, 1),
-                        padding=(get_padding(kernel_size, 1), 0),
+                        padding=(commons.get_padding(kernel_size, 1), 0),
                     )
                 ),
                 norm_f(
@@ -593,7 +644,7 @@ class DiscriminatorP(torch.nn.Module):
                         512,
                         (kernel_size, 1),
                         (stride, 1),
-                        padding=(get_padding(kernel_size, 1), 0),
+                        padding=(commons.get_padding(kernel_size, 1), 0),
                     )
                 ),
                 norm_f(
@@ -602,7 +653,7 @@ class DiscriminatorP(torch.nn.Module):
                         1024,
                         (kernel_size, 1),
                         (stride, 1),
-                        padding=(get_padding(kernel_size, 1), 0),
+                        padding=(commons.get_padding(kernel_size, 1), 0),
                     )
                 ),
                 norm_f(
@@ -611,14 +662,14 @@ class DiscriminatorP(torch.nn.Module):
                         1024,
                         (kernel_size, 1),
                         1,
-                        padding=(get_padding(kernel_size, 1), 0),
+                        padding=(commons.get_padding(kernel_size, 1), 0),
                     )
                 ),
             ]
         )
         self.conv_post = norm_f(Conv2d(1024, 1, (3, 1), 1, padding=(1, 0)))
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, list[torch.Tensor]]:
         fmap = []
 
         # 1d to 2d
@@ -641,7 +692,7 @@ class DiscriminatorP(torch.nn.Module):
 
 
 class DiscriminatorS(torch.nn.Module):
-    def __init__(self, use_spectral_norm=False):
+    def __init__(self, use_spectral_norm: bool = False) -> None:
         super(DiscriminatorS, self).__init__()
         norm_f = weight_norm if use_spectral_norm is False else spectral_norm
         self.convs = nn.ModuleList(
@@ -656,7 +707,7 @@ class DiscriminatorS(torch.nn.Module):
         )
         self.conv_post = norm_f(Conv1d(1024, 1, 3, 1, padding=1))
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, list[torch.Tensor]]:
         fmap = []
 
         for layer in self.convs:
@@ -671,7 +722,7 @@ class DiscriminatorS(torch.nn.Module):
 
 
 class MultiPeriodDiscriminator(torch.nn.Module):
-    def __init__(self, use_spectral_norm=False):
+    def __init__(self, use_spectral_norm: bool = False) -> None:
         super(MultiPeriodDiscriminator, self).__init__()
         periods = [2, 3, 5, 7, 11]
 
@@ -681,7 +732,11 @@ class MultiPeriodDiscriminator(torch.nn.Module):
         ]
         self.discriminators = nn.ModuleList(discs)
 
-    def forward(self, y, y_hat):
+    def forward(
+        self,
+        y: torch.Tensor,
+        y_hat: torch.Tensor,
+    ) -> tuple[list[torch.Tensor], list[torch.Tensor], list[torch.Tensor], list[torch.Tensor]]:
         y_d_rs = []
         y_d_gs = []
         fmap_rs = []
@@ -703,7 +758,7 @@ class ReferenceEncoder(nn.Module):
     outputs --- [N, ref_enc_gru_size]
     """
 
-    def __init__(self, spec_channels, gin_channels=0):
+    def __init__(self, spec_channels: int, gin_channels: int = 0) -> None:
         super().__init__()
         self.spec_channels = spec_channels
         ref_enc_filters = [32, 32, 64, 64, 128, 128]
@@ -732,7 +787,7 @@ class ReferenceEncoder(nn.Module):
         )
         self.proj = nn.Linear(128, gin_channels)
 
-    def forward(self, inputs, mask=None):
+    def forward(self, inputs: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor:
         N = inputs.size(0)
         out = inputs.view(N, 1, -1, self.spec_channels)  # [N, 1, Ty, n_freqs]
         for conv in self.convs:
@@ -750,7 +805,7 @@ class ReferenceEncoder(nn.Module):
 
         return self.proj(out.squeeze(0))
 
-    def calculate_channels(self, L, kernel_size, stride, pad, n_convs):
+    def calculate_channels(self, L: int, kernel_size: int, stride: int, pad: int, n_convs: int) -> int:
         for i in range(n_convs):
             L = (L - kernel_size + 2 * pad) // stride + 1
         return L
@@ -763,31 +818,31 @@ class SynthesizerTrn(nn.Module):
 
     def __init__(
         self,
-        n_vocab,
+        n_vocab: int,
-        spec_channels,
+        spec_channels: int,
-        segment_size,
+        segment_size: int,
-        inter_channels,
+        inter_channels: int,
-        hidden_channels,
+        hidden_channels: int,
-        filter_channels,
+        filter_channels: int,
-        n_heads,
+        n_heads: int,
-        n_layers,
+        n_layers: int,
-        kernel_size,
+        kernel_size: int,
-        p_dropout,
+        p_dropout: float,
-        resblock,
+        resblock: str,
-        resblock_kernel_sizes,
+        resblock_kernel_sizes: list[int],
-        resblock_dilation_sizes,
+        resblock_dilation_sizes: list[list[int]],
-        upsample_rates,
+        upsample_rates: list[int],
-        upsample_initial_channel,
+        upsample_initial_channel: int,
-        upsample_kernel_sizes,
+        upsample_kernel_sizes: list[int],
-        n_speakers=256,
+        n_speakers: int = 256,
-        gin_channels=256,
+        gin_channels: int = 256,
-        use_sdp=True,
+        use_sdp: bool = True,
-        n_flow_layer=4,
+        n_flow_layer: int = 4,
-        n_layers_trans_flow=4,
+        n_layers_trans_flow: int = 4,
-        flow_share_parameter=False,
+        flow_share_parameter: bool = False,
-        use_transformer_flow=True,
+        use_transformer_flow: bool = True,
-        **kwargs,
+        **kwargs: Any,
-    ):
+    ) -> None:
         super().__init__()
         self.n_vocab = n_vocab
         self.spec_channels = spec_channels
@@ -885,18 +940,27 @@ class SynthesizerTrn(nn.Module):
 
     def forward(
         self,
-        x,
+        x: torch.Tensor,
-        x_lengths,
+        x_lengths: torch.Tensor,
-        y,
+        y: torch.Tensor,
-        y_lengths,
+        y_lengths: torch.Tensor,
-        sid,
+        sid: torch.Tensor,
-        tone,
+        tone: torch.Tensor,
-        language,
+        language: torch.Tensor,
-        bert,
+        bert: torch.Tensor,
-        ja_bert,
+        ja_bert: torch.Tensor,
-        en_bert,
+        en_bert: torch.Tensor,
-        style_vec,
+        style_vec: torch.Tensor,
-    ):
+    ) -> tuple[
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        tuple[torch.Tensor, ...],
+        tuple[torch.Tensor, ...],
+    ]:
         if self.n_speakers > 0:
             g = self.emb_g(sid).unsqueeze(-1)  # [b, h, 1]
         else:
@@ -933,7 +997,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()
             )
@@ -974,27 +1038,28 @@ class SynthesizerTrn(nn.Module):
 
     def infer(
         self,
-        x,
+        x: torch.Tensor,
-        x_lengths,
+        x_lengths: torch.Tensor,
-        sid,
+        sid: torch.Tensor,
-        tone,
+        tone: torch.Tensor,
-        language,
+        language: torch.Tensor,
-        bert,
+        bert: torch.Tensor,
-        ja_bert,
+        ja_bert: torch.Tensor,
-        en_bert,
+        en_bert: torch.Tensor,
-        style_vec,
+        style_vec: torch.Tensor,
-        noise_scale=0.667,
+        noise_scale: float = 0.667,
-        length_scale=1,
+        length_scale: float = 1.0,
-        noise_scale_w=0.8,
+        noise_scale_w: float = 0.8,
-        max_len=None,
+        max_len: Optional[int] = None,
-        sdp_ratio=0,
+        sdp_ratio: float = 0.0,
-        y=None,
+        y: Optional[torch.Tensor] = None,
-    ):
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, tuple[torch.Tensor, ...]]:
         # x, m_p, logs_p, x_mask = self.enc_p(x, x_lengths, tone, language, bert)
         # g = self.gst(y)
         if self.n_speakers > 0:
             g = self.emb_g(sid).unsqueeze(-1)  # [b, h, 1]
         else:
+            assert y is not None
             g = self.ref_enc(y.transpose(1, 2)).unsqueeze(-1)
         x, m_p, logs_p, x_mask = self.enc_p(
             x, x_lengths, tone, language, bert, ja_bert, en_bert, style_vec, sid, g=g

style_bert_vits2/models/models_jp_extra.py

@@ -1,24 +1,28 @@
 import math
+from typing import Any, Optional
+
 import torch
 from torch import nn
+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 commons
+from style_bert_vits2.models import attentions
-import modules
+from style_bert_vits2.models import commons
-import attentions
+from style_bert_vits2.models import modules
-import monotonic_align
+from style_bert_vits2.models import monotonic_alignment
-
+from style_bert_vits2.nlp.symbols import NUM_LANGUAGES, NUM_TONES, SYMBOLS
-from torch.nn import Conv1d, ConvTranspose1d, Conv2d
-from torch.nn.utils import weight_norm, remove_weight_norm, spectral_norm
-
-from commons import init_weights, get_padding
-from text import symbols, num_tones, num_languages
 
 
 class DurationDiscriminator(nn.Module):  # vits2
     def __init__(
-        self, in_channels, filter_channels, kernel_size, p_dropout, gin_channels=0
+        self,
-    ):
+        in_channels: int,
+        filter_channels: int,
+        kernel_size: int,
+        p_dropout: float,
+        gin_channels: int = 0
+    ) -> None:
         super().__init__()
 
         self.in_channels = in_channels
@@ -49,7 +53,7 @@ class DurationDiscriminator(nn.Module):  # vits2
             nn.Linear(2 * filter_channels, 1), nn.Sigmoid()
         )
 
-    def forward_probability(self, x, dur):
+    def forward_probability(self, x: torch.Tensor, dur: torch.Tensor) -> torch.Tensor:
         dur = self.dur_proj(dur)
         x = torch.cat([x, dur], dim=1)
         x = x.transpose(1, 2)
@@ -57,7 +61,14 @@ class DurationDiscriminator(nn.Module):  # vits2
         output_prob = self.output_layer(x)
         return output_prob
 
-    def forward(self, x, x_mask, dur_r, dur_hat, g=None):
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        dur_r: torch.Tensor,
+        dur_hat: torch.Tensor,
+        g: Optional[torch.Tensor] = None,
+    ) -> list[torch.Tensor]:
         x = torch.detach(x)
         if g is not None:
             g = torch.detach(g)
@@ -82,17 +93,17 @@ class DurationDiscriminator(nn.Module):  # vits2
 class TransformerCouplingBlock(nn.Module):
     def __init__(
         self,
-        channels,
+        channels: int,
-        hidden_channels,
+        hidden_channels: int,
-        filter_channels,
+        filter_channels: int,
-        n_heads,
+        n_heads: int,
-        n_layers,
+        n_layers: int,
-        kernel_size,
+        kernel_size: int,
-        p_dropout,
+        p_dropout: float,
-        n_flows=4,
+        n_flows: int = 4,
-        gin_channels=0,
+        gin_channels: int = 0,
-        share_parameter=False,
+        share_parameter: bool = False,
-    ):
+    ) -> None:
         super().__init__()
         self.channels = channels
         self.hidden_channels = hidden_channels
@@ -104,16 +115,17 @@ class TransformerCouplingBlock(nn.Module):
         self.flows = nn.ModuleList()
 
         self.wn = (
-            attentions.FFT(
+            # attentions.FFT(
-                hidden_channels,
+            #     hidden_channels,
-                filter_channels,
+            #     filter_channels,
-                n_heads,
+            #     n_heads,
-                n_layers,
+            #     n_layers,
-                kernel_size,
+            #     kernel_size,
-                p_dropout,
+            #     p_dropout,
-                isflow=True,
+            #     isflow=True,
-                gin_channels=self.gin_channels,
+            #     gin_channels=self.gin_channels,
-            )
+            # )
+            None
             if share_parameter
             else None
         )
@@ -135,7 +147,13 @@ class TransformerCouplingBlock(nn.Module):
             )
             self.flows.append(modules.Flip())
 
-    def forward(self, x, x_mask, g=None, reverse=False):
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        g: Optional[torch.Tensor] = None,
+        reverse: bool = False,
+    ) -> torch.Tensor:
         if not reverse:
             for flow in self.flows:
                 x, _ = flow(x, x_mask, g=g, reverse=reverse)
@@ -148,13 +166,13 @@ class TransformerCouplingBlock(nn.Module):
 class StochasticDurationPredictor(nn.Module):
     def __init__(
         self,
-        in_channels,
+        in_channels: int,
-        filter_channels,
+        filter_channels: int,
-        kernel_size,
+        kernel_size: int,
-        p_dropout,
+        p_dropout: float,
-        n_flows=4,
+        n_flows: int = 4,
-        gin_channels=0,
+        gin_channels: int = 0,
-    ):
+    ) -> None:
         super().__init__()
         filter_channels = in_channels  # it needs to be removed from future version.
         self.in_channels = in_channels
@@ -194,7 +212,15 @@ class StochasticDurationPredictor(nn.Module):
         if gin_channels != 0:
             self.cond = nn.Conv1d(gin_channels, filter_channels, 1)
 
-    def forward(self, x, x_mask, w=None, g=None, reverse=False, noise_scale=1.0):
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        w: Optional[torch.Tensor] = None,
+        g: Optional[torch.Tensor] = None,
+        reverse: bool = False,
+        noise_scale: float = 1.0,
+    ) -> torch.Tensor:
         x = torch.detach(x)
         x = self.pre(x)
         if g is not None:
@@ -258,8 +284,13 @@ class StochasticDurationPredictor(nn.Module):
 
 class DurationPredictor(nn.Module):
     def __init__(
-        self, in_channels, filter_channels, kernel_size, p_dropout, gin_channels=0
+        self,
-    ):
+        in_channels: int,
+        filter_channels: int,
+        kernel_size: int,
+        p_dropout: float,
+        gin_channels: int = 0,
+    ) -> None:
         super().__init__()
 
         self.in_channels = in_channels
@@ -282,7 +313,7 @@ class DurationPredictor(nn.Module):
         if gin_channels != 0:
             self.cond = nn.Conv1d(gin_channels, in_channels, 1)
 
-    def forward(self, x, x_mask, g=None):
+    def forward(self, x: torch.Tensor, x_mask: torch.Tensor, g: Optional[torch.Tensor] = None) -> torch.Tensor:
         x = torch.detach(x)
         if g is not None:
             g = torch.detach(g)
@@ -300,14 +331,14 @@ class DurationPredictor(nn.Module):
 
 
 class Bottleneck(nn.Sequential):
-    def __init__(self, in_dim, hidden_dim):
+    def __init__(self, in_dim: int, hidden_dim: int) -> None:
         c_fc1 = nn.Linear(in_dim, hidden_dim, bias=False)
         c_fc2 = nn.Linear(in_dim, hidden_dim, bias=False)
-        super().__init__(*[c_fc1, c_fc2])
+        super().__init__(c_fc1, c_fc2)
 
 
 class Block(nn.Module):
-    def __init__(self, in_dim, hidden_dim) -> None:
+    def __init__(self, in_dim: int, hidden_dim: int) -> None:
         super().__init__()
         self.norm = nn.LayerNorm(in_dim)
         self.mlp = MLP(in_dim, hidden_dim)
@@ -318,13 +349,13 @@ class Block(nn.Module):
 
 
 class MLP(nn.Module):
-    def __init__(self, in_dim, hidden_dim):
+    def __init__(self, in_dim: int, hidden_dim: int) -> None:
         super().__init__()
         self.c_fc1 = nn.Linear(in_dim, hidden_dim, bias=False)
         self.c_fc2 = nn.Linear(in_dim, hidden_dim, bias=False)
         self.c_proj = nn.Linear(hidden_dim, in_dim, bias=False)
 
-    def forward(self, x: torch.Tensor):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         x = F.silu(self.c_fc1(x)) * self.c_fc2(x)
         x = self.c_proj(x)
         return x
@@ -333,16 +364,16 @@ class MLP(nn.Module):
 class TextEncoder(nn.Module):
     def __init__(
         self,
-        n_vocab,
+        n_vocab: int,
-        out_channels,
+        out_channels: int,
-        hidden_channels,
+        hidden_channels: int,
-        filter_channels,
+        filter_channels: int,
-        n_heads,
+        n_heads: int,
-        n_layers,
+        n_layers: int,
-        kernel_size,
+        kernel_size: int,
-        p_dropout,
+        p_dropout: float,
-        gin_channels=0,
+        gin_channels: int = 0,
-    ):
+    ) -> None:
         super().__init__()
         self.n_vocab = n_vocab
         self.out_channels = out_channels
@@ -353,11 +384,11 @@ class TextEncoder(nn.Module):
         self.kernel_size = kernel_size
         self.p_dropout = p_dropout
         self.gin_channels = gin_channels
-        self.emb = nn.Embedding(len(symbols), hidden_channels)
+        self.emb = nn.Embedding(len(SYMBOLS), hidden_channels)
         nn.init.normal_(self.emb.weight, 0.0, hidden_channels**-0.5)
-        self.tone_emb = nn.Embedding(num_tones, hidden_channels)
+        self.tone_emb = nn.Embedding(NUM_TONES, hidden_channels)
         nn.init.normal_(self.tone_emb.weight, 0.0, hidden_channels**-0.5)
-        self.language_emb = nn.Embedding(num_languages, hidden_channels)
+        self.language_emb = nn.Embedding(NUM_LANGUAGES, hidden_channels)
         nn.init.normal_(self.language_emb.weight, 0.0, hidden_channels**-0.5)
         self.bert_proj = nn.Conv1d(1024, hidden_channels, 1)
 
@@ -375,7 +406,16 @@ class TextEncoder(nn.Module):
         )
         self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
 
-    def forward(self, x, x_lengths, tone, language, bert, style_vec, g=None):
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_lengths: torch.Tensor,
+        tone: torch.Tensor,
+        language: torch.Tensor,
+        bert: torch.Tensor,
+        style_vec: torch.Tensor,
+        g: Optional[torch.Tensor] = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
         bert_emb = self.bert_proj(bert).transpose(1, 2)
         style_emb = self.style_proj(style_vec.unsqueeze(1))
         x = (
@@ -402,14 +442,14 @@ class TextEncoder(nn.Module):
 class ResidualCouplingBlock(nn.Module):
     def __init__(
         self,
-        channels,
+        channels: int,
-        hidden_channels,
+        hidden_channels: int,
-        kernel_size,
+        kernel_size: int,
-        dilation_rate,
+        dilation_rate: int,
-        n_layers,
+        n_layers: int,
-        n_flows=4,
+        n_flows: int = 4,
-        gin_channels=0,
+        gin_channels: int = 0,
-    ):
+    ) -> None:
         super().__init__()
         self.channels = channels
         self.hidden_channels = hidden_channels
@@ -434,7 +474,13 @@ class ResidualCouplingBlock(nn.Module):
             )
             self.flows.append(modules.Flip())
 
-    def forward(self, x, x_mask, g=None, reverse=False):
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        g: Optional[torch.Tensor] = None,
+        reverse: bool = False,
+    ) -> torch.Tensor:
         if not reverse:
             for flow in self.flows:
                 x, _ = flow(x, x_mask, g=g, reverse=reverse)
@@ -447,14 +493,14 @@ class ResidualCouplingBlock(nn.Module):
 class PosteriorEncoder(nn.Module):
     def __init__(
         self,
-        in_channels,
+        in_channels: int,
-        out_channels,
+        out_channels: int,
-        hidden_channels,
+        hidden_channels: int,
-        kernel_size,
+        kernel_size: int,
-        dilation_rate,
+        dilation_rate: int,
-        n_layers,
+        n_layers: int,
-        gin_channels=0,
+        gin_channels: int = 0,
-    ):
+    ) -> None:
         super().__init__()
         self.in_channels = in_channels
         self.out_channels = out_channels
@@ -474,7 +520,12 @@ class PosteriorEncoder(nn.Module):
         )
         self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
 
-    def forward(self, x, x_lengths, g=None):
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_lengths: torch.Tensor,
+        g: Optional[torch.Tensor] = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
         x_mask = torch.unsqueeze(commons.sequence_mask(x_lengths, x.size(2)), 1).to(
             x.dtype
         )
@@ -489,22 +540,22 @@ class PosteriorEncoder(nn.Module):
 class Generator(torch.nn.Module):
     def __init__(
         self,
-        initial_channel,
+        initial_channel: int,
-        resblock,
+        resblock_str: str,
-        resblock_kernel_sizes,
+        resblock_kernel_sizes: list[int],
-        resblock_dilation_sizes,
+        resblock_dilation_sizes: list[list[int]],
-        upsample_rates,
+        upsample_rates: list[int],
-        upsample_initial_channel,
+        upsample_initial_channel: int,
-        upsample_kernel_sizes,
+        upsample_kernel_sizes: list[int],
-        gin_channels=0,
+        gin_channels: int = 0,
-    ):
+    ) -> None:
         super(Generator, self).__init__()
         self.num_kernels = len(resblock_kernel_sizes)
         self.num_upsamples = len(upsample_rates)
         self.conv_pre = Conv1d(
             initial_channel, upsample_initial_channel, 7, 1, padding=3
         )
-        resblock = modules.ResBlock1 if resblock == "1" else modules.ResBlock2
+        resblock = modules.ResBlock1 if resblock_str == "1" else modules.ResBlock2
 
         self.ups = nn.ModuleList()
         for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
@@ -521,20 +572,22 @@ class Generator(torch.nn.Module):
             )
 
         self.resblocks = nn.ModuleList()
+        ch = None
         for i in range(len(self.ups)):
             ch = upsample_initial_channel // (2 ** (i + 1))
             for j, (k, d) in enumerate(
                 zip(resblock_kernel_sizes, resblock_dilation_sizes)
             ):
-                self.resblocks.append(resblock(ch, k, d))
+                self.resblocks.append(resblock(ch, k, d))  # type: ignore
 
+        assert ch is not None
         self.conv_post = Conv1d(ch, 1, 7, 1, padding=3, bias=False)
-        self.ups.apply(init_weights)
+        self.ups.apply(commons.init_weights)
 
         if gin_channels != 0:
             self.cond = nn.Conv1d(gin_channels, upsample_initial_channel, 1)
 
-    def forward(self, x, g=None):
+    def forward(self, x: torch.Tensor, g: Optional[torch.Tensor] = None) -> torch.Tensor:
         x = self.conv_pre(x)
         if g is not None:
             x = x + self.cond(g)
@@ -548,6 +601,7 @@ class Generator(torch.nn.Module):
                     xs = self.resblocks[i * self.num_kernels + j](x)
                 else:
                     xs += self.resblocks[i * self.num_kernels + j](x)
+            assert xs is not None
             x = xs / self.num_kernels
         x = F.leaky_relu(x)
         x = self.conv_post(x)
@@ -555,7 +609,7 @@ class Generator(torch.nn.Module):
 
         return x
 
-    def remove_weight_norm(self):
+    def remove_weight_norm(self) -> None:
         print("Removing weight norm...")
         for layer in self.ups:
             remove_weight_norm(layer)
@@ -564,7 +618,7 @@ class Generator(torch.nn.Module):
 
 
 class DiscriminatorP(torch.nn.Module):
-    def __init__(self, period, kernel_size=5, stride=3, use_spectral_norm=False):
+    def __init__(self, period: int, kernel_size: int = 5, stride: int = 3, use_spectral_norm: bool = False) -> None:
         super(DiscriminatorP, self).__init__()
         self.period = period
         self.use_spectral_norm = use_spectral_norm
@@ -577,7 +631,7 @@ class DiscriminatorP(torch.nn.Module):
                         32,
                         (kernel_size, 1),
                         (stride, 1),
-                        padding=(get_padding(kernel_size, 1), 0),
+                        padding=(commons.get_padding(kernel_size, 1), 0),
                     )
                 ),
                 norm_f(
@@ -586,7 +640,7 @@ class DiscriminatorP(torch.nn.Module):
                         128,
                         (kernel_size, 1),
                         (stride, 1),
-                        padding=(get_padding(kernel_size, 1), 0),
+                        padding=(commons.get_padding(kernel_size, 1), 0),
                     )
                 ),
                 norm_f(
@@ -595,7 +649,7 @@ class DiscriminatorP(torch.nn.Module):
                         512,
                         (kernel_size, 1),
                         (stride, 1),
-                        padding=(get_padding(kernel_size, 1), 0),
+                        padding=(commons.get_padding(kernel_size, 1), 0),
                     )
                 ),
                 norm_f(
@@ -604,7 +658,7 @@ class DiscriminatorP(torch.nn.Module):
                         1024,
                         (kernel_size, 1),
                         (stride, 1),
-                        padding=(get_padding(kernel_size, 1), 0),
+                        padding=(commons.get_padding(kernel_size, 1), 0),
                     )
                 ),
                 norm_f(
@@ -613,14 +667,14 @@ class DiscriminatorP(torch.nn.Module):
                         1024,
                         (kernel_size, 1),
                         1,
-                        padding=(get_padding(kernel_size, 1), 0),
+                        padding=(commons.get_padding(kernel_size, 1), 0),
                     )
                 ),
             ]
         )
         self.conv_post = norm_f(Conv2d(1024, 1, (3, 1), 1, padding=(1, 0)))
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, list[torch.Tensor]]:
         fmap = []
 
         # 1d to 2d
@@ -643,7 +697,7 @@ class DiscriminatorP(torch.nn.Module):
 
 
 class DiscriminatorS(torch.nn.Module):
-    def __init__(self, use_spectral_norm=False):
+    def __init__(self, use_spectral_norm: bool = False) -> None:
         super(DiscriminatorS, self).__init__()
         norm_f = weight_norm if use_spectral_norm is False else spectral_norm
         self.convs = nn.ModuleList(
@@ -658,7 +712,7 @@ class DiscriminatorS(torch.nn.Module):
         )
         self.conv_post = norm_f(Conv1d(1024, 1, 3, 1, padding=1))
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, list[torch.Tensor]]:
         fmap = []
 
         for layer in self.convs:
@@ -673,7 +727,7 @@ class DiscriminatorS(torch.nn.Module):
 
 
 class MultiPeriodDiscriminator(torch.nn.Module):
-    def __init__(self, use_spectral_norm=False):
+    def __init__(self, use_spectral_norm: bool = False) -> None:
         super(MultiPeriodDiscriminator, self).__init__()
         periods = [2, 3, 5, 7, 11]
 
@@ -683,7 +737,11 @@ class MultiPeriodDiscriminator(torch.nn.Module):
         ]
         self.discriminators = nn.ModuleList(discs)
 
-    def forward(self, y, y_hat):
+    def forward(
+        self,
+        y: torch.Tensor,
+        y_hat: torch.Tensor,
+    ) -> tuple[list[torch.Tensor], list[torch.Tensor], list[torch.Tensor], list[torch.Tensor]]:
         y_d_rs = []
         y_d_gs = []
         fmap_rs = []
@@ -703,8 +761,12 @@ class WavLMDiscriminator(nn.Module):
     """docstring for Discriminator."""
 
     def __init__(
-        self, slm_hidden=768, slm_layers=13, initial_channel=64, use_spectral_norm=False
+        self,
-    ):
+        slm_hidden: int = 768,
+        slm_layers: int = 13,
+        initial_channel: int = 64,
+        use_spectral_norm: bool = False,
+    ) -> None:
         super(WavLMDiscriminator, self).__init__()
         norm_f = weight_norm if use_spectral_norm == False else spectral_norm
         self.pre = norm_f(
@@ -734,7 +796,7 @@ class WavLMDiscriminator(nn.Module):
 
         self.conv_post = norm_f(Conv1d(initial_channel * 4, 1, 3, 1, padding=1))
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         x = self.pre(x)
 
         fmap = []
@@ -754,7 +816,7 @@ class ReferenceEncoder(nn.Module):
     outputs --- [N, ref_enc_gru_size]
     """
 
-    def __init__(self, spec_channels, gin_channels=0):
+    def __init__(self, spec_channels: int, gin_channels: int = 0) -> None:
         super().__init__()
         self.spec_channels = spec_channels
         ref_enc_filters = [32, 32, 64, 64, 128, 128]
@@ -783,7 +845,7 @@ class ReferenceEncoder(nn.Module):
         )
         self.proj = nn.Linear(128, gin_channels)
 
-    def forward(self, inputs, mask=None):
+    def forward(self, inputs: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor:
         N = inputs.size(0)
         out = inputs.view(N, 1, -1, self.spec_channels)  # [N, 1, Ty, n_freqs]
         for conv in self.convs:
@@ -801,7 +863,7 @@ class ReferenceEncoder(nn.Module):
 
         return self.proj(out.squeeze(0))
 
-    def calculate_channels(self, L, kernel_size, stride, pad, n_convs):
+    def calculate_channels(self, L: int, kernel_size: int, stride: int, pad: int, n_convs: int) -> int:
         for i in range(n_convs):
             L = (L - kernel_size + 2 * pad) // stride + 1
         return L
@@ -814,31 +876,31 @@ class SynthesizerTrn(nn.Module):
 
     def __init__(
         self,
-        n_vocab,
+        n_vocab: int,
-        spec_channels,
+        spec_channels: int,
-        segment_size,
+        segment_size: int,
-        inter_channels,
+        inter_channels: int,
-        hidden_channels,
+        hidden_channels: int,
-        filter_channels,
+        filter_channels: int,
-        n_heads,
+        n_heads: int,
-        n_layers,
+        n_layers: int,
-        kernel_size,
+        kernel_size: int,
-        p_dropout,
+        p_dropout: float,
-        resblock,
+        resblock: str,
-        resblock_kernel_sizes,
+        resblock_kernel_sizes: list[int],
-        resblock_dilation_sizes,
+        resblock_dilation_sizes: list[list[int]],
-        upsample_rates,
+        upsample_rates: list[int],
-        upsample_initial_channel,
+        upsample_initial_channel: int,
-        upsample_kernel_sizes,
+        upsample_kernel_sizes: list[int],
-        n_speakers=256,
+        n_speakers: int = 256,
-        gin_channels=256,
+        gin_channels: int = 256,
-        use_sdp=True,
+        use_sdp: bool = True,
-        n_flow_layer=4,
+        n_flow_layer: int = 4,
-        n_layers_trans_flow=6,
+        n_layers_trans_flow: int = 6,
-        flow_share_parameter=False,
+        flow_share_parameter: bool = False,
-        use_transformer_flow=True,
+        use_transformer_flow: bool = True,
-        **kwargs
+        **kwargs: Any,
-    ):
+    ) -> None:
         super().__init__()
         self.n_vocab = n_vocab
         self.spec_channels = spec_channels
@@ -935,16 +997,26 @@ class SynthesizerTrn(nn.Module):
 
     def forward(
         self,
-        x,
+        x: torch.Tensor,
-        x_lengths,
+        x_lengths: torch.Tensor,
-        y,
+        y: torch.Tensor,
-        y_lengths,
+        y_lengths: torch.Tensor,
-        sid,
+        sid: torch.Tensor,
-        tone,
+        tone: torch.Tensor,
-        language,
+        language: torch.Tensor,
-        bert,
+        bert: torch.Tensor,
-        style_vec,
+        style_vec: torch.Tensor,
-    ):
+    ) -> tuple[
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        tuple[torch.Tensor, ...],
+        tuple[torch.Tensor, ...],
+    ]:
         if self.n_speakers > 0:
             g = self.emb_g(sid).unsqueeze(-1)  # [b, h, 1]
         else:
@@ -981,7 +1053,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()
             )
@@ -1016,32 +1088,33 @@ class SynthesizerTrn(nn.Module):
             ids_slice,
             x_mask,
             y_mask,
-            (z, z_p, m_p, logs_p, m_q, logs_q),
+            (z, z_p, m_p, logs_p, m_q, logs_q),  # type: ignore
             (x, logw, logw_),  # , logw_sdp),
             g,
         )
 
     def infer(
         self,
-        x,
+        x: torch.Tensor,
-        x_lengths,
+        x_lengths: torch.Tensor,
-        sid,
+        sid: torch.Tensor,
-        tone,
+        tone: torch.Tensor,
-        language,
+        language: torch.Tensor,
-        bert,
+        bert: torch.Tensor,
-        style_vec,
+        style_vec: torch.Tensor,
-        noise_scale=0.667,
+        noise_scale: float = 0.667,
-        length_scale=1,
+        length_scale: float = 1.0,
-        noise_scale_w=0.8,
+        noise_scale_w: float = 0.8,
-        max_len=None,
+        max_len: Optional[int] = None,
-        sdp_ratio=0,
+        sdp_ratio: float = 0.0,
-        y=None,
+        y: Optional[torch.Tensor] = None,
-    ):
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, tuple[torch.Tensor, ...]]:
         # x, m_p, logs_p, x_mask = self.enc_p(x, x_lengths, tone, language, bert)
         # g = self.gst(y)
         if self.n_speakers > 0:
             g = self.emb_g(sid).unsqueeze(-1)  # [b, h, 1]
         else:
+            assert y is not None
             g = self.ref_enc(y.transpose(1, 2)).unsqueeze(-1)
         x, m_p, logs_p, x_mask = self.enc_p(
             x, x_lengths, tone, language, bert, style_vec, g=g

style_bert_vits2/models/modules.py

@@ -1,5 +1,5 @@
 import math
-import warnings
+from typing import Any, Optional, Union
 
 import torch
 from torch import nn
@@ -7,16 +7,16 @@ from torch.nn import Conv1d
 from torch.nn import functional as F
 from torch.nn.utils import remove_weight_norm, weight_norm
 
-import commons
+from style_bert_vits2.models import commons
-from attentions import Encoder
+from style_bert_vits2.models.attentions import Encoder
-from commons import get_padding, init_weights
+from style_bert_vits2.models.transforms import piecewise_rational_quadratic_transform
-from transforms import piecewise_rational_quadratic_transform
+
 
 LRELU_SLOPE = 0.1
 
 
 class LayerNorm(nn.Module):
-    def __init__(self, channels, eps=1e-5):
+    def __init__(self, channels: int, eps: float = 1e-5) -> None:
         super().__init__()
         self.channels = channels
         self.eps = eps
@@ -24,7 +24,7 @@ class LayerNorm(nn.Module):
         self.gamma = nn.Parameter(torch.ones(channels))
         self.beta = nn.Parameter(torch.zeros(channels))
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         x = x.transpose(1, -1)
         x = F.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps)
         return x.transpose(1, -1)
@@ -33,13 +33,13 @@ class LayerNorm(nn.Module):
 class ConvReluNorm(nn.Module):
     def __init__(
         self,
-        in_channels,
+        in_channels: int,
-        hidden_channels,
+        hidden_channels: int,
-        out_channels,
+        out_channels: int,
-        kernel_size,
+        kernel_size: int,
-        n_layers,
+        n_layers: int,
-        p_dropout,
+        p_dropout: float,
-    ):
+    ) -> None:
         super().__init__()
         self.in_channels = in_channels
         self.hidden_channels = hidden_channels
@@ -70,9 +70,10 @@ class ConvReluNorm(nn.Module):
             self.norm_layers.append(LayerNorm(hidden_channels))
         self.proj = nn.Conv1d(hidden_channels, out_channels, 1)
         self.proj.weight.data.zero_()
+        assert self.proj.bias is not None
         self.proj.bias.data.zero_()
 
-    def forward(self, x, x_mask):
+    def forward(self, x: torch.Tensor, x_mask: torch.Tensor) -> torch.Tensor:
         x_org = x
         for i in range(self.n_layers):
             x = self.conv_layers[i](x * x_mask)
@@ -87,7 +88,7 @@ class DDSConv(nn.Module):
     Dialted and Depth-Separable Convolution
     """
 
-    def __init__(self, channels, kernel_size, n_layers, p_dropout=0.0):
+    def __init__(self, channels: int, kernel_size: int, n_layers: int, p_dropout: float = 0.0) -> None:
         super().__init__()
         self.channels = channels
         self.kernel_size = kernel_size
@@ -116,7 +117,7 @@ class DDSConv(nn.Module):
             self.norms_1.append(LayerNorm(channels))
             self.norms_2.append(LayerNorm(channels))
 
-    def forward(self, x, x_mask, g=None):
+    def forward(self, x: torch.Tensor, x_mask: torch.Tensor, g: Optional[torch.Tensor] = None) -> torch.Tensor:
         if g is not None:
             x = x + g
         for i in range(self.n_layers):
@@ -134,13 +135,13 @@ class DDSConv(nn.Module):
 class WN(torch.nn.Module):
     def __init__(
         self,
-        hidden_channels,
+        hidden_channels: int,
-        kernel_size,
+        kernel_size: int,
-        dilation_rate,
+        dilation_rate: int,
-        n_layers,
+        n_layers: int,
-        gin_channels=0,
+        gin_channels: int = 0,
-        p_dropout=0,
+        p_dropout: float = 0,
-    ):
+    ) -> None:
         super(WN, self).__init__()
         assert kernel_size % 2 == 1
         self.hidden_channels = hidden_channels
@@ -183,7 +184,7 @@ class WN(torch.nn.Module):
             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):
+    def forward(self, x: torch.Tensor, x_mask: torch.Tensor, g: Optional[torch.Tensor] = None, **kwargs: Any) -> torch.Tensor:
         output = torch.zeros_like(x)
         n_channels_tensor = torch.IntTensor([self.hidden_channels])
 
@@ -210,7 +211,7 @@ class WN(torch.nn.Module):
                 output = output + res_skip_acts
         return output * x_mask
 
-    def remove_weight_norm(self):
+    def remove_weight_norm(self) -> None:
         if self.gin_channels != 0:
             torch.nn.utils.remove_weight_norm(self.cond_layer)
         for l in self.in_layers:
@@ -220,7 +221,7 @@ class WN(torch.nn.Module):
 
 
 class ResBlock1(torch.nn.Module):
-    def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5)):
+    def __init__(self, channels: int, kernel_size: int = 3, dilation: tuple[int, int, int] = (1, 3, 5)) -> None:
         super(ResBlock1, self).__init__()
         self.convs1 = nn.ModuleList(
             [
@@ -231,7 +232,7 @@ class ResBlock1(torch.nn.Module):
                         kernel_size,
                         1,
                         dilation=dilation[0],
-                        padding=get_padding(kernel_size, dilation[0]),
+                        padding=commons.get_padding(kernel_size, dilation[0]),
                     )
                 ),
                 weight_norm(
@@ -241,7 +242,7 @@ class ResBlock1(torch.nn.Module):
                         kernel_size,
                         1,
                         dilation=dilation[1],
-                        padding=get_padding(kernel_size, dilation[1]),
+                        padding=commons.get_padding(kernel_size, dilation[1]),
                     )
                 ),
                 weight_norm(
@@ -251,12 +252,12 @@ class ResBlock1(torch.nn.Module):
                         kernel_size,
                         1,
                         dilation=dilation[2],
-                        padding=get_padding(kernel_size, dilation[2]),
+                        padding=commons.get_padding(kernel_size, dilation[2]),
                     )
                 ),
             ]
         )
-        self.convs1.apply(init_weights)
+        self.convs1.apply(commons.init_weights)
 
         self.convs2 = nn.ModuleList(
             [
@@ -267,7 +268,7 @@ class ResBlock1(torch.nn.Module):
                         kernel_size,
                         1,
                         dilation=1,
-                        padding=get_padding(kernel_size, 1),
+                        padding=commons.get_padding(kernel_size, 1),
                     )
                 ),
                 weight_norm(
@@ -277,7 +278,7 @@ class ResBlock1(torch.nn.Module):
                         kernel_size,
                         1,
                         dilation=1,
-                        padding=get_padding(kernel_size, 1),
+                        padding=commons.get_padding(kernel_size, 1),
                     )
                 ),
                 weight_norm(
@@ -287,14 +288,14 @@ class ResBlock1(torch.nn.Module):
                         kernel_size,
                         1,
                         dilation=1,
-                        padding=get_padding(kernel_size, 1),
+                        padding=commons.get_padding(kernel_size, 1),
                     )
                 ),
             ]
         )
-        self.convs2.apply(init_weights)
+        self.convs2.apply(commons.init_weights)
 
-    def forward(self, x, x_mask=None):
+    def forward(self, x: torch.Tensor, x_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
         for c1, c2 in zip(self.convs1, self.convs2):
             xt = F.leaky_relu(x, LRELU_SLOPE)
             if x_mask is not None:
@@ -309,7 +310,7 @@ class ResBlock1(torch.nn.Module):
             x = x * x_mask
         return x
 
-    def remove_weight_norm(self):
+    def remove_weight_norm(self) -> None:
         for l in self.convs1:
             remove_weight_norm(l)
         for l in self.convs2:
@@ -317,7 +318,7 @@ class ResBlock1(torch.nn.Module):
 
 
 class ResBlock2(torch.nn.Module):
-    def __init__(self, channels, kernel_size=3, dilation=(1, 3)):
+    def __init__(self, channels: int, kernel_size: int = 3, dilation: tuple[int, int] = (1, 3)) -> None:
         super(ResBlock2, self).__init__()
         self.convs = nn.ModuleList(
             [
@@ -328,7 +329,7 @@ class ResBlock2(torch.nn.Module):
                         kernel_size,
                         1,
                         dilation=dilation[0],
-                        padding=get_padding(kernel_size, dilation[0]),
+                        padding=commons.get_padding(kernel_size, dilation[0]),
                     )
                 ),
                 weight_norm(
@@ -338,14 +339,14 @@ class ResBlock2(torch.nn.Module):
                         kernel_size,
                         1,
                         dilation=dilation[1],
-                        padding=get_padding(kernel_size, dilation[1]),
+                        padding=commons.get_padding(kernel_size, dilation[1]),
                     )
                 ),
             ]
         )
-        self.convs.apply(init_weights)
+        self.convs.apply(commons.init_weights)
 
-    def forward(self, x, x_mask=None):
+    def forward(self, x: torch.Tensor, x_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
         for c in self.convs:
             xt = F.leaky_relu(x, LRELU_SLOPE)
             if x_mask is not None:
@@ -356,13 +357,19 @@ class ResBlock2(torch.nn.Module):
             x = x * x_mask
         return x
 
-    def remove_weight_norm(self):
+    def remove_weight_norm(self) -> None:
         for l in self.convs:
             remove_weight_norm(l)
 
 
 class Log(nn.Module):
-    def forward(self, x, x_mask, reverse=False, **kwargs):
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        reverse: bool = False,
+        **kwargs: Any,
+    ) -> Union[tuple[torch.Tensor, torch.Tensor], torch.Tensor]:
         if not reverse:
             y = torch.log(torch.clamp_min(x, 1e-5)) * x_mask
             logdet = torch.sum(-y, [1, 2])
@@ -373,7 +380,13 @@ class Log(nn.Module):
 
 
 class Flip(nn.Module):
-    def forward(self, x, *args, reverse=False, **kwargs):
+    def forward(
+        self,
+        x: torch.Tensor,
+        *args: Any,
+        reverse: bool = False,
+        **kwargs: Any,
+    ) -> Union[tuple[torch.Tensor, torch.Tensor], torch.Tensor]:
         x = torch.flip(x, [1])
         if not reverse:
             logdet = torch.zeros(x.size(0)).to(dtype=x.dtype, device=x.device)
@@ -383,13 +396,19 @@ class Flip(nn.Module):
 
 
 class ElementwiseAffine(nn.Module):
-    def __init__(self, channels):
+    def __init__(self, channels: int) -> None:
         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):
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        reverse: bool = False,
+        **kwargs: Any,
+    ) -> Union[tuple[torch.Tensor, torch.Tensor], torch.Tensor]:
         if not reverse:
             y = self.m + torch.exp(self.logs) * x
             y = y * x_mask
@@ -403,15 +422,15 @@ class ElementwiseAffine(nn.Module):
 class ResidualCouplingLayer(nn.Module):
     def __init__(
         self,
-        channels,
+        channels: int,
-        hidden_channels,
+        hidden_channels: int,
-        kernel_size,
+        kernel_size: int,
-        dilation_rate,
+        dilation_rate: int,
-        n_layers,
+        n_layers: int,
-        p_dropout=0,
+        p_dropout: float = 0,
-        gin_channels=0,
+        gin_channels: int = 0,
-        mean_only=False,
+        mean_only: bool = False,
-    ):
+    ) -> None:
         assert channels % 2 == 0, "channels should be divisible by 2"
         super().__init__()
         self.channels = channels
@@ -433,9 +452,16 @@ class ResidualCouplingLayer(nn.Module):
         )
         self.post = nn.Conv1d(hidden_channels, self.half_channels * (2 - mean_only), 1)
         self.post.weight.data.zero_()
+        assert self.post.bias is not None
         self.post.bias.data.zero_()
 
-    def forward(self, x, x_mask, g=None, reverse=False):
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        g: Optional[torch.Tensor] = None,
+        reverse: bool = False,
+    ) -> Union[tuple[torch.Tensor, torch.Tensor], torch.Tensor]:
         x0, x1 = torch.split(x, [self.half_channels] * 2, 1)
         h = self.pre(x0) * x_mask
         h = self.enc(h, x_mask, g=g)
@@ -460,13 +486,13 @@ class ResidualCouplingLayer(nn.Module):
 class ConvFlow(nn.Module):
     def __init__(
         self,
-        in_channels,
+        in_channels: int,
-        filter_channels,
+        filter_channels: int,
-        kernel_size,
+        kernel_size: int,
-        n_layers,
+        n_layers: int,
-        num_bins=10,
+        num_bins: int = 10,
-        tail_bound=5.0,
+        tail_bound: float = 5.0,
-    ):
+    ) -> None:
         super().__init__()
         self.in_channels = in_channels
         self.filter_channels = filter_channels
@@ -482,9 +508,16 @@ class ConvFlow(nn.Module):
             filter_channels, self.half_channels * (num_bins * 3 - 1), 1
         )
         self.proj.weight.data.zero_()
+        assert self.proj.bias is not None
         self.proj.bias.data.zero_()
 
-    def forward(self, x, x_mask, g=None, reverse=False):
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        g: Optional[torch.Tensor] = None,
+        reverse: bool = False,
+    ) -> Union[tuple[torch.Tensor, torch.Tensor], torch.Tensor]:
         x0, x1 = torch.split(x, [self.half_channels] * 2, 1)
         h = self.pre(x0)
         h = self.convs(h, x_mask, g=g)
@@ -520,17 +553,17 @@ class ConvFlow(nn.Module):
 class TransformerCouplingLayer(nn.Module):
     def __init__(
         self,
-        channels,
+        channels: int,
-        hidden_channels,
+        hidden_channels: int,
-        kernel_size,
+        kernel_size: int,
-        n_layers,
+        n_layers: int,
-        n_heads,
+        n_heads: int,
-        p_dropout=0,
+        p_dropout: float = 0,
-        filter_channels=0,
+        filter_channels: int = 0,
-        mean_only=False,
+        mean_only: bool = False,
-        wn_sharing_parameter=None,
+        wn_sharing_parameter: Optional[nn.Module] = None,
-        gin_channels=0,
+        gin_channels: int = 0,
-    ):
+    ) -> None:
         assert channels % 2 == 0, "channels should be divisible by 2"
         super().__init__()
         self.channels = channels
@@ -557,9 +590,16 @@ class TransformerCouplingLayer(nn.Module):
         )
         self.post = nn.Conv1d(hidden_channels, self.half_channels * (2 - mean_only), 1)
         self.post.weight.data.zero_()
+        assert self.post.bias is not None
         self.post.bias.data.zero_()
 
-    def forward(self, x, x_mask, g=None, reverse=False):
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        g: Optional[torch.Tensor] = None,
+        reverse: bool = False,
+    ) -> Union[tuple[torch.Tensor, torch.Tensor], torch.Tensor]:
         x0, x1 = torch.split(x, [self.half_channels] * 2, 1)
         h = self.pre(x0) * x_mask
         h = self.enc(h, x_mask, g=g)

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

style_bert_vits2/models/transforms.py

@@ -1,7 +1,8 @@
-import torch
+from typing import Optional
-from torch.nn import functional as F
 
 import numpy as np
+import torch
+from torch.nn import functional as F
 
 
 DEFAULT_MIN_BIN_WIDTH = 1e-3
@@ -10,17 +11,18 @@ DEFAULT_MIN_DERIVATIVE = 1e-3
 
 
 def piecewise_rational_quadratic_transform(
-    inputs,
+    inputs: torch.Tensor,
-    unnormalized_widths,
+    unnormalized_widths: torch.Tensor,
-    unnormalized_heights,
+    unnormalized_heights: torch.Tensor,
-    unnormalized_derivatives,
+    unnormalized_derivatives: torch.Tensor,
-    inverse=False,
+    inverse: bool = False,
-    tails=None,
+    tails: Optional[str] = None,
-    tail_bound=1.0,
+    tail_bound: float = 1.0,
-    min_bin_width=DEFAULT_MIN_BIN_WIDTH,
+    min_bin_width: float = DEFAULT_MIN_BIN_WIDTH,
-    min_bin_height=DEFAULT_MIN_BIN_HEIGHT,
+    min_bin_height: float = DEFAULT_MIN_BIN_HEIGHT,
-    min_derivative=DEFAULT_MIN_DERIVATIVE,
+    min_derivative: float = DEFAULT_MIN_DERIVATIVE,
-):
+) -> tuple[torch.Tensor, torch.Tensor]:
+
     if tails is None:
         spline_fn = rational_quadratic_spline
         spline_kwargs = {}
@@ -37,28 +39,29 @@ def piecewise_rational_quadratic_transform(
         min_bin_width=min_bin_width,
         min_bin_height=min_bin_height,
         min_derivative=min_derivative,
-        **spline_kwargs
+        **spline_kwargs  # type: ignore
     )
     return outputs, logabsdet
 
 
-def searchsorted(bin_locations, inputs, eps=1e-6):
+def searchsorted(bin_locations: torch.Tensor, inputs: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
     bin_locations[..., -1] += eps
     return torch.sum(inputs[..., None] >= bin_locations, dim=-1) - 1
 
 
 def unconstrained_rational_quadratic_spline(
-    inputs,
+    inputs: torch.Tensor,
-    unnormalized_widths,
+    unnormalized_widths: torch.Tensor,
-    unnormalized_heights,
+    unnormalized_heights: torch.Tensor,
-    unnormalized_derivatives,
+    unnormalized_derivatives: torch.Tensor,
-    inverse=False,
+    inverse: bool = False,
-    tails="linear",
+    tails: str = "linear",
-    tail_bound=1.0,
+    tail_bound: float = 1.0,
-    min_bin_width=DEFAULT_MIN_BIN_WIDTH,
+    min_bin_width: float = DEFAULT_MIN_BIN_WIDTH,
-    min_bin_height=DEFAULT_MIN_BIN_HEIGHT,
+    min_bin_height: float = DEFAULT_MIN_BIN_HEIGHT,
-    min_derivative=DEFAULT_MIN_DERIVATIVE,
+    min_derivative: float = DEFAULT_MIN_DERIVATIVE,
-):
+) -> tuple[torch.Tensor, torch.Tensor]:
+
     inside_interval_mask = (inputs >= -tail_bound) & (inputs <= tail_bound)
     outside_interval_mask = ~inside_interval_mask
 
@@ -74,7 +77,7 @@ def unconstrained_rational_quadratic_spline(
         outputs[outside_interval_mask] = inputs[outside_interval_mask]
         logabsdet[outside_interval_mask] = 0
     else:
-        raise RuntimeError("{} tails are not implemented.".format(tails))
+        raise RuntimeError(f"{tails} tails are not implemented.")
 
     (
         outputs[inside_interval_mask],
@@ -98,19 +101,20 @@ def unconstrained_rational_quadratic_spline(
 
 
 def rational_quadratic_spline(
-    inputs,
+    inputs: torch.Tensor,
-    unnormalized_widths,
+    unnormalized_widths: torch.Tensor,
-    unnormalized_heights,
+    unnormalized_heights: torch.Tensor,
-    unnormalized_derivatives,
+    unnormalized_derivatives: torch.Tensor,
-    inverse=False,
+    inverse: bool = False,
-    left=0.0,
+    left: float = 0.0,
-    right=1.0,
+    right: float = 1.0,
-    bottom=0.0,
+    bottom: float = 0.0,
-    top=1.0,
+    top: float = 1.0,
-    min_bin_width=DEFAULT_MIN_BIN_WIDTH,
+    min_bin_width: float = DEFAULT_MIN_BIN_WIDTH,
-    min_bin_height=DEFAULT_MIN_BIN_HEIGHT,
+    min_bin_height: float = DEFAULT_MIN_BIN_HEIGHT,
-    min_derivative=DEFAULT_MIN_DERIVATIVE,
+    min_derivative: float = DEFAULT_MIN_DERIVATIVE,
-):
+) -> tuple[torch.Tensor, torch.Tensor]:
+
     if torch.min(inputs) < left or torch.max(inputs) > right:
         raise ValueError("Input to a transform is not within its domain")
 

style_bert_vits2/models/utils/__init__.py

@@ -0,0 +1,253 @@
+import glob
+import logging
+import os
+import re
+import subprocess
+from pathlib import Path
+from typing import Any, Optional, Union, TYPE_CHECKING
+
+import numpy as np
+import torch
+from numpy.typing import NDArray
+from scipy.io.wavfile import read
+
+from style_bert_vits2.logging import logger
+from style_bert_vits2.models.utils import checkpoints  # type: ignore
+from style_bert_vits2.models.utils import safetensors  # type: ignore
+
+if TYPE_CHECKING:
+    # tensorboard はライブラリとしてインストールされている場合は依存関係に含まれないため、型チェック時のみインポートする
+    from torch.utils.tensorboard import SummaryWriter
+
+
+__is_matplotlib_imported = False
+
+
+def summarize(
+    writer: "SummaryWriter",
+    global_step: int,
+    scalars: dict[str, float] = {},
+    histograms: dict[str, Any] = {},
+    images: dict[str, Any] = {},
+    audios: dict[str, Any] = {},
+    audio_sampling_rate: int = 22050,
+) -> None:
+    """
+    指定されたデータを TensorBoard にまとめて追加する
+
+    Args:
+        writer (SummaryWriter): TensorBoard への書き込みを行うオブジェクト
+        global_step (int): グローバルステップ数
+        scalars (dict[str, float]): スカラー値の辞書
+        histograms (dict[str, Any]): ヒストグラムの辞書
+        images (dict[str, Any]): 画像データの辞書
+        audios (dict[str, Any]): 音声データの辞書
+        audio_sampling_rate (int): 音声データのサンプリングレート
+    """
+    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")
+    for k, v in audios.items():
+        writer.add_audio(k, v, global_step, audio_sampling_rate)
+
+
+def is_resuming(dir_path: Union[str, Path]) -> bool:
+    """
+    指定されたディレクトリパスに再開可能なモデルが存在するかどうかを返す
+
+    Args:
+        dir_path: チェックするディレクトリのパス
+
+    Returns:
+        bool: 再開可能なモデルが存在するかどうか
+    """
+    # JP-ExtraバージョンではDURがなくWDがあったり変わるため、Gのみで判断する
+    g_list = glob.glob(os.path.join(dir_path, "G_*.pth"))
+    # d_list = glob.glob(os.path.join(dir_path, "D_*.pth"))
+    # dur_list = glob.glob(os.path.join(dir_path, "DUR_*.pth"))
+    return len(g_list) > 0
+
+
+def plot_spectrogram_to_numpy(spectrogram: NDArray[Any]) -> NDArray[Any]:
+    """
+    指定されたスペクトログラムを画像データに変換する
+
+    Args:
+        spectrogram (NDArray[Any]): スペクトログラム
+
+    Returns:
+        NDArray[Any]: 画像データ
+    """
+
+    global __is_matplotlib_imported
+    if not __is_matplotlib_imported:
+        import matplotlib
+
+        matplotlib.use("Agg")
+        __is_matplotlib_imported = True
+        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")
+    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="")  # type: ignore
+    data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
+    plt.close()
+    return data
+
+
+def plot_alignment_to_numpy(alignment: NDArray[Any], info: Optional[str] = None) -> NDArray[Any]:
+    """
+    指定されたアライメントを画像データに変換する
+
+    Args:
+        alignment (NDArray[Any]): アライメント
+        info (Optional[str]): 画像に追加する情報
+
+    Returns:
+        NDArray[Any]: 画像データ
+    """
+
+    global __is_matplotlib_imported
+    if not __is_matplotlib_imported:
+        import matplotlib
+
+        matplotlib.use("Agg")
+        __is_matplotlib_imported = True
+        mpl_logger = logging.getLogger("matplotlib")
+        mpl_logger.setLevel(logging.WARNING)
+    import matplotlib.pylab as plt
+
+    fig, ax = plt.subplots(figsize=(6, 4))
+    im = ax.imshow(
+        alignment.transpose(), aspect="auto", origin="lower", interpolation="none"
+    )
+    fig.colorbar(im, ax=ax)
+    xlabel = "Decoder timestep"
+    if info is not None:
+        xlabel += "\n\n" + info
+    plt.xlabel(xlabel)
+    plt.ylabel("Encoder timestep")
+    plt.tight_layout()
+
+    fig.canvas.draw()
+    data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep="")  # type: ignore
+    data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
+    plt.close()
+    return data
+
+
+def load_wav_to_torch(full_path: Union[str, Path]) -> tuple[torch.FloatTensor, int]:
+    """
+    指定された音声ファイルを読み込み、PyTorch のテンソルに変換して返す
+
+    Args:
+        full_path (Union[str, Path]): 音声ファイルのパス
+
+    Returns:
+        tuple[torch.FloatTensor, int]: 音声データのテンソルとサンプリングレート
+    """
+
+    sampling_rate, data = read(full_path)
+    return torch.FloatTensor(data.astype(np.float32)), sampling_rate
+
+
+def load_filepaths_and_text(filename: Union[str, Path], split: str = "|") -> list[list[str]]:
+    """
+    指定されたファイルからファイルパスとテキストを読み込む
+
+    Args:
+        filename (Union[str, Path]): ファイルのパス
+        split (str): ファイルの区切り文字 (デフォルト: "|")
+
+    Returns:
+        list[list[str]]: ファイルパスとテキストのリスト
+    """
+
+    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_logger(model_dir_path: Union[str, Path], filename: str = "train.log") -> logging.Logger:
+    """
+    ロガーを取得する
+
+    Args:
+        model_dir_path (Union[str, Path]): ログを保存するディレクトリのパス
+        filename (str): ログファイルの名前 (デフォルト: "train.log")
+
+    Returns:
+        logging.Logger: ロガー
+    """
+
+    global logger
+    logger = logging.getLogger(os.path.basename(model_dir_path))
+    logger.setLevel(logging.DEBUG)
+
+    formatter = logging.Formatter("%(asctime)s\t%(name)s\t%(levelname)s\t%(message)s")
+    if not os.path.exists(model_dir_path):
+        os.makedirs(model_dir_path)
+    h = logging.FileHandler(os.path.join(model_dir_path, filename))
+    h.setLevel(logging.DEBUG)
+    h.setFormatter(formatter)
+    logger.addHandler(h)
+    return logger
+
+
+def get_steps(model_path: Union[str, Path]) -> Optional[int]:
+    """
+    モデルのパスからイテレーション番号を取得する
+
+    Args:
+        model_path (Union[str, Path]): モデルのパス
+
+    Returns:
+        Optional[int]: イテレーション番号
+    """
+
+    matches = re.findall(r"\d+", model_path)  # type: ignore
+    return matches[-1] if matches else None
+
+
+def check_git_hash(model_dir_path: Union[str, Path]) -> None:
+    """
+    モデルのディレクトリに .git ディレクトリが存在する場合、ハッシュ値を比較する
+
+    Args:
+        model_dir_path (Union[str, Path]): モデルのディレクトリのパス
+    """
+
+    source_dir = os.path.dirname(os.path.realpath(__file__))
+    if not os.path.exists(os.path.join(source_dir, ".git")):
+        logger.warning(
+            "{} is not a git repository, therefore hash value comparison will be ignored.".format(
+                source_dir
+            )
+        )
+        return
+
+    cur_hash = subprocess.getoutput("git rev-parse HEAD")
+
+    path = os.path.join(model_dir_path, "githash")
+    if os.path.exists(path):
+        saved_hash = open(path).read()
+        if saved_hash != cur_hash:
+            logger.warning(
+                "git hash values are different. {}(saved) != {}(current)".format(
+                    saved_hash[:8], cur_hash[:8]
+                )
+            )
+    else:
+        open(path, "w").write(cur_hash)

style_bert_vits2/models/utils/checkpoints.py

@@ -0,0 +1,194 @@
+import glob
+import os
+import re
+from pathlib import Path
+from typing import Any, Optional, Union
+
+import torch
+
+from style_bert_vits2.logging import logger
+
+
+def load_checkpoint(
+    checkpoint_path: Union[str, Path],
+    model: torch.nn.Module,
+    optimizer: Optional[torch.optim.Optimizer] = None,
+    skip_optimizer: bool = False,
+    for_infer: bool = False
+) -> tuple[torch.nn.Module, Optional[torch.optim.Optimizer], float, int]:
+    """
+    指定されたパスからチェックポイントを読み込み、モデルとオプティマイザーを更新する。
+
+    Args:
+        checkpoint_path (Union[str, Path]): チェックポイントファイルのパス
+        model (torch.nn.Module): 更新するモデル
+        optimizer (Optional[torch.optim.Optimizer]): 更新するオプティマイザー。None の場合は更新しない
+        skip_optimizer (bool): オプティマイザーの更新をスキップするかどうかのフラグ
+        for_infer (bool): 推論用に読み込むかどうかのフラグ
+
+    Returns:
+        tuple[torch.nn.Module, Optional[torch.optim.Optimizer], float, int]: 更新されたモデルとオプティマイザー、学習率、イテレーション番号
+    """
+
+    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"]
+    logger.info(
+        f"Loading model and optimizer at iteration {iteration} from {checkpoint_path}"
+    )
+    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()  # type: ignore
+        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)  # type: ignore
+
+    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
+            new_state_dict[k] = saved_state_dict[k]
+            assert saved_state_dict[k].shape == v.shape, (
+                saved_state_dict[k].shape,
+                v.shape,
+            )
+        except:
+            # For upgrading from the old version
+            if "ja_bert_proj" in k:
+                v = torch.zeros_like(v)
+                logger.warning(
+                    f"Seems you are using the old version of the model, the {k} is automatically set to zero for backward compatibility"
+                )
+            elif "enc_q" in k and for_infer:
+                continue
+            else:
+                logger.error(f"{k} is not in the checkpoint {checkpoint_path}")
+
+            new_state_dict[k] = v
+
+    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 '{}' (iteration {})".format(checkpoint_path, iteration))
+
+    return model, optimizer, learning_rate, iteration
+
+
+def save_checkpoint(
+    model: torch.nn.Module,
+    optimizer: Union[torch.optim.Optimizer, torch.optim.AdamW],
+    learning_rate: float,
+    iteration: int,
+    checkpoint_path: Union[str, Path],
+) -> None:
+    """
+    モデルとオプティマイザーの状態を指定されたパスに保存する。
+
+    Args:
+        model (torch.nn.Module): 保存するモデル
+        optimizer (Union[torch.optim.Optimizer, torch.optim.AdamW]): 保存するオプティマイザー
+        learning_rate (float): 学習率
+        iteration (int): イテレーション数
+        checkpoint_path (Union[str, Path]): 保存先のパス
+    """
+    logger.info(f"Saving model and optimizer state at iteration {iteration} to {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,
+    )
+
+
+def clean_checkpoints(model_dir_path: Union[str, Path] = "logs/44k/", n_ckpts_to_keep: int = 2, sort_by_time: bool = True) -> None:
+    """
+    指定されたディレクトリから古いチェックポイントを削除して空き容量を確保する
+
+    Args:
+        model_dir_path (Union[str, Path]): モデルが保存されているディレクトリのパス
+        n_ckpts_to_keep (int): 保持するチェックポイントの数（G_0.pth と D_0.pth を除く）
+        sort_by_time (bool): True の場合、時間順に削除。False の場合、名前順に削除
+    """
+
+    ckpts_files = [
+        f
+        for f in os.listdir(model_dir_path)
+        if os.path.isfile(os.path.join(model_dir_path, f))
+    ]
+
+    def name_key(_f: str) -> int:
+        return int(re.compile("._(\\d+)\\.pth").match(_f).group(1))  # type: ignore
+
+    def time_key(_f: str) -> float:
+        return os.path.getmtime(os.path.join(model_dir_path, _f))
+
+    sort_key = time_key if sort_by_time else name_key
+
+    def x_sorted(_x: str) -> list[str]:
+        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(model_dir_path, fn)
+        for fn in (
+            x_sorted("G_")[:-n_ckpts_to_keep]
+            + x_sorted("D_")[:-n_ckpts_to_keep]
+            + x_sorted("WD_")[:-n_ckpts_to_keep]
+            + x_sorted("DUR_")[:-n_ckpts_to_keep]
+        )
+    ]
+
+    def del_info(fn: str) -> None:
+        return logger.info(f"Free up space by deleting ckpt {fn}")
+
+    def del_routine(x: str) -> list[Any]:
+        return [os.remove(x), del_info(x)]
+
+    [del_routine(fn) for fn in to_del]
+
+
+def get_latest_checkpoint_path(model_dir_path: Union[str, Path], regex: str = "G_*.pth") -> str:
+    """
+    指定されたディレクトリから最新のチェックポイントのパスを取得する
+
+    Args:
+        model_dir_path (Union[str, Path]): モデルが保存されているディレクトリのパス
+        regex (str): チェックポイントのファイル名の正規表現
+
+    Returns:
+        str: 最新のチェックポイントのパス
+    """
+
+    f_list = glob.glob(os.path.join(str(model_dir_path), regex))
+    f_list.sort(key=lambda f: int("".join(filter(str.isdigit, f))))
+    try:
+        x = f_list[-1]
+    except IndexError:
+        raise ValueError(f"No checkpoint found in {model_dir_path} with regex {regex}")
+
+    return x

style_bert_vits2/models/utils/safetensors.py

@@ -0,0 +1,91 @@
+from pathlib import Path
+from typing import Any, Optional, Union
+
+import torch
+from safetensors import safe_open
+from safetensors.torch import save_file
+
+from style_bert_vits2.logging import logger
+
+
+def load_safetensors(
+    checkpoint_path: Union[str, Path],
+    model: torch.nn.Module,
+    for_infer: bool = False,
+) -> tuple[torch.nn.Module, Optional[int]]:
+    """
+    指定されたパスから safetensors モデルを読み込み、モデルとイテレーションを返す。
+
+    Args:
+        checkpoint_path (Union[str, Path]): モデルのチェックポイントファイルのパス
+        model (torch.nn.Module): 読み込む対象のモデル
+        for_infer (bool): 推論用に読み込むかどうかのフラグ
+
+    Returns:
+        tuple[torch.nn.Module, Optional[int]]: 読み込まれたモデルとイテレーション番号（存在する場合）
+    """
+
+    tensors: dict[str, Any] = {}
+    iteration: Optional[int] = None
+    with safe_open(str(checkpoint_path), framework="pt", device="cpu") as f:  # type: ignore
+        for key in f.keys():
+            if key == "iteration":
+                iteration = f.get_tensor(key).item()
+            tensors[key] = f.get_tensor(key)
+    if hasattr(model, "module"):
+        result = model.module.load_state_dict(tensors, strict=False)
+    else:
+        result = model.load_state_dict(tensors, strict=False)
+    for key in result.missing_keys:
+        if key.startswith("enc_q") and for_infer:
+            continue
+        logger.warning(f"Missing key: {key}")
+    for key in result.unexpected_keys:
+        if key == "iteration":
+            continue
+        logger.warning(f"Unexpected key: {key}")
+    if iteration is None:
+        logger.info(f"Loaded '{checkpoint_path}'")
+    else:
+        logger.info(f"Loaded '{checkpoint_path}' (iteration {iteration})")
+
+    return model, iteration
+
+
+def save_safetensors(
+    model: torch.nn.Module,
+    iteration: int,
+    checkpoint_path: Union[str, Path],
+    is_half: bool = False,
+    for_infer: bool = False,
+) -> None:
+    """
+    モデルを safetensors 形式で保存する。
+
+    Args:
+        model (torch.nn.Module): 保存するモデル
+        iteration (int): イテレーション番号
+        checkpoint_path (Union[str, Path]): 保存先のパス
+        is_half (bool): モデルを半精度で保存するかどうかのフラグ
+        for_infer (bool): 推論用に保存するかどうかのフラグ
+    """
+
+    if hasattr(model, "module"):
+        state_dict = model.module.state_dict()
+    else:
+        state_dict = model.state_dict()
+    keys = []
+    for k in state_dict:
+        if "enc_q" in k and for_infer:
+            continue  # noqa: E701
+        keys.append(k)
+
+    new_dict = (
+        {k: state_dict[k].half() for k in keys}
+        if is_half
+        else {k: state_dict[k] for k in keys}
+    )
+    new_dict["iteration"] = torch.LongTensor([iteration])
+    logger.info(f"Saved safetensors to {checkpoint_path}")
+
+    save_file(new_dict, checkpoint_path)

style_bert_vits2/nlp/__init__.py

@@ -0,0 +1,114 @@
+from typing import Optional, TYPE_CHECKING
+
+from style_bert_vits2.constants import Languages
+from style_bert_vits2.nlp.symbols import (
+    LANGUAGE_ID_MAP,
+    LANGUAGE_TONE_START_MAP,
+    SYMBOLS,
+)
+
+# __init__.py は配下のモジュールをインポートした時点で実行される
+# PyTorch のインポートは重いので、型チェック時以外はインポートしない
+if TYPE_CHECKING:
+    import torch
+
+
+__symbol_to_id = {s: i for i, s in enumerate(SYMBOLS)}
+
+
+def extract_bert_feature(
+    text: str,
+    word2ph: list[int],
+    language: Languages,
+    device: str,
+    assist_text: Optional[str] = None,
+    assist_text_weight: float = 0.7,
+) -> "torch.Tensor":
+    """
+    テキストから BERT の特徴量を抽出する
+
+    Args:
+        text (str): テキスト
+        word2ph (list[int]): 元のテキストの各文字に音素が何個割り当てられるかを表すリスト
+        language (Languages): テキストの言語
+        device (str): 推論に利用するデバイス
+        assist_text (Optional[str], optional): 補助テキスト (デフォルト: None)
+        assist_text_weight (float, optional): 補助テキストの重み (デフォルト: 0.7)
+
+    Returns:
+        torch.Tensor: BERT の特徴量
+    """
+
+    if language == Languages.JP:
+        from style_bert_vits2.nlp.japanese.bert_feature import extract_bert_feature
+    elif language == Languages.EN:
+        from style_bert_vits2.nlp.english.bert_feature import extract_bert_feature
+    elif language == Languages.ZH:
+        from style_bert_vits2.nlp.chinese.bert_feature import extract_bert_feature
+    else:
+        raise ValueError(f"Language {language} not supported")
+
+    return extract_bert_feature(text, word2ph, device, assist_text, assist_text_weight)
+
+
+def clean_text(
+    text: str,
+    language: Languages,
+    use_jp_extra: bool = True,
+    raise_yomi_error: bool = False,
+) -> tuple[str, list[str], list[int], list[int]]:
+    """
+    テキストをクリーニングし、音素に変換する
+
+    Args:
+        text (str): クリーニングするテキスト
+        language (Languages): テキストの言語
+        use_jp_extra (bool, optional): テキストが日本語の場合に JP-Extra モデルを利用するかどうか。Defaults to True.
+        raise_yomi_error (bool, optional): False の場合、読めない文字が消えたような扱いとして処理される。Defaults to False.
+
+    Returns:
+        tuple[str, list[str], list[int], list[int]]: クリーニングされたテキストと、音素・アクセント・元のテキストの各文字に音素が何個割り当てられるかのリスト
+    """
+
+    # Changed to import inside if condition to avoid unnecessary import
+    if language == Languages.JP:
+        from style_bert_vits2.nlp.japanese.g2p import g2p
+        from style_bert_vits2.nlp.japanese.normalizer import normalize_text
+        norm_text = normalize_text(text)
+        phones, tones, word2ph = g2p(norm_text, use_jp_extra, raise_yomi_error)
+    elif language == Languages.EN:
+        from style_bert_vits2.nlp.english.g2p import g2p
+        from style_bert_vits2.nlp.english.normalizer import normalize_text
+        norm_text = normalize_text(text)
+        phones, tones, word2ph = g2p(norm_text)
+    elif language == Languages.ZH:
+        from style_bert_vits2.nlp.chinese.g2p import g2p
+        from style_bert_vits2.nlp.chinese.normalizer import normalize_text
+        norm_text = normalize_text(text)
+        phones, tones, word2ph = g2p(norm_text)
+    else:
+        raise ValueError(f"Language {language} not supported")
+
+    return norm_text, phones, tones, word2ph
+
+
+def cleaned_text_to_sequence(cleaned_phones: list[str], tones: list[int], language: Languages) -> tuple[list[int], list[int], list[int]]:
+    """
+    テキスト文字列を、テキスト内の記号に対応する一連の ID に変換する
+
+    Args:
+        cleaned_phones (list[str]): clean_text() でクリーニングされた音素のリスト (?)
+        tones (list[int]): 各音素のアクセント
+        language (Languages): テキストの言語
+
+    Returns:
+        tuple[list[int], list[int], list[int]]: List of integers corresponding to the symbols in the text
+    """
+
+    phones = [__symbol_to_id[symbol] for symbol in cleaned_phones]
+    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

style_bert_vits2/nlp/bert_models.py

@@ -0,0 +1,177 @@
+"""
+Style-Bert-VITS2 の学習・推論に必要な各言語ごとの BERT モデルをロード/取得するためのモジュール。
+
+オリジナルの Bert-VITS2 では各言語ごとの BERT モデルが初回インポート時にハードコードされたパスから「暗黙的に」ロードされているが、
+場合によっては多重にロードされて非効率なほか、BERT モデルのロード元のパスがハードコードされているためライブラリ化ができない。
+
+そこで、ライブラリの利用前に、音声合成に利用する言語の BERT モデルだけを「明示的に」ロードできるようにした。
+一度 load_model/tokenizer() で当該言語の BERT モデルがロードされていれば、ライブラリ内部のどこからでもロード済みのモデル/トークナイザーを取得できる。
+"""
+
+import gc
+from typing import cast, Optional, Union
+
+import torch
+from transformers import (
+    AutoModelForMaskedLM,
+    AutoTokenizer,
+    DebertaV2Model,
+    DebertaV2Tokenizer,
+    PreTrainedModel,
+    PreTrainedTokenizer,
+    PreTrainedTokenizerFast,
+)
+
+from style_bert_vits2.constants import DEFAULT_BERT_TOKENIZER_PATHS, Languages
+from style_bert_vits2.logging import logger
+
+
+# 各言語ごとのロード済みの BERT モデルを格納する辞書
+__loaded_models: dict[Languages, Union[PreTrainedModel, DebertaV2Model]] = {}
+
+# 各言語ごとのロード済みの BERT トークナイザーを格納する辞書
+__loaded_tokenizers: dict[Languages, Union[PreTrainedTokenizer, PreTrainedTokenizerFast, DebertaV2Tokenizer]] = {}
+
+
+def load_model(
+    language: Languages,
+    pretrained_model_name_or_path: Optional[str] = None,
+) -> Union[PreTrainedModel, DebertaV2Model]:
+    """
+    指定された言語の BERT モデルをロードし、ロード済みの BERT モデルを返す。
+    一度ロードされていれば、ロード済みの BERT モデルを即座に返す。
+    ライブラリ利用時は常に pretrain_model_name_or_path (Hugging Face のリポジトリ名 or ローカルのファイルパス) を指定する必要がある。
+    ロードにはそれなりに時間がかかるため、ライブラリ利用前に明示的に pretrained_model_name_or_path を指定してロードしておくべき。
+
+    Style-Bert-VITS2 では、BERT モデルに下記の 3 つが利用されている。
+    これ以外の BERT モデルを指定した場合は正常に動作しない可能性が高い。
+    - 日本語: ku-nlp/deberta-v2-large-japanese-char-wwm
+    - 英語: microsoft/deberta-v3-large
+    - 中国語: hfl/chinese-roberta-wwm-ext-large
+
+    Args:
+        language (Languages): ロードする学習済みモデルの対象言語
+        pretrained_model_name_or_path (Optional[str]): ロードする学習済みモデルの名前またはパス。指定しない場合はデフォルトのパスが利用される (デフォルト: None)
+
+    Returns:
+        Union[PreTrainedModel, DebertaV2Model]: ロード済みの BERT モデル
+    """
+
+    # すでにロード済みの場合はそのまま返す
+    if language in __loaded_models:
+        return __loaded_models[language]
+
+    # pretrained_model_name_or_path が指定されていない場合はデフォルトのパスを利用
+    if pretrained_model_name_or_path is None:
+        assert DEFAULT_BERT_TOKENIZER_PATHS[language].exists(), \
+            f"The default {language} BERT model does not exist on the file system. Please specify the path to the pre-trained model."
+        pretrained_model_name_or_path = str(DEFAULT_BERT_TOKENIZER_PATHS[language])
+
+    # BERT モデルをロードし、辞書に格納して返す
+    ## 英語のみ DebertaV2Model でロードする必要がある
+    if language == Languages.EN:
+        model = cast(DebertaV2Model, DebertaV2Model.from_pretrained(pretrained_model_name_or_path))
+    else:
+        model = AutoModelForMaskedLM.from_pretrained(pretrained_model_name_or_path)
+    __loaded_models[language] = model
+    logger.info(f"Loaded the {language} BERT model from {pretrained_model_name_or_path}")
+
+    return model
+
+
+def load_tokenizer(
+    language: Languages,
+    pretrained_model_name_or_path: Optional[str] = None,
+) -> Union[PreTrainedTokenizer, PreTrainedTokenizerFast, DebertaV2Tokenizer]:
+    """
+    指定された言語の BERT モデルをロードし、ロード済みの BERT トークナイザーを返す。
+    一度ロードされていれば、ロード済みの BERT トークナイザーを即座に返す。
+    ライブラリ利用時は常に pretrain_model_name_or_path (Hugging Face のリポジトリ名 or ローカルのファイルパス) を指定する必要がある。
+    ロードにはそれなりに時間がかかるため、ライブラリ利用前に明示的に pretrained_model_name_or_path を指定してロードしておくべき。
+
+    Style-Bert-VITS2 では、BERT モデルに下記の 3 つが利用されている。
+    これ以外の BERT モデルを指定した場合は正常に動作しない可能性が高い。
+    - 日本語: ku-nlp/deberta-v2-large-japanese-char-wwm
+    - 英語: microsoft/deberta-v3-large
+    - 中国語: hfl/chinese-roberta-wwm-ext-large
+
+    Args:
+        language (Languages): ロードする学習済みモデルの対象言語
+        pretrained_model_name_or_path (Optional[str]): ロードする学習済みモデルの名前またはパス。指定しない場合はデフォルトのパスが利用される (デフォルト: None)
+
+    Returns:
+        Union[PreTrainedTokenizer, PreTrainedTokenizerFast, DebertaV2Tokenizer]: ロード済みの BERT トークナイザー
+    """
+
+    # すでにロード済みの場合はそのまま返す
+    if language in __loaded_tokenizers:
+        return __loaded_tokenizers[language]
+
+    # pretrained_model_name_or_path が指定されていない場合はデフォルトのパスを利用
+    if pretrained_model_name_or_path is None:
+        assert DEFAULT_BERT_TOKENIZER_PATHS[language].exists(), \
+            f"The default {language} BERT tokenizer does not exist on the file system. Please specify the path to the pre-trained model."
+        pretrained_model_name_or_path = str(DEFAULT_BERT_TOKENIZER_PATHS[language])
+
+    # BERT トークナイザーをロードし、辞書に格納して返す
+    ## 英語のみ DebertaV2Tokenizer でロードする必要がある
+    if language == Languages.EN:
+        tokenizer = DebertaV2Tokenizer.from_pretrained(pretrained_model_name_or_path)
+    else:
+        tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path)
+    __loaded_tokenizers[language] = tokenizer
+    logger.info(f"Loaded the {language} BERT tokenizer from {pretrained_model_name_or_path}")
+
+    return tokenizer
+
+
+def unload_model(language: Languages) -> None:
+    """
+    指定された言語の BERT モデルをアンロードする。
+
+    Args:
+        language (Languages): アンロードする BERT モデルの言語
+    """
+
+    if language in __loaded_models:
+        del __loaded_models[language]
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        logger.info(f"Unloaded the {language} BERT model")
+
+
+def unload_tokenizer(language: Languages) -> None:
+    """
+    指定された言語の BERT トークナイザーをアンロードする。
+
+    Args:
+        language (Languages): アンロードする BERT トークナイザーの言語
+    """
+
+    if language in __loaded_tokenizers:
+        del __loaded_tokenizers[language]
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        logger.info(f"Unloaded the {language} BERT tokenizer")
+
+
+def unload_all_models() -> None:
+    """
+    すべての BERT モデルをアンロードする。
+    """
+
+    for language in list(__loaded_models.keys()):
+        unload_model(language)
+    logger.info("Unloaded all BERT models")
+
+
+def unload_all_tokenizers() -> None:
+    """
+    すべての BERT トークナイザーをアンロードする。
+    """
+
+    for language in list(__loaded_tokenizers.keys()):
+        unload_tokenizer(language)
+    logger.info("Unloaded all BERT tokenizers")

style_bert_vits2/nlp/chinese/bert_feature.py

@@ -1,54 +1,58 @@
-import sys
+from typing import Optional
 
 import torch
-from transformers import AutoModelForMaskedLM, AutoTokenizer
 
-from config import config
+from style_bert_vits2.constants import Languages
-
+from style_bert_vits2.nlp import bert_models
-LOCAL_PATH = "./bert/chinese-roberta-wwm-ext-large"
-
-tokenizer = AutoTokenizer.from_pretrained(LOCAL_PATH)
-
-models = dict()
 
 
-def get_bert_feature(
+def extract_bert_feature(
-    text,
+    text: str,
-    word2ph,
+    word2ph: list[int],
-    device=config.bert_gen_config.device,
+    device: str,
-    assist_text=None,
+    assist_text: Optional[str] = None,
-    assist_text_weight=0.7,
+    assist_text_weight: float = 0.7,
-):
+) -> torch.Tensor:
-    if (
+    """
-        sys.platform == "darwin"
+    中国語のテキストから BERT の特徴量を抽出する
-        and torch.backends.mps.is_available()
+
-        and device == "cpu"
+    Args:
-    ):
+        text (str): 中国語のテキスト
-        device = "mps"
+        word2ph (list[int]): 元のテキストの各文字に音素が何個割り当てられるかを表すリスト
-    if not device:
+        device (str): 推論に利用するデバイス
-        device = "cuda"
+        assist_text (Optional[str], optional): 補助テキスト (デフォルト: None)
+        assist_text_weight (float, optional): 補助テキストの重み (デフォルト: 0.7)
+
+    Returns:
+        torch.Tensor: BERT の特徴量
+    """
+
     if device == "cuda" and not torch.cuda.is_available():
         device = "cpu"
-    if device not in models.keys():
+    model = bert_models.load_model(Languages.ZH).to(device)  # type: ignore
-        models[device] = AutoModelForMaskedLM.from_pretrained(LOCAL_PATH).to(device)
+
+    style_res_mean = None
     with torch.no_grad():
+        tokenizer = bert_models.load_tokenizer(Languages.ZH)
         inputs = tokenizer(text, return_tensors="pt")
         for i in inputs:
-            inputs[i] = inputs[i].to(device)
+            inputs[i] = inputs[i].to(device)  # type: ignore
-        res = models[device](**inputs, output_hidden_states=True)
+        res = model(**inputs, output_hidden_states=True)
         res = torch.cat(res["hidden_states"][-3:-2], -1)[0].cpu()
         if assist_text:
             style_inputs = tokenizer(assist_text, return_tensors="pt")
             for i in style_inputs:
-                style_inputs[i] = style_inputs[i].to(device)
+                style_inputs[i] = style_inputs[i].to(device)  # type: ignore
-            style_res = models[device](**style_inputs, output_hidden_states=True)
+            style_res = model(**style_inputs, output_hidden_states=True)
             style_res = torch.cat(style_res["hidden_states"][-3:-2], -1)[0].cpu()
             style_res_mean = style_res.mean(0)
+
     assert len(word2ph) == len(text) + 2
     word2phone = word2ph
     phone_level_feature = []
     for i in range(len(word2phone)):
         if assist_text:
+            assert style_res_mean is not None
             repeat_feature = (
                 res[i].repeat(word2phone[i], 1) * (1 - assist_text_weight)
                 + style_res_mean.repeat(word2phone[i], 1) * assist_text_weight

style_bert_vits2/nlp/chinese/g2p.py

@@ -1,75 +1,23 @@
-import os
 import re
-
+from pathlib import Path
-import cn2an
-from pypinyin import lazy_pinyin, Style
-
-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()
-}
 
 import jieba.posseg as psg
+from pypinyin import lazy_pinyin, Style
+
+from style_bert_vits2.nlp.chinese.tone_sandhi import ToneSandhi
+from style_bert_vits2.nlp.symbols import PUNCTUATIONS
 
 
-rep_map = {
+__PINYIN_TO_SYMBOL_MAP = {
-    "：": ",",
+    line.split("\t")[0]: line.strip().split("\t")[1]
-    "；": ",",
+    for line in open(Path(__file__).parent / "opencpop-strict.txt").readlines()
-    "，": ",",
-    "。": ".",
-    "！": "!",
-    "？": "?",
-    "\n": ".",
-    "·": ",",
-    "、": ",",
-    "...": "…",
-    "$": ".",
-    "“": "'",
-    "”": "'",
-    '"': "'",
-    "‘": "'",
-    "’": "'",
-    "（": "'",
-    "）": "'",
-    "(": "'",
-    ")": "'",
-    "《": "'",
-    "》": "'",
-    "【": "'",
-    "】": "'",
-    "[": "'",
-    "]": "'",
-    "—": "-",
-    "～": "-",
-    "~": "-",
-    "「": "'",
-    "」": "'",
 }
 
-tone_modifier = ToneSandhi()
 
-
+def g2p(text: str) -> tuple[list[str], list[int], list[int]]:
-def replace_punctuation(text):
+    pattern = r"(?<=[{0}])\s*".format("".join(PUNCTUATIONS))
-    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
-    )
-
-    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() != ""]
-    phones, tones, word2ph = _g2p(sentences)
+    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 + ["_"]
@@ -78,34 +26,22 @@ def g2p(text):
     return phones, tones, word2ph
 
 
-def _get_initials_finals(word):
+def __g2p(segments: list[str]) -> tuple[list[str], list[int], list[int]]:
-    initials = []
-    finals = []
-    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
-    )
-    for c, v in zip(orig_initials, orig_finals):
-        initials.append(c)
-        finals.append(v)
-    return initials, finals
-
-
-def _g2p(segments):
     phones_list = []
     tones_list = []
     word2ph = []
+    tone_modifier = ToneSandhi()
     for seg in segments:
         # Replace all English words in the sentence
         seg = re.sub("[a-zA-Z]+", "", seg)
         seg_cut = psg.lcut(seg)
         initials = []
         finals = []
-        seg_cut = tone_modifier.pre_merge_for_modify(seg_cut)
+        seg_cut = tone_modifier.pre_merge_for_modify(seg_cut)  # type: ignore
         for word, pos in seg_cut:
             if pos == "eng":
                 continue
-            sub_initials, sub_finals = _get_initials_finals(word)
+            sub_initials, sub_finals = __get_initials_finals(word)
             sub_finals = tone_modifier.modified_tone(word, pos, sub_finals)
             initials.append(sub_initials)
             finals.append(sub_finals)
@@ -119,7 +55,7 @@ def _g2p(segments):
             # NOTE: post process for pypinyin outputs
             # we discriminate i, ii and iii
             if c == v:
-                assert c in punctuation
+                assert c in PUNCTUATIONS
                 phone = [c]
                 tone = "0"
                 word2ph.append(1)
@@ -159,8 +95,8 @@ def _g2p(segments):
                         if pinyin[0] in single_rep_map.keys():
                             pinyin = single_rep_map[pinyin[0]] + pinyin[1:]
 
-                assert pinyin in pinyin_to_symbol_map.keys(), (pinyin, seg, raw_pinyin)
+                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
@@ -168,32 +104,32 @@ def _g2p(segments):
     return phones_list, tones_list, word2ph
 
 
-def text_normalize(text):
+def __get_initials_finals(word: str) -> tuple[list[str], list[str]]:
-    numbers = re.findall(r"\d+(?:\.?\d+)?", text)
+    initials = []
-    for number in numbers:
+    finals = []
-        text = text.replace(number, cn2an.an2cn(number), 1)
+    orig_initials = lazy_pinyin(word, neutral_tone_with_five=True, style=Style.INITIALS)
-    text = replace_punctuation(text)
+    orig_finals = lazy_pinyin(
-    return text
+        word, neutral_tone_with_five=True, style=Style.FINALS_TONE3
-
+    )
-
+    for c, v in zip(orig_initials, orig_finals):
-def get_bert_feature(text, word2ph):
+        initials.append(c)
-    from text import chinese_bert
+        finals.append(v)
-
+    return initials, finals
-    return chinese_bert.get_bert_feature(text, word2ph)
 
 
 if __name__ == "__main__":
-    from text.chinese_bert import get_bert_feature
+    from style_bert_vits2.nlp.chinese.bert_feature import extract_bert_feature
+    from style_bert_vits2.nlp.chinese.normalizer import normalize_text
 
-    text = "啊！但是《原神》是由,米哈\游自主，  [研发]的一款全.新开放世界.冒险游戏"
+    text = "啊！但是《原神》是由,米哈游自主，  [研发]的一款全.新开放世界.冒险游戏"
-    text = text_normalize(text)
+    text = normalize_text(text)
     print(text)
     phones, tones, word2ph = g2p(text)
-    bert = get_bert_feature(text, word2ph)
+    bert = extract_bert_feature(text, word2ph, 'cuda')
 
     print(phones, tones, word2ph, bert.shape)
 
 
-# # 示例用法
+# 示例用法
 # text = "这是一个示例文本：,你好！这是一个测试...."
 # print(g2p_paddle(text))  # 输出: 这是一个示例文本你好这是一个测试

style_bert_vits2/nlp/chinese/normalizer.py

@@ -0,0 +1,61 @@
+import re
+
+import cn2an
+
+from style_bert_vits2.nlp.symbols import PUNCTUATIONS
+
+
+def normalize_text(text: str) -> str:
+    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 replace_punctuation(text: str) -> str:
+
+    REPLACE_MAP = {
+        "：": ",",
+        "；": ",",
+        "，": ",",
+        "。": ".",
+        "！": "!",
+        "？": "?",
+        "\n": ".",
+        "·": ",",
+        "、": ",",
+        "...": "…",
+        "$": ".",
+        "“": "'",
+        "”": "'",
+        '"': "'",
+        "‘": "'",
+        "’": "'",
+        "（": "'",
+        "）": "'",
+        "(": "'",
+        ")": "'",
+        "《": "'",
+        "》": "'",
+        "【": "'",
+        "】": "'",
+        "[": "'",
+        "]": "'",
+        "—": "-",
+        "～": "-",
+        "~": "-",
+        "「": "'",
+        "」": "'",
+    }
+
+    text = text.replace("嗯", "恩").replace("呣", "母")
+    pattern = re.compile("|".join(re.escape(p) for p in REPLACE_MAP.keys()))
+
+    replaced_text = pattern.sub(lambda x: REPLACE_MAP[x.group()], text)
+
+    replaced_text = re.sub(
+        r"[^\u4e00-\u9fa5" + "".join(PUNCTUATIONS) + r"]+", "", replaced_text
+    )
+
+    return replaced_text

style_bert_vits2/nlp/chinese/tone_sandhi.py

@@ -11,8 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import List
-from typing import Tuple
 
 import jieba
 from pypinyin import lazy_pinyin
@@ -463,7 +461,7 @@ 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 (
@@ -522,7 +520,7 @@ class ToneSandhi:
         finals = sum(finals_list, [])
         return finals
 
-    def _bu_sandhi(self, word: str, finals: List[str]) -> List[str]:
+    def _bu_sandhi(self, word: str, finals: list[str]) -> list[str]:
         # e.g. 看不懂
         if len(word) == 3 and word[1] == "不":
             finals[1] = finals[1][:-1] + "5"
@@ -533,7 +531,7 @@ class ToneSandhi:
                     finals[i] = finals[i][:-1] + "2"
         return finals
 
-    def _yi_sandhi(self, word: str, finals: List[str]) -> List[str]:
+    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 != "一"]
@@ -558,9 +556,9 @@ class ToneSandhi:
                             finals[i] = finals[i][:-1] + "4"
         return finals
 
-    def _split_word(self, word: str) -> List[str]:
+    def _split_word(self, word: str) -> list[str]:
         word_list = jieba.cut_for_search(word)
-        word_list = sorted(word_list, key=lambda i: len(i), reverse=False)
+        word_list = sorted(word_list, key=lambda i: len(i), reverse=False)  # type: ignore
         first_subword = word_list[0]
         first_begin_idx = word.find(first_subword)
         if first_begin_idx == 0:
@@ -571,7 +569,7 @@ class ToneSandhi:
             new_word_list = [second_subword, first_subword]
         return new_word_list
 
-    def _three_sandhi(self, word: str, finals: List[str]) -> List[str]:
+    def _three_sandhi(self, word: str, finals: list[str]) -> list[str]:
         if len(word) == 2 and self._all_tone_three(finals):
             finals[0] = finals[0][:-1] + "2"
         elif len(word) == 3:
@@ -611,12 +609,12 @@ class ToneSandhi:
 
         return finals
 
-    def _all_tone_three(self, finals: List[str]) -> bool:
+    def _all_tone_three(self, finals: list[str]) -> bool:
         return all(x[-1] == "3" for x in finals)
 
     # merge "不" and the word behind it
     # if don't merge, "不" sometimes appears alone according to jieba, which may occur sandhi error
-    def _merge_bu(self, seg: List[Tuple[str, str]]) -> List[Tuple[str, str]]:
+    def _merge_bu(self, seg: list[tuple[str, str]]) -> list[tuple[str, str]]:
         new_seg = []
         last_word = ""
         for word, pos in seg:
@@ -636,7 +634,7 @@ class ToneSandhi:
     # e.g.
     # input seg: [('听', 'v'), ('一', 'm'), ('听', 'v')]
     # output seg: [['听一听', 'v']]
-    def _merge_yi(self, seg: List[Tuple[str, str]]) -> List[Tuple[str, str]]:
+    def _merge_yi(self, seg: list[tuple[str, str]]) -> list[tuple[str, str]]:
         new_seg = [] * len(seg)
         # function 1
         i = 0
@@ -674,8 +672,8 @@ class ToneSandhi:
 
     # the first and the second words are all_tone_three
     def _merge_continuous_three_tones(
-        self, seg: List[Tuple[str, str]]
+        self, seg: list[tuple[str, str]]
-    ) -> 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)
@@ -709,8 +707,8 @@ 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]]
+        self, seg: list[tuple[str, str]]
-    ) -> 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)
@@ -738,7 +736,7 @@ class ToneSandhi:
                 new_seg.append([word, pos])
         return new_seg
 
-    def _merge_er(self, seg: List[Tuple[str, str]]) -> List[Tuple[str, str]]:
+    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] != "#":
@@ -747,7 +745,7 @@ class ToneSandhi:
                 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]:
@@ -756,7 +754,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)
@@ -768,7 +766,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)

style_bert_vits2/nlp/english/bert_feature.py

@@ -0,0 +1,66 @@
+from typing import Optional
+
+import torch
+
+from style_bert_vits2.constants import Languages
+from style_bert_vits2.nlp import bert_models
+
+
+def extract_bert_feature(
+    text: str,
+    word2ph: list[int],
+    device: str,
+    assist_text: Optional[str] = None,
+    assist_text_weight: float = 0.7,
+) -> torch.Tensor:
+    """
+    英語のテキストから BERT の特徴量を抽出する
+
+    Args:
+        text (str): 英語のテキスト
+        word2ph (list[int]): 元のテキストの各文字に音素が何個割り当てられるかを表すリスト
+        device (str): 推論に利用するデバイス
+        assist_text (Optional[str], optional): 補助テキスト (デフォルト: None)
+        assist_text_weight (float, optional): 補助テキストの重み (デフォルト: 0.7)
+
+    Returns:
+        torch.Tensor: BERT の特徴量
+    """
+
+    if device == "cuda" and not torch.cuda.is_available():
+        device = "cpu"
+    model = bert_models.load_model(Languages.EN).to(device)  # type: ignore
+
+    style_res_mean = None
+    with torch.no_grad():
+        tokenizer = bert_models.load_tokenizer(Languages.EN)
+        inputs = tokenizer(text, return_tensors="pt")
+        for i in inputs:
+            inputs[i] = inputs[i].to(device)  # type: ignore
+        res = model(**inputs, output_hidden_states=True)
+        res = torch.cat(res["hidden_states"][-3:-2], -1)[0].cpu()
+        if assist_text:
+            style_inputs = tokenizer(assist_text, return_tensors="pt")
+            for i in style_inputs:
+                style_inputs[i] = style_inputs[i].to(device)  # type: ignore
+            style_res = model(**style_inputs, output_hidden_states=True)
+            style_res = torch.cat(style_res["hidden_states"][-3:-2], -1)[0].cpu()
+            style_res_mean = style_res.mean(0)
+
+    assert len(word2ph) == res.shape[0], (text, res.shape[0], len(word2ph))
+    word2phone = word2ph
+    phone_level_feature = []
+    for i in range(len(word2phone)):
+        if assist_text:
+            assert style_res_mean is not None
+            repeat_feature = (
+                res[i].repeat(word2phone[i], 1) * (1 - assist_text_weight)
+                + style_res_mean.repeat(word2phone[i], 1) * assist_text_weight
+            )
+        else:
+            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

style_bert_vits2/nlp/english/cmudict.py

@@ -0,0 +1,46 @@
+import pickle
+from pathlib import Path
+
+
+CMU_DICT_PATH = Path(__file__).parent / "cmudict.rep"
+CACHE_PATH = Path(__file__).parent / "cmudict_cache.pickle"
+
+
+def get_dict() -> dict[str, list[list[str]]]:
+    if CACHE_PATH.exists():
+        with open(CACHE_PATH, "rb") as pickle_file:
+            g2p_dict = pickle.load(pickle_file)
+    else:
+        g2p_dict = read_dict()
+        cache_dict(g2p_dict, CACHE_PATH)
+
+    return g2p_dict
+
+
+def read_dict() -> dict[str, list[list[str]]]:
+    g2p_dict = {}
+    start_line = 49
+    with open(CMU_DICT_PATH) as f:
+        line = f.readline()
+        line_index = 1
+        while line:
+            if line_index >= start_line:
+                line = line.strip()
+                word_split = line.split("  ")
+                word = word_split[0]
+
+                syllable_split = word_split[1].split(" - ")
+                g2p_dict[word] = []
+                for syllable in syllable_split:
+                    phone_split = syllable.split(" ")
+                    g2p_dict[word].append(phone_split)
+
+            line_index = line_index + 1
+            line = f.readline()
+
+    return g2p_dict
+
+
+def cache_dict(g2p_dict: dict[str, list[list[str]]], file_path: Path) -> None:
+    with open(file_path, "wb") as pickle_file:
+        pickle.dump(g2p_dict, pickle_file)

style_bert_vits2/nlp/english/g2p.py

@@ -0,0 +1,240 @@
+import re
+
+from g2p_en import G2p
+
+from style_bert_vits2.constants import Languages
+from style_bert_vits2.nlp import bert_models
+from style_bert_vits2.nlp.english.cmudict import get_dict
+from style_bert_vits2.nlp.symbols import PUNCTUATIONS, SYMBOLS
+
+
+def g2p(text: str) -> tuple[list[str], list[int], list[int]]:
+
+    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",
+    }
+
+    _g2p = G2p()
+
+    phones = []
+    tones = []
+    phone_len = []
+    # tokens = [tokenizer.tokenize(i) for i in words]
+    words = __text_to_words(text)
+    eng_dict = get_dict()
+
+    for word in words:
+        temp_phones, temp_tones = [], []
+        if len(word) > 1:
+            if "'" in word:
+                word = ["".join(word)]
+        for w in word:
+            if w in PUNCTUATIONS:
+                temp_phones.append(w)
+                temp_tones.append(0)
+                continue
+            if w.upper() in eng_dict:
+                phns, tns = __refine_syllables(eng_dict[w.upper()])
+                temp_phones += [__post_replace_ph(i) for i in phns]
+                temp_tones += tns
+                # w2ph.append(len(phns))
+            else:
+                phone_list = list(filter(lambda p: p != " ", _g2p(w)))  # type: ignore
+                phns = []
+                tns = []
+                for ph in phone_list:
+                    if ph in ARPA:
+                        ph, tn = __refine_ph(ph)
+                        phns.append(ph)
+                        tns.append(tn)
+                    else:
+                        phns.append(ph)
+                        tns.append(0)
+                temp_phones += [__post_replace_ph(i) for i in phns]
+                temp_tones += tns
+        phones += temp_phones
+        tones += temp_tones
+        phone_len.append(len(temp_phones))
+        # phones = [post_replace_ph(i) for i in phones]
+
+    word2ph = []
+    for token, pl in zip(words, phone_len):
+        word_len = len(token)
+
+        aaa = __distribute_phone(pl, word_len)
+        word2ph += aaa
+
+    phones = ["_"] + phones + ["_"]
+    tones = [0] + tones + [0]
+    word2ph = [1] + word2ph + [1]
+    assert len(phones) == len(tones), text
+    assert len(phones) == sum(word2ph), text
+
+    return phones, tones, word2ph
+
+
+def __post_replace_ph(ph: str) -> str:
+    REPLACE_MAP = {
+        "：": ",",
+        "；": ",",
+        "，": ",",
+        "。": ".",
+        "！": "!",
+        "？": "?",
+        "\n": ".",
+        "·": ",",
+        "、": ",",
+        "…": "...",
+        "···": "...",
+        "・・・": "...",
+        "v": "V",
+    }
+    if ph in REPLACE_MAP.keys():
+        ph = REPLACE_MAP[ph]
+    if ph in SYMBOLS:
+        return ph
+    if ph not in SYMBOLS:
+        ph = "UNK"
+    return ph
+
+
+def __refine_ph(phn: str) -> tuple[str, int]:
+    tone = 0
+    if re.search(r"\d$", phn):
+        tone = int(phn[-1]) + 1
+        phn = phn[:-1]
+    else:
+        tone = 3
+    return phn.lower(), tone
+
+
+def __refine_syllables(syllables: list[list[str]]) -> tuple[list[str], list[int]]:
+    tones = []
+    phonemes = []
+    for phn_list in syllables:
+        for i in range(len(phn_list)):
+            phn = phn_list[i]
+            phn, tone = __refine_ph(phn)
+            phonemes.append(phn)
+            tones.append(tone)
+    return phonemes, tones
+
+
+def __distribute_phone(n_phone: int, n_word: int) -> list[int]:
+    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
+
+
+def __text_to_words(text: str) -> list[list[str]]:
+    tokenizer = bert_models.load_tokenizer(Languages.EN)
+    tokens = tokenizer.tokenize(text)
+    words = []
+    for idx, t in enumerate(tokens):
+        if t.startswith("▁"):
+            words.append([t[1:]])
+        else:
+            if t in PUNCTUATIONS:
+                if idx == len(tokens) - 1:
+                    words.append([f"{t}"])
+                else:
+                    if (
+                        not tokens[idx + 1].startswith("▁")
+                        and tokens[idx + 1] not in PUNCTUATIONS
+                    ):
+                        if idx == 0:
+                            words.append([])
+                        words[-1].append(f"{t}")
+                    else:
+                        words.append([f"{t}"])
+            else:
+                if idx == 0:
+                    words.append([])
+                words[-1].append(f"{t}")
+    return words
+
+
+if __name__ == "__main__":
+    # print(get_dict())
+    # print(eng_word_to_phoneme("hello"))
+    print(g2p("In this paper, we propose 1 DSPGAN, a GAN-based universal vocoder."))
+    # all_phones = set()
+    # eng_dict = get_dict()
+    # for k, syllables in eng_dict.items():
+    #     for group in syllables:
+    #         for ph in group:
+    #             all_phones.add(ph)
+    # print(all_phones)

style_bert_vits2/nlp/english/normalizer.py

@@ -0,0 +1,130 @@
+import re
+
+import inflect
+
+
+__INFLECT = inflect.engine()
+__COMMA_NUMBER_PATTERN = re.compile(r"([0-9][0-9\,]+[0-9])")
+__DECIMAL_NUMBER_PATTERN = re.compile(r"([0-9]+\.[0-9]+)")
+__POUNDS_PATTERN = re.compile(r"£([0-9\,]*[0-9]+)")
+__DOLLARS_PATTERN = re.compile(r"\$([0-9\.\,]*[0-9]+)")
+__ORDINAL_PATTERN = re.compile(r"[0-9]+(st|nd|rd|th)")
+__NUMBER_PATTERN = re.compile(r"[0-9]+")
+
+
+def normalize_text(text: str) -> str:
+    text = __normalize_numbers(text)
+    text = replace_punctuation(text)
+    text = re.sub(r"([,;.\?\!])([\w])", r"\1 \2", text)
+    return text
+
+
+def replace_punctuation(text: str) -> str:
+    REPLACE_MAP = {
+        "：": ",",
+        "；": ",",
+        "，": ",",
+        "。": ".",
+        "！": "!",
+        "？": "?",
+        "\n": ".",
+        "．": ".",
+        "…": "...",
+        "···": "...",
+        "・・・": "...",
+        "·": ",",
+        "・": ",",
+        "、": ",",
+        "$": ".",
+        "“": "'",
+        "”": "'",
+        '"': "'",
+        "‘": "'",
+        "’": "'",
+        "（": "'",
+        "）": "'",
+        "(": "'",
+        ")": "'",
+        "《": "'",
+        "》": "'",
+        "【": "'",
+        "】": "'",
+        "[": "'",
+        "]": "'",
+        "—": "-",
+        "−": "-",
+        "～": "-",
+        "~": "-",
+        "「": "'",
+        "」": "'",
+    }
+    pattern = re.compile("|".join(re.escape(p) for p in REPLACE_MAP.keys()))
+    replaced_text = pattern.sub(lambda x: REPLACE_MAP[x.group()], text)
+    # replaced_text = re.sub(
+    #     r"[^\u3040-\u309F\u30A0-\u30FF\u4E00-\u9FFF\u3400-\u4DBF\u3005"
+    #     + "".join(punctuation)
+    #     + r"]+",
+    #     "",
+    #     replaced_text,
+    # )
+    return replaced_text
+
+
+def __normalize_numbers(text: str) -> str:
+    text = re.sub(__COMMA_NUMBER_PATTERN, __remove_commas, text)
+    text = re.sub(__POUNDS_PATTERN, r"\1 pounds", text)
+    text = re.sub(__DOLLARS_PATTERN, __expand_dollars, text)
+    text = re.sub(__DECIMAL_NUMBER_PATTERN, __expand_decimal_point, text)
+    text = re.sub(__ORDINAL_PATTERN, __expand_ordinal, text)
+    text = re.sub(__NUMBER_PATTERN, __expand_number, text)
+    return text
+
+
+def __expand_dollars(m: re.Match[str]) -> str:
+    match = m.group(1)
+    parts = match.split(".")
+    if len(parts) > 2:
+        return match + " dollars"  # Unexpected format
+    dollars = int(parts[0]) if parts[0] else 0
+    cents = int(parts[1]) if len(parts) > 1 and parts[1] else 0
+    if dollars and cents:
+        dollar_unit = "dollar" if dollars == 1 else "dollars"
+        cent_unit = "cent" if cents == 1 else "cents"
+        return "%s %s, %s %s" % (dollars, dollar_unit, cents, cent_unit)
+    elif dollars:
+        dollar_unit = "dollar" if dollars == 1 else "dollars"
+        return "%s %s" % (dollars, dollar_unit)
+    elif cents:
+        cent_unit = "cent" if cents == 1 else "cents"
+        return "%s %s" % (cents, cent_unit)
+    else:
+        return "zero dollars"
+
+
+def __remove_commas(m: re.Match[str]) -> str:
+    return m.group(1).replace(",", "")
+
+
+def __expand_ordinal(m: re.Match[str]) -> str:
+    return __INFLECT.number_to_words(m.group(0))  # type: ignore
+
+
+def __expand_number(m: re.Match[str]) -> str:
+    num = int(m.group(0))
+    if num > 1000 and num < 3000:
+        if num == 2000:
+            return "two thousand"
+        elif num > 2000 and num < 2010:
+            return "two thousand " + __INFLECT.number_to_words(num % 100)  # type: ignore
+        elif num % 100 == 0:
+            return __INFLECT.number_to_words(num // 100) + " hundred"  # type: ignore
+        else:
+            return __INFLECT.number_to_words(
+                num, andword="", zero="oh", group=2  # type: ignore
+            ).replace(", ", " ")  # type: ignore
+    else:
+        return __INFLECT.number_to_words(num, andword="")  # type: ignore
+
+
+def __expand_decimal_point(m: re.Match[str]) -> str:
+    return m.group(1).replace(".", " point ")

style_bert_vits2/nlp/japanese/bert_feature.py

@@ -0,0 +1,73 @@
+from typing import Optional
+
+import torch
+
+from style_bert_vits2.constants import Languages
+from style_bert_vits2.nlp import bert_models
+from style_bert_vits2.nlp.japanese.g2p import text_to_sep_kata
+
+
+def extract_bert_feature(
+    text: str,
+    word2ph: list[int],
+    device: str,
+    assist_text: Optional[str] = None,
+    assist_text_weight: float = 0.7,
+) -> torch.Tensor:
+    """
+    日本語のテキストから BERT の特徴量を抽出する
+
+    Args:
+        text (str): 日本語のテキスト
+        word2ph (list[int]): 元のテキストの各文字に音素が何個割り当てられるかを表すリスト
+        device (str): 推論に利用するデバイス
+        assist_text (Optional[str], optional): 補助テキスト (デフォルト: None)
+        assist_text_weight (float, optional): 補助テキストの重み (デフォルト: 0.7)
+
+    Returns:
+        torch.Tensor: BERT の特徴量
+    """
+
+    # 各単語が何文字かを作る `word2ph` を使う必要があるので、読めない文字は必ず無視する
+    # でないと `word2ph` の結果とテキストの文字数結果が整合性が取れない
+    text = "".join(text_to_sep_kata(text, raise_yomi_error=False)[0])
+    if assist_text:
+        assist_text = "".join(text_to_sep_kata(assist_text, raise_yomi_error=False)[0])
+
+    if device == "cuda" and not torch.cuda.is_available():
+        device = "cpu"
+    model = bert_models.load_model(Languages.JP).to(device)  # type: ignore
+
+    style_res_mean = None
+    with torch.no_grad():
+        tokenizer = bert_models.load_tokenizer(Languages.JP)
+        inputs = tokenizer(text, return_tensors="pt")
+        for i in inputs:
+            inputs[i] = inputs[i].to(device)  # type: ignore
+        res = model(**inputs, output_hidden_states=True)
+        res = torch.cat(res["hidden_states"][-3:-2], -1)[0].cpu()
+        if assist_text:
+            style_inputs = tokenizer(assist_text, return_tensors="pt")
+            for i in style_inputs:
+                style_inputs[i] = style_inputs[i].to(device)  # type: ignore
+            style_res = model(**style_inputs, output_hidden_states=True)
+            style_res = torch.cat(style_res["hidden_states"][-3:-2], -1)[0].cpu()
+            style_res_mean = style_res.mean(0)
+
+    assert len(word2ph) == len(text) + 2, text
+    word2phone = word2ph
+    phone_level_feature = []
+    for i in range(len(word2phone)):
+        if assist_text:
+            assert style_res_mean is not None
+            repeat_feature = (
+                res[i].repeat(word2phone[i], 1) * (1 - assist_text_weight)
+                + style_res_mean.repeat(word2phone[i], 1) * assist_text_weight
+            )
+        else:
+            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

style_bert_vits2/nlp/japanese/g2p.py

@@ -0,0 +1,491 @@
+import re
+
+from style_bert_vits2.constants import Languages
+from style_bert_vits2.logging import logger
+from style_bert_vits2.nlp import bert_models
+from style_bert_vits2.nlp.japanese import pyopenjtalk_worker as pyopenjtalk
+from style_bert_vits2.nlp.japanese.mora_list import MORA_KATA_TO_MORA_PHONEMES
+from style_bert_vits2.nlp.japanese.normalizer import replace_punctuation
+from style_bert_vits2.nlp.symbols import PUNCTUATIONS
+
+
+def g2p(
+    norm_text: str,
+    use_jp_extra: bool = True,
+    raise_yomi_error: bool = False
+) -> tuple[list[str], list[int], list[int]]:
+    """
+    他で使われるメインの関数。`normalize_text()` で正規化された `norm_text` を受け取り、
+    - phones: 音素のリスト（ただし `!` や `,` や `.` など punctuation が含まれうる）
+    - tones: アクセントのリスト、0（低）と1（高）からなり、phones と同じ長さ
+    - word2ph: 元のテキストの各文字に音素が何個割り当てられるかを表すリスト
+    のタプルを返す。
+    ただし `phones` と `tones` の最初と終わりに `_` が入り、応じて `word2ph` の最初と最後に 1 が追加される。
+
+    Args:
+        norm_text (str): 正規化されたテキスト
+        use_jp_extra (bool, optional): False の場合、「ん」の音素を「N」ではなく「n」とする。Defaults to True.
+        raise_yomi_error (bool, optional): False の場合、読めない文字が消えたような扱いとして処理される。Defaults to False.
+
+    Returns:
+        tuple[list[str], list[int], list[int]]: 音素のリスト、アクセントのリスト、word2ph のリスト
+    """
+
+    # pyopenjtalk のフルコンテキストラベルを使ってアクセントを取り出すと、punctuation の位置が消えてしまい情報が失われてしまう：
+    # 「こんにちは、世界。」と「こんにちは！世界。」と「こんにちは！！！？？？世界……。」は全て同じになる。
+    # よって、まず punctuation 無しの音素とアクセントのリストを作り、
+    # それとは別に pyopenjtalk.run_frontend() で得られる音素リスト（こちらは punctuation が保持される）を使い、
+    # アクセント割当をしなおすことによって punctuation を含めた音素とアクセントのリストを作る。
+
+    # punctuation がすべて消えた、音素とアクセントのタプルのリスト（「ん」は「N」）
+    phone_tone_list_wo_punct = __g2phone_tone_wo_punct(norm_text)
+
+    # sep_text: 単語単位の単語のリスト、読めない文字があったら raise_yomi_error なら例外、そうでないなら読めない文字が消えて返ってくる
+    # sep_kata: 単語単位の単語のカタカナ読みのリスト
+    sep_text, sep_kata = text_to_sep_kata(norm_text, raise_yomi_error=raise_yomi_error)
+
+    # sep_phonemes: 各単語ごとの音素のリストのリスト
+    sep_phonemes = __handle_long([__kata_to_phoneme_list(i) for i in sep_kata])
+
+    # phone_w_punct: sep_phonemes を結合した、punctuation を元のまま保持した音素列
+    phone_w_punct: list[str] = []
+    for i in sep_phonemes:
+        phone_w_punct += i
+
+    # punctuation 無しのアクセント情報を使って、punctuation を含めたアクセント情報を作る
+    phone_tone_list = __align_tones(phone_w_punct, phone_tone_list_wo_punct)
+    # logger.debug(f"phone_tone_list:\n{phone_tone_list}")
+
+    # word2ph は厳密な解答は不可能なので（「今日」「眼鏡」等の熟字訓が存在）、
+    # Bert-VITS2 では、単語単位の分割を使って、単語の文字ごとにだいたい均等に音素を分配する
+
+    # sep_text から、各単語を1文字1文字分割して、文字のリスト（のリスト）を作る
+    sep_tokenized: list[list[str]] = []
+    for i in sep_text:
+        if i not in PUNCTUATIONS:
+            sep_tokenized.append(
+                bert_models.load_tokenizer(Languages.JP).tokenize(i)
+            )  # ここでおそらく`i`が文字単位に分割される
+        else:
+            sep_tokenized.append([i])
+
+    # 各単語について、音素の数と文字の数を比較して、均等っぽく分配する
+    word2ph = []
+    for token, phoneme in zip(sep_tokenized, sep_phonemes):
+        phone_len = len(phoneme)
+        word_len = len(token)
+        word2ph += __distribute_phone(phone_len, word_len)
+
+    # 最初と最後に `_` 記号を追加、アクセントは 0（低）、word2ph もそれに合わせて追加
+    phone_tone_list = [("_", 0)] + phone_tone_list + [("_", 0)]
+    word2ph = [1] + word2ph + [1]
+
+    phones = [phone for phone, _ in phone_tone_list]
+    tones = [tone for _, tone in phone_tone_list]
+
+    assert len(phones) == sum(word2ph), f"{len(phones)} != {sum(word2ph)}"
+
+    # use_jp_extra でない場合は「N」を「n」に変換
+    if not use_jp_extra:
+        phones = [phone if phone != "N" else "n" for phone in phones]
+
+    return phones, tones, word2ph
+
+
+def text_to_sep_kata(
+    norm_text: str,
+    raise_yomi_error: bool = False
+) -> tuple[list[str], list[str]]:
+    """
+    `normalize_text` で正規化済みの `norm_text` を受け取り、それを単語分割し、
+    分割された単語リストとその読み（カタカナ or 記号1文字）のリストのタプルを返す。
+    単語分割結果は、`g2p()` の `word2ph` で1文字あたりに割り振る音素記号の数を決めるために使う。
+    例:
+    `私はそう思う!って感じ?` →
+    ["私", "は", "そう", "思う", "!", "って", "感じ", "?"], ["ワタシ", "ワ", "ソー", "オモウ", "!", "ッテ", "カンジ", "?"]
+
+    Args:
+        norm_text (str): 正規化されたテキスト
+        raise_yomi_error (bool, optional): False の場合、読めない文字が消えたような扱いとして処理される。Defaults to False.
+
+    Returns:
+        tuple[list[str], list[str]]: 分割された単語リストと、その読み（カタカナ or 記号1文字）のリスト
+    """
+
+    # parsed: OpenJTalkの解析結果
+    parsed = pyopenjtalk.run_frontend(norm_text)
+    sep_text: list[str] = []
+    sep_kata: list[str] = []
+
+    for parts in parsed:
+        # word: 実際の単語の文字列
+        # yomi: その読み、但し無声化サインの`’`は除去
+        word, yomi = replace_punctuation(parts["string"]), parts["pron"].replace(
+            "’", ""
+        )
+        """
+        ここで `yomi` の取りうる値は以下の通りのはず。
+        - `word` が通常単語 → 通常の読み（カタカナ）
+            （カタカナからなり、長音記号も含みうる、`アー` 等）
+        - `word` が `ー` から始まる → `ーラー` や `ーーー` など
+        - `word` が句読点や空白等 → `、`
+        - `word` が punctuation の繰り返し → 全角にしたもの
+        基本的に punctuation は1文字ずつ分かれるが、何故かある程度連続すると1つにまとまる。
+        他にも `word` が読めないキリル文字アラビア文字等が来ると `、` になるが、正規化でこの場合は起きないはず。
+        また元のコードでは `yomi` が空白の場合の処理があったが、これは起きないはず。
+        処理すべきは `yomi` が `、` の場合のみのはず。
+        """
+        assert yomi != "", f"Empty yomi: {word}"
+        if yomi == "、":
+            # word は正規化されているので、`.`, `,`, `!`, `'`, `-`, `--` のいずれか
+            if not set(word).issubset(set(PUNCTUATIONS)):  # 記号繰り返しか判定
+                # ここは pyopenjtalk が読めない文字等のときに起こる
+                if raise_yomi_error:
+                    raise YomiError(f"Cannot read: {word} in:\n{norm_text}")
+                logger.warning(f"Ignoring unknown: {word} in:\n{norm_text}")
+                continue
+            # yomi は元の記号のままに変更
+            yomi = word
+        elif yomi == "？":
+            assert word == "?", f"yomi `？` comes from: {word}"
+            yomi = "?"
+        sep_text.append(word)
+        sep_kata.append(yomi)
+
+    return sep_text, sep_kata
+
+
+def __g2phone_tone_wo_punct(text: str) -> list[tuple[str, int]]:
+    """
+    テキストに対して、音素とアクセント（0か1）のペアのリストを返す。
+    ただし「!」「.」「?」等の非音素記号 (punctuation) は全て消える（ポーズ記号も残さない）。
+    非音素記号を含める処理は `align_tones()` で行われる。
+    また「っ」は「q」に、「ん」は「N」に変換される。
+    例: "こんにちは、世界ー。。元気？！" →
+    [('k', 0), ('o', 0), ('N', 1), ('n', 1), ('i', 1), ('ch', 1), ('i', 1), ('w', 1), ('a', 1), ('s', 1), ('e', 1), ('k', 0), ('a', 0), ('i', 0), ('i', 0), ('g', 1), ('e', 1), ('N', 0), ('k', 0), ('i', 0)]
+
+    Args:
+        text (str): テキスト
+
+    Returns:
+        list[tuple[str, int]]: 音素とアクセントのペアのリスト
+    """
+
+    prosodies = __pyopenjtalk_g2p_prosody(text, drop_unvoiced_vowels=True)
+    # logger.debug(f"prosodies: {prosodies}")
+    result: list[tuple[str, int]] = []
+    current_phrase: list[tuple[str, int]] = []
+    current_tone = 0
+
+    for i, letter in enumerate(prosodies):
+        # 特殊記号の処理
+
+        # 文頭記号、無視する
+        if letter == "^":
+            assert i == 0, "Unexpected ^"
+        # アクセント句の終わりに来る記号
+        elif letter in ("$", "?", "_", "#"):
+            # 保持しているフレーズを、アクセント数値を 0-1 に修正し結果に追加
+            result.extend(__fix_phone_tone(current_phrase))
+            # 末尾に来る終了記号、無視（文中の疑問文は `_` になる）
+            if letter in ("$", "?"):
+                assert i == len(prosodies) - 1, f"Unexpected {letter}"
+            # あとは "_"（ポーズ）と "#"（アクセント句の境界）のみ
+            # これらは残さず、次のアクセント句に備える。
+            current_phrase = []
+            # 0 を基準点にしてそこから上昇・下降する（負の場合は上の `fix_phone_tone` で直る）
+            current_tone = 0
+        # アクセント上昇記号
+        elif letter == "[":
+            current_tone = current_tone + 1
+        # アクセント下降記号
+        elif letter == "]":
+            current_tone = current_tone - 1
+        # それ以外は通常の音素
+        else:
+            if letter == "cl":  # 「っ」の処理
+                letter = "q"
+            # elif letter == "N":  # 「ん」の処理
+            #     letter = "n"
+            current_phrase.append((letter, current_tone))
+
+    return result
+
+
+def __pyopenjtalk_g2p_prosody(text: str, drop_unvoiced_vowels: bool = True) -> list[str]:
+    """
+    ESPnet の実装から引用、変更点無し。「ん」は「N」なことに注意。
+    ref: https://github.com/espnet/espnet/blob/master/espnet2/text/phoneme_tokenizer.py
+    ------------------------------------------------------------------------------------------
+
+    Extract phoneme + prosoody symbol sequence from input full-context labels.
+
+    The algorithm is based on `Prosodic features control by symbols as input of
+    sequence-to-sequence acoustic modeling for neural TTS`_ with some r9y9's tweaks.
+
+    Args:
+        text (str): Input text.
+        drop_unvoiced_vowels (bool): whether to drop unvoiced vowels.
+
+    Returns:
+        List[str]: List of phoneme + prosody symbols.
+
+    Examples:
+        >>> from espnet2.text.phoneme_tokenizer import pyopenjtalk_g2p_prosody
+        >>> pyopenjtalk_g2p_prosody("こんにちは。")
+        ['^', 'k', 'o', '[', 'N', 'n', 'i', 'ch', 'i', 'w', 'a', '$']
+
+    .. _`Prosodic features control by symbols as input of sequence-to-sequence acoustic
+        modeling for neural TTS`: https://doi.org/10.1587/transinf.2020EDP7104
+    """
+
+    def _numeric_feature_by_regex(regex: str, s: str) -> int:
+        match = re.search(regex, s)
+        if match is None:
+            return -50
+        return int(match.group(1))
+
+    labels = pyopenjtalk.make_label(pyopenjtalk.run_frontend(text))
+    N = len(labels)
+
+    phones = []
+    for n in range(N):
+        lab_curr = labels[n]
+
+        # current phoneme
+        p3 = re.search(r"\-(.*?)\+", lab_curr).group(1)  # type: ignore
+        # deal unvoiced vowels as normal vowels
+        if drop_unvoiced_vowels and p3 in "AEIOU":
+            p3 = p3.lower()
+
+        # deal with sil at the beginning and the end of text
+        if p3 == "sil":
+            assert n == 0 or n == N - 1
+            if n == 0:
+                phones.append("^")
+            elif n == N - 1:
+                # check question form or not
+                e3 = _numeric_feature_by_regex(r"!(\d+)_", lab_curr)
+                if e3 == 0:
+                    phones.append("$")
+                elif e3 == 1:
+                    phones.append("?")
+            continue
+        elif p3 == "pau":
+            phones.append("_")
+            continue
+        else:
+            phones.append(p3)
+
+        # accent type and position info (forward or backward)
+        a1 = _numeric_feature_by_regex(r"/A:([0-9\-]+)\+", lab_curr)
+        a2 = _numeric_feature_by_regex(r"\+(\d+)\+", lab_curr)
+        a3 = _numeric_feature_by_regex(r"\+(\d+)/", lab_curr)
+
+        # number of mora in accent phrase
+        f1 = _numeric_feature_by_regex(r"/F:(\d+)_", lab_curr)
+
+        a2_next = _numeric_feature_by_regex(r"\+(\d+)\+", labels[n + 1])
+        # accent phrase border
+        if a3 == 1 and a2_next == 1 and p3 in "aeiouAEIOUNcl":
+            phones.append("#")
+        # pitch falling
+        elif a1 == 0 and a2_next == a2 + 1 and a2 != f1:
+            phones.append("]")
+        # pitch rising
+        elif a2 == 1 and a2_next == 2:
+            phones.append("[")
+
+    return phones
+
+
+def __fix_phone_tone(phone_tone_list: list[tuple[str, int]]) -> list[tuple[str, int]]:
+    """
+    `phone_tone_list` の tone（アクセントの値）を 0 か 1 の範囲に修正する。
+    例: [(a, 0), (i, -1), (u, -1)] → [(a, 1), (i, 0), (u, 0)]
+
+    Args:
+        phone_tone_list (list[tuple[str, int]]): 音素とアクセントのペアのリスト
+
+    Returns:
+        list[tuple[str, int]]: 修正された音素とアクセントのペアのリスト
+    """
+
+    tone_values = set(tone for _, tone in phone_tone_list)
+    if len(tone_values) == 1:
+        assert tone_values == {0}, tone_values
+        return phone_tone_list
+    elif len(tone_values) == 2:
+        if tone_values == {0, 1}:
+            return phone_tone_list
+        elif tone_values == {-1, 0}:
+            return [
+                (letter, 0 if tone == -1 else 1) for letter, tone in phone_tone_list
+            ]
+        else:
+            raise ValueError(f"Unexpected tone values: {tone_values}")
+    else:
+        raise ValueError(f"Unexpected tone values: {tone_values}")
+
+
+def __handle_long(sep_phonemes: list[list[str]]) -> list[list[str]]:
+    """
+    フレーズごとに分かれた音素（長音記号がそのまま）のリストのリスト `sep_phonemes` を受け取り、
+    その長音記号を処理して、音素のリストのリストを返す。
+    基本的には直前の音素を伸ばすが、直前の音素が母音でない場合もしくは冒頭の場合は、
+    おそらく長音記号とダッシュを勘違いしていると思われるので、ダッシュに対応する音素 `-` に変換する。
+
+    Args:
+        sep_phonemes (list[list[str]]): フレーズごとに分かれた音素のリストのリスト
+
+    Returns:
+        list[list[str]]: 長音記号を処理した音素のリストのリスト
+    """
+
+    # 母音の集合 (便宜上「ん」を含める)
+    VOWELS = {"a", "i", "u", "e", "o", "N"}
+
+    for i in range(len(sep_phonemes)):
+        if len(sep_phonemes[i]) == 0:
+            # 空白文字等でリストが空の場合
+            continue
+        if sep_phonemes[i][0] == "ー":
+            if i != 0:
+                prev_phoneme = sep_phonemes[i - 1][-1]
+                if prev_phoneme in VOWELS:
+                    # 母音と「ん」のあとの伸ばし棒なので、その母音に変換
+                    sep_phonemes[i][0] = sep_phonemes[i - 1][-1]
+                else:
+                    # 「。ーー」等おそらく予期しない長音記号
+                    # ダッシュの勘違いだと思われる
+                    sep_phonemes[i][0] = "-"
+            else:
+                # 冒頭に長音記号が来ていおり、これはダッシュの勘違いと思われる
+                sep_phonemes[i][0] = "-"
+        if "ー" in sep_phonemes[i]:
+            for j in range(len(sep_phonemes[i])):
+                if sep_phonemes[i][j] == "ー":
+                    sep_phonemes[i][j] = sep_phonemes[i][j - 1][-1]
+
+    return sep_phonemes
+
+
+def __kata_to_phoneme_list(text: str) -> list[str]:
+    """
+    原則カタカナの `text` を受け取り、それをそのままいじらずに音素記号のリストに変換。
+    注意点：
+    - punctuation かその繰り返しが来た場合、punctuation たちをそのままリストにして返す。
+    - 冒頭に続く「ー」はそのまま「ー」のままにする（`handle_long()` で処理される）
+    - 文中の「ー」は前の音素記号の最後の音素記号に変換される。
+    例：
+    `ーーソーナノカーー` → ["ー", "ー", "s", "o", "o", "n", "a", "n", "o", "k", "a", "a", "a"]
+    `?` → ["?"]
+    `!?!?!?!?!` → ["!", "?", "!", "?", "!", "?", "!", "?", "!"]
+
+    Args:
+        text (str): カタカナのテキスト
+
+    Returns:
+        list[str]: 音素記号のリスト
+    """
+
+    if set(text).issubset(set(PUNCTUATIONS)):
+        return list(text)
+    # `text` がカタカナ（`ー`含む）のみからなるかどうかをチェック
+    if re.fullmatch(r"[\u30A0-\u30FF]+", text) is None:
+        raise ValueError(f"Input must be katakana only: {text}")
+    sorted_keys = sorted(MORA_KATA_TO_MORA_PHONEMES.keys(), key=len, reverse=True)
+    pattern = "|".join(map(re.escape, sorted_keys))
+
+    def mora2phonemes(mora: str) -> str:
+        cosonant, vowel = MORA_KATA_TO_MORA_PHONEMES[mora]
+        if cosonant is None:
+            return f" {vowel}"
+        return f" {cosonant} {vowel}"
+
+    spaced_phonemes = re.sub(pattern, lambda m: mora2phonemes(m.group()), text)
+
+    # 長音記号「ー」の処理
+    long_pattern = r"(\w)(ー*)"
+    long_replacement = lambda m: m.group(1) + (" " + m.group(1)) * len(m.group(2))  # type: ignore
+    spaced_phonemes = re.sub(long_pattern, long_replacement, spaced_phonemes)
+
+    return spaced_phonemes.strip().split(" ")
+
+
+def __align_tones(
+    phones_with_punct: list[str],
+    phone_tone_list: list[tuple[str, int]]
+) -> list[tuple[str, int]]:
+    """
+    例: …私は、、そう思う。
+    phones_with_punct:
+        [".", ".", ".", "w", "a", "t", "a", "sh", "i", "w", "a", ",", ",", "s", "o", "o", "o", "m", "o", "u", "."]
+    phone_tone_list:
+        [("w", 0), ("a", 0), ("t", 1), ("a", 1), ("sh", 1), ("i", 1), ("w", 1), ("a", 1), ("_", 0), ("s", 0), ("o", 0), ("o", 1), ("o", 1), ("m", 1), ("o", 1), ("u", 0))]
+    Return:
+        [(".", 0), (".", 0), (".", 0), ("w", 0), ("a", 0), ("t", 1), ("a", 1), ("sh", 1), ("i", 1), ("w", 1), ("a", 1), (",", 0), (",", 0), ("s", 0), ("o", 0), ("o", 1), ("o", 1), ("m", 1), ("o", 1), ("u", 0), (".", 0)]
+
+    Args:
+        phones_with_punct (list[str]): punctuation を含む音素のリスト
+        phone_tone_list (list[tuple[str, int]]): punctuation を含まない音素とアクセントのペアのリスト
+
+    Returns:
+        list[tuple[str, int]]: punctuation を含む音素とアクセントのペアのリスト
+    """
+
+    result: list[tuple[str, int]] = []
+    tone_index = 0
+    for phone in phones_with_punct:
+        if tone_index >= len(phone_tone_list):
+            # 余った punctuation がある場合 → (punctuation, 0) を追加
+            result.append((phone, 0))
+        elif phone == phone_tone_list[tone_index][0]:
+            # phone_tone_list の現在の音素と一致する場合 → tone をそこから取得、(phone, tone) を追加
+            result.append((phone, phone_tone_list[tone_index][1]))
+            # 探す index を1つ進める
+            tone_index += 1
+        elif phone in PUNCTUATIONS:
+            # phone が punctuation の場合 → (phone, 0) を追加
+            result.append((phone, 0))
+        else:
+            logger.debug(f"phones: {phones_with_punct}")
+            logger.debug(f"phone_tone_list: {phone_tone_list}")
+            logger.debug(f"result: {result}")
+            logger.debug(f"tone_index: {tone_index}")
+            logger.debug(f"phone: {phone}")
+            raise ValueError(f"Unexpected phone: {phone}")
+
+    return result
+
+
+def __distribute_phone(n_phone: int, n_word: int) -> list[int]:
+    """
+    左から右に 1 ずつ振り分け、次にまた左から右に1ずつ増やし、というふうに、
+    音素の数 `n_phone` を単語の数 `n_word` に分配する。
+
+    Args:
+        n_phone (int): 音素の数
+        n_word (int): 単語の数
+
+    Returns:
+        list[int]: 単語ごとの音素の数のリスト
+    """
+
+    phones_per_word = [0] * n_word
+    for _ 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
+
+
+class YomiError(Exception):
+    """
+    OpenJTalk で、読みが正しく取得できない箇所があるときに発生する例外。
+    基本的に「学習の前処理のテキスト処理時」には発生させ、そうでない場合は、
+    ignore_yomi_error=True にしておいて、この例外を発生させないようにする。
+    """
+
+    pass

style_bert_vits2/nlp/japanese/g2p_utils.py

@@ -0,0 +1,90 @@
+from style_bert_vits2.nlp.japanese.g2p import g2p
+from style_bert_vits2.nlp.japanese.mora_list import (
+    MORA_KATA_TO_MORA_PHONEMES,
+    MORA_PHONEMES_TO_MORA_KATA,
+)
+from style_bert_vits2.nlp.symbols import PUNCTUATIONS
+
+
+def g2kata_tone(norm_text: str) -> list[tuple[str, int]]:
+    """
+    テキストからカタカナとアクセントのペアのリストを返す。
+    推論時のみに使われる関数のため、常に `raise_yomi_error=False` を指定して g2p() を呼ぶ仕様になっている。
+
+    Args:
+        norm_text: 正規化されたテキスト。
+
+    Returns:
+        カタカナと音高のリスト。
+    """
+
+    phones, tones, _ = g2p(norm_text, use_jp_extra=True, raise_yomi_error=False)
+    return phone_tone2kata_tone(list(zip(phones, tones)))
+
+
+def phone_tone2kata_tone(phone_tone: list[tuple[str, int]]) -> list[tuple[str, int]]:
+    """
+    phone_tone の phone 部分をカタカナに変換する。ただし最初と最後の ("_", 0) は無視する。
+
+    Args:
+        phone_tone: 音素と音高のリスト。
+
+    Returns:
+        カタカナと音高のリスト。
+    """
+
+    # 子音の集合
+    CONSONANTS = set([
+        consonant
+        for consonant, _ in MORA_KATA_TO_MORA_PHONEMES.values()
+        if consonant is not None
+    ])
+
+    phone_tone = phone_tone[1:]  # 最初の("_", 0)を無視
+    phones = [phone for phone, _ in phone_tone]
+    tones = [tone for _, tone in phone_tone]
+    result: list[tuple[str, int]] = []
+    current_mora = ""
+    for phone, next_phone, tone, next_tone in zip(phones, phones[1:], tones, tones[1:]):
+        # zip の関係で最後の ("_", 0) は無視されている
+        if phone in PUNCTUATIONS:
+            result.append((phone, tone))
+            continue
+        if phone in CONSONANTS:  # n以外の子音の場合
+            assert current_mora == "", f"Unexpected {phone} after {current_mora}"
+            assert tone == next_tone, f"Unexpected {phone} tone {tone} != {next_tone}"
+            current_mora = phone
+        else:
+            # phoneが母音もしくは「N」
+            current_mora += phone
+            result.append((MORA_PHONEMES_TO_MORA_KATA[current_mora], tone))
+            current_mora = ""
+
+    return result
+
+
+def kata_tone2phone_tone(kata_tone: list[tuple[str, int]]) -> list[tuple[str, int]]:
+    """
+    `phone_tone2kata_tone()` の逆の変換を行う。
+
+    Args:
+        kata_tone: カタカナと音高のリスト。
+
+    Returns:
+        音素と音高のリスト。
+    """
+
+    result: list[tuple[str, int]] = [("_", 0)]
+    for mora, tone in kata_tone:
+        if mora in PUNCTUATIONS:
+            result.append((mora, tone))
+        else:
+            consonant, vowel = MORA_KATA_TO_MORA_PHONEMES[mora]
+            if consonant is None:
+                result.append((vowel, tone))
+            else:
+                result.append((consonant, tone))
+                result.append((vowel, tone))
+    result.append(("_", 0))
+
+    return result

style_bert_vits2/nlp/japanese/mora_list.py

@@ -1,10 +1,10 @@
 """
-VOICEVOXのソースコードからお借りして最低限に改造したコード。
+以下のコードは VOICEVOX のソースコードからお借りし最低限の改造を行ったもの。
 https://github.com/VOICEVOX/voicevox_engine/blob/master/voicevox_engine/tts_pipeline/mora_list.py
 """
 
 """
-以下のモーラ対応表はOpenJTalkのソースコードから取得し、
+以下のモーラ対応表は OpenJTalk のソースコードから取得し、
 カタカナ表記とモーラが一対一対応するように改造した。
 ライセンス表記：
 -----------------------------------------------------------------
@@ -46,13 +46,15 @@ OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
 OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 POSSIBILITY OF SUCH DAMAGE.
 """
+
 from typing import Optional
 
-# (カタカナ, 子音, 母音)の順。子音がない場合はNoneを入れる。
+
+# (カタカナ, 子音, 母音)の順。子音がない場合は None を入れる。
 # 但し「ン」と「ッ」は母音のみという扱いで、「ン」は「N」、「ッ」は「q」とする。
 # （元々「ッ」は「cl」）
-# また「デェ = dy e」はpyopenjtalkの出力（de e）と合わないため削除
+# また「デェ = dy e」は pyopenjtalk の出力（de e）と合わないため削除
-_mora_list_minimum: list[tuple[str, Optional[str], str]] = [
+__MORA_LIST_MINIMUM: list[tuple[str, Optional[str], str]] = [
     ("ヴォ", "v", "o"),
     ("ヴェ", "v", "e"),
     ("ヴィ", "v", "i"),
@@ -199,7 +201,7 @@ _mora_list_minimum: list[tuple[str, Optional[str], str]] = [
     ("イ", None, "i"),
     ("ア", None, "a"),
 ]
-_mora_list_additional: list[tuple[str, Optional[str], str]] = [
+__MORA_LIST_ADDITIONAL: list[tuple[str, Optional[str], str]] = [
     ("ヴョ", "by", "o"),
     ("ヴュ", "by", "u"),
     ("ヴャ", "by", "a"),
@@ -220,13 +222,15 @@ _mora_list_additional: list[tuple[str, Optional[str], str]] = [
     ("ァ", None, "a"),
 ]
 
+# モーラの音素表記とカタカナの対応表
 # 例: "vo" -> "ヴォ", "a" -> "ア"
-mora_phonemes_to_mora_kata: dict[str, str] = {
+MORA_PHONEMES_TO_MORA_KATA: dict[str, str] = {
-    (consonant or "") + vowel: kana for [kana, consonant, vowel] in _mora_list_minimum
+    (consonant or "") + vowel: kana for [kana, consonant, vowel] in __MORA_LIST_MINIMUM
 }
 
+# モーラのカタカナ表記と音素の対応表
 # 例: "ヴォ" -> ("v", "o"), "ア" -> (None, "a")
-mora_kata_to_mora_phonemes: dict[str, tuple[Optional[str], str]] = {
+MORA_KATA_TO_MORA_PHONEMES: dict[str, tuple[Optional[str], str]] = {
     kana: (consonant, vowel)
-    for [kana, consonant, vowel] in _mora_list_minimum + _mora_list_additional
+    for [kana, consonant, vowel] in __MORA_LIST_MINIMUM + __MORA_LIST_ADDITIONAL
 }

style_bert_vits2/nlp/japanese/normalizer.py

@@ -0,0 +1,161 @@
+import re
+import unicodedata
+from num2words import num2words
+
+from style_bert_vits2.nlp.symbols import PUNCTUATIONS
+
+
+def normalize_text(text: str) -> str:
+    """
+    日本語のテキストを正規化する。
+    結果は、ちょうど次の文字のみからなる：
+    - ひらがな
+    - カタカナ（全角長音記号「ー」が入る！）
+    - 漢字
+    - 半角アルファベット（大文字と小文字）
+    - ギリシャ文字
+    - `.` （句点`。`や`…`の一部や改行等）
+    - `,` （読点`、`や`:`等）
+    - `?` （疑問符`？`）
+    - `!` （感嘆符`！`）
+    - `'` （`「`や`」`等）
+    - `-` （`―`（ダッシュ、長音記号ではない）や`-`等）
+
+    注意点:
+    - 三点リーダー`…`は`...`に変換される（`なるほど…。` → `なるほど....`）
+    - 数字は漢字に変換される（`1,100円` → `千百円`、`52.34` → `五十二点三四`）
+    - 読点や疑問符等の位置・個数等は保持される（`??あ、、！！！` → `??あ,,!!!`）
+
+    Args:
+        text (str): 正規化するテキスト
+
+    Returns:
+        str: 正規化されたテキスト
+    """
+
+    res = unicodedata.normalize("NFKC", text)  # ここでアルファベットは半角になる
+    res = __convert_numbers_to_words(res)  # 「100円」→「百円」等
+    # 「～」と「~」も長音記号として扱う
+    res = res.replace("~", "ー")
+    res = res.replace("～", "ー")
+
+    res = replace_punctuation(res)  # 句読点等正規化、読めない文字を削除
+
+    # 結合文字の濁点・半濁点を削除
+    # 通常の「ば」等はそのままのこされる、「あ゛」は上で「あ゙」になりここで「あ」になる
+    res = res.replace("\u3099", "")  # 結合文字の濁点を削除、る゙ → る
+    res = res.replace("\u309A", "")  # 結合文字の半濁点を削除、な゚ → な
+    return res
+
+
+def replace_punctuation(text: str) -> str:
+    """
+    句読点等を「.」「,」「!」「?」「'」「-」に正規化し、OpenJTalk で読みが取得できるもののみ残す：
+    漢字・平仮名・カタカナ、アルファベット、ギリシャ文字
+
+    Args:
+        text (str): 正規化するテキスト
+
+    Returns:
+        str: 正規化されたテキスト
+    """
+
+    # 記号類の正規化変換マップ
+    REPLACE_MAP = {
+        "：": ",",
+        "；": ",",
+        "，": ",",
+        "。": ".",
+        "！": "!",
+        "？": "?",
+        "\n": ".",
+        "．": ".",
+        "…": "...",
+        "···": "...",
+        "・・・": "...",
+        "·": ",",
+        "・": ",",
+        "、": ",",
+        "$": ".",
+        "“": "'",
+        "”": "'",
+        '"': "'",
+        "‘": "'",
+        "’": "'",
+        "（": "'",
+        "）": "'",
+        "(": "'",
+        ")": "'",
+        "《": "'",
+        "》": "'",
+        "【": "'",
+        "】": "'",
+        "[": "'",
+        "]": "'",
+        # NFKC 正規化後のハイフン・ダッシュの変種を全て通常半角ハイフン - \u002d に変換
+        "\u02d7": "\u002d",  # ˗, Modifier Letter Minus Sign
+        "\u2010": "\u002d",  # ‐, Hyphen,
+        # "\u2011": "\u002d",  # ‑, Non-Breaking Hyphen, NFKC により \u2010 に変換される
+        "\u2012": "\u002d",  # ‒, Figure Dash
+        "\u2013": "\u002d",  # –, En Dash
+        "\u2014": "\u002d",  # —, Em Dash
+        "\u2015": "\u002d",  # ―, Horizontal Bar
+        "\u2043": "\u002d",  # ⁃, Hyphen Bullet
+        "\u2212": "\u002d",  # −, Minus Sign
+        "\u23af": "\u002d",  # ⎯, Horizontal Line Extension
+        "\u23e4": "\u002d",  # ⏤, Straightness
+        "\u2500": "\u002d",  # ─, Box Drawings Light Horizontal
+        "\u2501": "\u002d",  # ━, Box Drawings Heavy Horizontal
+        "\u2e3a": "\u002d",  # ⸺, Two-Em Dash
+        "\u2e3b": "\u002d",  # ⸻, Three-Em Dash
+        # "～": "-",  # これは長音記号「ー」として扱うよう変更
+        # "~": "-",  # これも長音記号「ー」として扱うよう変更
+        "「": "'",
+        "」": "'",
+    }
+
+    pattern = re.compile("|".join(re.escape(p) for p in REPLACE_MAP.keys()))
+
+    # 句読点を辞書で置換
+    replaced_text = pattern.sub(lambda x: REPLACE_MAP[x.group()], text)
+
+    replaced_text = re.sub(
+        # ↓ ひらがな、カタカナ、漢字
+        r"[^\u3040-\u309F\u30A0-\u30FF\u4E00-\u9FFF\u3400-\u4DBF\u3005"
+        # ↓ 半角アルファベット（大文字と小文字）
+        + r"\u0041-\u005A\u0061-\u007A"
+        # ↓ 全角アルファベット（大文字と小文字）
+        + r"\uFF21-\uFF3A\uFF41-\uFF5A"
+        # ↓ ギリシャ文字
+        + r"\u0370-\u03FF\u1F00-\u1FFF"
+        # ↓ "!", "?", "…", ",", ".", "'", "-", 但し`…`はすでに`...`に変換されている
+        + "".join(PUNCTUATIONS) + r"]+",
+        # 上述以外の文字を削除
+        "",
+        replaced_text,
+    )
+
+    return replaced_text
+
+
+def __convert_numbers_to_words(text: str) -> str:
+    """
+    記号や数字を日本語の文字表現に変換する。
+
+    Args:
+        text (str): 変換するテキスト
+
+    Returns:
+        str: 変換されたテキスト
+    """
+
+    NUMBER_WITH_SEPARATOR_PATTERN = re.compile("[0-9]{1,3}(,[0-9]{3})+")
+    CURRENCY_MAP = {"$": "ドル", "¥": "円", "£": "ポンド", "€": "ユーロ"}
+    CURRENCY_PATTERN = re.compile(r"([$¥£€])([0-9.]*[0-9])")
+    NUMBER_PATTERN = re.compile(r"[0-9]+(\.[0-9]+)?")
+
+    res = NUMBER_WITH_SEPARATOR_PATTERN.sub(lambda m: m[0].replace(",", ""), text)
+    res = CURRENCY_PATTERN.sub(lambda m: m[2] + CURRENCY_MAP.get(m[1], m[1]), res)
+    res = NUMBER_PATTERN.sub(lambda m: num2words(m[0], lang="ja"), res)
+
+    return res

style_bert_vits2/nlp/japanese/pyopenjtalk_worker/__init__.py

@@ -0,0 +1,150 @@
+"""
+Run the pyopenjtalk worker in a separate process
+to avoid user dictionary access error
+"""
+
+from typing import Any, Optional
+
+from style_bert_vits2.logging import logger
+from style_bert_vits2.nlp.japanese.pyopenjtalk_worker.worker_client import WorkerClient
+from style_bert_vits2.nlp.japanese.pyopenjtalk_worker.worker_common import WORKER_PORT
+
+
+WORKER_CLIENT: Optional[WorkerClient] = None
+
+
+# pyopenjtalk interface
+# g2p(): not used
+
+
+def run_frontend(text: str) -> list[dict[str, Any]]:
+    if WORKER_CLIENT is not None:
+        ret = WORKER_CLIENT.dispatch_pyopenjtalk("run_frontend", text)
+        assert isinstance(ret, list)
+        return ret
+    else:
+        # without worker
+        import pyopenjtalk
+        return pyopenjtalk.run_frontend(text)
+
+
+def make_label(njd_features: Any) -> list[str]:
+    if WORKER_CLIENT is not None:
+        ret = WORKER_CLIENT.dispatch_pyopenjtalk("make_label", njd_features)
+        assert isinstance(ret, list)
+        return ret
+    else:
+        # without worker
+        import pyopenjtalk
+        return pyopenjtalk.make_label(njd_features)
+
+
+def mecab_dict_index(path: str, out_path: str, dn_mecab: Optional[str] = None) -> None:
+    if WORKER_CLIENT is not None:
+        WORKER_CLIENT.dispatch_pyopenjtalk("mecab_dict_index", path, out_path, dn_mecab)
+    else:
+        # without worker
+        import pyopenjtalk
+        pyopenjtalk.mecab_dict_index(path, out_path, dn_mecab)
+
+
+def update_global_jtalk_with_user_dict(path: str) -> None:
+    if WORKER_CLIENT is not None:
+        WORKER_CLIENT.dispatch_pyopenjtalk("update_global_jtalk_with_user_dict", path)
+    else:
+        # without worker
+        import pyopenjtalk
+        pyopenjtalk.update_global_jtalk_with_user_dict(path)
+
+
+def unset_user_dict() -> None:
+    if WORKER_CLIENT is not None:
+        WORKER_CLIENT.dispatch_pyopenjtalk("unset_user_dict")
+    else:
+        # without worker
+        import pyopenjtalk
+        pyopenjtalk.unset_user_dict()
+
+
+# initialize module when imported
+
+
+def initialize_worker(port: int = WORKER_PORT) -> None:
+    import atexit
+    import signal
+    import socket
+    import sys
+    import time
+
+    global WORKER_CLIENT
+    if WORKER_CLIENT:
+        return
+
+    client = None
+    try:
+        client = WorkerClient(port)
+    except (socket.timeout, socket.error):
+        logger.debug("try starting pyopenjtalk worker server")
+        import os
+        import subprocess
+
+        worker_pkg_path = os.path.relpath(
+            os.path.dirname(__file__), os.getcwd()
+        ).replace(os.sep, ".")
+        args = [sys.executable, "-m", worker_pkg_path, "--port", str(port)]
+        # new session, new process group
+        if sys.platform.startswith("win"):
+            cf = subprocess.CREATE_NEW_CONSOLE | subprocess.CREATE_NEW_PROCESS_GROUP  # type: ignore
+            si = subprocess.STARTUPINFO()  # type: ignore
+            si.dwFlags |= subprocess.STARTF_USESHOWWINDOW  # type: ignore
+            si.wShowWindow = subprocess.SW_HIDE  # type: ignore
+            subprocess.Popen(args, creationflags=cf, startupinfo=si)
+        else:
+            # align with Windows behavior
+            # start_new_session is same as specifying setsid in preexec_fn
+            subprocess.Popen(args, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, start_new_session=True)
+
+        # wait until server listening
+        count = 0
+        while True:
+            try:
+                client = WorkerClient(port)
+                break
+            except socket.error:
+                time.sleep(0.5)
+                count += 1
+                # 20: max number of retries
+                if count == 20:
+                    raise TimeoutError("サーバーに接続できませんでした")
+
+    logger.debug("pyopenjtalk worker server started")
+    WORKER_CLIENT = client
+    atexit.register(terminate_worker)
+
+    # when the process is killed
+    def signal_handler(signum: int, frame: Any):
+        terminate_worker()
+
+    try:
+        signal.signal(signal.SIGTERM, signal_handler)
+    except ValueError:
+        # signal only works in main thread
+        pass
+
+
+# top-level declaration
+def terminate_worker() -> None:
+    logger.debug("pyopenjtalk worker server terminated")
+    global WORKER_CLIENT
+    if not WORKER_CLIENT:
+        return
+
+    # prepare for unexpected errors
+    try:
+        if WORKER_CLIENT.status() == 1:
+            WORKER_CLIENT.quit_server()
+    except Exception as e:
+        logger.error(e)
+
+    WORKER_CLIENT.close()
+    WORKER_CLIENT = None

style_bert_vits2/nlp/japanese/pyopenjtalk_worker/__main__.py

@@ -1,16 +1,16 @@
-import argparse
+import argparse
-
+
-from .worker_server import WorkerServer
+from style_bert_vits2.nlp.japanese.pyopenjtalk_worker.worker_common import WORKER_PORT
-from .worker_common import WORKER_PORT
+from style_bert_vits2.nlp.japanese.pyopenjtalk_worker.worker_server import WorkerServer
-
+
-
+
-def main():
+def main() -> None:
-    parser = argparse.ArgumentParser()
+    parser = argparse.ArgumentParser()
-    parser.add_argument("--port", type=int, default=WORKER_PORT)
+    parser.add_argument("--port", type=int, default=WORKER_PORT)
-    args = parser.parse_args()
+    args = parser.parse_args()
-    server = WorkerServer()
+    server = WorkerServer()
-    server.start_server(port=args.port)
+    server.start_server(port=args.port)
-
+
-
+
-if __name__ == "__main__":
+if __name__ == "__main__":
-    main()
+    main()

style_bert_vits2/nlp/japanese/pyopenjtalk_worker/worker_client.py

@@ -1,55 +1,63 @@
-from typing import Any
+import socket
-import socket
+from typing import Any, cast
-
+
-from .worker_common import RequestType, receive_data, send_data
+from style_bert_vits2.logging import logger
-
+from style_bert_vits2.nlp.japanese.pyopenjtalk_worker.worker_common import RequestType, receive_data, send_data
-from common.log import logger
+
-
+
-
+class WorkerClient:
-class WorkerClient:
+    """ pyopenjtalk worker client """
-    def __init__(self, port: int) -> None:
+
-        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+
-        # 5: timeout
+    def __init__(self, port: int) -> None:
-        sock.settimeout(5)
+        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-        sock.connect((socket.gethostname(), port))
+        # timeout: 60 seconds
-        self.sock = sock
+        sock.settimeout(60)
-
+        sock.connect((socket.gethostname(), port))
-    def __enter__(self):
+        self.sock = sock
-        return self
+
-
+
-    def __exit__(self, exc_type, exc_value, traceback):
+    def __enter__(self) -> "WorkerClient":
-        self.close()
+        return self
-
+
-    def close(self):
+
-        self.sock.close()
+    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
-
+        self.close()
-    def dispatch_pyopenjtalk(self, func: str, *args, **kwargs):
+
-        data = {
+
-            "request-type": RequestType.PYOPENJTALK,
+    def close(self) -> None:
-            "func": func,
+        self.sock.close()
-            "args": args,
+
-            "kwargs": kwargs,
+
-        }
+    def dispatch_pyopenjtalk(self, func: str, *args: Any, **kwargs: Any) -> Any:
-        logger.trace(f"client sends request: {data}")
+        data = {
-        send_data(self.sock, data)
+            "request-type": RequestType.PYOPENJTALK,
-        logger.trace("client sent request successfully")
+            "func": func,
-        response = receive_data(self.sock)
+            "args": args,
-        logger.trace(f"client received response: {response}")
+            "kwargs": kwargs,
-        return response.get("return")
+        }
-
+        logger.trace(f"client sends request: {data}")
-    def status(self):
+        send_data(self.sock, data)
-        data = {"request-type": RequestType.STATUS}
+        logger.trace("client sent request successfully")
-        logger.trace(f"client sends request: {data}")
+        response = receive_data(self.sock)
-        send_data(self.sock, data)
+        logger.trace(f"client received response: {response}")
-        logger.trace("client sent request successfully")
+        return response.get("return")
-        response = receive_data(self.sock)
+
-        logger.trace(f"client received response: {response}")
+
-        return response.get("client-count")
+    def status(self) -> int:
-
+        data = {"request-type": RequestType.STATUS}
-    def quit_server(self):
+        logger.trace(f"client sends request: {data}")
-        data = {"request-type": RequestType.QUIT_SERVER}
+        send_data(self.sock, data)
-        logger.trace(f"client sends request: {data}")
+        logger.trace("client sent request successfully")
-        send_data(self.sock, data)
+        response = receive_data(self.sock)
-        logger.trace("client sent request successfully")
+        logger.trace(f"client received response: {response}")
-        response = receive_data(self.sock)
+        return cast(int, response.get("client-count"))
-        logger.trace(f"client received response: {response}")
+
+
+    def quit_server(self) -> None:
+        data = {"request-type": RequestType.QUIT_SERVER}
+        logger.trace(f"client sends request: {data}")
+        send_data(self.sock, data)
+        logger.trace("client sent request successfully")
+        response = receive_data(self.sock)
+        logger.trace(f"client received response: {response}")

style_bert_vits2/nlp/japanese/pyopenjtalk_worker/worker_common.py

@@ -1,44 +1,45 @@
-from typing import Any, Optional, Final
+import json
-from enum import IntEnum, auto
+import socket
-import socket
+from enum import IntEnum, auto
-import json
+from typing import Any, Final
-
+
-WORKER_PORT: Final[int] = 7861
+
-HEADER_SIZE: Final[int] = 4
+WORKER_PORT: Final[int] = 7861
-
+HEADER_SIZE: Final[int] = 4
-
+
-class RequestType(IntEnum):
+
-    STATUS = auto()
+class RequestType(IntEnum):
-    QUIT_SERVER = auto()
+    STATUS = auto()
-    PYOPENJTALK = auto()
+    QUIT_SERVER = auto()
-
+    PYOPENJTALK = auto()
-
+
-class ConnectionClosedException(Exception):
+
-    pass
+class ConnectionClosedException(Exception):
-
+    pass
-
+
-# socket communication
+
-
+# socket communication
-
+
-def send_data(sock: socket.socket, data: dict[str, Any]):
+
-    json_data = json.dumps(data).encode()
+def send_data(sock: socket.socket, data: dict[str, Any]):
-    header = len(json_data).to_bytes(HEADER_SIZE, byteorder="big")
+    json_data = json.dumps(data).encode()
-    sock.sendall(header + json_data)
+    header = len(json_data).to_bytes(HEADER_SIZE, byteorder="big")
-
+    sock.sendall(header + json_data)
-
+
-def _receive_until(sock: socket.socket, size: int):
+
-    data = b""
+def __receive_until(sock: socket.socket, size: int):
-    while len(data) < size:
+    data = b""
-        part = sock.recv(size - len(data))
+    while len(data) < size:
-        if part == b"":
+        part = sock.recv(size - len(data))
-            raise ConnectionClosedException("接続が閉じられました")
+        if part == b"":
-        data += part
+            raise ConnectionClosedException("接続が閉じられました")
-
+        data += part
-    return data
+
-
+    return data
-
+
-def receive_data(sock: socket.socket) -> dict[str, Any]:
+
-    header = _receive_until(sock, HEADER_SIZE)
+def receive_data(sock: socket.socket) -> dict[str, Any]:
-    data_length = int.from_bytes(header, byteorder="big")
+    header = __receive_until(sock, HEADER_SIZE)
-    body = _receive_until(sock, data_length)
+    data_length = int.from_bytes(header, byteorder="big")
-    return json.loads(body.decode())
+    body = __receive_until(sock, data_length)
+    return json.loads(body.decode())

style_bert_vits2/nlp/japanese/pyopenjtalk_worker/worker_server.py

@@ -1,118 +1,126 @@
-import pyopenjtalk
+import select
-import socket
+import socket
-import select
+import time
-import time
+from typing import Any, cast
-
+
-from .worker_common import (
+import pyopenjtalk
-    ConnectionClosedException,
+
-    RequestType,
+from style_bert_vits2.logging import logger
-    receive_data,
+from style_bert_vits2.nlp.japanese.pyopenjtalk_worker.worker_common import (
-    send_data,
+    ConnectionClosedException,
-)
+    RequestType,
-
+    receive_data,
-from common.log import logger
+    send_data,
-
+)
-# To make it as fast as possible
+
-# Probably faster than calling getattr every time
+
-_PYOPENJTALK_FUNC_DICT = {
+# To make it as fast as possible
-    "run_frontend": pyopenjtalk.run_frontend,
+# Probably faster than calling getattr every time
-    "make_label": pyopenjtalk.make_label,
+PYOPENJTALK_FUNC_DICT = {
-    "mecab_dict_index": pyopenjtalk.mecab_dict_index,
+    "run_frontend": pyopenjtalk.run_frontend,
-    "update_global_jtalk_with_user_dict": pyopenjtalk.update_global_jtalk_with_user_dict,
+    "make_label": pyopenjtalk.make_label,
-    "unset_user_dict": pyopenjtalk.unset_user_dict,
+    "mecab_dict_index": pyopenjtalk.mecab_dict_index,
-}
+    "update_global_jtalk_with_user_dict": pyopenjtalk.update_global_jtalk_with_user_dict,
-
+    "unset_user_dict": pyopenjtalk.unset_user_dict,
-
+}
-class WorkerServer:
+
-    def __init__(self) -> None:
+
-        self.client_count: int = 0
+class WorkerServer:
-        self.quit: bool = False
+    """ pyopenjtalk worker server """
-
+
-    def handle_request(self, request):
+
-        request_type = None
+    def __init__(self) -> None:
-        try:
+        self.client_count: int = 0
-            request_type = RequestType(request.get("request-type"))
+        self.quit: bool = False
-        except Exception:
+
-            return {
+
-                "success": False,
+    def handle_request(self, request: dict[str, Any]) -> dict[str, Any]:
-                "reason": "request-type is invalid",
+        request_type = None
-            }
+        try:
-
+            request_type = RequestType(cast(int, request.get("request-type")))
-        if request_type:
+        except Exception:
-            if request_type == RequestType.STATUS:
+            return {
-                response = {
+                "success": False,
-                    "success": True,
+                "reason": "request-type is invalid",
-                    "client-count": self.client_count,
+            }
-                }
+
-            elif request_type == RequestType.QUIT_SERVER:
+        response: dict[str, Any] = {}
-                self.quit = True
+        if request_type:
-                response = {"success": True}
+            if request_type == RequestType.STATUS:
-            elif request_type == RequestType.PYOPENJTALK:
+                response = {
-                func_name = request.get("func")
+                    "success": True,
-                assert isinstance(func_name, str)
+                    "client-count": self.client_count,
-                func = _PYOPENJTALK_FUNC_DICT[func_name]
+                }
-                args = request.get("args")
+            elif request_type == RequestType.QUIT_SERVER:
-                kwargs = request.get("kwargs")
+                self.quit = True
-                assert isinstance(args, list)
+                response = {"success": True}
-                assert isinstance(kwargs, dict)
+            elif request_type == RequestType.PYOPENJTALK:
-                ret = func(*args, **kwargs)
+                func_name = request.get("func")
-                response = {"success": True, "return": ret}
+                assert isinstance(func_name, str)
-            else:
+                func = PYOPENJTALK_FUNC_DICT[func_name]
-                # NOT REACHED
+                args = request.get("args")
-                response = request
+                kwargs = request.get("kwargs")
-
+                assert isinstance(args, list)
-        return response
+                assert isinstance(kwargs, dict)
-
+                ret = func(*args, **kwargs)
-    def start_server(self, port: int, no_client_timeout: int = 30):
+                response = {"success": True, "return": ret}
-        logger.info("start pyopenjtalk worker server")
+            else:
-        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as server_socket:
+                # NOT REACHED
-            server_socket.bind((socket.gethostname(), port))
+                response = request
-            server_socket.listen()
+
-            sockets = [server_socket]
+        return response
-            no_client_since = time.time()
+
-            while True:
+
-                if self.client_count == 0:
+    def start_server(self, port: int, no_client_timeout: int = 30) -> None:
-                    if no_client_since is None:
+        logger.info("start pyopenjtalk worker server")
-                        no_client_since = time.time()
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as server_socket:
-                    elif (time.time() - no_client_since) > no_client_timeout:
+            server_socket.bind((socket.gethostname(), port))
-                        logger.info("quit because there is no client")
+            server_socket.listen()
-                        return
+            sockets = [server_socket]
-                else:
+            no_client_since = time.time()
-                    no_client_since = None
+            while True:
-
+                if self.client_count == 0:
-                ready_sockets, _, _ = select.select(sockets, [], [], 0.1)
+                    if no_client_since is None:
-                for sock in ready_sockets:
+                        no_client_since = time.time()
-                    if sock is server_socket:
+                    elif (time.time() - no_client_since) > no_client_timeout:
-                        logger.info("new client connected")
+                        logger.info("quit because there is no client")
-                        client_socket, _ = server_socket.accept()
+                        return
-                        sockets.append(client_socket)
+                else:
-                        self.client_count += 1
+                    no_client_since = None
-                    else:
+
-                        # client
+                ready_sockets, _, _ = select.select(sockets, [], [], 0.1)
-                        try:
+                for sock in ready_sockets:
-                            request = receive_data(sock)
+                    if sock is server_socket:
-                        except Exception as e:
+                        logger.info("new client connected")
-                            sock.close()
+                        client_socket, _ = server_socket.accept()
-                            sockets.remove(sock)
+                        sockets.append(client_socket)
-                            self.client_count -= 1
+                        self.client_count += 1
-                            # unexpected disconnections
+                    else:
-                            if not isinstance(e, ConnectionClosedException):
+                        # client
-                                logger.error(e)
+                        try:
-
+                            request = receive_data(sock)
-                            logger.info("close connection")
+                        except Exception as e:
-                            continue
+                            sock.close()
-
+                            sockets.remove(sock)
-                        logger.trace(f"server received request: {request}")
+                            self.client_count -= 1
-
+                            # unexpected disconnections
-                        response = self.handle_request(request)
+                            if not isinstance(e, ConnectionClosedException):
-                        logger.trace(f"server sends response: {response}")
+                                logger.error(e)
-                        try:
+
-                            send_data(sock, response)
+                            logger.info("close connection")
-                            logger.trace("server sent response successfully")
+                            continue
-                        except Exception:
+
-                            logger.warning(
+                        logger.trace(f"server received request: {request}")
-                                "an exception occurred during sending responce"
+
-                            )
+                        response = self.handle_request(request)
-                        if self.quit:
+                        logger.trace(f"server sends response: {response}")
-                            logger.info("quit pyopenjtalk worker server")
+                        try:
-                            return
+                            send_data(sock, response)
+                            logger.trace("server sent response successfully")
+                        except Exception:
+                            logger.warning(
+                                "an exception occurred during sending responce"
+                            )
+                        if self.quit:
+                            logger.info("quit pyopenjtalk worker server")
+                            return

style_bert_vits2/nlp/japanese/user_dict/README.md

@@ -0,0 +1,27 @@
+
+## ユーザー辞書関連のコードについて
+
+このフォルダに含まれるユーザー辞書関連のコードは、[VOICEVOX ENGINE](https://github.com/VOICEVOX/voicevox_engine) プロジェクトのコードを改変したものを使用しています。  
+VOICEVOX プロジェクトのチームに深く感謝し、その貢献を尊重します。
+
+### 元のコード
+
+- [voicevox_engine/user_dict/](https://github.com/VOICEVOX/voicevox_engine/tree/f181411ec69812296989d9cc583826c22eec87ae/voicevox_engine/user_dict)
+- [voicevox_engine/model.py](https://github.com/VOICEVOX/voicevox_engine/blob/f181411ec69812296989d9cc583826c22eec87ae/voicevox_engine/model.py#L207)
+
+### 改変の詳細
+
+- ファイル名の書き換えおよびそれに伴う import 文の書き換え。
+- VOICEVOX 固有の部分をコメントアウト。
+- mutex を使用している部分をコメントアウト。
+- 参照している pyopenjtalk の違いによるメソッド名の書き換え。
+- UserDictWord の mora_count のデフォルト値を None に指定。
+- `model.py` のうち、必要な Pydantic モデルのみを抽出。
+
+### ライセンス
+
+元の VOICEVOX ENGINE のリポジトリのコードは、LGPL v3 と、ソースコードの公開が不要な別ライセンスのデュアルライセンスの下で使用されています。  
+当プロジェクトにおけるこのモジュールも LGPL ライセンスの下にあります。
+
+詳細については、プロジェクトのルートディレクトリにある [LGPL_LICENSE](/LGPL_LICENSE) ファイルをご参照ください。  
+また、元の VOICEVOX ENGINE プロジェクトのライセンスについては、[こちら](https://github.com/VOICEVOX/voicevox_engine/blob/master/LICENSE) をご覧ください。

style_bert_vits2/nlp/japanese/user_dict/__init__.py

@@ -1,41 +1,34 @@
-# このファイルは、VOICEVOXプロジェクトのVOICEVOX engineからお借りしています。
+"""
-# 引用元:
+このファイルは、VOICEVOX プロジェクトの VOICEVOX ENGINE からお借りしています。
-# https://github.com/VOICEVOX/voicevox_engine/blob/f181411ec69812296989d9cc583826c22eec87ae/voicevox_engine/user_dict/user_dict.py
+引用元: https://github.com/VOICEVOX/voicevox_engine/blob/f181411ec69812296989d9cc583826c22eec87ae/voicevox_engine/user_dict/user_dict.py
-# ライセンス: LGPL-3.0
+ライセンス: LGPL-3.0
-# 詳しくは、このファイルと同じフォルダにあるREADME.mdを参照してください。
+詳しくは、このファイルと同じフォルダにある README.md を参照してください。
+"""
+
 import json
 import sys
-import threading
 import traceback
 from pathlib import Path
 from typing import Dict, List, Optional
 from uuid import UUID, uuid4
 
 import numpy as np
-from .. import pyopenjtalk_worker as pyopenjtalk
-
-pyopenjtalk.initialize()
 from fastapi import HTTPException
 
-from .word_model import UserDictWord, WordTypes
+from style_bert_vits2.constants import DEFAULT_USER_DICT_DIR
-
+from style_bert_vits2.nlp.japanese import pyopenjtalk_worker as pyopenjtalk
-# from ..utility.mutex_utility import mutex_wrapper
+from style_bert_vits2.nlp.japanese.user_dict.word_model import UserDictWord, WordTypes
-# from ..utility.path_utility import engine_root, get_save_dir
+from style_bert_vits2.nlp.japanese.user_dict.part_of_speech_data import MAX_PRIORITY, MIN_PRIORITY, part_of_speech_data
-from .part_of_speech_data import MAX_PRIORITY, MIN_PRIORITY, part_of_speech_data
-from common.constants import USER_DICT_DIR
 
 # root_dir = engine_root()
 # save_dir = get_save_dir()
-root_dir = Path(USER_DICT_DIR)
-save_dir = Path(USER_DICT_DIR)
 
+# if not save_dir.is_dir():
+#     save_dir.mkdir(parents=True)
 
-if not save_dir.is_dir():
+default_dict_path = DEFAULT_USER_DICT_DIR / "default.csv"  # VOICEVOXデフォルト辞書ファイルのパス
-    save_dir.mkdir(parents=True)
+user_dict_path = DEFAULT_USER_DICT_DIR / "user_dict.json"  # ユーザー辞書ファイルのパス
-
+compiled_dict_path = DEFAULT_USER_DICT_DIR / "user.dic"  # コンパイル済み辞書ファイルのパス
-default_dict_path = root_dir / "default.csv"  # VOICEVOXデフォルト辞書ファイルのパス
-user_dict_path = save_dir / "user_dict.json"  # ユーザー辞書ファイルのパス
-compiled_dict_path = save_dir / "user.dic"  # コンパイル済み辞書ファイルのパス
 
 
 # # 同時書き込みの制御
@@ -56,7 +49,7 @@ def _write_to_json(user_dict: Dict[str, UserDictWord], user_dict_path: Path) ->
     """
     converted_user_dict = {}
     for word_uuid, word in user_dict.items():
-        word_dict = word.dict()
+        word_dict = word.model_dump()
         word_dict["cost"] = _priority2cost(
             word_dict["context_id"], word_dict["priority"]
         )
@@ -86,6 +79,7 @@ def update_dict(
     compiled_dict_path : Path
         コンパイル済み辞書ファイルのパス
     """
+
     random_string = uuid4()
     tmp_csv_path = compiled_dict_path.with_suffix(
         f".dict_csv-{random_string}.tmp"

style_bert_vits2/nlp/japanese/user_dict/part_of_speech_data.py

@@ -1,12 +1,13 @@
-# このファイルは、VOICEVOXプロジェクトのVOICEVOX engineからお借りしています。
+"""
-# 引用元:
+このファイルは、VOICEVOX プロジェクトの VOICEVOX ENGINE からお借りしています。
-# https://github.com/VOICEVOX/voicevox_engine/blob/f181411ec69812296989d9cc583826c22eec87ae/voicevox_engine/user_dict/part_of_speech_data.py
+引用元: https://github.com/VOICEVOX/voicevox_engine/blob/f181411ec69812296989d9cc583826c22eec87ae/voicevox_engine/user_dict/part_of_speech_data.py
-# ライセンス: LGPL-3.0
+ライセンス: LGPL-3.0
-# 詳しくは、このファイルと同じフォルダにあるREADME.mdを参照してください。
+詳しくは、このファイルと同じフォルダにある README.md を参照してください。
+"""
 
 from typing import Dict
 
-from .word_model import (
+from style_bert_vits2.nlp.japanese.user_dict.word_model import (
     USER_DICT_MAX_PRIORITY,
     USER_DICT_MIN_PRIORITY,
     PartOfSpeechDetail,

style_bert_vits2/nlp/japanese/user_dict/word_model.py

@@ -1,8 +1,10 @@
-# このファイルは、VOICEVOXプロジェクトのVOICEVOX engineからお借りしています。
+"""
-# 引用元:
+このファイルは、VOICEVOX プロジェクトの VOICEVOX ENGINE からお借りしています。
-# https://github.com/VOICEVOX/voicevox_engine/blob/f181411ec69812296989d9cc583826c22eec87ae/voicevox_engine/model.py#L207
+引用元: https://github.com/VOICEVOX/voicevox_engine/blob/f181411ec69812296989d9cc583826c22eec87ae/voicevox_engine/model.py#L207
-# ライセンス: LGPL-3.0
+ライセンス: LGPL-3.0
-# 詳しくは、このファイルと同じフォルダにあるREADME.mdを参照してください。
+詳しくは、このファイルと同じフォルダにある README.md を参照してください。
+"""
+
 from enum import Enum
 from re import findall, fullmatch
 from typing import List, Optional

style_bert_vits2/nlp/symbols.py

@@ -1,9 +1,14 @@
-punctuation = ["!", "?", "…", ",", ".", "'", "-"]
+# Punctuations
-pu_symbols = punctuation + ["SP", "UNK"]
+PUNCTUATIONS = ["!", "?", "…", ",", ".", "'", "-"]
-pad = "_"
 
-# chinese
+# Punctuations and special tokens
-zh_symbols = [
+PUNCTUATION_SYMBOLS = PUNCTUATIONS + ["SP", "UNK"]
+
+# Padding
+PAD = "_"
+
+# Chinese symbols
+ZH_SYMBOLS = [
     "E",
     "En",
     "a",
@@ -70,10 +75,10 @@ zh_symbols = [
     "EE",
     "OO",
 ]
-num_zh_tones = 6
+NUM_ZH_TONES = 6
 
-# japanese
+# Japanese
-ja_symbols = [
+JP_SYMBOLS = [
     "N",
     "a",
     "a:",
@@ -117,10 +122,10 @@ ja_symbols = [
     "z",
     "zy",
 ]
-num_ja_tones = 2
+NUM_JP_TONES = 2
 
 # English
-en_symbols = [
+EN_SYMBOLS = [
     "aa",
     "ae",
     "ah",
@@ -161,27 +166,29 @@ en_symbols = [
     "z",
     "zh",
 ]
-num_en_tones = 4
+NUM_EN_TONES = 4
 
-# combine all symbols
+# Combine all symbols
-normal_symbols = sorted(set(zh_symbols + ja_symbols + en_symbols))
+NORMAL_SYMBOLS = sorted(set(ZH_SYMBOLS + JP_SYMBOLS + EN_SYMBOLS))
-symbols = [pad] + normal_symbols + pu_symbols
+SYMBOLS = [PAD] + NORMAL_SYMBOLS + PUNCTUATION_SYMBOLS
-sil_phonemes_ids = [symbols.index(i) for i in pu_symbols]
+SIL_PHONEMES_IDS = [SYMBOLS.index(i) for i in PUNCTUATION_SYMBOLS]
 
-# combine all tones
+# Combine all tones
-num_tones = num_zh_tones + num_ja_tones + num_en_tones
+NUM_TONES = NUM_ZH_TONES + NUM_JP_TONES + NUM_EN_TONES
 
-# language maps
+# Language maps
-language_id_map = {"ZH": 0, "JP": 1, "EN": 2}
+LANGUAGE_ID_MAP = {"ZH": 0, "JP": 1, "EN": 2}
-num_languages = len(language_id_map.keys())
+NUM_LANGUAGES = len(LANGUAGE_ID_MAP.keys())
 
-language_tone_start_map = {
+# Language tone start map
+LANGUAGE_TONE_START_MAP = {
     "ZH": 0,
-    "JP": num_zh_tones,
+    "JP": NUM_ZH_TONES,
-    "EN": num_zh_tones + num_ja_tones,
+    "EN": NUM_ZH_TONES + NUM_JP_TONES,
 }
 
+
 if __name__ == "__main__":
-    a = set(zh_symbols)
+    a = set(ZH_SYMBOLS)
-    b = set(en_symbols)
+    b = set(EN_SYMBOLS)
     print(sorted(a & b))

style_bert_vits2/tts_model.py

@@ -0,0 +1,428 @@
+import warnings
+from pathlib import Path
+from typing import Any, Optional, Union
+
+import gradio as gr
+import numpy as np
+import pyannote.audio
+import torch
+from gradio.processing_utils import convert_to_16_bit_wav
+from numpy.typing import NDArray
+from pydantic import BaseModel
+
+from style_bert_vits2.constants import (
+    DEFAULT_ASSIST_TEXT_WEIGHT,
+    DEFAULT_LENGTH,
+    DEFAULT_LINE_SPLIT,
+    DEFAULT_NOISE,
+    DEFAULT_NOISEW,
+    DEFAULT_SDP_RATIO,
+    DEFAULT_SPLIT_INTERVAL,
+    DEFAULT_STYLE,
+    DEFAULT_STYLE_WEIGHT,
+    Languages,
+)
+from style_bert_vits2.models.hyper_parameters import HyperParameters
+from style_bert_vits2.models.infer import get_net_g, infer
+from style_bert_vits2.models.models import SynthesizerTrn
+from style_bert_vits2.models.models_jp_extra import SynthesizerTrn as SynthesizerTrnJPExtra
+from style_bert_vits2.logging import logger
+from style_bert_vits2.voice import adjust_voice
+
+
+class TTSModel:
+    """
+    Style-Bert-Vits2 の音声合成モデルを操作するクラス。
+    モデル/ハイパーパラメータ/スタイルベクトルのパスとデバイスを指定して初期化し、model.infer() メソッドを呼び出すと音声合成を行える。
+    """
+
+
+    def __init__(
+        self,
+        model_path: Path,
+        config_path: Path,
+        style_vec_path: Path,
+        device: str,
+    ) -> None:
+        """
+        Style-Bert-Vits2 の音声合成モデルを初期化する。
+        この時点ではモデルはロードされていない (明示的にロードしたい場合は model.load() を呼び出す)。
+
+        Args:
+            model_path (Path): モデル (.safetensors) のパス
+            config_path (Path): ハイパーパラメータ (config.json) のパス
+            style_vec_path (Path): スタイルベクトル (style_vectors.npy) のパス
+            device (str): 音声合成時に利用するデバイス (cpu, cuda, mps など)
+        """
+
+        self.model_path: Path = model_path
+        self.config_path: Path = config_path
+        self.style_vec_path: Path = style_vec_path
+        self.device: str = device
+        self.hyper_parameters: HyperParameters = HyperParameters.load_from_json(self.config_path)
+        self.spk2id: dict[str, int] = self.hyper_parameters.data.spk2id
+        self.id2spk: dict[int, str] = {v: k for k, v in self.spk2id.items()}
+
+        num_styles: int = self.hyper_parameters.data.num_styles
+        if hasattr(self.hyper_parameters.data, "style2id"):
+            self.style2id: dict[str, int] = self.hyper_parameters.data.style2id
+        else:
+            self.style2id: dict[str, int] = {str(i): i for i in range(num_styles)}
+        if len(self.style2id) != num_styles:
+            raise ValueError(
+                f"Number of styles ({num_styles}) does not match the number of style2id ({len(self.style2id)})"
+            )
+
+        self.__style_vector_inference: Optional[pyannote.audio.Inference] = None
+        self.__style_vectors: NDArray[Any] = np.load(self.style_vec_path)
+        if self.__style_vectors.shape[0] != num_styles:
+            raise ValueError(
+                f"The number of styles ({num_styles}) does not match the number of style vectors ({self.__style_vectors.shape[0]})"
+            )
+
+        self.__net_g: Union[SynthesizerTrn, SynthesizerTrnJPExtra, None] = None
+
+
+    def load(self) -> None:
+        """
+        音声合成モデルをデバイスにロードする。
+        """
+        self.__net_g = get_net_g(
+            model_path = str(self.model_path),
+            version = self.hyper_parameters.version,
+            device = self.device,
+            hps = self.hyper_parameters,
+        )
+
+
+    def __get_style_vector(self, style_id: int, weight: float = 1.0) -> NDArray[Any]:
+        """
+        スタイルベクトルを取得する。
+
+        Args:
+            style_id (int): スタイル ID (0 から始まるインデックス)
+            weight (float, optional): スタイルベクトルの重み. Defaults to 1.0.
+
+        Returns:
+            NDArray[Any]: スタイルベクトル
+        """
+        mean = self.__style_vectors[0]
+        style_vec = self.__style_vectors[style_id]
+        style_vec = mean + (style_vec - mean) * weight
+        return style_vec
+
+
+    def __get_style_vector_from_audio(self, audio_path: str, weight: float = 1.0) -> NDArray[Any]:
+        """
+        音声からスタイルベクトルを推論する。
+
+        Args:
+            audio_path (str): 音声ファイルのパス
+            weight (float, optional): スタイルベクトルの重み. Defaults to 1.0.
+        Returns:
+            NDArray[Any]: スタイルベクトル
+        """
+
+        # スタイルベクトルを取得するための推論モデルを初期化
+        if self.__style_vector_inference is None:
+            self.__style_vector_inference = pyannote.audio.Inference(
+                model = pyannote.audio.Model.from_pretrained("pyannote/wespeaker-voxceleb-resnet34-LM"),
+                window = "whole",
+            )
+            self.__style_vector_inference.to(torch.device(self.device))
+
+        # 音声からスタイルベクトルを推論
+        xvec = self.__style_vector_inference(audio_path)
+        mean = self.__style_vectors[0]
+        xvec = mean + (xvec - mean) * weight
+        return xvec
+
+
+    def infer(
+        self,
+        text: str,
+        language: Languages = Languages.JP,
+        speaker_id: int = 0,
+        reference_audio_path: Optional[str] = None,
+        sdp_ratio: float = DEFAULT_SDP_RATIO,
+        noise: float = DEFAULT_NOISE,
+        noise_w: float = DEFAULT_NOISEW,
+        length: float = DEFAULT_LENGTH,
+        line_split: bool = DEFAULT_LINE_SPLIT,
+        split_interval: float = DEFAULT_SPLIT_INTERVAL,
+        assist_text: Optional[str] = None,
+        assist_text_weight: float = DEFAULT_ASSIST_TEXT_WEIGHT,
+        use_assist_text: bool = False,
+        style: str = DEFAULT_STYLE,
+        style_weight: float = DEFAULT_STYLE_WEIGHT,
+        given_tone: Optional[list[int]] = None,
+        pitch_scale: float = 1.0,
+        intonation_scale: float = 1.0,
+    ) -> tuple[int, NDArray[Any]]:
+        """
+        テキストから音声を合成する。
+
+        Args:
+            text (str): 読み上げるテキスト
+            language (Languages, optional): 言語. Defaults to Languages.JP.
+            speaker_id (int, optional): 話者 ID. Defaults to 0.
+            reference_audio_path (Optional[str], optional): 音声スタイルの参照元の音声ファイルのパス. Defaults to None.
+            sdp_ratio (float, optional): SDP レシオ (値を大きくするとより感情豊かになる傾向がある). Defaults to DEFAULT_SDP_RATIO.
+            noise (float, optional): ノイズの大きさ. Defaults to DEFAULT_NOISE.
+            noise_w (float, optional): ノイズの大きさの重み. Defaults to DEFAULT_NOISEW.
+            length (float, optional): 長さ. Defaults to DEFAULT_LENGTH.
+            line_split (bool, optional): テキストを改行ごとに分割して生成するかどうか. Defaults to DEFAULT_LINE_SPLIT.
+            split_interval (float, optional): 改行ごとに分割する場合の無音 (秒). Defaults to DEFAULT_SPLIT_INTERVAL.
+            assist_text (Optional[str], optional): 感情表現の参照元の補助テキスト. Defaults to None.
+            assist_text_weight (float, optional): 感情表現の補助テキストを適用する強さ. Defaults to DEFAULT_ASSIST_TEXT_WEIGHT.
+            use_assist_text (bool, optional): 音声合成時に感情表現の補助テキストを使用するかどうか. Defaults to False.
+            style (str, optional): 音声スタイル (Neutral, Happy など). Defaults to DEFAULT_STYLE.
+            style_weight (float, optional): 音声スタイルを適用する強さ. Defaults to DEFAULT_STYLE_WEIGHT.
+            given_tone (Optional[list[int]], optional): アクセントのトーンのリスト. Defaults to None.
+            pitch_scale (float, optional): ピッチの高さ (1.0 から変更すると若干音質が低下する). Defaults to 1.0.
+            intonation_scale (float, optional): イントネーションの高さ (1.0 から変更すると若干音質が低下する). Defaults to 1.0.
+
+        Returns:
+            tuple[int, NDArray[Any]]: サンプリングレートと音声データ (16bit PCM)
+        """
+
+        logger.info(f"Start generating audio data from text:\n{text}")
+        if language != "JP" and self.hyper_parameters.version.endswith("JP-Extra"):
+            raise ValueError(
+                "The model is trained with JP-Extra, but the language is not JP"
+            )
+        if reference_audio_path == "":
+            reference_audio_path = None
+        if assist_text == "" or not use_assist_text:
+            assist_text = None
+
+        if self.__net_g is None:
+            self.load()
+        assert self.__net_g is not None
+        if reference_audio_path is None:
+            style_id = self.style2id[style]
+            style_vector = self.__get_style_vector(style_id, style_weight)
+        else:
+            style_vector = self.__get_style_vector_from_audio(
+                reference_audio_path, style_weight
+            )
+        if not line_split:
+            with torch.no_grad():
+                audio = infer(
+                    text = text,
+                    sdp_ratio = sdp_ratio,
+                    noise_scale = noise,
+                    noise_scale_w = noise_w,
+                    length_scale = length,
+                    sid = speaker_id,
+                    language = language,
+                    hps = self.hyper_parameters,
+                    net_g = self.__net_g,
+                    device = self.device,
+                    assist_text = assist_text,
+                    assist_text_weight = assist_text_weight,
+                    style_vec = style_vector,
+                    given_tone = given_tone,
+                )
+        else:
+            texts = text.split("\n")
+            texts = [t for t in texts if t != ""]
+            audios = []
+            with torch.no_grad():
+                for i, t in enumerate(texts):
+                    audios.append(
+                        infer(
+                            text = t,
+                            sdp_ratio = sdp_ratio,
+                            noise_scale = noise,
+                            noise_scale_w = noise_w,
+                            length_scale = length,
+                            sid = speaker_id,
+                            language = language,
+                            hps = self.hyper_parameters,
+                            net_g = self.__net_g,
+                            device = self.device,
+                            assist_text = assist_text,
+                            assist_text_weight = assist_text_weight,
+                            style_vec = style_vector,
+                        )
+                    )
+                    if i != len(texts) - 1:
+                        audios.append(np.zeros(int(44100 * split_interval)))
+                audio = np.concatenate(audios)
+        logger.info("Audio data generated successfully")
+        if not (pitch_scale == 1.0 and intonation_scale == 1.0):
+            _, audio = adjust_voice(
+                fs = self.hyper_parameters.data.sampling_rate,
+                wave = audio,
+                pitch_scale = pitch_scale,
+                intonation_scale = intonation_scale,
+            )
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            audio = convert_to_16_bit_wav(audio)
+        return (self.hyper_parameters.data.sampling_rate, audio)
+
+
+class TTSModelInfo(BaseModel):
+    name: str
+    files: list[str]
+    styles: list[str]
+    speakers: list[str]
+
+
+class TTSModelHolder:
+    """
+    Style-Bert-Vits2 の音声合成モデルを管理するクラス。
+    model_holder.models_info から指定されたディレクトリ内にある音声合成モデルの一覧を取得できる。
+    """
+
+
+    def __init__(self, model_root_dir: Path, device: str) -> None:
+        """
+        Style-Bert-Vits2 の音声合成モデルを管理するクラスを初期化する。
+        音声合成モデルは下記のように配置されていることを前提とする (.safetensors のファイル名は自由) 。
+        ```
+        model_root_dir
+        ├── model-name-1
+        │   ├── config.json
+        │   ├── model-name-1_e160_s14000.safetensors
+        │   └── style_vectors.npy
+        ├── model-name-2
+        │   ├── config.json
+        │   ├── model-name-2_e160_s14000.safetensors
+        │   └── style_vectors.npy
+        └── ...
+        ```
+
+        Args:
+            model_root_dir (Path): 音声合成モデルが配置されているディレクトリのパス
+            device (str): 音声合成時に利用するデバイス (cpu, cuda, mps など)
+        """
+
+        self.root_dir: Path = model_root_dir
+        self.device: str = device
+        self.model_files_dict: dict[str, list[Path]] = {}
+        self.current_model: Optional[TTSModel] = None
+        self.model_names: list[str] = []
+        self.models_info: list[TTSModelInfo] = []
+        self.refresh()
+
+
+    def refresh(self) -> None:
+        """
+        音声合成モデルの一覧を更新する。
+        """
+
+        self.model_files_dict = {}
+        self.model_names = []
+        self.current_model = None
+        self.models_info = []
+
+        model_dirs = [d for d in self.root_dir.iterdir() if d.is_dir()]
+        for model_dir in model_dirs:
+            model_files = [
+                f
+                for f in model_dir.iterdir()
+                if f.suffix in [".pth", ".pt", ".safetensors"]
+            ]
+            if len(model_files) == 0:
+                logger.warning(f"No model files found in {model_dir}, so skip it")
+                continue
+            config_path = model_dir / "config.json"
+            if not config_path.exists():
+                logger.warning(
+                    f"Config file {config_path} not found, so skip {model_dir}"
+                )
+                continue
+            self.model_files_dict[model_dir.name] = model_files
+            self.model_names.append(model_dir.name)
+            hyper_parameters = HyperParameters.load_from_json(config_path)
+            style2id: dict[str, int] = hyper_parameters.data.style2id
+            styles = list(style2id.keys())
+            spk2id: dict[str, int] = hyper_parameters.data.spk2id
+            speakers = list(spk2id.keys())
+            self.models_info.append(TTSModelInfo(
+                name = model_dir.name,
+                files = [str(f) for f in model_files],
+                styles = styles,
+                speakers = speakers,
+            ))
+
+
+    def get_model(self, model_name: str, model_path_str: str) -> TTSModel:
+        """
+        指定された音声合成モデルのインスタンスを取得する。
+        この時点ではモデルはロードされていない (明示的にロードしたい場合は model.load() を呼び出す)。
+
+        Args:
+            model_name (str): 音声合成モデルの名前
+            model_path_str (str): 音声合成モデルのファイルパス (.safetensors)
+
+        Returns:
+            TTSModel: 音声合成モデルのインスタンス
+        """
+
+        model_path = Path(model_path_str)
+        if model_name not in self.model_files_dict:
+            raise ValueError(f"Model `{model_name}` is not found")
+        if model_path not in self.model_files_dict[model_name]:
+            raise ValueError(f"Model file `{model_path}` is not found")
+        if self.current_model is None or self.current_model.model_path != model_path:
+            self.current_model = TTSModel(
+                model_path = model_path,
+                config_path = self.root_dir / model_name / "config.json",
+                style_vec_path = self.root_dir / model_name / "style_vectors.npy",
+                device = self.device,
+            )
+
+        return self.current_model
+
+
+    def get_model_for_gradio(self, model_name: str, model_path_str: str) -> tuple[gr.Dropdown, gr.Button, gr.Dropdown]:
+        model_path = Path(model_path_str)
+        if model_name not in self.model_files_dict:
+            raise ValueError(f"Model `{model_name}` is not found")
+        if model_path not in self.model_files_dict[model_name]:
+            raise ValueError(f"Model file `{model_path}` is not found")
+        if (
+            self.current_model is not None
+            and self.current_model.model_path == model_path
+        ):
+            # Already loaded
+            speakers = list(self.current_model.spk2id.keys())
+            styles = list(self.current_model.style2id.keys())
+            return (
+                gr.Dropdown(choices=styles, value=styles[0]),  # type: ignore
+                gr.Button(interactive=True, value="音声合成"),
+                gr.Dropdown(choices=speakers, value=speakers[0]),  # type: ignore
+            )
+        self.current_model = TTSModel(
+            model_path = model_path,
+            config_path = self.root_dir / model_name / "config.json",
+            style_vec_path = self.root_dir / model_name / "style_vectors.npy",
+            device = self.device,
+        )
+        speakers = list(self.current_model.spk2id.keys())
+        styles = list(self.current_model.style2id.keys())
+        return (
+            gr.Dropdown(choices=styles, value=styles[0]),  # type: ignore
+            gr.Button(interactive=True, value="音声合成"),
+            gr.Dropdown(choices=speakers, value=speakers[0]),  # type: ignore
+        )
+
+
+    def update_model_files_for_gradio(self, model_name: str) -> gr.Dropdown:
+        model_files = self.model_files_dict[model_name]
+        return gr.Dropdown(choices=model_files, value=model_files[0])  # type: ignore
+
+
+    def update_model_names_for_gradio(self) -> tuple[gr.Dropdown, gr.Dropdown, gr.Button]:
+        self.refresh()
+        initial_model_name = self.model_names[0]
+        initial_model_files = self.model_files_dict[initial_model_name]
+        return (
+            gr.Dropdown(choices=self.model_names, value=initial_model_name),  # type: ignore
+            gr.Dropdown(choices=initial_model_files, value=initial_model_files[0]),  # type: ignore
+            gr.Button(interactive=False),  # For tts_button
+        )

style_bert_vits2/utils/stdout_wrapper.py

@@ -1,40 +1,47 @@
-"""
-`sys.stdout` wrapper for both Google Colab and local environment.
-"""
-
 import sys
 import tempfile
+from typing import TextIO
 
 
-class StdoutWrapper:
+class StdoutWrapper(TextIO):
-    def __init__(self):
+    """
+    `sys.stdout` wrapper for both Google Colab and local environment.
+    """
+
+
+    def __init__(self) -> None:
         self.temp_file = tempfile.NamedTemporaryFile(
             mode="w+", delete=False, encoding="utf-8"
         )
         self.original_stdout = sys.stdout
 
-    def write(self, message: str):
+
-        self.temp_file.write(message)
+    def write(self, message: str) -> int:
+        result = self.temp_file.write(message)
         self.temp_file.flush()
         print(message, end="", file=self.original_stdout)
+        return result
 
-    def flush(self):
+
+    def flush(self) -> None:
         self.temp_file.flush()
 
-    def read(self):
-        self.temp_file.seek(0)
-        return self.temp_file.read()
 
-    def close(self):
+    def read(self, n: int = -1) -> str:
+        self.temp_file.seek(0)
+        return self.temp_file.read(n)
+
+
+    def close(self) -> None:
         self.temp_file.close()
 
-    def fileno(self):
+
+    def fileno(self) -> int:
         return self.temp_file.fileno()
 
 
 try:
-    import google.colab
+    import google.colab  # type: ignore
-
     SAFE_STDOUT = StdoutWrapper()
 except ImportError:
     SAFE_STDOUT = sys.stdout

style_bert_vits2/utils/strenum.py

@@ -0,0 +1,36 @@
+import enum
+
+
+class StrEnum(str, enum.Enum):
+    """
+    Enum where members are also (and must be) strings (backported from Python 3.11).
+    """
+
+    def __new__(cls, *values: str) -> "StrEnum":
+        "values must already be of type `str`"
+        if len(values) > 3:
+            raise TypeError('too many arguments for str(): %r' % (values, ))
+        if len(values) == 1:
+            # it must be a string
+            if not isinstance(values[0], str):  # type: ignore
+                raise TypeError('%r is not a string' % (values[0], ))
+        if len(values) >= 2:
+            # check that encoding argument is a string
+            if not isinstance(values[1], str):  # type: ignore
+                raise TypeError('encoding must be a string, not %r' % (values[1], ))
+        if len(values) == 3:
+            # check that errors argument is a string
+            if not isinstance(values[2], str):  # type: ignore
+                raise TypeError('errors must be a string, not %r' % (values[2]))
+        value = str(*values)
+        member = str.__new__(cls, value)
+        member._value_ = value
+        return member
+
+
+    @staticmethod
+    def _generate_next_value_(name: str, start: int, count: int, last_values: list[str]) -> str:
+        """
+        Return the lower-cased version of the member name.
+        """
+        return name.lower()

style_bert_vits2/utils/subprocess.py

@@ -0,0 +1,54 @@
+import subprocess
+import sys
+from typing import Any, Callable
+
+from style_bert_vits2.logging import logger
+from style_bert_vits2.utils.stdout_wrapper import SAFE_STDOUT
+
+
+def run_script_with_log(cmd: list[str], ignore_warning: bool = False) -> tuple[bool, str]:
+    """
+    指定されたコマンドを実行し、そのログを記録する。
+
+    Args:
+        cmd: 実行するコマンドのリスト
+        ignore_warning: 警告を無視するかどうかのフラグ
+
+    Returns:
+        tuple[bool, str]: 実行が成功したかどうかのブール値と、エラーまたは警告のメッセージ（ある場合）
+    """
+
+    logger.info(f"Running: {' '.join(cmd)}")
+    result = subprocess.run(
+        [sys.executable] + cmd,
+        stdout = SAFE_STDOUT,
+        stderr = subprocess.PIPE,
+        text = True,
+        encoding = "utf-8",
+    )
+    if result.returncode != 0:
+        logger.error(f"Error: {' '.join(cmd)}\n{result.stderr}")
+        return False, result.stderr
+    elif result.stderr and not ignore_warning:
+        logger.warning(f"Warning: {' '.join(cmd)}\n{result.stderr}")
+        return True, result.stderr
+    logger.success(f"Success: {' '.join(cmd)}")
+
+    return True, ""
+
+
+def second_elem_of(original_function: Callable[..., tuple[Any, Any]]) -> Callable[..., Any]:
+    """
+    与えられた関数をラップし、その戻り値の 2 番目の要素のみを返す関数を生成する。
+
+    Args:
+        original_function (Callable[..., tuple[Any, Any]])): ラップする元の関数
+
+    Returns:
+        Callable[..., Any]: 元の関数の戻り値の 2 番目の要素のみを返す関数
+    """
+
+    def inner_function(*args, **kwargs) -> Any:  # type: ignore
+        return original_function(*args, **kwargs)[1]
+
+    return inner_function

style_bert_vits2/voice.py

@@ -0,0 +1,52 @@
+from typing import Any
+
+import numpy as np
+import pyworld
+from numpy.typing import NDArray
+
+
+def adjust_voice(
+    fs: int,
+    wave: NDArray[Any],
+    pitch_scale: float = 1.0,
+    intonation_scale: float = 1.0,
+) -> tuple[int, NDArray[Any]]:
+    """
+    音声のピッチとイントネーションを調整する。
+    変更すると若干音質が劣化するので、どちらも初期値のままならそのまま返す。
+
+    Args:
+        fs (int): 音声のサンプリング周波数
+        wave (NDArray[Any]): 音声データ
+        pitch_scale (float, optional): ピッチの高さ. Defaults to 1.0.
+        intonation_scale (float, optional): イントネーションの高さ. Defaults to 1.0.
+
+    Returns:
+        tuple[int, NDArray[Any]]: 調整後の音声データのサンプリング周波数と音声データ
+    """
+
+    if pitch_scale == 1.0 and intonation_scale == 1.0:
+        # 初期値の場合は、音質劣化を避けるためにそのまま返す
+        return fs, wave
+
+    # pyworld で f0 を加工して合成
+    # pyworld よりもよいのがあるかもしれないが……
+
+    wave = wave.astype(np.double)
+
+    # 質が高そうだしとりあえずharvestにしておく
+    f0, t = pyworld.harvest(wave, fs)
+
+    sp = pyworld.cheaptrick(wave, f0, t, fs)
+    ap = pyworld.d4c(wave, f0, t, fs)
+
+    non_zero_f0 = [f for f in f0 if f != 0]
+    f0_mean = sum(non_zero_f0) / len(non_zero_f0)
+
+    for i, f in enumerate(f0):
+        if f == 0:
+            continue
+        f0[i] = pitch_scale * f0_mean + intonation_scale * (f - f0_mean)
+
+    wave = pyworld.synthesize(f0, sp, ap, fs)
+    return fs, wave

style_gen.py

@@ -6,10 +6,10 @@ import numpy as np
 import torch
 from tqdm import tqdm
 
-import utils
-from common.log import logger
-from common.stdout_wrapper import SAFE_STDOUT
 from config import config
+from style_bert_vits2.logging import logger
+from style_bert_vits2.models.hyper_parameters import HyperParameters
+from style_bert_vits2.utils.stdout_wrapper import SAFE_STDOUT
 
 warnings.filterwarnings("ignore", category=UserWarning)
 from pyannote.audio import Inference, Model
@@ -72,7 +72,7 @@ if __name__ == "__main__":
     config_path = args.config
     num_processes = args.num_processes
 
-    hps = utils.get_hparams_from_file(config_path)
+    hps = HyperParameters.load_from_json(config_path)
 
     device = config.style_gen_config.device
 

tests/.gitignore

@@ -0,0 +1 @@
+*.wav

tests/test_main.py

@@ -0,0 +1,55 @@
+import pytest
+from scipy.io import wavfile
+
+from style_bert_vits2.constants import BASE_DIR, Languages
+from style_bert_vits2.tts_model import TTSModelHolder
+
+
+def synthesize(device: str = 'cpu'):
+
+    # 音声合成モデルが配置されていれば、音声合成を実行
+    model_holder = TTSModelHolder(BASE_DIR / 'model_assets', device)
+    if len(model_holder.models_info) > 0:
+
+        # jvnv-F2-jp モデルを探す
+        for model_info in model_holder.models_info:
+            if model_info.name == 'jvnv-F2-jp':
+                # すべてのスタイルに対して音声合成を実行
+                for style in model_info.styles:
+
+                    # 音声合成を実行
+                    model = model_holder.get_model(model_info.name, model_info.files[0])
+                    model.load()
+                    sample_rate, audio_data = model.infer(
+                        "あらゆる現実を、すべて自分のほうへねじ曲げたのだ。",
+                        # 言語 (JP, EN, ZH / JP-Extra モデルの場合は JP のみ)
+                        language = Languages.JP,
+                        # 話者 ID (音声合成モデルに複数の話者が含まれる場合のみ必須、単一話者のみの場合は 0)
+                        speaker_id = 0,
+                        # 感情表現の強さ (0.0 〜 1.0)
+                        sdp_ratio = 0.4,
+                        # スタイル (Neutral, Happy など)
+                        style = style,
+                        # スタイルの強さ (0.0 〜 100.0)
+                        style_weight = 6.0,
+                    )
+
+                    # 音声データを保存
+                    (BASE_DIR / 'tests/wavs').mkdir(exist_ok=True, parents=True)
+                    wav_file_path = BASE_DIR / f'tests/wavs/{style}.wav'
+                    with open(wav_file_path, 'wb') as f:
+                        wavfile.write(f, sample_rate, audio_data)
+
+                    # 音声データが保存されたことを確認
+                    assert wav_file_path.exists()
+                    # wav_file_path.unlink()
+    else:
+        pytest.skip("音声合成モデルが見つかりませんでした。")
+
+
+def test_synthesize_cpu():
+    synthesize(device='cpu')
+
+
+def test_synthesize_cuda():
+    synthesize(device='cuda')

text/__init__.py

@@ -1,18 +0,0 @@
-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.
-    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

text/cleaner.py

@@ -1,26 +0,0 @@
-def clean_text(text, language, use_jp_extra=True, raise_yomi_error=False):
-    # Changed to import inside if condition to avoid unnecessary import
-    if language == "ZH":
-        from . import chinese as language_module
-
-        norm_text = language_module.text_normalize(text)
-        phones, tones, word2ph = language_module.g2p(norm_text)
-    elif language == "EN":
-        from . import english as language_module
-
-        norm_text = language_module.text_normalize(text)
-        phones, tones, word2ph = language_module.g2p(norm_text)
-    elif language == "JP":
-        from . import japanese as language_module
-
-        norm_text = language_module.text_normalize(text)
-        phones, tones, word2ph = language_module.g2p(
-            norm_text, use_jp_extra, raise_yomi_error=raise_yomi_error
-        )
-    else:
-        raise ValueError(f"Language {language} not supported")
-    return norm_text, phones, tones, word2ph
-
-
-if __name__ == "__main__":
-    pass

text/english.py

@@ -1,495 +0,0 @@
-import pickle
-import os
-import re
-from g2p_en import G2p
-from transformers import DebertaV2Tokenizer
-
-from text import symbols
-from text.symbols import punctuation
-
-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")
-_g2p = G2p()
-LOCAL_PATH = "./bert/deberta-v3-large"
-tokenizer = DebertaV2Tokenizer.from_pretrained(LOCAL_PATH)
-
-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": ".",
-        "·": ",",
-        "、": ",",
-        "…": "...",
-        "···": "...",
-        "・・・": "...",
-        "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
-
-
-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\u3005"
-    #     + "".join(punctuation)
-    #     + r"]+",
-    #     "",
-    #     replaced_text,
-    # )
-
-    return replaced_text
-
-
-def read_dict():
-    g2p_dict = {}
-    start_line = 49
-    with open(CMU_DICT_PATH) as f:
-        line = f.readline()
-        line_index = 1
-        while line:
-            if line_index >= start_line:
-                line = line.strip()
-                word_split = line.split("  ")
-                word = word_split[0]
-
-                syllable_split = word_split[1].split(" - ")
-                g2p_dict[word] = []
-                for syllable in syllable_split:
-                    phone_split = syllable.split(" ")
-                    g2p_dict[word].append(phone_split)
-
-            line_index = line_index + 1
-            line = f.readline()
-
-    return g2p_dict
-
-
-def cache_dict(g2p_dict, file_path):
-    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:
-            g2p_dict = pickle.load(pickle_file)
-    else:
-        g2p_dict = read_dict()
-        cache_dict(g2p_dict, CACHE_PATH)
-
-    return g2p_dict
-
-
-eng_dict = get_dict()
-
-
-def refine_ph(phn):
-    tone = 0
-    if re.search(r"\d$", phn):
-        tone = int(phn[-1]) + 1
-        phn = phn[:-1]
-    else:
-        tone = 3
-    return phn.lower(), tone
-
-
-def refine_syllables(syllables):
-    tones = []
-    phonemes = []
-    for phn_list in syllables:
-        for i in range(len(phn_list)):
-            phn = phn_list[i]
-            phn, tone = refine_ph(phn)
-            phonemes.append(phn)
-            tones.append(tone)
-    return phonemes, tones
-
-
-import re
-import inflect
-
-_inflect = inflect.engine()
-_comma_number_re = re.compile(r"([0-9][0-9\,]+[0-9])")
-_decimal_number_re = re.compile(r"([0-9]+\.[0-9]+)")
-_pounds_re = re.compile(r"£([0-9\,]*[0-9]+)")
-_dollars_re = re.compile(r"\$([0-9\.\,]*[0-9]+)")
-_ordinal_re = re.compile(r"[0-9]+(st|nd|rd|th)")
-_number_re = re.compile(r"[0-9]+")
-
-# List of (regular expression, replacement) pairs for abbreviations:
-_abbreviations = [
-    (re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1])
-    for x in [
-        ("mrs", "misess"),
-        ("mr", "mister"),
-        ("dr", "doctor"),
-        ("st", "saint"),
-        ("co", "company"),
-        ("jr", "junior"),
-        ("maj", "major"),
-        ("gen", "general"),
-        ("drs", "doctors"),
-        ("rev", "reverend"),
-        ("lt", "lieutenant"),
-        ("hon", "honorable"),
-        ("sgt", "sergeant"),
-        ("capt", "captain"),
-        ("esq", "esquire"),
-        ("ltd", "limited"),
-        ("col", "colonel"),
-        ("ft", "fort"),
-    ]
-]
-
-
-# List of (ipa, lazy ipa) pairs:
-_lazy_ipa = [
-    (re.compile("%s" % x[0]), x[1])
-    for x in [
-        ("r", "ɹ"),
-        ("æ", "e"),
-        ("ɑ", "a"),
-        ("ɔ", "o"),
-        ("ð", "z"),
-        ("θ", "s"),
-        ("ɛ", "e"),
-        ("ɪ", "i"),
-        ("ʊ", "u"),
-        ("ʒ", "ʥ"),
-        ("ʤ", "ʥ"),
-        ("ˈ", "↓"),
-    ]
-]
-
-# List of (ipa, lazy ipa2) pairs:
-_lazy_ipa2 = [
-    (re.compile("%s" % x[0]), x[1])
-    for x in [
-        ("r", "ɹ"),
-        ("ð", "z"),
-        ("θ", "s"),
-        ("ʒ", "ʑ"),
-        ("ʤ", "dʑ"),
-        ("ˈ", "↓"),
-    ]
-]
-
-# List of (ipa, ipa2) pairs
-_ipa_to_ipa2 = [
-    (re.compile("%s" % x[0]), x[1]) for x in [("r", "ɹ"), ("ʤ", "dʒ"), ("ʧ", "tʃ")]
-]
-
-
-def _expand_dollars(m):
-    match = m.group(1)
-    parts = match.split(".")
-    if len(parts) > 2:
-        return match + " dollars"  # Unexpected format
-    dollars = int(parts[0]) if parts[0] else 0
-    cents = int(parts[1]) if len(parts) > 1 and parts[1] else 0
-    if dollars and cents:
-        dollar_unit = "dollar" if dollars == 1 else "dollars"
-        cent_unit = "cent" if cents == 1 else "cents"
-        return "%s %s, %s %s" % (dollars, dollar_unit, cents, cent_unit)
-    elif dollars:
-        dollar_unit = "dollar" if dollars == 1 else "dollars"
-        return "%s %s" % (dollars, dollar_unit)
-    elif cents:
-        cent_unit = "cent" if cents == 1 else "cents"
-        return "%s %s" % (cents, cent_unit)
-    else:
-        return "zero dollars"
-
-
-def _remove_commas(m):
-    return m.group(1).replace(",", "")
-
-
-def _expand_ordinal(m):
-    return _inflect.number_to_words(m.group(0))
-
-
-def _expand_number(m):
-    num = int(m.group(0))
-    if num > 1000 and num < 3000:
-        if num == 2000:
-            return "two thousand"
-        elif num > 2000 and num < 2010:
-            return "two thousand " + _inflect.number_to_words(num % 100)
-        elif num % 100 == 0:
-            return _inflect.number_to_words(num // 100) + " hundred"
-        else:
-            return _inflect.number_to_words(
-                num, andword="", zero="oh", group=2
-            ).replace(", ", " ")
-    else:
-        return _inflect.number_to_words(num, andword="")
-
-
-def _expand_decimal_point(m):
-    return m.group(1).replace(".", " point ")
-
-
-def normalize_numbers(text):
-    text = re.sub(_comma_number_re, _remove_commas, text)
-    text = re.sub(_pounds_re, r"\1 pounds", text)
-    text = re.sub(_dollars_re, _expand_dollars, text)
-    text = re.sub(_decimal_number_re, _expand_decimal_point, text)
-    text = re.sub(_ordinal_re, _expand_ordinal, text)
-    text = re.sub(_number_re, _expand_number, text)
-    return text
-
-
-def text_normalize(text):
-    text = normalize_numbers(text)
-    text = replace_punctuation(text)
-    text = re.sub(r"([,;.\?\!])([\w])", r"\1 \2", text)
-    return text
-
-
-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
-
-
-def sep_text(text):
-    words = re.split(r"([,;.\?\!\s+])", text)
-    words = [word for word in words if word.strip() != ""]
-    return words
-
-
-def text_to_words(text):
-    tokens = tokenizer.tokenize(text)
-    words = []
-    for idx, t in enumerate(tokens):
-        if t.startswith("▁"):
-            words.append([t[1:]])
-        else:
-            if t in punctuation:
-                if idx == len(tokens) - 1:
-                    words.append([f"{t}"])
-                else:
-                    if (
-                        not tokens[idx + 1].startswith("▁")
-                        and tokens[idx + 1] not in punctuation
-                    ):
-                        if idx == 0:
-                            words.append([])
-                        words[-1].append(f"{t}")
-                    else:
-                        words.append([f"{t}"])
-            else:
-                if idx == 0:
-                    words.append([])
-                words[-1].append(f"{t}")
-    return words
-
-
-def g2p(text):
-    phones = []
-    tones = []
-    phone_len = []
-    # words = sep_text(text)
-    # tokens = [tokenizer.tokenize(i) for i in words]
-    words = text_to_words(text)
-
-    for word in words:
-        temp_phones, temp_tones = [], []
-        if len(word) > 1:
-            if "'" in word:
-                word = ["".join(word)]
-        for w in word:
-            if w in punctuation:
-                temp_phones.append(w)
-                temp_tones.append(0)
-                continue
-            if w.upper() in eng_dict:
-                phns, tns = refine_syllables(eng_dict[w.upper()])
-                temp_phones += [post_replace_ph(i) for i in phns]
-                temp_tones += tns
-                # w2ph.append(len(phns))
-            else:
-                phone_list = list(filter(lambda p: p != " ", _g2p(w)))
-                phns = []
-                tns = []
-                for ph in phone_list:
-                    if ph in arpa:
-                        ph, tn = refine_ph(ph)
-                        phns.append(ph)
-                        tns.append(tn)
-                    else:
-                        phns.append(ph)
-                        tns.append(0)
-                temp_phones += [post_replace_ph(i) for i in phns]
-                temp_tones += tns
-        phones += temp_phones
-        tones += temp_tones
-        phone_len.append(len(temp_phones))
-        # phones = [post_replace_ph(i) for i in phones]
-
-    word2ph = []
-    for token, pl in zip(words, phone_len):
-        word_len = len(token)
-
-        aaa = distribute_phone(pl, word_len)
-        word2ph += aaa
-
-    phones = ["_"] + phones + ["_"]
-    tones = [0] + tones + [0]
-    word2ph = [1] + word2ph + [1]
-    assert len(phones) == len(tones), text
-    assert len(phones) == sum(word2ph), text
-
-    return phones, tones, word2ph
-
-
-def get_bert_feature(text, word2ph):
-    from text import english_bert_mock
-
-    return english_bert_mock.get_bert_feature(text, word2ph)
-
-
-if __name__ == "__main__":
-    # print(get_dict())
-    # print(eng_word_to_phoneme("hello"))
-    print(g2p("In this paper, we propose 1 DSPGAN, a GAN-based universal vocoder."))
-    # all_phones = set()
-    # for k, syllables in eng_dict.items():
-    #     for group in syllables:
-    #         for ph in group:
-    #             all_phones.add(ph)
-    # print(all_phones)

text/english_bert_mock.py

@@ -1,63 +0,0 @@
-import sys
-
-import torch
-from transformers import DebertaV2Model, DebertaV2Tokenizer
-
-from config import config
-
-
-LOCAL_PATH = "./bert/deberta-v3-large"
-
-tokenizer = DebertaV2Tokenizer.from_pretrained(LOCAL_PATH)
-
-models = dict()
-
-
-def get_bert_feature(
-    text,
-    word2ph,
-    device=config.bert_gen_config.device,
-    assist_text=None,
-    assist_text_weight=0.7,
-):
-    if (
-        sys.platform == "darwin"
-        and torch.backends.mps.is_available()
-        and device == "cpu"
-    ):
-        device = "mps"
-    if not device:
-        device = "cuda"
-    if device == "cuda" and not torch.cuda.is_available():
-        device = "cpu"
-    if device not in models.keys():
-        models[device] = DebertaV2Model.from_pretrained(LOCAL_PATH).to(device)
-    with torch.no_grad():
-        inputs = tokenizer(text, return_tensors="pt")
-        for i in inputs:
-            inputs[i] = inputs[i].to(device)
-        res = models[device](**inputs, output_hidden_states=True)
-        res = torch.cat(res["hidden_states"][-3:-2], -1)[0].cpu()
-        if assist_text:
-            style_inputs = tokenizer(assist_text, return_tensors="pt")
-            for i in style_inputs:
-                style_inputs[i] = style_inputs[i].to(device)
-            style_res = models[device](**style_inputs, output_hidden_states=True)
-            style_res = torch.cat(style_res["hidden_states"][-3:-2], -1)[0].cpu()
-            style_res_mean = style_res.mean(0)
-    assert len(word2ph) == res.shape[0], (text, res.shape[0], len(word2ph))
-    word2phone = word2ph
-    phone_level_feature = []
-    for i in range(len(word2phone)):
-        if assist_text:
-            repeat_feature = (
-                res[i].repeat(word2phone[i], 1) * (1 - assist_text_weight)
-                + style_res_mean.repeat(word2phone[i], 1) * assist_text_weight
-            )
-        else:
-            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/get_bert.py

@@ -1,16 +0,0 @@
-from .japanese_bert import get_bert_feature as get_japanese_bert_feature
-
-
-def get_bert(text, word2ph, language, device, assist_text=None, assist_text_weight=0.7):
-    if language == "ZH":
-        from .chinese_bert import get_bert_feature
-    elif language == "EN":
-        from .english_bert_mock import get_bert_feature
-    elif language == "JP":
-        # pyopenjtalkのworkerを1度だけ起動するため、ここでのimportは避ける
-        # 他言語のようにimportすると、get_bertが呼ばれるたびにpyopenjtalkのworkerが起動してしまう
-        get_bert_feature = get_japanese_bert_feature
-    else:
-        raise ValueError(f"Language {language} not supported")
-
-    return get_bert_feature(text, word2ph, device, assist_text, assist_text_weight)

text/japanese.py

@@ -1,644 +0,0 @@
-# Convert Japanese text to phonemes which is
-# compatible with Julius https://github.com/julius-speech/segmentation-kit
-import re
-import unicodedata
-from pathlib import Path
-
-from num2words import num2words
-from transformers import AutoTokenizer
-
-from common.log import logger
-from text import punctuation
-from text.japanese_mora_list import (
-    mora_kata_to_mora_phonemes,
-    mora_phonemes_to_mora_kata,
-)
-from text.user_dict import update_dict
-
-from . import pyopenjtalk_worker as pyopenjtalk
-
-# 最初にpyopenjtalkの辞書を更新
-update_dict()
-
-# 子音の集合
-COSONANTS = set(
-    [
-        cosonant
-        for cosonant, _ in mora_kata_to_mora_phonemes.values()
-        if cosonant is not None
-    ]
-)
-
-# 母音の集合、便宜上「ん」を含める
-VOWELS = {"a", "i", "u", "e", "o", "N"}
-
-
-class YomiError(Exception):
-    """
-    OpenJTalkで、読みが正しく取得できない箇所があるときに発生する例外。
-    基本的に「学習の前処理のテキスト処理時」には発生させ、そうでない場合は、
-    ignore_yomi_error=Trueにしておいて、この例外を発生させないようにする。
-    """
-
-    pass
-
-
-# 正規化で記号を変換するための辞書
-rep_map = {
-    "：": ",",
-    "；": ",",
-    "，": ",",
-    "。": ".",
-    "！": "!",
-    "？": "?",
-    "\n": ".",
-    "．": ".",
-    "…": "...",
-    "···": "...",
-    "・・・": "...",
-    "·": ",",
-    "・": ",",
-    "、": ",",
-    "$": ".",
-    "“": "'",
-    "”": "'",
-    '"': "'",
-    "‘": "'",
-    "’": "'",
-    "（": "'",
-    "）": "'",
-    "(": "'",
-    ")": "'",
-    "《": "'",
-    "》": "'",
-    "【": "'",
-    "】": "'",
-    "[": "'",
-    "]": "'",
-    # NFKC正規化後のハイフン・ダッシュの変種を全て通常半角ハイフン - \u002d に変換
-    "\u02d7": "\u002d",  # ˗, Modifier Letter Minus Sign
-    "\u2010": "\u002d",  # ‐, Hyphen,
-    # "\u2011": "\u002d",  # ‑, Non-Breaking Hyphen, NFKCにより\u2010に変換される
-    "\u2012": "\u002d",  # ‒, Figure Dash
-    "\u2013": "\u002d",  # –, En Dash
-    "\u2014": "\u002d",  # —, Em Dash
-    "\u2015": "\u002d",  # ―, Horizontal Bar
-    "\u2043": "\u002d",  # ⁃, Hyphen Bullet
-    "\u2212": "\u002d",  # −, Minus Sign
-    "\u23af": "\u002d",  # ⎯, Horizontal Line Extension
-    "\u23e4": "\u002d",  # ⏤, Straightness
-    "\u2500": "\u002d",  # ─, Box Drawings Light Horizontal
-    "\u2501": "\u002d",  # ━, Box Drawings Heavy Horizontal
-    "\u2e3a": "\u002d",  # ⸺, Two-Em Dash
-    "\u2e3b": "\u002d",  # ⸻, Three-Em Dash
-    # "～": "-",  # これは長音記号「ー」として扱うよう変更
-    # "~": "-",  # これも長音記号「ー」として扱うよう変更
-    "「": "'",
-    "」": "'",
-}
-
-
-def text_normalize(text):
-    """
-    日本語のテキストを正規化する。
-    結果は、ちょうど次の文字のみからなる：
-    - ひらがな
-    - カタカナ（全角長音記号「ー」が入る！）
-    - 漢字
-    - 半角アルファベット（大文字と小文字）
-    - ギリシャ文字
-    - `.` （句点`。`や`…`の一部や改行等）
-    - `,` （読点`、`や`:`等）
-    - `?` （疑問符`？`）
-    - `!` （感嘆符`！`）
-    - `'` （`「`や`」`等）
-    - `-` （`―`（ダッシュ、長音記号ではない）や`-`等）
-
-    注意点:
-    - 三点リーダー`…`は`...`に変換される（`なるほど…。` → `なるほど....`）
-    - 数字は漢字に変換される（`1,100円` → `千百円`、`52.34` → `五十二点三四`）
-    - 読点や疑問符等の位置・個数等は保持される（`??あ、、！！！` → `??あ,,!!!`）
-    """
-    res = unicodedata.normalize("NFKC", text)  # ここでアルファベットは半角になる
-    res = japanese_convert_numbers_to_words(res)  # 「100円」→「百円」等
-    # 「～」と「~」も長音記号として扱う
-    res = res.replace("~", "ー")
-    res = res.replace("～", "ー")
-
-    res = replace_punctuation(res)  # 句読点等正規化、読めない文字を削除
-
-    # 結合文字の濁点・半濁点を削除
-    # 通常の「ば」等はそのままのこされる、「あ゛」は上で「あ゙」になりここで「あ」になる
-    res = res.replace("\u3099", "")  # 結合文字の濁点を削除、る゙ → る
-    res = res.replace("\u309A", "")  # 結合文字の半濁点を削除、な゚ → な
-    return res
-
-
-def replace_punctuation(text: str) -> str:
-    """句読点等を「.」「,」「!」「?」「'」「-」に正規化し、OpenJTalkで読みが取得できるもののみ残す：
-    漢字・平仮名・カタカナ、アルファベット、ギリシャ文字
-    """
-    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\u3005"
-        # ↓ 半角アルファベット（大文字と小文字）
-        + r"\u0041-\u005A\u0061-\u007A"
-        # ↓ 全角アルファベット（大文字と小文字）
-        + r"\uFF21-\uFF3A\uFF41-\uFF5A"
-        # ↓ ギリシャ文字
-        + r"\u0370-\u03FF\u1F00-\u1FFF"
-        # ↓ "!", "?", "…", ",", ".", "'", "-", 但し`…`はすでに`...`に変換されている
-        + "".join(punctuation) + r"]+",
-        # 上述以外の文字を削除
-        "",
-        replaced_text,
-    )
-
-    return replaced_text
-
-
-_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 g2p(
-    norm_text: str, use_jp_extra: bool = True, raise_yomi_error: bool = False
-) -> tuple[list[str], list[int], list[int]]:
-    """
-    他で使われるメインの関数。`text_normalize()`で正規化された`norm_text`を受け取り、
-    - phones: 音素のリスト（ただし`!`や`,`や`.`等punctuationが含まれうる）
-    - tones: アクセントのリスト、0（低）と1（高）からなり、phonesと同じ長さ
-    - word2ph: 元のテキストの各文字に音素が何個割り当てられるかを表すリスト
-    のタプルを返す。
-    ただし`phones`と`tones`の最初と終わりに`_`が入り、応じて`word2ph`の最初と最後に1が追加される。
-
-    use_jp_extra: Falseの場合、「ん」の音素を「N」ではなく「n」とする。
-    raise_yomi_error: Trueの場合、読めない文字があるときに例外を発生させる。
-    Falseの場合は読めない文字が消えたような扱いとして処理される。
-    """
-    # pyopenjtalkのフルコンテキストラベルを使ってアクセントを取り出すと、punctuationの位置が消えてしまい情報が失われてしまう：
-    # 「こんにちは、世界。」と「こんにちは！世界。」と「こんにちは！！！？？？世界……。」は全て同じになる。
-    # よって、まずpunctuation無しの音素とアクセントのリストを作り、
-    # それとは別にpyopenjtalk.run_frontend()で得られる音素リスト（こちらはpunctuationが保持される）を使い、
-    # アクセント割当をしなおすことによってpunctuationを含めた音素とアクセントのリストを作る。
-
-    # punctuationがすべて消えた、音素とアクセントのタプルのリスト（「ん」は「N」）
-    phone_tone_list_wo_punct = g2phone_tone_wo_punct(norm_text)
-
-    # sep_text: 単語単位の単語のリスト、読めない文字があったらraise_yomi_errorなら例外、そうでないなら読めない文字が消えて返ってくる
-    # sep_kata: 単語単位の単語のカタカナ読みのリスト
-    sep_text, sep_kata = text2sep_kata(norm_text, raise_yomi_error=raise_yomi_error)
-
-    # sep_phonemes: 各単語ごとの音素のリストのリスト
-    sep_phonemes = handle_long([kata2phoneme_list(i) for i in sep_kata])
-
-    # phone_w_punct: sep_phonemesを結合した、punctuationを元のまま保持した音素列
-    phone_w_punct: list[str] = []
-    for i in sep_phonemes:
-        phone_w_punct += i
-
-    # punctuation無しのアクセント情報を使って、punctuationを含めたアクセント情報を作る
-    phone_tone_list = align_tones(phone_w_punct, phone_tone_list_wo_punct)
-    # logger.debug(f"phone_tone_list:\n{phone_tone_list}")
-    # word2phは厳密な解答は不可能なので（「今日」「眼鏡」等の熟字訓が存在）、
-    # Bert-VITS2では、単語単位の分割を使って、単語の文字ごとにだいたい均等に音素を分配する
-
-    # sep_textから、各単語を1文字1文字分割して、文字のリスト（のリスト）を作る
-    sep_tokenized: list[list[str]] = []
-    for i in sep_text:
-        if i not in punctuation:
-            sep_tokenized.append(
-                tokenizer.tokenize(i)
-            )  # ここでおそらく`i`が文字単位に分割される
-        else:
-            sep_tokenized.append([i])
-
-    # 各単語について、音素の数と文字の数を比較して、均等っぽく分配する
-    word2ph = []
-    for token, phoneme in zip(sep_tokenized, sep_phonemes):
-        phone_len = len(phoneme)
-        word_len = len(token)
-        word2ph += distribute_phone(phone_len, word_len)
-
-    # 最初と最後に`_`記号を追加、アクセントは0（低）、word2phもそれに合わせて追加
-    phone_tone_list = [("_", 0)] + phone_tone_list + [("_", 0)]
-    word2ph = [1] + word2ph + [1]
-
-    phones = [phone for phone, _ in phone_tone_list]
-    tones = [tone for _, tone in phone_tone_list]
-
-    assert len(phones) == sum(word2ph), f"{len(phones)} != {sum(word2ph)}"
-
-    # use_jp_extraでない場合は「N」を「n」に変換
-    if not use_jp_extra:
-        phones = [phone if phone != "N" else "n" for phone in phones]
-
-    return phones, tones, word2ph
-
-
-def g2kata_tone(norm_text: str) -> list[tuple[str, int]]:
-    """
-    テキストからカタカナとアクセントのペアのリストを返す。
-    推論時のみに使われるので、常に`raise_yomi_error=False`でg2pを呼ぶ。
-    """
-    phones, tones, _ = g2p(norm_text, use_jp_extra=True, raise_yomi_error=False)
-    return phone_tone2kata_tone(list(zip(phones, tones)))
-
-
-def phone_tone2kata_tone(phone_tone: list[tuple[str, int]]) -> list[tuple[str, int]]:
-    """phone_toneをのphone部分をカタカナに変換する。ただし最初と最後の("_", 0)は無視"""
-    phone_tone = phone_tone[1:]  # 最初の("_", 0)を無視
-    phones = [phone for phone, _ in phone_tone]
-    tones = [tone for _, tone in phone_tone]
-    result: list[tuple[str, int]] = []
-    current_mora = ""
-    for phone, next_phone, tone, next_tone in zip(phones, phones[1:], tones, tones[1:]):
-        # zipの関係で最後の("_", 0)は無視されている
-        if phone in punctuation:
-            result.append((phone, tone))
-            continue
-        if phone in COSONANTS:  # n以外の子音の場合
-            assert current_mora == "", f"Unexpected {phone} after {current_mora}"
-            assert tone == next_tone, f"Unexpected {phone} tone {tone} != {next_tone}"
-            current_mora = phone
-        else:
-            # phoneが母音もしくは「N」
-            current_mora += phone
-            result.append((mora_phonemes_to_mora_kata[current_mora], tone))
-            current_mora = ""
-    return result
-
-
-def kata_tone2phone_tone(kata_tone: list[tuple[str, int]]) -> list[tuple[str, int]]:
-    """`phone_tone2kata_tone()`の逆。"""
-    result: list[tuple[str, int]] = [("_", 0)]
-    for mora, tone in kata_tone:
-        if mora in punctuation:
-            result.append((mora, tone))
-        else:
-            cosonant, vowel = mora_kata_to_mora_phonemes[mora]
-            if cosonant is None:
-                result.append((vowel, tone))
-            else:
-                result.append((cosonant, tone))
-                result.append((vowel, tone))
-    result.append(("_", 0))
-    return result
-
-
-def g2phone_tone_wo_punct(text: str) -> list[tuple[str, int]]:
-    """
-    テキストに対して、音素とアクセント（0か1）のペアのリストを返す。
-    ただし「!」「.」「?」等の非音素記号(punctuation)は全て消える（ポーズ記号も残さない）。
-    非音素記号を含める処理は`align_tones()`で行われる。
-    また「っ」は「q」に、「ん」は「N」に変換される。
-    例: "こんにちは、世界ー。。元気？！" →
-    [('k', 0), ('o', 0), ('N', 1), ('n', 1), ('i', 1), ('ch', 1), ('i', 1), ('w', 1), ('a', 1), ('s', 1), ('e', 1), ('k', 0), ('a', 0), ('i', 0), ('i', 0), ('g', 1), ('e', 1), ('N', 0), ('k', 0), ('i', 0)]
-    """
-    prosodies = pyopenjtalk_g2p_prosody(text, drop_unvoiced_vowels=True)
-    # logger.debug(f"prosodies: {prosodies}")
-    result: list[tuple[str, int]] = []
-    current_phrase: list[tuple[str, int]] = []
-    current_tone = 0
-    for i, letter in enumerate(prosodies):
-        # 特殊記号の処理
-
-        # 文頭記号、無視する
-        if letter == "^":
-            assert i == 0, "Unexpected ^"
-        # アクセント句の終わりに来る記号
-        elif letter in ("$", "?", "_", "#"):
-            # 保持しているフレーズを、アクセント数値を0-1に修正し結果に追加
-            result.extend(fix_phone_tone(current_phrase))
-            # 末尾に来る終了記号、無視（文中の疑問文は`_`になる）
-            if letter in ("$", "?"):
-                assert i == len(prosodies) - 1, f"Unexpected {letter}"
-            # あとは"_"（ポーズ）と"#"（アクセント句の境界）のみ
-            # これらは残さず、次のアクセント句に備える。
-            current_phrase = []
-            # 0を基準点にしてそこから上昇・下降する（負の場合は上の`fix_phone_tone`で直る）
-            current_tone = 0
-        # アクセント上昇記号
-        elif letter == "[":
-            current_tone = current_tone + 1
-        # アクセント下降記号
-        elif letter == "]":
-            current_tone = current_tone - 1
-        # それ以外は通常の音素
-        else:
-            if letter == "cl":  # 「っ」の処理
-                letter = "q"
-            # elif letter == "N":  # 「ん」の処理
-            #     letter = "n"
-            current_phrase.append((letter, current_tone))
-    return result
-
-
-def text2sep_kata(
-    norm_text: str, raise_yomi_error: bool = False
-) -> tuple[list[str], list[str]]:
-    """
-    `text_normalize`で正規化済みの`norm_text`を受け取り、それを単語分割し、
-    分割された単語リストとその読み（カタカナor記号1文字）のリストのタプルを返す。
-    単語分割結果は、`g2p()`の`word2ph`で1文字あたりに割り振る音素記号の数を決めるために使う。
-    例:
-    `私はそう思う!って感じ?` →
-    ["私", "は", "そう", "思う", "!", "って", "感じ", "?"], ["ワタシ", "ワ", "ソー", "オモウ", "!", "ッテ", "カンジ", "?"]
-
-    raise_yomi_error: Trueの場合、読めない文字があるときに例外を発生させる。
-    Falseの場合は読めない文字が消えたような扱いとして処理される。
-    """
-    # parsed: OpenJTalkの解析結果
-    parsed = pyopenjtalk.run_frontend(norm_text)
-    sep_text: list[str] = []
-    sep_kata: list[str] = []
-    for parts in parsed:
-        # word: 実際の単語の文字列
-        # yomi: その読み、但し無声化サインの`’`は除去
-        word, yomi = replace_punctuation(parts["string"]), parts["pron"].replace(
-            "’", ""
-        )
-        """
-        ここで`yomi`の取りうる値は以下の通りのはず。
-        - `word`が通常単語 → 通常の読み（カタカナ）
-            （カタカナからなり、長音記号も含みうる、`アー` 等）
-        - `word`が`ー` から始まる → `ーラー` や `ーーー` など
-        - `word`が句読点や空白等 → `、`
-        - `word`がpunctuationの繰り返し → 全角にしたもの
-        基本的にpunctuationは1文字ずつ分かれるが、何故かある程度連続すると1つにまとまる。
-        他にも`word`が読めないキリル文字アラビア文字等が来ると`、`になるが、正規化でこの場合は起きないはず。
-        また元のコードでは`yomi`が空白の場合の処理があったが、これは起きないはず。
-        処理すべきは`yomi`が`、`の場合のみのはず。
-        """
-        assert yomi != "", f"Empty yomi: {word}"
-        if yomi == "、":
-            # wordは正規化されているので、`.`, `,`, `!`, `'`, `-`, `--` のいずれか
-            if not set(word).issubset(set(punctuation)):  # 記号繰り返しか判定
-                # ここはpyopenjtalkが読めない文字等のときに起こる
-                if raise_yomi_error:
-                    raise YomiError(f"Cannot read: {word} in:\n{norm_text}")
-                logger.warning(f"Ignoring unknown: {word} in:\n{norm_text}")
-                continue
-            # yomiは元の記号のままに変更
-            yomi = word
-        elif yomi == "？":
-            assert word == "?", f"yomi `？` comes from: {word}"
-            yomi = "?"
-        sep_text.append(word)
-        sep_kata.append(yomi)
-    return sep_text, sep_kata
-
-
-# ESPnetの実装から引用、変更点無し。「ん」は「N」なことに注意。
-# https://github.com/espnet/espnet/blob/master/espnet2/text/phoneme_tokenizer.py
-def pyopenjtalk_g2p_prosody(text: str, drop_unvoiced_vowels: bool = True) -> list[str]:
-    """Extract phoneme + prosoody symbol sequence from input full-context labels.
-
-    The algorithm is based on `Prosodic features control by symbols as input of
-    sequence-to-sequence acoustic modeling for neural TTS`_ with some r9y9's tweaks.
-
-    Args:
-        text (str): Input text.
-        drop_unvoiced_vowels (bool): whether to drop unvoiced vowels.
-
-    Returns:
-        List[str]: List of phoneme + prosody symbols.
-
-    Examples:
-        >>> from espnet2.text.phoneme_tokenizer import pyopenjtalk_g2p_prosody
-        >>> pyopenjtalk_g2p_prosody("こんにちは。")
-        ['^', 'k', 'o', '[', 'N', 'n', 'i', 'ch', 'i', 'w', 'a', '$']
-
-    .. _`Prosodic features control by symbols as input of sequence-to-sequence acoustic
-        modeling for neural TTS`: https://doi.org/10.1587/transinf.2020EDP7104
-
-    """
-    labels = pyopenjtalk.make_label(pyopenjtalk.run_frontend(text))
-    N = len(labels)
-
-    phones = []
-    for n in range(N):
-        lab_curr = labels[n]
-
-        # current phoneme
-        p3 = re.search(r"\-(.*?)\+", lab_curr).group(1)
-        # deal unvoiced vowels as normal vowels
-        if drop_unvoiced_vowels and p3 in "AEIOU":
-            p3 = p3.lower()
-
-        # deal with sil at the beginning and the end of text
-        if p3 == "sil":
-            assert n == 0 or n == N - 1
-            if n == 0:
-                phones.append("^")
-            elif n == N - 1:
-                # check question form or not
-                e3 = _numeric_feature_by_regex(r"!(\d+)_", lab_curr)
-                if e3 == 0:
-                    phones.append("$")
-                elif e3 == 1:
-                    phones.append("?")
-            continue
-        elif p3 == "pau":
-            phones.append("_")
-            continue
-        else:
-            phones.append(p3)
-
-        # accent type and position info (forward or backward)
-        a1 = _numeric_feature_by_regex(r"/A:([0-9\-]+)\+", lab_curr)
-        a2 = _numeric_feature_by_regex(r"\+(\d+)\+", lab_curr)
-        a3 = _numeric_feature_by_regex(r"\+(\d+)/", lab_curr)
-
-        # number of mora in accent phrase
-        f1 = _numeric_feature_by_regex(r"/F:(\d+)_", lab_curr)
-
-        a2_next = _numeric_feature_by_regex(r"\+(\d+)\+", labels[n + 1])
-        # accent phrase border
-        if a3 == 1 and a2_next == 1 and p3 in "aeiouAEIOUNcl":
-            phones.append("#")
-        # pitch falling
-        elif a1 == 0 and a2_next == a2 + 1 and a2 != f1:
-            phones.append("]")
-        # pitch rising
-        elif a2 == 1 and a2_next == 2:
-            phones.append("[")
-
-    return phones
-
-
-def _numeric_feature_by_regex(regex, s):
-    match = re.search(regex, s)
-    if match is None:
-        return -50
-    return int(match.group(1))
-
-
-def fix_phone_tone(phone_tone_list: list[tuple[str, int]]) -> list[tuple[str, int]]:
-    """
-    `phone_tone_list`のtone（アクセントの値）を0か1の範囲に修正する。
-    例: [(a, 0), (i, -1), (u, -1)] → [(a, 1), (i, 0), (u, 0)]
-    """
-    tone_values = set(tone for _, tone in phone_tone_list)
-    if len(tone_values) == 1:
-        assert tone_values == {0}, tone_values
-        return phone_tone_list
-    elif len(tone_values) == 2:
-        if tone_values == {0, 1}:
-            return phone_tone_list
-        elif tone_values == {-1, 0}:
-            return [
-                (letter, 0 if tone == -1 else 1) for letter, tone in phone_tone_list
-            ]
-        else:
-            raise ValueError(f"Unexpected tone values: {tone_values}")
-    else:
-        raise ValueError(f"Unexpected tone values: {tone_values}")
-
-
-def distribute_phone(n_phone: int, n_word: int) -> list[int]:
-    """
-    左から右に1ずつ振り分け、次にまた左から右に1ずつ増やし、というふうに、
-    音素の数`n_phone`を単語の数`n_word`に分配する。
-    """
-    phones_per_word = [0] * n_word
-    for _ 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
-
-
-def handle_long(sep_phonemes: list[list[str]]) -> list[list[str]]:
-    """
-    フレーズごとに分かれた音素（長音記号がそのまま）のリストのリスト`sep_phonemes`を受け取り、
-    その長音記号を処理して、音素のリストのリストを返す。
-    基本的には直前の音素を伸ばすが、直前の音素が母音でない場合もしくは冒頭の場合は、
-    おそらく長音記号とダッシュを勘違いしていると思われるので、ダッシュに対応する音素`-`に変換する。
-    """
-    for i in range(len(sep_phonemes)):
-        if len(sep_phonemes[i]) == 0:
-            # 空白文字等でリストが空の場合
-            continue
-        if sep_phonemes[i][0] == "ー":
-            if i != 0:
-                prev_phoneme = sep_phonemes[i - 1][-1]
-                if prev_phoneme in VOWELS:
-                    # 母音と「ん」のあとの伸ばし棒なので、その母音に変換
-                    sep_phonemes[i][0] = sep_phonemes[i - 1][-1]
-                else:
-                    # 「。ーー」等おそらく予期しない長音記号
-                    # ダッシュの勘違いだと思われる
-                    sep_phonemes[i][0] = "-"
-            else:
-                # 冒頭に長音記号が来ていおり、これはダッシュの勘違いと思われる
-                sep_phonemes[i][0] = "-"
-        if "ー" in sep_phonemes[i]:
-            for j in range(len(sep_phonemes[i])):
-                if sep_phonemes[i][j] == "ー":
-                    sep_phonemes[i][j] = sep_phonemes[i][j - 1][-1]
-    return sep_phonemes
-
-
-tokenizer = AutoTokenizer.from_pretrained("./bert/deberta-v2-large-japanese-char-wwm")
-
-
-def align_tones(
-    phones_with_punct: list[str], phone_tone_list: list[tuple[str, int]]
-) -> list[tuple[str, int]]:
-    """
-    例:
-    …私は、、そう思う。
-    phones_with_punct:
-    [".", ".", ".", "w", "a", "t", "a", "sh", "i", "w", "a", ",", ",", "s", "o", "o", "o", "m", "o", "u", "."]
-    phone_tone_list:
-    [("w", 0), ("a", 0), ("t", 1), ("a", 1), ("sh", 1), ("i", 1), ("w", 1), ("a", 1), ("_", 0), ("s", 0), ("o", 0), ("o", 1), ("o", 1), ("m", 1), ("o", 1), ("u", 0))]
-    Return:
-    [(".", 0), (".", 0), (".", 0), ("w", 0), ("a", 0), ("t", 1), ("a", 1), ("sh", 1), ("i", 1), ("w", 1), ("a", 1), (",", 0), (",", 0), ("s", 0), ("o", 0), ("o", 1), ("o", 1), ("m", 1), ("o", 1), ("u", 0), (".", 0)]
-    """
-    result: list[tuple[str, int]] = []
-    tone_index = 0
-    for phone in phones_with_punct:
-        if tone_index >= len(phone_tone_list):
-            # 余ったpunctuationがある場合 → (punctuation, 0)を追加
-            result.append((phone, 0))
-        elif phone == phone_tone_list[tone_index][0]:
-            # phone_tone_listの現在の音素と一致する場合 → toneをそこから取得、(phone, tone)を追加
-            result.append((phone, phone_tone_list[tone_index][1]))
-            # 探すindexを1つ進める
-            tone_index += 1
-        elif phone in punctuation:
-            # phoneがpunctuationの場合 → (phone, 0)を追加
-            result.append((phone, 0))
-        else:
-            logger.debug(f"phones: {phones_with_punct}")
-            logger.debug(f"phone_tone_list: {phone_tone_list}")
-            logger.debug(f"result: {result}")
-            logger.debug(f"tone_index: {tone_index}")
-            logger.debug(f"phone: {phone}")
-            raise ValueError(f"Unexpected phone: {phone}")
-    return result
-
-
-def kata2phoneme_list(text: str) -> list[str]:
-    """
-    原則カタカナの`text`を受け取り、それをそのままいじらずに音素記号のリストに変換。
-    注意点：
-    - punctuationかその繰り返しが来た場合、punctuationたちをそのままリストにして返す。
-    - 冒頭に続く「ー」はそのまま「ー」のままにする（`handle_long()`で処理される）
-    - 文中の「ー」は前の音素記号の最後の音素記号に変換される。
-    例：
-    `ーーソーナノカーー` → ["ー", "ー", "s", "o", "o", "n", "a", "n", "o", "k", "a", "a", "a"]
-    `?` → ["?"]
-    `!?!?!?!?!` → ["!", "?", "!", "?", "!", "?", "!", "?", "!"]
-    """
-    if set(text).issubset(set(punctuation)):
-        return list(text)
-    # `text`がカタカナ（`ー`含む）のみからなるかどうかをチェック
-    if re.fullmatch(r"[\u30A0-\u30FF]+", text) is None:
-        raise ValueError(f"Input must be katakana only: {text}")
-    sorted_keys = sorted(mora_kata_to_mora_phonemes.keys(), key=len, reverse=True)
-    pattern = "|".join(map(re.escape, sorted_keys))
-
-    def mora2phonemes(mora: str) -> str:
-        cosonant, vowel = mora_kata_to_mora_phonemes[mora]
-        if cosonant is None:
-            return f" {vowel}"
-        return f" {cosonant} {vowel}"
-
-    spaced_phonemes = re.sub(pattern, lambda m: mora2phonemes(m.group()), text)
-
-    # 長音記号「ー」の処理
-    long_pattern = r"(\w)(ー*)"
-    long_replacement = lambda m: m.group(1) + (" " + m.group(1)) * len(m.group(2))
-    spaced_phonemes = re.sub(long_pattern, long_replacement, spaced_phonemes)
-    return spaced_phonemes.strip().split(" ")
-
-
-if __name__ == "__main__":
-    tokenizer = AutoTokenizer.from_pretrained(
-        "./bert/deberta-v2-large-japanese-char-wwm"
-    )
-    text = "こんにちは、世界。"
-    from text.japanese_bert import get_bert_feature
-
-    text = text_normalize(text)
-
-    phones, tones, word2ph = g2p(text)
-    bert = get_bert_feature(text, word2ph)
-
-    print(phones, tones, word2ph, bert.shape)

text/japanese_bert.py

@@ -1,70 +0,0 @@
-import sys
-
-import torch
-from transformers import AutoModelForMaskedLM, AutoTokenizer
-
-from config import config
-from text.japanese import text2sep_kata, text_normalize
-
-LOCAL_PATH = "./bert/deberta-v2-large-japanese-char-wwm"
-
-tokenizer = AutoTokenizer.from_pretrained(LOCAL_PATH)
-
-models = dict()
-
-
-def get_bert_feature(
-    text,
-    word2ph,
-    device=config.bert_gen_config.device,
-    assist_text=None,
-    assist_text_weight=0.7,
-):
-    # 各単語が何文字かを作る`word2ph`を使う必要があるので、読めない文字は必ず無視する
-    # でないと`word2ph`の結果とテキストの文字数結果が整合性が取れない
-    text = "".join(text2sep_kata(text, raise_yomi_error=False)[0])
-
-    if assist_text:
-        assist_text = "".join(text2sep_kata(assist_text, raise_yomi_error=False)[0])
-    if (
-        sys.platform == "darwin"
-        and torch.backends.mps.is_available()
-        and device == "cpu"
-    ):
-        device = "mps"
-    if not device:
-        device = "cuda"
-    if device == "cuda" and not torch.cuda.is_available():
-        device = "cpu"
-    if device not in models.keys():
-        models[device] = AutoModelForMaskedLM.from_pretrained(LOCAL_PATH).to(device)
-    with torch.no_grad():
-        inputs = tokenizer(text, return_tensors="pt")
-        for i in inputs:
-            inputs[i] = inputs[i].to(device)
-        res = models[device](**inputs, output_hidden_states=True)
-        res = torch.cat(res["hidden_states"][-3:-2], -1)[0].cpu()
-        if assist_text:
-            style_inputs = tokenizer(assist_text, return_tensors="pt")
-            for i in style_inputs:
-                style_inputs[i] = style_inputs[i].to(device)
-            style_res = models[device](**style_inputs, output_hidden_states=True)
-            style_res = torch.cat(style_res["hidden_states"][-3:-2], -1)[0].cpu()
-            style_res_mean = style_res.mean(0)
-
-    assert len(word2ph) == len(text) + 2, f"word2ph: {word2ph}, text: {text}"
-    word2phone = word2ph
-    phone_level_feature = []
-    for i in range(len(word2phone)):
-        if assist_text:
-            repeat_feature = (
-                res[i].repeat(word2phone[i], 1) * (1 - assist_text_weight)
-                + style_res_mean.repeat(word2phone[i], 1) * assist_text_weight
-            )
-        else:
-            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/pyopenjtalk_worker/__init__.py

@@ -1,126 +0,0 @@
-"""
-Run the pyopenjtalk worker in a separate process
-to avoid user dictionary access error
-"""
-
-from typing import Optional, Any
-
-from .worker_common import WORKER_PORT
-from .worker_client import WorkerClient
-
-from common.log import logger
-
-WORKER_CLIENT: Optional[WorkerClient] = None
-
-# pyopenjtalk interface
-
-# g2p: not used
-
-
-def run_frontend(text: str) -> list[dict[str, Any]]:
-    assert WORKER_CLIENT
-    ret = WORKER_CLIENT.dispatch_pyopenjtalk("run_frontend", text)
-    assert isinstance(ret, list)
-    return ret
-
-
-def make_label(njd_features) -> list[str]:
-    assert WORKER_CLIENT
-    ret = WORKER_CLIENT.dispatch_pyopenjtalk("make_label", njd_features)
-    assert isinstance(ret, list)
-    return ret
-
-
-def mecab_dict_index(path: str, out_path: str, dn_mecab: Optional[str] = None):
-    assert WORKER_CLIENT
-    WORKER_CLIENT.dispatch_pyopenjtalk("mecab_dict_index", path, out_path, dn_mecab)
-
-
-def update_global_jtalk_with_user_dict(path: str):
-    assert WORKER_CLIENT
-    WORKER_CLIENT.dispatch_pyopenjtalk("update_global_jtalk_with_user_dict", path)
-
-
-def unset_user_dict():
-    assert WORKER_CLIENT
-    WORKER_CLIENT.dispatch_pyopenjtalk("unset_user_dict")
-
-
-# initialize module when imported
-
-
-def initialize(port: int = WORKER_PORT):
-    import time
-    import socket
-    import sys
-    import atexit
-    import signal
-
-    logger.debug("initialize")
-    global WORKER_CLIENT
-    if WORKER_CLIENT:
-        return
-
-    client = None
-    try:
-        client = WorkerClient(port)
-    except (socket.timeout, socket.error):
-        logger.debug("try starting pyopenjtalk worker server")
-        import os
-        import subprocess
-
-        worker_pkg_path = os.path.relpath(
-            os.path.dirname(__file__), os.getcwd()
-        ).replace(os.sep, ".")
-        args = [sys.executable, "-m", worker_pkg_path, "--port", str(port)]
-        # new session, new process group
-        if sys.platform.startswith("win"):
-            cf = subprocess.CREATE_NEW_CONSOLE | subprocess.CREATE_NEW_PROCESS_GROUP  # type: ignore
-            si = subprocess.STARTUPINFO()  # type: ignore
-            si.dwFlags |= subprocess.STARTF_USESHOWWINDOW  # type: ignore
-            si.wShowWindow = subprocess.SW_HIDE  # type: ignore
-            subprocess.Popen(args, creationflags=cf, startupinfo=si)
-        else:
-            # align with Windows behavior
-            # start_new_session is same as specifying setsid in preexec_fn
-            subprocess.Popen(args, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, start_new_session=True)  # type: ignore
-
-        # wait until server listening
-        count = 0
-        while True:
-            try:
-                client = WorkerClient(port)
-                break
-            except socket.error:
-                time.sleep(1)
-                count += 1
-                # 10: max number of retries
-                if count == 10:
-                    raise TimeoutError("サーバーに接続できませんでした")
-
-    WORKER_CLIENT = client
-    atexit.register(terminate)
-
-    # when the process is killed
-    def signal_handler(signum, frame):
-        terminate()
-
-    signal.signal(signal.SIGTERM, signal_handler)
-
-
-# top-level declaration
-def terminate():
-    logger.debug("terminate")
-    global WORKER_CLIENT
-    if not WORKER_CLIENT:
-        return
-
-    # repare for unexpected errors
-    try:
-        if WORKER_CLIENT.status() == 1:
-            WORKER_CLIENT.quit_server()
-    except Exception as e:
-        logger.error(e)
-
-    WORKER_CLIENT.close()
-    WORKER_CLIENT = None

text/user_dict/README.md

@@ -1,19 +0,0 @@
-このフォルダに含まれるユーザー辞書関連のコードは、[VOICEVOX engine](https://github.com/VOICEVOX/voicevox_engine)プロジェクトのコードを改変したものを使用しています。VOICEVOXプロジェクトのチームに深く感謝し、その貢献を尊重します。
-
-**元のコード**:
-
-- [voicevox_engine/user_dict/](https://github.com/VOICEVOX/voicevox_engine/tree/f181411ec69812296989d9cc583826c22eec87ae/voicevox_engine/user_dict)
-- [voicevox_engine/model.py](https://github.com/VOICEVOX/voicevox_engine/blob/f181411ec69812296989d9cc583826c22eec87ae/voicevox_engine/model.py#L207)
-
-**改変の詳細**:
-
-- ファイル名の書き換えおよびそれに伴うimport文の書き換え。
-- VOICEVOX固有の部分をコメントアウト。
-- mutexを使用している部分をコメントアウト。
-- 参照しているpyopenjtalkの違いによるメソッド名の書き換え。
-- UserDictWordのmora_countのデフォルト値をNoneに指定。
-- Pydanticのモデルで必要な箇所のみを抽出。
-
-**ライセンス**:
-
-元のVOICEVOX engineのリポジトリのコードは、LGPL v3 と、ソースコードの公開が不要な別ライセンスのデュアルライセンスの下で使用されています。当プロジェクトにおけるこのモジュールもLGPLライセンスの下にあります。詳細については、プロジェクトのルートディレクトリにある[LGPL_LICENSE](/LGPL_LICENSE)ファイルをご参照ください。また、元のVOICEVOX engineプロジェクトのライセンスについては、[こちら](https://github.com/VOICEVOX/voicevox_engine/blob/master/LICENSE)をご覧ください。

train_ms.py

@@ -1,6 +1,5 @@
 import argparse
 import datetime
-import gc
 import os
 import platform
 
@@ -15,11 +14,7 @@ from torch.utils.tensorboard import SummaryWriter
 from tqdm import tqdm
 
 # logging.getLogger("numba").setLevel(logging.WARNING)
-import commons
 import default_style
-import utils
-from common.log import logger
-from common.stdout_wrapper import SAFE_STDOUT
 from config import config
 from data_utils import (
     DistributedBucketSampler,
@@ -28,8 +23,17 @@ from data_utils import (
 )
 from losses import discriminator_loss, feature_loss, generator_loss, kl_loss
 from mel_processing import mel_spectrogram_torch, spec_to_mel_torch
-from models import DurationDiscriminator, MultiPeriodDiscriminator, SynthesizerTrn
+from style_bert_vits2.logging import logger
-from text.symbols import symbols
+from style_bert_vits2.models import commons
+from style_bert_vits2.models import utils
+from style_bert_vits2.models.hyper_parameters import HyperParameters
+from style_bert_vits2.models.models import (
+    DurationDiscriminator,
+    MultiPeriodDiscriminator,
+    SynthesizerTrn,
+)
+from style_bert_vits2.nlp.symbols import SYMBOLS
+from style_bert_vits2.utils.stdout_wrapper import SAFE_STDOUT
 
 torch.backends.cuda.matmul.allow_tf32 = True
 torch.backends.cudnn.allow_tf32 = (
@@ -127,7 +131,7 @@ def run():
     local_rank = int(os.environ["LOCAL_RANK"])
     n_gpus = dist.get_world_size()
 
-    hps = utils.get_hparams_from_file(args.config)
+    hps = HyperParameters.load_from_json(args.config)
     # This is needed because we have to pass values to `train_and_evaluate()`
     hps.model_dir = model_dir
     hps.speedup = args.speedup
@@ -244,10 +248,7 @@ def run():
             drop_last=False,
             collate_fn=collate_fn,
         )
-    if (
+    if hps.model.use_noise_scaled_mas is True:
-        "use_noise_scaled_mas" in hps.model.keys()
-        and hps.model.use_noise_scaled_mas is True
-    ):
         logger.info("Using noise scaled MAS for VITS2")
         mas_noise_scale_initial = 0.01
         noise_scale_delta = 2e-6
@@ -255,10 +256,7 @@ def run():
         logger.info("Using normal MAS for VITS1")
         mas_noise_scale_initial = 0.0
         noise_scale_delta = 0.0
-    if (
+    if hps.model.use_duration_discriminator is True:
-        "use_duration_discriminator" in hps.model.keys()
-        and hps.model.use_duration_discriminator is True
-    ):
         logger.info("Using duration discriminator for VITS2")
         net_dur_disc = DurationDiscriminator(
             hps.model.hidden_channels,
@@ -267,10 +265,7 @@ def run():
             0.1,
             gin_channels=hps.model.gin_channels if hps.data.n_speakers != 0 else 0,
         ).cuda(local_rank)
-    if (
+    if hps.model.use_spk_conditioned_encoder is True:
-        "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"
@@ -279,13 +274,35 @@ def run():
         logger.info("Using normal encoder for VITS1")
 
     net_g = SynthesizerTrn(
-        len(symbols),
+        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,
+        # hps.model 以下のすべての値を引数に渡す
+        use_spk_conditioned_encoder = hps.model.use_spk_conditioned_encoder,
+        use_noise_scaled_mas = hps.model.use_noise_scaled_mas,
+        use_mel_posterior_encoder = hps.model.use_mel_posterior_encoder,
+        use_duration_discriminator = hps.model.use_duration_discriminator,
+        use_wavlm_discriminator = hps.model.use_wavlm_discriminator,
+        inter_channels = hps.model.inter_channels,
+        hidden_channels = hps.model.hidden_channels,
+        filter_channels = hps.model.filter_channels,
+        n_heads = hps.model.n_heads,
+        n_layers = hps.model.n_layers,
+        kernel_size = hps.model.kernel_size,
+        p_dropout = hps.model.p_dropout,
+        resblock = hps.model.resblock,
+        resblock_kernel_sizes = hps.model.resblock_kernel_sizes,
+        resblock_dilation_sizes = hps.model.resblock_dilation_sizes,
+        upsample_rates = hps.model.upsample_rates,
+        upsample_initial_channel = hps.model.upsample_initial_channel,
+        upsample_kernel_sizes = hps.model.upsample_kernel_sizes,
+        n_layers_q = hps.model.n_layers_q,
+        use_spectral_norm = hps.model.use_spectral_norm,
+        gin_channels = hps.model.gin_channels,
+        slm = hps.model.slm,
     ).cuda(local_rank)
 
     if getattr(hps.train, "freeze_ZH_bert", False):
@@ -344,31 +361,25 @@ def run():
 
     if utils.is_resuming(model_dir):
         if net_dur_disc is not None:
-            _, _, dur_resume_lr, epoch_str = utils.load_checkpoint(
+            _, _, dur_resume_lr, epoch_str = utils.checkpoints.load_checkpoint(
-                utils.latest_checkpoint_path(model_dir, "DUR_*.pth"),
+                utils.checkpoints.get_latest_checkpoint_path(model_dir, "DUR_*.pth"),
                 net_dur_disc,
                 optim_dur_disc,
-                skip_optimizer=(
+                skip_optimizer=hps.train.skip_optimizer,
-                    hps.train.skip_optimizer if "skip_optimizer" in hps.train else True
-                ),
             )
             if not optim_dur_disc.param_groups[0].get("initial_lr"):
                 optim_dur_disc.param_groups[0]["initial_lr"] = dur_resume_lr
-        _, optim_g, g_resume_lr, epoch_str = utils.load_checkpoint(
+        _, optim_g, g_resume_lr, epoch_str = utils.checkpoints.load_checkpoint(
-            utils.latest_checkpoint_path(model_dir, "G_*.pth"),
+            utils.checkpoints.get_latest_checkpoint_path(model_dir, "G_*.pth"),
             net_g,
             optim_g,
-            skip_optimizer=(
+            skip_optimizer=hps.train.skip_optimizer,
-                hps.train.skip_optimizer if "skip_optimizer" in hps.train else True
-            ),
         )
-        _, optim_d, d_resume_lr, epoch_str = utils.load_checkpoint(
+        _, optim_d, d_resume_lr, epoch_str = utils.checkpoints.load_checkpoint(
-            utils.latest_checkpoint_path(model_dir, "D_*.pth"),
+            utils.checkpoints.get_latest_checkpoint_path(model_dir, "D_*.pth"),
             net_d,
             optim_d,
-            skip_optimizer=(
+            skip_optimizer=hps.train.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
@@ -378,21 +389,21 @@ def run():
         epoch_str = max(epoch_str, 1)
         # global_step = (epoch_str - 1) * len(train_loader)
         global_step = int(
-            utils.get_steps(utils.latest_checkpoint_path(model_dir, "G_*.pth"))
+            utils.get_steps(utils.checkpoints.get_latest_checkpoint_path(model_dir, "G_*.pth"))
         )
         logger.info(
             f"******************Found the model. Current epoch is {epoch_str}, gloabl step is {global_step}*********************"
         )
     else:
         try:
-            _ = utils.load_safetensors(
+            _ = utils.safetensors.load_safetensors(
                 os.path.join(model_dir, "G_0.safetensors"), net_g
             )
-            _ = utils.load_safetensors(
+            _ = utils.safetensors.load_safetensors(
                 os.path.join(model_dir, "D_0.safetensors"), net_d
             )
             if net_dur_disc is not None:
-                _ = utils.load_safetensors(
+                _ = utils.safetensors.load_safetensors(
                     os.path.join(model_dir, "DUR_0.safetensors"), net_dur_disc
                 )
             logger.info("Loaded the pretrained models.")
@@ -485,14 +496,16 @@ def run():
 
         if epoch == hps.train.epochs:
             # Save the final models
-            utils.save_checkpoint(
+            assert optim_g is not None
+            utils.checkpoints.save_checkpoint(
                 net_g,
                 optim_g,
                 hps.train.learning_rate,
                 epoch,
                 os.path.join(model_dir, "G_{}.pth".format(global_step)),
             )
-            utils.save_checkpoint(
+            assert optim_d is not None
+            utils.checkpoints.save_checkpoint(
                 net_d,
                 optim_d,
                 hps.train.learning_rate,
@@ -500,14 +513,15 @@ def run():
                 os.path.join(model_dir, "D_{}.pth".format(global_step)),
             )
             if net_dur_disc is not None:
-                utils.save_checkpoint(
+                assert optim_dur_disc is not None
+                utils.checkpoints.save_checkpoint(
                     net_dur_disc,
                     optim_dur_disc,
                     hps.train.learning_rate,
                     epoch,
                     os.path.join(model_dir, "DUR_{}.pth".format(global_step)),
                 )
-            utils.save_safetensors(
+            utils.safetensors.save_safetensors(
                 net_g,
                 epoch,
                 os.path.join(
@@ -544,7 +558,7 @@ def train_and_evaluate(
     rank,
     local_rank,
     epoch,
-    hps,
+    hps: HyperParameters,
     nets,
     optims,
     schedulers,
@@ -778,14 +792,15 @@ def train_and_evaluate(
             ):
                 if not hps.speedup:
                     evaluate(hps, net_g, eval_loader, writer_eval)
-                utils.save_checkpoint(
+                assert hps.model_dir is not None
+                utils.checkpoints.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(
+                utils.checkpoints.save_checkpoint(
                     net_d,
                     optim_d,
                     hps.train.learning_rate,
@@ -793,7 +808,7 @@ def train_and_evaluate(
                     os.path.join(hps.model_dir, "D_{}.pth".format(global_step)),
                 )
                 if net_dur_disc is not None:
-                    utils.save_checkpoint(
+                    utils.checkpoints.save_checkpoint(
                         net_dur_disc,
                         optim_dur_disc,
                         hps.train.learning_rate,
@@ -802,13 +817,13 @@ def train_and_evaluate(
                     )
                 keep_ckpts = config.train_ms_config.keep_ckpts
                 if keep_ckpts > 0:
-                    utils.clean_checkpoints(
+                    utils.checkpoints.clean_checkpoints(
-                        path_to_models=hps.model_dir,
+                        model_dir_path=hps.model_dir,
                         n_ckpts_to_keep=keep_ckpts,
                         sort_by_time=True,
                     )
                 # Save safetensors (for inference) to `model_assets/{model_name}`
-                utils.save_safetensors(
+                utils.safetensors.save_safetensors(
                     net_g,
                     epoch,
                     os.path.join(

train_ms_jp_extra.py

@@ -1,25 +1,20 @@
 import argparse
 import datetime
-import gc
 import os
 import platform
 
 import torch
 import torch.distributed as dist
+from huggingface_hub import HfApi
 from torch.cuda.amp import GradScaler, autocast
 from torch.nn import functional as F
 from torch.nn.parallel import DistributedDataParallel as DDP
 from torch.utils.data import DataLoader
 from torch.utils.tensorboard import SummaryWriter
 from tqdm import tqdm
-from huggingface_hub import HfApi
 
 # logging.getLogger("numba").setLevel(logging.WARNING)
-import commons
 import default_style
-import utils
-from common.log import logger
-from common.stdout_wrapper import SAFE_STDOUT
 from config import config
 from data_utils import (
     DistributedBucketSampler,
@@ -28,13 +23,18 @@ from data_utils import (
 )
 from losses import WavLMLoss, discriminator_loss, feature_loss, generator_loss, kl_loss
 from mel_processing import mel_spectrogram_torch, spec_to_mel_torch
-from models_jp_extra import (
+from style_bert_vits2.logging import logger
+from style_bert_vits2.models import commons
+from style_bert_vits2.models import utils
+from style_bert_vits2.models.hyper_parameters import HyperParameters
+from style_bert_vits2.models.models_jp_extra import (
     DurationDiscriminator,
     MultiPeriodDiscriminator,
     SynthesizerTrn,
     WavLMDiscriminator,
 )
-from text.symbols import symbols
+from style_bert_vits2.nlp.symbols import SYMBOLS
+from style_bert_vits2.utils.stdout_wrapper import SAFE_STDOUT
 
 torch.backends.cuda.matmul.allow_tf32 = True
 torch.backends.cudnn.allow_tf32 = (
@@ -130,7 +130,7 @@ def run():
     local_rank = int(os.environ["LOCAL_RANK"])
     n_gpus = dist.get_world_size()
 
-    hps = utils.get_hparams_from_file(args.config)
+    hps = HyperParameters.load_from_json(args.config)
     # This is needed because we have to pass values to `train_and_evaluate()
     hps.model_dir = model_dir
     hps.speedup = args.speedup
@@ -247,10 +247,7 @@ def run():
             drop_last=False,
             collate_fn=collate_fn,
         )
-    if (
+    if hps.model.use_noise_scaled_mas is True:
-        "use_noise_scaled_mas" in hps.model.keys()
-        and hps.model.use_noise_scaled_mas is True
-    ):
         logger.info("Using noise scaled MAS for VITS2")
         mas_noise_scale_initial = 0.01
         noise_scale_delta = 2e-6
@@ -258,10 +255,7 @@ def run():
         logger.info("Using normal MAS for VITS1")
         mas_noise_scale_initial = 0.0
         noise_scale_delta = 0.0
-    if (
+    if hps.model.use_duration_discriminator is True:
-        "use_duration_discriminator" in hps.model.keys()
-        and hps.model.use_duration_discriminator is True
-    ):
         logger.info("Using duration discriminator for VITS2")
         net_dur_disc = DurationDiscriminator(
             hps.model.hidden_channels,
@@ -272,19 +266,13 @@ def run():
         ).cuda(local_rank)
     else:
         net_dur_disc = None
-    if (
+    if hps.model.use_wavlm_discriminator is True:
-        "use_wavlm_discriminator" in hps.model.keys()
-        and hps.model.use_wavlm_discriminator is True
-    ):
         net_wd = WavLMDiscriminator(
             hps.model.slm.hidden, hps.model.slm.nlayers, hps.model.slm.initial_channel
         ).cuda(local_rank)
     else:
         net_wd = None
-    if (
+    if hps.model.use_spk_conditioned_encoder is True:
-        "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"
@@ -293,13 +281,35 @@ def run():
         logger.info("Using normal encoder for VITS1")
 
     net_g = SynthesizerTrn(
-        len(symbols),
+        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,
+        # hps.model 以下のすべての値を引数に渡す
+        use_spk_conditioned_encoder = hps.model.use_spk_conditioned_encoder,
+        use_noise_scaled_mas = hps.model.use_noise_scaled_mas,
+        use_mel_posterior_encoder = hps.model.use_mel_posterior_encoder,
+        use_duration_discriminator = hps.model.use_duration_discriminator,
+        use_wavlm_discriminator = hps.model.use_wavlm_discriminator,
+        inter_channels = hps.model.inter_channels,
+        hidden_channels = hps.model.hidden_channels,
+        filter_channels = hps.model.filter_channels,
+        n_heads = hps.model.n_heads,
+        n_layers = hps.model.n_layers,
+        kernel_size = hps.model.kernel_size,
+        p_dropout = hps.model.p_dropout,
+        resblock = hps.model.resblock,
+        resblock_kernel_sizes = hps.model.resblock_kernel_sizes,
+        resblock_dilation_sizes = hps.model.resblock_dilation_sizes,
+        upsample_rates = hps.model.upsample_rates,
+        upsample_initial_channel = hps.model.upsample_initial_channel,
+        upsample_kernel_sizes = hps.model.upsample_kernel_sizes,
+        n_layers_q = hps.model.n_layers_q,
+        use_spectral_norm = hps.model.use_spectral_norm,
+        gin_channels = hps.model.gin_channels,
+        slm = hps.model.slm,
     ).cuda(local_rank)
     if getattr(hps.train, "freeze_JP_bert", False):
         logger.info("Freezing (JP) bert encoder !!!")
@@ -372,15 +382,11 @@ def run():
     if utils.is_resuming(model_dir):
         if net_dur_disc is not None:
             try:
-                _, _, dur_resume_lr, epoch_str = utils.load_checkpoint(
+                _, _, dur_resume_lr, epoch_str = utils.checkpoints.load_checkpoint(
-                    utils.latest_checkpoint_path(model_dir, "DUR_*.pth"),
+                    utils.checkpoints.get_latest_checkpoint_path(model_dir, "DUR_*.pth"),
                     net_dur_disc,
                     optim_dur_disc,
-                    skip_optimizer=(
+                    skip_optimizer=hps.train.skip_optimizer,
-                        hps.train.skip_optimizer
-                        if "skip_optimizer" in hps.train
-                        else True
-                    ),
                 )
                 if not optim_dur_disc.param_groups[0].get("initial_lr"):
                     optim_dur_disc.param_groups[0]["initial_lr"] = dur_resume_lr
@@ -390,15 +396,11 @@ def run():
                 print("Initialize dur_disc")
         if net_wd is not None:
             try:
-                _, optim_wd, wd_resume_lr, epoch_str = utils.load_checkpoint(
+                _, optim_wd, wd_resume_lr, epoch_str = utils.checkpoints.load_checkpoint(
-                    utils.latest_checkpoint_path(model_dir, "WD_*.pth"),
+                    utils.checkpoints.get_latest_checkpoint_path(model_dir, "WD_*.pth"),
                     net_wd,
                     optim_wd,
-                    skip_optimizer=(
+                    skip_optimizer=hps.train.skip_optimizer,
-                        hps.train.skip_optimizer
-                        if "skip_optimizer" in hps.train
-                        else True
-                    ),
                 )
                 if not optim_wd.param_groups[0].get("initial_lr"):
                     optim_wd.param_groups[0]["initial_lr"] = wd_resume_lr
@@ -408,21 +410,17 @@ def run():
                 logger.info("Initialize wavlm")
 
         try:
-            _, optim_g, g_resume_lr, epoch_str = utils.load_checkpoint(
+            _, optim_g, g_resume_lr, epoch_str = utils.checkpoints.load_checkpoint(
-                utils.latest_checkpoint_path(model_dir, "G_*.pth"),
+                utils.checkpoints.get_latest_checkpoint_path(model_dir, "G_*.pth"),
                 net_g,
                 optim_g,
-                skip_optimizer=(
+                skip_optimizer=hps.train.skip_optimizer,
-                    hps.train.skip_optimizer if "skip_optimizer" in hps.train else True
-                ),
             )
-            _, optim_d, d_resume_lr, epoch_str = utils.load_checkpoint(
+            _, optim_d, d_resume_lr, epoch_str = utils.checkpoints.load_checkpoint(
-                utils.latest_checkpoint_path(model_dir, "D_*.pth"),
+                utils.checkpoints.get_latest_checkpoint_path(model_dir, "D_*.pth"),
                 net_d,
                 optim_d,
-                skip_optimizer=(
+                skip_optimizer=hps.train.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
@@ -432,7 +430,7 @@ def run():
             epoch_str = max(epoch_str, 1)
             # global_step = (epoch_str - 1) * len(train_loader)
             global_step = int(
-                utils.get_steps(utils.latest_checkpoint_path(model_dir, "G_*.pth"))
+                utils.get_steps(utils.checkpoints.get_latest_checkpoint_path(model_dir, "G_*.pth"))
             )
             logger.info(
                 f"******************Found the model. Current epoch is {epoch_str}, gloabl step is {global_step}*********************"
@@ -446,18 +444,18 @@ def run():
             global_step = 0
     else:
         try:
-            _ = utils.load_safetensors(
+            _ = utils.safetensors.load_safetensors(
                 os.path.join(model_dir, "G_0.safetensors"), net_g
             )
-            _ = utils.load_safetensors(
+            _ = utils.safetensors.load_safetensors(
                 os.path.join(model_dir, "D_0.safetensors"), net_d
             )
             if net_dur_disc is not None:
-                _ = utils.load_safetensors(
+                _ = utils.safetensors.load_safetensors(
                     os.path.join(model_dir, "DUR_0.safetensors"), net_dur_disc
                 )
             if net_wd is not None:
-                _ = utils.load_safetensors(
+                _ = utils.safetensors.load_safetensors(
                     os.path.join(model_dir, "WD_0.safetensors"), net_wd
                 )
             logger.info("Loaded the pretrained models.")
@@ -564,14 +562,16 @@ def run():
             scheduler_wd.step()
         if epoch == hps.train.epochs:
             # Save the final models
-            utils.save_checkpoint(
+            assert optim_g is not None
+            utils.checkpoints.save_checkpoint(
                 net_g,
                 optim_g,
                 hps.train.learning_rate,
                 epoch,
                 os.path.join(model_dir, "G_{}.pth".format(global_step)),
             )
-            utils.save_checkpoint(
+            assert optim_d is not None
+            utils.checkpoints.save_checkpoint(
                 net_d,
                 optim_d,
                 hps.train.learning_rate,
@@ -579,7 +579,8 @@ def run():
                 os.path.join(model_dir, "D_{}.pth".format(global_step)),
             )
             if net_dur_disc is not None:
-                utils.save_checkpoint(
+                assert optim_dur_disc is not None
+                utils.checkpoints.save_checkpoint(
                     net_dur_disc,
                     optim_dur_disc,
                     hps.train.learning_rate,
@@ -587,14 +588,15 @@ def run():
                     os.path.join(model_dir, "DUR_{}.pth".format(global_step)),
                 )
             if net_wd is not None:
-                utils.save_checkpoint(
+                assert optim_wd is not None
+                utils.checkpoints.save_checkpoint(
                     net_wd,
                     optim_wd,
                     hps.train.learning_rate,
                     epoch,
                     os.path.join(model_dir, "WD_{}.pth".format(global_step)),
                 )
-            utils.save_safetensors(
+            utils.safetensors.save_safetensors(
                 net_g,
                 epoch,
                 os.path.join(
@@ -927,14 +929,14 @@ def train_and_evaluate(
             ):
                 if not hps.speedup:
                     evaluate(hps, net_g, eval_loader, writer_eval)
-                utils.save_checkpoint(
+                utils.checkpoints.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(
+                utils.checkpoints.save_checkpoint(
                     net_d,
                     optim_d,
                     hps.train.learning_rate,
@@ -942,7 +944,7 @@ def train_and_evaluate(
                     os.path.join(hps.model_dir, "D_{}.pth".format(global_step)),
                 )
                 if net_dur_disc is not None:
-                    utils.save_checkpoint(
+                    utils.checkpoints.save_checkpoint(
                         net_dur_disc,
                         optim_dur_disc,
                         hps.train.learning_rate,
@@ -950,7 +952,7 @@ def train_and_evaluate(
                         os.path.join(hps.model_dir, "DUR_{}.pth".format(global_step)),
                     )
                 if net_wd is not None:
-                    utils.save_checkpoint(
+                    utils.checkpoints.save_checkpoint(
                         net_wd,
                         optim_wd,
                         hps.train.learning_rate,
@@ -959,13 +961,13 @@ def train_and_evaluate(
                     )
                 keep_ckpts = config.train_ms_config.keep_ckpts
                 if keep_ckpts > 0:
-                    utils.clean_checkpoints(
+                    utils.checkpoints.clean_checkpoints(
-                        path_to_models=hps.model_dir,
+                        model_dir_path=hps.model_dir,
                         n_ckpts_to_keep=keep_ckpts,
                         sort_by_time=True,
                     )
                 # Save safetensors (for inference) to `model_assets/{model_name}`
-                utils.save_safetensors(
+                utils.safetensors.save_safetensors(
                     net_g,
                     epoch,
                     os.path.join(

transcribe.py

@@ -7,9 +7,9 @@ import yaml
 from faster_whisper import WhisperModel
 from tqdm import tqdm
 
-from common.constants import Languages
+from style_bert_vits2.constants import Languages
-from common.log import logger
+from style_bert_vits2.logging import logger
-from common.stdout_wrapper import SAFE_STDOUT
+from style_bert_vits2.utils.stdout_wrapper import SAFE_STDOUT
 
 
 def transcribe(wav_path: Path, initial_prompt=None, language="ja"):

update_status.py

@@ -1,93 +0,0 @@
-import os
-import gradio as gr
-
-lang_dict = {"EN(英文)": "_en", "ZH(中文)": "_zh", "JP(日语)": "_jp"}
-
-
-def raw_dir_convert_to_path(target_dir: str, lang):
-    res = target_dir.rstrip("/").rstrip("\\")
-    if (not target_dir.startswith("raw")) and (not target_dir.startswith("./raw")):
-        res = os.path.join("./raw", res)
-    if (
-        (not res.endswith("_zh"))
-        and (not res.endswith("_jp"))
-        and (not res.endswith("_en"))
-    ):
-        res += lang_dict[lang]
-    return res
-
-
-def update_g_files():
-    g_files = []
-    cnt = 0
-    for root, dirs, files in os.walk(os.path.abspath("./logs")):
-        for file in files:
-            if file.startswith("G_") and file.endswith(".pth"):
-                g_files.append(os.path.join(root, file))
-                cnt += 1
-    print(g_files)
-    return f"更新模型列表完成, 共找到{cnt}个模型", gr.Dropdown.update(choices=g_files)
-
-
-def update_c_files():
-    c_files = []
-    cnt = 0
-    for root, dirs, files in os.walk(os.path.abspath("./logs")):
-        for file in files:
-            if file.startswith("config.json"):
-                c_files.append(os.path.join(root, file))
-                cnt += 1
-    print(c_files)
-    return f"更新模型列表完成, 共找到{cnt}个配置文件", gr.Dropdown.update(
-        choices=c_files
-    )
-
-
-def update_model_folders():
-    subdirs = []
-    cnt = 0
-    for root, dirs, files in os.walk(os.path.abspath("./logs")):
-        for dir_name in dirs:
-            if os.path.basename(dir_name) != "eval":
-                subdirs.append(os.path.join(root, dir_name))
-                cnt += 1
-    print(subdirs)
-    return f"更新模型文件夹列表完成, 共找到{cnt}个文件夹", gr.Dropdown.update(
-        choices=subdirs
-    )
-
-
-def update_wav_lab_pairs():
-    wav_count = tot_count = 0
-    for root, _, files in os.walk("./raw"):
-        for file in files:
-            # print(file)
-            file_path = os.path.join(root, file)
-            if file.lower().endswith(".wav"):
-                lab_file = os.path.splitext(file_path)[0] + ".lab"
-                if os.path.exists(lab_file):
-                    wav_count += 1
-                tot_count += 1
-    return f"{wav_count} / {tot_count}"
-
-
-def update_raw_folders():
-    subdirs = []
-    cnt = 0
-    script_path = os.path.dirname(os.path.abspath(__file__))  # 获取当前脚本的绝对路径
-    raw_path = os.path.join(script_path, "raw")
-    print(raw_path)
-    os.makedirs(raw_path, exist_ok=True)
-    for root, dirs, files in os.walk(raw_path):
-        for dir_name in dirs:
-            relative_path = os.path.relpath(
-                os.path.join(root, dir_name), script_path
-            )  # 获取相对路径
-            subdirs.append(relative_path)
-            cnt += 1
-    print(subdirs)
-    return (
-        f"更新raw音频文件夹列表完成, 共找到{cnt}个文件夹",
-        gr.Dropdown.update(choices=subdirs),
-        gr.Textbox.update(value=update_wav_lab_pairs()),
-    )

utils.py

@@ -1,501 +0,0 @@
-import argparse
-import glob
-import json
-import logging
-import os
-import re
-import subprocess
-
-import numpy as np
-import torch
-from huggingface_hub import hf_hub_download
-from safetensors import safe_open
-from safetensors.torch import save_file
-from scipy.io.wavfile import read
-
-from common.log import logger
-
-MATPLOTLIB_FLAG = False
-
-
-def download_checkpoint(
-    dir_path, repo_config, token=None, regex="G_*.pth", mirror="openi"
-):
-    repo_id = repo_config["repo_id"]
-    f_list = glob.glob(os.path.join(dir_path, regex))
-    if f_list:
-        print("Use existed model, skip downloading.")
-        return
-    for file in ["DUR_0.pth", "D_0.pth", "G_0.pth"]:
-        hf_hub_download(repo_id, file, local_dir=dir_path, local_dir_use_symlinks=False)
-
-
-def load_checkpoint(
-    checkpoint_path, model, optimizer=None, skip_optimizer=False, for_infer=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"]
-    logger.info(
-        f"Loading model and optimizer at iteration {iteration} from {checkpoint_path}"
-    )
-    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
-        optimizer.load_state_dict(new_opt_dict)
-
-    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
-            new_state_dict[k] = saved_state_dict[k]
-            assert saved_state_dict[k].shape == v.shape, (
-                saved_state_dict[k].shape,
-                v.shape,
-            )
-        except:
-            # For upgrading from the old version
-            if "ja_bert_proj" in k:
-                v = torch.zeros_like(v)
-                logger.warning(
-                    f"Seems you are using the old version of the model, the {k} is automatically set to zero for backward compatibility"
-                )
-            elif "enc_q" in k and for_infer:
-                continue
-            else:
-                logger.error(f"{k} is not in the checkpoint {checkpoint_path}")
-
-            new_state_dict[k] = v
-
-    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 '{}' (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"):
-        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,
-    )
-
-
-def save_safetensors(model, iteration, checkpoint_path, is_half=False, for_infer=False):
-    """
-    Save model with safetensors.
-    """
-    if hasattr(model, "module"):
-        state_dict = model.module.state_dict()
-    else:
-        state_dict = model.state_dict()
-    keys = []
-    for k in state_dict:
-        if "enc_q" in k and for_infer:
-            continue  # noqa: E701
-        keys.append(k)
-
-    new_dict = (
-        {k: state_dict[k].half() for k in keys}
-        if is_half
-        else {k: state_dict[k] for k in keys}
-    )
-    new_dict["iteration"] = torch.LongTensor([iteration])
-    logger.info(f"Saved safetensors to {checkpoint_path}")
-    save_file(new_dict, checkpoint_path)
-
-
-def load_safetensors(checkpoint_path, model, for_infer=False):
-    """
-    Load safetensors model.
-    """
-
-    tensors = {}
-    iteration = None
-    with safe_open(checkpoint_path, framework="pt", device="cpu") as f:
-        for key in f.keys():
-            if key == "iteration":
-                iteration = f.get_tensor(key).item()
-            tensors[key] = f.get_tensor(key)
-    if hasattr(model, "module"):
-        result = model.module.load_state_dict(tensors, strict=False)
-    else:
-        result = model.load_state_dict(tensors, strict=False)
-    for key in result.missing_keys:
-        if key.startswith("enc_q") and for_infer:
-            continue
-        logger.warning(f"Missing key: {key}")
-    for key in result.unexpected_keys:
-        if key == "iteration":
-            continue
-        logger.warning(f"Unexpected key: {key}")
-    if iteration is None:
-        logger.info(f"Loaded '{checkpoint_path}'")
-    else:
-        logger.info(f"Loaded '{checkpoint_path}' (iteration {iteration})")
-    return model, iteration
-
-
-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")
-    for k, v in audios.items():
-        writer.add_audio(k, v, global_step, audio_sampling_rate)
-
-
-def is_resuming(dir_path):
-    # JP-ExtraバージョンではDURがなくWDがあったり変わるため、Gのみで判断する
-    g_list = glob.glob(os.path.join(dir_path, "G_*.pth"))
-    # d_list = glob.glob(os.path.join(dir_path, "D_*.pth"))
-    # dur_list = glob.glob(os.path.join(dir_path, "DUR_*.pth"))
-    return len(g_list) > 0
-
-
-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))))
-    try:
-        x = f_list[-1]
-    except IndexError:
-        raise ValueError(f"No checkpoint found in {dir_path} with regex {regex}")
-    return x
-
-
-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.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")
-    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 = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
-    plt.close()
-    return data
-
-
-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.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"
-    )
-    fig.colorbar(im, ax=ax)
-    xlabel = "Decoder timestep"
-    if info is not None:
-        xlabel += "\n\n" + info
-    plt.xlabel(xlabel)
-    plt.ylabel("Encoder timestep")
-    plt.tight_layout()
-
-    fig.canvas.draw()
-    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
-
-
-def load_wav_to_torch(full_path):
-    sampling_rate, data = read(full_path)
-    return torch.FloatTensor(data.astype(np.float32)), sampling_rate
-
-
-def load_filepaths_and_text(filename, split="|"):
-    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")
-
-    args = parser.parse_args()
-    model_dir = os.path.join("./logs", args.model)
-
-    if not os.path.exists(model_dir):
-        os.makedirs(model_dir)
-
-    config_path = args.config
-    config_save_path = os.path.join(model_dir, "config.json")
-    if init:
-        with open(config_path, "r", encoding="utf-8") as f:
-            data = f.read()
-        with open(config_save_path, "w", encoding="utf-8") as f:
-            f.write(data)
-    else:
-        with open(config_save_path, "r", vencoding="utf-8") as f:
-            data = f.read()
-    config = json.loads(data)
-    hparams = HParams(**config)
-    hparams.model_dir = model_dir
-    return hparams
-
-
-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
-    """
-    import re
-
-    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
-
-    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]
-            + x_sorted("WD_")[:-n_ckpts_to_keep]
-            + x_sorted("DUR_")[:-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:
-        data = f.read()
-    config = json.loads(data)
-
-    hparams = HParams(**config)
-    hparams.model_dir = model_dir
-    return hparams
-
-
-def get_hparams_from_file(config_path):
-    # print("config_path: ", config_path)
-    with open(config_path, "r", encoding="utf-8") as f:
-        data = f.read()
-    config = json.loads(data)
-
-    hparams = HParams(**config)
-    return hparams
-
-
-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.warning(
-            "{} is not a git repository, therefore hash value comparison will be ignored.".format(
-                source_dir
-            )
-        )
-        return
-
-    cur_hash = subprocess.getoutput("git rev-parse HEAD")
-
-    path = os.path.join(model_dir, "githash")
-    if os.path.exists(path):
-        saved_hash = open(path).read()
-        if saved_hash != cur_hash:
-            logger.warning(
-                "git hash values are different. {}(saved) != {}(current)".format(
-                    saved_hash[:8], cur_hash[:8]
-                )
-            )
-    else:
-        open(path, "w").write(cur_hash)
-
-
-def get_logger(model_dir, filename="train.log"):
-    global logger
-    logger = logging.getLogger(os.path.basename(model_dir))
-    logger.setLevel(logging.DEBUG)
-
-    formatter = logging.Formatter("%(asctime)s\t%(name)s\t%(levelname)s\t%(message)s")
-    if not os.path.exists(model_dir):
-        os.makedirs(model_dir)
-    h = logging.FileHandler(os.path.join(model_dir, filename))
-    h.setLevel(logging.DEBUG)
-    h.setFormatter(formatter)
-    logger.addHandler(h)
-    return logger
-
-
-class HParams:
-    def __init__(self, **kwargs):
-        for k, v in kwargs.items():
-            if type(v) == dict:
-                v = HParams(**v)
-            self[k] = v
-
-    def keys(self):
-        return self.__dict__.keys()
-
-    def items(self):
-        return self.__dict__.items()
-
-    def values(self):
-        return self.__dict__.values()
-
-    def __len__(self):
-        return len(self.__dict__)
-
-    def __getitem__(self, key):
-        return getattr(self, key)
-
-    def __setitem__(self, key, value):
-        return setattr(self, key, value)
-
-    def __contains__(self, key):
-        return key in self.__dict__
-
-    def __repr__(self):
-        return self.__dict__.__repr__()
-
-
-def load_model(model_path, config_path):
-    hps = get_hparams_from_file(config_path)
-    net = SynthesizerTrn(
-        # len(symbols),
-        108,
-        hps.data.filter_length // 2 + 1,
-        hps.train.segment_size // hps.data.hop_length,
-        n_speakers=hps.data.n_speakers,
-        **hps.model,
-    ).to("cpu")
-    _ = net.eval()
-    _ = load_checkpoint(model_path, net, None, skip_optimizer=True)
-    return net
-
-
-def mix_model(
-    network1, network2, output_path, voice_ratio=(0.5, 0.5), tone_ratio=(0.5, 0.5)
-):
-    if hasattr(network1, "module"):
-        state_dict1 = network1.module.state_dict()
-        state_dict2 = network2.module.state_dict()
-    else:
-        state_dict1 = network1.state_dict()
-        state_dict2 = network2.state_dict()
-    for k in state_dict1.keys():
-        if k not in state_dict2.keys():
-            continue
-        if "enc_p" in k:
-            state_dict1[k] = (
-                state_dict1[k].clone() * tone_ratio[0]
-                + state_dict2[k].clone() * tone_ratio[1]
-            )
-        else:
-            state_dict1[k] = (
-                state_dict1[k].clone() * voice_ratio[0]
-                + state_dict2[k].clone() * voice_ratio[1]
-            )
-    for k in state_dict2.keys():
-        if k not in state_dict1.keys():
-            state_dict1[k] = state_dict2[k].clone()
-    torch.save(
-        {"model": state_dict1, "iteration": 0, "optimizer": None, "learning_rate": 0},
-        output_path,
-    )
-
-
-def get_steps(model_path):
-    matches = re.findall(r"\d+", model_path)
-    return matches[-1] if matches else None