NeoAnd/Docling
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: Secure torch model inits with global locks (#1884)

Browse Source 
Signed-off-by: Christoph Auer <cau@zurich.ibm.com>
This commit is contained in:
Christoph Auer 2025-07-04 07:27:26 +02:00 committed by GitHub
parent 13865c06f5
commit 598c9c53d4
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
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}"