diff --git a/unstructured_inference/models/chipper.py b/unstructured_inference/models/chipper.py index 9128897c..06aa64e7 100644 --- a/unstructured_inference/models/chipper.py +++ b/unstructured_inference/models/chipper.py @@ -9,7 +9,6 @@ import torch import transformers from cv2.typing import MatLike -from huggingface_hub import hf_hub_download from PIL.Image import Image from transformers import DonutProcessor, VisionEncoderDecoderModel from transformers.generation.logits_process import LogitsProcessor @@ -22,7 +21,7 @@ from unstructured_inference.models.unstructuredmodel import ( UnstructuredElementExtractionModel, ) -from unstructured_inference.utils import LazyDict, strip_tags +from unstructured_inference.utils import LazyDict, strip_tags, download_if_needed_and_get_local_path MODEL_TYPES: Dict[str, Union[LazyDict, dict]] = { "chipperv1": { @@ -115,7 +114,7 @@ def initialize( token=auth_token, ) if swap_head: - lm_head_file = hf_hub_download( + lm_head_file = download_if_needed_and_get_local_path( repo_id=pre_trained_model_repo, filename="lm_head.pth", token=auth_token, diff --git a/unstructured_inference/models/detectron2.py b/unstructured_inference/models/detectron2.py index c38f6848..520298ce 100644 --- a/unstructured_inference/models/detectron2.py +++ b/unstructured_inference/models/detectron2.py @@ -3,7 +3,6 @@ from pathlib import Path from typing import Any, Dict, Final, List, Optional, Union -from huggingface_hub import hf_hub_download from layoutparser.models.detectron2.layoutmodel import ( Detectron2LayoutModel, is_detectron2_available, @@ -17,7 +16,11 @@ from unstructured_inference.models.unstructuredmodel import ( UnstructuredObjectDetectionModel, ) -from unstructured_inference.utils import LazyDict, LazyEvaluateInfo +from unstructured_inference.utils import ( + LazyDict, + LazyEvaluateInfo, + download_if_needed_and_get_local_path, +) DETECTRON_CONFIG: Final = "lp://PubLayNet/faster_rcnn_R_50_FPN_3x/config" DEFAULT_LABEL_MAP: Final[Dict[int, str]] = { @@ -35,12 +38,12 @@ MODEL_TYPES = { "detectron2_lp": LazyDict( model_path=LazyEvaluateInfo( - hf_hub_download, + download_if_needed_and_get_local_path, "layoutparser/detectron2", "PubLayNet/faster_rcnn_R_50_FPN_3x/model_final.pth", ), config_path=LazyEvaluateInfo( - hf_hub_download, + download_if_needed_and_get_local_path, "layoutparser/detectron2", "PubLayNet/faster_rcnn_R_50_FPN_3x/config.yml", ), @@ -49,12 +52,12 @@ ), "checkbox": LazyDict( model_path=LazyEvaluateInfo( - hf_hub_download, + download_if_needed_and_get_local_path, "unstructuredio/oer-checkbox", "detectron2_finetuned_oer_checkbox.pth", ), config_path=LazyEvaluateInfo( - hf_hub_download, + download_if_needed_and_get_local_path, "unstructuredio/oer-checkbox", "detectron2_oer_checkbox.json", ), diff --git a/unstructured_inference/models/detectron2onnx.py b/unstructured_inference/models/detectron2onnx.py index 3def8ced..79cd0a1a 100644 --- a/unstructured_inference/models/detectron2onnx.py +++ b/unstructured_inference/models/detectron2onnx.py @@ -4,7 +4,6 @@ import cv2 import numpy as np import onnxruntime -from huggingface_hub import hf_hub_download from huggingface_hub.constants import HUGGINGFACE_HUB_CACHE from onnxruntime.capi import _pybind_state as C from onnxruntime.quantization import QuantType, quantize_dynamic @@ -16,7 +15,11 @@ from unstructured_inference.models.unstructuredmodel import ( UnstructuredObjectDetectionModel, ) -from unstructured_inference.utils import LazyDict, LazyEvaluateInfo +from unstructured_inference.utils import ( + LazyDict, + LazyEvaluateInfo, + download_if_needed_and_get_local_path, +) onnxruntime.set_default_logger_severity(logger_onnx.getEffectiveLevel()) @@ -34,7 +37,7 @@ MODEL_TYPES: Dict[str, Union[LazyDict, dict]] = { "detectron2_onnx": LazyDict( model_path=LazyEvaluateInfo( - hf_hub_download, + download_if_needed_and_get_local_path, "unstructuredio/detectron2_faster_rcnn_R_50_FPN_3x", "model.onnx", ), @@ -52,7 +55,7 @@ }, "detectron2_mask_rcnn": LazyDict( model_path=LazyEvaluateInfo( - hf_hub_download, + download_if_needed_and_get_local_path, "unstructuredio/detectron2_mask_rcnn_X_101_32x8d_FPN_3x", "model.onnx", ), diff --git a/unstructured_inference/models/yolox.py b/unstructured_inference/models/yolox.py index 852e15b3..0acd93f3 100644 --- a/unstructured_inference/models/yolox.py +++ b/unstructured_inference/models/yolox.py @@ -8,14 +8,17 @@ import cv2 import numpy as np import onnxruntime -from huggingface_hub import hf_hub_download from onnxruntime.capi import _pybind_state as C from PIL import Image as PILImage from unstructured_inference.constants import ElementType, Source from unstructured_inference.inference.layoutelement import LayoutElement from unstructured_inference.models.unstructuredmodel import UnstructuredObjectDetectionModel -from unstructured_inference.utils import LazyDict, LazyEvaluateInfo +from unstructured_inference.utils import ( + LazyDict, + LazyEvaluateInfo, + download_if_needed_and_get_local_path, +) YOLOX_LABEL_MAP = { 0: ElementType.CAPTION, @@ -34,7 +37,7 @@ MODEL_TYPES = { "yolox": LazyDict( model_path=LazyEvaluateInfo( - hf_hub_download, + download_if_needed_and_get_local_path, "unstructuredio/yolo_x_layout", "yolox_l0.05.onnx", ), @@ -42,7 +45,7 @@ ), "yolox_tiny": LazyDict( model_path=LazyEvaluateInfo( - hf_hub_download, + download_if_needed_and_get_local_path, "unstructuredio/yolo_x_layout", "yolox_tiny.onnx", ), @@ -50,7 +53,7 @@ ), "yolox_quantized": LazyDict( model_path=LazyEvaluateInfo( - hf_hub_download, + download_if_needed_and_get_local_path, "unstructuredio/yolo_x_layout", "yolox_l0.05_quantized.onnx", ),