Browse Source

fix: add backend argument handling and logging for lmdeploy backend configuration

myhloli 2 weeks ago
parent
commit
a2136c22a5

+ 0 - 1
mineru/backend/vlm/utils.py

@@ -66,7 +66,6 @@ def set_lmdeploy_backend(device_type:str) -> str:
             raise ValueError("Unsupported operating system.")
     else:
         raise ValueError(f"Unsupported device type: {device_type}")
-    logger.info(f"Set lmdeploy_backend to: {lmdeploy_backend}")
     return lmdeploy_backend
 
 

+ 1 - 0
mineru/backend/vlm/vlm_analyze.py

@@ -138,6 +138,7 @@ class ModelSingleton:
                             device_type = device_type.split(":")[0]
 
                         lm_backend = set_lmdeploy_backend(device_type)
+                        logger.info(f"Set lmdeploy_backend to: {lm_backend}")
 
                         if lm_backend == "pytorch":
                             kwargs["device_type"] = device_type

+ 17 - 3
mineru/model/vlm/lmdeploy_server.py

@@ -1,6 +1,8 @@
 import os
 import sys
 
+from loguru import logger
+
 from mineru.backend.vlm.utils import set_lmdeploy_backend
 from mineru.utils.models_download_utils import auto_download_and_get_model_root_path
 
@@ -11,7 +13,9 @@ def main():
     has_port_arg = False
     has_gpu_memory_utilization_arg = False
     has_log_level_arg = False
+    has_backend_arg = False
     device_type = "cuda"
+    lm_backend = ""
 
     # 检查现有参数
     for i, arg in enumerate(args):
@@ -21,10 +25,17 @@ def main():
             has_gpu_memory_utilization_arg = True
         if arg == "--log-level" or arg.startswith("--log-level="):
             has_log_level_arg = True
+        if arg == "--backend":
+            has_backend_arg = True
+            if i + 1 < len(args):
+                lm_backend = args[i + 1]
+        if arg.startswith("--backend="):
+            has_backend_arg = True
+            lm_backend = arg.split("=", 1)[1]
         if arg == "--device":
             if i + 1 < len(args):
                 device_type = args[i + 1]
-        elif arg.startswith("--device="):
+        if arg.startswith("--device="):
             device_type = arg.split("=", 1)[1]
 
     # 添加默认参数
@@ -37,9 +48,12 @@ def main():
 
     if ":" in device_type:
         device_type = device_type.split(":")[0]
-    lm_backend = set_lmdeploy_backend(device_type)
+    if lm_backend == "":
+        lm_backend = set_lmdeploy_backend(device_type)
+    logger.info(f"Set lmdeploy_backend to: {lm_backend}")
+
     # args中如果有--backend参数,则不设置
-    if not any(arg == "--backend" or arg.startswith("--backend=") for arg in args):
+    if not has_backend_arg:
         args.extend(["--backend", lm_backend])
 
     model_path = auto_download_and_get_model_root_path("/", "vlm")