app.py

#!/usr/bin/env python3

import logging
import os
import subprocess
import tempfile
import time
from typing import Any, List

import gradio as gr
import httpx
import yaml

log = logging.getLogger("glmocr_selfhosted_space")
logging.basicConfig(level=logging.INFO)

MODEL_ID = "zai-org/GLM-OCR"
OCR_API_PORT = 8080
OCR_API_BASE = f"http://127.0.0.1:{OCR_API_PORT}"
# Fixed runtime settings (no environment variables required).
VLLM_MAX_MODEL_LEN = "8192"
VLLM_GPU_MEMORY_UTIL = "0.75"
VLLM_EXTRA_ARGS = ""
# Keep output budget below context length to leave room for prompt tokens.
GLMOCR_MAX_OUTPUT_TOKENS = 2048

_parser = None
_vllm_proc = None


def _render_pdf_pages_to_images(pdf_path: str) -> List[str]:
    import pymupdf as fitz

    doc = fitz.open(pdf_path)
    page_images: List[str] = []
    try:
        for i in range(len(doc)):
            page = doc[i]
            pix = page.get_pixmap(matrix=fitz.Matrix(2.5, 2.5), alpha=False)
            img_path = os.path.join(
                tempfile.gettempdir(), f"glmocr_page_{os.getpid()}_{i}_{int(time.time()*1000)}.png"
            )
            pix.save(img_path)
            page_images.append(img_path)
    finally:
        doc.close()
    return page_images


def _start_vllm_if_needed() -> None:
    global _vllm_proc
    if _vllm_proc is not None and _vllm_proc.poll() is None:
        return

    cmd: List[str] = [
        "python",
        "-m",
        "vllm.entrypoints.openai.api_server",
        "--model",
        MODEL_ID,
        "--served-model-name",
        "glm-ocr",
        "--port",
        str(OCR_API_PORT),
        "--max-model-len",
        VLLM_MAX_MODEL_LEN,
        "--gpu-memory-utilization",
        VLLM_GPU_MEMORY_UTIL,
    ]
    if VLLM_EXTRA_ARGS:
        cmd.extend(VLLM_EXTRA_ARGS.split())

    log.info("Starting vLLM server: %s", " ".join(cmd))
    _vllm_proc = subprocess.Popen(cmd)

    deadline = time.time() + 420
    last_err: str = ""
    while time.time() < deadline:
        if _vllm_proc.poll() is not None:
            raise RuntimeError("vLLM server exited early. Check Space logs for details.")
        try:
            resp = httpx.get(f"{OCR_API_BASE}/v1/models", timeout=8.0)
            if resp.status_code == 200:
                log.info("vLLM is ready.")
                return
        except Exception as e:
            last_err = str(e)
        time.sleep(3)
    raise RuntimeError(f"Timed out waiting for vLLM startup. Last error: {last_err}")


def _configure_glmocr_to_keep_header_footer() -> None:
    import glmocr

    config_path = os.path.join(os.path.dirname(glmocr.__file__), "config.yaml")
    with open(config_path, "r", encoding="utf-8") as f:
        cfg = yaml.safe_load(f) or {}

    pipeline = cfg.setdefault("pipeline", {})

    maas = pipeline.setdefault("maas", {})
    maas["enabled"] = False

    ocr_api = pipeline.setdefault("ocr_api", {})
    ocr_api["api_host"] = "127.0.0.1"
    ocr_api["api_port"] = OCR_API_PORT
    ocr_api["model"] = "glm-ocr"
    ocr_api["api_path"] = "/v1/chat/completions"
    ocr_api["api_mode"] = "openai"
    ocr_api["verify_ssl"] = False
    page_loader = pipeline.setdefault("page_loader", {})
    page_loader["max_tokens"] = GLMOCR_MAX_OUTPUT_TOKENS

    layout = pipeline.setdefault("layout", {})
    label_task_mapping = layout.setdefault("label_task_mapping", {})
    text_labels = set(label_task_mapping.get("text", []) or [])
    abandon_labels = set(label_task_mapping.get("abandon", []) or [])
    skip_labels = set(label_task_mapping.get("skip", []) or [])

    text_labels.update({"header", "footer"})
    abandon_labels.discard("header")
    abandon_labels.discard("footer")

    # Keep image-only decorative zones out of OCR by default.
    # If you want logo OCR too, move these to text_labels.
    skip_labels.update({"header_image", "footer_image"})
    abandon_labels.discard("header_image")
    abandon_labels.discard("footer_image")

    label_task_mapping["text"] = sorted(text_labels)
    label_task_mapping["abandon"] = sorted(abandon_labels)
    label_task_mapping["skip"] = sorted(skip_labels)

    with open(config_path, "w", encoding="utf-8") as f:
        yaml.safe_dump(cfg, f, sort_keys=False)


def get_parser():
    global _parser
    if _parser is None:
        _start_vllm_if_needed()
        _configure_glmocr_to_keep_header_footer()
        from glmocr import GlmOcr

        # Force local self-hosted pipeline and avoid MaaS fallback.
        _parser = GlmOcr(mode="selfhosted")
    return _parser


def _extract_md(result: Any) -> str:
    if result is None:
        return ""
    if isinstance(result, list):
        chunks = []
        for item in result:
            md = getattr(item, "markdown_result", "")
            if md:
                chunks.append(str(md).strip())
        return "\n\n---page-separator---\n\n".join([c for c in chunks if c]).strip()
    md = getattr(result, "markdown_result", "")
    return str(md).strip() if md else ""


def run_ocr(file_obj):
    if file_obj is None:
        return "Please upload a file."
    path = file_obj.name if hasattr(file_obj, "name") else str(file_obj)
    page_images: List[str] = []
    try:
        parser = get_parser()
        is_pdf = path.lower().endswith(".pdf")
        if is_pdf:
            page_images = _render_pdf_pages_to_images(path)
            # Parse one rendered page per request item so we can preserve page boundaries.
            result = parser.parse(page_images)
        else:
            result = parser.parse(path)
        md = _extract_md(result) or "(No content)"
        return md
    except Exception as e:
        import traceback

        log.exception("run_ocr failed: %s", e)
        return f"Error: {e}\n\n{traceback.format_exc()}"
    finally:
        for p in page_images:
            try:
                if isinstance(p, str) and p.endswith(".png") and "glmocr_page_" in os.path.basename(p):
                    os.unlink(p)
            except Exception:
                pass


def build_demo():
    with gr.Blocks(title="GLM-OCR Self-hosted (Header/Footer enabled)") as demo:
        gr.Markdown("# GLM-OCR Self-hosted (Header/Footer enabled)")
        gr.Markdown(
            "Runs a local vLLM server inside the Space and configures GLM-OCR "
            "to include `header` and `footer` regions in OCR output."
        )
        file_in = gr.File(
            label="Upload PDF or image",
            file_types=[".pdf", ".png", ".jpg", ".jpeg", ".tiff", ".bmp"],
        )
        run_btn = gr.Button("Run OCR", variant="primary")
        out_md = gr.Textbox(label="Markdown", lines=30)
        run_btn.click(fn=run_ocr, inputs=file_in, outputs=out_md)
    return demo


if __name__ == "__main__":
    build_demo().launch()