tiiuae

/

Falcon-OCR

Like 175

Follow

Technology Innovation Institute 2.23k

Problem

_load_layout_model() guards on self._layout_model, an attribute it never
assigns, so the early return is unreachable and generate_with_layout() reloads
the PP-DocLayoutV3 detector (weights + image processor) from disk on every
call.

Current code (modeling_falcon_ocr.py, lines 485–496 on main):

def _load_layout_model(self, layout_model: str = "PaddlePaddle/PP-DocLayoutV3_safetensors"):
    if hasattr(self, "_layout_model"):          # never true
        return
    ...
    self._layout_processor = PPDocLayoutV3ImageProcessorFast.from_pretrained(layout_model)
    self._layout_det_model = AutoModelForObjectDetection.from_pretrained(...)
    self._layout_id2label = self._layout_det_model.config.id2label
    self._tvF = tvF

grep -n "_layout_model" modeling_falcon_ocr.py returns three hits: the
parameter name, the guard on 486, and the call site on 756. Nothing assigns it.
generate_with_layout() calls _load_layout_model() unconditionally (line 756),
so the detector is re-instantiated on every invocation.

Reproduction

import torch
from transformers import AutoModelForCausalLM
from PIL import Image

model = AutoModelForCausalLM.from_pretrained(
    "tiiuae/Falcon-OCR", trust_remote_code=True,
    torch_dtype=torch.bfloat16, device_map="cuda",
)

img = Image.open("page.png").convert("RGB")
for _ in range(3):
    with torch.inference_mode():
        model.generate_with_layout([img])

Before: three separate Loading weights: 100%|...| 858/858 progress bars,
plus repeated HEAD https://huggingface.co/.../config.json and
preprocessor_config.json requests — one full PP-DocLayoutV3 load per iteration.

After: one load; iterations 2 and 3 reuse it and emit nothing further.

Impact

• The full detector is re-read from disk and re-allocated on the GPU per call.
• Hub metadata round-trips on every call — fragile under rate limiting, awkward offline.
• For multi-page pipelines calling generate_with_layout() once per page, this
  load dominates per-call cost and makes per-page timing measurements meaningless.

The change

Assign the attribute the guard checks, keyed on the requested layout_model so
a caller passing a different detector on a later call still gets a reload:

-        if hasattr(self, "_layout_model"):
+        if getattr(self, "_layout_model", None) == layout_model:
             return
...
         self._tvF = tvF
+        self._layout_model = layout_model

_layout_model is assigned last, so a load that raises partway through
leaves the guard false and the next call retries cleanly.

No device handling is needed: _layout_det_model is assigned to an nn.Module
attribute, so it is registered as a child module and model.to(...) moves the
cached detector along with the rest of the model (verified below).

Verified

On transformers 5.9.0, torch 2.12.0+cu130, revision
ccc96cb86610dc5976f786bacd7bd8a5bbe567e2, counting
AutoModelForObjectDetection.from_pretrained calls and comparing detector
object identity across calls:

• Before: 3 calls to _load_layout_model() → 3 loads, 3 distinct detector objects.
• After: 3 calls → 1 load, the same detector object returned each time.
• '_layout_det_model' in model._modules is True — it is a registered submodule.
• After model.to("cpu"): the cached detector is the same object, now on cpu, and
  the guard correctly does not reload. _run_layout_detection() runs fine on cpu
  and again after moving back to cuda.
• Not run: the repository's own test suite. Device moves were exercised
  cuda→cpu→cuda on a single GPU, not across two GPUs.