refactor: enhance model source handling and improve environment variable setup

mineru/backend/vlm/vlm_analyze.py

@@ -8,6 +8,7 @@ from mineru.utils.pdf_image_tools import load_images_from_pdf
 from .base_predictor import BasePredictor
 from .predictor import get_predictor
 from .token_to_middle_json import result_to_middle_json
+from ...utils.enum_class import ModelPath
 
 
 class ModelSingleton:
@@ -40,7 +41,7 @@ def doc_analyze(
     image_writer: DataWriter | None,
     predictor: BasePredictor | None = None,
     backend="huggingface",
-    model_path="jinzhenj/OEEzRkQ3RTAtMDMx-0415",  # TODO: change to formal path after release.
+    model_path=ModelPath.vlm_root_hf,
     server_url: str | None = None,
 ):
     if predictor is None:
@@ -66,7 +67,7 @@ async def aio_doc_analyze(
     image_writer: DataWriter | None,
     predictor: BasePredictor | None = None,
     backend="huggingface",
-    model_path="jinzhenj/OEEzRkQ3RTAtMDMx-0415",  # TODO: change to formal path after release.
+    model_path=ModelPath.vlm_root_hf,
     server_url: str | None = None,
 ):
     if predictor is None:

mineru/cli/client.py

@@ -113,9 +113,19 @@ from .common import do_parse, read_fn, pdf_suffixes, image_suffixes
     help='Device mode for model inference, e.g., "cpu", "cuda", "cuda:0", "npu", "npu:0", "mps". Default is "cpu". Adapted only for the case where the backend is set to "pipeline". ',
     default=None,
 )
+@click.option(
+    '-r',
+    '--repo',
+    'model_source',
+    type=click.Choice(['huggingface', 'modelscope']),
+    help="""
+    The source of the model repository. Default is 'huggingface'.
+    """,
+    default='huggingface',
+)
 
 
-def main(input_path, output_dir, backend, lang, server_url, start_page_id, end_page_id, formula_enable, table_enable, device_mode, virtual_vram):
+def main(input_path, output_dir, backend, lang, server_url, start_page_id, end_page_id, formula_enable, table_enable, device_mode, virtual_vram, model_source):
 
     os.environ['MINERU_FORMULA_ENABLE'] = str(formula_enable).lower()
     os.environ['MINERU_TABLE_ENABLE'] = str(table_enable).lower()
@@ -138,6 +148,8 @@ def main(input_path, output_dir, backend, lang, server_url, start_page_id, end_p
 
     os.environ['MINERU_VIRTUAL_VRAM_SIZE']= str(get_virtual_vram_size())
 
+    os.environ['MINERU_MODEL_SOURCE'] = model_source
+
     os.makedirs(output_dir, exist_ok=True)
 
     def parse_doc(path_list: list[Path]):

mineru/cli/common.py

@@ -16,6 +16,7 @@ from mineru.backend.pipeline.pipeline_analyze import doc_analyze as pipeline_doc
 from mineru.data.data_reader_writer import FileBasedDataWriter
 from mineru.utils.draw_bbox import draw_layout_bbox, draw_span_bbox
 from mineru.utils.enum_class import MakeMode
+from mineru.utils.models_download_utils import get_file_from_repos
 from mineru.utils.pdf_image_tools import images_bytes_to_pdf_bytes
 
 pdf_suffixes = [".pdf"]
@@ -77,7 +78,6 @@ def do_parse(
     pdf_bytes_list: list[bytes],
     p_lang_list: list[str],
     backend="pipeline",
-    model_path="jinzhenj/OEEzRkQ3RTAtMDMx-0415",  # TODO: change to formal path after release.
     parse_method="auto",
     p_formula_enable=True,
     p_table_enable=True,
@@ -165,6 +165,7 @@ def do_parse(
             pdf_bytes = convert_pdf_bytes_to_bytes_by_pypdfium2(pdf_bytes, start_page_id, end_page_id)
             local_image_dir, local_md_dir = prepare_env(output_dir, pdf_file_name, parse_method)
             image_writer, md_writer = FileBasedDataWriter(local_image_dir), FileBasedDataWriter(local_md_dir)
+            model_path = get_file_from_repos('/','vlm')
             middle_json, infer_result = vlm_doc_analyze(pdf_bytes, image_writer=image_writer, backend=backend, model_path=model_path, server_url=server_url)
 
             pdf_info = middle_json["pdf_info"]
@@ -215,8 +216,8 @@ def do_parse(
 
 
 if __name__ == "__main__":
-    pdf_path = "../../demo/pdfs/demo2.pdf"
-    # pdf_path = "C:/Users/zhaoxiaomeng/Downloads/数学新星问题征解第一期(2014.03).pdf"
+    # pdf_path = "../../demo/pdfs/demo2.pdf"
+    pdf_path = "C:/Users/zhaoxiaomeng/Downloads/input_img_0.jpg"
 
     try:
        do_parse("./output", [Path(pdf_path).stem], [read_fn(Path(pdf_path))],["ch"], end_page_id=20,)