fix: honor picture description batching and scale options (#3132)

* fix: honor picture description batching and scale options

Signed-off-by: Hassan Raza <raihassanraza10@gmail.com>

* fix: address picture description review feedback

Signed-off-by: Hassan Raza <raihassanraza10@gmail.com>

* test: fix picture description vlm init on py314

Signed-off-by: Hassan Raza <raihassanraza10@gmail.com>

* test: simplify picture description vlm init stub

Signed-off-by: Hassan Raza <raihassanraza10@gmail.com>

---------

Signed-off-by: Hassan Raza <raihassanraza10@gmail.com>

Author: M-Hassan-Raza

docling/datamodel/pipeline_options.py

@@ -536,19 +536,21 @@ class PictureDescriptionBaseOptions(BaseOptions):
     batch_size: Annotated[
         int,
         Field(
+            ge=1,
             description=(
                 "Number of images to process in a single batch during picture description. Higher values improve "
                 "throughput but increase memory usage. Adjust based on available GPU/CPU memory."
-            )
+            ),
         ),
     ] = 8
     scale: Annotated[
         float,
         Field(
+            gt=0,
             description=(
                 "Scaling factor for image resolution before processing. Higher values (e.g., 2.0) provide more detail "
                 "for the vision model but increase processing time and memory. Range: 0.5-4.0 typical."
-            )
+            ),
         ),
     ] = 2.0
     picture_area_threshold: Annotated[
@@ -715,6 +717,15 @@ class PictureDescriptionVlmOptions(PictureDescriptionBaseOptions):
             )
         ),
     ] = {"max_new_tokens": 200, "do_sample": False}
+    padding_side: Annotated[
+        Literal["left", "right"],
+        Field(
+            description=(
+                "Tokenizer padding side used for batched generation. Defaults to left to preserve the legacy "
+                "behavior, but can be overridden for models that require right padding."
+            )
+        ),
+    ] = "left"
 
     @property
     def repo_cache_folder(self) -> str:

docling/models/picture_description_base_model.py

@@ -39,9 +39,16 @@ def __init__(
         options: PictureDescriptionBaseOptions,
         accelerator_options: AcceleratorOptions,
     ):
+        if options.batch_size < 1:
+            raise ValueError("Picture description batch_size must be >= 1")
+        if options.scale <= 0:
+            raise ValueError("Picture description scale must be > 0")
+
         self.enabled = enabled
         self.options = options
         self.provenance = "not-implemented"
+        self.elements_batch_size = options.batch_size
+        self.images_scale = options.scale
 
     def is_processable(self, doc: DoclingDocument, element: NodeItem) -> bool:
         return self.enabled and isinstance(element, PictureItem)

docling/models/stages/picture_description/picture_description_vlm_model.py

