tuna2134/sbv2-v2
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

Improve: Use the device_map option in transformers to load BERT models directly to the GPU

Browse Source
This commit is contained in:
tsukumi
2024-08-10 16:56:19 +09:00
parent bdf20e8911
commit 80dd6dbc22
7 changed files with 42 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
server_editor.py
View File

@@ -156,12 +156,6 @@ pyopenjtalk.initialize_worker()
# pyopenjtalk の辞書を更新
update_dict()
# 事前に BERT モデル/トークナイザーをロードしておく
## ここでロードしなくても必要になった際に自動ロードされるが、時間がかかるため事前にロードしておいた方が体験が良い
## server_editor.py は日本語にしか対応していないため、日本語の BERT モデル/トークナイザーのみロードする
bert_models.load_model(Languages.JP)
bert_models.load_tokenizer(Languages.JP)
class AudioResponse(Response):
media_type = "audio/wav"
@@ -194,6 +188,12 @@ port = int(args.port)
# download_default_models()
skip_static_files = bool(args.skip_static_files)
# 事前に BERT モデル/トークナイザーをロードしておく
## ここでロードしなくても必要になった際に自動ロードされるが、時間がかかるため事前にロードしておいた方が体験が良い
## server_editor.py は日本語にしか対応していないため、日本語の BERT モデル/トークナイザーのみロードする
bert_models.load_model(Languages.JP, device_map=device)
bert_models.load_tokenizer(Languages.JP)
model_holder = TTSModelHolder(model_dir, device)
if len(model_holder.model_names) == 0:
logger.error(f"Models not found in {model_dir}.")