mirrors/docling
1
0
Fork 0
mirror of https://github.com/DS4SD/docling.git synced 2025-07-27 04:24:45 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: Secure torch model inits with global locks

Browse Source 
Signed-off-by: Christoph Auer <cau@zurich.ibm.com>
This commit is contained in:
Christoph Auer 2025-07-02 17:20:38 +02:00
parent 3089cf2d26
commit c0ef74d9cc
1 changed files with 16 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

27
docling/models/picture_description_vlm_model.py
View File

@ -1,3 +1,4 @@
import threading
from collections.abc import Iterable from collections.abc import Iterable
from pathlib import Path from pathlib import Path
from typing import Optional, Type, Union from typing import Optional, Type, Union
@ -15,6 +16,9 @@ from docling.models.utils.hf_model_download import (
) )
from docling.utils.accelerator_utils import decide_device from docling.utils.accelerator_utils import decide_device
# Global lock for model initialization to prevent threading issues
_model_init_lock = threading.Lock()
class PictureDescriptionVlmModel( class PictureDescriptionVlmModel(
PictureDescriptionBaseModel, HuggingFaceModelDownloadMixin PictureDescriptionBaseModel, HuggingFaceModelDownloadMixin
@ -57,17 +61,18 @@ class PictureDescriptionVlmModel(
) )
# Initialize processor and model # Initialize processor and model
self.processor = AutoProcessor.from_pretrained(artifacts_path) with _model_init_lock:
self.model = AutoModelForVision2Seq.from_pretrained( self.processor = AutoProcessor.from_pretrained(artifacts_path)
artifacts_path, self.model = AutoModelForVision2Seq.from_pretrained(
torch_dtype=torch.bfloat16, artifacts_path,
_attn_implementation=( torch_dtype=torch.bfloat16,
"flash_attention_2" _attn_implementation=(
if self.device.startswith("cuda") "flash_attention_2"
and accelerator_options.cuda_use_flash_attention2 if self.device.startswith("cuda")
else "eager" and accelerator_options.cuda_use_flash_attention2
), else "eager"
).to(self.device) ),
).to(self.device)
self.provenance = f"{self.options.repo_id}" self.provenance = f"{self.options.repo_id}"