Enhance device compatibility by auto-detecting available hardware (CUDA/MPS) and updating model loading functions accordingly

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
@@ -24,7 +34,7 @@ class LuxTTS:
         self.transcriber = transcriber
         self.device = device
         self.vocos.freq_range = 12000
-        
+
 
 
     def encode_prompt(self, prompt_audio, duration=5, rms=0.001):
@@ -33,17 +43,17 @@ class LuxTTS:
         encode_dict = {"prompt_tokens": prompt_tokens, 'prompt_features_lens': prompt_features_lens, 'prompt_features': prompt_features, 'prompt_rms': prompt_rms}
 
         return encode_dict
-      
+
     def generate_speech(self, text, encode_dict, num_steps=4, guidance_scale=3.0, t_shift=0.5, speed=1.0, return_smooth=False):
         """encodes text and generates speech using flow matching model according to steps, guidance scale, and t_shift(like temp)"""
-      
+
         prompt_tokens, prompt_features_lens, prompt_features, prompt_rms = encode_dict.values()
 
         if return_smooth == True:
             self.vocos.return_48k = False
         else:
             self.vocos.return_48k = True
-            
+
         if self.device == 'cpu':
             final_wav = generate_cpu(prompt_tokens, prompt_features_lens, prompt_features, prompt_rms, text, self.model, self.vocos, self.tokenizer, num_step=num_steps, guidance_scale=guidance_scale, t_shift=t_shift, speed=speed)
         else:

zipvoice/modeling_utils.py

@@ -52,7 +52,7 @@ def process_audio(audio, transcriber, tokenizer, feature_extractor, device, targ
 
     prompt_wav = torch.from_numpy(prompt_wav).unsqueeze(0)
     prompt_wav, prompt_rms = rms_norm(prompt_wav, target_rms)
-    
+
     prompt_features = feature_extractor.extract(
         prompt_wav, sampling_rate=24000
     ).to(device)
@@ -90,16 +90,16 @@ 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")
-    
+
     token_file = f"{model_path}/tokens.txt"
     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,16 +111,16 @@ 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
 
@@ -136,15 +136,15 @@ def load_models_cpu(model_path = None, num_thread=2):
     model_config  = f"{model_path}/config.json"
 
     transcriber = pipeline("automatic-speech-recognition", model="openai/whisper-tiny", device='cpu')
-    
+
     tokenizer = EmiliaTokenizer(token_file=token_file)
     tokenizer_config = {"vocab_size": tokenizer.vocab_size, "pad_id": tokenizer.pad_id}
-    
+
     with open(model_config, "r") as f:
         model_config = json.load(f)
 
     model = OnnxModel(text_encoder_path, fm_decoder_path, num_thread=num_thread)
-    
+
     vocos = Vocos.from_hparams(f'{model_path}/vocoder/config.yaml').eval()
     parametrize.remove_parametrizations(vocos.upsampler.upsample_layers[0], "weight")
     parametrize.remove_parametrizations(vocos.upsampler.upsample_layers[1], "weight")

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)