feat: add experimental MPS support to PyTorch backend. (#639)

python/rapidocr/config.yaml

@@ -92,6 +92,8 @@ EngineConfig:
         npu_ep_cfg:
             device_id: 0
 
+        use_mps: false
+
     tensorrt:
         device_id: 0
         use_fp16: true

python/rapidocr/inference_engine/pytorch/device_config.py

@@ -17,6 +17,7 @@ model_dir = root_dir / "models"
 class DeviceConfig:
     def __init__(self, cfg):
         self.use_cuda = cfg.engine_cfg.use_cuda
+        self.use_mps = cfg.engine_cfg.use_mps
         self.use_npu = cfg.engine_cfg.use_npu
 
         self.cfg = cfg
@@ -26,6 +27,9 @@ class DeviceConfig:
             device_id = self.cfg.engine_cfg.cuda_ep_cfg.device_id
             return self.get_device(DeviceType.CUDA, device_id)
 
+        if self.use_mps:
+            return self.get_device(DeviceType.MPS)
+
         if self.use_npu:
             device_id = self.cfg.engine_cfg.npu_ep_cfg.device_id
             return self.get_device(DeviceType.NPU, device_id)
@@ -41,6 +45,9 @@ class DeviceConfig:
         if device_type == DeviceType.CUDA:
             return self.config_cuda(device_id)
 
+        if device_type == DeviceType.MPS:
+            return self.config_mps()
+
         if device_type == DeviceType.NPU:
             return self.config_npu(device_id)
 
@@ -57,6 +64,13 @@ class DeviceConfig:
         logger.info(f"Using GPU device with ID: {device_id}")
         return torch.device(f"cuda:{device_id}")
 
+    def config_mps(self) -> torch.device:
+        if not torch.backends.mps.is_available():
+            raise DeviceConfigError("MPS is not available.")
+
+        logger.info("Using MPS device")
+        return torch.device("mps")
+
     def config_npu(self, device_id: int) -> torch.device:
         try:
             import torch_npu

python/rapidocr/inference_engine/pytorch/main.py

@@ -21,6 +21,9 @@ class TorchInferSession(InferSession):
             inp = torch.from_numpy(img)
             inp = inp.to(self.device)
             outputs = self.predictor(inp).cpu().numpy()
+            if self.device.type=="mps":
+                del inp
+                torch.mps.empty_cache()
             return outputs
 
     def have_key(self, key: str = "character") -> bool:

python/rapidocr/utils/typings.py

@@ -12,6 +12,7 @@ class DeviceType(Enum):
     MLU = "mlu"
     DCU = "dcu"
     GCU = "gcu"
+    MPS = "mps"
 
 
 class LangDet(Enum):