Merge pull request #78 from yrom/feat/mac-support

Add a new Command-Line Interface and support for mps device (Apple Silicon)

.gitignore

@@ -1,4 +1,10 @@
 venv/
 __pycache__
+*.egg-info
 *.DS_Store
-.idea/
+.idea/
+checkpoints/*.pth
+checkpoints/*.vocab
+checkpoints/*.model
+checkpoints/.cache
+outputs/

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

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()

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__":

requirements.txt

@@ -17,9 +17,9 @@ accelerate==0.25.0
 tensorboard==2.9.1
 omegaconf
 sentencepiece
-pypinyin
 librosa
 gradio
 tqdm
-WeTextProcessing # arm机器如果安装失败，请注释此行
-wetext
+
+WeTextProcessing; platform_machine != "Darwin"
+wetext; platform_system == "Darwin"

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",
+    ],
+)

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")