Upload 4 files

.gitignore

@@ -0,0 +1,4 @@
+__pycache__/
+*.pyc
+.gradio/
+/tmp/

README.md

@@ -1,16 +1,81 @@
 ---
 title: Chandra OCR 2
-emoji: 💬
-colorFrom: yellow
+emoji: 📄
+colorFrom: indigo
 colorTo: purple
 sdk: gradio
-sdk_version: 6.5.1
 app_file: app.py
 pinned: false
-hf_oauth: true
-hf_oauth_scopes:
-- inference-api
 license: apache-2.0
+short_description: Layout-aware document OCR to markdown, HTML, or JSON
+models:
+  - datalab-to/chandra-ocr-2
+tags:
+  - ocr
+  - document-ai
+  - vision-language
+suggested_hardware: zero-a10g
 ---
 
-An example chatbot using [Gradio](https://gradio.app), [`huggingface_hub`](https://huggingface.co/docs/huggingface_hub/v0.22.2/en/index), and the [Hugging Face Inference API](https://huggingface.co/docs/api-inference/index).
+# 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.
+
+## 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 |
+| 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.
+- Output is offered three ways: rendered markdown, markdown source, and the raw
+  model string (useful when a prompt type returns HTML or JSON instead).
+
+The app prefers the official `chandra` package (`generate_hf` + `BatchInputItem`
++ `parse_markdown`). If that import fails, it falls back to driving the chat
+template through plain `transformers` so the Space still boots.
+
+## Configuration
+
+Edit the constants at the top of `app.py`:
+
+| Constant | Default | Purpose |
+|---|---|---|
+| `PROMPT_TYPES` | `ocr_layout`, … | Prompt types offered in the dropdown |
+| `MAX_PAGES` | 20 | Per-run page cap |
+| `PAGES_PER_GPU_CALL` | 3 | Pages per ZeroGPU allocation |
+| `GPU_DURATION` | 180 | Seconds requested per allocation |
+
+Only `ocr_layout` is confirmed from the model card Quickstart. The others are
+exposed on the assumption that the card's markdown/HTML/JSON output modes map to
+prompt types; if one errors, remove it or check the `chandra` package docs.
+
+## 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

@@ -1,69 +1,412 @@
+"""
+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.
+The model is ~20 GB in bf16 and will NOT run on the free CPU tier.
+"""
+
+import inspect
+import json
+import os
+import time
+import zipfile
+from pathlib import Path
+
 import gradio as gr
-from huggingface_hub import InferenceClient
-
-
-def respond(
-    message,
-    history: list[dict[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-    hf_token: gr.OAuthToken,
-):
-    """
-    For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
+import torch
+from PIL import Image
+
+# ---------------------------------------------------------------------------
+# Config
+# ---------------------------------------------------------------------------
+
+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
+
+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"}
+
+# ---------------------------------------------------------------------------
+# 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:
+    import spaces
+
+    gpu = spaces.GPU
+except ImportError:  # local / non-ZeroGPU deploy
+
+    def gpu(*args, **kwargs):
+        if args and callable(args[0]):
+            return args[0]
+
+        def deco(fn):
+            return fn
+
+        return deco
+
+
+# ---------------------------------------------------------------------------
+# Chandra package — preferred path. Falls back to plain transformers.
+# ---------------------------------------------------------------------------
+
+CHANDRA_ERR = None
+try:
+    from chandra.model import generate_hf
+    from chandra.model.schema import BatchInputItem
+
+    try:
+        from chandra.output import parse_markdown
+    except ImportError:
+        from chandra.model.output import parse_markdown
+
+    USE_CHANDRA = True
+except Exception as e:  # noqa: BLE001
+    USE_CHANDRA = False
+    CHANDRA_ERR = f"{type(e).__name__}: {e}"
+
+
+# ---------------------------------------------------------------------------
+# Model
+# ---------------------------------------------------------------------------
+
+def _load():
+    from transformers import AutoProcessor
+
+    try:
+        from transformers import AutoModelForImageTextToText as VLM
+    except ImportError:
+        from transformers import AutoModelForVision2Seq as VLM
+
+    # 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)
+    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
+        m = VLM.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16, **kw)
+
+    if ON_ZERO:
+        m = m.to("cuda")
+
+    m.eval()
+    p = AutoProcessor.from_pretrained(MODEL_ID)
+    p.tokenizer.padding_side = "left"   # required by chandra's batched generate
+    m.processor = p
+    return m, p
+
+
+print(f"Loading {MODEL_ID} ...")
+_t0 = time.time()
+model, processor = _load()
+print(f"Loaded in {time.time() - _t0:.0f}s | chandra pkg: {USE_CHANDRA} ({CHANDRA_ERR or 'ok'})")
+
+
+# ---------------------------------------------------------------------------
+# Page extraction
+# ---------------------------------------------------------------------------
+
+def as_path(f) -> Path:
+    """Normalise str / Path / 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.
     """
-    client = InferenceClient(token=hf_token.token, model="openai/gpt-oss-20b")
+    if isinstance(f, (str, os.PathLike)):
+        return Path(f)
+    return Path(getattr(f, "name", str(f)))
+
+
+def pdf_to_images(path: Path, dpi: int):
+    try:
+        import pymupdf
+    except ImportError:
+        import fitz as pymupdf
 
-    messages = [{"role": "system", "content": system_message}]
+    doc = pymupdf.open(str(path))
+    pages = []
+    for i, page in enumerate(doc):
+        pix = page.get_pixmap(dpi=dpi)
+        pages.append((f"{path.stem}_p{i + 1:03d}",
+                      Image.frombytes("RGB", (pix.width, pix.height), pix.samples)))
+    doc.close()
+    return pages
 
-    messages.extend(history)
 
-    messages.append({"role": "user", "content": message})
+def collect_pages(files, dpi: int):
+    pages = []
+    for f in files:
+        p = as_path(f)
+        ext = p.suffix.lower()
+        if ext == ".pdf":
+            pages.extend(pdf_to_images(p, dpi))
+        elif ext in IMAGE_EXTS:
+            pages.append((p.stem, Image.open(p).convert("RGB")))
+    return pages
 
-    response = ""
 
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        choices = message.choices
-        token = ""
-        if len(choices) and choices[0].delta.content:
-            token = choices[0].delta.content
+def fit(img: Image.Image, max_side: int) -> Image.Image:
+    """Cap the long edge — visual token count scales with area, so this is the
+    single biggest lever on VRAM and latency."""
+    img = img.convert("RGB")
+    if max(img.size) > max_side:
+        s = max_side / max(img.size)
+        img = img.resize((max(1, int(img.width * s)), max(1, int(img.height * s))),
+                         Image.LANCZOS)
+    return img
 
-        response += token
-        yield response
 
+# ---------------------------------------------------------------------------
+# Inference
+# ---------------------------------------------------------------------------
 
+@gpu(duration=GPU_DURATION)
+@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 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)
+        return [getattr(r, "raw", None) or getattr(r, "markdown", "") or str(r)
+                for r in results]
+
+    # ---- fallback: drive the chat template directly ----
+    outs = []
+    for im in images:
+        msgs = [{"role": "user", "content": [
+            {"type": "image", "image": im},
+            {"type": "text", "text": prompt_type},
+        ]}]
+        inputs = processor.apply_chat_template(
+            msgs, tokenize=True, add_generation_prompt=True,
+            return_dict=True, return_tensors="pt",
+        ).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)
+        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())
+    return outs
+
+
+def to_markdown(raw: str) -> str:
+    if USE_CHANDRA:
+        try:
+            return parse_markdown(raw)
+        except Exception:  # noqa: BLE001
+            pass
+    return raw
+
+
+# ---------------------------------------------------------------------------
+# Orchestration
+# ---------------------------------------------------------------------------
+
+def run(files, prompt_type, dpi, max_side, max_new_tokens,
+        progress=gr.Progress()):
+
+    def msg(m):
+        return "", "", "", m, None, None
+
+    if not files:
+        yield msg("Upload a PDF or some images first.")
+        return
+
+    try:
+        pages = collect_pages(files, int(dpi))
+    except Exception as e:  # noqa: BLE001
+        yield msg(f"Could not read those files:\n{type(e).__name__}: {e}")
+        return
+
+    if not pages:
+        yield msg("No PDF or image files found in that upload.")
+        return
+
+    truncated = ""
+    if len(pages) > MAX_PAGES:
+        truncated = f"  (truncated from {len(pages)})"
+        pages = pages[:MAX_PAGES]
+
+    md_parts, raw_parts, per_page = [], [], []
+    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]
+
+        progress(start / len(pages),
+                 desc=f"{names[0]} … ({start + 1}-{start + len(chunk)}/{len(pages)})")
+
+        try:
+            raws = _infer_chunk(imgs, prompt_type, int(max_new_tokens))
+        except torch.cuda.OutOfMemoryError:
+            torch.cuda.empty_cache()
+            raws = ["[OUT OF MEMORY — lower 'Max image side']"] * len(chunk)
+        except Exception as e:  # noqa: BLE001
+            raws = [f"[FAILED: {type(e).__name__}: {e}]"] * len(chunk)
+
+        for name, raw in zip(names, raws):
+            md = to_markdown(raw)
+            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))
+
+        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)",
+               None, None)
+
+    # ---- artefacts ----
+    stamp = time.strftime("%Y%m%d_%H%M%S")
+    md_path = OUT_DIR / f"chandra_{stamp}.md"
+    md_path.write_text("\n".join(md_parts), encoding="utf-8")
+
+    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:
+            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)
+        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],
+        }, indent=2))
+
+    total = time.time() - t_start
+    joined = "\n".join(md_parts)
+    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)",
+           str(md_path), str(zip_path))
+
+
+# ---------------------------------------------------------------------------
+# UI
+# ---------------------------------------------------------------------------
+
+GR_MAJOR = int(gr.__version__.split(".")[0])
+
+
+def C(cls, **kw):
+    """Build a component, dropping kwargs this Gradio version rejects.
+
+    Gradio 6 removed Textbox.show_copy_button and moved theme/css from
+    Blocks() to launch(). This keeps one file working across 4/5/6.
+    """
+    try:
+        allowed = set(inspect.signature(cls.__init__).parameters)
+        if "kwargs" not in allowed:
+            kw = {k: v for k, v in kw.items() if k in allowed}
+    except (TypeError, ValueError):
+        pass
+    return cls(**kw)
+
+
+CSS = """
+#raw_out textarea { font-family: ui-monospace, monospace; font-size: 12px; }
+#status textarea  { font-family: ui-monospace, monospace; font-size: 12px; }
+.md_pane { max-height: 640px; overflow-y: auto; }
 """
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
+
+_STYLE = dict(theme=gr.themes.Soft(), css=CSS)
+_BLOCKS_KW = {} if GR_MAJOR >= 6 else _STYLE
+_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
+
+Layout-aware OCR from [Datalab](https://datalab.to). Handles tables, math,
+forms, handwriting and 90+ languages, preserving reading order and structure.
+
+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.
+
+Model: [`{MODEL_ID}`](https://huggingface.co/{MODEL_ID}) · weights are under a
+modified OpenRAIL-M licence (free for research, personal use, and companies
+under $2M funding/revenue — **not** for building a competitor to Datalab's API).
 """
-chatbot = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
-    ],
-)
-
-with gr.Blocks() as demo:
-    with gr.Sidebar():
-        gr.LoginButton()
-    chatbot.render()
+    )
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            files = C(gr.Files, label="PDFs / images",
+                      file_types=[".pdf", ".png", ".jpg", ".jpeg", ".webp",
+                                  ".bmp", ".tif", ".tiff"],
+                      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.")
+
+            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.")
+                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.")
+                max_new = C(gr.Slider, minimum=512, maximum=8192, value=4096, step=256,
+                            label="Max new tokens",
+                            info="Layout output is verbose — keep this high.")
+
+            status = C(gr.Textbox, label="Status", lines=3, elem_id="status")
+
+        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)
+                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.Row():
+                md_file = C(gr.File, label="Combined .md")
+                zip_file = C(gr.File, label="All pages .zip")
+
+    go.click(
+        run,
+        inputs=[files, prompt_type, dpi, max_side, max_new],
+        outputs=[md_view, md_src, raw_view, status, md_file, zip_file],
+    )
 
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.queue(max_size=12).launch(show_error=True, **_LAUNCH_KW)

requirements.txt

@@ -0,0 +1,11 @@
+# Chandra OCR 2 Space
+# torch is preinstalled on ZeroGPU/GPU Spaces — do not pin it here.
+
+chandra-ocr
+transformers>=4.57.0
+accelerate>=0.34.0
+huggingface_hub>=0.26.0
+pymupdf>=1.24.0
+pillow>=10.0.0
+gradio>=5.0.0
+spaces