Browse Source

feat: add lmdeploy backend support and refactor related components

myhloli 2 weeks ago
parent
commit
85558061ff

+ 36 - 23
mineru/backend/vlm/vlm_analyze.py

@@ -49,7 +49,7 @@ class ModelSingleton:
             for param in ["batch_size", "max_concurrency", "http_timeout"]:
                 if param in kwargs:
                     del kwargs[param]
-            if backend in ['transformers', 'vllm-engine', "vllm-async-engine", "mlx-engine", "lmdeploy-engine", "lmdeploy-async-engine"] and not model_path:
+            if backend in ['transformers', 'vllm-engine', "vllm-async-engine", "mlx-engine", "lmdeploy-engine"] and not model_path:
                 model_path = auto_download_and_get_model_root_path("/","vlm")
                 if backend == "transformers":
                     try:
@@ -86,28 +86,6 @@ class ModelSingleton:
                     except ImportError:
                         raise ImportError("Please install mlx-vlm to use the mlx-engine backend.")
                     model, processor = mlx_load(model_path)
-                elif backend == "lmdeploy-engine":
-                        try:
-                            from lmdeploy.serve.vl_async_engine import VLAsyncEngine
-                            from lmdeploy import PytorchEngineConfig, GenerationConfig
-                        except ImportError:
-                            raise ImportError("Please install vllm to use the vllm-engine backend.")
-                        lmdeploy_engine = VLAsyncEngine(model_path, backend='pytorch',
-                                     backend_config=PytorchEngineConfig(tp=1, block_size=128,
-                                     cache_max_entry_count=0.8, max_batch_size=256,
-                                     device_type="ascend", session_len=16382))
-
-
-                elif backend == "lmdeploy-async-engine":
-                        try:
-                            from lmdeploy.serve.vl_async_engine import VLAsyncEngine
-                            from lmdeploy import PytorchEngineConfig, GenerationConfig
-                        except ImportError:
-                            raise ImportError("Please install vllm to use the vllm-async-engine backend.")
-                        lmdeploy_engine = VLAsyncEngine(model_path, backend='pytorch',
-                                     backend_config=PytorchEngineConfig(tp=1, block_size=128,
-                                     cache_max_entry_count=0.8, max_batch_size=256,
-                                     device_type="ascend", session_len=16384))
                 else:
                     if os.getenv('OMP_NUM_THREADS') is None:
                         os.environ["OMP_NUM_THREADS"] = "1"
@@ -141,6 +119,41 @@ class ModelSingleton:
                             kwargs["logits_processors"] = [MinerULogitsProcessor]
                         # 使用kwargs为 vllm初始化参数
                         vllm_async_llm = AsyncLLM.from_engine_args(AsyncEngineArgs(**kwargs))
