Merge pull request #15 from builtbybasit/mlx-support

Add Mac MPS support

pyproject.toml

@@ -5,7 +5,25 @@ description = "Linacodec is a highly compressive and rapid audio tokenizer for a
 readme = "README.md"
 requires-python = ">=3.10"
 dependencies = [
-  "huggingface_hub",
+  "cn2an",
+  "huggingface-hub",
+  "inflect",
+  "jieba",
+  "lhotse",
+  "librosa",
+  "linacodec",
+  "numpy",
+  "onnxruntime",
+  "piper-phonemize",
+  "pydub",
+  "pypinyin",
+  "safetensors",
+  "setuptools<81",
+  "tensorboard",
+  "torch",
+  "torchaudio",
+  "transformers<=4.57.6",
+  "vocos",
 ]
 
 [project.optional-dependencies]
@@ -19,3 +37,12 @@ build-backend = "uv_build"
 
 [tool.ruff]
 line-length = 120
+
+[tool.uv]
+find-links = ["https://k2-fsa.github.io/icefall/piper_phonemize.html"]
+
+[tool.uv.sources]
+linacodec = { git = "https://github.com/ysharma3501/LinaCodec.git" }
+
+[tool.uv.build-backend]
+module-root = "" # The root module of the project

zipvoice/luxvoice.py

@@ -1,20 +1,30 @@
+import torch
 from zipvoice.modeling_utils import process_audio, generate, load_models_gpu, load_models_cpu
 from zipvoice.onnx_modeling import generate_cpu
 
 class LuxTTS:
     """
-    LuxTTS class for encoding prompt and generating speech on cpu/cuda.
+    LuxTTS class for encoding prompt and generating speech on cpu/cuda/mps.
     """
 
     def __init__(self, model_path='YatharthS/LuxTTS', device='cuda', threads=4):
         if model_path == 'YatharthS/LuxTTS':
             model_path = None
 
+        # Auto-detect better device if cuda is requested but not available
+        if device == 'cuda' and not torch.cuda.is_available():
+            if torch.backends.mps.is_available():
+                print("CUDA not available, switching to MPS")
+                device = 'mps'
+            else:
+                print("CUDA not available, switching to CPU")
+                device = 'cpu'
+
         if device == 'cpu':
             model, feature_extractor, vocos, tokenizer, transcriber = load_models_cpu(model_path, threads)
             print("Loading model on CPU")
         else:
-            model, feature_extractor, vocos, tokenizer, transcriber = load_models_gpu(model_path)
+            model, feature_extractor, vocos, tokenizer, transcriber = load_models_gpu(model_path, device=device)
             print("Loading model on GPU")
 
         self.model = model

zipvoice/modeling_utils.py

@@ -90,7 +90,7 @@ def generate(prompt_tokens, prompt_features_lens, prompt_features, prompt_rms, t
 
     return wav
 
-def load_models_gpu(model_path=None):
+def load_models_gpu(model_path=None, device="cuda"):
     params = LuxTTSConfig()
     if model_path is None:
         model_path = snapshot_download("YatharthS/LuxTTS")
@@ -99,7 +99,7 @@ def load_models_gpu(model_path=None):
     model_ckpt = f"{model_path}/model.pt"
     model_config = f"{model_path}/config.json"
 
-    transcriber = pipeline("automatic-speech-recognition", model="openai/whisper-base", device='cuda')
+    transcriber = pipeline("automatic-speech-recognition", model="openai/whisper-base", device=device)
     tokenizer = EmiliaTokenizer(token_file=token_file)
     tokenizer_config = {"vocab_size": tokenizer.vocab_size, "pad_id": tokenizer.pad_id}
 
@@ -111,15 +111,15 @@ def load_models_gpu(model_path=None):
         **tokenizer_config,
     )
     load_checkpoint(filename=model_ckpt, model=model, strict=True)
-    params.device = torch.device("cuda", 0)
+    params.device = torch.device(device, 0)
 
     model = model.to(params.device).eval()
     feature_extractor = VocosFbank()
 
-    vocos = Vocos.from_hparams(f'{model_path}/vocoder/config.yaml').cuda()
+    vocos = Vocos.from_hparams(f'{model_path}/vocoder/config.yaml').to(device)
     parametrize.remove_parametrizations(vocos.upsampler.upsample_layers[0], "weight")
     parametrize.remove_parametrizations(vocos.upsampler.upsample_layers[1], "weight")
-    vocos.load_state_dict(torch.load(f'{model_path}/vocoder/vocos.bin'))
+    vocos.load_state_dict(torch.load(f'{model_path}/vocoder/vocos.bin', map_location=params.device))
 
     params.sampling_rate = model_config["feature"]["sampling_rate"]
     return model, feature_extractor, vocos, tokenizer, transcriber

zipvoice/models/modules/scaling.py

@@ -33,10 +33,24 @@ import torch
 import torch.nn as nn
 from torch import Tensor
 
