Browse Source

fix: import MinerULogitsProcessor conditionally for vllm-engine and vllm-async-engine backends

myhloli 5 days ago
parent
commit
64586d03ea
1 changed files with 2 additions and 2 deletions
  1. 2 2
      mineru/backend/vlm/vlm_analyze.py

+ 2 - 2
mineru/backend/vlm/vlm_analyze.py

@@ -97,7 +97,6 @@ class ModelSingleton:
                 if backend == "vllm-engine":
                     try:
                         import vllm
-                        from mineru_vl_utils import MinerULogitsProcessor
                     except ImportError:
                         raise ImportError("Please install vllm to use the vllm-engine backend.")
                     if "gpu_memory_utilization" not in kwargs:
@@ -105,6 +104,7 @@ class ModelSingleton:
                     if "model" not in kwargs:
                         kwargs["model"] = model_path
                     if enable_custom_logits_processors() and ("logits_processors" not in kwargs):
+                        from mineru_vl_utils import MinerULogitsProcessor
                         kwargs["logits_processors"] = [MinerULogitsProcessor]
                     # 使用kwargs为 vllm初始化参数
                     vllm_llm = vllm.LLM(**kwargs)
@@ -112,7 +112,6 @@ class ModelSingleton:
                     try:
                         from vllm.engine.arg_utils import AsyncEngineArgs
                         from vllm.v1.engine.async_llm import AsyncLLM
-                        from mineru_vl_utils import MinerULogitsProcessor
                     except ImportError:
                         raise ImportError("Please install vllm to use the vllm-async-engine backend.")
                     if "gpu_memory_utilization" not in kwargs:
@@ -120,6 +119,7 @@ class ModelSingleton:
                     if "model" not in kwargs:
                         kwargs["model"] = model_path
                     if enable_custom_logits_processors() and ("logits_processors" not in kwargs):
+                        from mineru_vl_utils import MinerULogitsProcessor
                         kwargs["logits_processors"] = [MinerULogitsProcessor]
                     # 使用kwargs为 vllm初始化参数
                     vllm_async_llm = AsyncLLM.from_engine_args(AsyncEngineArgs(**kwargs))