Granite for Docling (#49012)
* feat: add GraniteForDocling
Add the GraniteForDocling vision-language model for document conversion.
Pages are tiled at 512x512 with a pixel-shuffle connector, DeepStack
injection of intermediate vision features, and a dense Granite-style
text decoder that generates DocLang.
The same modeling code loads different GraniteForDocling sizes and
configurations via text_config and vision_config, and different
projectors (high-resolution or density router). The fine path can be
selected per request with fine_route=True; checkpoints trained with the
coarse path only set use_fine_route=False and do not build it. Optional
multi-token prediction heads (num_mtp_layers) add an auxiliary training
loss. Serving engines such as vLLM can use the same heads as draft heads
for speculative decoding; generate does not.
The model is written as a modular file; the configuration, image
processors (torchvision and PIL backends), processor and modeling files
are generated from it. The image processors tile pages on the grid
closest to their aspect ratio, capped at 16 tiles per side to match the
tokenizer's tile position markers, and pad batches on the tile dimension.
Co-authored-by: Matteo Omenetti <omenetti.matteo@gmail.com>
* Fix GraniteForDocling MTP and DeepStack feature reuse (#3)
Keep multi-token prediction spans based on the original sequence length so later heads retain valid positions.
Preserve projected DeepStack features when reusing image embeddings, including cached and beam generation. Resolve the fake image token ID after special-token registration and preserve DocLang markup in the pipeline example.
Regenerate the standalone model and processor from the modular source and add regression coverage. The per-call image_seq_len override remains a documented known issue.
* docs: drop redundant dtype/attn_implementation and padding_side tips
Review feedback: the config already sets the dtype and sdpa is the default;
padding_side=left ships in the checkpoint's tokenizer config instead.
* Address review: vision model owns the connector, standard Granite MLP, no MTP module
- The connector and the DeepStack projections move into GraniteForDoclingVisionModel,
which returns the image tokens of all tiles in tile order plus the DeepStack
features (qwen3-vl layout); position embeddings and projections are modules.
- The decoder uses the standard Granite MLP; the fused checkpoint weights are
split and the connector renamed through the conversion mapping.
- The multi-token prediction heads are no longer built or loaded by transformers
(serving engines use them); their config keys stay for those engines.
- Drop the image feature reuse inputs, the runtime special-token registration,
the processing_utils change and the manual gradient-checkpointing warning.
- Image processors tile grouped by shape, keep fine_route as a saved attribute
and return rows/cols without tensor conversion.
* Native checkpoint layout, no MTP config, modeling-structure fixes
- Checkpoints ship in the model's own layout (gate/up/down projections,
connector under the vision model), so no conversion mapping is needed.
- Drop the multi-token prediction config fields; heads stored under mtp.*
are ignored on load.
- Position embedding buffer as an attribute, no single-linear wrapper module,
@auto_docstring on the vision model forward.
- Docs: drop the processor_kwargs note and the MTP mentions.
* Image processor: reuse GotOcr2's tiled-canvas selection
Only the grid enumeration stays local: it caps the tiles per side at 16, the
largest tile position marker in the tokenizer. The cap has to be part of the
enumeration because crop_image_to_patches selects its grid through the same
function; a clamp at the call site would tile with one grid and label the
prompt with another, and lowering max_image_tiles to 16 would forbid grids
such as 4x6 or 2x16.
* Image processors: inherit GotOcr2's patch count, mirror the PIL backend on the torchvision one
get_number_of_image_patches is GotOcr2's; the tile grid the processor needs
for the prompt comes from a small get_tile_grid. The PIL _preprocess follows
the torchvision one step by step (per-image loop, since grouping by shape is
a torchvision utility), padding on the tile dimension the same way.
* DeepStack injection as in Qwen3-VL, do_pad for tile padding
- The text model adds DeepStack feature i after decoder layer i through
_deepstack_process, as Qwen3VLTextModel does; deepstack_attn_layers is gone
(the released checkpoints used layers 0..n-1, so outputs are unchanged).
- Both image processors pad samples to the same number of tiles through a
do_pad-guarded pad(), with all-zero tiles the model discards.
* Address review: check sub-configs, accept precomputed vision features, and group connector routes.
* Address review: keep connector settings on the vision config and simplify the forward path.
* Address review: reuse tile padding for the fine mask and cover beam expansion.
* Rename the model to GraniteDocling and finish the review.
The model type and classes are granite_docling / GraniteDocling. The connector is a coarse path and a fine path, vision tiles use Siglip embeddings, and the image processor only caps GotOCR's tile grid. The fine route is one bool per sample, with padding kept on the pixel attention mask.
* Point the GraniteDocling docs at the IBM Granite Hub org.
The Docling project created the models. The published weights and the example ids live under ibm-granite.
* Add the spread rule to the image processor and update for main.
A checkpoint can now carry a spread rule in its image processor config: `dense_max_patches` tiles instead of
`max_patches` for a wide landscape page (`dense_min_aspect` < width/height <= `dense_max_aspect`) that the
`max_patches` grid would shrink below `dense_max_scale`. Off by default. `get_page_tile_grid` applies it, so
cropping, the prompt's image tokens and `get_number_of_image_patches` agree.
After merging main: regenerated from the modular file (type annotations only in the modeling file), the processor's
`__call__` has its own docstring (the inherited one documented an argument that is not in the signature), and the
model card date is refreshed.
* Address review: one connector route per call, GotOcr2's crop, Hub test fixtures.
- `get_image_features` and `predict_fine_route` run one vision pass on the real tiles. The processor never mixes
routes in a batch, so a batch takes one connector path: fine when the processor returned `tile_fine_mask`.
`pixel_attention_mask` is required with `pixel_values`; the all-zero-tile fallback and the `encoder_outputs`
reuse are gone.
- Each connector path is a module with its own forward; the connector keeps the shared norms and picks the path.
Same state dict keys.
- The image processors call GotOcr2's `crop_image_to_patches` instead of overriding it. `get_page_max_patches`
gives it the tile budget of the page (the spread rule's when it applies), which is the budget `get_tile_grid`
uses, so the tiles and the prompt agree.
- `return_mm_token_type_ids=True` works: the processor uses `ProcessorMixin.create_mm_token_type_ids`.
- Tests load `hf-internal-testing/tiny-processor-granite-docling` and `ibm-granite/granite-for-docling-500m`
instead of local paths; docs and code examples name the `ibm-granite` checkpoint.
- The tester passes `pixel_attention_mask` and takes one route; the mixed-route tests are replaced by a
padding-tile test and a processor-built padded batch in the integration tests.
`use_fine_route` and the `dense_*` spread rule stay: the 500M checkpoint has the fine path and the 2B does not,
and the 2B image processor config sets the spread rule.
* Regenerate the processor after merging main.
The input validation inherited from Idefics3 changed upstream; the generated processor follows it. The two
sub-model entries in `IGNORE_NON_TESTED` are gone: `check_repo` now recognises models built from a sub-config
as backbones.
* sync with changes on main
* fix
* fix
* skip
* fixes
* last one?
---------
Co-authored-by: Matteo Omenetti <omenetti.matteo@gmail.com>
Co-authored-by: Matteo <43417658+Matteo-Omenetti@users.noreply.github.com>
Co-authored-by: Ahmed Nassar AHN@zurich.ibm.com <AHN@zurich.ibm.com>
Co-authored-by: vasqu <antonprogamer@gmail.com>