+if torch.cuda.is_available():
+    DEVICE_TYPE = "cuda"
+elif torch.backends.mps.is_available():
+    DEVICE_TYPE = "mps"
+else:
+    DEVICE_TYPE = "cpu"
+
+def get_memory_allocated():
+    if DEVICE_TYPE == "cuda":
+        return torch.cuda.memory_allocated()
+    elif DEVICE_TYPE == "mps":
+        return torch.mps.current_allocated_memory()
+    else:
+        return 0
 
 def custom_amp_decorator(dec, cuda_amp_deprecated):
     def decorator(func):
-        return dec(func) if not cuda_amp_deprecated else dec(device_type="cuda")(func)
+        return dec(func) if not cuda_amp_deprecated else dec(device_type=DEVICE_TYPE)(func)
 
     return decorator
 
@@ -319,7 +333,7 @@ class SoftmaxFunction(torch.autograd.Function):
     @staticmethod
     def backward(ctx, ans_grad: Tensor):
         (ans,) = ctx.saved_tensors
-        with torch.amp.autocast("cuda", enabled=False):
+        with torch.amp.autocast(DEVICE_TYPE, enabled=False):
             ans_grad = ans_grad.to(torch.float32)
             ans = ans.to(torch.float32)
             x_grad = ans_grad * ans
@@ -535,7 +549,7 @@ class BalancerFunction(torch.autograd.Function):
 
         try:
             with torch.enable_grad():
-                with torch.amp.autocast("cuda", enabled=False):
+                with torch.amp.autocast(DEVICE_TYPE, enabled=False):
                     x = x.to(torch.float32)
                     x = x.detach()
                     x.requires_grad = True
@@ -648,7 +662,7 @@ class Balancer(torch.nn.Module):
         if (
             torch.jit.is_scripting()
             or not x.requires_grad
-            or (x.is_cuda and self.mem_cutoff(torch.cuda.memory_allocated()))
+            or ((x.is_cuda or x.device.type == "mps") and self.mem_cutoff(get_memory_allocated()))
         ):
             return _no_op(x)
 
@@ -802,7 +816,7 @@ class WhiteningPenaltyFunction(torch.autograd.Function):
 
         try:
             with torch.enable_grad():
-                with torch.amp.autocast("cuda", enabled=False):
+                with torch.amp.autocast(DEVICE_TYPE, enabled=False):
                     x_detached = x_orig.to(torch.float32).detach()
                     x_detached.requires_grad = True
 
@@ -1046,7 +1060,7 @@ class SwooshLFunction(torch.autograd.Function):
 
         coeff = -0.08
 
-        with torch.amp.autocast("cuda", enabled=False):
+        with torch.amp.autocast(DEVICE_TYPE, enabled=False):
             with torch.enable_grad():
                 x = x.detach()
                 x.requires_grad = True
@@ -1124,7 +1138,7 @@ class SwooshRFunction(torch.autograd.Function):
 
         zero = torch.tensor(0.0, dtype=x.dtype, device=x.device)
 
-        with torch.amp.autocast("cuda", enabled=False):
+        with torch.amp.autocast(DEVICE_TYPE, enabled=False):
             with torch.enable_grad():
                 x = x.detach()
                 x.requires_grad = True
@@ -1187,7 +1201,7 @@ class SwooshROnnx(torch.nn.Module):
 # simple version of SwooshL that does not redefine the backprop, used in
 # ActivationDropoutAndLinearFunction.
 def SwooshLForward(x: Tensor):
-    with torch.amp.autocast("cuda", enabled=False):
+    with torch.amp.autocast(DEVICE_TYPE, enabled=False):
         x = x.to(torch.float32)
         x_offset = x - 4.0
         log_sum = (1.0 + x_offset.exp()).log().to(x.dtype)
@@ -1198,7 +1212,7 @@ def SwooshLForward(x: Tensor):
 # simple version of SwooshR that does not redefine the backprop, used in
 # ActivationDropoutAndLinearFunction.
 def SwooshRForward(x: Tensor):
-    with torch.amp.autocast("cuda", enabled=False):
+    with torch.amp.autocast(DEVICE_TYPE, enabled=False):
         x = x.to(torch.float32)
         x_offset = x - 1.0
         log_sum = (1.0 + x_offset.exp()).log().to(x.dtype)

zipvoice/models/modules/zipformer.py

@@ -27,6 +27,13 @@ from typing import Optional, Tuple, Union
 import torch
 from torch import Tensor, nn
 
+if torch.cuda.is_available():
+    DEVICE_TYPE = "cuda"
+elif torch.backends.mps.is_available():
+    DEVICE_TYPE = "mps"
+else:
+    DEVICE_TYPE = "cpu"
+
 from zipvoice.models.modules.scaling import (
     ActivationDropoutAndLinear,
     Balancer,
@@ -1310,7 +1317,7 @@ class RelPositionMultiheadAttentionWeights(nn.Module):
         (num_heads, batch_size, seq_len, seq_len) = attn_weights.shape
 
         with torch.no_grad():
-            with torch.amp.autocast("cuda", enabled=False):
+            with torch.amp.autocast(DEVICE_TYPE, enabled=False):
                 attn_weights = attn_weights.to(torch.float32)
                 attn_weights_entropy = (
                     -((attn_weights + 1.0e-20).log() * attn_weights)