From 4de7611bda5e2f8a10f63829ad9d6706fa94858a Mon Sep 17 00:00:00 2001 From: yrom Date: Sat, 17 May 2025 14:40:01 +0800 Subject: [PATCH] =?UTF-8?q?fix=20=E6=89=B9=E9=87=8F=E6=8E=A8=E7=90=861.5?= =?UTF-8?q?=E7=89=88=E6=9C=AC=E6=A8=A1=E5=9E=8B=E9=97=AE=E9=A2=98=EF=BC=8C?= =?UTF-8?q?=E8=B0=83=E6=95=B4=E5=88=86=E5=8F=A5=E9=80=BB=E8=BE=91=E5=92=8C?= =?UTF-8?q?=E5=8F=82=E6=95=B0=E8=AE=BE=E7=BD=AE?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 将pad 改为全 eos token - 优化bucket_sentences 算法 --- indextts/infer.py | 158 ++++++++++++++++++++++++++++++++-------------- 1 file changed, 109 insertions(+), 49 deletions(-) diff --git a/indextts/infer.py b/indextts/infer.py index 7bacde8..59e5c90 100644 --- a/indextts/infer.py +++ b/indextts/infer.py @@ -2,7 +2,7 @@ import os import re import time from subprocess import CalledProcessError -from typing import List +from typing import Dict, List, Tuple import numpy as np import sentencepiece as spm @@ -125,8 +125,13 @@ class IndexTTS: self.cache_cond_mel = None # 进度引用显示（可选） self.gr_progress = None + self.model_version = self.cfg.version if hasattr(self.cfg, "version") else None def remove_long_silence(self, codes: torch.Tensor, silent_token=52, max_consecutive=30): + """ + Shrink special tokens (silent_token and stop_mel_token) in codes + codes: [B, T] + """ code_lens = [] codes_list = [] device = codes.device @@ -134,59 +139,92 @@ class IndexTTS: isfix = False for i in range(0, codes.shape[0]): code = codes[i] - if self.cfg.gpt.stop_mel_token not in code: - code_lens.append(len(code)) - len_ = len(code) + if not torch.any(code == self.stop_mel_token).item(): + len_ = code.size(0) else: - # len_ = code.cpu().tolist().index(8193)+1 - len_ = (code == self.stop_mel_token).nonzero(as_tuple=False)[0] + 1 - len_ = len_ - 2 + stop_mel_idx = (code == self.stop_mel_token).nonzero(as_tuple=False) + len_ = stop_mel_idx[0].item() if len(stop_mel_idx) > 0 else code.size(0) count = torch.sum(code == silent_token).item() if count > max_consecutive: - code = code.cpu().tolist() - ncode = [] + # code = code.cpu().tolist() + ncode_idx = [] n = 0 - for k in range(0, len_): + for k in range(len_): + assert code[k] != self.stop_mel_token, f"stop_mel_token {self.stop_mel_token} should be shrinked here" if code[k] != silent_token: - ncode.append(code[k]) + ncode_idx.append(k) n = 0 elif code[k] == silent_token and n < 10: - ncode.append(code[k]) + ncode_idx.append(k) n += 1 # if (k == 0 and code[k] == 52) or (code[k] == 52 and code[k-1] == 52): # n += 1 - len_ = len(ncode) - ncode = torch.LongTensor(ncode) - codes_list.append(ncode.to(device, dtype=dtype)) + # new code + len_ = len(ncode_idx) + codes_list.append(code[ncode_idx]) isfix = True - # codes[i] = self.stop_mel_token - # codes[i, 0:len_] = ncode else: - codes_list.append(codes[i]) + # shrink to len_ + codes_list.append(code[:len_]) code_lens.append(len_) - - codes = pad_sequence(codes_list, batch_first=True) if isfix else codes[:, :-2] - code_lens = torch.LongTensor(code_lens).to(device, dtype=dtype) + if isfix: + if len(codes_list) > 1: + codes = pad_sequence(codes_list, batch_first=True, padding_value=self.stop_mel_token) + else: + codes = codes_list[0].unsqueeze(0) + else: + # unchanged + pass + # clip codes to max length + max_len = max(code_lens) + if max_len < codes.shape[1]: + codes = codes[:, :max_len] + code_lens = torch.tensor(code_lens, dtype=torch.long, device=device) return codes, code_lens - def bucket_sentences(self, sentences, enable=False): + def bucket_sentences(self, sentences, bucket_max_size=4) -> List[List[Dict]]: """ - Sentence data bucketing + Sentence data bucketing. + if ``bucket_max_size=1``, return all sentences in one bucket. """ - max_len = max(len(s) for s in sentences) - half = max_len // 2 - outputs = [[], []] + outputs: List = [] for idx, sent in enumerate(sentences): - if enable is False or len(sent) <= half: - outputs[0].append({"idx": idx, "sent": sent}) - else: - outputs[1].append({"idx": idx, "sent": sent}) - return [item for item in outputs if item] + outputs.append({"idx": idx, "sent": sent, "len": len(sent)}) + + if len(outputs) > bucket_max_size: + # split sentences into buckets by sentence length + buckets = [] + factor = 1.5 + last_bucket_sent_len_median = 0 + for sent in sorted(outputs, key=lambda x: x["len"]): + current_sent_len = sent["len"] + if current_sent_len == 0: + print(">> skip empty sentence") + continue + if last_bucket_sent_len_median == 0 \ + or current_sent_len > last_bucket_sent_len_median * factor \ + or len(buckets[-1]) > bucket_max_size: + # new bucket + buckets.append([sent]) + last_bucket_sent_len_median = current_sent_len + else: + # current bucket can hold more sentences + buckets[-1].append(sent) # sorted + mid = len(buckets[-1]) // 2 + last_bucket_sent_len_median = buckets[-1][mid]["len"] + return buckets + return [outputs] - def pad_tokens_cat(self, tokens: List[torch.Tensor]): + + def pad_tokens_cat(self, tokens: List[torch.Tensor]) -> torch.Tensor: if len(tokens) <= 1: return tokens[-1] + if self.model_version and self.model_version >= 1.5: + # 1.5版本以上，直接使用stop_text_token 右侧填充 + # [1, N] -> [N,] + tokens = [t.squeeze(0) for t in tokens] + return pad_sequence(tokens, batch_first=True, padding_value=self.cfg.gpt.stop_text_token, padding_side="right") max_len = max(t.size(1) for t in tokens) outputs = [] for tensor in tokens: @@ -214,8 +252,18 @@ class IndexTTS: 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): + def infer_fast(self, audio_prompt, text, output_path, verbose=False, max_text_tokens_per_sentence=100, sentences_bucket_max_size=4): + """ + Args: + ``max_text_tokens_per_sentence``: 分句的最大token数，默认``100``，可以根据GPU硬件情况调整 + - 越小，batch 越多，推理速度越*快*，占用内存更多，可能影响质量 + - 越大，batch 越少，推理速度越*慢*，占用内存和质量更接近于非快速推理 + ``sentences_bucket_max_size``: 分句分桶的最大容量，默认``4``，可以根据GPU内存调整 + - 越大，bucket数量越少，batch越多，推理速度越*快*，占用内存更多，可能影响质量 + - 越小，bucket数量越多，batch越少，推理速度越*慢*，占用内存和质量更接近于非快速推理 + """ print(">> start fast inference...") + self._set_gr_progress(0, "start fast inference...") if verbose: print(f"origin text:{text}") @@ -245,10 +293,12 @@ class IndexTTS: # text_tokens text_tokens_list = self.tokenizer.tokenize(text) - sentences = self.tokenizer.split_sentences(text_tokens_list) + + sentences = self.tokenizer.split_sentences(text_tokens_list, max_tokens_per_sentence=max_text_tokens_per_sentence) if verbose: - print("text token count:", len(text_tokens_list)) - print("sentences count:", len(sentences)) + print(">> text token count:", len(text_tokens_list)) + print(" splited sentences count:", len(sentences)) + print(" max_text_tokens_per_sentence:", max_text_tokens_per_sentence) print(*sentences, sep="\n") top_p = 0.8 @@ -270,8 +320,11 @@ class IndexTTS: # text processing 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) + bucket_max_size = sentences_bucket_max_size if self.device != "cpu" else 1 + all_sentences = self.bucket_sentences(sentences, bucket_max_size=bucket_max_size) + bucket_count = len(all_sentences) + if verbose: + print(">> sentences bucket_count:", bucket_count, [len(s) for s in all_sentences]) for sentences in all_sentences: temp_tokens: List[torch.Tensor] = [] all_text_tokens.append(temp_tokens) @@ -294,8 +347,8 @@ class IndexTTS: for item_tokens in all_text_tokens: batch_num = len(item_tokens) batch_text_tokens = self.pad_tokens_cat(item_tokens) - batch_cond_mel_lengths = torch.cat([cond_mel_lengths] * batch_num, dim=0) - batch_auto_conditioning = torch.cat([auto_conditioning] * batch_num, dim=0) + batch_cond_mel_lengths = cond_mel_lengths.expand(batch_num) # [batch_num] + batch_auto_conditioning = auto_conditioning.expand(batch_num, -1, -1) # [batch_num, n_mels, L] all_batch_num += batch_num # gpt speech @@ -325,11 +378,15 @@ class IndexTTS: 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 = codes[codes != self.cfg.gpt.stop_mel_token] - codes, _ = torch.unique_consecutive(codes, return_inverse=True) codes = codes.unsqueeze(0) # [x] -> [1, x] - code_lens = torch.tensor([codes.shape[-1]], device=codes.device, dtype=codes.dtype) + if verbose: + print("codes:", codes.shape) + print(codes) codes, code_lens = self.remove_long_silence(codes, silent_token=52, max_consecutive=30) + if verbose: + print("fix codes:", codes.shape) + print(codes) + print("code_lens:", code_lens) text_tokens = batch_tokens[i] all_idxs.append(batch_sentences[i]["idx"]) m_start_time = time.perf_counter() @@ -343,14 +400,16 @@ class IndexTTS: return_latent=True, clip_inputs=False) gpt_forward_time += time.perf_counter() - m_start_time all_latents.append(latent) - + del all_batch_codes, all_text_tokens, all_sentences # bigvgan chunk chunk_size = 2 all_latents = [all_latents[all_idxs.index(i)] for i in range(len(all_latents))] + if verbose: + print(">> all_latents:", len(all_latents)) + print(*[l.shape for l in all_latents], sep=", ") 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...") @@ -370,7 +429,7 @@ class IndexTTS: # clear cache tqdm_progress.close() # 确保进度条被关闭 - chunk_latents.clear() + del all_latents, chunk_latents end_time = time.perf_counter() self.torch_empty_cache() @@ -385,7 +444,7 @@ class IndexTTS: print(f">> Total fast inference time: {end_time - start_time:.2f} seconds") print(f">> Generated audio length: {wav_length:.2f} seconds") print(f">> [fast] bigvgan chunk_length: {chunk_length}") - print(f">> [fast] batch_num: {all_batch_num} bucket_enable: {bucket_enable}") + print(f">> [fast] batch_num: {all_batch_num} bucket_max_size: {bucket_max_size}", f"bucket_count: {bucket_count}" if bucket_max_size > 1 else "") print(f">> [fast] RTF: {(end_time - start_time) / wav_length:.4f}") # save audio @@ -403,7 +462,7 @@ class IndexTTS: return (sampling_rate, wav_data) # 原始推理模式 - def infer(self, audio_prompt, text, output_path, verbose=False): + def infer(self, audio_prompt, text, output_path, verbose=False, max_text_tokens_per_sentence=120): print(">> start inference...") self._set_gr_progress(0, "start inference...") if verbose: @@ -431,10 +490,11 @@ class IndexTTS: auto_conditioning = cond_mel text_tokens_list = self.tokenizer.tokenize(text) - sentences = self.tokenizer.split_sentences(text_tokens_list) + sentences = self.tokenizer.split_sentences(text_tokens_list, max_text_tokens_per_sentence) if verbose: print("text token count:", len(text_tokens_list)) print("sentences count:", len(sentences)) + print("max_text_tokens_per_sentence:", max_text_tokens_per_sentence) print(*sentences, sep="\n") top_p = 0.8 top_k = 30