From 879e270d39f6b816bed8b830a1360d74fe2f2e09 Mon Sep 17 00:00:00 2001
From: Yrom <yrom@qq.com>
Date: Fri, 11 Apr 2025 21:19:21 +0800
Subject: [PATCH] Adds MPS support for Apple Silicon

---
 indextts/cli.py   | 25 ++++++++++++++++++-------
 indextts/infer.py | 19 ++++++++++++++++---
 2 files changed, 34 insertions(+), 10 deletions(-)

diff --git a/indextts/cli.py b/indextts/cli.py
index 18f12da..e0e2f51 100644
--- a/indextts/cli.py
+++ b/indextts/cli.py
@@ -14,8 +14,12 @@ def main():
     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()
@@ -33,19 +37,26 @@ def main():
             sys.exit(1)
         else:
             os.remove(output_path)
-
+    
     try:
         import torch
-        if not torch.cuda.is_available():
-            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)
-            
+
+    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)
-    tts.infer(audio_prompt=args.voice, text=args.text, output_path=output_path)
+    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 db89333..1019242 100644
--- a/indextts/infer.py
+++ b/indextts/infer.py
@@ -17,16 +17,29 @@ 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)
-        if torch.cuda.is_available():
+    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.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