fix: replace mispronounced words in TTS (#350)

* [Fix] Replace mispronounced words in TTS using hack method

* [Refactor] Move homophones_map.json to res folder

* [Refactor] Cache homophones_map.json and document source

* [Refactor] document process in class docstring.

ChatTTS/core.py

@@ -1,5 +1,6 @@
 
 import os
+import json
 import logging
 from functools import partial
 from omegaconf import OmegaConf
@@ -9,7 +10,7 @@ from vocos import Vocos
 from .model.dvae import DVAE
 from .model.gpt import GPT_warpper
 from .utils.gpu_utils import select_device
-from .utils.infer_utils import count_invalid_characters, detect_language, apply_character_map, apply_half2full_map
+from .utils.infer_utils import count_invalid_characters, detect_language, apply_character_map, apply_half2full_map, HomophonesReplacer
 from .utils.io_utils import get_latest_modified_file
 from .infer.api import refine_text, infer_code
 
@@ -22,6 +23,7 @@ class Chat:
     def __init__(self, ):
         self.pretrain_models = {}
         self.normalizer = {}
+        self.homophones_replacer = None
         self.logger = logging.getLogger(__name__)
         
     def check_model(self, level = logging.INFO, use_decoder = False):
@@ -136,6 +138,7 @@ class Chat:
         use_decoder=True,
         do_text_normalization=True,
         lang=None,
+        do_homophone_replacement=True
     ):
         
         assert self.check_model(use_decoder=use_decoder)
@@ -156,6 +159,10 @@ class Chat:
             if len(invalid_characters):
                 self.logger.log(logging.WARNING, f'Invalid characters found! : {invalid_characters}')
                 text[i] = apply_character_map(t)
+            if do_homophone_replacement and self.init_homophones_replacer():
+                text[i] = self.homophones_replacer.replace(t)
+                if t != text[i]:
+                    self.logger.log(logging.INFO, f'Homophones replace: {t} -> {text[i]}')
 
         if not skip_refine_text:
             text_tokens = refine_text(self.pretrain_models, text, **params_refine_text)['ids']
@@ -219,3 +226,17 @@ class Chat:
                     'Run: conda install -c conda-forge pynini=2.1.5 && pip install nemo_text_processing',
                 )
         return False
+
+    def init_homophones_replacer(self):
+        if self.homophones_replacer:
+            return True
+        else:
+            try:
+                self.homophones_replacer = HomophonesReplacer(os.path.join(os.path.dirname(__file__), 'res', 'homophones_map.json'))
+                self.logger.log(logging.INFO, 'homophones_replacer loaded.')
+                return True
+            except (IOError, json.JSONDecodeError) as e:
+                self.logger.log(logging.WARNING, f'Error loading homophones map: {e}')
+            except Exception as e:
+                self.logger.log(logging.WARNING, f'Error loading homophones_replacer: {e}')
+        return False

ChatTTS/utils/infer_utils.py

@@ -2,6 +2,8 @@
 import re
 import torch
 import torch.nn.functional as F
+import os
+import json
 
     
 class CustomRepetitionPenaltyLogitsProcessorRepeat():
@@ -44,7 +46,42 @@ class CustomRepetitionPenaltyLogitsProcessor():
         scores.scatter_(1, input_ids, score)
         
         return scores
-    
+
+class HomophonesReplacer:
+    """
+    Homophones Replacer
+
+    Replace the mispronounced characters with correctly pronounced ones.
+
+    Creation process of homophones_map.json:
+
+    1. Establish a word corpus using the [Tencent AI Lab Embedding Corpora v0.2.0 large] with 12 million entries. After cleaning, approximately 1.8 million entries remain. Use ChatTTS to infer the text.
+    2. Record discrepancies between the inferred and input text, identifying about 180,000 misread words.
+    3. Create a pinyin to common characters mapping using correctly read characters by ChatTTS.
+    4. For each discrepancy, extract the correct pinyin using [python-pinyin] and find homophones with the correct pronunciation from the mapping.
+
+    Thanks to:
+    [Tencent AI Lab Embedding Corpora for Chinese and English Words and Phrases](https://ai.tencent.com/ailab/nlp/en/embedding.html)
+    [python-pinyin](https://github.com/mozillazg/python-pinyin)
+
+    """
+    def __init__(self, map_file_path):
+        self.homophones_map = self.load_homophones_map(map_file_path)
+
+    def load_homophones_map(self, map_file_path):
+        with open(map_file_path, 'r', encoding='utf-8') as f:
+            homophones_map = json.load(f)
+        return homophones_map
+
+    def replace(self, text):
+        result = []
+        for char in text:
+            if char in self.homophones_map:
+                result.append(self.homophones_map[char])
+            else:
+                result.append(char)
+        return ''.join(result)
+
 def count_invalid_characters(s):
     
     s = re.sub(r'\[uv_break\]|\[laugh\]|\[lbreak\]', '', s)
@@ -138,4 +175,4 @@ def apply_half2full_map(text):
 
 def apply_character_map(text):
     translation_table = str.maketrans(character_map)
-    return text.translate(translation_table)
+    return text.translate(translation_table)