#!/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()