Refactor: Use bert_models.transfer_model()

pyproject.toml

@@ -80,10 +80,12 @@ cov = ["test-cov", "cov-report"]
 detached = true
 dependencies = ["black[jupyter]", "isort"]
 [tool.hatch.envs.style.scripts]
+# Usage: `hatch run style:check`
 check = [
     "black --check --diff .",
     "isort --check-only --diff --profile black --gitignore --lai 2 . --sg \"Data/*\" --sg \"inputs/*\" --sg \"model_assets/*\" --sg \"static/*\"",
 ]
+# Usage: `hatch run style:fmt`
 fmt = [
     "black .",
     "isort --profile black --gitignore --lai 2 . --sg \"Data/*\" --sg \"inputs/*\" --sg \"model_assets/*\" --sg \"static/*\"",

style_bert_vits2/nlp/bert_models.py

@@ -159,6 +159,8 @@ def load_tokenizer(
 def transfer_model(language: Languages, device: str) -> None:
     """
     指定された言語の BERT モデルを、指定されたデバイスに移動する。
+    モデルのロード後に推論デバイスを変更したい場合に利用する。
+    既に指定されたデバイスにモデルがロードされている場合は何も行われない。
 
     Args:
         language (Languages): モデルを移動する言語
@@ -171,7 +173,9 @@ def transfer_model(language: Languages, device: str) -> None:
     current_device = str(__loaded_models[language].device)
     if current_device != device:
         __loaded_models[language].to(device)  # type: ignore
-        logger.info(f"Transferred the {language} BERT model from {current_device} to {device}")
+        logger.info(
+            f"Transferred the {language} BERT model from {current_device} to {device}"
+        )
 
 
 def unload_model(language: Languages) -> None:

style_bert_vits2/nlp/chinese/bert_feature.py

@@ -29,7 +29,8 @@ def extract_bert_feature(
 
     if device == "cuda" and not torch.cuda.is_available():
         device = "cpu"
-    model = bert_models.load_model(Languages.ZH).to(device)  # type: ignore
+    model = bert_models.load_model(Languages.ZH)
+    bert_models.transfer_model(Languages.ZH, device)
 
     style_res_mean = None
     with torch.no_grad():

style_bert_vits2/nlp/english/bert_feature.py

@@ -29,7 +29,8 @@ def extract_bert_feature(
 
     if device == "cuda" and not torch.cuda.is_available():
         device = "cpu"
-    model = bert_models.load_model(Languages.EN).to(device)  # type: ignore
+    model = bert_models.load_model(Languages.EN)
+    bert_models.transfer_model(Languages.EN, device)
 
     style_res_mean = None
     with torch.no_grad():

style_bert_vits2/nlp/japanese/bert_feature.py

@@ -36,7 +36,8 @@ def extract_bert_feature(
 
     if device == "cuda" and not torch.cuda.is_available():
         device = "cpu"
-    model = bert_models.load_model(Languages.JP).to(device)  # type: ignore
+    model = bert_models.load_model(Languages.JP)
+    bert_models.transfer_model(Languages.JP, device)
 
     style_res_mean = None
     with torch.no_grad():