@@ -57,7 +57,6 @@ def __init__(
                 import torch
                 from transformers import (
                     AutoModelForImageTextToText,
-                    AutoModelForVision2Seq,
                     AutoProcessor,
                 )
             except ImportError:
@@ -68,6 +67,9 @@ def __init__(
             # Initialize processor and model
             with _model_init_lock:
                 self.processor = AutoProcessor.from_pretrained(artifacts_path)
+                tokenizer = getattr(self.processor, "tokenizer", None)
+                if tokenizer is not None:
+                    tokenizer.padding_side = self.options.padding_side
                 self.model = AutoModelForImageTextToText.from_pretrained(
                     artifacts_path,
                     device_map=self.device,
@@ -89,6 +91,10 @@ def __init__(
     def _annotate_images(self, images: Iterable[Image.Image]) -> Iterable[str]:
         from transformers import GenerationConfig
 
+        image_batch = list(images)
+        if not image_batch:
+            return
+
         # Create input messages
         messages = [
             {
@@ -100,24 +106,25 @@ def _annotate_images(self, images: Iterable[Image.Image]) -> Iterable[str]:
             },
         ]
 
-        # TODO: do batch generation
-
-        for image in images:
-            # Prepare inputs
-            prompt = self.processor.apply_chat_template(
-                messages, add_generation_prompt=True
-            )
-            inputs = self.processor(text=prompt, images=[image], return_tensors="pt")
-            inputs = inputs.to(self.device)
-
-            # Generate outputs
-            generated_ids = self.model.generate(
-                **inputs,
-                generation_config=GenerationConfig(**self.options.generation_config),
-            )
-            generated_texts = self.processor.batch_decode(
-                generated_ids[:, inputs["input_ids"].shape[1] :],
-                skip_special_tokens=True,
-            )
-
-            yield generated_texts[0].strip()
+        prompt = self.processor.apply_chat_template(
+            messages, add_generation_prompt=True
+        )
+        inputs = self.processor(
+            text=[prompt] * len(image_batch),
+            images=image_batch,
+            return_tensors="pt",
+            padding=True,
+        )
+        inputs = inputs.to(self.device)
+
+        generated_ids = self.model.generate(
+            **inputs,
+            generation_config=GenerationConfig(**self.options.generation_config),
+        )
+        generated_texts = self.processor.batch_decode(
+            generated_ids[:, inputs["input_ids"].shape[1] :],
+            skip_special_tokens=True,
+        )
+
+        for text in generated_texts:
+            yield text.strip()

tests/test_picture_description_base_model.py

@@ -0,0 +1,179 @@
+from collections.abc import Iterable
+from types import SimpleNamespace
+from typing import ClassVar, List, Type
+
+import pytest
+from docling_core.types.doc import (
+    DoclingDocument,
+    ImageRef,
+    PictureItem,
+    ProvenanceItem,
+)
+from docling_core.types.doc.base import BoundingBox, Size
+from PIL import Image
+
+from docling.datamodel.accelerator_options import AcceleratorOptions
+from docling.datamodel.base_models import ItemAndImageEnrichmentElement
+from docling.datamodel.pipeline_options import (
+    PictureDescriptionBaseOptions,
+    PictureDescriptionVlmEngineOptions,
+    PipelineOptions,
+)
+from docling.models.picture_description_base_model import PictureDescriptionBaseModel
+from docling.pipeline.base_pipeline import BasePipeline
+
+
+class _TestOptions(PictureDescriptionBaseOptions):
+    kind: ClassVar[str] = "test"
+
+
+class _ConfiguredPictureDescriptionModel(PictureDescriptionBaseModel):
+    def __init__(self, options: PictureDescriptionBaseOptions) -> None:
+        super().__init__(
+            enabled=True,
+            enable_remote_services=False,
+            artifacts_path=None,
+            options=options,
+            accelerator_options=AcceleratorOptions(),
+        )
+
+    @classmethod
+    def get_options_type(cls) -> Type[PictureDescriptionBaseOptions]:
+        return _TestOptions
+
+    def _annotate_images(self, images: Iterable[Image.Image]) -> Iterable[str]:
+        for _image in images:
+            yield "test description"
+
+
+class _BatchRecordingPictureDescriptionModel(_ConfiguredPictureDescriptionModel):
+    def __init__(self, options: PictureDescriptionBaseOptions) -> None:
+        super().__init__(options)
+        self.batch_sizes: List[int] = []
+
+    def __call__(
+        self,
+        doc: DoclingDocument,
+        element_batch: Iterable[ItemAndImageEnrichmentElement],
+    ) -> Iterable[PictureItem]:
+        element_list = list(element_batch)
+        self.batch_sizes.append(len(element_list))
+        for element in element_list:
+            assert isinstance(element.item, PictureItem)
+            yield element.item
+
+
+class _PictureDescriptionPipeline(BasePipeline):
+    def _build_document(self, conv_res):
+        return conv_res
+
+    def _determine_status(self, conv_res):
+        return conv_res.status
+
+    @classmethod
+    def get_default_options(cls) -> PipelineOptions:
+        return PipelineOptions()
+
+    @classmethod
+    def is_backend_supported(cls, backend) -> bool:
+        return True
+
+
+def _make_picture_doc(*, count: int, embed_images: bool = True) -> DoclingDocument:
+    doc = DoclingDocument(name="test")
+    for _ in range(count):
+        image = (
+            ImageRef.from_pil(Image.new("RGB", (20, 20), "red"), dpi=72)
+            if embed_images
+            else None
+        )
+        doc.add_picture(image=image)
+    return doc
+
+
+def test_picture_description_options_control_batch_size_and_scale() -> None:
+    model = _ConfiguredPictureDescriptionModel(_TestOptions(batch_size=3, scale=1.5))
+
+    assert model.elements_batch_size == 3
+    assert model.images_scale == 1.5
+
+
+def test_picture_description_batch_size_controls_pipeline_chunking() -> None:
+    pipeline = _PictureDescriptionPipeline(PipelineOptions())
+    model = _BatchRecordingPictureDescriptionModel(_TestOptions(batch_size=2))
+    pipeline.enrichment_pipe = [model]
+    conv_res = SimpleNamespace(
+        document=_make_picture_doc(count=5),
+        timings={},
+        status="success",
+    )
+
+    pipeline._enrich_document(conv_res)
+
+    assert model.batch_sizes == [2, 2, 1]
+
+
+def test_picture_description_scale_is_used_for_cropping() -> None:
+    model = _ConfiguredPictureDescriptionModel(_TestOptions(scale=1.5))
+    doc = DoclingDocument(name="test")
+    doc.add_page(page_no=1, size=Size(width=100, height=100))
+    picture = doc.add_picture(
+        prov=ProvenanceItem(
+            page_no=1,
+            bbox=BoundingBox(l=10, t=10, r=30, b=30),
+            charspan=(0, 0),
+        )
+    )
+
+    class _PageSpy:
+        def __init__(self):
+            self.page_no = 1
+            self.calls = []
+
+        def get_image(self, *, scale, cropbox):
+            self.calls.append({"scale": scale, "cropbox": cropbox})
+            return Image.new("RGB", (5, 5), "blue")
+
+    page = _PageSpy()
+    conv_res = SimpleNamespace(document=doc, pages=[page])
+
+    prepared = model.prepare_element(conv_res=conv_res, element=picture)
+
+    assert prepared is not None
+    assert page.calls[0]["scale"] == 1.5
+
+
+def test_picture_description_embedded_images_keep_original_size() -> None:
+    model = _ConfiguredPictureDescriptionModel(_TestOptions(scale=1.5))
+    doc = _make_picture_doc(count=1, embed_images=True)
+
+    prepared = model.prepare_element(
+        conv_res=SimpleNamespace(document=doc, pages=[]), element=doc.pictures[0]
+    )
+
+    assert prepared is not None
+    assert prepared.image.size == (20, 20)
+
+
+def test_picture_description_batch_size_must_be_positive() -> None:
+    with pytest.raises(ValueError):
+        _TestOptions(batch_size=0)
+
+
+def test_picture_description_scale_must_be_positive() -> None:
+    with pytest.raises(ValueError):
+        _TestOptions(scale=0)
+
+
+def test_picture_description_preset_batch_size_must_be_positive() -> None:
+    options = PictureDescriptionVlmEngineOptions.from_preset("smolvlm", batch_size=0)
+
+    with pytest.raises(ValueError, match="batch_size"):
+        _ConfiguredPictureDescriptionModel(options)
+
+
+def test_picture_description_preset_scale_must_be_positive() -> None:
+    options = PictureDescriptionVlmEngineOptions.from_preset("smolvlm", scale=0)
+
+    with pytest.raises(ValueError, match="scale"):
+        _ConfiguredPictureDescriptionModel(options)