Update app.py

app.py

@@ -3,25 +3,26 @@ import json
 from huggingface_hub import snapshot_download
 import gradio as gr
 
-# 1) Pre-download only the MFD/YOLO weights
-MODEL_REPO = "opendatalab/pdf-extract-kit-1.0"
-LOCAL_MODELS = "./models"
+# 1) Pre‐download only the MFD/YOLO weights into ./models
+MODEL_REPO    = "opendatalab/pdf-extract-kit-1.0"
+LOCAL_MODELS  = "./models"
 snapshot_download(
-    repo_id=MODEL_REPO,
-    local_dir=LOCAL_MODELS,
-    allow_patterns="models/MFD/YOLO/*",
-    max_workers=4
+    repo_id      = MODEL_REPO,
+    local_dir    = LOCAL_MODELS,
+    allow_patterns = "models/MFD/YOLO/*",
+    max_workers  = 4
 )
 
-# 2) Write a minimal magic-pdf.json pointing to our models
+# 2) Write magic-pdf.json pointing at the actual weights path
 CFG_PATH = os.path.expanduser("~/magic-pdf.json")
 if not os.path.exists(CFG_PATH):
     cfg = {
-        "device": "cpu",               # CPU fallback
-        "models-dir": LOCAL_MODELS,    # where we downloaded yolo_v8_ft.pt
-        "layout-model": "layoutlmv3",  
+        "device":         "cpu",
+        # ← TWEAKED: point into the nested 'models' folder
+        "models-dir":     os.path.join(LOCAL_MODELS, "models"),
+        "layout-model":   "layoutlmv3",
         "formula-enable": True,
-        "table-enable": True
+        "table-enable":   True
     }
     with open(CFG_PATH, "w", encoding="utf-8") as f:
         json.dump(cfg, f, ensure_ascii=False, indent=2)
@@ -36,19 +37,18 @@ def convert_with_mineru(pdf_file, out_fmt):
     datasets = read_local_pdfs(pdf_file.name)
     tmp, img_dir = "output", os.path.join("output", "images")
     os.makedirs(img_dir, exist_ok=True)
-    md_writer = FileBasedDataWriter(tmp)
+    md_writer  = FileBasedDataWriter(tmp)
     img_writer = FileBasedDataWriter(img_dir)
 
     results = []
     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)
-        )
-        base = os.path.splitext(os.path.basename(pdf_file.name))[0]
+        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))
+
+        base    = os.path.splitext(os.path.basename(pdf_file.name))[0]
         md_name = f"{base}.md"
         pipe.dump_md(md_writer, md_name, os.path.basename(img_dir))
         with open(os.path.join(tmp, md_name), encoding="utf-8") as f:
@@ -71,10 +71,10 @@ demo = gr.Interface(
     inputs=[gr.File(label="Upload PDF"), gr.Radio(["markdown", "json"], label="Format")],
     outputs=gr.Code(label="Result"),
     title="MinerU PDF → Markdown/JSON (Fixed)",
-    description="Pre-downloads YOLO weights and configures MinerU for Spaces."
+    description="Downloads YOLO weights and points magic-pdf at the correct folder."
 )
 
 if __name__ == "__main__":
-    # Recommended: ensure HF_HUB_CACHE points to ./models
+    # Ensure HF_HUB_CACHE is consistent with our models folder
     os.environ.setdefault("HF_HUB_CACHE", LOCAL_MODELS)
     demo.launch(server_name="0.0.0.0", server_port=7860)