feat: enhance OCR model initialization and add caching for snapshot downloads

mineru/backend/pipeline/model_init.py

@@ -17,6 +17,7 @@ from ...model.table.rec.unet_table.main import UnetTableModel
 from ...utils.config_reader import get_device
 from ...utils.enum_class import ModelPath
 from ...utils.models_download_utils import auto_download_and_get_model_root_path
+from ...utils.ocr_language import normalize_ocr_model_lang
 
 PIPELINE_MODEL_INIT_LOCK = threading.RLock()
 # 这些锁保护 pipeline 与 hybrid 共享的 atom model/native 模型推理调用，避免多线程同时进入同一个模型对象。
@@ -128,34 +129,20 @@ def pp_doclayout_v2_model_init(weight, device='cpu'):
     model = PPDocLayoutV2LayoutModel(weight, device)
     return model
 
+
 def ocr_model_init(det_db_box_thresh=0.5,
                    lang=None,
                    det_db_unclip_ratio=1.5,
                    enable_merge_det_boxes=True
                    ):
+    ocr_kwargs = {
+        "lang": normalize_ocr_model_lang(lang),
+        "det_db_box_thresh": det_db_box_thresh,
+        "det_db_unclip_ratio": det_db_unclip_ratio,
+        "enable_merge_det_boxes": enable_merge_det_boxes,
+    }
 
-    if lang in [None, "ch"]:
-        use_dilation = True
-        det_db_unclip_ratio = 1.8
-    else:
-        use_dilation = False
-
-    if lang is not None and lang != '':
-        model = PytorchPaddleOCR(
-            det_db_box_thresh=det_db_box_thresh,
-            lang=lang,
-            use_dilation=use_dilation,
-            det_db_unclip_ratio=det_db_unclip_ratio,
-            enable_merge_det_boxes=enable_merge_det_boxes,
-        )
-    else:
-        model = PytorchPaddleOCR(
-            det_db_box_thresh=det_db_box_thresh,
-            use_dilation=use_dilation,
-            det_db_unclip_ratio=det_db_unclip_ratio,
-            enable_merge_det_boxes=enable_merge_det_boxes,
-        )
-    return model
+    return PytorchPaddleOCR(**ocr_kwargs)
 
 
 class AtomModelSingleton:
@@ -172,17 +159,17 @@ class AtomModelSingleton:
     def get_atom_model(self, atom_model_name: str, **kwargs):
 
         lang = kwargs.get('lang', None)
-
+        ocr_singleton_lang = normalize_ocr_model_lang(lang)
         if atom_model_name in [AtomicModel.WiredTable, AtomicModel.WirelessTable]:
             key = (
                 atom_model_name,
-                lang
+                ocr_singleton_lang
             )
         elif atom_model_name in [AtomicModel.OCR]:
             key = (
                 atom_model_name,
                 kwargs.get('det_db_box_thresh', 0.5),
-                lang,
+                ocr_singleton_lang,
                 kwargs.get('det_db_unclip_ratio', 1.5),
                 kwargs.get('enable_merge_det_boxes', True)
             )

mineru/utils/models_download_utils.py

@@ -1,11 +1,36 @@
 # Copyright (c) Opendatalab. All rights reserved.
 import os
+from functools import lru_cache
+
 from huggingface_hub import snapshot_download as hf_snapshot_download
 from modelscope import snapshot_download as ms_snapshot_download
 
 from mineru.utils.config_reader import get_local_models_dir
 from mineru.utils.enum_class import ModelPath
 
+
+@lru_cache(maxsize=None)
+def _snapshot_download_cached(model_source: str, repo_mode: str, repo: str, relative_path: str) -> str:
+    """按进程缓存远端 snapshot_download 结果，减少重复缓存检查和 Fetching 日志。"""
+    if model_source == "huggingface":
+        snapshot_download = hf_snapshot_download
+    elif model_source == "modelscope":
+        snapshot_download = ms_snapshot_download
+    else:
+        raise ValueError(f"未知的仓库类型: {model_source}")
+
+    if repo_mode == 'pipeline':
+        return snapshot_download(repo, allow_patterns=[relative_path, relative_path + "/*"])
+
+    if repo_mode == 'vlm':
+        # VLM 整仓下载和局部路径下载都参与缓存，但保持原有 allow_patterns 行为。
+        if relative_path == "/":
+            return snapshot_download(repo)
+        return snapshot_download(repo, allow_patterns=[relative_path, relative_path + "/*"])
+
+    raise ValueError(f"Unsupported repo_mode: {repo_mode}, must be 'pipeline' or 'vlm'")
+
+
 def auto_download_and_get_model_root_path(relative_path: str, repo_mode='pipeline') -> str:
     """
     支持文件或目录的可靠下载。
@@ -45,25 +70,14 @@ def auto_download_and_get_model_root_path(relative_path: str, repo_mode='pipelin
     repo = repo_mapping[repo_mode].get(model_source, repo_mapping[repo_mode]['default'])
 
 
-    if model_source == "huggingface":
-        snapshot_download = hf_snapshot_download
-    elif model_source == "modelscope":
-        snapshot_download = ms_snapshot_download
-    else:
-        raise ValueError(f"未知的仓库类型: {model_source}")
-
-    cache_dir = None
-
     if repo_mode == 'pipeline':
         relative_path = relative_path.strip('/')
-        cache_dir = snapshot_download(repo, allow_patterns=[relative_path, relative_path+"/*"])
     elif repo_mode == 'vlm':
         # VLM 模式下，根据 relative_path 的不同处理方式
-        if relative_path == "/":
-            cache_dir = snapshot_download(repo)
-        else:
+        if relative_path != "/":
             relative_path = relative_path.strip('/')
-            cache_dir = snapshot_download(repo, allow_patterns=[relative_path, relative_path+"/*"])
+
+    cache_dir = _snapshot_download_cached(model_source, repo_mode, repo, relative_path)
 
     if not cache_dir:
         raise FileNotFoundError(f"Failed to download model: {relative_path} from {repo}")
@@ -73,4 +87,4 @@ def auto_download_and_get_model_root_path(relative_path: str, repo_mode='pipelin
 if __name__ == '__main__':
     path1 = "models/README.md"
     root = auto_download_and_get_model_root_path(path1)
-    print("本地文件绝对路径:", os.path.join(root, path1))
+    print("本地文件绝对路径:", os.path.join(root, path1))