app.py

"""
MinerU on ZeroGPU.

We keep mineru's official Gradio UI intact (i18n, status panel, PDF viewer,
examples, locale, latex delimiters) by importing `mineru.cli.gradio_app`
and invoking its Click command in-process.

But we replace the local mineru-api FastAPI subprocess + HTTP round-trip with
a direct in-process call to `aio_do_parse`, wrapped in @spaces.GPU so it runs
on ZeroGPU. The model is preloaded at module scope into the snapshot.
"""
import os
import sys
import time
import uuid
import shutil
import asyncio
import zipfile
import subprocess
from pathlib import Path
from typing import Callable

# Set before any mineru import
os.environ.setdefault("MINERU_MODEL_SOURCE", "huggingface")
os.environ.setdefault("MINERU_VLM_FORMULA_ENABLE", "true")
os.environ.setdefault("MINERU_VLM_TABLE_ENABLE", "true")

# CUDA hook before any torch / cuda touch
import spaces  # noqa: E402


def _ensure_mineru():
    try:
        import mineru  # noqa: F401
        return
    except ImportError:
        pass
    subprocess.check_call(
        [sys.executable, "-m", "pip", "install", "-q",
         "mineru[vlm,core]>=3.1.0", "gradio-pdf>=0.0.22"],
    )

_ensure_mineru()

# Download all weights (idempotent if already present)
print("Downloading MinerU model weights...")
subprocess.check_call(["mineru-models-download", "-s", "huggingface", "-m", "all"])

# ---------------------------------------------------------------------------
# Preload the transformers backend into the ZeroGPU snapshot
# ---------------------------------------------------------------------------
from mineru.backend.vlm.vlm_analyze import ModelSingleton
from mineru.cli.common import aio_do_parse

print("Preloading MinerU VLM (transformers backend) into ZeroGPU snapshot...")
_t0 = time.time()
ModelSingleton().get_model(
    backend="transformers",
    model_path=None,
    server_url=None,
)
print(f"Preloaded in {time.time() - _t0:.1f}s")


# ---------------------------------------------------------------------------
# Monkey-patch mineru-gradio: skip the local mineru-api subprocess entirely
# and route the inference call directly to aio_do_parse, wrapped in @spaces.GPU.
# ---------------------------------------------------------------------------
# Patch pdf_image_tools to use a thread pool instead of process pool —
# we're inside a daemon @spaces.GPU worker which can't spawn child processes.
from mineru.utils import pdf_image_tools as _pit
from concurrent.futures import ThreadPoolExecutor

def _thread_pdf_render_executor(max_workers=None):
    return ThreadPoolExecutor(max_workers=max_workers or 4)
_pit._create_pdf_render_executor = _thread_pdf_render_executor
_pit._pdf_render_executor = ThreadPoolExecutor(max_workers=4)
_pit._get_pdf_render_executor = lambda: _pit._pdf_render_executor

from mineru.cli import gradio_app as mineru_gradio
from mineru.cli.gradio_app import (
    STATUS_PREPARING_REQUEST,
    STATUS_CHECKING_SERVER,
    STATUS_SUBMITTING_TASK,
    STATUS_QUEUED_ON_SERVER,
    STATUS_PROCESSING_ON_SERVER,
    STATUS_DOWNLOADING_RESULT,
    STATUS_PROCESSING_OUTPUT,
    STATUS_COMPLETED,
    create_gradio_run_paths,
    resolve_parse_method,
    resolve_parse_dir,
    replace_image_with_gradio_file_urls,
    maybe_generate_local_preview,
    compress_directory_to_zip,
    office_suffixes,
    normalize_language,
    build_gradio_upload_name,
)
from mineru.cli import api_client as _api_client


from mineru.cli.common import do_parse as _sync_do_parse

@spaces.GPU(duration=600)
def _gpu_do_parse_sync(
    output_dir,
    pdf_file_names,
    pdf_bytes_list,
    p_lang_list,
    backend,
    formula_enable,
    table_enable,
    end_page_id,
    f_dump_md,
    f_dump_content_list,
    f_dump_orig_pdf,
    f_dump_middle_json,
    f_dump_model_output,
    image_analysis,
):
    # Use the SYNC do_parse — aio_do_parse pulls in ProcessPoolExecutor for PDF
    # rendering, and ZeroGPU's daemon worker can't spawn children.
    return _sync_do_parse(
        output_dir=output_dir,
        pdf_file_names=pdf_file_names,
        pdf_bytes_list=pdf_bytes_list,
        p_lang_list=p_lang_list,
        backend=backend,
        formula_enable=formula_enable,
        table_enable=table_enable,
        end_page_id=end_page_id,
        f_draw_layout_bbox=True,   # produce the layout-annotated PDF for the doc-preview pane
        f_draw_span_bbox=False,
        f_dump_md=f_dump_md,
        f_dump_content_list=f_dump_content_list,
        f_dump_orig_pdf=f_dump_orig_pdf,
        f_dump_middle_json=f_dump_middle_json,
        f_dump_model_output=f_dump_model_output,
        image_analysis=image_analysis,
    )


