From 0eb065e9b6e4619d4c412ed98bc7408915ca3f95 Mon Sep 17 00:00:00 2001 From: Nikos Livathinos <100353117+nikos-livathinos@users.noreply.github.com> Date: Fri, 8 Nov 2024 16:48:41 +0100 Subject: [PATCH] fix(EasyOcrModel): Support the use_gpu pipeline parameter in EasyOcrModel. Initialize easyocr (#282) fix(EasyOcrModel): Support the use_gpu pipeline parameter in EasyOcrModel. Initialize easyocr without GPU if MPS is available. Signed-off-by: Nikos Livathinos --- docling/models/easyocr_model.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/docling/models/easyocr_model.py b/docling/models/easyocr_model.py index 8dff0fff8..6c72be5b8 100644 --- a/docling/models/easyocr_model.py +++ b/docling/models/easyocr_model.py @@ -2,6 +2,7 @@ from typing import Iterable import numpy +import torch from docling_core.types.doc import BoundingBox, CoordOrigin from docling.datamodel.base_models import OcrCell, Page @@ -30,8 +31,12 @@ def __init__(self, enabled: bool, options: EasyOcrOptions): "Alternatively, Docling has support for other OCR engines. See the documentation." ) + use_gpu = ( + False if torch.backends.mps.is_available() else self.options.use_gpu + ) self.reader = easyocr.Reader( lang_list=self.options.lang, + gpu=use_gpu, model_storage_directory=self.options.model_storage_directory, download_enabled=self.options.download_enabled, )