add mac m1 mps support (#2477)

Author: xmx0632

funasr/auto/auto_model.py

@@ -184,6 +184,7 @@ def build_model(**kwargs):
         device = kwargs.get("device", "cuda")
         if ((device =="cuda" and not torch.cuda.is_available())
             or (device == "xpu" and not torch.xpu.is_available())
+            or (device == "mps" and not torch.backends.mps.is_available())
             or kwargs.get("ngpu", 1) == 0):
             device = "cpu"
             kwargs["batch_size"] = 1

funasr/frontends/fused.py

@@ -80,6 +80,8 @@ def __init__(self, frontends=None, align_method="linear_projection", proj_dim=10
             dev = "cuda"
         elif torch.xpu.is_available():
             dev = "xpu"
+        elif torch.backends.mps.is_available():
+            dev = "mps"
         else:
             dev = "cpu"
         if self.align_method == "linear_projection":

funasr/utils/export_utils.py

@@ -28,12 +28,12 @@ def export(
                 **kwargs,
             )
         elif type == "torchscript":
-            device = "cuda" if torch.cuda.is_available() else "xpu" if torch.xpu.is_available() else "cpu"
+            device = "cuda" if torch.cuda.is_available() else "xpu" if torch.xpu.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
             print("Exporting torchscripts on device {}".format(device))
             _torchscripts(m, path=export_dir, device=device)
         elif type == "bladedisc":
             assert (
-                torch.cuda.is_available() or torch.xpu.is_available()
+                torch.cuda.is_available() or torch.xpu.is_available() or torch.backends.mps.is_available()
             ), "Currently bladedisc optimization for FunASR only supports GPU"
             # bladedisc only optimizes encoder/decoder modules
             if hasattr(m, "encoder") and hasattr(m, "decoder"):
@@ -44,7 +44,7 @@ def export(
 
         elif type == "onnx_fp16":
             assert (
-                torch.cuda.is_available() or torch.xpu.is_available()
+                torch.cuda.is_available() or torch.xpu.is_available() or torch.backends.mps.is_available()
             ), "Currently onnx_fp16 optimization for FunASR only supports GPU"
 
             if hasattr(m, "encoder") and hasattr(m, "decoder"):