Browse Source

fix: improve device type and backend handling in lmdeploy configuration

myhloli 2 weeks ago
parent
commit
39e7aa52a2

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

@@ -64,7 +64,7 @@ def set_lmdeploy_backend(device_type:str) -> str:
         else:
         else:
             raise ValueError("Unsupported operating system.")
             raise ValueError("Unsupported operating system.")
     else:
     else:
-        raise ValueError(f"Unsupported device type: {device_type}")
+        raise ValueError(f"Unsupported lmdeploy device type: {device_type}")
     return lmdeploy_backend
     return lmdeploy_backend
 
 
 
 

+ 11 - 7
mineru/backend/vlm/vlm_analyze.py

@@ -129,15 +129,19 @@ class ModelSingleton:
                         if "cache_max_entry_count" not in kwargs:
                         if "cache_max_entry_count" not in kwargs:
                             kwargs["cache_max_entry_count"] = 0.5
                             kwargs["cache_max_entry_count"] = 0.5
 
 
-                        if "device" in kwargs:
-                            device_type = kwargs.pop("device")
+                        if "lmdeploy_device" in kwargs:
+                            device_type = kwargs.pop("lmdeploy_device")
+                            if device_type not in ["cuda", "ascend", "maca", "camb"]:
+                                raise ValueError(f"Unsupported lmdeploy device type: {device_type}")
                         else:
                         else:
-                            device_type = os.getenv('MINERU_DEVICE_MODE', "cuda").lower()
-                        # device_type 如果有则去除":"
-                        if ":" in device_type:
-                            device_type = device_type.split(":")[0]
+                            device_type = "cuda"
 
 
-                        lm_backend = set_lmdeploy_backend(device_type)
+                        if "lmdeploy_backend" in kwargs:
+                            lm_backend = kwargs.pop("lmdeploy_backend")
+                            if lm_backend not in ["pytorch", "turbomind"]:
+                                raise ValueError(f"Unsupported lmdeploy backend: {lm_backend}")
+                        else:
+                            lm_backend = set_lmdeploy_backend(device_type)
                         logger.info(f"Set lmdeploy_backend to: {lm_backend}")
                         logger.info(f"Set lmdeploy_backend to: {lm_backend}")
 
 
                         if lm_backend == "pytorch":
                         if lm_backend == "pytorch":

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

@@ -14,7 +14,7 @@ def main():
     has_gpu_memory_utilization_arg = False
     has_gpu_memory_utilization_arg = False
     has_log_level_arg = False
     has_log_level_arg = False
     has_backend_arg = False
     has_backend_arg = False
-    device_type = "cuda"
+    device_type = ""
     lm_backend = ""
     lm_backend = ""
 
 
     # 检查现有参数
     # 检查现有参数
@@ -46,10 +46,14 @@ def main():
     if not has_log_level_arg:
     if not has_log_level_arg:
         args.extend(["--log-level", "ERROR"])
         args.extend(["--log-level", "ERROR"])
 
 
-    if ":" in device_type:
-        device_type = device_type.split(":")[0]
     if lm_backend == "":
     if lm_backend == "":
+        if device_type == "":
+            device_type = "cuda"
+        elif device_type not in ["cuda", "ascend", "maca", "camb"]:
+            raise ValueError(f"Unsupported lmdeploy device type: {device_type}")
         lm_backend = set_lmdeploy_backend(device_type)
         lm_backend = set_lmdeploy_backend(device_type)
+    elif lm_backend not in ["pytorch", "turbomind"]:
+        raise ValueError(f"Unsupported lmdeploy backend: {lm_backend}")
     logger.info(f"Set lmdeploy_backend to: {lm_backend}")
     logger.info(f"Set lmdeploy_backend to: {lm_backend}")
 
 
     if lm_backend == "pytorch":
     if lm_backend == "pytorch":