Update pdf_server.py

magic_pdf/tools/pdf_server.py

@@ -16,8 +16,7 @@ from pydantic import BaseModel
 import uvicorn
 import time
 import configparser
-from magic_pdf.dict2md.ocr_vllm_client import PredictClient,compress_image
-# from magic_pdf.dict2md.ocr_client import PredictClient,compress_image
+
 from magic_pdf.parse.ofd_parse import *
 from magic_pdf.tools.ofd_parser import OFDParser
 from magic_pdf.parse.pdf_client import ocrPdfClient
@@ -82,6 +81,11 @@ def ocr_pdf_serve(args: str):
     config.read(args.config_path)
     # host = config.get('server', 'pdf_host')
     # port = int(config.get('server', 'pdf_port'))
+    vllm_able = config.get('vllm', 'vllm_able')
+    if vllm_able:
+        from magic_pdf.dict2md.ocr_vllm_client import PredictClient, compress_image
+    else:
+        from magic_pdf.dict2md.ocr_client import PredictClient,compress_image
     pdf_server = config.get('server', 'pdf_server').split('://')[1]
     host, port = pdf_server.split(':')[0], int(pdf_server.split(':')[1])
     global config_path