Add suppport for Thai

melo/data_utils.py

@@ -173,7 +173,7 @@ def get_text(self, text, word2ph, phone, tone, language_str, wav_path):
             if language_str in ["ZH"]:
                 bert = bert
                 ja_bert = torch.zeros(768, len(phone))
-            elif language_str in ["JP", "EN", "ZH_MIX_EN", "KR", 'SP', 'ES', 'FR', 'DE', 'RU']:
+            elif language_str in ["JP", "EN", "ZH_MIX_EN", "KR", 'SP', 'ES', 'FR', 'DE', 'RU', 'TH']:
                 ja_bert = bert
                 bert = torch.zeros(1024, len(phone))
             else:

melo/models.py

@@ -830,6 +830,23 @@ def __init__(
             num_languages=num_languages,
             num_tones=num_tones,
         )
+        self.enc_p = TextEncoder(
+            219,  # Initialize with the original symbol size
+            inter_channels,
+            hidden_channels,
+            filter_channels,
+            n_heads,
+            n_layers,
+            kernel_size,
+            p_dropout,
+            gin_channels=self.enc_gin_channels,
+            num_languages=num_languages,
+            num_tones=num_tones,
+        )
+        if n_vocab != 219:
+            old_embeddings = self.enc_p.emb
+            new_num_tokens = n_vocab
+            self.enc_p.emb = self.get_resized_embeddings(old_embeddings, new_num_tokens)
         self.dec = Generator(
             inter_channels,
             resblock,
@@ -884,6 +901,23 @@ def __init__(
             self.ref_enc = ReferenceEncoder(spec_channels, gin_channels, layernorm=norm_refenc)
         self.use_vc = use_vc
 
+    def get_resized_embeddings(self, old_embeddings, new_num_tokens):
+        old_num_tokens, old_embedding_dim = old_embeddings.weight.size()
+        if old_num_tokens == new_num_tokens:
+            return old_embeddings
+
+        if not isinstance(old_embeddings, nn.Embedding):
+            raise TypeError(
+                f"Old embeddings are of type {type(old_embeddings)}, which is not an instance of {nn.Embedding}. "
+                f"You should either use a different resize function or make sure that `old_embeddings` are an instance of {nn.Embedding}."
+            )
+
+        new_embeddings = nn.Embedding(new_num_tokens, old_embedding_dim).to(
+            device=old_embeddings.weight.device, dtype=old_embeddings.weight.dtype
+        )
+        new_embeddings.weight.data[:old_num_tokens, :] = old_embeddings.weight.data[:old_num_tokens, :]
+
+        return new_embeddings
 
     def forward(self, x, x_lengths, y, y_lengths, sid, tone, language, bert, ja_bert):
         if self.n_speakers > 0:
@@ -998,7 +1032,7 @@ def infer(
             sdp_ratio
         ) + self.dp(x, x_mask, g=g) * (1 - sdp_ratio)
         w = torch.exp(logw) * x_mask * length_scale
-        
+
         w_ceil = torch.ceil(w)
         y_lengths = torch.clamp_min(torch.sum(w_ceil, [1, 2]), 1).long()
         y_mask = torch.unsqueeze(commons.sequence_mask(y_lengths, None), 1).to(
@@ -1020,7 +1054,7 @@ def infer(
         # print('max/min of o:', o.max(), o.min())
         return o, attn, y_mask, (z, z_p, m_p, logs_p)
 
-    def voice_conversion(self, y, y_lengths, sid_src, sid_tgt, tau=1.0):        
+    def voice_conversion(self, y, y_lengths, sid_src, sid_tgt, tau=1.0):
         g_src = sid_src
         g_tgt = sid_tgt
         z, m_q, logs_q, y_mask = self.enc_q(y, y_lengths, g=g_src, tau=tau)

melo/text/__init__.py

@@ -28,8 +28,9 @@ def get_bert(norm_text, word2ph, language, device):
     from .spanish_bert import get_bert_feature as sp_bert
     from .french_bert import get_bert_feature as fr_bert
     from .korean import get_bert_feature as kr_bert
+    from .thai import get_bert_feature as th_bert
 
-    lang_bert_func_map = {"ZH": zh_bert, "EN": en_bert, "JP": jp_bert, 'ZH_MIX_EN': zh_mix_en_bert, 
-                          'FR': fr_bert, 'SP': sp_bert, 'ES': sp_bert, "KR": kr_bert}
+    lang_bert_func_map = {"ZH": zh_bert, "EN": en_bert, "JP": jp_bert, 'ZH_MIX_EN': zh_mix_en_bert,
+                          'FR': fr_bert, 'SP': sp_bert, 'ES': sp_bert, "KR": kr_bert, "TH": th_bert}
     bert = lang_bert_func_map[language](norm_text, word2ph, device)
     return bert

melo/text/cleaner.py

@@ -1,9 +1,9 @@
-from . import chinese, japanese, english, chinese_mix, korean, french, spanish
+from . import chinese, japanese, english, chinese_mix, korean, french, spanish, thai
 from . import cleaned_text_to_sequence
 import copy
 
 language_module_map = {"ZH": chinese, "JP": japanese, "EN": english, 'ZH_MIX_EN': chinese_mix, 'KR': korean,
-                    'FR': french, 'SP': spanish, 'ES': spanish}
+                    'FR': french, 'SP': spanish, 'ES': spanish, 'TH': thai}
 
 
 def clean_text(text, language):
@@ -17,13 +17,13 @@ def clean_text_bert(text, language, device=None):
     language_module = language_module_map[language]
     norm_text = language_module.text_normalize(text)
     phones, tones, word2ph = language_module.g2p(norm_text)
-    
+
     word2ph_bak = copy.deepcopy(word2ph)
     for i in range(len(word2ph)):
         word2ph[i] = word2ph[i] * 2
     word2ph[0] += 1
     bert = language_module.get_bert_feature(norm_text, word2ph, device=device)
-    
+
     return norm_text, phones, tones, word2ph_bak, bert
 
 
@@ -33,4 +33,4 @@ def text_to_sequence(text, language):
 
 
 if __name__ == "__main__":
-    pass
+    pass

melo/text/english_bert.py

@@ -2,12 +2,13 @@
 from transformers import AutoTokenizer, AutoModelForMaskedLM
 import sys
 
-model_id = 'bert-base-uncased'
-tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = None
+models = {}
+tokenizers = {}
 
-def get_bert_feature(text, word2ph, device=None):
+def get_bert_feature(text, word2ph, device=None, model_id='airesearch/wangchanberta-base-att-spm-uncased'):
     global model
+    global tokenizer
+
     if (
         sys.platform == "darwin"
         and torch.backends.mps.is_available()
@@ -16,24 +17,33 @@ def get_bert_feature(text, word2ph, device=None):
         device = "mps"
     if not device:
         device = "cuda"
-    if model is None:
+
+    if model_id not in models:
         model = AutoModelForMaskedLM.from_pretrained(model_id).to(
             device
         )
+        models[model_id] = model
+        tokenizer = AutoTokenizer.from_pretrained(model_id)
+        tokenizers[model_id] = tokenizer
+    else:
+        model = models[model_id]
+        tokenizer = tokenizers[model_id]
+
     with torch.no_grad():
         inputs = tokenizer(text, return_tensors="pt")
+        import pdb; pdb.set_trace();
         for i in inputs:
             inputs[i] = inputs[i].to(device)
         res = model(**inputs, output_hidden_states=True)
         res = torch.cat(res["hidden_states"][-3:-2], -1)[0].cpu()
-        
-    assert inputs["input_ids"].shape[-1] == len(word2ph)
-    word2phone = word2ph
-    phone_level_feature = []
-    for i in range(len(word2phone)):
-        repeat_feature = res[i].repeat(word2phone[i], 1)
-        phone_level_feature.append(repeat_feature)
-
-    phone_level_feature = torch.cat(phone_level_feature, dim=0)
+
+        assert inputs["input_ids"].shape[-1] == len(word2ph), f"{inputs['input_ids'].shape[-1]}/{len(word2ph)}"
+
+        word2phone = word2ph[1:-1]
+        phone_level_feature = []
+        for i in range(len(word2phone)):
+            repeat_feature = res[i].repeat(word2phone[i], 1)
+            phone_level_feature.append(repeat_feature)
+        phone_level_feature = torch.cat(phone_level_feature, dim=0)
 
     return phone_level_feature.T