Browse Source

fix: enhance device and backend configuration handling in lmdeploy and vlm modules

myhloli 1 week ago
parent
commit
cf1fbd2923
3 changed files with 36 additions and 14 deletions
  1. 16 13
      mineru/backend/vlm/vlm_analyze.py
  2. 18 1
      mineru/cli/vlm_server.py
  3. 2 0
      mineru/model/vlm/lmdeploy_server.py

+ 16 - 13
mineru/backend/vlm/vlm_analyze.py

@@ -132,19 +132,22 @@ class ModelSingleton:
                     if "cache_max_entry_count" not in kwargs:
                         kwargs["cache_max_entry_count"] = 0.5
 
-                    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:
-                        device_type = "cuda"
-
-                    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)
+                    device_type = os.getenv("MINERU_LMDEPLOY_DEVICE", "")
+                    if device_type == "":
+                        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:
+                            device_type = "cuda"
+                    lm_backend = os.getenv("MINERU_LMDEPLOY_BACKEND", "")
+                    if lm_backend == "":
+                        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"lmdeploy device is: {device_type}, lmdeploy backend is: {lm_backend}")
 
                     if lm_backend == "pytorch":

+ 18 - 1
mineru/cli/vlm_server.py

@@ -32,11 +32,28 @@ def openai_server(ctx, inference_engine):
             inference_engine = 'vllm'
             logger.info("Using vLLM as the inference engine for VLM server.")
         except ImportError:
-            inference_engine = 'lmdeploy'
             logger.info("vLLM not found, falling back to LMDeploy as the inference engine for VLM server.")
+            try:
+                import lmdeploy
+                inference_engine = 'lmdeploy'
+                logger.info("Using LMDeploy as the inference engine for VLM server.")
+            except ImportError:
+                logger.error("Neither vLLM nor LMDeploy is installed. Please install at least one of them.")
+                sys.exit(1)
+
     if inference_engine == 'vllm':
+        try:
+            import vllm
+        except ImportError:
+            logger.error("vLLM is not installed. Please install vLLM or choose LMDeploy as the inference engine.")
+            sys.exit(1)
         vllm_server()
     elif inference_engine == 'lmdeploy':
+        try:
+            import lmdeploy
+        except ImportError:
+            logger.error("LMDeploy is not installed. Please install LMDeploy or choose vLLM as the inference engine.")
+            sys.exit(1)
         lmdeploy_server()
 
 if __name__ == "__main__":

+ 2 - 0
mineru/model/vlm/lmdeploy_server.py

@@ -53,10 +53,12 @@ def main():
     if not has_log_level_arg:
         args.extend(["--log-level", "ERROR"])
 
+    device_type = os.getenv("MINERU_LMDEPLOY_DEVICE", device_type)
     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 = os.getenv("MINERU_LMDEPLOY_BACKEND", lm_backend)
     if lm_backend == "":
         lm_backend = set_lmdeploy_backend(device_type)
     elif lm_backend not in ["pytorch", "turbomind"]: