Added support for metal performance shaders (MPS by apple)

Author: tsepuri

src/icatcher/cli.py

@@ -192,11 +192,19 @@ def load_models(opt, download_only=False):
             face_detector_model_file = file_paths[
                 file_names.index("Resnet50_Final.pth")
             ]
-            face_detector_model = RetinaFace(
+            if opt.device.startswith("mps"):
+                face_detector_model = RetinaFace(
+                    gpu_id=opt.gpu_id,
+                    model_path=face_detector_model_file,
+                    network="resnet50",
+                    device="mps",
+                )
+            else:
+                face_detector_model = RetinaFace(
                 gpu_id=opt.gpu_id,
                 model_path=face_detector_model_file,
                 network="resnet50",
-            )
+                )
         elif opt.fd_model == "opencv_dnn":
             face_detector_model_file = file_paths[
                 file_names.index("face_model.caffemodel")
@@ -215,6 +223,10 @@ def load_models(opt, download_only=False):
             state_dict = torch.load(
                 str(path_to_gaze_model), map_location=torch.device(opt.device)
             )
+        elif opt.device.startswith("mps"):
+            state_dict = torch.load(
+                str(path_to_gaze_model), map_location=torch.device("mps")
+            )
         else:
             state_dict = torch.load(str(path_to_gaze_model))
         try:

src/icatcher/options.py

@@ -272,13 +272,17 @@ def parse_arguments(my_string=None):
         args.device = "cpu"
     else:
         import os
-
-        os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu_id)
-        args.device = "cuda:{}".format(0)
         import torch
 
-        if not torch.cuda.is_available():
-            raise ValueError("GPU is not available. Was torch compiled with CUDA?")
+        if torch.cuda.is_available():
+            os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu_id)
+            args.device = f"cuda:{args.gpu_id}"
+        else:
+            if torch.backends.mps.is_available():
+                args.device = f"mps:{args.gpu_id}"
+            else:
+                raise ValueError("GPU is not available. Was torch compiled with CUDA or MPS?")
+
     # figure out how many cpus can be used
     use_cpu = True if args.gpu_id == -1 else False
     if use_cpu: