Update app.py

app.py

@@ -3,22 +3,26 @@ import json
 from huggingface_hub import snapshot_download
 import gradio as gr
 
-# 1) Pre‐download only the MFD/YOLO weights into ./models
-MODEL_REPO    = "opendatalab/pdf-extract-kit-1.0"
-LOCAL_MODELS  = "./models"
+# 1) Pre-download all MinerU model weights under models/
+MODEL_REPO   = "opendatalab/pdf-extract-kit-1.0"
+LOCAL_MODELS = "./models"
+
+# Grab both YOLO and MFR weights in one go
 snapshot_download(
     repo_id      = MODEL_REPO,
     local_dir    = LOCAL_MODELS,
-    allow_patterns = "models/MFD/YOLO/*",
+    allow_patterns = [
+        "models/MFD/YOLO/*",
+        "models/MFR/*"
+    ],
     max_workers  = 4
 )
 
-# 2) Write magic-pdf.json pointing at the actual weights path
+# 2) Write magic-pdf.json pointing at the nested 'models' directory
 CFG_PATH = os.path.expanduser("~/magic-pdf.json")
 if not os.path.exists(CFG_PATH):
     cfg = {
         "device":         "cpu",
-        # ← TWEAKED: point into the nested 'models' folder
         "models-dir":     os.path.join(LOCAL_MODELS, "models"),
         "layout-model":   "layoutlmv3",
         "formula-enable": True,
@@ -44,9 +48,11 @@ def convert_with_mineru(pdf_file, out_fmt):
     for ds in datasets:
         method = ds.classify()
         infer  = ds.apply(doc_analyze, ocr=(method == SupportedPdfParseMethod.OCR))
-        pipe   = (infer.pipe_ocr_mode(img_writer)
-                  if method == SupportedPdfParseMethod.OCR
-                  else infer.pipe_txt_mode(img_writer))
+        pipe   = (
+            infer.pipe_ocr_mode(img_writer)
+            if method == SupportedPdfParseMethod.OCR
+            else infer.pipe_txt_mode(img_writer)
+        )
 
         base    = os.path.splitext(os.path.basename(pdf_file.name))[0]
         md_name = f"{base}.md"
@@ -70,11 +76,11 @@ demo = gr.Interface(
     fn=convert_with_mineru,
     inputs=[gr.File(label="Upload PDF"), gr.Radio(["markdown", "json"], label="Format")],
     outputs=gr.Code(label="Result"),
-    title="MinerU PDF → Markdown/JSON (Fixed)",
-    description="Downloads YOLO weights and points magic-pdf at the correct folder."
+    title="MinerU PDF → Markdown/JSON (Fully Fixed)",
+    description="Pre-downloads all necessary YOLO and MFR weights and configures magic-pdf correctly."
 )
 
 if __name__ == "__main__":
-    # Ensure HF_HUB_CACHE is consistent with our models folder
+    # Ensure HF_HUB_CACHE aligns with our models folder
     os.environ.setdefault("HF_HUB_CACHE", LOCAL_MODELS)
     demo.launch(server_name="0.0.0.0", server_port=7860)