From 99ccb69a47aa544bb96eefb81765dc5986363e46 Mon Sep 17 00:00:00 2001 From: Nikos Livathinos Date: Tue, 10 Dec 2024 14:46:21 +0000 Subject: [PATCH] fix: Do proper check to set the device in EasyOCR, RapidOCR. Signed-off-by: Nikos Livathinos Signed-off-by: Christoph Auer --- docling/models/easyocr_model.py | 11 +++++++---- docling/models/rapid_ocr_model.py | 2 +- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/docling/models/easyocr_model.py b/docling/models/easyocr_model.py index 7ad79b96..68beb2fa 100644 --- a/docling/models/easyocr_model.py +++ b/docling/models/easyocr_model.py @@ -45,10 +45,13 @@ class EasyOcrModel(BaseOcrModel): 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) in device, + [AcceleratorDevice.CUDA, AcceleratorDevice.MPS], + ) + ) + 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 index f476b067..7cf8c7d3 100644 --- a/docling/models/rapid_ocr_model.py +++ b/docling/models/rapid_ocr_model.py @@ -42,7 +42,7 @@ class RapidOcrModel(BaseOcrModel): # Decide the accelerator devices device = decide_device(accelerator_options.device) - use_cuda = device == AcceleratorDevice.CUDA + use_cuda = AcceleratorDevice.CUDA in device use_dml = accelerator_options.device == AcceleratorDevice.AUTO intra_op_num_threads = accelerator_options.num_threads