diff --git a/pyproject.toml b/pyproject.toml index 9ab0564..795fd76 100644 --- a/pyproject.toml +++ b/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/*\"", diff --git a/style_bert_vits2/nlp/bert_models.py b/style_bert_vits2/nlp/bert_models.py index be78141..1dcadb5 100644 --- a/style_bert_vits2/nlp/bert_models.py +++ b/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: diff --git a/style_bert_vits2/nlp/chinese/bert_feature.py b/style_bert_vits2/nlp/chinese/bert_feature.py index adc0788..d004593 100644 --- a/style_bert_vits2/nlp/chinese/bert_feature.py +++ b/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(): diff --git a/style_bert_vits2/nlp/english/bert_feature.py b/style_bert_vits2/nlp/english/bert_feature.py index 7692c9e..ede208e 100644 --- a/style_bert_vits2/nlp/english/bert_feature.py +++ b/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(): diff --git a/style_bert_vits2/nlp/japanese/bert_feature.py b/style_bert_vits2/nlp/japanese/bert_feature.py index d6e214b..0e2d9eb 100644 --- a/style_bert_vits2/nlp/japanese/bert_feature.py +++ b/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():