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: enable use_cuda_flash_attention2 for PictureDescriptionVlmModel

Browse Source 
Signed-off-by: Zach Cox <zach.s.cox@gmail.com>
This commit is contained in:
Zach Cox 2025-04-29 19:33:58 -04:00
parent 976e92e289
commit 0961cda5fb
1 changed files with 4 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
docling/models/picture_description_vlm_model.py
View File

@ -57,7 +57,10 @@ class PictureDescriptionVlmModel(PictureDescriptionBaseModel):
artifacts_path,
torch_dtype=torch.bfloat16,
_attn_implementation=(
"flash_attention_2" if self.device.startswith("cuda") else "eager"
"flash_attention_2"
if self.device.startswith("cuda")
and accelerator_options.cuda_use_flash_attention2
else "eager"
),
).to(self.device)