Refactor: add style_bert_vits2/text_processing/bert_models.py to hold loaded BERT models/tokenizer and replace all from_pretrained() to load_model/load_tokenizer

text/__init__.py

@@ -1,12 +1,17 @@
+from style_bert_vits2.constants import Languages
 from style_bert_vits2.text_processing.symbols import *
 
+
 _symbol_to_id = {s: i for i, s in enumerate(SYMBOLS)}
 
 
-def cleaned_text_to_sequence(cleaned_text, tones, language):
-    """Converts a string of text to a sequence of IDs corresponding to the symbols in the text.
+def cleaned_text_to_sequence(cleaned_text: str, tones: list[int], language: Languages):
+    """
+    Converts a string of text to a sequence of IDs corresponding to the symbols in the text.
+
     Args:
       text: string to convert to a sequence
+
     Returns:
       List of integers corresponding to the symbols in the text
     """
@@ -18,12 +23,19 @@ def cleaned_text_to_sequence(cleaned_text, tones, language):
     return phones, tones, lang_ids
 
 
-def get_bert(text, word2ph, language, device, assist_text=None, assist_text_weight=0.7):
-    if language == "ZH":
+def get_bert(
+    text: str,
+    word2ph,
+    language: Languages,
+    device: str,
+    assist_text: str | None = None,
+    assist_text_weight: float = 0.7,
+):
+    if language == Languages.ZH:
         from .chinese_bert import get_bert_feature
-    elif language == "EN":
+    elif language == Languages.EN:
         from .english_bert_mock import get_bert_feature
-    elif language == "JP":
+    elif language == Languages.JP:
         from .japanese_bert import get_bert_feature
     else:
         raise ValueError(f"Language {language} not supported")