|
|
@@ -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")
|