Detect if mps is available across python backends

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

Author: mudler

backend/python/chatterbox/backend.py

@@ -41,7 +41,9 @@ def LoadModel(self, request, context):
         else:
             print("CUDA is not available", file=sys.stderr)
             device = "cpu"
-
+        mps_available = hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
+        if mps_available:
+            device = "mps"
         if not torch.cuda.is_available() and request.CUDA:
             return backend_pb2.Result(success=False, message="CUDA is not available")
 

backend/python/coqui/backend.py

@@ -40,7 +40,9 @@ def LoadModel(self, request, context):
         else:
             print("CUDA is not available", file=sys.stderr)
             device = "cpu"
-
+        mps_available = hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
+        if mps_available:
+            device = "mps"
         if not torch.cuda.is_available() and request.CUDA:
             return backend_pb2.Result(success=False, message="CUDA is not available")
 

backend/python/diffusers/backend.py

@@ -368,6 +368,9 @@ def LoadModel(self, request, context):
             device = "cpu" if not request.CUDA else "cuda"
             if XPU:
                 device = "xpu"
+            mps_available = hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
+            if mps_available:
+                device = "mps"
             self.device = device
             if request.LoraAdapter:
                 # Check if its a local file and not a directory ( we load lora differently for a safetensor file )

backend/python/faster-whisper/backend.py

@@ -10,7 +10,7 @@
 import os
 import backend_pb2
 import backend_pb2_grpc
-
+import torch
 from faster_whisper import WhisperModel
 
 import grpc
@@ -35,7 +35,9 @@ def LoadModel(self, request, context):
         # device = "cuda" if request.CUDA else "cpu"
         if request.CUDA:
             device = "cuda"
-
+        mps_available = hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
+        if mps_available:
+            device = "mps"
         try:
             print("Preparing models, please wait", file=sys.stderr)
             self.model = WhisperModel(request.Model, device=device, compute_type="float16")

backend/python/kitten-tts/backend.py

@@ -33,18 +33,6 @@ def Health(self, request, context):
         return backend_pb2.Reply(message=bytes("OK", 'utf-8'))
     def LoadModel(self, request, context):
 
-        # Get device
-        # device = "cuda" if request.CUDA else "cpu"
-        if torch.cuda.is_available():
-            print("CUDA is available", file=sys.stderr)
-            device = "cuda"
-        else:
-            print("CUDA is not available", file=sys.stderr)
-            device = "cpu"
-
-        if not torch.cuda.is_available() and request.CUDA:
-            return backend_pb2.Result(success=False, message="CUDA is not available")
-
         self.AudioPath = None
         # List available KittenTTS models
         print("Available KittenTTS voices: expr-voice-2-m, expr-voice-2-f, expr-voice-3-m, expr-voice-3-f, expr-voice-4-m, expr-voice-4-f, expr-voice-5-m, expr-voice-5-f")

backend/python/kokoro/backend.py

@@ -33,17 +33,6 @@ def Health(self, request, context):
         return backend_pb2.Reply(message=bytes("OK", 'utf-8'))
 
     def LoadModel(self, request, context):
-        # Get device
-        if torch.cuda.is_available():
-            print("CUDA is available", file=sys.stderr)
-            device = "cuda"
-        else:
-            print("CUDA is not available", file=sys.stderr)
-            device = "cpu"
-
-        if not torch.cuda.is_available() and request.CUDA:
-            return backend_pb2.Result(success=False, message="CUDA is not available")
-
         try:
             print("Preparing Kokoro TTS pipeline, please wait", file=sys.stderr)
             # empty dict

backend/python/transformers/backend.py

@@ -94,7 +94,9 @@ def LoadModel(self, request, context):
         self.SentenceTransformer = False
 
         device_map="cpu"
-
+        mps_available = hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
+        if mps_available:
+            device_map = "mps"
         quantization = None
         autoTokenizer = True