Merge pull request #95 from zunan-islands/dev-refactor

細かな修正 (#92 を参照)

bert_gen.py

@@ -73,10 +73,10 @@ if __name__ == "__main__":
     config_path = args.config
     hps = HyperParameters.load_from_json(config_path)
     lines = []
-    with open(hps.data.training_files, encoding="utf-8") as f:
+    with open(hps.data.training_files, "r", encoding="utf-8") as f:
         lines.extend(f.readlines())
 
-    with open(hps.data.validation_files, encoding="utf-8") as f:
+    with open(hps.data.validation_files, "r", encoding="utf-8") as f:
         lines.extend(f.readlines())
     add_blank = [hps.data.add_blank] * len(lines)
 

config.py

@@ -238,7 +238,7 @@ class Config:
                 "If you have no special needs, please do not modify default_config.yml."
             )
             # sys.exit(0)
-        with open(file=config_path, mode="r", encoding="utf-8") as file:
+        with open(config_path, "r", encoding="utf-8") as file:
             yaml_config: Dict[str, any] = yaml.safe_load(file.read())
             model_name: str = yaml_config["model_name"]
             self.model_name: str = model_name

pyproject.toml

@@ -51,6 +51,27 @@ Source = "https://github.com/litagin02/Style-Bert-VITS2"
 [tool.hatch.version]
 path = "style_bert_vits2/constants.py"
 
