Merge pull request #1436 from roboflow/feature/proper-mps-handling-for-macos

Apply changes such that MPS is correctly supported for `rfdetr`, yet still remain non-default for MacOS

Author: PawelPeczek-Roboflow

.github/workflows/integration_tests_inference_experimental_gpu.yml

@@ -2,8 +2,6 @@ name: INTEGRATION TESTS - inference_experimental (GPU)
 permissions:
   contents: read
 on:
-  pull_request:
-    branches: [main]
   push:
     branches: [main]
   workflow_dispatch:

inference_experimental/inference_exp/__init__.py

@@ -1,3 +1,8 @@
+import os
+
+if os.environ.get("PYTORCH_ENABLE_MPS_FALLBACK") is None:
+    os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
+
 from inference_exp.entities import ColorFormat
 from inference_exp.models.auto_loaders.core import AutoModel
 from inference_exp.models.base.classification import (

inference_experimental/inference_exp/models/common/onnx.py

@@ -196,9 +196,9 @@ def run_session_via_iobinding(
     )
     device = get_input_device(inputs=inputs)
     if device.type != "cuda":
-        inputs_np = {name: value.numpy() for name, value in inputs.items()}
+        inputs_np = {name: value.cpu().numpy() for name, value in inputs.items()}
         results = session.run(None, inputs_np)
-        return [torch.from_numpy(element) for element in results]
+        return [torch.from_numpy(element).to(device=device) for element in results]
     try:
         import pycuda.driver as cuda
         from inference_exp.models.common.cuda import use_primary_cuda_context

inference_experimental/tests/integration_tests/models/test_clip_predictions.py

@@ -7,10 +7,11 @@
 import torch.nn.functional as F
 from clip import clip
 from filelock import FileLock
+from inference_exp.configuration import DEFAULT_DEVICE
 from PIL.Image import Image
 from torch import nn
 
-EXPECTED_DOG_IMAGE_EMBEDDING = torch.Tensor(
+EXPECTED_DOG_IMAGE_EMBEDDING = torch.tensor(
     [
         [
             -0.031097251921892166,
@@ -1038,7 +1039,8 @@
             0.006933738477528095,
             -0.0005487799644470215,
         ]
-    ]
+    ],
+    device=DEFAULT_DEVICE,
 )
 
 

inference_experimental/tests/integration_tests/models/test_perception_encoder_predictions.py

@@ -1,8 +1,9 @@
 import numpy as np
 import pytest
 import torch
+from inference_exp.configuration import DEFAULT_DEVICE
 
-EXPECTED_DOG_IMAGE_EMBEDDING = torch.Tensor(
+EXPECTED_DOG_IMAGE_EMBEDDING = torch.tensor(
     [
         [
             -0.04075119271874428,
@@ -1030,7 +1031,8 @@
             0.027445748448371887,
             0.03301752358675003,
         ]
-    ]
+    ],
+    device=DEFAULT_DEVICE,
 )