fix: Do proper check to set the device in EasyOCR, RapidOCR.

Signed-off-by: Nikos Livathinos <[email protected]>
DS4SD · Dec 13, 2024 · 30dbab5 · 30dbab5
1 parent bb1774d
commit 30dbab5
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 5 deletions.
diff --git a/docling/models/easyocr_model.py b/docling/models/easyocr_model.py
@@ -45,10 +45,13 @@ def __init__(
             if self.options.use_gpu:
                 device = decide_device(accelerator_options.device)
                 # Enable easyocr GPU if running on CUDA, MPS
-                use_gpu = device in [
-                    AcceleratorDevice.CUDA,
-                    AcceleratorDevice.MPS,
-                ]
+                use_gpu = any(
+                    filter(
+                        lambda x: str(x).lower() in device,
+                        [AcceleratorDevice.CUDA.value, AcceleratorDevice.MPS.value],
+                    )
+                )
+
             self.reader = easyocr.Reader(
                 lang_list=self.options.lang,
                 gpu=use_gpu,

diff --git a/docling/models/rapid_ocr_model.py b/docling/models/rapid_ocr_model.py
@@ -42,7 +42,7 @@ def __init__(
 
             # Decide the accelerator devices
             device = decide_device(accelerator_options.device)
-            use_cuda = device == AcceleratorDevice.CUDA
+            use_cuda = str(AcceleratorDevice.CUDA.value).lower() in device
             use_dml = accelerator_options.device == AcceleratorDevice.AUTO
             intra_op_num_threads = accelerator_options.num_threads