Quellcode durchsuchen

fix: update model path handling in model.py and models_download_utils.py

myhloli vor 5 Monaten
Ursprung
Commit
fa9aaaa7b7
2 geänderte Dateien mit 7 neuen und 3 gelöschten Zeilen
  1. 1 1
      mineru/model/vlm_sglang_model/model.py
  2. 6 2
      mineru/utils/models_download_utils.py

+ 1 - 1
mineru/model/vlm_sglang_model/model.py

@@ -62,7 +62,7 @@ class Mineru2QwenForCausalLM(nn.Module):
 
         # load vision tower
         mm_vision_tower = self.config.mm_vision_tower
-        model_root_path = auto_download_and_get_model_root_path("/", "vlm")
+        model_root_path = auto_download_and_get_model_root_path(mm_vision_tower, "vlm")
         mm_vision_tower = f"{model_root_path}/{mm_vision_tower}"
 
         if "clip" in mm_vision_tower:

+ 6 - 2
mineru/utils/models_download_utils.py

@@ -57,8 +57,12 @@ def auto_download_and_get_model_root_path(relative_path: str, repo_mode='pipelin
         relative_path = relative_path.strip('/')
         cache_dir = snapshot_download(repo, allow_patterns=[relative_path, relative_path+"/*"])
     elif repo_mode == 'vlm':
-        # VLM 模式下,直接下载整个模型目录
-        cache_dir = snapshot_download(repo)
+        # VLM 模式下,根据 relative_path 的不同处理方式
+        if relative_path == "/":
+            cache_dir = snapshot_download(repo)
+        else:
+            relative_path = relative_path.strip('/')
+            cache_dir = snapshot_download(repo, allow_patterns=[relative_path, relative_path+"/*"])
 
     if not cache_dir:
         raise FileNotFoundError(f"Failed to download model: {relative_path} from {repo}")