From ebb37ad7bb7e54a1757b9da0820e88a2349b476e Mon Sep 17 00:00:00 2001
From: litagin02 <litagin02@gmail.com>
Date: Fri, 8 Mar 2024 18:18:38 +0900
Subject: [PATCH] Fix multiprocessing bert_gen error of pyopenjtalk_worker

---
 bert_gen.py           |  5 ++---
 infer.py              |  3 ++-
 preprocess_all.py     |  7 +++++--
 text/__init__.py      | 13 -------------
 text/japanese.py      |  6 ++----
 text/japanese_bert.py |  2 +-
 webui/merge.py        |  4 ++--
 7 files changed, 14 insertions(+), 26 deletions(-)

diff --git a/bert_gen.py b/bert_gen.py
index 1e4fb61..19dbbb6 100644
--- a/bert_gen.py
+++ b/bert_gen.py
@@ -6,14 +6,13 @@ import torch.multiprocessing as mp
 from tqdm import tqdm
 
 import commons
+from text.get_bert import get_bert
 import text.pyopenjtalk_worker as pyopenjtalk
 import utils
 from common.log import logger
 from common.stdout_wrapper import SAFE_STDOUT
 from config import config
-from text import cleaned_text_to_sequence, get_bert
-
-pyopenjtalk.initialize()
+from text import cleaned_text_to_sequence
 
 
 def process_line(x):
diff --git a/infer.py b/infer.py
index 3707df1..436ea8b 100644
--- a/infer.py
+++ b/infer.py
@@ -1,10 +1,11 @@
 import torch
 
 import commons
+from text.get_bert import get_bert
 import utils
 from models import SynthesizerTrn
 from models_jp_extra import SynthesizerTrn as SynthesizerTrnJPExtra
-from text import cleaned_text_to_sequence, get_bert
+from text import cleaned_text_to_sequence
 from text.cleaner import clean_text
 from text.symbols import symbols
 from common.log import logger
diff --git a/preprocess_all.py b/preprocess_all.py
index fd3dcd5..62c0b4f 100644
--- a/preprocess_all.py
+++ b/preprocess_all.py
@@ -1,5 +1,5 @@
 import argparse
-from webui_train import preprocess_all
+from webui.train import preprocess_all
 from multiprocessing import cpu_count
 
 if __name__ == "__main__":
@@ -75,7 +75,10 @@ if __name__ == "__main__":
         default=200,
     )
     parser.add_argument(
-        "--yomi_error", type=str, help="Yomi error. raise, skip, use", default="raise"
+        "--yomi_error",
+        type=str,
+        help="Yomi error. Options: raise, skip, use",
+        default="raise",
     )
 
     args = parser.parse_args()
diff --git a/text/__init__.py b/text/__init__.py
index d8ae88d..7ba6e20 100644
--- a/text/__init__.py
+++ b/text/__init__.py
@@ -16,16 +16,3 @@ def cleaned_text_to_sequence(cleaned_text, tones, language):
     lang_id = language_id_map[language]
     lang_ids = [lang_id for i in phones]
     return phones, tones, lang_ids
-
-
-def get_bert(text, word2ph, language, device, assist_text=None, assist_text_weight=0.7):
-    if language == "ZH":
-        from .chinese_bert import get_bert_feature
-    elif language == "EN":
-        from .english_bert_mock import get_bert_feature
-    elif language == "JP":
-        from .japanese_bert import get_bert_feature
-    else:
-        raise ValueError(f"Language {language} not supported")
-
-    return get_bert_feature(text, word2ph, device, assist_text, assist_text_weight)
diff --git a/text/japanese.py b/text/japanese.py
index fea0eaa..f7bfcae 100644
--- a/text/japanese.py
+++ b/text/japanese.py
@@ -4,9 +4,6 @@ import re
 import unicodedata
 from pathlib import Path
 
-from . import pyopenjtalk_worker as pyopenjtalk
-
-pyopenjtalk.initialize()
 from num2words import num2words
 from transformers import AutoTokenizer
 
@@ -16,9 +13,10 @@ from text.japanese_mora_list import (
     mora_kata_to_mora_phonemes,
     mora_phonemes_to_mora_kata,
 )
-
 from text.user_dict import update_dict
 
+from . import pyopenjtalk_worker as pyopenjtalk
+
 # 最初にpyopenjtalkの辞書を更新
 update_dict()
 
diff --git a/text/japanese_bert.py b/text/japanese_bert.py
index dcee0f3..9ebc64a 100644
--- a/text/japanese_bert.py
+++ b/text/japanese_bert.py
@@ -52,7 +52,7 @@ def get_bert_feature(
             style_res = torch.cat(style_res["hidden_states"][-3:-2], -1)[0].cpu()
             style_res_mean = style_res.mean(0)
 
-    assert len(word2ph) == len(text) + 2, text
+    assert len(word2ph) == len(text) + 2, f"word2ph: {word2ph}, text: {text}"
     word2phone = word2ph
     phone_level_feature = []
     for i in range(len(word2phone)):
diff --git a/webui/merge.py b/webui/merge.py
index c9386ef..aeab08d 100644
--- a/webui/merge.py
+++ b/webui/merge.py
@@ -254,7 +254,7 @@ def simple_tts(model_name, text, style=DEFAULT_STYLE, style_weight=1.0):
     return model.infer(text, style=style, style_weight=style_weight)
 
 
-def update_two_model_names_dropdown():
+def update_two_model_names_dropdown(model_holder: ModelHolder):
     new_names, new_files, _ = model_holder.update_model_names_gr()
     return new_names, new_files, new_names, new_files
 
@@ -455,7 +455,7 @@ def create_merge_app(model_holder: ModelHolder) -> gr.Blocks:
         )
 
         refresh_button.click(
-            update_two_model_names_dropdown,
+            lambda: update_two_model_names_dropdown(model_holder),
             outputs=[model_name_a, model_path_a, model_name_b, model_path_b],
         )