+                    elif backend == "lmdeploy-engine":
+                        try:
+                            from lmdeploy import PytorchEngineConfig, TurbomindEngineConfig
+                            from lmdeploy.serve.vl_async_engine import VLAsyncEngine
+                        except ImportError:
+                            raise ImportError("Please install lmdeploy to use the lmdeploy-engine backend.")
+                        if "cache_max_entry_count" not in kwargs:
+                            kwargs["cache_max_entry_count"] = 0.5
+
+                        device = kwargs.get("device", "").lower()
+                        # 特定设备强制使用 pytorch backend
+                        if device in ["ascend", "maca", "camb"]:
+                            lm_backend = "pytorch"
+                            backend_config = PytorchEngineConfig(**kwargs)
+                        else:
+                            # 其他情况根据 lm_backend 参数决定,默认使用 turbomind
+                            lm_backend = kwargs.get("lm_backend", "turbomind")
+                            if lm_backend == "pytorch":
+                                backend_config = PytorchEngineConfig(**kwargs)
+                            else:
+                                lm_backend = "turbomind"  # 确保非 pytorch 时使用 turbomind
+                                backend_config = TurbomindEngineConfig(**kwargs)
+
+                        log_level = 'ERROR'
+                        from lmdeploy.utils import get_logger
+                        logger = get_logger('lmdeploy')
+                        logger.setLevel(log_level)
+                        if os.getenv('TM_LOG_LEVEL') is None:
+                            os.environ['TM_LOG_LEVEL'] = log_level
+
+                        lmdeploy_engine = VLAsyncEngine(
+                            model_path,
+                            backend=lm_backend,
+                            backend_config=backend_config,
+                        )
             self._models[key] = MinerUClient(
                 backend=backend,
                 model=model,

+ 0 - 6
mineru/cli/common.py

@@ -331,9 +331,6 @@ def do_parse(
         if backend == "vllm-async-engine":
             raise Exception("vlm-vllm-async-engine backend is not supported in sync mode, please use vlm-vllm-engine backend")
 
-        if backend == "lmdeploy-async-engine":
-            raise Exception("vlm-lmdeploy-async-engine backend is not supported in sync mode, please use vlm-lmdeploy-engine backend")
-
         os.environ['MINERU_VLM_FORMULA_ENABLE'] = str(formula_enable)
         os.environ['MINERU_VLM_TABLE_ENABLE'] = str(table_enable)
 
@@ -385,9 +382,6 @@ async def aio_do_parse(
         if backend == "vllm-engine":
             raise Exception("vlm-vllm-engine backend is not supported in async mode, please use vlm-vllm-async-engine backend")
 
-        if backend == "lmdeploy-engine":
-            raise Exception("vlm-lmdeploy-engine backend is not supported in async mode, please use vlm-lmdeploy-async-engine backend")
-
         os.environ['MINERU_VLM_FORMULA_ENABLE'] = str(formula_enable)
         os.environ['MINERU_VLM_TABLE_ENABLE'] = str(table_enable)
 

+ 4 - 4
mineru/cli/gradio_app.py

@@ -274,7 +274,7 @@ def to_pdf(file_path):
 
 # 更新界面函数
 def update_interface(backend_choice):
-    if backend_choice in ["vlm-transformers", "vlm-vllm-async-engine", "vlm-lmdeploy-async-engine", "vlm-mlx-engine"]:
+    if backend_choice in ["vlm-transformers", "vlm-vllm-async-engine", "vlm-lmdeploy-engine", "vlm-mlx-engine"]:
         return gr.update(visible=False), gr.update(visible=False)
     elif backend_choice in ["vlm-http-client"]:
         return gr.update(visible=True), gr.update(visible=False)
@@ -380,7 +380,7 @@ def main(ctx,
             from mineru.backend.vlm.vlm_analyze import ModelSingleton
             model_singleton = ModelSingleton()
             predictor = model_singleton.get_model(
-                "lmdeploy-async-engine",
+                "lmdeploy-engine",
                 None,
                 None,
                 **kwargs
@@ -402,8 +402,8 @@ def main(ctx,
                         drop_list = ["pipeline", "vlm-vllm-async-engine"]
                         preferred_option = "vlm-vllm-async-engine"
                     elif lmdeploy_engine_enable:
-                        drop_list = ["pipeline", "vlm-lmdeploy-async-engine"]
-                        preferred_option = "vlm-lmdeploy-async-engine"
+                        drop_list = ["pipeline", "vlm-lmdeploy-engine"]
+                        preferred_option = "vlm-lmdeploy-engine"
                     else:
                         drop_list = ["pipeline", "vlm-transformers", "vlm-http-client"]
                         if is_mac_os_version_supported():

+ 10 - 0
mineru/cli/vlm_server.py

@@ -0,0 +1,10 @@
+
+
+def vllm_server():
+    from mineru.model.vlm.vllm_server import main
+    main()
+
+
+def lmdeploy_server():
+    from mineru.model.vlm.lmdeploy_server import main
+    main()

+ 0 - 4
mineru/cli/vlm_vllm_server.py

@@ -1,4 +0,0 @@
-from mineru.model.vlm_vllm_model.server import main
-
-if __name__ == "__main__":
-    main()

+ 0 - 0
mineru/model/vlm_vllm_model/__init__.py → mineru/model/vlm/__init__.py


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

@@ -0,0 +1,56 @@
+import os
+import sys
+
+from mineru.utils.models_download_utils import auto_download_and_get_model_root_path
+
+
+def main():
+    args = sys.argv[1:]
+
+    has_port_arg = False
+    has_gpu_memory_utilization_arg = False
+    has_log_level_arg = False
+    has_device_arg = False
+
+    # 检查现有参数
+    for i, arg in enumerate(args):
+        if arg == "--server-port" or arg.startswith("--server-port="):
+            has_port_arg = True
+        if arg == "--cache-max-entry-count" or arg.startswith("--cache-max-entry-count="):
+            has_gpu_memory_utilization_arg = True
+        if arg == "--log-level" or arg.startswith("--log-level="):
+            has_log_level_arg = True
+        if arg == "--device":
+            if i + 1 < len(args):
+                device_type = args[i + 1]
+        elif arg.startswith("--device="):
+            device_type = arg.split("=", 1)[1]
+
+    # 添加默认参数
+    if not has_port_arg:
+        args.extend(["--server-port", "30000"])
+    if not has_gpu_memory_utilization_arg:
+        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"])
+
+    model_path = auto_download_and_get_model_root_path("/", "vlm")
+
+    # 重构参数,将模型路径作为位置参数
+    sys.argv = [sys.argv[0]] + ["serve", "api_server", model_path] + args
+
+    if os.getenv('OMP_NUM_THREADS') is None:
+        os.environ["OMP_NUM_THREADS"] = "1"
+
+    # 启动vllm服务器
+    print(f"start lmdeploy server: {sys.argv}")
+
+    # 使用os.system调用启动lmdeploy服务器
+    os.system("lmdeploy " + " ".join(sys.argv[1:]))
+
+
+if __name__ == "__main__":
+    main()

+ 0 - 0
mineru/model/vlm_vllm_model/server.py → mineru/model/vlm/vllm_server.py


+ 6 - 1
pyproject.toml

@@ -58,6 +58,10 @@ vlm = [
 vllm = [
     "vllm>=0.10.1.1,<0.12",
 ]
+lmdeploy = [
+    "lmdeploy>=0.10.2,<0.11",
+    "qwen_vl_utils>=0.0.14,<0.1",
+]
 mlx = [
     "mlx-vlm>=0.3.3,<0.4",
 ]
@@ -105,7 +109,8 @@ issues = "https://github.com/opendatalab/MinerU/issues"
 
 [project.scripts]
 mineru = "mineru.cli:client.main"
-mineru-vllm-server = "mineru.cli.vlm_vllm_server:main"
+mineru-vllm-server = "mineru.cli.vlm_server:vllm_server"
+mineru-lmdeploy-server = "mineru.cli.vlm_server:lmdeploy_server"
 mineru-models-download = "mineru.cli.models_download:download_models"
 mineru-api = "mineru.cli.fast_api:main"
 mineru-gradio = "mineru.cli.gradio_app:main"