diff --git a/.gitignore b/.gitignore index 42f0240..9a078a3 100644 --- a/.gitignore +++ b/.gitignore @@ -9,3 +9,4 @@ checkpoints/*.vocab checkpoints/*.model checkpoints/.cache outputs/ +build/ diff --git a/MANIFEST.in b/MANIFEST.in new file mode 100644 index 0000000..ad8e99e --- /dev/null +++ b/MANIFEST.in @@ -0,0 +1,3 @@ +global-exclude *~ *.py[cod] +include *.cu *.cpp +include *.h *.hpp diff --git a/indextts/infer.py b/indextts/infer.py index b2233b4..7bacde8 100644 --- a/indextts/infer.py +++ b/indextts/infer.py @@ -2,6 +2,7 @@ import os import re import time from subprocess import CalledProcessError +from typing import List import numpy as np import sentencepiece as spm @@ -12,6 +13,7 @@ from omegaconf import OmegaConf from tqdm import tqdm import warnings + warnings.filterwarnings("ignore", category=FutureWarning) warnings.filterwarnings("ignore", category=UserWarning) @@ -19,9 +21,9 @@ from indextts.BigVGAN.models import BigVGAN as Generator from indextts.gpt.model import UnifiedVoice from indextts.utils.checkpoint import load_checkpoint from indextts.utils.feature_extractors import MelSpectrogramFeatures -from indextts.utils.common import tokenize_by_CJK_char -from indextts.utils.front import TextNormalizer +from indextts.utils.front import TextNormalizer, TextTokenizer + class IndexTTS: def __init__( @@ -43,9 +45,9 @@ class IndexTTS: self.device = "cuda:0" self.is_fp16 = is_fp16 self.use_cuda_kernel = use_cuda_kernel is None or use_cuda_kernel - elif torch.mps.is_available(): + elif hasattr(torch, "mps") and torch.backends.mps.is_available(): self.device = "mps" - self.is_fp16 = is_fp16 + self.is_fp16 = False # Use float16 on MPS is overhead than float32 self.use_cuda_kernel = False else: self.device = "cpu" @@ -57,7 +59,7 @@ class IndexTTS: self.model_dir = model_dir self.dtype = torch.float16 if self.is_fp16 else None self.stop_mel_token = self.cfg.gpt.stop_mel_token - + # Comment-off to load the VQ-VAE model for debugging tokenizer # https://github.com/index-tts/index-tts/issues/34 # @@ -83,19 +85,21 @@ class IndexTTS: if self.is_fp16: try: import deepspeed + use_deepspeed = True - except (ImportError, OSError,CalledProcessError) as e: + except (ImportError, OSError, CalledProcessError) as e: use_deepspeed = False print(f">> DeepSpeed加载失败，回退到标准推理: {e}") self.gpt.post_init_gpt2_config(use_deepspeed=use_deepspeed, kv_cache=True, half=True) else: self.gpt.post_init_gpt2_config(use_deepspeed=False, kv_cache=False, half=False) - + if self.use_cuda_kernel: # preload the CUDA kernel for BigVGAN try: from indextts.BigVGAN.alias_free_activation.cuda import load + anti_alias_activation_cuda = load.load() print(">> Preload custom CUDA kernel for BigVGAN", anti_alias_activation_cuda) except: @@ -110,29 +114,18 @@ class IndexTTS: self.bigvgan.remove_weight_norm() self.bigvgan.eval() print(">> bigvgan weights restored from:", self.bigvgan_path) - self.bpe_path = os.path.join(self.model_dir, self.cfg.dataset['bpe_model']) - self.tokenizer = spm.SentencePieceProcessor(model_file=self.bpe_path) - print(">> bpe model loaded from:", self.bpe_path) + self.bpe_path = os.path.join(self.model_dir, self.cfg.dataset["bpe_model"]) self.normalizer = TextNormalizer() self.normalizer.load() print(">> TextNormalizer loaded") + self.tokenizer = TextTokenizer(self.bpe_path, self.normalizer) + print(">> bpe model loaded from:", self.bpe_path) # 缓存参考音频mel： self.cache_audio_prompt = None self.cache_cond_mel = None # 进度引用显示（可选） self.gr_progress = None - def preprocess_text(self, text): - # chinese_punctuation = "，。！？；：“”‘’（）【】《》" - # english_punctuation = ",.!?;:\"\"''()[]<>" - # - # # 创建一个映射字典 - # punctuation_map = str.maketrans(chinese_punctuation, english_punctuation) - - # 使用translate方法替换标点符号 - # return text.translate(punctuation_map) - return self.normalizer.infer(text) - def remove_long_silence(self, codes: torch.Tensor, silent_token=52, max_consecutive=30): code_lens = [] codes_list = [] @@ -167,8 +160,8 @@ class IndexTTS: ncode = torch.LongTensor(ncode) codes_list.append(ncode.to(device, dtype=dtype)) isfix = True - #codes[i] = self.stop_mel_token - #codes[i, 0:len_] = ncode + # codes[i] = self.stop_mel_token + # codes[i, 0:len_] = ncode else: codes_list.append(codes[i]) code_lens.append(len_) @@ -177,53 +170,36 @@ class IndexTTS: code_lens = torch.LongTensor(code_lens).to(device, dtype=dtype) return codes, code_lens - def split_sentences(self, text): - """ - Split the text into sentences based on punctuation marks. - """ - # 匹配标点符号（包括中英文标点） - pattern = r'(?<=[.!?;。！？；])\s*' - sentences = re.split(pattern, text) - # 过滤掉空字符串和仅包含标点符号的字符串 - return [ - sentence.strip() for sentence in sentences if sentence.strip() and sentence.strip() not in {"'", ".", ","} - ] - - def bucket_sentences(self, sentences, enable): + def bucket_sentences(self, sentences, enable=False): """ Sentence data bucketing """ max_len = max(len(s) for s in sentences) half = max_len // 2 - outputs = [[],[]] + outputs = [[], []] for idx, sent in enumerate(sentences): - if enable == False or len(sent) <= half: - outputs[0].append({"idx":idx,"sent":sent}) + if enable is False or len(sent) <= half: + outputs[0].append({"idx": idx, "sent": sent}) else: - outputs[1].append({"idx":idx,"sent":sent}) + outputs[1].append({"idx": idx, "sent": sent}) return [item for item in outputs if item] - - def pad_tokens_cat(self, tokens): - if len(tokens) <= 1:return tokens[-1] + + def pad_tokens_cat(self, tokens: List[torch.Tensor]): + if len(tokens) <= 1: + return tokens[-1] max_len = max(t.size(1) for t in tokens) outputs = [] for tensor in tokens: pad_len = max_len - tensor.size(1) if pad_len > 0: n = min(8, pad_len) - tensor = torch.nn.functional.pad(tensor, - (0, n), - value=self.cfg.gpt.stop_text_token - ) - tensor = torch.nn.functional.pad(tensor, - (0, pad_len - n), - value=self.cfg.gpt.start_text_token - ) - tensor = tensor[:,:max_len] + tensor = torch.nn.functional.pad(tensor, (0, n), value=self.cfg.gpt.stop_text_token) + tensor = torch.nn.functional.pad(tensor, (0, pad_len - n), value=self.cfg.gpt.start_text_token) + tensor = tensor[:, :max_len] outputs.append(tensor) tokens = torch.cat(outputs, dim=0) return tokens - + def torch_empty_cache(self): try: if "cuda" in str(self.device): @@ -231,13 +207,12 @@ class IndexTTS: elif "mps" in str(self.device): torch.mps.empty_cache() except Exception as e: - pass - + pass + def _set_gr_progress(self, value, desc): - if self.gr_progress is not None:self.gr_progress(value, desc=desc) - - - + if self.gr_progress is not None: + self.gr_progress(value, desc=desc) + # 快速推理：对于“多句长文本”，可实现至少 2~10 倍以上的速度提升~ （First modified by sunnyboxs 2025-04-16） def infer_fast(self, audio_prompt, text, output_path, verbose=False): print(">> start fast inference...") @@ -245,9 +220,6 @@ class IndexTTS: if verbose: print(f"origin text:{text}") start_time = time.perf_counter() - normalized_text = self.preprocess_text(text) - print(f"normalized text:{normalized_text}") - # 如果参考音频改变了，才需要重新生成 cond_mel, 提升速度 if self.cache_cond_mel is None or self.cache_audio_prompt != audio_prompt: @@ -260,23 +232,26 @@ class IndexTTS: cond_mel_frame = cond_mel.shape[-1] if verbose: print(f"cond_mel shape: {cond_mel.shape}", "dtype:", cond_mel.dtype) - + self.cache_audio_prompt = audio_prompt self.cache_cond_mel = cond_mel else: cond_mel = self.cache_cond_mel cond_mel_frame = cond_mel.shape[-1] pass - + auto_conditioning = cond_mel - cond_mel_lengths = torch.tensor([cond_mel_frame],device=self.device) - + cond_mel_lengths = torch.tensor([cond_mel_frame], device=self.device) + # text_tokens - sentences = self.split_sentences(normalized_text) + text_tokens_list = self.tokenizer.tokenize(text) + sentences = self.tokenizer.split_sentences(text_tokens_list) if verbose: - print("sentences:", sentences) - - top_p = .8 + print("text token count:", len(text_tokens_list)) + print("sentences count:", len(sentences)) + print(*sentences, sep="\n") + + top_p = 0.8 top_k = 30 temperature = 1.0 autoregressive_batch_size = 1 @@ -293,32 +268,23 @@ class IndexTTS: bigvgan_time = 0 # text processing - all_text_tokens = [] + all_text_tokens: List[List[torch.Tensor]] = [] self._set_gr_progress(0.1, "text processing...") bucket_enable = True # 预分桶开关，优先保证质量=True。优先保证速度=False。 all_sentences = self.bucket_sentences(sentences, enable=bucket_enable) for sentences in all_sentences: - temp_tokens = [] + temp_tokens: List[torch.Tensor] = [] all_text_tokens.append(temp_tokens) for item in sentences: sent = item["sent"] - # sent = " ".join([char for char in sent.upper()]) if lang == "ZH" else sent.upper() - cleand_text = tokenize_by_CJK_char(sent) - # cleand_text = "他 那 像 HONG3 小 孩 似 的 话 , 引 得 人 们 HONG1 堂 大 笑 , 大 家 听 了 一 HONG3 而 散 ." - if verbose: - print("cleand_text:", cleand_text) - - text_tokens = torch.tensor(self.tokenizer.EncodeAsIds(cleand_text),dtype=torch.int32, device=self.device).unsqueeze(0) - # text_tokens = F.pad(text_tokens, (0, 1)) # This may not be necessary. - # text_tokens = F.pad(text_tokens, (1, 0), value=0) - # text_tokens = F.pad(text_tokens, (0, 1), value=1) + text_tokens = self.tokenizer.convert_tokens_to_ids(sent) + text_tokens = torch.tensor(text_tokens, dtype=torch.int32, device=self.device).unsqueeze(0) if verbose: print(text_tokens) print(f"text_tokens shape: {text_tokens.shape}, text_tokens type: {text_tokens.dtype}") # debug tokenizer - text_token_syms = self.tokenizer.IdToPiece(text_tokens[0].tolist()) - print(text_token_syms) - + text_token_syms = self.tokenizer.convert_ids_to_tokens(text_tokens[0].tolist()) + print("text_token_syms is same as sentence tokens", text_token_syms == sent) temp_tokens.append(text_tokens) @@ -331,12 +297,12 @@ class IndexTTS: batch_cond_mel_lengths = torch.cat([cond_mel_lengths] * batch_num, dim=0) batch_auto_conditioning = torch.cat([auto_conditioning] * batch_num, dim=0) all_batch_num += batch_num - + # gpt speech self._set_gr_progress(0.2, "gpt inference speech...") m_start_time = time.perf_counter() with torch.no_grad(): - with torch.amp.autocast(self.device, enabled=self.dtype is not None, dtype=self.dtype): + with torch.amp.autocast(batch_text_tokens.device.type, enabled=self.dtype is not None, dtype=self.dtype): temp_codes = self.gpt.inference_speech(batch_auto_conditioning, batch_text_tokens, cond_mel_lengths=batch_cond_mel_lengths, # text_lengths=text_len, @@ -351,25 +317,24 @@ class IndexTTS: max_generate_length=max_mel_tokens) all_batch_codes.append(temp_codes) gpt_gen_time += time.perf_counter() - m_start_time - - + # gpt latent self._set_gr_progress(0.5, "gpt inference latents...") all_idxs = [] all_latents = [] for batch_codes, batch_tokens, batch_sentences in zip(all_batch_codes, all_text_tokens, all_sentences): for i in range(batch_codes.shape[0]): - codes = batch_codes[i] # [x] + codes = batch_codes[i] # [x] codes = codes[codes != self.cfg.gpt.stop_mel_token] codes, _ = torch.unique_consecutive(codes, return_inverse=True) - codes = codes.unsqueeze(0) # [x] -> [1, x] + codes = codes.unsqueeze(0) # [x] -> [1, x] code_lens = torch.tensor([codes.shape[-1]], device=codes.device, dtype=codes.dtype) codes, code_lens = self.remove_long_silence(codes, silent_token=52, max_consecutive=30) text_tokens = batch_tokens[i] all_idxs.append(batch_sentences[i]["idx"]) m_start_time = time.perf_counter() with torch.no_grad(): - with torch.amp.autocast(self.device, enabled=self.dtype is not None, dtype=self.dtype): + with torch.amp.autocast(text_tokens.device.type, enabled=self.dtype is not None, dtype=self.dtype): latent = \ self.gpt(auto_conditioning, text_tokens, torch.tensor([text_tokens.shape[-1]], device=text_tokens.device), codes, @@ -378,16 +343,15 @@ class IndexTTS: return_latent=True, clip_inputs=False) gpt_forward_time += time.perf_counter() - m_start_time all_latents.append(latent) - - + # bigvgan chunk - chunk_size = 2 + chunk_size = 2 all_latents = [all_latents[all_idxs.index(i)] for i in range(len(all_latents))] - chunk_latents = [all_latents[i:i + chunk_size] for i in range(0, len(all_latents), chunk_size)] + chunk_latents = [all_latents[i : i + chunk_size] for i in range(0, len(all_latents), chunk_size)] chunk_length = len(chunk_latents) latent_length = len(all_latents) all_latents = None - + # bigvgan chunk decode self._set_gr_progress(0.7, "bigvgan decode...") tqdm_progress = tqdm(total=latent_length, desc="bigvgan") @@ -395,26 +359,26 @@ class IndexTTS: tqdm_progress.update(len(items)) latent = torch.cat(items, dim=1) with torch.no_grad(): - with torch.amp.autocast(self.device, enabled=self.dtype is not None, dtype=self.dtype): + with torch.amp.autocast(latent.device.type, enabled=self.dtype is not None, dtype=self.dtype): m_start_time = time.perf_counter() wav, _ = self.bigvgan(latent, auto_conditioning.transpose(1, 2)) bigvgan_time += time.perf_counter() - m_start_time wav = wav.squeeze(1) pass wav = torch.clamp(32767 * wav, -32767.0, 32767.0) - wavs.append(wav) - - # clear cache + wavs.append(wav.cpu()) # to cpu before saving + + # clear cache tqdm_progress.close() # 确保进度条被关闭 chunk_latents.clear() end_time = time.perf_counter() self.torch_empty_cache() - + # wav audio output self._set_gr_progress(0.9, "save audio...") wav = torch.cat(wavs, dim=1) wav_length = wav.shape[-1] / sampling_rate - print(f">> Reference audio length: {cond_mel_frame*256 / sampling_rate:.2f} seconds") + print(f">> Reference audio length: {cond_mel_frame * 256 / sampling_rate:.2f} seconds") print(f">> gpt_gen_time: {gpt_gen_time:.2f} seconds") print(f">> gpt_forward_time: {gpt_forward_time:.2f} seconds") print(f">> bigvgan_time: {bigvgan_time:.2f} seconds") @@ -425,21 +389,19 @@ class IndexTTS: print(f">> [fast] RTF: {(end_time - start_time) / wav_length:.4f}") # save audio - wav = wav.cpu() # to cpu + wav = wav.cpu() # to cpu if output_path: # 直接保存音频到指定路径中 - os.makedirs(os.path.dirname(output_path),exist_ok=True) + os.makedirs(os.path.dirname(output_path), exist_ok=True) torchaudio.save(output_path, wav.type(torch.int16), sampling_rate) print(">> wav file saved to:", output_path) return output_path else: # 返回以符合Gradio的格式要求 wav_data = wav.type(torch.int16) - wav_data = wav_data.numpy().T + wav_data = wav_data.numpy().T return (sampling_rate, wav_data) - - - + # 原始推理模式 def infer(self, audio_prompt, text, output_path, verbose=False): print(">> start inference...") @@ -447,9 +409,6 @@ class IndexTTS: if verbose: print(f"origin text:{text}") start_time = time.perf_counter() - normalized_text = self.preprocess_text(text) - print(f"normalized text:{normalized_text}") - # 如果参考音频改变了，才需要重新生成 cond_mel, 提升速度 if self.cache_cond_mel is None or self.cache_audio_prompt != audio_prompt: @@ -462,22 +421,22 @@ class IndexTTS: cond_mel_frame = cond_mel.shape[-1] if verbose: print(f"cond_mel shape: {cond_mel.shape}", "dtype:", cond_mel.dtype) - + self.cache_audio_prompt = audio_prompt self.cache_cond_mel = cond_mel else: cond_mel = self.cache_cond_mel cond_mel_frame = cond_mel.shape[-1] pass - auto_conditioning = cond_mel - - sentences = self.split_sentences(normalized_text) + text_tokens_list = self.tokenizer.tokenize(text) + sentences = self.tokenizer.split_sentences(text_tokens_list) if verbose: - print("sentences:", sentences) - - top_p = .8 + print("text token count:", len(text_tokens_list)) + print("sentences count:", len(sentences)) + print(*sentences, sep="\n") + top_p = 0.8 top_k = 30 temperature = 1.0 autoregressive_batch_size = 1 @@ -494,13 +453,8 @@ class IndexTTS: bigvgan_time = 0 for sent in sentences: - # sent = " ".join([char for char in sent.upper()]) if lang == "ZH" else sent.upper() - cleand_text = tokenize_by_CJK_char(sent) - # cleand_text = "他 那 像 HONG3 小 孩 似 的 话 , 引 得 人 们 HONG1 堂 大 笑 , 大 家 听 了 一 HONG3 而 散 ." - if verbose: - print("cleand_text:", cleand_text) - - text_tokens = torch.tensor(self.tokenizer.EncodeAsIds(cleand_text),dtype=torch.int32, device=self.device).unsqueeze(0) + text_tokens = self.tokenizer.convert_tokens_to_ids(sent) + text_tokens = torch.tensor(text_tokens, dtype=torch.int32, device=self.device).unsqueeze(0) # text_tokens = F.pad(text_tokens, (0, 1)) # This may not be necessary. # text_tokens = F.pad(text_tokens, (1, 0), value=0) # text_tokens = F.pad(text_tokens, (0, 1), value=1) @@ -508,15 +462,15 @@ class IndexTTS: print(text_tokens) print(f"text_tokens shape: {text_tokens.shape}, text_tokens type: {text_tokens.dtype}") # debug tokenizer - text_token_syms = self.tokenizer.IdToPiece(text_tokens[0].tolist()) - print(text_token_syms) + text_token_syms = self.tokenizer.convert_ids_to_tokens(text_tokens[0].tolist()) + print("text_token_syms is same as sentence tokens", text_token_syms == sent) # text_len = torch.IntTensor([text_tokens.size(1)], device=text_tokens.device) # print(text_len) m_start_time = time.perf_counter() with torch.no_grad(): - with torch.amp.autocast(self.device, enabled=self.dtype is not None, dtype=self.dtype): + with torch.amp.autocast(text_tokens.device.type, enabled=self.dtype is not None, dtype=self.dtype): codes = self.gpt.inference_speech(auto_conditioning, text_tokens, cond_mel_lengths=torch.tensor([auto_conditioning.shape[-1]], device=text_tokens.device), @@ -531,7 +485,7 @@ class IndexTTS: repetition_penalty=repetition_penalty, max_generate_length=max_mel_tokens) gpt_gen_time += time.perf_counter() - m_start_time - #codes = codes[:, :-2] + # codes = codes[:, :-2] code_lens = torch.tensor([codes.shape[-1]], device=codes.device, dtype=codes.dtype) if verbose: print(codes, type(codes)) @@ -548,7 +502,7 @@ class IndexTTS: m_start_time = time.perf_counter() # latent, text_lens_out, code_lens_out = \ - with torch.amp.autocast(self.device, enabled=self.dtype is not None, dtype=self.dtype): + with torch.amp.autocast(text_tokens.device.type, enabled=self.dtype is not None, dtype=self.dtype): latent = \ self.gpt(auto_conditioning, text_tokens, torch.tensor([text_tokens.shape[-1]], device=text_tokens.device), codes, @@ -565,12 +519,12 @@ class IndexTTS: wav = torch.clamp(32767 * wav, -32767.0, 32767.0) print(f"wav shape: {wav.shape}", "min:", wav.min(), "max:", wav.max()) # wavs.append(wav[:, :-512]) - wavs.append(wav) + wavs.append(wav.cpu()) # to cpu before saving end_time = time.perf_counter() wav = torch.cat(wavs, dim=1) wav_length = wav.shape[-1] / sampling_rate - print(f">> Reference audio length: {cond_mel_frame*256 / sampling_rate:.2f} seconds") + print(f">> Reference audio length: {cond_mel_frame * 256 / sampling_rate:.2f} seconds") print(f">> gpt_gen_time: {gpt_gen_time:.2f} seconds") print(f">> gpt_forward_time: {gpt_forward_time:.2f} seconds") print(f">> bigvgan_time: {bigvgan_time:.2f} seconds") @@ -578,25 +532,22 @@ class IndexTTS: print(f">> Generated audio length: {wav_length:.2f} seconds") print(f">> RTF: {(end_time - start_time) / wav_length:.4f}") - # torchaudio.save(output_path, wav.cpu().type(torch.int16), sampling_rate) - # print(">> wav file saved to:", output_path) - # save audio - wav = wav.cpu() # to cpu + wav = wav.cpu() # to cpu if output_path: # 直接保存音频到指定路径中 if os.path.isfile(output_path): os.remove(output_path) print(">> remove old wav file:", output_path) if os.path.dirname(output_path) != "": - os.makedirs(os.path.dirname(output_path),exist_ok=True) + os.makedirs(os.path.dirname(output_path), exist_ok=True) torchaudio.save(output_path, wav.type(torch.int16), sampling_rate) print(">> wav file saved to:", output_path) return output_path else: # 返回以符合Gradio的格式要求 wav_data = wav.type(torch.int16) - wav_data = wav_data.numpy().T + wav_data = wav_data.numpy().T return (sampling_rate, wav_data) diff --git a/indextts/utils/common.py b/indextts/utils/common.py index b22d064..eb2f89a 100644 --- a/indextts/utils/common.py +++ b/indextts/utils/common.py @@ -10,7 +10,7 @@ MATPLOTLIB_FLAG = False def load_audio(audiopath, sampling_rate): audio, sr = torchaudio.load(audiopath) - #print(f"wave shape: {audio.shape}, sample_rate: {sr}") + # print(f"wave shape: {audio.shape}, sample_rate: {sr}") if audio.size(0) > 1: # mix to mono audio = audio[0].unsqueeze(0) @@ -26,13 +26,13 @@ def load_audio(audiopath, sampling_rate): return audio -def tokenize_by_CJK_char(line: str) -> str: - """ +def tokenize_by_CJK_char(line: str, do_upper_case=True) -> str: + """ Tokenize a line of text with CJK char. Note: All return charaters will be upper case. - Example: + Example: input = "你好世界是 hello world 的中文" output = "你 好 世 界 是 HELLO WORLD 的 中 文" @@ -44,11 +44,41 @@ def tokenize_by_CJK_char(line: str) -> str: A new string tokenize by CJK char. """ # The CJK ranges is from https://github.com/alvations/nltk/blob/79eed6ddea0d0a2c212c1060b477fc268fec4d4b/nltk/tokenize/util.py - pattern = re.compile( + CJK_RANGE_PATTERN = ( r"([\u1100-\u11ff\u2e80-\ua4cf\ua840-\uD7AF\uF900-\uFAFF\uFE30-\uFE4F\uFF65-\uFFDC\U00020000-\U0002FFFF])" - ) - chars = pattern.split(line.strip().upper()) - return " ".join([w.strip() for w in chars if w.strip()]) + ) + chars = re.split(CJK_RANGE_PATTERN, line.strip()) + return " ".join([w.strip().upper() if do_upper_case else w.strip() for w in chars if w.strip()]) + + +def de_tokenized_by_CJK_char(line: str, do_lower_case=False) -> str: + """ + Example: + input = "你 好 世 界 是 HELLO WORLD 的 中 文" + output = "你好世界是 hello world 的中文" + + do_lower_case: + input = "SEE YOU!" + output = "see you!" + """ + # replace english words in the line with placeholders + english_word_pattern = re.compile(r"([A-Z]+(?:[\s-][A-Z-]+)*)", re.IGNORECASE) + english_sents = english_word_pattern.findall(line) + for i, sent in enumerate(english_sents): + line = line.replace(sent, f"") + + words = line.split() + # restore english sentences + sent_placeholder_pattern = re.compile(r"^.*?()") + for i in range(len(words)): + m = sent_placeholder_pattern.match(words[i]) + if m: + # restore the english word + placeholder_index = int(m.group(2)) + words[i] = words[i].replace(m.group(1), english_sents[placeholder_index]) + if do_lower_case: + words[i] = words[i].lower() + return "".join(words) def make_pad_mask(lengths: torch.Tensor, max_len: int = 0) -> torch.Tensor: @@ -70,10 +100,7 @@ def make_pad_mask(lengths: torch.Tensor, max_len: int = 0) -> torch.Tensor: """ batch_size = lengths.size(0) max_len = max_len if max_len > 0 else lengths.max().item() - seq_range = torch.arange(0, - max_len, - dtype=torch.int64, - device=lengths.device) + seq_range = torch.arange(0, max_len, dtype=torch.int64, device=lengths.device) seq_range_expand = seq_range.unsqueeze(0).expand(batch_size, max_len) seq_length_expand = lengths.unsqueeze(-1) mask = seq_range_expand >= seq_length_expand diff --git a/indextts/utils/front.py b/indextts/utils/front.py index 93df6f0..d7ff8c7 100644 --- a/indextts/utils/front.py +++ b/indextts/utils/front.py @@ -1,6 +1,12 @@ # -*- coding: utf-8 -*- +import os import traceback import re +from typing import List, Union, overload +import warnings +from indextts.utils.common import tokenize_by_CJK_char, de_tokenized_by_CJK_char +from sentencepiece import SentencePieceProcessor + class TextNormalizer: def __init__(self): @@ -18,8 +24,9 @@ class TextNormalizer: "·": "-", "、": ",", "...": "…", + ",,,": "…", + "，，，": "…", "……": "…", - "$": ".", "“": "'", "”": "'", '"': "'", @@ -42,69 +49,129 @@ class TextNormalizer: "」": "'", ":": ",", } + self.zh_char_rep_map = { + "$": ".", + **self.char_rep_map, + } def match_email(self, email): # 正则表达式匹配邮箱格式：数字英文@数字英文.英文 - pattern = r'^[a-zA-Z0-9]+@[a-zA-Z0-9]+\.[a-zA-Z]+$' + pattern = r"^[a-zA-Z0-9]+@[a-zA-Z0-9]+\.[a-zA-Z]+$" return re.match(pattern, email) is not None + """ 匹配拼音声调格式：pinyin+数字，声调1-5，5表示轻声 例如：xuan4, jve2, ying1, zhong4, shang5 """ PINYIN_TONE_PATTERN = r"([bmnpqdfghjklzcsxwy]?h?[aeiouüv]{1,2}[ng]*|ng)([1-5])" + """ + 匹配人名，格式：中文·中文，中文·中文-中文 + 例如：克里斯托弗·诺兰，约瑟夫·高登-莱维特 + """ + NAME_PATTERN = r"[\u4e00-\u9fff]+([-·—][\u4e00-\u9fff]+){1,2}" + def use_chinese(self, s): - has_chinese = bool(re.search(r'[\u4e00-\u9fff]', s)) - has_alpha = bool(re.search(r'[a-zA-Z]', s)) + has_chinese = bool(re.search(r"[\u4e00-\u9fff]", s)) + has_alpha = bool(re.search(r"[a-zA-Z]", s)) is_email = self.match_email(s) if has_chinese or not has_alpha or is_email: return True - has_pinyin = bool(re.search(self.PINYIN_TONE_PATTERN, s, re.IGNORECASE)) + has_pinyin = bool(re.search(TextNormalizer.PINYIN_TONE_PATTERN, s, re.IGNORECASE)) return has_pinyin def load(self): # print(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..")) # sys.path.append(model_dir) import platform + if platform.system() == "Darwin": from wetext import Normalizer - self.zh_normalizer = Normalizer(remove_erhua=False,lang="zh",operator="tn") - self.en_normalizer = Normalizer(lang="en",operator="tn") + + self.zh_normalizer = Normalizer(remove_erhua=False, lang="zh", operator="tn") + self.en_normalizer = Normalizer(lang="en", operator="tn") else: from tn.chinese.normalizer import Normalizer as NormalizerZh from tn.english.normalizer import Normalizer as NormalizerEn - self.zh_normalizer = NormalizerZh(remove_interjections=False, remove_erhua=False,overwrite_cache=False) + + self.zh_normalizer = NormalizerZh(remove_interjections=False, remove_erhua=False, overwrite_cache=False) self.en_normalizer = NormalizerEn(overwrite_cache=False) - def infer(self, text: str): + def normalize(self, text: str) -> str: if not self.zh_normalizer or not self.en_normalizer: print("Error, text normalizer is not initialized !!!") return "" - replaced_text, pinyin_list = self.save_pinyin_tones(text.rstrip()) - - try: - normalizer = self.zh_normalizer if self.use_chinese(replaced_text) else self.en_normalizer - result = normalizer.normalize(replaced_text) - except Exception: - result = "" - print(traceback.format_exc()) - result = self.restore_pinyin_tones(result, pinyin_list) - pattern = re.compile("|".join(re.escape(p) for p in self.char_rep_map.keys())) - result = pattern.sub(lambda x: self.char_rep_map[x.group()], result) + if self.use_chinese(text): + replaced_text, pinyin_list = self.save_pinyin_tones(text.rstrip()) + + replaced_text, original_name_list = self.save_names(replaced_text) + try: + result = self.zh_normalizer.normalize(replaced_text) + except Exception: + result = "" + print(traceback.format_exc()) + # 恢复人名 + result = self.restore_names(result, original_name_list) + # 恢复拼音声调 + result = self.restore_pinyin_tones(result, pinyin_list) + pattern = re.compile("|".join(re.escape(p) for p in self.zh_char_rep_map.keys())) + result = pattern.sub(lambda x: self.zh_char_rep_map[x.group()], result) + else: + try: + result = self.en_normalizer.normalize(text) + except Exception: + result = text + print(traceback.format_exc()) + pattern = re.compile("|".join(re.escape(p) for p in self.char_rep_map.keys())) + result = pattern.sub(lambda x: self.char_rep_map[x.group()], result) return result - def correct_pinyin(self, pinyin): + def correct_pinyin(self, pinyin: str): """ 将 jqx 的韵母为 u/ü 的拼音转换为 v 如：ju -> jv , que -> qve, xün -> xvn """ - if pinyin[0] not in "jqx": + if pinyin[0] not in "jqxJQX": return pinyin # 匹配 jqx 的韵母为 u/ü 的拼音 pattern = r"([jqx])[uü](n|e|an)*(\d)" repl = r"\g<1>v\g<2>\g<3>" - pinyin = re.sub(pattern, repl, pinyin) - return pinyin + pinyin = re.sub(pattern, repl, pinyin, flags=re.IGNORECASE) + return pinyin.upper() + + def save_names(self, original_text): + """ + 替换人名为占位符 、 , ... + 例如：克里斯托弗·诺兰 -> + """ + # 人名 + name_pattern = re.compile(TextNormalizer.NAME_PATTERN, re.IGNORECASE) + original_name_list = re.findall(name_pattern, original_text) + if len(original_name_list) == 0: + return (original_text, None) + original_name_list = list(set("".join(n) for n in original_name_list)) + transformed_text = original_text + # 替换占位符 、 , ... + for i, name in enumerate(original_name_list): + number = chr(ord("a") + i) + transformed_text = transformed_text.replace(name, f"") + + return transformed_text, original_name_list + + def restore_names(self, normalized_text, original_name_list): + """ + 恢复人名为原来的文字 + 例如： -> original_name_list[0] + """ + if not original_name_list or len(original_name_list) == 0: + return normalized_text + + transformed_text = normalized_text + # 替换为占位符 、 , ... + for i, name in enumerate(original_name_list): + number = chr(ord("a") + i) + transformed_text = transformed_text.replace(f"", name) + return transformed_text def save_pinyin_tones(self, original_text): """ @@ -112,17 +179,17 @@ class TextNormalizer: 例如：xuan4 -> """ # 声母韵母+声调数字 - origin_pinyin_pattern = re.compile(self.PINYIN_TONE_PATTERN, re.IGNORECASE) + origin_pinyin_pattern = re.compile(TextNormalizer.PINYIN_TONE_PATTERN, re.IGNORECASE) original_pinyin_list = re.findall(origin_pinyin_pattern, original_text) if len(original_pinyin_list) == 0: return (original_text, None) - original_pinyin_list = list(set(''.join(p) for p in original_pinyin_list)) + original_pinyin_list = list(set("".join(p) for p in original_pinyin_list)) transformed_text = original_text # 替换为占位符 , , ... for i, pinyin in enumerate(original_pinyin_list): number = chr(ord("a") + i) transformed_text = transformed_text.replace(pinyin, f"") - + # print("original_text: ", original_text) # print("transformed_text: ", transformed_text) return transformed_text, original_pinyin_list @@ -136,7 +203,7 @@ class TextNormalizer: return normalized_text transformed_text = normalized_text - # 替换为占位符 , , ... + # 替换占位符 , , ... for i, pinyin in enumerate(original_pinyin_list): number = chr(ord("a") + i) pinyin = self.correct_pinyin(pinyin) @@ -145,14 +212,214 @@ class TextNormalizer: # print("transformed_text: ", transformed_text) return transformed_text -if __name__ == '__main__': + +class TextTokenizer: + def __init__(self, vocab_file: str, normalizer: TextNormalizer = None): + self.vocab_file = vocab_file + self.normalizer = normalizer + + if self.vocab_file is None: + raise ValueError("vocab_file is None") + if not os.path.exists(self.vocab_file): + raise ValueError(f"vocab_file {self.vocab_file} does not exist") + if self.normalizer: + self.normalizer.load() + # 加载词表 + self.sp_model = SentencePieceProcessor(model_file=self.vocab_file) + + self.pre_tokenizers = [ + # 预处理器 + tokenize_by_CJK_char, + ] + + @property + def vocab_size(self): + return self.sp_model.GetPieceSize() + + @property + def unk_token(self): + return "" + + @property + def pad_token(self): + return None + + @property + def bos_token(self): + return "" + + @property + def eos_token(self): + return "" + + @property + def pad_token_id(self): + return -1 + + @property + def bos_token_id(self): + return 0 + + @property + def eos_token_id(self): + return 1 + + @property + def unk_token_id(self): + return self.sp_model.unk_id() + + @property + def special_tokens_map(self): + return { + "unk_token": self.unk_token, + "pad_token": self.pad_token, + "bos_token": self.bos_token, + "eos_token": self.eos_token, + } + + def get_vocab(self): + vocab = {self.convert_ids_to_tokens(i): i for i in range(self.vocab_size)} + return vocab + + @overload + def convert_ids_to_tokens(self, ids: int) -> str: ... + + @overload + def convert_ids_to_tokens(self, ids: List[int]) -> List[str]: ... + + def convert_ids_to_tokens(self, ids: Union[List[int], int]): + return self.sp_model.IdToPiece(ids) + + def convert_tokens_to_ids(self, tokens: Union[List[str], str]) -> List[int]: + if isinstance(tokens, str): + tokens = [tokens] + return [self.sp_model.PieceToId(token) for token in tokens] + + def tokenize(self, text: str) -> List[str]: + return self.encode(text, out_type=str) + + def encode(self, text: str, **kwargs): + if len(text) == 0: + return [] + if len(text.strip()) == 1: + return self.sp_model.Encode(text, out_type=kwargs.pop("out_type", int), **kwargs) + # 预处理 + if self.normalizer: + text = self.normalizer.normalize(text) + if len(self.pre_tokenizers) > 0: + for pre_tokenizer in self.pre_tokenizers: + text = pre_tokenizer(text) + return self.sp_model.Encode(text, out_type=kwargs.pop("out_type", int), **kwargs) + + def batch_encode(self, texts: List[str], **kwargs): + # 预处理 + if self.normalizer: + texts = [self.normalizer.normalize(text) for text in texts] + if len(self.pre_tokenizers) > 0: + for pre_tokenizer in self.pre_tokenizers: + texts = [pre_tokenizer(text) for text in texts] + return self.sp_model.Encode(texts, out_type=kwargs.pop("out_type", int), **kwargs) + + def decode(self, ids: Union[List[int], int], do_lower_case=False, **kwargs): + if isinstance(ids, int): + ids = [ids] + decoded = self.sp_model.Decode(ids, out_type=kwargs.pop("out_type", str), **kwargs) + return de_tokenized_by_CJK_char(decoded, do_lower_case=do_lower_case) + + @staticmethod + def split_sentences_by_token( + tokenized_str: List[str], split_tokens: List[str], max_tokens_per_sentence: int + ) -> List[List[str]]: + """ + 将tokenize后的结果按特定token进一步分割 + """ + sentences: List[List[str]] = [] + current_sentence = [] + for i in range(len(tokenized_str)): + token = tokenized_str[i] + current_sentence.append(token) + if token in split_tokens: + if len(current_sentence) == 1: + # 如果当前tokens只有一个，且是切分符号，则忽略这条句子 + pass + elif len(current_sentence) == 2 and current_sentence[0] == '▁': + # 如果当前tokens只有两个，且仅有切分符号，则忽略这条句子 + pass + elif len(current_sentence) <= max_tokens_per_sentence: + if i < len(tokenized_str) - 1: + if tokenized_str[i + 1] in ["'", "▁'"]: + # 后续token是'，则不切分 + current_sentence.append(tokenized_str[i + 1]) + i += 1 + + sentences.append(current_sentence) + else: + # 如果当前tokens的长度超过最大限制 + if "," in current_sentence or "▁," in current_sentence: + # 如果当前tokens中有,，则按,分割 + sub_sentences = TextTokenizer.split_sentences_by_token( + current_sentence, [",", "▁,"], max_tokens_per_sentence=max_tokens_per_sentence + ) + elif "-" in current_sentence: + # 没有,，则按-分割 + sub_sentences = TextTokenizer.split_sentences_by_token( + current_sentence, ["-"], max_tokens_per_sentence=max_tokens_per_sentence + ) + else: + # 按照长度分割 + sub_sentences = [ + current_sentence[:max_tokens_per_sentence], + current_sentence[max_tokens_per_sentence:], + ] + warnings.warn( + f"The tokens length of sentence exceeds limit: {max_tokens_per_sentence}, " + f"Tokens in sentence: {current_sentence}." + "Maybe unexpected behavior", + RuntimeWarning, + ) + sentences.extend(sub_sentences) + current_sentence = [] + if len(current_sentence) > 0: + sentences.append(current_sentence) + # 如果相邻的句子加起来长度小于最大限制，则合并 + merged_sentences = [] + for sentence in sentences: + if len(sentence) == 0: + continue + if len(merged_sentences) == 0: + merged_sentences.append(sentence) + elif len(merged_sentences[-1]) + len(sentence) <= max_tokens_per_sentence: + merged_sentences[-1] = merged_sentences[-1] + sentence + else: + merged_sentences.append(sentence) + return merged_sentences + + punctuation_marks_tokens = [ + ".", + "!", + "?", + "▁.", + # "▁!", # unk + "▁?", + "▁...", # ellipsis + ] + def split_sentences(self, tokenized: List[str], max_tokens_per_sentence=120) -> List[List[str]]: + return TextTokenizer.split_sentences_by_token( + tokenized, self.punctuation_marks_tokens, max_tokens_per_sentence=max_tokens_per_sentence + ) + + +if __name__ == "__main__": # 测试程序 + text_normalizer = TextNormalizer() - text_normalizer.load() + cases = [ + "IndexTTS 正式发布1.0版本了，效果666", + "晕XUAN4是一种GAN3觉", "我爱你！", "I love you!", - "我爱你的英语是”I love you“", + "“我爱你”的英语是“I love you”", "2.5平方电线", "共465篇，约315万字", "2002年的第一场雪，下在了2003年", @@ -164,14 +431,62 @@ if __name__ == '__main__': "他这条视频点赞3000+，评论1000+，收藏500+", "这是1024元的手机，你要吗？", "受不liao3你了", - "”衣裳“不读衣chang2，而是读衣shang5", + "“衣裳”不读衣chang2，而是读衣shang5", "最zhong4要的是：不要chong2蹈覆辙", - "IndexTTS 正式发布1.0版本了，效果666", + "不zuo1死就不会死", "See you at 8:00 AM", "8:00 AM 开会", + "Couting down 3, 2, 1, go!", + "数到3就开始：1、2、3", + "This sales for 2.5% off, only $12.5.", "苹果于2030/1/2发布新 iPhone 2X 系列手机，最低售价仅 ¥12999", + "这酒...里...有毒...", + # 异常case + "只有,,,才是最好的", + # 人名 + "约瑟夫·高登-莱维特（Joseph Gordon-Levitt is an American actor）", + "蒂莫西·唐纳德·库克（英文名：Timothy Donald Cook），通称蒂姆·库克（Tim Cook），美国商业经理、工业工程师和工业开发商，现任苹果公司首席执行官。", + # 长句子 + "《盗梦空间》是由美国华纳兄弟影片公司出品的电影，由克里斯托弗·诺兰执导并编剧，莱昂纳多·迪卡普里奥、玛丽昂·歌迪亚、约瑟夫·高登-莱维特、艾利奥特·佩吉、汤姆·哈迪等联袂主演，2010年7月16日在美国上映，2010年9月1日在中国内地上映，2020年8月28日在中国内地重映。影片剧情游走于梦境与现实之间，被定义为“发生在意识结构内的当代动作科幻片”，讲述了由莱昂纳多·迪卡普里奥扮演的造梦师，带领特工团队进入他人梦境，从他人的潜意识中盗取机密，并重塑他人梦境的故事。", ] - for case in cases: - print(f"原始文本: {case}") - print(f"处理后文本: {text_normalizer.infer(case)}") + # 测试分词器 + tokenizer = TextTokenizer( + vocab_file="checkpoints/bpe.model", + normalizer=text_normalizer, + ) + + codes = tokenizer.batch_encode( + cases, + out_type=int, + ) + + print(f"vocab_size: {tokenizer.vocab_size}") + # print(f"pad_token: {tokenizer.pad_token}, pad_token_id: {tokenizer.pad_token_id}") + print(f"bos_token: {tokenizer.bos_token}, bos_token_id: {tokenizer.bos_token_id}") + print(f"eos_token: {tokenizer.eos_token}, eos_token_id: {tokenizer.eos_token_id}") + print(f"unk_token: {tokenizer.unk_token}, unk_token_id: {tokenizer.unk_token_id}") + # 不应该有 unk_token_id + for t in set([*TextTokenizer.punctuation_marks_tokens, ",", "▁,", "-", "▁..."]): + tokens = tokenizer.convert_tokens_to_ids(t) + if tokenizer.unk_token_id in tokens: + print(f"Warning: {t} is unknown token") + print(f"`{t}`", "->", tokens, "->", tokenizer.convert_ids_to_tokens(tokens)) + for ch in set(tokenizer.normalizer.zh_char_rep_map.values()): + # 测试 normalize后的字符能被分词器识别 + print(f"`{ch}`", "->", tokenizer.sp_model.Encode(ch, out_type=str)) + print(f"` {ch}`", "->", tokenizer.sp_model.Encode(f" {ch}", out_type=str)) + for i in range(len(cases)): + print(f"原始文本: {cases[i]}") + print(f"Normalized: {text_normalizer.normalize(cases[i])}") + tokens = tokenizer.tokenize(cases[i]) + print(f"Tokenzied: {tokens}") + sentences = tokenizer.split_sentences(tokens, max_tokens_per_sentence=100) + print("Splitted sentences count:", len(sentences)) + if len(sentences) > 1: + for j in range(len(sentences)): + print(f" {j}, count:", len(sentences[j]), ", tokens:", "".join(sentences[j])) + #print(f"Token IDs (first 10): {codes[i][:10]}") + if tokenizer.unk_token in codes[i]: + print(f"Warning: `{cases[i]}` contains UNKNOWN token") + print(f"Decoded: {tokenizer.decode(codes[i], do_lower_case=True)}") print("-" * 50) diff --git a/setup.py b/setup.py index 635d7f3..333c559 100644 --- a/setup.py +++ b/setup.py @@ -6,7 +6,7 @@ from setuptools import find_packages, setup setup( name="indextts", - version="0.1.0", + version="0.1.1", author="Index SpeechTeam", author_email="xuanwu@bilibili.com", long_description=open("README.md", encoding="utf8").read(), diff --git a/tests/regression_test.py b/tests/regression_test.py index 58d5be0..aaf5423 100644 --- a/tests/regression_test.py +++ b/tests/regression_test.py @@ -10,7 +10,14 @@ if __name__ == "__main__": tts.infer(audio_prompt=prompt_wav, text=text, output_path=f"outputs/{text[:20]}.wav", verbose=True) text="There is a vehicle arriving in dock number 7?" tts.infer(audio_prompt=prompt_wav, text=text, output_path=f"outputs/{text[:20]}.wav", verbose=True) - + text = "“我爱你！”的英语是“I love you!”" + tts.infer(audio_prompt=prompt_wav, text=text, output_path=f"outputs/{text[:20]}.wav", verbose=True) + text = "Joseph Gordon-Levitt is an American actor" + tts.infer(audio_prompt=prompt_wav, text=text, output_path=f"outputs/{text[:20]}.wav", verbose=True) + text = "约瑟夫·高登-莱维特是美国演员" + tts.infer(audio_prompt=prompt_wav, text=text, output_path=f"outputs/{text[:20]}.wav", verbose=True) + text = "蒂莫西·唐纳德·库克（英文名：Timothy Donald Cook），通称蒂姆·库克（Tim Cook），现任苹果公司首席执行官。" + tts.infer(audio_prompt=prompt_wav, text=text, output_path="outputs/蒂莫西·唐纳德·库克.wav", verbose=True) # 并行推理测试 text="亲爱的伙伴们，大家好！每一次的努力都是为了更好的未来，要善于从失败中汲取经验，让我们一起勇敢前行,迈向更加美好的明天！" tts.infer_fast(audio_prompt=prompt_wav, text=text, output_path=f"outputs/{text[:20]}.wav", verbose=True) @@ -25,4 +32,11 @@ if __name__ == "__main__": 感谢您的收听，下期再见！ '''.replace("\n", "") tts.infer_fast(audio_prompt=prompt_wav, text=text, output_path=f"outputs/{text[:20]}.wav", verbose=True) - + # 长文本推理测试 + text = """《盗梦空间》是由美国华纳兄弟影片公司出品的电影，由克里斯托弗·诺兰执导并编剧， +莱昂纳多·迪卡普里奥、玛丽昂·歌迪亚、约瑟夫·高登-莱维特、艾利奥特·佩吉、汤姆·哈迪等联袂主演， +2010年7月16日在美国上映，2010年9月1日在中国内地上映，2020年8月28日在中国内地重映。 +影片剧情游走于梦境与现实之间，被定义为“发生在意识结构内的当代动作科幻片”， +讲述了由莱昂纳多·迪卡普里奥扮演的造梦师，带领特工团队进入他人梦境，从他人的潜意识中盗取机密，并重塑他人梦境的故事。 +""".replace("\n", "") + tts.infer_fast(audio_prompt=prompt_wav, text=text, output_path=f"outputs/{text[:20]}.wav", verbose=True)