Upload app.py with huggingface_hub

app.py

@@ -1,136 +1,207 @@
-
 import gradio as gr
 import yaml
 import re
+import json
 import os
 import torch
 
-# Paths from installed glmocr package (required on HF Space)
 import glmocr
-GLMOCR_BASE = os.path.dirname(glmocr.__file__)
-config_path = os.path.join(GLMOCR_BASE, "config.yaml")
+GLMOCR_BASE    = os.path.dirname(glmocr.__file__)
+config_path    = os.path.join(GLMOCR_BASE, "config.yaml")
 formatter_path = os.path.join(GLMOCR_BASE, "postprocess", "result_formatter.py")
 
-# ── STEP 1: Fix config — keep header & footer (do NOT add them to abandon) ──
+# ── STEP 1: Fix config.yaml ───────────────────────────────────
 with open(config_path, "r") as f:
     config = yaml.safe_load(f)
+
 config["pipeline"]["result_formatter"]["abandon"] = [
-    "number", "footnote", "aside_text", "reference",
-    "footer_image", "header_image",
+    "number", "footnote", "aside_text",
+    "reference", "footer_image", "header_image",
 ]
 config["pipeline"]["enable_layout"] = True
+
 with open(config_path, "w") as f:
     yaml.dump(config, f, default_flow_style=False, sort_keys=False)
-print("✅ config.yaml fixed (header & footer kept in output)")
 
-# ── STEP 2: Fix result_formatter.py (remove hardcoded header/footer) ───────
+print("✅ config.yaml fixed")
+
+# ── STEP 2: Fix result_formatter.py ──────────────────────────
 with open(formatter_path, "r") as f:
     source = f.read()
+
 labels_to_remove = [
-    '"header"', "'header'", '"footer"', "'footer'",
-    '"doc_header"', "'doc_header'", '"doc_footer"', "'doc_footer'"
+    '"header"', "'header'",
+    '"footer"', "'footer'",
+    '"doc_header"', "'doc_header'",
+    '"doc_footer"', "'doc_footer'"
 ]
 for label in labels_to_remove:
     source = re.sub(r',\s*' + re.escape(label), '', source)
     source = re.sub(re.escape(label) + r'\s*,', '', source)
     source = re.sub(re.escape(label), '', source)
+
 with open(formatter_path, "w") as f:
     f.write(source)
+
 print("✅ result_formatter.py fixed")
 
-# ── STEP 3: Load model ────────────────────────────────────────────────────
-from transformers import AutoProcessor, GlmOcrForConditionalGeneration
-print("Loading model... (~2GB first run)")
-processor = AutoProcessor.from_pretrained("zai-org/GLM-OCR")
-model = GlmOcrForConditionalGeneration.from_pretrained(
-    "zai-org/GLM-OCR",
-    torch_dtype=torch.bfloat16,
-    device_map="auto",
-)
-print("✅ Model ready on", next(model.parameters()).device)
-ABANDON = set(config["pipeline"]["result_formatter"]["abandon"])
-
-# ── STEP 4: OCR (PDF → image then run model) ───────────────────────────────
+# ── STEP 3: Lazy model load ───────────────────────────────────
+processor = None
+model     = None
+ABANDON   = set(config["pipeline"]["result_formatter"]["abandon"])
+
+def load_model():
+    global processor, model
+    if model is not None:
+        return
+    from transformers import AutoProcessor, GlmOcrForConditionalGeneration
+    print("Loading model...")
+    processor = AutoProcessor.from_pretrained("zai-org/GLM-OCR")
+    model = GlmOcrForConditionalGeneration.from_pretrained(
+        "zai-org/GLM-OCR",
+        torch_dtype=torch.bfloat16,
+        device_map="auto",
+    )
+    print(f"✅ Model ready on {next(model.parameters()).device}")
+
+# ── STEP 4: OCR one image ─────────────────────────────────────
+def ocr_image(img_path):
+    messages = [
+        {"role": "user", "content": [
+            {"type": "image", "url": img_path},
+            {"type": "text",  "text": "Document Parsing:"}
+        ]}
+    ]
+    inputs = processor.apply_chat_template(
+        messages, tokenize=True, add_generation_prompt=True,
+        return_dict=True, return_tensors="pt"
+    ).to(model.device)
+    inputs.pop("token_type_ids", None)
+
+    with torch.no_grad():
+        output_ids = model.generate(**inputs, max_new_tokens=2048)
+
+    raw = processor.decode(
+        output_ids[0][inputs["input_ids"].shape[1]:],
+        skip_special_tokens=False
+    )
+
+    for token in ["<|assistant|>", "<|user|>", "<|system|>",
+                  "<|endoftext|>", "</s>", "<s>"]:
+        raw = raw.replace(token, "")
+    raw = raw.strip()
+
+    json_match = re.search(r'\[.*\]', raw, re.DOTALL)
+    if json_match:
+        try:
+            regions = json.loads(json_match.group())
+            return regions, raw, "json"
+        except:
+            pass
+
+    return [], raw, "raw"
+
+# ── STEP 5: Main OCR function ─────────────────────────────────
 def run_ocr(uploaded_file):
     if uploaded_file is None:
         return "Please upload a file.", "No regions detected."
