From 7f0b2528066e3b7e02bfed100dcef3ca5806dca3 Mon Sep 17 00:00:00 2001
From: tsukumi <tsukumijima@users.noreply.github.com>
Date: Fri, 8 Mar 2024 12:49:11 +0000
Subject: [PATCH] Refactor: add Pydantic model representing hyper-parameters of
 Style-Bert-VITS2 model

---
 requirements.txt                            |   1 +
 style_bert_vits2/models/hyper_parameters.py | 116 ++++++++++++++++++++
 style_bert_vits2/models/utils.py            |  33 ------
 3 files changed, 117 insertions(+), 33 deletions(-)
 create mode 100644 style_bert_vits2/models/hyper_parameters.py

diff --git a/requirements.txt b/requirements.txt
index 6bf329f..669af2f 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -14,6 +14,7 @@ numba
 numpy
 psutil
 pyannote.audio>=3.1.0
+pydantic
 pyloudnorm
 # pyopenjtalk-prebuilt  # Should be manually uninstalled
 pyopenjtalk-dict
diff --git a/style_bert_vits2/models/hyper_parameters.py b/style_bert_vits2/models/hyper_parameters.py
new file mode 100644
index 0000000..cee7924
--- /dev/null
+++ b/style_bert_vits2/models/hyper_parameters.py
@@ -0,0 +1,116 @@
+"""
+Style-Bert-VITS2 モデルのハイパーパラメータを表す Pydantic モデル。
+デフォルト値は configs/configs_jp_extra.json 内の定義と同一で、
+万が一ロードした config.json に存在しないキーがあった際のフェイルセーフとして適用される。
+"""
+
+from pathlib import Path
+from typing import Optional, Union
+
+from pydantic import BaseModel
+
+
+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 __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: dict[str, Union[int, str]] = {
+        "model": "./slm/wavlm-base-plus",
+        "sr": 16000,
+        "hidden": 768,
+        "nlayers": 13,
+        "initial_channel": 64
+    }
+
+class HyperParameters(BaseModel):
+    version: str = "2.0-JP-Extra"
+    model_name: str = 'dummy'
+    train: __HyperParametersTrain
+    data: __HyperParametersData
+    model: __HyperParametersModel
+
+
+    @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())
diff --git a/style_bert_vits2/models/utils.py b/style_bert_vits2/models/utils.py
index 8c7e842..9a7947f 100644
--- a/style_bert_vits2/models/utils.py
+++ b/style_bert_vits2/models/utils.py
@@ -357,39 +357,6 @@ def check_git_hash(model_dir):
         open(path, "w").write(cur_hash)
 
 
-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 get_hparams_from_file(config_path):
     # print("config_path: ", config_path)
     with open(config_path, "r", encoding="utf-8") as f: