mirrors/docling
1
0
Fork 0
mirror of https://github.com/DS4SD/docling.git synced 2025-07-31 14:34:40 +00:00
Code Issues Packages Projects Releases Wiki Activity

use default device until we enable global management

Browse Source 
Signed-off-by: Michele Dolfi <dol@zurich.ibm.com>
This commit is contained in:
Michele Dolfi 2024-11-27 11:39:23 +01:00
parent 74e005df63
commit c1b6442670
2 changed files with 40 additions and 31 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
docling/datamodel/pipeline_options.py
View File

@ -46,13 +46,13 @@ class RapidOcrOptions(OcrOptions):
use_cls: Optional[bool] = None # same default as rapidocr
use_rec: Optional[bool] = None # same default as rapidocr
class Device(Enum):
CPU = "CPU"
CUDA = "CUDA"
DIRECTML = "DIRECTML"
AUTO = "AUTO"
# class Device(Enum):
# CPU = "CPU"
# CUDA = "CUDA"
# DIRECTML = "DIRECTML"
# AUTO = "AUTO"
device: Device = Device.AUTO # Default value is AUTO
# device: Device = Device.AUTO # Default value is AUTO
print_verbose: bool = False # same default as rapidocr

59
docling/models/rapid_ocr_model.py
View File

@ -30,34 +30,43 @@ class RapidOcrModel(BaseOcrModel):
"Alternatively, Docling has support for other OCR engines. See the documentation."
)
# Same as Defaults in RapidOCR
cls_use_cuda = False
rec_use_cuda = False
det_use_cuda = False
det_use_dml = False
cls_use_dml = False
rec_use_dml = False
# This configuration option will be revamped while introducing device settings for all models.
# For the moment we will default to auto and let onnx-runtime pick the best.
cls_use_cuda = True
rec_use_cuda = True
det_use_cuda = True
det_use_dml = True
cls_use_dml = True
rec_use_dml = True
# If we set everything to true onnx-runtime would automatically choose the fastest accelerator
if self.options.device == self.options.Device.AUTO:
cls_use_cuda = True
rec_use_cuda = True
det_use_cuda = True
det_use_dml = True
cls_use_dml = True
rec_use_dml = True
# # Same as Defaults in RapidOCR
# cls_use_cuda = False
# rec_use_cuda = False
# det_use_cuda = False
# det_use_dml = False
# cls_use_dml = False
# rec_use_dml = False
# If we set use_cuda to true onnx would use the cuda device available in runtime if no cuda device is available it would run on CPU.
elif self.options.device == self.options.Device.CUDA:
cls_use_cuda = True
rec_use_cuda = True
det_use_cuda = True
# # If we set everything to true onnx-runtime would automatically choose the fastest accelerator
# if self.options.device == self.options.Device.AUTO:
# cls_use_cuda = True
# rec_use_cuda = True
# det_use_cuda = True
# det_use_dml = True
# cls_use_dml = True
# rec_use_dml = True
# If we set use_dml to true onnx would use the dml device available in runtime if no dml device is available it would work on CPU.
elif self.options.device == self.options.Device.DIRECTML:
det_use_dml = True
cls_use_dml = True
rec_use_dml = True
# # If we set use_cuda to true onnx would use the cuda device available in runtime if no cuda device is available it would run on CPU.
# elif self.options.device == self.options.Device.CUDA:
# cls_use_cuda = True
# rec_use_cuda = True
# det_use_cuda = True
# # If we set use_dml to true onnx would use the dml device available in runtime if no dml device is available it would work on CPU.
# elif self.options.device == self.options.Device.DIRECTML:
# det_use_dml = True
# cls_use_dml = True
# rec_use_dml = True
self.reader = RapidOCR(
text_score=self.options.text_score,