Parcourir la source

refactor: enhance model downloading with logging and new entry points

myhloli il y a 5 mois
Parent
commit
20790663b1
3 fichiers modifiés avec 7 ajouts et 0 suppressions
  1. 1 0
      mineru/cli/models_download.py
  2. 4 0
      mineru/model/vlm_sglang_model/server.py
  3. 2 0
      setup.py

+ 1 - 0
mineru/cli/models_download.py

@@ -85,6 +85,7 @@ def download_models():
         download_finish_path = ""
         if model_type == 'pipeline':
             for model_path in [ModelPath.doclayout_yolo, ModelPath.yolo_v8_mfd, ModelPath.unimernet_small, ModelPath.pytorch_paddle, ModelPath.layout_reader, ModelPath.slanet_plus]:
+                click.echo(f"Downloading model: {model_path}")
                 download_finish_path = auto_download_and_get_model_root_path(model_path, repo_mode=model_type)
         elif model_type == 'vlm':
             download_finish_path = auto_download_and_get_model_root_path("/", repo_mode=model_type)

+ 4 - 0
mineru/model/vlm_sglang_model/server.py

@@ -7,6 +7,7 @@ from sglang.srt.managers.io_struct import GenerateReqInput
 from sglang.srt.server_args import prepare_server_args
 from sglang.srt.utils import kill_process_tree
 
+from mineru.utils.models_download_utils import auto_download_and_get_model_root_path
 from .logit_processor import Mineru2LogitProcessor
 
 _custom_logit_processor_str = Mineru2LogitProcessor().to_str()
@@ -33,6 +34,9 @@ def main():
 
     server_args.enable_custom_logit_processor = True
 
+    if server_args.model_path is None:
+        server_args.model_path = auto_download_and_get_model_root_path("/","vlm")
+
     try:
         launch_server(server_args)
     finally:

+ 2 - 0
setup.py

@@ -85,6 +85,8 @@ if __name__ == '__main__':
         entry_points={
             "console_scripts": [
                 "mineru = mineru.cli:client.main",  # 命令行入口点,mineru命令将调用mineru.cli.client.main函数
+                "mineru-sglang-server = mineru.cli.vlm-sglang_server:main",  # sglang服务器入口点
+                "mineru-models-download = mineru.cli.models_download:download_models",  # 模型下载入口点
             ],
         },  # 项目提供的可执行命令
         include_package_data=True,  # 是否包含非代码文件,如数据文件、配置文件等