async def _run_to_markdown_job_local(
    file_path,
    end_pages=10,
    is_ocr=False,
    formula_enable=True,
    table_enable=True,
    image_analysis=True,
    language="ch",
    backend="vlm-transformers",
    url=None,
    api_url=None,
    status_callback: Callable[[str], None] | None = None,
):
    """In-process replacement for mineru-gradio's _run_to_markdown_job.

    Original calls a local mineru-api over HTTP. We call aio_do_parse directly
    inside a @spaces.GPU window so ZeroGPU treats this as a single GPU task."""

    def emit(msg: str) -> None:
        if status_callback is not None:
            status_callback(msg)

    if file_path is None:
        return "", "", None, None

    # Force the transformers backend regardless of what the UI dropdown set,
    # since we only support that path on ZeroGPU.
    if backend == "pipeline":
        backend = "pipeline"
    elif backend.startswith("vlm-") or backend == "auto":
        backend = "vlm-transformers"

    normalized_language = normalize_language(language)
    file_path = str(file_path)
    file_suffix = Path(file_path).suffix.lower().lstrip(".")
    parse_method = resolve_parse_method(file_path, is_ocr, backend)
    run_root, extract_root, archive_zip_path = create_gradio_run_paths(file_path)
    run_root.mkdir(parents=True, exist_ok=True)
    extract_root.mkdir(parents=True, exist_ok=True)

    upload_name = build_gradio_upload_name(file_path)
    file_stem = Path(upload_name).stem

    task_id = uuid.uuid4().hex[:12]
    emit(STATUS_PREPARING_REQUEST)
    emit(STATUS_CHECKING_SERVER)
    emit(STATUS_SUBMITTING_TASK)
    emit(f"Task submitted: task_id={task_id}")
    # Local execution skips the server queue; advance the UI past it
    emit(STATUS_QUEUED_ON_SERVER)
    emit(STATUS_PROCESSING_ON_SERVER)

    # Load the file bytes
    from mineru.cli.common import read_fn
    pdf_bytes = read_fn(Path(file_path))

    # Run the parser inside a single @spaces.GPU window
    end_page_id = (int(end_pages) - 1) if end_pages else None
    # Inference happens inside _gpu_do_parse_sync, which is decorated.
    # Run it in a thread so we don't block the asyncio loop.
    await asyncio.to_thread(
        _gpu_do_parse_sync,
        output_dir=str(extract_root),
        pdf_file_names=[file_stem],
        pdf_bytes_list=[pdf_bytes],
        p_lang_list=[normalized_language],
        backend=backend,
        formula_enable=bool(formula_enable),
        table_enable=bool(table_enable),
        end_page_id=end_page_id,
        f_dump_md=True,
        f_dump_content_list=True,
        f_dump_orig_pdf=True,
        f_dump_middle_json=True,
        f_dump_model_output=True,
        image_analysis=bool(image_analysis),
    )

    file_name = file_stem
    local_md_dir = resolve_parse_dir(
        extract_root,
        file_name,
        backend,
        parse_method,
        allow_office_fallback=True,
    )
    preview_pdf_path = maybe_generate_local_preview(
        extract_root=extract_root,
        file_name=file_name,
        file_suffix=file_suffix,
        backend=backend,
        parse_method=parse_method,
    )

    emit(STATUS_DOWNLOADING_RESULT)
    emit(STATUS_PROCESSING_OUTPUT)

    # Zip the per-document output dir (same shape as what mineru-api ships back)
    if compress_directory_to_zip(local_md_dir, archive_zip_path) != 0:
        # Fallback: zip with stdlib
        with zipfile.ZipFile(archive_zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
            for p in Path(local_md_dir).rglob("*"):
                if p.is_file():
                    zf.write(p, p.relative_to(local_md_dir))

    md_path = Path(local_md_dir) / f"{file_name}.md"
    if not md_path.exists():
        # office_suffixes may have produced a different name
        candidates = list(Path(local_md_dir).glob("*.md"))
        if candidates:
            md_path = candidates[0]
    txt_content = md_path.read_text(encoding="utf-8", errors="replace") if md_path.exists() else ""
    md_content = replace_image_with_gradio_file_urls(txt_content, str(local_md_dir))

    if file_suffix in office_suffixes:
        preview_pdf_path = None

    emit(STATUS_COMPLETED)
    return md_content, txt_content, str(archive_zip_path), preview_pdf_path


# Patch the original mineru-gradio job runner
mineru_gradio._run_to_markdown_job = _run_to_markdown_job_local

# Disable the local-API startup hook
mineru_gradio.maybe_prepare_local_api_for_gradio_startup = lambda **kw: None
async def _noop_ensure(*a, **k):
    return None
mineru_gradio.ensure_local_api_ready_for_gradio_startup = _noop_ensure

# Replace resolve_server_health to return a stub: avoids the upstream HTTP probe
class _StubServerHealth:
    max_concurrent_requests = 1
    base_url = "http://disabled/"

async def _resolve_server_health_stub(http_client, api_url):
    return _StubServerHealth()
mineru_gradio.resolve_server_health = _resolve_server_health_stub
mineru_gradio.resolve_gradio_max_concurrent_requests = lambda api_url, server_health: 1


# ---------------------------------------------------------------------------
# Launch mineru-gradio in-process with our patches active
# ---------------------------------------------------------------------------
sys.argv = [
    "mineru-gradio",
    "--enable-api", "false",
    "--enable-http-client", "false",
    "--enable-vlm-preload", "false",  # already preloaded above
    "--max-convert-pages", "25",
    "--latex-delimiters-type", "b",
]

from mineru.cli.gradio_app import main as _mineru_main
_mineru_main(standalone_mode=False)