Update inference.py

inference.py

@@ -1,476 +1,3 @@
-# import re
-# import yaml
-# from munch import Munch
-# import numpy as np
-# import librosa
-# import noisereduce as nr
-# from meldataset import TextCleaner
-# import torch
-# import torchaudio
-# from nltk.tokenize import word_tokenize
-# import nltk
-# nltk.download('punkt_tab')
-
-# from models import ProsodyPredictor, TextEncoder, StyleEncoder
-# from Modules.hifigan import Decoder
-
-# import sys
-# import phonemizer
-# if sys.platform.startswith("win"):
-#     try:
-#         from phonemizer.backend.espeak.wrapper import EspeakWrapper
-#         import espeakng_loader
-#         EspeakWrapper.set_library(espeakng_loader.get_library_path())
-#     except Exception as e:
-#         print(e)
-
-# def espeak_phn(text, lang):
-#     try:
-#         my_phonemizer = phonemizer.backend.EspeakBackend(language=lang, preserve_punctuation=True,  with_stress=True, language_switch='remove-flags')
-#         return my_phonemizer.phonemize([text])[0]
-#     except Exception as e:
-#         print(e)
-
-# class Preprocess:
-#     def __text_normalize(self, text):
-#         punctuation = ["，", "、", "،", ";", "(", "．", "。", "…", "!", "–", ":", "?"]
-#         map_to = "."
-#         punctuation_pattern = re.compile(f"[{''.join(re.escape(p) for p in punctuation)}]")
-#         #replace punctuation that acts like a comma or period
-#         text = punctuation_pattern.sub(map_to, text)
-#         #replace consecutive whitespace chars with a single space and strip leading/trailing spaces
-#         text = re.sub(r'\s+', ' ', text).strip()
-#         return text
-#     def __merge_fragments(self, texts, n):
-#         merged = []
-#         i = 0
-#         while i < len(texts):
-#             fragment = texts[i]
-#             j = i + 1
-#             while len(fragment.split()) < n and j < len(texts):
-#                 fragment += ", " + texts[j]
-#                 j += 1
-#             merged.append(fragment)
-#             i = j
-#         if len(merged[-1].split()) < n and len(merged) > 1: #handle last sentence
-#             merged[-2] = merged[-2] + ", " + merged[-1]
-#             del merged[-1]
-#         else:
-#             merged[-1] = merged[-1]
-#         return merged
-#     def wave_preprocess(self, wave):
-#         to_mel = torchaudio.transforms.MelSpectrogram(n_mels=80, n_fft=2048, win_length=1200, hop_length=300)
-#         mean, std = -4, 4
-#         wave_tensor = torch.from_numpy(wave).float()
-#         mel_tensor = to_mel(wave_tensor)
-#         mel_tensor = (torch.log(1e-5 + mel_tensor.unsqueeze(0)) - mean) / std
-#         return mel_tensor
-#     def text_preprocess(self, text, n_merge=12):
-#         text_norm = self.__text_normalize(text).split(".")#split by sentences.
-#         text_norm = [s.strip() for s in text_norm]
-#         text_norm = list(filter(lambda x: x != '', text_norm)) #filter empty index
-#         text_norm = self.__merge_fragments(text_norm, n=n_merge) #merge if a sentence has less that n 
-#         return text_norm
-#     def length_to_mask(self, lengths):
-#         mask = torch.arange(lengths.max()).unsqueeze(0).expand(lengths.shape[0], -1).type_as(lengths)
-#         mask = torch.gt(mask+1, lengths.unsqueeze(1))
-#         return mask
-
-
-# import re
-# import sys
-# import yaml
-# import nltk
-# import numpy as np
-# import librosa
-# import torch
-# import phonemizer
-# import noisereduce as nr
-# from munch import Munch
-# from nltk.tokenize import word_tokenize
-
-# from meldataset import TextCleaner
-# from models import ProsodyPredictor, TextEncoder, StyleEncoder
-# from Modules.hifigan import Decoder
-
-# # Không download ở runtime trên Space (dễ treo / fail do network)
-# # nltk.download('punkt_tab')
-# # Nếu bạn cần, chuyển sang packages/requirements hoặc chạy local build step.
-# # Trên Space, khuyến nghị bỏ phụ thuộc NLTK hoặc thay bằng tokenizer đơn giản.
-
-# if sys.platform.startswith("win"):
-#     try:
-#         from phonemizer.backend.espeak.wrapper import EspeakWrapper
-#         import espeakng_loader
-#         EspeakWrapper.set_library(espeakng_loader.get_library_path())
-#     except Exception as e:
-#         print(e)
-
-
-# def espeak_phn(text, lang):
-#     try:
-#         my_phonemizer = phonemizer.backend.EspeakBackend(
-#             language=lang,
-#             preserve_punctuation=True,
-#             with_stress=True,
-#             language_switch="remove-flags",
-#         )
-#         return my_phonemizer.phonemize([text])[0]
-#     except Exception as e:
-#         print(e)
-#         return text
-
-
-# class Preprocess:
-#     def __text_normalize(self, text):
-#         punctuation = ["，", "、", "،", ";", "(", "．", "。", "…", "!", "–", ":", "?"]
-#         map_to = "."
-#         punctuation_pattern = re.compile(f"[{''.join(re.escape(p) for p in punctuation)}]")
-#         text = punctuation_pattern.sub(map_to, text)
-#         text = re.sub(r"\s+", " ", text).strip()
-#         return text
-
-#     def __merge_fragments(self, texts, n):
-#         merged = []
-#         i = 0
-#         while i < len(texts):
-#             fragment = texts[i]
-#             j = i + 1
-#             while len(fragment.split()) < n and j < len(texts):
-#                 fragment += ", " + texts[j]
-#                 j += 1
-#             merged.append(fragment)
-#             i = j
-
-#         if len(merged) > 1 and len(merged[-1].split()) < n:
-#             merged[-2] = merged[-2] + ", " + merged[-1]
-#             del merged[-1]
-#         return merged
-
-#     def wave_preprocess(self, wave, sr=24000):
-#         """
-#         Thay torchaudio bằng librosa để tránh dependency torchaudio trên HF Space.
-#         Output giống shape cũ: (1, 80, T)
-#         """
-#         if wave is None:
-#             raise ValueError("wave is None")
-#         wave = np.asarray(wave)
-#         if wave.ndim != 1:
-#             wave = wave.squeeze()
-#         wave = wave.astype(np.float32)
-
-#         # Mel spectrogram (power). Nếu muốn khớp torchaudio default power=2.0, để power=2.0.
-#         mel = librosa.feature.melspectrogram(
-#             y=wave,
-#             sr=sr,
-#             n_fft=2048,
-#             win_length=1200,
-#             hop_length=300,
-#             n_mels=80,
-#             power=2.0,
-#         )  # (80, T)
-
-#         mean, std = -4, 4
-#         mel = np.log(1e-5 + mel)
-#         mel = (mel - mean) / std
-
-#         mel_tensor = torch.from_numpy(mel).float().unsqueeze(0)  # (1, 80, T)
-#         return mel_tensor
-
-#     def text_preprocess(self, text, n_merge=12):
-#         text_norm = self.__text_normalize(text).split(".")
-#         text_norm = [s.strip() for s in text_norm]
-#         text_norm = list(filter(lambda x: x != "", text_norm))
-#         text_norm = self.__merge_fragments(text_norm, n=n_merge)
-#         return text_norm
-
-#     def length_to_mask(self, lengths):
-#         mask = torch.arange(lengths.max()).unsqueeze(0).expand(lengths.shape[0], -1).type_as(lengths)
-#         mask = torch.gt(mask + 1, lengths.unsqueeze(1))
-#         return mask
-
-
-# #For inference only
-# class StyleTTS2(torch.nn.Module):
-#     def __init__(self, config_path, models_path):
-#         super().__init__()
-#         self.register_buffer("get_device", torch.empty(0))
-#         self.preprocess = Preprocess()
-#         self.ref_s = None
-#         config = yaml.safe_load(open(config_path, "r", encoding="utf-8"))
-        
-#         try:
-#             symbols = (
-#                             list(config['symbol']['pad']) +
-#                             list(config['symbol']['punctuation']) +
-#                             list(config['symbol']['letters']) +
-#                             list(config['symbol']['letters_ipa']) +
-#                             list(config['symbol']['extend'])
-#                         )
-#             symbol_dict = {}
-#             for i in range(len((symbols))):
-#                 symbol_dict[symbols[i]] = i
-            
-#             n_token = len(symbol_dict) + 1
-#             print("\nFound:", n_token, "symbols")
-#         except Exception as e:
-#             print(f"\nERROR: Cannot find {e} in config file!\nYour config file is likely outdated, please download updated version from the repository.")
-#             raise SystemExit(1)
-        
-#         args = self.__recursive_munch(config['model_params'])
-#         args['n_token'] = n_token
-        
-#         self.cleaner = TextCleaner(symbol_dict, debug=False)
-
-#         assert args.decoder.type in ['hifigan'], 'Decoder type unknown'
-
-#         self.decoder            = Decoder(dim_in=args.hidden_dim, style_dim=args.style_dim, dim_out=args.n_mels,
-#                                         resblock_kernel_sizes = args.decoder.resblock_kernel_sizes,
-#                                         upsample_rates = args.decoder.upsample_rates,
-#                                         upsample_initial_channel=args.decoder.upsample_initial_channel,
-#                                         resblock_dilation_sizes=args.decoder.resblock_dilation_sizes,
-#                                         upsample_kernel_sizes=args.decoder.upsample_kernel_sizes)
-#         self.predictor           = ProsodyPredictor(style_dim=args.style_dim, d_hid=args.hidden_dim, nlayers=args.n_layer, max_dur=args.max_dur, dropout=args.dropout)
-#         self.text_encoder        = TextEncoder(channels=args.hidden_dim, kernel_size=5, depth=args.n_layer, n_symbols=args.n_token)
-#         self.style_encoder       = StyleEncoder(dim_in=args.dim_in, style_dim=args.style_dim, max_conv_dim=args.hidden_dim)# acoustic style encoder
-
-#         self.__load_models(models_path)
-    
-#     def __recursive_munch(self, d):
-#         if isinstance(d, dict):
-#             return Munch((k, self.__recursive_munch(v)) for k, v in d.items())
-#         elif isinstance(d, list):
-#             return [self.__recursive_munch(v) for v in d]
-#         else:
-#             return d
-        
-#     def __init_replacement_func(self, replacements):
-#         replacement_iter = iter(replacements)
-#         def replacement(match):
-#             return next(replacement_iter)
-#         return replacement
-    
-#     def __replace_outliers_zscore(self, tensor, threshold=3.0, factor=0.95):
-#         mean = tensor.mean()
-#         std = tensor.std()
-#         z = (tensor - mean) / std
-
-#         # Identify outliers
-#         outlier_mask = torch.abs(z) > threshold
-#         # Compute replacement value, respecting sign
-#         sign = torch.sign(tensor - mean)
-#         replacement = mean + sign * (threshold * std * factor)
-
-#         result = tensor.clone()
-#         result[outlier_mask] = replacement[outlier_mask]
-
-#         return result
-        
-#     def __load_models(self, models_path):
-#         module_params = []
-#         model = {'decoder':self.decoder, 'predictor':self.predictor, 'text_encoder':self.text_encoder, 'style_encoder':self.style_encoder}
-
-#         params_whole = torch.load(models_path, map_location='cpu')
-#         params = params_whole['net']
-#         params = {key: value for key, value in params.items() if key in model.keys()}
-
-#         for key in model:
-#             try:
-#                 model[key].load_state_dict(params[key])
-#             except:
-#                 from collections import OrderedDict
-#                 state_dict = params[key]
-#                 new_state_dict = OrderedDict()
-#                 for k, v in state_dict.items():
-#                     name = k[7:] # remove `module.`
-#                     new_state_dict[name] = v
-#                 model[key].load_state_dict(new_state_dict, strict=False)
-
-#             total_params = sum(p.numel() for p in model[key].parameters())
-#             print(key,":",total_params)
-#             module_params.append(total_params)
-
-#         print('\nTotal',":",sum(module_params))
-
-#     def __compute_style(self, path, denoise, split_dur):
-#         device = self.get_device.device
-#         denoise = min(denoise, 1)
-#         if split_dur != 0: split_dur = max(int(split_dur), 1)
-#         max_samples = 24000*20 #max 20 seconds ref audio
-#         print("Computing the style for:", path)
-        
-#         wave, sr = librosa.load(path, sr=24000)
-#         audio, index = librosa.effects.trim(wave, top_db=30)
-#         if sr != 24000:
-#             audio = librosa.resample(audio, sr, 24000)
-#         if len(audio) > max_samples:
-#             audio = audio[:max_samples]
-        
-#         if denoise > 0.0:
-#             audio_denoise = nr.reduce_noise(y=audio, sr=sr, n_fft=2048, win_length=1200, hop_length=300)
-#             audio = audio*(1-denoise) + audio_denoise*denoise
-
-#         with torch.no_grad():
-#             if split_dur>0 and len(audio)/sr>=4: #Only effective if audio length is >= 4s
-#                 #This option will split the ref audio to multiple parts, calculate styles and average them
-#                 count = 0
-#                 ref_s = None
-#                 jump = sr*split_dur
-#                 total_len = len(audio)
-                
-#                 #Need to init before the loop
-#                 mel_tensor = self.preprocess.wave_preprocess(audio[0:jump]).to(device)
-#                 ref_s = self.style_encoder(mel_tensor.unsqueeze(1))
-#                 count += 1
-#                 for i in range(jump, total_len, jump):
-#                     if i+jump >= total_len:
-#                         left_dur = (total_len-i)/sr
-#                         if left_dur >= 1: #Still count if left over dur is >= 1s
-#                             mel_tensor = self.preprocess.wave_preprocess(audio[i:total_len]).to(device)
-#                             ref_s += self.style_encoder(mel_tensor.unsqueeze(1))
-#                             count += 1
-#                         continue
-#                     mel_tensor = self.preprocess.wave_preprocess(audio[i:i+jump]).to(device)
-#                     ref_s += self.style_encoder(mel_tensor.unsqueeze(1))
-#                     count += 1
-#                 ref_s /= count
-#             else:
-#                 mel_tensor = self.preprocess.wave_preprocess(audio).to(device)
-#                 ref_s = self.style_encoder(mel_tensor.unsqueeze(1))
-
-#         return ref_s
-        
-#     def __inference(self, phonem, ref_s, speed=1, prev_d_mean=0, t=0.1):
-#         device = self.get_device.device
-#         speed = min(max(speed, 0.0001), 2) #speed range [0, 2]
-        
-#         phonem = ' '.join(word_tokenize(phonem))
-#         tokens = self.cleaner(phonem)
-#         tokens.insert(0, 0)
-#         tokens.append(0)
-#         tokens = torch.LongTensor(tokens).to(device).unsqueeze(0)
-        
-#         with torch.no_grad():
-#             input_lengths = torch.LongTensor([tokens.shape[-1]]).to(device)
-#             text_mask = self.preprocess.length_to_mask(input_lengths).to(device)
-
-#             # encode
-#             t_en = self.text_encoder(tokens, input_lengths, text_mask)
-#             s = ref_s.to(device)
-        
-#             # cal alignment
-#             d = self.predictor.text_encoder(t_en, s, input_lengths, text_mask)
-#             x, _ = self.predictor.lstm(d)
-#             duration = self.predictor.duration_proj(x) 
-#             duration = torch.sigmoid(duration).sum(axis=-1)
-
-#             if prev_d_mean != 0:#Stabilize speaking speed between splits
-#                 dur_stats = torch.empty(duration.shape).normal_(mean=prev_d_mean, std=duration.std()).to(device)
-#             else:
-#                 dur_stats = torch.empty(duration.shape).normal_(mean=duration.mean(), std=duration.std()).to(device)
-#             duration = duration*(1-t) + dur_stats*t
-#             duration[:,1:-2] = self.__replace_outliers_zscore(duration[:,1:-2]) #Normalize outlier
-
-#             duration /= speed
-                
-#             pred_dur = torch.round(duration.squeeze()).clamp(min=1)
-#             pred_aln_trg = torch.zeros(input_lengths, int(pred_dur.sum().data))
-#             c_frame = 0
-#             for i in range(pred_aln_trg.size(0)):
-#                 pred_aln_trg[i, c_frame:c_frame + int(pred_dur[i].data)] = 1
-#                 c_frame += int(pred_dur[i].data)
-#             alignment = pred_aln_trg.unsqueeze(0).to(device)
-
-#             # encode prosody
-#             en = (d.transpose(-1, -2) @ alignment)
-#             F0_pred, N_pred = self.predictor.F0Ntrain(en, s)
-#             asr = (t_en @ pred_aln_trg.unsqueeze(0).to(device))
-
-#             out = self.decoder(asr, F0_pred, N_pred, s)
-        
-#         return out.squeeze().cpu().numpy(), duration.mean()
-    
-#     def get_styles(self, speakers, denoise=0.3, avg_style=True):
-#         if avg_style:   split_dur = 2
-#         else:           split_dur = 0
-#         styles = {}
-#         for id in speakers:
-#             ref_s = self.__compute_style(speakers[id]['path'], denoise=denoise, split_dur=split_dur)
-#             styles[id] = {
-#                 'style': ref_s,
-#                 'path': speakers[id]['path'],
-#                 'lang': speakers[id]['lang'],
-#                 'speed': speakers[id]['speed'],
-#             }
-#         return styles
-
-#     def generate(self, text, styles, stabilize=True, n_merge=16, default_speaker= "[id_1]"):
-#         if stabilize:   smooth_value=0.2
-#         else:           smooth_value=0    
-        
-#         list_wav        = []
-#         prev_d_mean     = 0
-#         lang_pattern    = r'\[([^\]]+)\]\{([^}]+)\}'
-
-#         text = re.sub(r'[\n\r\t\f\v]', '', text)
-#         #fix lang tokens span to multiple sents
-#         find_lang_tokens = re.findall(lang_pattern, text)
-#         if find_lang_tokens:
-#             cus_text = []
-#             for lang, t in find_lang_tokens:
-#                 parts = self.preprocess.text_preprocess(t, n_merge=0)
-#                 parts = ".".join([f"[{lang}]" + f"{{{p}}}"for p in parts])
-#                 cus_text.append(parts)
-#             replacement_func = self.__init_replacement_func(cus_text)
-#             text = re.sub(lang_pattern, replacement_func, text)
-
-#         texts = re.split(r'(\[id_\d+\])', text) #split the text by speaker ids while keeping the ids.
-#         if len(texts) <= 1 or bool(re.match(r'(\[id_\d+\])', texts[0]) == False): #Add a default speaker
-#             texts.insert(0, default_speaker)
-#         curr_id = None
-#         for i in range(len(texts)): #remove consecutive ids
-#             if bool(re.match(r'(\[id_\d+\])', texts[i])):
-#                 if texts[i]!=curr_id:
-#                     curr_id = texts[i]
-#                 else:
-#                     texts[i] = ''
-#         del curr_id
-#         texts = list(filter(lambda x: x != '', texts))
-
-#         print("Generating Audio...")
-#         for i in texts:
-#             if bool(re.match(r'(\[id_\d+\])', i)):
-#                 #Set up env for matched speaker
-#                 speaker_id = i.strip('[]')
-#                 current_ref_s = styles[speaker_id]['style']
-#                 speed = styles[speaker_id]['speed']
-#                 continue
-#             text_norm = self.preprocess.text_preprocess(i, n_merge=n_merge)
-#             for sentence in text_norm:
-#                 cus_phonem = []
-#                 find_lang_tokens = re.findall(lang_pattern, sentence)
-#                 if find_lang_tokens:
-#                     for lang, t in find_lang_tokens:
-#                         try:
-#                             phonem = espeak_phn(t, lang)
-#                             cus_phonem.append(phonem)
-#                         except Exception as e:
-#                             print(e)
-                        
-#                 replacement_func = self.__init_replacement_func(cus_phonem)
-#                 phonem =  espeak_phn(sentence, styles[speaker_id]['lang'])
-#                 phonem = re.sub(lang_pattern, replacement_func, phonem)
-
-#                 wav, prev_d_mean = self.__inference(phonem, current_ref_s, speed=speed, prev_d_mean=prev_d_mean, t=smooth_value)
-#                 wav = wav[4000:-4000] #Remove weird pulse and silent tokens
-#                 list_wav.append(wav)
-        
-#         final_wav = np.concatenate(list_wav)
-#         final_wav = np.concatenate([np.zeros([4000]), final_wav, np.zeros([4000])], axis=0) # add padding
-#         return final_wav
-
-
-
 import re
 import sys
 import yaml
