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 pathlib import Path
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
# Global lock for model initialization to prevent threading issues
_model_init_lock = threading.Lock()
class PictureDescriptionVlmModel(
PictureDescriptionBaseModel, HuggingFaceModelDownloadMixin
@ -57,17 +61,18 @@ class PictureDescriptionVlmModel(
)
# Initialize processor and model
self.processor = AutoProcessor.from_pretrained(artifacts_path)
self.model = AutoModelForVision2Seq.from_pretrained(
artifacts_path,
torch_dtype=torch.bfloat16,
_attn_implementation=(
"flash_attention_2"
if self.device.startswith("cuda")
and accelerator_options.cuda_use_flash_attention2
else "eager"
),
).to(self.device)
with _model_init_lock:
self.processor = AutoProcessor.from_pretrained(artifacts_path)
self.model = AutoModelForVision2Seq.from_pretrained(
artifacts_path,
torch_dtype=torch.bfloat16,
_attn_implementation=(
"flash_attention_2"
if self.device.startswith("cuda")
and accelerator_options.cuda_use_flash_attention2
else "eager"
),
).to(self.device)
self.provenance = f"{self.options.repo_id}"