+[tool.hatch.build.targets.sdist]
+only-include = [
+    ".vscode",
+    "dict_data/default.csv",
+    "docs",
+    "style_bert_vits2",
+    "tests",
+    "LGPL_LICENSE",
+    "LICENSE",
+    "pyproject.toml",
+    "README.md",
+]
+exclude = [
+    ".git",
+    ".gitignore",
+    ".gitattributes",
+]
+
+[tool.hatch.build.targets.wheel]
+packages = ["style_bert_vits2"]
+
 [tool.hatch.envs.test]
 dependencies = [
     "coverage[toml]>=6.5",
@@ -75,18 +96,18 @@ cov = [
 [tool.hatch.envs.style]
 detached = true
 dependencies = [
-  "black",
-  "isort",
+    "black",
+    "isort",
 ]
 [tool.hatch.envs.style.scripts]
 check = [
-  "black --check --diff .",
-  "isort --check-only --diff --profile black --gitignore --lai 2 .",
+    "black --check --diff .",
+    "isort --check-only --diff --profile black --gitignore --lai 2 .",
 ]
 fmt = [
-  "black .",
-  "isort --profile black --gitignore --lai 2 .",
-  "check",
+    "black .",
+    "isort --profile black --gitignore --lai 2 .",
+    "check",
 ]
 
 [[tool.hatch.envs.test.matrix]]

style_bert_vits2/models/utils/__init__.py

@@ -180,7 +180,7 @@ def load_filepaths_and_text(
         list[list[str]]: ファイルパスとテキストのリスト
     """
 
-    with open(filename, encoding="utf-8") as f:
+    with open(filename, "r", encoding="utf-8") as f:
         filepaths_and_text = [line.strip().split(split) for line in f]
     return filepaths_and_text
 
@@ -249,7 +249,8 @@ def check_git_hash(model_dir_path: Union[str, Path]) -> None:
 
     path = os.path.join(model_dir_path, "githash")
     if os.path.exists(path):
-        saved_hash = open(path).read()
+        with open(path, "r", encoding="utf-8") as f:
+            saved_hash = f.read()
         if saved_hash != cur_hash:
             logger.warning(
                 "git hash values are different. {}(saved) != {}(current)".format(
@@ -257,4 +258,5 @@ def check_git_hash(model_dir_path: Union[str, Path]) -> None:
                 )
             )
     else:
-        open(path, "w").write(cur_hash)
+        with open(path, "w", encoding="utf-8") as f:
+            f.write(cur_hash)

style_bert_vits2/models/utils/checkpoints.py

@@ -85,7 +85,7 @@ def load_checkpoint(
     else:
         model.load_state_dict(new_state_dict, strict=False)
 
-    logger.info("Loaded '{}' (iteration {})".format(checkpoint_path, iteration))
+    logger.info(f"Loaded '{checkpoint_path}' (iteration {iteration})")
 
     return model, optimizer, learning_rate, iteration
 

style_bert_vits2/nlp/bert_models.py

@@ -38,12 +38,15 @@ __loaded_tokenizers: dict[
 def load_model(
     language: Languages,
     pretrained_model_name_or_path: Optional[str] = None,
+    cache_dir: Optional[str] = None,
+    revision: str = "main",
 ) -> Union[PreTrainedModel, DebertaV2Model]:
     """
     指定された言語の BERT モデルをロードし、ロード済みの BERT モデルを返す。
     一度ロードされていれば、ロード済みの BERT モデルを即座に返す。
-    ライブラリ利用時は常に pretrain_model_name_or_path (Hugging Face のリポジトリ名 or ローカルのファイルパス) を指定する必要がある。
+    ライブラリ利用時は常に必ず pretrain_model_name_or_path (Hugging Face のリポジトリ名 or ローカルのファイルパス) を指定する必要がある。
     ロードにはそれなりに時間がかかるため、ライブラリ利用前に明示的に pretrained_model_name_or_path を指定してロードしておくべき。
+    cache_dir と revision は pretrain_model_name_or_path がリポジトリ名の場合のみ有効。
 
     Style-Bert-VITS2 では、BERT モデルに下記の 3 つが利用されている。
     これ以外の BERT モデルを指定した場合は正常に動作しない可能性が高い。
@@ -54,6 +57,8 @@ def load_model(
     Args:
         language (Languages): ロードする学習済みモデルの対象言語
         pretrained_model_name_or_path (Optional[str]): ロードする学習済みモデルの名前またはパス。指定しない場合はデフォルトのパスが利用される (デフォルト: None)
+        cache_dir (Optional[str]): モデルのキャッシュディレクトリ。指定しない場合はデフォルトのキャッシュディレクトリが利用される (デフォルト: None)
+        revision (str): モデルの Hugging Face 上の Git リビジョン。指定しない場合は最新の main ブランチの内容が利用される (デフォルト: None)
 
     Returns:
         Union[PreTrainedModel, DebertaV2Model]: ロード済みの BERT モデル
@@ -75,10 +80,14 @@ def load_model(
     if language == Languages.EN:
         model = cast(
             DebertaV2Model,
-            DebertaV2Model.from_pretrained(pretrained_model_name_or_path),
+            DebertaV2Model.from_pretrained(
+                pretrained_model_name_or_path, cache_dir=cache_dir, revision=revision
+            ),
         )
     else:
-        model = AutoModelForMaskedLM.from_pretrained(pretrained_model_name_or_path)
+        model = AutoModelForMaskedLM.from_pretrained(
+            pretrained_model_name_or_path, cache_dir=cache_dir, revision=revision
+        )
     __loaded_models[language] = model
     logger.info(
         f"Loaded the {language} BERT model from {pretrained_model_name_or_path}"
@@ -90,12 +99,15 @@ def load_model(
 def load_tokenizer(
     language: Languages,
     pretrained_model_name_or_path: Optional[str] = None,
+    cache_dir: Optional[str] = None,
+    revision: str = "main",
 ) -> Union[PreTrainedTokenizer, PreTrainedTokenizerFast, DebertaV2Tokenizer]:
     """
     指定された言語の BERT モデルをロードし、ロード済みの BERT トークナイザーを返す。
     一度ロードされていれば、ロード済みの BERT トークナイザーを即座に返す。
-    ライブラリ利用時は常に pretrain_model_name_or_path (Hugging Face のリポジトリ名 or ローカルのファイルパス) を指定する必要がある。
+    ライブラリ利用時は常に必ず pretrain_model_name_or_path (Hugging Face のリポジトリ名 or ローカルのファイルパス) を指定する必要がある。
     ロードにはそれなりに時間がかかるため、ライブラリ利用前に明示的に pretrained_model_name_or_path を指定してロードしておくべき。
+    cache_dir と revision は pretrain_model_name_or_path がリポジトリ名の場合のみ有効。
 
     Style-Bert-VITS2 では、BERT モデルに下記の 3 つが利用されている。
     これ以外の BERT モデルを指定した場合は正常に動作しない可能性が高い。
@@ -106,6 +118,8 @@ def load_tokenizer(
     Args:
         language (Languages): ロードする学習済みモデルの対象言語
         pretrained_model_name_or_path (Optional[str]): ロードする学習済みモデルの名前またはパス。指定しない場合はデフォルトのパスが利用される (デフォルト: None)
+        cache_dir (Optional[str]): モデルのキャッシュディレクトリ。指定しない場合はデフォルトのキャッシュディレクトリが利用される (デフォルト: None)
+        revision (str): モデルの Hugging Face 上の Git リビジョン。指定しない場合は最新の main ブランチの内容が利用される (デフォルト: None)
 
     Returns:
         Union[PreTrainedTokenizer, PreTrainedTokenizerFast, DebertaV2Tokenizer]: ロード済みの BERT トークナイザー
@@ -125,9 +139,17 @@ def load_tokenizer(
     # BERT トークナイザーをロードし、辞書に格納して返す
     ## 英語のみ DebertaV2Tokenizer でロードする必要がある
     if language == Languages.EN:
-        tokenizer = DebertaV2Tokenizer.from_pretrained(pretrained_model_name_or_path)
+        tokenizer = DebertaV2Tokenizer.from_pretrained(
+            pretrained_model_name_or_path,
+            cache_dir=cache_dir,
+            revision=revision,
+        )
     else:
-        tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path)
+        tokenizer = AutoTokenizer.from_pretrained(
+            pretrained_model_name_or_path,
+            cache_dir=cache_dir,
+            revision=revision,
+        )
     __loaded_tokenizers[language] = tokenizer
     logger.info(
         f"Loaded the {language} BERT tokenizer from {pretrained_model_name_or_path}"

style_bert_vits2/nlp/chinese/g2p.py

@@ -8,10 +8,10 @@ from style_bert_vits2.nlp.chinese.tone_sandhi import ToneSandhi
 from style_bert_vits2.nlp.symbols import PUNCTUATIONS
 
 
-__PINYIN_TO_SYMBOL_MAP = {
-    line.split("\t")[0]: line.strip().split("\t")[1]
-    for line in open(Path(__file__).parent / "opencpop-strict.txt").readlines()
-}
+with open(Path(__file__).parent / "opencpop-strict.txt", "r", encoding="utf-8") as f:
+    __PINYIN_TO_SYMBOL_MAP = {
+        line.split("\t")[0]: line.strip().split("\t")[1] for line in f.readlines()
+    }
 
 
 def g2p(text: str) -> tuple[list[str], list[int], list[int]]:

style_bert_vits2/nlp/english/cmudict.py

@@ -20,7 +20,7 @@ def get_dict() -> dict[str, list[list[str]]]:
 def read_dict() -> dict[str, list[list[str]]]:
     g2p_dict = {}
     start_line = 49
-    with open(CMU_DICT_PATH) as f:
+    with open(CMU_DICT_PATH, "r", encoding="utf-8") as f:
         line = f.readline()
         line_index = 1
         while line:

style_gen.py

@@ -73,7 +73,7 @@ if __name__ == "__main__":
     device = config.style_gen_config.device
 
     training_lines: list[str] = []
-    with open(hps.data.training_files, encoding="utf-8") as f:
+    with open(hps.data.training_files, "r", encoding="utf-8") as f:
         training_lines.extend(f.readlines())
     with ThreadPoolExecutor(max_workers=num_processes) as executor:
         training_results = list(
@@ -94,7 +94,7 @@ if __name__ == "__main__":
         )
 
     val_lines: list[str] = []
-    with open(hps.data.validation_files, encoding="utf-8") as f:
+    with open(hps.data.validation_files, "r", encoding="utf-8") as f:
         val_lines.extend(f.readlines())
 
     with ThreadPoolExecutor(max_workers=num_processes) as executor:

webui/merge.py

@@ -47,11 +47,11 @@ def merge_style(model_name_a, model_name_b, weight, output_name, style_triple_li
         os.path.join(assets_root, model_name_b, "style_vectors.npy")
     )  # (style_num_b, 256)
     with open(
-        os.path.join(assets_root, model_name_a, "config.json"), encoding="utf-8"
+        os.path.join(assets_root, model_name_a, "config.json"), "r", encoding="utf-8"
     ) as f:
         config_a = json.load(f)
     with open(
-        os.path.join(assets_root, model_name_b, "config.json"), encoding="utf-8"
+        os.path.join(assets_root, model_name_b, "config.json"), "r", encoding="utf-8"
     ) as f:
         config_b = json.load(f)
     style2id_a = config_a["data"]["style2id"]
@@ -88,7 +88,7 @@ def merge_style(model_name_a, model_name_b, weight, output_name, style_triple_li
     # recipe.jsonを読み込んで、style_triple_listを追記
     info_path = os.path.join(assets_root, output_name, "recipe.json")
     if os.path.exists(info_path):
-        with open(info_path, encoding="utf-8") as f:
+        with open(info_path, "r", encoding="utf-8") as f:
             info = json.load(f)
     else:
         info = {}
@@ -261,12 +261,12 @@ def update_two_model_names_dropdown(model_holder: TTSModelHolder):
 
 def load_styles_gr(model_name_a, model_name_b):
     config_path_a = os.path.join(assets_root, model_name_a, "config.json")
-    with open(config_path_a, encoding="utf-8") as f:
+    with open(config_path_a, "r", encoding="utf-8") as f:
         config_a = json.load(f)
     styles_a = list(config_a["data"]["style2id"].keys())
 
     config_path_b = os.path.join(assets_root, model_name_b, "config.json")
-    with open(config_path_b, encoding="utf-8") as f:
+    with open(config_path_b, "r", encoding="utf-8") as f:
         config_b = json.load(f)
     styles_b = list(config_b["data"]["style2id"].keys())