Upload 2 files

Readme.md

@@ -0,0 +1,105 @@
+---
+title: Chandra OCR 2
+emoji: 📄
+colorFrom: indigo
+colorTo: purple
+sdk: gradio
+app_file: app.py
+pinned: false
+license: apache-2.0
+short_description: Layout-aware document OCR to markdown + extracted figures
+models:
+  - datalab-to/chandra-ocr-2
+tags:
+  - ocr
+  - document-ai
+  - vision-language
+suggested_hardware: zero-a10g
+---
+
+# Chandra OCR 2 — Space demo
+
+Gradio demo for [`datalab-to/chandra-ocr-2`](https://huggingface.co/datalab-to/chandra-ocr-2),
+Datalab's layout-aware document OCR model. Upload PDFs or images and get back
+markdown with tables, math, forms, and reading order preserved — **plus** a
+separate figure-extraction stage that crops charts, diagrams, photos and stamps
+at full resolution and reports the model's caption and any structured data it
+read from them.
+
+## Hardware
+
+**This will not run on the free CPU tier.** The model is ~10B parameters,
+roughly 20 GB in bf16.
+
+| Hardware | Works? |
+|---|---|
+| CPU basic (free) | No — OOM at load |
+| ZeroGPU (H200 slice) | Yes — recommended |
+| RTX PRO 6000 (Blackwell, sm_120) | Yes |
+| L40S / A100 | Yes |
+| T4 (16 GB) | No, unless you add 4-bit quantisation |
+
+Set this under **Settings → Hardware** after creating the Space.
+`suggested_hardware` in the frontmatter is only a hint to visitors; it does not
+provision anything.
+
+## How it works
+
+- PDFs are rasterised page by page with PyMuPDF (no poppler/apt needed).
+- Each page is capped on its long edge before inference — visual token count
+  scales with area, so this is the main lever on latency and memory.
+- Pages are batched `PAGES_PER_GPU_CALL` at a time so each ZeroGPU allocation
+  finishes inside its duration budget.
+- **Stage 1 (text):** raw layout output → markdown via `parse_markdown`.
+- **Stage 2 (figures):** `parse_chunks` locates figure-like blocks, which are
+  cropped from the **full-resolution** source page (the bbox is per-axis
+  normalised, so it maps cleanly onto any resolution with the same aspect
+  ratio). Captions and structured data come from the model's layout output and
+  are labelled **degraded** when the fallback path was used.
+
+Output is offered as rendered markdown, markdown source, the raw model string,
+a figure gallery + metadata table, a combined `.md`, and a `.zip` containing
+per-page markdown, per-page raw output, extracted figure images, and a
+`manifest.json`.
+
+The app prefers the official `chandra` package (`generate_hf` + `BatchInputItem`
++ `parse_markdown` + `parse_chunks`). If that import fails, it falls back to
+driving the chat template through plain `transformers` so the Space still boots
+— and reports the degradation in the status panel and diagnostics.
+
+## Configuration
+
+Edit the constants at the top of `app.py`:
+
+| Constant | Default | Purpose |
+|---|---|---|
+| `PAGES_PER_GPU_CALL` | 3 | Pages per ZeroGPU allocation |
+| `GPU_DURATION` | 180 | Seconds requested per allocation |
+| `MAX_PAGES` | 20 | Per-run page cap |
+
+Prompt types are **enumerated from the package** (`PROMPT_MAPPING`), not
+hardcoded. Only `ocr_layout` and `ocr` are currently exposed; the dropdown is
+editable so you can try undocumented values.
+
+## Troubleshooting
+
+| Symptom | Cause / fix |
+|---|---|
+| Space won't start, opaque "upload failed" | Model load failed at import — the diagnostics panel auto-opens with the exact error. |
+| `CUDA out of memory` per page | Lower **Max image side**; it's the biggest VRAM lever. |
+| No figures in the gallery | The model returned no figure bboxes for this prompt type, or the `chandra` package is unavailable (outputs marked degraded). |
+| `flash-attn` build error | Do **not** add flash-attn — Blackwell (sm_120) has no prebuilt wheels. SDPA is used. |
+| `torch` version mismatch | Do **not** pin torch — the Space image ships a build matched to the driver/CUDA. |
+
+## Licence — read before making this public
+
+The Space **code** here is Apache-2.0. The **model weights** are not:
+
+> Code is Apache 2.0. Model weights are under a modified OpenRAIL-M license.
+> Free for research, personal use, and startups under $2M funding/revenue.
+> Cannot be used competitively with our API.
+
+A public, free, hosted OCR endpoint is plausibly "competitive with our API."
+If you are past the revenue threshold, or intend this as a product rather than a
+demo, check with Datalab first. Setting the Space to **private** avoids the
+question entirely.

app.py

@@ -4,8 +4,25 @@ Chandra OCR 2 — Hugging Face Space demo.
 Model:  datalab-to/chandra-ocr-2  (~10B, bf16)
 Docs:   https://huggingface.co/datalab-to/chandra-ocr-2
 
-Hardware: needs ZeroGPU (H200 slice) or a paid A100/L40S.
+Two pipeline stages:
+  1. Text & layout extraction — all text in reading order, with tables, math,
+     forms and headings preserved as markdown (chandra's parse_markdown).
+  2. Visual understanding — locate figures/charts/diagrams/photos/stamps in the
+     layout output, crop them at full resolution, and surface the model's
+     caption + any structured data it read from them.
+
+Hardware: needs ZeroGPU (H200 slice) or a paid A100/L40S/RTX PRO 6000.
 The model is ~20 GB in bf16 and will NOT run on the free CPU tier.
+
+API facts verified against the installed `chandra` package (not guessed):
+  - generate_hf(batch, model, max_output_tokens=...) -> list[GenerationResult]
+      GenerationResult has .raw, .token_count, .error
+  - PROMPT_MAPPING contains exactly: "ocr_layout", "ocr"
+  - parse_markdown(raw, include_headers_footers=False, include_images=True)
+  - parse_chunks(raw, image, bbox_scale=1000) -> list of
+      {"bbox": [x0, y0, x1, y1] (pixels), "label": str, "content": str}
+  - settings.BBOX_SCALE == 1000 (bbox is normalised per-axis, so it is
+      resolution-independent and can be mapped onto the full-res source image)
 """
 
 import inspect
@@ -25,13 +42,6 @@ from PIL import Image
 
 MODEL_ID = "datalab-to/chandra-ocr-2"
 
-# Chandra is prompt-type driven rather than free-text prompted. 'ocr_layout' is
-# the one shown in the model card Quickstart; the others are exposed because the
-# card advertises markdown / HTML / JSON output. If one errors, the package
-# doesn't support that name in your installed version — check `chandra` docs.
-PROMPT_TYPES = ["ocr_layout", "ocr", "ocr_html", "ocr_json", "layout"]
-DEFAULT_PROMPT_TYPE = "ocr_layout"
-
 PAGES_PER_GPU_CALL = 3      # keep each ZeroGPU allocation inside its duration budget
 GPU_DURATION = 180          # seconds requested per allocation
 MAX_PAGES = 20              # guard against someone uploading a 500-page PDF
@@ -40,12 +50,20 @@ OUT_DIR = Path(os.environ.get("CHANDRA_OUT_DIR", "/tmp/chandra_out"))
 OUT_DIR.mkdir(parents=True, exist_ok=True)
 
 IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".webp", ".bmp", ".tif", ".tiff"}
+PDF_EXTS = {".pdf"}
+
+# Labels the layout model can attach to blocks that count as "figures" for
+# stage 2 (visual understanding). Kept lowercase for matching.
+FIGURE_LABELS = {
+    "figure", "chart", "diagram", "photo", "image", "picture", "stamp",
+    "graph", "plot", "illustration", "logo", "icon", "drawing", "map",
+}
+CAPTION_LABELS = {"caption", "figure-caption", "figcaption", "figure_caption"}
 
 # ---------------------------------------------------------------------------
 # ZeroGPU shim — lets the same file run locally without the `spaces` package
 # ---------------------------------------------------------------------------
 
-# Set by the platform on ZeroGPU hardware.
 ON_ZERO = os.environ.get("SPACES_ZERO_GPU") == "true"
 
 try:
@@ -69,25 +87,40 @@ except ImportError:  # local / non-ZeroGPU deploy
 # ---------------------------------------------------------------------------
 
 CHANDRA_ERR = None
+USE_CHANDRA = False
+HAS_PARSE_CHUNKS = False
+PROMPT_TYPES = ["ocr_layout"]  # safe default; replaced below if package present
+
 try:
     from chandra.model import generate_hf
     from chandra.model.schema import BatchInputItem
+    from chandra.model import PROMPT_MAPPING  # enumerate real prompt types
 
     try:
-        from chandra.output import parse_markdown
+        from chandra.output import parse_markdown, parse_chunks
     except ImportError:
-        from chandra.model.output import parse_markdown
+        from chandra.model.output import parse_markdown, parse_chunks
 
     USE_CHANDRA = True
+    HAS_PARSE_CHUNKS = True
+    PROMPT_TYPES = list(PROMPT_MAPPING.keys())
 except Exception as e:  # noqa: BLE001
-    USE_CHANDRA = False
     CHANDRA_ERR = f"{type(e).__name__}: {e}"
+    USE_CHANDRA = False
+    HAS_PARSE_CHUNKS = False
+
+DEFAULT_PROMPT_TYPE = "ocr_layout" if "ocr_layout" in PROMPT_TYPES else PROMPT_TYPES[0]
 
 
 # ---------------------------------------------------------------------------
-# Model
+# Model — loaded at module scope, guarded so a failure never kills the Space
 # ---------------------------------------------------------------------------
 
+MODEL = None
+PROCESSOR = None
+MODEL_ERROR = None
+
+
 def _load():
     from transformers import AutoProcessor
 
@@ -99,13 +132,15 @@ def _load():
     # On ZeroGPU there is no GPU visible at import time, so accelerate's
     # device_map="auto" would strand the model on CPU. ZeroGPU instead
     # intercepts .to("cuda") at global scope. Elsewhere, device_map is fine.
-    kw = dict(low_cpu_mem_usage=True)
+    # Blackwell (sm_120) -> SDPA, never flash-attn (no prebuilt wheels).
+    kw = dict(low_cpu_mem_usage=True, attn_implementation="sdpa")
     if not ON_ZERO:
         kw["device_map"] = "auto"
 
     try:
         m = VLM.from_pretrained(MODEL_ID, dtype=torch.bfloat16, **kw)
     except TypeError:  # transformers < 4.56 spelled it torch_dtype
+        kw.pop("attn_implementation", None)
         m = VLM.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16, **kw)
 
     if ON_ZERO:
@@ -118,40 +153,46 @@ def _load():
     return m, p
 
 
-model = processor = None
-MODEL_ERR = None
-
-print(f"Loading {MODEL_ID} ...")
-_t0 = time.time()
 try:
-    model, processor = _load()
-    print(f"Loaded in {time.time() - _t0:.0f}s")
+    print(f"Loading {MODEL_ID} ...")
+    _t0 = time.time()
+    MODEL, PROCESSOR = _load()
+    print(f"Loaded in {time.time() - _t0:.0f}s | chandra pkg: {USE_CHANDRA} "
+          f"({CHANDRA_ERR or 'ok'})")
 except Exception as e:  # noqa: BLE001
-    # Do NOT re-raise. A crash here kills the container and every interaction
-    # afterwards fails with an opaque error, including the file picker.
-    import traceback
-    MODEL_ERR = f"{type(e).__name__}: {e}"
-    traceback.print_exc()
-    print(f"!! MODEL FAILED TO LOAD: {MODEL_ERR}")
+    MODEL_ERROR = f"{type(e).__name__}: {e}"
+    print(f"MODEL LOAD FAILED: {MODEL_ERROR}")
+
+MODEL_LOADED = MODEL is not None
 
-print(f"chandra pkg: {USE_CHANDRA} ({CHANDRA_ERR or 'ok'})")
 
+# ---------------------------------------------------------------------------
+# Diagnostics — surfaced in the UI, not just logs
+# ---------------------------------------------------------------------------
 
-def diagnostics() -> str:
-    import platform
+def _transformers_version():
     try:
-        import transformers as _tf
-        tfv = _tf.__version__
+        import transformers
+        return transformers.__version__
     except Exception:  # noqa: BLE001
-        tfv = "?"
-    gpu_name = torch.cuda.get_device_name(0) if torch.cuda.is_available() else "none visible"
-    return (
-        f"python {platform.python_version()} | torch {torch.__version__} | "
-        f"transformers {tfv} | gradio {gr.__version__}\n"
-        f"ZeroGPU: {ON_ZERO} | cuda now: {torch.cuda.is_available()} ({gpu_name})\n"
-        f"chandra package: {'ok' if USE_CHANDRA else CHANDRA_ERR}\n"
-        f"model: {'loaded' if model is not None else 'FAILED — ' + str(MODEL_ERR)}"
-    )
+        return "n/a"
+
+
+def diagnostics_md() -> str:
+    cuda = torch.cuda.is_available()
+    gpu_name = torch.cuda.get_device_name(0) if cuda else "none visible"
+    rows = [
+        ("torch", torch.__version__),
+        ("transformers", _transformers_version()),
+        ("gradio", gr.__version__),
+        ("CUDA visible", str(cuda)),
+        ("GPU", gpu_name),
+        ("chandra package", "imported" if USE_CHANDRA else f"FAILED — {CHANDRA_ERR}"),
+        ("prompt types", ", ".join(PROMPT_TYPES)),
+        ("model loaded", "yes" if MODEL_LOADED else "NO"),
+        ("model error", MODEL_ERROR or "—"),
+    ]
+    return "\n".join(f"- **{k}**: `{v}`" for k, v in rows)
 
 
 # ---------------------------------------------------------------------------
@@ -159,13 +200,15 @@ def diagnostics() -> str:
 # ---------------------------------------------------------------------------
 
 def as_path(f) -> Path:
-    """Normalise str / Path / Gradio file object into a real Path.
+    """Normalise str / os.PathLike / Gradio file object into a real Path.
 
-    Note pathlib.Path also has .name, but there it's the basename — checking
-    hasattr first would silently drop the directory.
+    pathlib.Path also has a .name attribute, but there it is the *basename* —
+    a hasattr(f, "name") check would silently drop the directory. Handle the
+    types distinctly instead.
     """
     if isinstance(f, (str, os.PathLike)):
         return Path(f)
+    # Gradio file object (or tempfile.NamedTemporaryFile) exposes .name as a path
     return Path(getattr(f, "name", str(f)))
 
 
@@ -186,21 +229,22 @@ def pdf_to_images(path: Path, dpi: int):
 
 
 def collect_pages(files, dpi: int):
-    """Returns (pages, skipped). Never raises on a single bad file."""
-    pages, skipped = [], []
+    """Return (pages, rejected) where pages is [(name, full_res_PIL)] and
+    rejected is [(name, reason)]. One bad file never fails the run."""
+    pages, rejected = [], []
     for f in files:
         p = as_path(f)
         ext = p.suffix.lower()
         try:
-            if ext == ".pdf":
+            if ext in PDF_EXTS:
                 pages.extend(pdf_to_images(p, dpi))
             elif ext in IMAGE_EXTS:
                 pages.append((p.stem, Image.open(p).convert("RGB")))
             else:
-                skipped.append(f"{p.name} (unsupported type '{ext or 'none'}')")
+                rejected.append((p.name, f"unsupported type '{ext or 'none'}'"))
         except Exception as e:  # noqa: BLE001
-            skipped.append(f"{p.name} ({type(e).__name__}: {e})")
-    return pages, skipped
+            rejected.append((p.name, f"{type(e).__name__}: {e}"))
+    return pages, rejected
 
 
 def fit(img: Image.Image, max_side: int) -> Image.Image:
@@ -222,37 +266,32 @@ def fit(img: Image.Image, max_side: int) -> Image.Image:
 @torch.inference_mode()
 def _infer_chunk(images, prompt_type: str, max_new_tokens: int):
     """OCR a small batch of PIL images. Returns a list of raw model strings."""
-    if model is None:
-        raise RuntimeError(
-            f"Model never loaded: {MODEL_ERR}. Check the Container logs and "
-            f"confirm Settings > Hardware is ZeroGPU or a GPU tier."
-        )
+    if MODEL is None:
+        raise RuntimeError("model not loaded")
 
     if USE_CHANDRA:
         batch = [BatchInputItem(image=im, prompt_type=prompt_type) for im in images]
-        try:
-            results = generate_hf(batch, model, max_tokens=max_new_tokens)
-        except TypeError:
-            results = generate_hf(batch, model)
+        results = generate_hf(batch, MODEL, max_output_tokens=max_new_tokens)
         return [getattr(r, "raw", None) or getattr(r, "markdown", "") or str(r)
                 for r in results]
 
-    # ---- fallback: drive the chat template directly ----
+    # ---- fallback: drive the chat template directly (degraded) ----
     outs = []
     for im in images:
         msgs = [{"role": "user", "content": [
             {"type": "image", "image": im},
             {"type": "text", "text": prompt_type},
         ]}]
-        inputs = processor.apply_chat_template(
+        inputs = PROCESSOR.apply_chat_template(
             msgs, tokenize=True, add_generation_prompt=True,
             return_dict=True, return_tensors="pt",
-        ).to(model.device)
+        ).to(MODEL.device)
         if "pixel_values" in inputs:
-            inputs["pixel_values"] = inputs["pixel_values"].to(model.dtype)
-        gen = model.generate(**inputs, max_new_tokens=int(max_new_tokens), do_sample=False)
+            inputs["pixel_values"] = inputs["pixel_values"].to(MODEL.dtype)
+        gen = MODEL.generate(**inputs, max_new_tokens=int(max_new_tokens),
+                             do_sample=False)
         trimmed = [o[len(i):] for i, o in zip(inputs["input_ids"], gen)]
-        outs.append(processor.batch_decode(trimmed, skip_special_tokens=True)[0].strip())
+        outs.append(PROCESSOR.batch_decode(trimmed, skip_special_tokens=True)[0].strip())
     return outs
 
 
@@ -266,39 +305,111 @@ def to_markdown(raw: str) -> str:
 
 
 # ---------------------------------------------------------------------------
-# Orchestration
+# Stage 2 — figure extraction (visual understanding)
+# ---------------------------------------------------------------------------
+
+def _find_caption(chunks, fig_idx: int):
+    """Look for a caption chunk immediately after (or before) the figure."""
+    for j in (fig_idx + 1, fig_idx - 1):
+        if 0 <= j < len(chunks):
+            lbl = (chunks[j].get("label") or "").lower()
+            if lbl in CAPTION_LABELS:
+                return chunks[j].get("content") or ""
+    return ""
+
+
+def extract_figures(raw: str, full_res: Image.Image):
+    """Locate figure-like blocks in the layout output and crop them from the
+    FULL-RESOLUTION source image (bbox is per-axis normalised, so it maps
+    cleanly onto any resolution with the same aspect ratio).
+
+    Returns (figures, degraded) where figures is a list of dicts and degraded
+    is True when we could not use the real chandra path (so captions/data are
+    NOT from the model's figure understanding).
+    """
+    if not USE_CHANDRA or not HAS_PARSE_CHUNKS:
+        return [], True
+
+    try:
+        chunks = parse_chunks(raw, full_res, bbox_scale=1000)
+    except Exception:  # noqa: BLE001
+        return [], True
+
+    figures = []
+    for i, ch in enumerate(chunks):
+        label = (ch.get("label") or "").lower()
+        if label not in FIGURE_LABELS:
+            continue
+        bbox = ch.get("bbox")
+        if not bbox or len(bbox) != 4:
+            continue
+        x0, y0, x1, y1 = (int(v) for v in bbox)
+        w, h = full_res.size
+        x0, y0 = max(0, x0), max(0, y0)
+        x1, y1 = min(w, x1), min(h, y1)
+        if x1 <= x0 or y1 <= y0:
+            continue
+        crop = full_res.crop((x0, y0, x1, y1))
+        figures.append({
+            "page": None,          # filled by caller
+            "label": label,
+            "bbox": [x0, y0, x1, y1],
+            "crop": crop,
+            "caption": _find_caption(chunks, i),
+            "structured": (ch.get("content") or "").strip(),
+        })
+    return figures, False
+
+
+# ---------------------------------------------------------------------------
+# Orchestration — streaming generator
 # ---------------------------------------------------------------------------
 
-def run(files, prompt_type, dpi, max_side, max_new_tokens,
-        progress=gr.Progress()):
+FIG_HEADERS = ["Page", "Label", "BBox (x0,y0,x1,y1)", "Caption", "Structured data"]
+
 
-    def msg(m):
-        return "", "", "", m, None, None
+def run(files, prompt_type, dpi, max_side, max_new_tokens, progress=gr.Progress()):
+
+    def empty(status):
+        return "", "", "", status, None, [], None, None
+
+    if not MODEL_LOADED:
+        yield empty(f"Model failed to load — see diagnostics.\n{MODEL_ERROR}")
+        return
 
     if not files:
-        yield msg("Upload a PDF or some images first.")
+        yield empty("Upload a PDF or some images first.")
         return
 
-    pages, skipped = collect_pages(files, int(dpi))
-    note = ("\nSkipped: " + "; ".join(skipped)) if skipped else ""
+    try:
+        pages, rejected = collect_pages(files, int(dpi))
+    except Exception as e:  # noqa: BLE001
+        yield empty(f"Could not read those files:\n{type(e).__name__}: {e}")
+        return
 
     if not pages:
-        yield msg("Nothing readable in that upload." + note +
-                  "\n\nSupported: .pdf, .png, .jpg, .jpeg, .webp, .bmp, .tif, .tiff")
+        yield empty("No readable PDF or image files found in that upload.")
         return
 
-    truncated = note
+    skip_note = ""
+    if rejected:
+        skip_note = "\nSkipped:\n" + "\n".join(f"  - {n}: {r}" for n, r in rejected)
+
+    truncated = ""
     if len(pages) > MAX_PAGES:
-        truncated += f"  (truncated from {len(pages)} pages)"
+        truncated = f"  (truncated from {len(pages)})"
         pages = pages[:MAX_PAGES]
 
-    md_parts, raw_parts, per_page = [], [], []
+    md_parts, raw_parts = [], []
+    all_figs = []          # flattened figure records for the gallery/table
+    per_page = []          # (name, md, raw, figs)
     t_start = time.time()
 
     for start in range(0, len(pages), PAGES_PER_GPU_CALL):
         chunk = pages[start:start + PAGES_PER_GPU_CALL]
         names = [n for n, _ in chunk]
-        imgs = [fit(im, int(max_side)) for _, im in chunk]
+        full_res = [im for _, im in chunk]
+        imgs = [fit(im, int(max_side)) for im in full_res]
 
         progress(start / len(pages),
                  desc=f"{names[0]} … ({start + 1}-{start + len(chunk)}/{len(pages)})")
@@ -311,18 +422,25 @@ def run(files, prompt_type, dpi, max_side, max_new_tokens,
         except Exception as e:  # noqa: BLE001
             raws = [f"[FAILED: {type(e).__name__}: {e}]"] * len(chunk)
 
-        for name, raw in zip(names, raws):
+        for name, fimg, raw in zip(names, full_res, raws):
             md = to_markdown(raw)
+            figs, degraded = extract_figures(raw, fimg)
+            for fg in figs:
+                fg["page"] = name
+                fg["degraded"] = degraded
+                all_figs.append(fg)
+
             md_parts.append(f"\n\n---\n\n## {name}\n\n{md}")
             raw_parts.append(f"===== {name} =====\n{raw}")
-            per_page.append((name, md, raw))
+            per_page.append((name, md, raw, figs))
 
         elapsed = time.time() - t_start
         joined = "\n".join(md_parts)
         yield (joined, joined,
                "\n\n".join(raw_parts),
                f"{len(per_page)}/{len(pages)} pages{truncated} · {elapsed:.0f}s "
-               f"({elapsed / max(1, len(per_page)):.1f}s/page)",
+               f"({elapsed / max(1, len(per_page)):.1f}s/page){skip_note}",
+               _gallery(all_figs), _fig_rows(all_figs),
                None, None)
 
     # ---- artefacts ----
@@ -332,16 +450,25 @@ def run(files, prompt_type, dpi, max_side, max_new_tokens,
 
     zip_path = OUT_DIR / f"chandra_{stamp}.zip"
     with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as z:
-        for name, md, raw in per_page:
+        for name, md, raw, figs in per_page:
             safe = "".join(c if c.isalnum() or c in "-_." else "_" for c in name)
             z.writestr(f"markdown/{safe}.md", md)
             z.writestr(f"raw/{safe}.txt", raw)
+            for k, fg in enumerate(figs):
+                buf = _to_png_bytes(fg["crop"])
+                if buf:
+                    z.writestr(f"figures/{safe}_{k:02d}_{fg['label']}.png", buf)
         z.writestr("manifest.json", json.dumps({
             "model": MODEL_ID,
             "prompt_type": prompt_type,
             "dpi": dpi,
             "max_side": max_side,
-            "pages": [n for n, _, _ in per_page],
+            "pages": [n for n, _, _, _ in per_page],
+            "figures": [
+                {"page": f["page"], "label": f["label"], "bbox": f["bbox"],
+                 "caption": f["caption"], "structured": f["structured"]}
+                for f in all_figs
+            ],
         }, indent=2))
 
     total = time.time() - t_start
@@ -349,12 +476,47 @@ def run(files, prompt_type, dpi, max_side, max_new_tokens,
     yield (joined, joined,
            "\n\n".join(raw_parts),
            f"Done — {len(per_page)} page(s){truncated} in {total:.0f}s "
-           f"({total / max(1, len(per_page)):.1f}s/page)",
+           f"({total / max(1, len(per_page)):.1f}s/page){skip_note}",
+           _gallery(all_figs), _fig_rows(all_figs),
            str(md_path), str(zip_path))
 
 
+def _to_png_bytes(img: Image.Image):
+    import io
+    buf = io.BytesIO()
+    try:
+        img.save(buf, format="PNG")
+        return buf.getvalue()
+    except Exception:  # noqa: BLE001
+        return None
+
+
+def _gallery(figs):
+    """Gallery entries: (crop, caption). Caption is clearly marked degraded when
+    it did not come from the model's figure-understanding path."""
+    out = []
+    for f in figs:
+        cap = f["caption"] or "(no caption from model)"
+        if f.get("degraded"):
+            cap = f"[degraded — not from model] {cap}"
+        out.append((f["crop"], f"{f['page']} · {f['label']} · {cap}"))
+    return out or None
+
+
+def _fig_rows(figs):
+    rows = []
+    for f in figs:
+        cap = f["caption"] or ""
+        if f.get("degraded"):
+            cap = f"[degraded] {cap}".strip()
+        rows.append([f["page"], f["label"],
+                     f"({f['bbox'][0]},{f['bbox'][1]},{f['bbox'][2]},{f['bbox'][3]})",
+                     cap, f["structured"]])
+    return rows
+
+
 # ---------------------------------------------------------------------------
-# UI
+# UI — Gradio 4/5/6 tolerant
 # ---------------------------------------------------------------------------
 
 GR_MAJOR = int(gr.__version__.split(".")[0])
@@ -388,11 +550,16 @@ _LAUNCH_KW = _STYLE if GR_MAJOR >= 6 else {}
 with gr.Blocks(title="Chandra OCR 2", **_BLOCKS_KW) as demo:
     gr.Markdown(
         f"""
-# Chandra OCR 2 — document → markdown / HTML / JSON
+# Chandra OCR 2 — document → markdown / HTML / JSON + figures
 
 Layout-aware OCR from [Datalab](https://datalab.to). Handles tables, math,
 forms, handwriting and 90+ languages, preserving reading order and structure.
 
+Two outputs per page:
+1. **Text & layout** — markdown with tables/math/forms/headings preserved.
+2. **Figures** — charts, diagrams, photos and stamps cropped at full resolution,
+   with the model's caption and any structured data it read from them.
+
 Upload **PDFs and/or images**; each page is processed separately and results
 stream in below. Capped at **{MAX_PAGES} pages** per run in this demo.
 
@@ -404,20 +571,19 @@ under $2M funding/revenue — **not** for building a competitor to Datalab's API
 
     with gr.Row():
         with gr.Column(scale=1):
-            # No file_types filter on purpose: it silently rejects valid files
-            # on some Gradio/browser combinations. Validation happens in
-            # collect_pages() instead, which can give a useful message.
+            # NOTE: no file_types filter — it silently rejects valid files on
+            # some browser/version combinations. Validation happens in Python.
             files = C(gr.Files, label="PDFs / images", file_count="multiple")
             go = C(gr.Button, value="Run OCR", variant="primary")
 
             prompt_type = C(gr.Dropdown, label="Prompt type", choices=PROMPT_TYPES,
                             value=DEFAULT_PROMPT_TYPE, allow_custom_value=True,
-                            info="'ocr_layout' is the documented default.")
+                            info="Enumerated from the chandra package; editable "
+                                 "so you can try undocumented values.")
 
             with gr.Accordion("Advanced", open=False):
                 dpi = C(gr.Slider, minimum=100, maximum=400, value=200, step=25,
-                        label="PDF render DPI",
-                        info="200-300 suits most scans.")
+                        label="PDF render DPI", info="200-300 suits most scans.")
                 max_side = C(gr.Slider, minimum=768, maximum=2560, value=1540, step=64,
                              label="Max image side (px)",
                              info="Biggest lever on speed and VRAM.")
@@ -425,25 +591,28 @@ under $2M funding/revenue — **not** for building a competitor to Datalab's API
                             label="Max new tokens",
                             info="Layout output is verbose — keep this high.")
 
-            status = C(gr.Textbox, label="Status", lines=3, elem_id="status",
-                       value=("" if model is not None else
-                              f"MODEL NOT LOADED — {MODEL_ERR}\n"
-                              f"Check Settings > Hardware (needs ZeroGPU or a GPU tier)."))
+            status = C(gr.Textbox, label="Status", lines=3, elem_id="status")
 
-            with gr.Accordion("Diagnostics", open=model is None):
-                C(gr.Textbox, label=None, lines=5, value=diagnostics(),
-                  interactive=False)
+            # Diagnostics — auto-open when the model failed to load.
+            diag = gr.Accordion("Diagnostics", open=not MODEL_LOADED)
+            with diag:
+                gr.Markdown(diagnostics_md())
 
         with gr.Column(scale=2):
             with gr.Tabs():
                 with gr.Tab("Rendered"):
                     md_view = C(gr.Markdown, value="", elem_classes=["md_pane"])
                 with gr.Tab("Markdown source"):
-                    md_src = C(gr.Textbox, label=None, lines=24,
-                               show_copy_button=True)
+                    md_src = C(gr.Textbox, label=None, lines=24, show_copy_button=True)
                 with gr.Tab("Raw model output"):
                     raw_view = C(gr.Textbox, label=None, lines=24,
                                  elem_id="raw_out", show_copy_button=True)
+            with gr.Tab("Figures"):
+                fig_gallery = C(gr.Gallery, label="Extracted figures",
+                                columns=3, height="auto",
+                                object_fit="contain")
+                fig_table = C(gr.Dataframe, headers=FIG_HEADERS,
+                              label="Figure metadata", interactive=False)
             with gr.Row():
                 md_file = C(gr.File, label="Combined .md")
                 zip_file = C(gr.File, label="All pages .zip")
@@ -451,9 +620,10 @@ under $2M funding/revenue — **not** for building a competitor to Datalab's API
     go.click(
         run,
         inputs=[files, prompt_type, dpi, max_side, max_new],
-        outputs=[md_view, md_src, raw_view, status, md_file, zip_file],
+        outputs=[md_view, md_src, raw_view, status, fig_gallery, fig_table,
+                 md_file, zip_file],
     )
 
 
 if __name__ == "__main__":
-    demo.queue(max_size=12).launch(show_error=True, **_LAUNCH_KW)
+    demo.queue(max_size=12).launch(show_error=True, **_LAUNCH_KW)