mirrors/docling
1
0
Fork 0
mirror of https://github.com/DS4SD/docling.git synced 2025-07-26 20:14:47 +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 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}"