+
     try:
+        load_model()
+
         path = uploaded_file.name if hasattr(uploaded_file, "name") else str(uploaded_file)
+
         if path.lower().endswith(".pdf"):
-            try:
-                import fitz
-                doc = fitz.open(path)
-                page = doc[0]
-                pix = page.get_pixmap(matrix=fitz.Matrix(1, 1), alpha=False)
-                img_path = path[:-4] + "_page0.png"
+            import fitz
+            doc = fitz.open(path)
+            page_images = []
+            for i in range(len(doc)):
+                pix      = doc[i].get_pixmap(matrix=fitz.Matrix(1.5, 1.5), alpha=False)
+                img_path = f"/tmp/page_{i}.png"
                 pix.save(img_path)
-                doc.close()
-                path = img_path
-            except Exception as e:
-                return "PDF conversion failed: " + str(e), "Failed."
-
-        messages = [
-            {"role": "user", "content": [
-                {"type": "image", "url": path},
-                {"type": "text", "text": "Document Parsing:"}
-            ]}
-        ]
-        inputs = processor.apply_chat_template(
-            messages, tokenize=True, add_generation_prompt=True,
-            return_dict=True, return_tensors="pt"
-        ).to(model.device)
-        inputs.pop("token_type_ids", None)
-        with torch.no_grad():
-            output_ids = model.generate(**inputs, max_new_tokens=2048)
-        raw = processor.decode(
-            output_ids[0][inputs["input_ids"].shape[1]:],
-            skip_special_tokens=False
-        )
-        raw = raw.replace("<|user|>", "").strip()
-        json_match = re.search(r'\[.*\]', raw, re.DOTALL)
-        regions = json.loads(json_match.group()) if json_match else []
-        header_count = footer_count = 0
-        region_lines = []
-        markdown_parts = []
-        for region in regions:
-            label = region.get("label", "text")
-            content = str(region.get("content", ""))
-            if label in ABANDON:
-                continue
-            if label == "header":
-                header_count += 1
-                region_lines.append("🔵 HEADER:\n" + content + "\n")
-                markdown_parts.append("<!-- HEADER -->\n" + content)
-            elif label == "footer":
-                footer_count += 1
-                region_lines.append("🟢 FOOTER:\n" + content + "\n")
-                markdown_parts.append("<!-- FOOTER -->\n" + content)
+                page_images.append(img_path)
+            doc.close()
+        else:
+            page_images = [path]
+
+        total_headers = 0
+        total_footers = 0
+        all_markdown  = []
+        all_summary   = []
+
+        for page_num, img_path in enumerate(page_images):
+            regions, raw, output_type = ocr_image(img_path)
+
+            page_md      = []
+            page_summary = [f"── PAGE {page_num + 1} of {len(page_images)} ──"]
+
+            if output_type == "raw":
+                page_md.append(raw)
+                page_summary.append("[raw HTML/text output]")
+                page_summary.append(f"Content length: {len(raw)} chars")
             else:
