diff --git a/.gitignore b/.gitignore index f9f2d87..23e0d7b 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,10 @@ venv/ __pycache__ +*.egg-info *.DS_Store -.idea/ \ No newline at end of file +.idea/ +checkpoints/*.pth +checkpoints/*.vocab +checkpoints/*.model +checkpoints/.cache +outputs/ diff --git a/README.md b/README.md index 2d5eea3..0d15775 100644 --- a/README.md +++ b/README.md @@ -103,7 +103,21 @@ conda activate index-tts pip install -r requirements.txt apt-get install ffmpeg ``` + 3. Download models: + +Download by `huggingface-cli`: + +```bash +# 如果下载速度慢，可以使用官方的镜像 +export HF_ENDPOINT="https://hf-mirror.com" +huggingface-cli download IndexTeam/Index-TTS \ + bigvgan_discriminator.pth bigvgan_generator.pth bpe.model dvae.pth gpt.pth unigram_12000.vocab \ + --local-dir checkpoints +``` + +Or by `wget`: + ```bash wget https://huggingface.co/IndexTeam/Index-TTS/resolve/main/bigvgan_discriminator.pth -P checkpoints wget https://huggingface.co/IndexTeam/Index-TTS/resolve/main/bigvgan_generator.pth -P checkpoints @@ -112,11 +126,32 @@ wget https://huggingface.co/IndexTeam/Index-TTS/resolve/main/dvae.pth -P checkpo wget https://huggingface.co/IndexTeam/Index-TTS/resolve/main/gpt.pth -P checkpoints wget https://huggingface.co/IndexTeam/Index-TTS/resolve/main/unigram_12000.vocab -P checkpoints ``` + 4. Run test script: + + ```bash # Please put your prompt audio in 'test_data' and rename it to 'input.wav' PYTHONPATH=. python indextts/infer.py ``` + +5. Use as command line tool: + +```bash +# Make sure pytorch has been installed before running this command +pip install -e . +indextts "大家好，我现在正在bilibili 体验 ai 科技，说实话，来之前我绝对想不到！AI技术已经发展到这样匪夷所思的地步了！" \ + --voice reference_voice.wav \ + --model_dir checkpoints \ + --config checkpoints/config.yaml \ + --output output.wav +``` + +Use `--help` to see more options. +```bash +indextts --help +``` + #### Web Demo ```bash python webui.py diff --git a/indextts/cli.py b/indextts/cli.py new file mode 100644 index 0000000..e0e2f51 --- /dev/null +++ b/indextts/cli.py @@ -0,0 +1,62 @@ +import os +import sys +import warnings +# Suppress warnings from tensorflow and other libraries +warnings.filterwarnings("ignore", category=UserWarning) +warnings.filterwarnings("ignore", category=FutureWarning) +def main(): + import argparse + parser = argparse.ArgumentParser(description="IndexTTS Command Line") + parser.add_argument("text", type=str, help="Text to be synthesized") + parser.add_argument("-v", "--voice", type=str, required=True, help="Path to the audio prompt file (wav format)") + parser.add_argument("-o", "--output_path", type=str, default="gen.wav", help="Path to the output wav file") + parser.add_argument("-c", "--config", type=str, default="checkpoints/config.yaml", help="Path to the config file. Default is 'checkpoints/config.yaml'") + parser.add_argument("--model_dir", type=str, default="checkpoints", help="Path to the model directory. Default is 'checkpoints'") + parser.add_argument("--fp16", action="store_true", default=True, help="Use FP16 for inference if available") + parser.add_argument("-f", "--force", action="store_true", default=False, help="Force to overwrite the output file if it exists") + parser.add_argument("-d", "--device", type=str, default=None, help="Device to run the model on (cpu, cuda, mps)." ) + args = parser.parse_args() + if len(args.text.strip()) == 0: + print("ERROR: Text is empty.") + parser.print_help() + sys.exit(1) + if not os.path.exists(args.voice): + print(f"Audio prompt file {args.voice} does not exist.") + parser.print_help() + sys.exit(1) + if not os.path.exists(args.config): + print(f"Config file {args.config} does not exist.") + parser.print_help() + sys.exit(1) + + output_path = args.output_path + if os.path.exists(output_path): + if not args.force: + print(f"ERROR: Output file {output_path} already exists. Use --force to overwrite.") + parser.print_help() + sys.exit(1) + else: + os.remove(output_path) + + try: + import torch + except ImportError: + print("ERROR: PyTorch is not installed. Please install it first.") + sys.exit(1) + + if args.device is None: + if torch.cuda.is_available(): + args.device = "cuda:0" + elif torch.mps.is_available(): + args.device = "mps" + else: + args.device = "cpu" + args.fp16 = False # Disable FP16 on CPU + print("WARNING: Running on CPU may be slow.") + + from indextts.infer import IndexTTS + tts = IndexTTS(cfg_path=args.config, model_dir=args.model_dir, is_fp16=args.fp16, device=args.device) + tts.infer(audio_prompt=args.voice, text=args.text.strip(), output_path=output_path) + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/indextts/infer.py b/indextts/infer.py index 59d1be3..1019242 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,35 @@ 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 + def __init__(self, cfg_path='checkpoints/config.yaml', model_dir='checkpoints', is_fp16=True, device=None): + """ + Args: + cfg_path (str): path to the config file. + model_dir (str): path to the model directory. + is_fp16 (bool): whether to use fp16. + device (str): device to use (e.g., 'cuda:0', 'cpu'). If None, it will be set automatically based on the availability of CUDA or MPS. + """ + if device is not None: + self.device = device + self.is_fp16 = False if device == 'cpu' else is_fp16 + elif torch.cuda.is_available(): + self.device = 'cuda:0' + self.is_fp16 = is_fp16 + elif torch.mps.is_available(): + self.device = 'mps' + 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.cfg = OmegaConf.load(cfg_path) + 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 +53,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 +90,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 +122,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 +131,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 +146,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 +156,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 +170,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 +194,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 +223,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 +240,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__": diff --git a/requirements.txt b/requirements.txt index 803d193..5491e63 100644 --- a/requirements.txt +++ b/requirements.txt @@ -17,9 +17,9 @@ accelerate==0.25.0 tensorboard==2.9.1 omegaconf sentencepiece -pypinyin librosa gradio tqdm -WeTextProcessing # arm机器如果安装失败，请注释此行 -wetext \ No newline at end of file + +WeTextProcessing; platform_machine != "Darwin" +wetext; platform_system == "Darwin" \ No newline at end of file diff --git a/setup.py b/setup.py new file mode 100644 index 0000000..a0e86fc --- /dev/null +++ b/setup.py @@ -0,0 +1,50 @@ + +import platform +from setuptools import find_packages, setup + + + +setup( + name="indextts", + version="0.1.0", + author="Index SpeechTeam", + author_email="xuanwu@bilibili.com", + long_description=open("README.md", encoding="utf8").read(), + long_description_content_type="text/markdown", + description="An Industrial-Level Controllable and Efficient Zero-Shot Text-To-Speech System", + url="https://github.com/index-tts/index-tts", + packages=find_packages(), + include_package_data=True, + install_requires=[ + "torch==2.6.0", + "torchaudio", + "transformers==4.36.2", + "accelerate", + "tokenizers==0.15.0", + "einops==0.8.1", + "matplotlib==3.8.2", + "omegaconf", + "sentencepiece", + "librosa", + "numpy", + "wetext" if platform.system() == "Darwin" else "WeTextProcessing", + ], + extras_require={ + "webui": ["gradio"], + }, + entry_points={ + "console_scripts": [ + "indextts = indextts.cli:main", + ] + }, + license="Apache-2.0", + python_requires=">=3.10", + classifiers=[ + "Programming Language :: Python :: 3", + "Operating System :: OS Independent", + "License :: OSI Approved :: Apache Software License", + "Intended Audience :: Science/Research", + "Topic :: Scientific/Engineering", + "Topic :: Scientific/Engineering :: Artificial Intelligence", + ], +) \ No newline at end of file diff --git a/webui.py b/webui.py index 1d3c7dd..5a4527f 100644 --- a/webui.py +++ b/webui.py @@ -68,4 +68,4 @@ with gr.Blocks() as demo: if __name__ == "__main__": demo.queue(20) - demo.launch(server_name="0.0.0.0") + demo.launch(server_name="127.0.0.1")