From ec65755fc874d8d408a6168918b7e8a1db93cf1f Mon Sep 17 00:00:00 2001 From: Yrom Date: Fri, 11 Apr 2025 20:52:31 +0800 Subject: [PATCH] Support inference on CPU --- indextts/cli.py | 2 +- indextts/infer.py | 128 +++++++++++++++++++--------------------------- 2 files changed, 55 insertions(+), 75 deletions(-) diff --git a/indextts/cli.py b/indextts/cli.py index 285b40c..18f12da 100644 --- a/indextts/cli.py +++ b/indextts/cli.py @@ -37,7 +37,7 @@ def main(): try: import torch if not torch.cuda.is_available(): - print("WARNING: CUDA is not available. Running in CPU mode.") + print("WARNING: CUDA is not available. Running on CPU may be slow.") except ImportError: print("ERROR: PyTorch is not installed. Please install it first.") sys.exit(1) diff --git a/indextts/infer.py b/indextts/infer.py index 59d1be3..db89333 100644 --- a/indextts/infer.py +++ b/indextts/infer.py @@ -1,7 +1,6 @@ import os import re -import sys - +import time import sentencepiece as spm import torch import torchaudio @@ -16,17 +15,22 @@ from indextts.utils.common import tokenize_by_CJK_char from indextts.vqvae.xtts_dvae import DiscreteVAE from indextts.utils.front import TextNormalizer + class IndexTTS: def __init__(self, cfg_path='checkpoints/config.yaml', model_dir='checkpoints', is_fp16=True): self.cfg = OmegaConf.load(cfg_path) - self.device = 'cuda:0' - self.model_dir = model_dir - self.is_fp16 = is_fp16 - self.stop_mel_token = self.cfg.gpt.stop_mel_token - if self.is_fp16: - self.dtype = torch.float16 + if torch.cuda.is_available(): + self.device = 'cuda:0' + self.is_fp16 = is_fp16 else: - self.dtype = None + self.device = 'cpu' + self.is_fp16 = False + print(">> Be patient, it may take a while to run in CPU mode.") + + 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 + self.dvae = DiscreteVAE(**self.cfg.vqvae) self.dvae_path = os.path.join(self.model_dir, self.cfg.dvae_checkpoint) load_checkpoint(self.dvae, self.dvae_path) @@ -36,7 +40,6 @@ class IndexTTS: else: self.dvae.eval() print(">> vqvae weights restored from:", self.dvae_path) - self.gpt = UnifiedVoice(**self.cfg.gpt) self.gpt_path = os.path.join(self.model_dir, self.cfg.gpt_checkpoint) load_checkpoint(self.gpt, self.gpt_path) @@ -74,9 +77,11 @@ class IndexTTS: # return text.translate(punctuation_map) return self.normalizer.infer(text) - def remove_long_silence(self, codes, silent_token=52, max_consecutive=30): + def remove_long_silence(self, codes: torch.Tensor, silent_token=52, max_consecutive=30): code_lens = [] codes_list = [] + device = codes.device + dtype = codes.dtype isfix = False for i in range(0, codes.shape[0]): code = codes[i] @@ -104,7 +109,7 @@ class IndexTTS: # n += 1 len_ = len(ncode) ncode = torch.LongTensor(ncode) - codes_list.append(ncode.cuda()) + codes_list.append(ncode.to(device, dtype=dtype)) isfix = True #codes[i] = self.stop_mel_token #codes[i, 0:len_] = ncode @@ -113,7 +118,7 @@ class IndexTTS: code_lens.append(len_) codes = pad_sequence(codes_list, batch_first=True) if isfix else codes[:, :-2] - code_lens = torch.LongTensor(code_lens).cuda() + code_lens = torch.LongTensor(code_lens).to(device, dtype=dtype) return codes, code_lens def infer(self, audio_prompt, text, output_path): @@ -128,7 +133,7 @@ class IndexTTS: audio = audio[0].unsqueeze(0) audio = torchaudio.transforms.Resample(sr, 24000)(audio) cond_mel = MelSpectrogramFeatures()(audio).to(self.device) - print(f"cond_mel shape: {cond_mel.shape}") + print(f"cond_mel shape: {cond_mel.shape}", "dtype:", cond_mel.dtype) auto_conditioning = cond_mel @@ -138,7 +143,7 @@ class IndexTTS: punctuation = ["!", "?", ".", ";", "！", "？", "。", "；"] pattern = r"(?<=[{0}])\s*".format("".join(punctuation)) sentences = [i for i in re.split(pattern, text) if i.strip() != ""] - print(sentences) + print("sentences:", sentences) top_p = .8 top_k = 30 @@ -152,20 +157,23 @@ class IndexTTS: lang = "EN" lang = "ZH" wavs = [] - wavs1 = [] + print(">> start inference...") + + start_time = time.time() for sent in sentences: print(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 而 散 ." - print(cleand_text) + print("cleand_text:", cleand_text) + text_tokens = torch.IntTensor(tokenizer.encode(cleand_text)).unsqueeze(0).to(self.device) # 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 = text_tokens.to(self.device) + # text_tokens = text_tokens.to(self.device) print(text_tokens) print(f"text_tokens shape: {text_tokens.shape}, text_tokens type: {text_tokens.dtype}") text_token_syms = [tokenizer.IdToPiece(idx) for idx in text_tokens[0].tolist()] @@ -173,38 +181,24 @@ class IndexTTS: text_len = [text_tokens.size(1)] text_len = torch.IntTensor(text_len).to(self.device) print(text_len) + with torch.no_grad(): - if self.is_fp16: - with torch.cuda.amp.autocast(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), - # text_lengths=text_len, - do_sample=True, - top_p=top_p, - top_k=top_k, - temperature=temperature, - num_return_sequences=autoregressive_batch_size, - length_penalty=length_penalty, - num_beams=num_beams, - repetition_penalty=repetition_penalty, - max_generate_length=max_mel_tokens) - else: + with torch.amp.autocast(self.device, 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), - # text_lengths=text_len, - do_sample=True, - top_p=top_p, - top_k=top_k, - temperature=temperature, - num_return_sequences=autoregressive_batch_size, - length_penalty=length_penalty, - num_beams=num_beams, - repetition_penalty=repetition_penalty, - max_generate_length=max_mel_tokens) + cond_mel_lengths=torch.tensor([auto_conditioning.shape[-1]], + device=text_tokens.device), + # text_lengths=text_len, + do_sample=True, + top_p=top_p, + top_k=top_k, + temperature=temperature, + num_return_sequences=autoregressive_batch_size, + length_penalty=length_penalty, + num_beams=num_beams, + repetition_penalty=repetition_penalty, + max_generate_length=max_mel_tokens) #codes = codes[:, :-2] - code_lens = torch.tensor([codes.shape[-1]]) + code_lens = torch.tensor([codes.shape[-1]], device=codes.device, dtype=codes.dtype) print(codes, type(codes)) print(f"codes shape: {codes.shape}, codes type: {codes.dtype}") print(f"code len: {code_lens}") @@ -216,35 +210,14 @@ class IndexTTS: print(f"code len: {code_lens}") # latent, text_lens_out, code_lens_out = \ - if self.is_fp16: - with torch.cuda.amp.autocast(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, - code_lens*self.gpt.mel_length_compression, - cond_mel_lengths=torch.tensor([auto_conditioning.shape[-1]], device=text_tokens.device), - return_latent=True, clip_inputs=False) - latent = latent.transpose(1, 2) - wav, _ = self.bigvgan(latent.transpose(1, 2), auto_conditioning.transpose(1, 2)) - wav = wav.squeeze(1).cpu() - - else: + with torch.amp.autocast(self.device, 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, - code_lens*self.gpt.mel_length_compression, - cond_mel_lengths=torch.tensor([auto_conditioning.shape[-1]], device=text_tokens.device), - return_latent=True, clip_inputs=False) + torch.tensor([text_tokens.shape[-1]], device=text_tokens.device), codes, + code_lens*self.gpt.mel_length_compression, + cond_mel_lengths=torch.tensor([auto_conditioning.shape[-1]], device=text_tokens.device), + return_latent=True, clip_inputs=False) latent = latent.transpose(1, 2) - ''' - latent_list = [] - for lat, t_len in zip(latent, text_lens_out): - lat = lat[:, t_len:] - latent_list.append(lat) - latent = torch.stack(latent_list) - print(f"latent shape: {latent.shape}") - ''' - wav, _ = self.bigvgan(latent.transpose(1, 2), auto_conditioning.transpose(1, 2)) wav = wav.squeeze(1).cpu() @@ -254,8 +227,15 @@ class IndexTTS: # wavs.append(wav[:, :-512]) wavs.append(wav) + end_time = time.time() + elapsed_time = end_time - start_time + minutes, seconds = divmod(int(elapsed_time), 60) + milliseconds = int((elapsed_time - int(elapsed_time)) * 1000) + print(f">> inference done. time: {minutes}:{seconds}.{milliseconds}") + print(">> saving wav file") wav = torch.cat(wavs, dim=1) - torchaudio.save(output_path, wav.type(torch.int16), 24000) + torchaudio.save(output_path, wav.type(torch.int16), sampling_rate) + print(">> wav file saved to:", output_path) if __name__ == "__main__":