Cloyir/index-tts2-ForDgxSpark
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Adds MPS support for Apple Silicon

Browse Source
This commit is contained in:
Yrom 2025-04-11 21:19:21 +08:00
parent ec65755fc8
commit 879e270d39
No known key found for this signature in database
2 changed files with 34 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

25
indextts/cli.py
View File

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

19
indextts/infer.py
View File

@ -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