Refactor: replace utils.HParams with HyperParameters Pydantic model

HyperParameters is largely a drop-in replacement for utils.HParams, which ensures type safety for hyper-parameters.

style_bert_vits2/models/hyper_parameters.py

@@ -1,16 +1,16 @@
 """
 Style-Bert-VITS2 モデルのハイパーパラメータを表す Pydantic モデル。
-デフォルト値は configs/configs_jp_extra.json 内の定義と同一で、
+デフォルト値は configs/configs_jp_extra.json 内の定義と概ね同一で、
 万が一ロードした config.json に存在しないキーがあった際のフェイルセーフとして適用される。
 """
 
 from pathlib import Path
 from typing import Optional, Union
 
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict
 
 
-class __HyperParametersTrain(BaseModel):
+class HyperParametersTrain(BaseModel):
     log_interval: int = 200
     eval_interval: int = 1000
     seed: int = 42
@@ -36,7 +36,8 @@ class __HyperParametersTrain(BaseModel):
     freeze_style: bool = False
     freeze_decoder: bool = False
 
-class __HyperParametersData(BaseModel):
+
+class HyperParametersData(BaseModel):
     use_jp_extra: bool = True
     training_files: str = "Data/dummy/train.list"
     validation_files: str = "Data/dummy/val.list"
@@ -59,7 +60,8 @@ class __HyperParametersData(BaseModel):
       "Neutral": 0,
     }
 
-class __HyperParametersModel(BaseModel):
+
+class HyperParametersModel(BaseModel):
     use_spk_conditioned_encoder: bool = True
     use_noise_scaled_mas: bool = True
     use_mel_posterior_encoder: bool = False
@@ -93,12 +95,21 @@ class __HyperParametersModel(BaseModel):
         "initial_channel": 64
     }
 
+
 class HyperParameters(BaseModel):
-    version: str = "2.0-JP-Extra"
     model_name: str = 'dummy'
-    train: __HyperParametersTrain
-    data: __HyperParametersData
-    model: __HyperParametersModel
+    version: str = "2.0-JP-Extra"
+    train: HyperParametersTrain
+    data: HyperParametersData
+    model: HyperParametersModel
+
+    # 以下は学習時にのみ動的に設定されるパラメータ (通常 config.json には存在しない)
+    model_dir: Optional[str] = None
+    speedup: bool = False
+    repo_id: Optional[str] = None
+
+    # model_ 以下を Pydantic の保護対象から除外する
+    model_config = ConfigDict(protected_namespaces=())
 
 
     @staticmethod
@@ -112,5 +123,6 @@ class HyperParameters(BaseModel):
         Returns:
             HyperParameters: ハイパーパラメータ
         """
+
         with open(json_path, "r") as f:
             return HyperParameters.model_validate_json(f.read())

style_bert_vits2/models/infer.py

@@ -1,34 +1,81 @@
+from typing import Any, cast, Optional, Union
+
 import torch
-from typing import Optional
+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):
+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(
-            len(SYMBOLS),
-            hps.data.filter_length // 2 + 1,
-            hps.train.segment_size // hps.data.hop_length,
-            n_speakers=hps.data.n_speakers,
-            **hps.model,
+            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(
-            len(SYMBOLS),
-            hps.data.filter_length // 2 + 1,
-            hps.train.segment_size // hps.data.hop_length,
+            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,
+            # 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()
@@ -44,7 +91,7 @@ def get_net_g(model_path: str, version: str, device: str, hps):
 def get_text(
     text: str,
     language_str: Languages,
-    hps,
+    hps: HyperParameters,
     device: str,
     assist_text: Optional[str] = None,
     assist_text_weight: float = 0.7,
@@ -111,15 +158,15 @@ def get_text(
 
 def infer(
     text: str,
-    style_vec,
+    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,
-    net_g,
+    hps: HyperParameters,
+    net_g: Union[SynthesizerTrn, SynthesizerTrnJPExtra],
     device: str,
     skip_start: bool = False,
     skip_end: bool = False,
@@ -159,25 +206,25 @@ def infer(
         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)
+        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 = net_g.infer(
+            output = cast(SynthesizerTrnJPExtra, net_g).infer(
                 x_tst,
                 x_tst_lengths,
                 sid_tensor,
                 tones,
                 lang_ids,
                 ja_bert,
-                style_vec=style_vec,
+                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 = net_g.infer(
+            output = cast(SynthesizerTrn, net_g).infer(
                 x_tst,
                 x_tst_lengths,
                 sid_tensor,
@@ -186,7 +233,7 @@ def infer(
                 bert,
                 ja_bert,
                 en_bert,
-                style_vec=style_vec,
+                style_vec=style_vec_tensor,
                 sdp_ratio=sdp_ratio,
                 noise_scale=noise_scale,
                 noise_scale_w=noise_scale_w,
@@ -209,110 +256,5 @@ def infer(
         return audio
 
 
-def infer_multilang(
-    text: str,
-    style_vec,
-    sdp_ratio: float,
-    noise_scale: float,
-    noise_scale_w: float,
-    length_scale: float,
-    sid: int,
-    language: Languages,
-    hps,
-    net_g,
-    device: str,
-    skip_start: bool = False,
-    skip_end: bool = 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)  # type: ignore
-        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
-
-
 class InvalidToneError(ValueError):
     pass

style_bert_vits2/models/models.py

@@ -983,10 +983,10 @@ class SynthesizerTrn(nn.Module):
         en_bert,
         style_vec,
         noise_scale=0.667,
-        length_scale=1,
+        length_scale=1.0,
         noise_scale_w=0.8,
         max_len=None,
-        sdp_ratio=0,
+        sdp_ratio=0.0,
         y=None,
     ):
         # x, m_p, logs_p, x_mask = self.enc_p(x, x_lengths, tone, language, bert)

style_bert_vits2/models/models_jp_extra.py

@@ -1029,10 +1029,10 @@ class SynthesizerTrn(nn.Module):
         bert,
         style_vec,
         noise_scale=0.667,
-        length_scale=1,
+        length_scale=1.0,
         noise_scale_w=0.8,
         max_len=None,
-        sdp_ratio=0,
+        sdp_ratio=0.0,
         y=None,
     ):
         # x, m_p, logs_p, x_mask = self.enc_p(x, x_lengths, tone, language, bert)

style_bert_vits2/models/utils.py

@@ -355,45 +355,3 @@ def check_git_hash(model_dir):
             )
     else:
         open(path, "w").write(cur_hash)
-
-
-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
-
-
-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__()