Update app.py

app.py

@@ -1,84 +1,80 @@
 import os
 import json
+from huggingface_hub import snapshot_download
 import gradio as gr
 
-# Ensure default config for magic-pdf
-CONFIG_PATH = os.path.expanduser("~/magic-pdf.json")
-if not os.path.exists(CONFIG_PATH):
-    default_cfg = {
-        "device": "cpu",      # force CPU inference
-        "layout_model": "layout/mobilenetv3",
-        "formula_enable": True,
-        "table_enable": True
+# 1) Pre-download only the MFD/YOLO weights
+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
+)
+
+# 2) Write a minimal magic-pdf.json pointing to our models
+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",  
+        "formula-enable": True,
+        "table-enable": True
     }
-    with open(CONFIG_PATH, "w", encoding="utf-8") as cfg:
-        json.dump(default_cfg, cfg, ensure_ascii=False, indent=2)
+    with open(CFG_PATH, "w", encoding="utf-8") as f:
+        json.dump(cfg, f, ensure_ascii=False, indent=2)
 
-# MinerU API imports
+# 3) MinerU imports
 from magic_pdf.data.read_api import read_local_pdfs
 from magic_pdf.data.data_reader_writer import FileBasedDataWriter
 from magic_pdf.model.doc_analyze_by_custom_model import doc_analyze
 from magic_pdf.config.enums import SupportedPdfParseMethod
 
-def convert_with_mineru(pdf_path, out_format):
-    # 1) Read file into MinerU dataset(s)
-    datasets = read_local_pdfs(pdf_path)
-
-    tmp_dir = "output"
-    img_dir = os.path.join(tmp_dir, "images")
+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_dir)
+    md_writer = FileBasedDataWriter(tmp)
     img_writer = FileBasedDataWriter(img_dir)
 
-    all_pages = []
-
+    results = []
     for ds in datasets:
-        # 2) Classify & infer, with OCR fallback
         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))
-
-        # 3) Dump per‐document Markdown
-        basename = os.path.splitext(os.path.basename(pdf_path))[0]
-        md_fname = f"{basename}.md"
-        pipe.dump_md(md_writer, md_fname, os.path.basename(img_dir))
-        with open(os.path.join(tmp_dir, md_fname), "r", encoding="utf-8") as f:
-            page_md = f.read()
+        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:
+            md_text = f.read()
 
-        # 4) Dump structured JSON
-        json_fname = f"{basename}_content_list.json"
-        pipe.dump_content_list(md_writer, json_fname, os.path.basename(img_dir))
-        with open(os.path.join(tmp_dir, json_fname), "r", encoding="utf-8") as f:
-            content_list = json.load(f)
+        json_name = f"{base}_content_list.json"
+        pipe.dump_content_list(md_writer, json_name, os.path.basename(img_dir))
+        with open(os.path.join(tmp, json_name), encoding="utf-8") as f:
+            content = json.load(f)
 
-        all_pages.append({
-            "markdown": page_md,
-            "content_list": content_list
-        })
+        results.append({"markdown": md_text, "content_list": content})
 
-    # 5) Return desired format
-    if out_format == "markdown":
-        return "\n\n---\n\n".join(p["markdown"] for p in all_pages)
-    else:
-        return json.dumps(all_pages, ensure_ascii=False, indent=2)
+    if out_fmt == "markdown":
+        return "\n\n---\n\n".join(r["markdown"] for r in results)
+    return json.dumps(results, ensure_ascii=False, indent=2)
 
-# Gradio interface
+# 4) Gradio UI
 demo = gr.Interface(
     fn=convert_with_mineru,
-    inputs=[
-        gr.File(label="Upload PDF", file_types=[".pdf"]),
-        gr.Radio(["markdown", "json"], value="markdown", label="Output format")
-    ],
+    inputs=[gr.File(label="Upload PDF"), gr.Radio(["markdown", "json"], label="Format")],
     outputs=gr.Code(label="Result"),
-    title="MinerU-Powered PDF → Markdown/JSON",
-    description=(
-        "Leverage the advanced MinerU engine to extract text, images, tables, "
-        "and formulas from your PDF into clean Markdown or structured JSON. "
-        "A default CPU-only config is auto-generated if none is found."
-    )
+    title="MinerU PDF → Markdown/JSON (Fixed)",
+    description="Pre-downloads YOLO weights and configures MinerU for Spaces."
 )
 
 if __name__ == "__main__":
+    # Recommended: ensure HF_HUB_CACHE points to ./models
+    os.environ.setdefault("HF_HUB_CACHE", LOCAL_MODELS)
     demo.launch(server_name="0.0.0.0", server_port=7860)