diff --git a/docling/models/rapid_ocr_model.py b/docling/models/rapid_ocr_model.py index 24c3ff73..2dd5feb4 100644 --- a/docling/models/rapid_ocr_model.py +++ b/docling/models/rapid_ocr_model.py @@ -110,6 +110,9 @@ class RapidOcrModel(BaseOcrModel): use_cuda = str(AcceleratorDevice.CUDA.value).lower() in device use_dml = accelerator_options.device == AcceleratorDevice.AUTO intra_op_num_threads = accelerator_options.num_threads + gpu_id = 0 + if use_cuda and ":" in device: + gpu_id = int(device.split(":")[1]) _ALIASES = { "onnxruntime": EngineType.ONNXRUNTIME, "openvino": EngineType.OPENVINO, @@ -184,6 +187,10 @@ class RapidOcrModel(BaseOcrModel): "Det.engine_type": backend_enum, "Cls.engine_type": backend_enum, "Rec.engine_type": backend_enum, + "EngineConfig.paddle.use_cuda": use_cuda, + "EngineConfig.paddle.gpu_id": gpu_id, + "EngineConfig.torch.use_cuda": use_cuda, + "EngineConfig.torch.gpu_id": gpu_id, } if self.options.rec_font_path is not None: diff --git a/docs/usage/gpu.md b/docs/usage/gpu.md index 7ca1576e..012f6568 100644 --- a/docs/usage/gpu.md +++ b/docs/usage/gpu.md @@ -43,6 +43,21 @@ Setting a higher `page_batch_size` will run the Docling models (in particular th For a complete example see [gpu_standard_pipeline.py](../examples/gpu_standard_pipeline.py). +#### OCR engines + +The current Docling OCR engines rely on third-party libraries, hence GPU support depends on the availability in the respective engines. + +The only setup which is known to work at the moment is RapidOCR with the torch backend, which can be enabled via + +```py +pipeline_options = PdfPipelineOptions() +pipeline_options.ocr_options = RapidOcrOptions( + backend="torch", +) +``` + +More details in the GitHub discussion [#2451](https://github.com/docling-project/docling/discussions/2451). + ### VLM Pipeline