-                region_lines.append("[" + label + "]: " + content[:150])
-                markdown_parts.append(content)
+                for region in regions:
+                    label   = region.get("label", "text")
+                    content = str(region.get("content", ""))
+
+                    if label in ABANDON:
+                        continue
+
+                    if label == "header":
+                        total_headers += 1
+                        page_summary.append(f"🔵 HEADER: {content[:100]}")
+                        page_md.append(f"<!-- HEADER -->\n{content}")
+                    elif label == "footer":
+                        total_footers += 1
+                        page_summary.append(f"🟢 FOOTER: {content[:100]}")
+                        page_md.append(f"<!-- FOOTER -->\n{content}")
+                    else:
+                        page_summary.append(f"[{label}]: {content[:100]}")
+                        page_md.append(content)
+
+            all_markdown.append("\n\n".join(page_md))
+            all_summary.extend(page_summary)
+            all_summary.append("")
+
         summary = (
-            "Headers found : " + str(header_count) + "\n"
-            "Footers found : " + str(footer_count) + "\n"
-            "Total regions : " + str(len(regions)) + "\n" + "─"*40 + "\n"
-            + "\n".join(region_lines)
+            f"Total pages   : {len(page_images)}\n"
+            f"Headers found : {total_headers}\n"
+            f"Footers found : {total_footers}\n"
+            f"{'─'*40}\n"
+            + "\n".join(all_summary)
         )
-        markdown = "\n\n".join(markdown_parts) if markdown_parts else raw
+
+        markdown = "\n\n---\n\n".join(all_markdown)
         return markdown, summary
+
     except Exception as e:
         import traceback
-        return "Error: " + str(e) + "\n\n" + traceback.format_exc(), "Failed."
+        return f"Error: {str(e)}\n\n{traceback.format_exc()}", "Failed."
 
-# ── STEP 5: Gradio UI ──────────────────────────────────────────────────────
-with gr.Blocks(title="GLM-OCR — Header & Footer Kept") as demo:
-    gr.Markdown("# 🔍 GLM-OCR — Header & Footer Kept\nUpload PDF or image. Headers 🔵 and Footers 🟢 are kept in output.")
-    file_input = gr.File(label="Upload PDF or Image", file_types=[".pdf", ".png", ".jpg", ".jpeg", ".tiff", ".bmp"])
+# ── STEP 6: Gradio UI ─────────────────────────────────────────
+with gr.Blocks(title="GLM-OCR — Header & Footer Fix") as demo:
+    gr.Markdown("""
+    # 🔍 GLM-OCR — Header & Footer Fix
+    Upload PDF or image. Headers 🔵 and Footers 🟢 now appear in output.
+    Multi-page PDFs fully supported.
+    > First request takes ~2 min to load model. After that it is fast.
+    """)
+
+    file_input = gr.File(
+        label="Upload PDF or Image",
+        file_types=[".pdf", ".png", ".jpg", ".jpeg", ".tiff", ".bmp"]
+    )
     run_btn = gr.Button("▶  Run OCR", variant="primary", size="lg")
+
     with gr.Row():
         with gr.Column():
             gr.Markdown("### 📄 Markdown Output")
-            markdown_out = gr.Textbox(lines=25, label="")
+            markdown_out = gr.Textbox(lines=30, label="")
         with gr.Column():
             gr.Markdown("### 🗂️ Detected Regions")
-            regions_out = gr.Textbox(lines=25, label="")
-    run_btn.click(fn=run_ocr, inputs=file_input, outputs=[markdown_out, regions_out])
+            regions_out = gr.Textbox(lines=30, label="")
+
+    run_btn.click(fn=run_ocr, inputs=file_input,
+                  outputs=[markdown_out, regions_out])
+
 demo.launch()