From bc0729d970b1e2d490ca67f6136f9b81532e66e7 Mon Sep 17 00:00:00 2001
From: litagin02 <litagin02@gmail.com>
Date: Tue, 12 Mar 2024 15:23:18 +0900
Subject: [PATCH] Fix: initialize worker and dict for use

---
 app.py             | 20 +++++++++++++-------
 bert_gen.py        |  9 +++++++++
 preprocess_all.py  |  8 ++++++++
 preprocess_text.py |  8 ++++++++
 server_fastapi.py  |  5 +++++
 webui/__init__.py  | 16 ----------------
 6 files changed, 43 insertions(+), 23 deletions(-)
 delete mode 100644 webui/__init__.py

diff --git a/app.py b/app.py
index 60b91a6..0869ca5 100644
--- a/app.py
+++ b/app.py
@@ -6,16 +6,22 @@ import torch
 import yaml
 
 from style_bert_vits2.constants import GRADIO_THEME, VERSION
+from style_bert_vits2.nlp.japanese import pyopenjtalk_worker
+from style_bert_vits2.nlp.japanese.user_dict import update_dict
 from style_bert_vits2.tts_model import TTSModelHolder
-from webui import (
-    create_dataset_app,
-    create_inference_app,
-    create_merge_app,
-    create_style_vectors_app,
-    create_train_app,
-)
+from webui.dataset import create_dataset_app
+from webui.inference import create_inference_app
+from webui.merge import create_merge_app
+from webui.style_vectors import create_style_vectors_app
+from webui.train import create_train_app
 
 
+# このプロセスからはワーカーを起動して辞書を使いたいので、ここで初期化
+pyopenjtalk_worker.initialize_worker()
+
+# dict_data/ 以下の辞書データを pyopenjtalk に適用
+update_dict()
+
 # Get path settings
 with Path("configs/paths.yml").open("r", encoding="utf-8") as f:
     path_config: dict[str, str] = yaml.safe_load(f.read())
diff --git a/bert_gen.py b/bert_gen.py
index 5a16af7..2ded369 100644
--- a/bert_gen.py
+++ b/bert_gen.py
@@ -10,9 +10,18 @@ from style_bert_vits2.logging import logger
 from style_bert_vits2.models import commons
 from style_bert_vits2.models.hyper_parameters import HyperParameters
 from style_bert_vits2.nlp import cleaned_text_to_sequence, extract_bert_feature
+from style_bert_vits2.nlp.japanese import pyopenjtalk_worker
+from style_bert_vits2.nlp.japanese.user_dict import update_dict
 from style_bert_vits2.utils.stdout_wrapper import SAFE_STDOUT
 
 
+# このプロセスからはワーカーを起動して辞書を使いたいので、ここで初期化
+pyopenjtalk_worker.initialize_worker()
+
+# dict_data/ 以下の辞書データを pyopenjtalk に適用
+update_dict()
+
+
 def process_line(x):
     line, add_blank = x
     device = config.bert_gen_config.device
diff --git a/preprocess_all.py b/preprocess_all.py
index c41159c..2eaf846 100644
--- a/preprocess_all.py
+++ b/preprocess_all.py
@@ -1,9 +1,17 @@
 import argparse
 from multiprocessing import cpu_count
 
+from style_bert_vits2.nlp.japanese import pyopenjtalk_worker
+from style_bert_vits2.nlp.japanese.user_dict import update_dict
 from webui.train import preprocess_all
 
 
+# このプロセスからはワーカーを起動して辞書を使いたいので、ここで初期化
+pyopenjtalk_worker.initialize_worker()
+
+# dict_data/ 以下の辞書データを pyopenjtalk に適用
+update_dict()
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument(
diff --git a/preprocess_text.py b/preprocess_text.py
index ef2e259..6bc8585 100644
--- a/preprocess_text.py
+++ b/preprocess_text.py
@@ -10,9 +10,17 @@ from tqdm import tqdm
 from config import config
 from style_bert_vits2.logging import logger
 from style_bert_vits2.nlp import clean_text
+from style_bert_vits2.nlp.japanese import pyopenjtalk_worker
+from style_bert_vits2.nlp.japanese.user_dict import update_dict
 from style_bert_vits2.utils.stdout_wrapper import SAFE_STDOUT
 
 
+# このプロセスからはワーカーを起動して辞書を使いたいので、ここで初期化
+pyopenjtalk_worker.initialize_worker()
+
+# dict_data/ 以下の辞書データを pyopenjtalk に適用
+update_dict()
+
 preprocess_text_config = config.preprocess_text_config
 
 
diff --git a/server_fastapi.py b/server_fastapi.py
index 4ac4a7e..a8e8c93 100644
--- a/server_fastapi.py
+++ b/server_fastapi.py
@@ -35,7 +35,9 @@ from style_bert_vits2.constants import (
 )
 from style_bert_vits2.logging import logger
 from style_bert_vits2.nlp import bert_models
+from style_bert_vits2.nlp.japanese import pyopenjtalk_worker
 from style_bert_vits2.nlp.japanese import pyopenjtalk_worker as pyopenjtalk
+from style_bert_vits2.nlp.japanese.user_dict import update_dict
 from style_bert_vits2.tts_model import TTSModel, TTSModelHolder
 
 
@@ -46,6 +48,9 @@ ln = config.server_config.language
 ## pyopenjtalk_worker は TCP ソケットサーバーのため、ここで起動する
 pyopenjtalk.initialize_worker()
 
+# dict_data/ 以下の辞書データを pyopenjtalk に適用
+update_dict()
+
 # 事前に BERT モデル/トークナイザーをロードしておく
 ## ここでロードしなくても必要になった際に自動ロードされるが、時間がかかるため事前にロードしておいた方が体験が良い
 bert_models.load_model(Languages.JP)
diff --git a/webui/__init__.py b/webui/__init__.py
deleted file mode 100644
index 4bc8efa..0000000
--- a/webui/__init__.py
+++ /dev/null
@@ -1,16 +0,0 @@
-from .dataset import create_dataset_app
-from .inference import create_inference_app
-from .merge import create_merge_app
-from .style_vectors import create_style_vectors_app
-from .train import create_train_app
-
-
-class TrainSettings:
-    def __init__(self, setting_json):
-        self.setting_json = setting_json
-
-    def __enter__(self):
-        return self
-
-    def __exit__(self, exc_type, exc_value, traceback):
-        pass