Update app.py

app.py

@@ -1,87 +1,184 @@
 #!/usr/bin/env python3
 
-import asyncio
+import json
 import logging
 import os
-from typing import Any
+import subprocess
+import time
+from typing import Any, List
 
-try:
-    _orig_close = asyncio.BaseEventLoop.close
+import gradio as gr
+import httpx
+import yaml
 
-    def _safe_close(self):
+log = logging.getLogger("glmocr_selfhosted_space")
+logging.basicConfig(level=logging.INFO)
+
+MODEL_ID = os.getenv("GLMOCR_MODEL_ID", "zai-org/GLM-OCR")
+OCR_API_PORT = int(os.getenv("OCR_API_PORT", "8080"))
+OCR_API_BASE = f"http://127.0.0.1:{OCR_API_PORT}"
+VLLM_MAX_MODEL_LEN = os.getenv("VLLM_MAX_MODEL_LEN", "8192")
+VLLM_GPU_MEMORY_UTIL = os.getenv("VLLM_GPU_MEMORY_UTIL", "0.92")
+VLLM_EXTRA_ARGS = os.getenv("VLLM_EXTRA_ARGS", "").strip()
+
+_parser = None
+_vllm_proc = None
+
+
+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:
-            _orig_close(self)
-        except (ValueError, OSError):
-            pass
+            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}")
 
-    asyncio.BaseEventLoop.close = _safe_close
-except Exception:
-    pass
 
-log = logging.getLogger("glmocr_raw_app")
-logging.basicConfig(level=logging.INFO)
+def _configure_glmocr_to_keep_header_footer() -> None:
+    import glmocr
 
-GLMOCR_API_KEY = "cee1d52dd91a4ab591b3f6e105f8ad89.LgbQTECuzX0zrito"
+    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 {}
 
-_parser = None
+    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
+
+    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
 
-        kwargs = {"mode": "maas"}
-        if GLMOCR_API_KEY:
-            kwargs["api_key"] = GLMOCR_API_KEY
-        _parser = GlmOcr(**kwargs)
+        _parser = GlmOcr()
     return _parser
 
 
-def _extract_markdown(result: Any) -> str:
+def _extract_md(result: Any) -> str:
     if result is None:
         return ""
     if isinstance(result, list):
-        parts = []
+        chunks = []
         for item in result:
             md = getattr(item, "markdown_result", "")
             if md:
-                parts.append(str(md).strip())
-        return "\n\n---page-separator---\n\n".join(p for p in parts if p).strip()
+                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_raw(uploaded_file):
-    if uploaded_file is None:
-        return "Please upload a file."
+def _extract_json(result: Any) -> str:
+    if result is None:
+        return "{}"
+    if isinstance(result, list):
+        payload = [getattr(item, "json_result", None) for item in result]
+        return json.dumps(payload, ensure_ascii=False, indent=2)
+    return json.dumps(getattr(result, "json_result", None), ensure_ascii=False, indent=2)
+
+
+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)
     try:
-        path = uploaded_file.name if hasattr(uploaded_file, "name") else str(uploaded_file)
         parser = get_parser()
         result = parser.parse(path)
-        markdown = _extract_markdown(result)
-        return markdown or "(No content)"
+        md = _extract_md(result) or "(No content)"
+        jr = _extract_json(result)
+        return md, jr
     except Exception as e:
         import traceback
 
-        log.exception("run_ocr_raw failed: %s", e)
-        return f"Error: {e}\n\n{traceback.format_exc()}"
+        log.exception("run_ocr failed: %s", e)
+        return f"Error: {e}\n\n{traceback.format_exc()}", "{}"
 
 
-def _create_gradio_demo():
-    import gradio as gr
-
-    with gr.Blocks(title="GLM-OCR Raw (no fixes)") as demo:
-        gr.Markdown("# GLM-OCR Raw (No Fixes)")
+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 = gr.Textbox(lines=40, label="Raw markdown_result")
-        run_btn.click(fn=run_ocr_raw, inputs=file_in, outputs=out)
+        out_md = gr.Textbox(label="Markdown", lines=30)
+        out_json = gr.Textbox(label="JSON (layout details)", lines=20)
+        run_btn.click(fn=run_ocr, inputs=file_in, outputs=[out_md, out_json])
     return demo
 
 
 if __name__ == "__main__":
-    _create_gradio_demo().launch()
+    build_demo().launch()