Update app.py

app.py

@@ -8,11 +8,11 @@ import time
 import zipfile
 from pathlib import Path
 
-os.system('pip uninstall -y magic-pdf')
 os.system('pip uninstall -y mineru')
-os.system('pip install git+https://github.com/opendatalab/MinerU.git@dev')
+os.system('pip install git+https://github.com/myhloli/Magic-PDF.git@dev')
+os.system('pip install sglang[all]==0.4.8')
 
-os.system('mineru-models-download -s huggingface -m pipeline')
+os.system('mineru-models-download -s huggingface -m all')
 
 with open('/home/user/mineru.json', 'r') as file:
     config = json.load(file)
@@ -35,24 +35,27 @@ from gradio_pdf import PDF
 import gradio as gr
 from loguru import logger
 
-from mineru.cli.common import prepare_env, do_parse, read_fn
+from mineru.cli.common import prepare_env, read_fn, aio_do_parse
 from mineru.utils.hash_utils import str_sha256
 
 os.environ['MINERU_MODEL_SOURCE'] = 'local'
 
 
-def parse_pdf(doc_path, output_dir, end_page_id, is_ocr, formula_enable, table_enable, language):
+async def parse_pdf(doc_path, output_dir, end_page_id, is_ocr, formula_enable, table_enable, language, backend, url):
     os.makedirs(output_dir, exist_ok=True)
 
     try:
-        file_name = f'{str(Path(doc_path).stem)}_{time.strftime("%y%m%d_%H%M%S")}'
+        file_name = f'{safe_stem(Path(doc_path).stem)}_{time.strftime("%y%m%d_%H%M%S")}'
         pdf_data = read_fn(doc_path)
         if is_ocr:
             parse_method = 'ocr'
         else:
             parse_method = 'auto'
+
+        if backend.startswith("vlm"):
+            parse_method = "vlm"
         local_image_dir, local_md_dir = prepare_env(output_dir, file_name, parse_method)