@@ -485,9 +12,8 @@ from meldataset import TextCleaner
 from models import ProsodyPredictor, TextEncoder, StyleEncoder
 from Modules.hifigan import Decoder
 
-
 # -------------------------
-# Windows-only espeak-ng loader (không ảnh hưởng Linux/Space)
+# Windows-only espeak-ng loader
 # -------------------------
 if sys.platform.startswith("win"):
     try:
@@ -497,34 +23,29 @@ if sys.platform.startswith("win"):
     except Exception as e:
         print(e)
 
+_TOKEN_RE = re.compile(r"\S+")
+
+def normalize_phonem_tokens(phonem: str) -> str:
+    return " ".join(_TOKEN_RE.findall((phonem or "").strip()))
 
-def espeak_phn(text, lang):
+def espeak_phn(text: str, lang: str) -> str:
     """
-    Trả về phoneme string từ espeak backend.
-    Nếu backend fail, trả về text gốc (để không crash).
+    Nếu phonemizer/espeak lỗi -> raise để bạn biết ngay thiếu espeak-ng / libespeak-ng1 / voice 'vi'
     """
     try:
-        my_phonemizer = phonemizer.backend.EspeakBackend(
+        backend = phonemizer.backend.EspeakBackend(
             language=lang,
             preserve_punctuation=True,
             with_stress=True,
             language_switch="remove-flags",
         )
-        return my_phonemizer.phonemize([text])[0]
+        out = backend.phonemize([text])[0]
+        out = (out or "").strip()
+        if len(out) == 0:
+            raise RuntimeError(f"phonemizer returned empty output for lang='{lang}', text='{text[:50]}'")
+        return out
     except Exception as e:
-        print("[espeak_phn error]", e)
-        return text
-
-
-# -------------------------
-# Tokenization thay cho nltk.word_tokenize
-# Với phoneme/IPA, normalize whitespace là đủ.
-# -------------------------
-_TOKEN_RE = re.compile(r"\S+")
-
-def normalize_phonem_tokens(phonem: str) -> str:
-    return " ".join(_TOKEN_RE.findall((phonem or "").strip()))
-
+        raise RuntimeError(f"espeak/phonemizer failed (lang={lang}). Error: {e}")
 
 class Preprocess:
     def __text_normalize(self, text):
@@ -553,17 +74,7 @@ class Preprocess:
         return merged
 
     def wave_preprocess(self, wave, sr=24000):
-        """
-        Không dùng torchaudio.
-        Tạo log-mel bằng librosa, output shape (1, 80, T) giống code gốc.
-        """
-        if wave is None:
-            raise ValueError("wave is None")
-        wave = np.asarray(wave)
-        if wave.ndim != 1:
-            wave = wave.squeeze()
-        wave = wave.astype(np.float32)
-
+        wave = np.asarray(wave, dtype=np.float32).squeeze()
         mel = librosa.feature.melspectrogram(
             y=wave,
             sr=sr,
@@ -577,58 +88,41 @@ class Preprocess:
         mean, std = -4, 4
         mel = np.log(1e-5 + mel)
         mel = (mel - mean) / std
-
-        mel_tensor = torch.from_numpy(mel).float().unsqueeze(0)  # (1, 80, T)
-        return mel_tensor
+        return torch.from_numpy(mel).float().unsqueeze(0)  # (1, 80, T)
 
     def text_preprocess(self, text, n_merge=12):
         text_norm = self.__text_normalize(text).split(".")
-        text_norm = [s.strip() for s in text_norm]
-        text_norm = list(filter(lambda x: x != "", text_norm))
-        text_norm = self.__merge_fragments(text_norm, n=n_merge)
-        return text_norm
+        text_norm = [s.strip() for s in text_norm if s.strip()]
+        return self.__merge_fragments(text_norm, n=n_merge)
 
     def length_to_mask(self, lengths):
         mask = torch.arange(lengths.max()).unsqueeze(0).expand(lengths.shape[0], -1).type_as(lengths)
-        mask = torch.gt(mask + 1, lengths.unsqueeze(1))
-        return mask
-
+        return torch.gt(mask + 1, lengths.unsqueeze(1))
 
-# For inference only
 class StyleTTS2(torch.nn.Module):
     def __init__(self, config_path, models_path):
         super().__init__()
         self.register_buffer("get_device", torch.empty(0))
         self.preprocess = Preprocess()
-        self.ref_s = None
 
         config = yaml.safe_load(open(config_path, "r", encoding="utf-8"))
 
-        try:
-            symbols = (
-                list(config["symbol"]["pad"])
-                + list(config["symbol"]["punctuation"])
-                + list(config["symbol"]["letters"])
-                + list(config["symbol"]["letters_ipa"])
-                + list(config["symbol"]["extend"])
-            )
-            symbol_dict = {symbols[i]: i for i in range(len(symbols))}
-            n_token = len(symbol_dict) + 1
-            print("\nFound:", n_token, "symbols")
-        except Exception as e:
-            print(
-                f"\nERROR: Cannot find {e} in config file!\n"
-                "Your config file is likely outdated, please download updated version from the repository."
-            )
-            raise SystemExit(1)
+        symbols = (
+            list(config["symbol"]["pad"])
+            + list(config["symbol"]["punctuation"])
+            + list(config["symbol"]["letters"])
+            + list(config["symbol"]["letters_ipa"])
+            + list(config["symbol"]["extend"])
+        )
+        symbol_dict = {s: i for i, s in enumerate(symbols)}
+        n_token = len(symbol_dict) + 1
+        print("\nFound:", n_token, "symbols")
 
         args = self.__recursive_munch(config["model_params"])
         args["n_token"] = n_token
 
         self.cleaner = TextCleaner(symbol_dict, debug=False)
 
-        assert args.decoder.type in ["hifigan"], "Decoder type unknown"
-
         self.decoder = Decoder(
             dim_in=args.hidden_dim,
             style_dim=args.style_dim,
@@ -663,21 +157,14 @@ class StyleTTS2(torch.nn.Module):
     def __recursive_munch(self, d):
         if isinstance(d, dict):
             return Munch((k, self.__recursive_munch(v)) for k, v in d.items())
-        elif isinstance(d, list):
+        if isinstance(d, list):
             return [self.__recursive_munch(v) for v in d]
-        else:
-            return d
-
-    def __init_replacement_func(self, replacements):
-        replacement_iter = iter(replacements)
-        def replacement(match):
-            return next(replacement_iter)
-        return replacement
+        return d
 
     def __replace_outliers_zscore(self, tensor, threshold=3.0, factor=0.95):
         mean = tensor.mean()
         std = tensor.std()
-        z = (tensor - mean) / std
+        z = (tensor - mean) / (std + 1e-8)
         outlier_mask = torch.abs(z) > threshold
         sign = torch.sign(tensor - mean)
         replacement = mean + sign * (threshold * std * factor)
@@ -686,7 +173,6 @@ class StyleTTS2(torch.nn.Module):
         return result
 
     def __load_models(self, models_path):
-        module_params = []
         model = {
             "decoder": self.decoder,
             "predictor": self.predictor,
@@ -696,45 +182,28 @@ class StyleTTS2(torch.nn.Module):
 
         params_whole = torch.load(models_path, map_location="cpu")
         params = params_whole["net"]
-        params = {key: value for key, value in params.items() if key in model.keys()}
+        params = {k: v for k, v in params.items() if k in model}
 
-        for key in model:
+        for k in model:
             try:
-                model[key].load_state_dict(params[key])
+                model[k].load_state_dict(params[k])
             except Exception:
                 from collections import OrderedDict
-                state_dict = params[key]
                 new_state_dict = OrderedDict()
-                for k, v in state_dict.items():
-                    name = k[7:]  # remove `module.`
-                    new_state_dict[name] = v
-                model[key].load_state_dict(new_state_dict, strict=False)
-
-            total_params = sum(p.numel() for p in model[key].parameters())
-            print(key, ":", total_params)
-            module_params.append(total_params)
+                for kk, vv in params[k].items():
+                    new_state_dict[kk[7:]] = vv  # strip "module."
+                model[k].load_state_dict(new_state_dict, strict=False)
 
-        print("\nTotal", ":", sum(module_params))
+            print(k, ":", sum(p.numel() for p in model[k].parameters()))
 
     def __compute_style(self, path, denoise, split_dur):
         device = self.get_device.device
         denoise = min(float(denoise), 1.0)
-        if split_dur != 0:
-            split_dur = max(int(split_dur), 1)
-
-        max_samples = 24000 * 20
-        print("Computing the style for:", path)
+        split_dur = int(split_dur) if split_dur else 0
 
         wave, sr = librosa.load(path, sr=24000)
         audio, _ = librosa.effects.trim(wave, top_db=30)
 
-        if sr != 24000:
-            audio = librosa.resample(audio, orig_sr=sr, target_sr=24000)
-            sr = 24000
-
-        if len(audio) > max_samples:
-            audio = audio[:max_samples]
-
         if denoise > 0.0:
             audio_denoise = nr.reduce_noise(
                 y=audio, sr=sr, n_fft=2048, win_length=1200, hop_length=300
@@ -743,49 +212,39 @@ class StyleTTS2(torch.nn.Module):
 
         with torch.no_grad():
             if split_dur > 0 and len(audio) / sr >= 4:
-                count = 0
                 jump = sr * split_dur
                 total_len = len(audio)
+                ref_s = None
+                count = 0
 
-                mel_tensor = self.preprocess.wave_preprocess(audio[0:jump]).to(device)
-                ref_s = self.style_encoder(mel_tensor.unsqueeze(1))
-                count += 1
-
-                for i in range(jump, total_len, jump):
-                    if i + jump >= total_len:
-                        left_dur = (total_len - i) / sr
-                        if left_dur >= 1:
-                            mel_tensor = self.preprocess.wave_preprocess(audio[i:total_len]).to(device)
-                            ref_s += self.style_encoder(mel_tensor.unsqueeze(1))
-                            count += 1
+                for i in range(0, total_len, jump):
+                    seg = audio[i : min(i + jump, total_len)]
+                    if len(seg) < sr:  # <1s thì bỏ
                         continue
-
-                    mel_tensor = self.preprocess.wave_preprocess(audio[i : i + jump]).to(device)
-                    ref_s += self.style_encoder(mel_tensor.unsqueeze(1))
+                    mel = self.preprocess.wave_preprocess(seg).to(device)
+                    s = self.style_encoder(mel.unsqueeze(1))
+                    ref_s = s if ref_s is None else (ref_s + s)
                     count += 1
 
-                ref_s /= count
+                if ref_s is None:
+                    mel = self.preprocess.wave_preprocess(audio).to(device)
+                    ref_s = self.style_encoder(mel.unsqueeze(1))
+                else:
+                    ref_s = ref_s / count
             else:
-                mel_tensor = self.preprocess.wave_preprocess(audio).to(device)
-                ref_s = self.style_encoder(mel_tensor.unsqueeze(1))
+                mel = self.preprocess.wave_preprocess(audio).to(device)
+                ref_s = self.style_encoder(mel.unsqueeze(1))
 
         return ref_s
 
-    def __inference(self, phonem, ref_s, speed=1, prev_d_mean=0, t=0.1):
+    def __inference(self, phonem, ref_s, speed=1.0, prev_d_mean=0.0, t=0.1):
         device = self.get_device.device
-        speed = min(max(float(speed), 0.0001), 2.0)
+        speed = float(np.clip(speed, 1e-4, 2.0))
 
         phonem = normalize_phonem_tokens(phonem)
-
         tokens = self.cleaner(phonem)
-        tokens.insert(0, 0)
-        tokens.append(0)
-
-        # Guard: nếu cleaner trả rỗng thì fail sớm thay vì tạo audio 0s
-        if len(tokens) <= 2:
-            return np.zeros((0,), dtype=np.float32), 0.0
-
-        tokens = torch.LongTensor(tokens).to(device).unsqueeze(0)
+        tokens = [0] + tokens + [0]
+        tokens = torch.LongTensor(tokens).unsqueeze(0).to(device)
 
         with torch.no_grad():
             input_lengths = torch.LongTensor([tokens.shape[-1]]).to(device)
@@ -797,41 +256,34 @@ class StyleTTS2(torch.nn.Module):
             d = self.predictor.text_encoder(t_en, s, input_lengths, text_mask)
             x, _ = self.predictor.lstm(d)
             duration = self.predictor.duration_proj(x)
-            duration = torch.sigmoid(duration).sum(axis=-1)
+            duration = torch.sigmoid(duration).sum(dim=-1)
 
             if prev_d_mean != 0:
-                dur_stats = torch.empty(duration.shape).normal_(
-                    mean=prev_d_mean, std=duration.std()
-                ).to(device)
+                dur_stats = torch.empty_like(duration).normal_(mean=prev_d_mean, std=duration.std() + 1e-8).to(device)
             else:
-                dur_stats = torch.empty(duration.shape).normal_(
-                    mean=duration.mean(), std=duration.std()
-                ).to(device)
+                dur_stats = torch.empty_like(duration).normal_(mean=duration.mean(), std=duration.std() + 1e-8).to(device)
 
             duration = duration * (1 - t) + dur_stats * t
-            if duration.shape[1] > 3:
-                duration[:, 1:-2] = self.__replace_outliers_zscore(duration[:, 1:-2])
-            duration /= speed
+            duration[:, 1:-2] = self.__replace_outliers_zscore(duration[:, 1:-2])
+            duration = duration / speed
 
-            pred_dur = torch.round(duration.squeeze()).clamp(min=1)
+            pred_dur = torch.round(duration.squeeze(0)).clamp(min=1)
 
             L = int(input_lengths.item())
             T = int(pred_dur.sum().item())
-            if T <= 0:
-                return np.zeros((0,), dtype=np.float32), float(duration.mean().item())
+            pred_aln_trg = torch.zeros((L, T), device=device)
 
-            pred_aln_trg = torch.zeros((L, T))
-            c_frame = 0
+            c = 0
             for i in range(L):
                 di = int(pred_dur[i].item())
-                pred_aln_trg[i, c_frame : c_frame + di] = 1
-                c_frame += di
+                pred_aln_trg[i, c : c + di] = 1
+                c += di
 
-            alignment = pred_aln_trg.unsqueeze(0).to(device)
+            alignment = pred_aln_trg.unsqueeze(0)
 
-            en = (d.transpose(-1, -2) @ alignment)
+            en = d.transpose(-1, -2) @ alignment
             F0_pred, N_pred = self.predictor.F0Ntrain(en, s)
-            asr = (t_en @ pred_aln_trg.unsqueeze(0).to(device))
+            asr = t_en @ pred_aln_trg.unsqueeze(0)
 
             out = self.decoder(asr, F0_pred, N_pred, s)
 
@@ -840,17 +292,13 @@ class StyleTTS2(torch.nn.Module):
     def get_styles(self, speakers, denoise=0.3, avg_style=True):
         split_dur = 2 if avg_style else 0
         styles = {}
-        for sid in speakers:
-            ref_s = self.__compute_style(
-                speakers[sid]["path"],
-                denoise=denoise,
-                split_dur=split_dur,
-            )
+        for sid, meta in speakers.items():
+            ref_s = self.__compute_style(meta["path"], denoise=denoise, split_dur=split_dur)
             styles[sid] = {
                 "style": ref_s,
-                "path": speakers[sid]["path"],
-                "lang": speakers[sid]["lang"],
-                "speed": speakers[sid]["speed"],
+                "path": meta["path"],
+                "lang": meta["lang"],
+                "speed": meta["speed"],
             }
         return styles
 
@@ -860,89 +308,45 @@ class StyleTTS2(torch.nn.Module):
         list_wav = []
         prev_d_mean = 0.0
         lang_pattern = r"\[([^\]]+)\]\{([^}]+)\}"
-
         text = re.sub(r"[\n\r\t\f\v]", "", text)
 
-        # fix lang tokens span to multiple sents
-        find_lang_tokens = re.findall(lang_pattern, text)
-        if find_lang_tokens:
-            cus_text = []
-            for lang, t in find_lang_tokens:
-                parts = self.preprocess.text_preprocess(t, n_merge=0)
-                parts = ".".join([f"[{lang}]{{{p}}}" for p in parts])
-                cus_text.append(parts)
-            replacement_func = self.__init_replacement_func(cus_text)
-            text = re.sub(lang_pattern, replacement_func, text)
-
-        texts = re.split(r"(\[id_\d+\])", text)
-        if len(texts) <= 1 or (re.match(r"(\[id_\d+\])", texts[0]) is None):
-            texts.insert(0, default_speaker)
-
-        # remove consecutive ids
-        curr_id = None
-        for i in range(len(texts)):
-            if re.match(r"(\[id_\d+\])", texts[i]):
-                if texts[i] != curr_id:
-                    curr_id = texts[i]
-                else:
-                    texts[i] = ""
-        texts = list(filter(lambda x: x != "", texts))
-
-        print("Generating Audio...")
+        # split by speaker tags
+        parts = re.split(r"(\[id_\d+\])", text)
+        if len(parts) <= 1 or re.match(r"(\[id_\d+\])", parts[0]) is None:
+            parts.insert(0, default_speaker)
 
         speaker_id = None
         current_ref_s = None
         speed = 1.0
 
-        for seg in texts:
-            if re.match(r"(\[id_\d+\])", seg):
-                speaker_id = seg.strip("[]")  # "id_1"
-                if speaker_id not in styles:
-                    raise KeyError(f"speaker_id '{speaker_id}' not found in styles keys={list(styles.keys())[:5]}...")
+        for p in parts:
+            if re.match(r"(\[id_\d+\])", p):
+                speaker_id = p.strip("[]")
                 current_ref_s = styles[speaker_id]["style"]
                 speed = styles[speaker_id]["speed"]
                 continue
 
-            if speaker_id is None or current_ref_s is None:
-                # input text không có speaker tag hợp lệ
-                speaker_id = default_speaker.strip("[]")
-                current_ref_s = styles[speaker_id]["style"]
-                speed = styles[speaker_id]["speed"]
-
-            text_norm = self.preprocess.text_preprocess(seg, n_merge=n_merge)
-            for sentence in text_norm:
-                cus_phonem = []
-                find_lang_tokens = re.findall(lang_pattern, sentence)
-                if find_lang_tokens:
-                    for lang, t in find_lang_tokens:
-                        cus_phonem.append(espeak_phn(t, lang))
+            if not p.strip():
+                continue
 
-                replacement_func = self.__init_replacement_func(cus_phonem)
+            for sentence in self.preprocess.text_preprocess(p, n_merge=n_merge):
+                # phonemize
                 phonem = espeak_phn(sentence, styles[speaker_id]["lang"])
-                phonem = re.sub(lang_pattern, replacement_func, phonem)
-
                 wav, prev_d_mean = self.__inference(
-                    phonem,
-                    current_ref_s,
-                    speed=speed,
-                    prev_d_mean=prev_d_mean,
-                    t=smooth_value,
+                    phonem, current_ref_s, speed=speed, prev_d_mean=prev_d_mean, t=smooth_value
                 )
 
-                if wav is None or wav.shape[0] == 0:
-                    continue
-
+                # trim an toàn
                 trim = 4000
                 if wav.shape[0] > 2 * trim:
                     wav = wav[trim:-trim]
 
-                # chỉ append 1 lần
-                list_wav.append(wav)
+                if wav.size > 0:
+                    list_wav.append(wav)
 
         if len(list_wav) == 0:
-            # trả một đoạn silence ngắn để tránh 0s file
-            return np.zeros((2400,), dtype=np.float32)
+            return np.zeros((2400,), dtype=np.float32)  # 0.1s silence để không crash
 
-        final_wav = np.concatenate(list_wav, axis=0)
-        final_wav = np.concatenate([np.zeros([4000]), final_wav, np.zeros([4000])], axis=0)
+        final_wav = np.concatenate(list_wav)
+        final_wav = np.concatenate([np.zeros((4000,), dtype=np.float32), final_wav, np.zeros((4000,), dtype=np.float32)])
         return final_wav