download_models.py

"""
Download MinerU pipeline models from Hugging Face Hub.

Called once during Docker build: python download_models.py

Optimizations vs original:
  - Skip-if-exists: if models are already present (Docker layer cache reuse),
    the entire download is skipped without re-downloading anything.
  - MFR excluded: formula recognition models (unimernet, ~1-2 GB) are not
    downloaded because formula recognition is disabled in the config. Excluding
    them cuts download time and image size significantly.
  - layoutreader optional: if download fails (network issue, repo unavailable),
    the script logs a warning and continues. MinerU falls back to its built-in
    layout ordering without layoutreader.

Models saved to /app/models/
Config written to /root/magic-pdf.json
"""

import json
import os
import sys

MODELS_DIR = "/app/models"
EXTRACT_KIT_DIR = os.path.join(MODELS_DIR, "PDF-Extract-Kit-1.0")
LAYOUTREADER_DIR = os.path.join(MODELS_DIR, "layoutreader")


def _models_present() -> bool:
    """Return True if the key layout-model directory already exists."""
    marker = os.path.join(EXTRACT_KIT_DIR, "models", "Layout")
    return os.path.isdir(marker)


def _write_config(layoutreader_dir: str) -> None:
    config = {
        "bucket_info": {},
        "models-dir": os.path.join(EXTRACT_KIT_DIR, "models"),
        "layoutreader-model-dir": layoutreader_dir,
        "device-mode": "cpu",
        "layout-config": {
            "model": "doclayout_yolo"
        },
        "formula-config": {
            "mfd_model": "yolo_v8_mfd",
            "mfr_model": "unimernet_small",
            "enable": False
        },
        "table-config": {
            "model": "rapid_table",
            "enable": True,
            "max_time": 400
        },
        "backend": "pipeline"
    }
    config_path = os.path.expanduser("~/magic-pdf.json")
    with open(config_path, "w") as f:
        json.dump(config, f, indent=2)
    print(f"Config written → {config_path}")
    return config_path


def download() -> None:
    try:
        from huggingface_hub import snapshot_download
    except ImportError:
        print("ERROR: huggingface_hub not installed", file=sys.stderr)
        sys.exit(1)

    # ── Skip-if-exists ────────────────────────────────────────────────────────
    if _models_present():
        print("Models already present — skipping download (Docker layer cache).")
        # Config may still need writing if this is a fresh container from cached layer
        lr_dir = LAYOUTREADER_DIR if os.path.isdir(LAYOUTREADER_DIR) else ""
        _write_config(lr_dir)
        return

    os.makedirs(MODELS_DIR, exist_ok=True)

    # ── PDF-Extract-Kit-1.0 ───────────────────────────────────────────────────
    # Excluded via ignore_patterns:
    #   models/MFR — formula recognition (unimernet). Disabled in config.
    #                Saves ~1-2 GB of download and disk.
    print("=" * 60)
    print("Downloading PDF-Extract-Kit-1.0 ...")
    print("  (MFR/formula-recognition excluded — disabled in config)")
    print("=" * 60)
    snapshot_download(
        repo_id="opendatalab/PDF-Extract-Kit-1.0",
        local_dir=EXTRACT_KIT_DIR,
        ignore_patterns=[
            "*.git*",
            ".gitattributes",
            "models/MFR*",
            "models/MFR/*",
        ],
    )
    print(f"  → {EXTRACT_KIT_DIR}")

    # ── layoutreader (optional) ───────────────────────────────────────────────
    # Improves reading-order accuracy. If unavailable, MinerU uses fallback.
    layoutreader_dir = ""
    print("=" * 60)
    print("Downloading layoutreader (optional, improves reading order) ...")
    print("=" * 60)
    try:
        snapshot_download(
            repo_id="hantian/layoutreader",
            local_dir=LAYOUTREADER_DIR,
            ignore_patterns=["*.git*", ".gitattributes"],
        )
        layoutreader_dir = LAYOUTREADER_DIR
        print(f"  → {LAYOUTREADER_DIR}")
    except Exception as exc:
        print(f"  WARNING: layoutreader download failed ({exc})")
        print("  Continuing without layoutreader — MinerU will use fallback ordering.")

    # ── Write config ──────────────────────────────────────────────────────────
    _write_config(layoutreader_dir)

    print("\n✓ Model setup complete.")
    print(f"  models-dir          : {os.path.join(EXTRACT_KIT_DIR, 'models')}")
    print(f"  layoutreader-dir    : {layoutreader_dir or '(not available)'}")
    print(f"  device-mode         : cpu")
    print(f"  formula recognition : disabled (MFR models excluded)")
    print(f"  table recognition   : enabled")


if __name__ == "__main__":
    download()