""" 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)