-        do_parse(
+        await aio_do_parse(
             output_dir=output_dir,
             pdf_file_names=[file_name],
             pdf_bytes_list=[pdf_data],
@@ -61,6 +64,8 @@ def parse_pdf(doc_path, output_dir, end_page_id, is_ocr, formula_enable, table_e
             end_page_id=end_page_id,
             p_formula_enable=formula_enable,
             p_table_enable=table_enable,
+            backend=backend,
+            server_url=url,
         )
         return local_md_dir, file_name
     except Exception as e:
@@ -112,12 +117,10 @@ def replace_image_with_base64(markdown_text, image_dir_path):
     return re.sub(pattern, replace, markdown_text)
 
 
-def to_markdown(file_path, end_pages, is_ocr, formula_enable, table_enable, language):
+async def to_markdown(file_path, end_pages=10, is_ocr=False, formula_enable=True, table_enable=True, language="ch", backend="pipeline", url=None):
     file_path = to_pdf(file_path)
-    if end_pages > 20:
-        end_pages = 20
     # 获取识别的md文件以及压缩包文件路径
-    local_md_dir, file_name = parse_pdf(file_path, './output', end_pages - 1, is_ocr, formula_enable, table_enable, language)
+    local_md_dir, file_name = await parse_pdf(file_path, './output', end_pages - 1, is_ocr, formula_enable, table_enable, language, backend, url)
     archive_zip_path = os.path.join('./output', str_sha256(local_md_dir) + '.zip')
     zip_archive_success = compress_directory_to_zip(local_md_dir, archive_zip_path)
     if zip_archive_success == 0:
@@ -134,8 +137,12 @@ def to_markdown(file_path, end_pages, is_ocr, formula_enable, table_enable, lang
     return md_content, txt_content, archive_zip_path, new_pdf_path
 
 
-latex_delimiters = [{"left": "\\[", "right": "\\]", "display": True},
-                    {"left": "\\(", "right": "\\)", "display": False}]
+latex_delimiters = [
+    {'left': '$$', 'right': '$$', 'display': True},
+    {'left': '$', 'right': '$', 'display': False},
+    {'left': '\\(', 'right': '\\)', 'display': False},
+    {'left': '\\[', 'right': '\\]', 'display': True},
+]
 
 
 with open("header.html", "r") as file:
@@ -192,19 +199,39 @@ def to_pdf(file_path):
     return tmp_file_path
     
 
-if __name__ == '__main__':
+def main():
+    example_enable = True
+
+    try:
+        print("预初始化SgLang引擎...")
+        from mineru.backend.vlm.vlm_analyze import ModelSingleton
+        modelsingleton = ModelSingleton()
+        predictor = modelsingleton.get_model(
+            "sglang-engine",
+            None,
+            'http://localhost:30000',
+            mem_fraction_static=0.5,
+            enable_torch_compile=True,
+        )
+        print("SgLang引擎初始化完成")
+    except Exception as e:
+        logger.exception(e)
+
     with gr.Blocks() as demo:
         gr.HTML(header)
         with gr.Row():
             with gr.Column(variant='panel', scale=5):
                 with gr.Row():
                     file = gr.File(label='Please upload a PDF or image', file_types=['.pdf', '.png', '.jpeg', '.jpg'])
-                with gr.Row(equal_height=True):
-                    with gr.Column(scale=4):
-                        max_pages = gr.Slider(1, 20, 10, step=1, label='Max convert pages')
-                    with gr.Column(scale=1):
-                        language = gr.Dropdown(all_lang, label='Language', value='ch')
                 with gr.Row():
+                    max_pages = gr.Slider(1, 20, 10, step=1, label='Max convert pages')
+                with gr.Row():
+                    backend = gr.Dropdown(["pipeline", "vlm-sglang-engine"], label="Backend", value="pipeline")
+                with gr.Row(visible=True) as ocr_options:
+                    language = gr.Dropdown(all_lang, label='Language', value='ch')
+                with gr.Row(visible=False) as client_options:
+                    url = gr.Textbox(label='Server URL', value='http://localhost:30000', placeholder='http://localhost:30000')
+                with gr.Row(visible=True) as pipeline_options:
                     is_ocr = gr.Checkbox(label='Force enable OCR', value=False)
                     formula_enable = gr.Checkbox(label='Enable formula recognition', value=True)
                     table_enable = gr.Checkbox(label='Enable table recognition(test)', value=True)
@@ -212,13 +239,15 @@ if __name__ == '__main__':
                     change_bu = gr.Button('Convert')
                     clear_bu = gr.ClearButton(value='Clear')
                 pdf_show = PDF(label='PDF preview', interactive=False, visible=True, height=800)
-                with gr.Accordion('Examples:'):
+                if example_enable:
                     example_root = os.path.join(os.path.dirname(__file__), 'examples')
-                    gr.Examples(
-                        examples=[os.path.join(example_root, _) for _ in os.listdir(example_root) if
-                                  _.endswith('pdf')],
-                        inputs=file
-                    )
+                    if os.path.exists(example_root):
+                        with gr.Accordion('Examples:'):
+                            gr.Examples(
+                                examples=[os.path.join(example_root, _) for _ in os.listdir(example_root) if
+                                          _.endswith('pdf')],
+                                inputs=file
+                            )
 
             with gr.Column(variant='panel', scale=5):
                 output_file = gr.File(label='convert result', interactive=False)
@@ -229,9 +258,35 @@ if __name__ == '__main__':
                                          line_breaks=True)
                     with gr.Tab('Markdown text'):
                         md_text = gr.TextArea(lines=45, show_copy_button=True)
-        file.change(fn=to_pdf, inputs=file, outputs=pdf_show)
-        change_bu.click(fn=to_markdown, inputs=[file, max_pages, is_ocr, formula_enable, table_enable, language],
-                        outputs=[md, md_text, output_file, pdf_show])
+
+
+        # 更新界面函数
+        def update_interface(backend_choice):
+            if backend_choice in ["vlm-transformers", "vlm-sglang-engine"]:
+                return gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
+            elif backend_choice in ["vlm-sglang-client"]:  # pipeline
+                return gr.update(visible=True), gr.update(visible=False), gr.update(visible=False)
+            elif backend_choice in ["pipeline"]:
+                return gr.update(visible=False), gr.update(visible=True), gr.update(visible=True)
+            else:
+                pass
+
+
+        # 添加事件处理
+        backend.change(
+            fn=update_interface,
+            inputs=[backend],
+            outputs=[client_options, ocr_options, pipeline_options],
+            api_name=False
+        )
+
+        file.change(fn=to_pdf, inputs=file, outputs=pdf_show, api_name=False)
+        change_bu.click(fn=to_markdown, inputs=[file, max_pages, is_ocr, formula_enable, table_enable, language, backend, url],
+                        outputs=[md, md_text, output_file, pdf_show], api_name=False)
         clear_bu.add([file, md, pdf_show, md_text, output_file, is_ocr])
 
-    demo.launch(ssr_mode=True)
+    demo.launch(server_name='localhost', ssr_mode=True)
+
+
+if __name__ == '__main__':
+    main()