Update pdf_server.py

magic_pdf/tools/pdf_server.py

@@ -9,6 +9,7 @@ import magic_pdf.model as model_config
 from magic_pdf.dict2md.ocr_client import PredictClient
 # from magic_pdf.dict2md.ocr_vllm_client import PredictClient
 from magic_pdf.libs.version import __version__
+from magic_pdf.model.doc_analyze_by_custom_model import ModelSingleton
 from magic_pdf.rw.AbsReaderWriter import AbsReaderWriter
 from magic_pdf.rw.DiskReaderWriter import DiskReaderWriter
 from magic_pdf.tools.common import do_parse, parse_pdf_methods
@@ -27,6 +28,7 @@ logger.add("parse.log", rotation="10 MB", level="INFO",
 config_path = None
 
 ocr_status = None
+custom_model = None
 
 class ocrRequest(BaseModel):
     path: str
@@ -80,6 +82,14 @@ def ocr_pdf_serve(args: str):
     ocr_client = PredictClient(ocr_server)
     global ocr_status
     ocr_status = ocr_client.check_health()
+
+    ocr = True
+    show_log = False
+    model_manager = ModelSingleton()
+    global custom_model
+    custom_model = model_manager.get_model(ocr, show_log)
+
+
     uvicorn.run(app, host=host, port=port)
 
 @app.get("/health")
@@ -117,6 +127,7 @@ async def pdf_ocr(request: ocrRequest):
                 [],
                 method,
                 debug_able,
+                model=custom_model,
                 start_page_id=start_page_id,
                 end_page_id=end_page_id,
             )