Added tentative Apple Silicon GPU support

Author: beveradb

audio_separator/separator.py

@@ -23,6 +23,7 @@ def __init__(
         model_file_dir="/tmp/audio-separator-models/",
         output_dir=None,
         use_cuda=False,
+        use_coreml=False,
         output_format="WAV",
         output_subtype=None,
         normalization_enabled=True,
@@ -50,6 +51,7 @@ def __init__(
         self.model_file_dir = model_file_dir
         self.output_dir = output_dir
         self.use_cuda = use_cuda
+        self.use_coreml = use_coreml
 
         # Create the model directory if it does not exist
         os.makedirs(self.model_file_dir, exist_ok=True)
@@ -116,6 +118,26 @@ def __init__(
             else:
                 raise Exception("CUDA requested but not available with current Torch installation. Do you have an Nvidia GPU?")
 
+        elif self.use_coreml:
+            self.logger.debug("Apple Silicon CoreML requested, checking Torch version")
+            self.logger.debug(f"Torch version: {str(torch.__version__)}")
+
+            mps_available = hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
+            self.logger.debug(f"Is Apple Silicon CoreML MPS available? {str(mps_available)}")
+
+            if mps_available:
+                self.logger.debug("Running in Apple Silicon MPS GPU mode")
+
+                # TODO: Change this to use MPS once FFTs are supported, see https://github.com/pytorch/pytorch/issues/78044
+                # self.device = torch.device("mps")
+
+                self.device = torch.device("cpu")
+                self.run_type = ["CoreMLExecutionProvider"]
+            else:
+                raise Exception(
+                    "Apple Silicon CoreML / MPS requested but not available with current Torch installation. Do you have an Apple Silicon GPU?"
+                )
+
         else:
             self.logger.debug("Running in CPU mode")
             self.device = torch.device("cpu")

audio_separator/utils/cli.py

@@ -52,6 +52,12 @@ def main():
         help="Optional: use Nvidia GPU with CUDA for separation (default: %(default)s). Example: --use_cuda=true",
     )
 
+    parser.add_argument(
+        "--use_coreml",
+        action="store_true",
+        help="Optional: use Apple Silicon GPU with CoreML for separation (default: %(default)s). Example: --use_coreml=true",
+    )
+
     parser.add_argument(
         "--output_format",
         default="FLAC",
@@ -97,6 +103,7 @@ def main():
         model_file_dir=args.model_file_dir,
         output_dir=args.output_dir,
         use_cuda=args.use_cuda,
+        use_coreml=args.use_coreml,
         output_format=args.output_format,
         denoise_enabled=args.denoise,
         normalization_enabled=args.normalize,