Support inference on CPU
This commit is contained in:
Yrom
parent
471a45435c
commit
ec65755fc8
@ -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)
|
||||
|
||||
@ -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__":
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user