fix: add Linux environment detection and set lmdeploy backend based on device type

mineru/backend/vlm/utils.py

@@ -3,6 +3,7 @@ import os
 from loguru import logger
 from packaging import version
 
+from mineru.utils.check_sys_env import is_windows_environment, is_linux_environment
 from mineru.utils.config_reader import get_device
 from mineru.utils.model_utils import get_vram
 
@@ -44,6 +45,33 @@ def enable_custom_logits_processors() -> bool:
         return True
 
 
+def set_lmdeploy_backend(device_type:str) -> str:
+    lmdeploy_backend = ""
+    if device_type.lower() in ["ascend", "maca", "camb"]:
+        lmdeploy_backend = "pytorch"
+    elif device_type.lower() in ["cuda"]:
+        import torch
+        if not torch.cuda.is_available():
+            raise ValueError("CUDA is not available.")
+        if is_windows_environment():
+            lmdeploy_backend = "turbomind"
+        elif is_linux_environment():
+            major, minor = torch.cuda.get_device_capability()
+            compute_capability = f"{major}.{minor}"
+            if version.parse(compute_capability) >= version.parse("8.0"):
+                lmdeploy_backend = "pytorch"
+            else:
+                lmdeploy_backend = "turbomind"
+        else:
+            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
+
+
+
+
 def set_default_gpu_memory_utilization() -> float:
     from vllm import __version__ as vllm_version
     if version.parse(vllm_version) >= version.parse("0.11.0"):

mineru/backend/vlm/vlm_analyze.py

@@ -4,7 +4,8 @@ import time
 
 from loguru import logger
 
-from .utils import enable_custom_logits_processors, set_default_gpu_memory_utilization, set_default_batch_size
+from .utils import enable_custom_logits_processors, set_default_gpu_memory_utilization, set_default_batch_size, \
+    set_lmdeploy_backend
 from .model_output_to_middle_json import result_to_middle_json
 from ...data.data_reader_writer import DataWriter
 from mineru.utils.pdf_image_tools import load_images_from_pdf
@@ -128,15 +129,23 @@ class ModelSingleton:
                         if "cache_max_entry_count" not in kwargs:
                             kwargs["cache_max_entry_count"] = 0.5
 
-                        # 默认使用 turbomind
-                        lm_backend = "turbomind"
-                        device = kwargs.get("device", "cuda").lower()
-                        # 特定设备强制使用 pytorch backend
-                        if device in ["ascend", "maca", "camb"]:
-                            lm_backend = "pytorch"
-                            backend_config = PytorchEngineConfig(**kwargs)
+                        if "device" in kwargs:
+                            device_type = kwargs.pop("device")
                         else:
+                            device_type = os.getenv('MINERU_DEVICE_MODE', "cuda").lower()
+                        # device_type 如果有则去除":"
+                        if ":" in device_type:
+                            device_type = device_type.split(":")[0]
+
+                        lm_backend = set_lmdeploy_backend(device_type)
+
+                        if lm_backend == "pytorch":
+                            kwargs["device_type"] = device_type
+                            backend_config = PytorchEngineConfig(**kwargs)
+                        elif lm_backend == "turbomind":
                             backend_config = TurbomindEngineConfig(**kwargs)
+                        else:
+                            raise ValueError(f"Unsupported lmdeploy backend: {lm_backend}")
 
                         log_level = 'ERROR'
                         from lmdeploy.utils import get_logger

mineru/model/vlm/lmdeploy_server.py

@@ -1,6 +1,7 @@
 import os
 import sys
 
+from mineru.backend.vlm.utils import set_lmdeploy_backend
 from mineru.utils.models_download_utils import auto_download_and_get_model_root_path
 
 
@@ -10,7 +11,6 @@ def main():
     has_port_arg = False
     has_gpu_memory_utilization_arg = False
     has_log_level_arg = False
-    has_device_arg = False
     device_type = "cuda"
 
     # 检查现有参数
@@ -22,11 +22,9 @@ def main():
         if arg == "--log-level" or arg.startswith("--log-level="):
             has_log_level_arg = True
         if arg == "--device":
-            has_device_arg = True
             if i + 1 < len(args):
                 device_type = args[i + 1]
         elif arg.startswith("--device="):
-            has_device_arg = True
             device_type = arg.split("=", 1)[1]
 
     # 添加默认参数
@@ -36,9 +34,13 @@ def main():
         args.extend(["--cache-max-entry-count", "0.5"])
     if not has_log_level_arg:
         args.extend(["--log-level", "ERROR"])
-    if has_device_arg:
-        if device_type.lower() in ["ascend", "maca", "camb"]:
-            args.extend(["--backend", "pytorch"])
+
+    if ":" in device_type:
+        device_type = device_type.split(":")[0]
+    lm_backend = set_lmdeploy_backend(device_type)
+    # args中如果有--backend参数，则不设置
+    if not any(arg == "--backend" or arg.startswith("--backend=") for arg in args):
+        args.extend(["--backend", lm_backend])
 
     model_path = auto_download_and_get_model_root_path("/", "vlm")
 

mineru/utils/check_sys_env.py

@@ -13,6 +13,10 @@ def is_mac_environment() -> bool:
     return platform.system() == "Darwin"
 
 
+def is_linux_environment() -> bool:
+    return platform.system() == "Linux"
+
+
 # Detect if CPU is Apple Silicon architecture
 def is_apple_silicon_cpu() -> bool:
     return platform.machine() in ["arm64", "aarch64"]