Model update

README.md

@@ -286,7 +286,9 @@ processor = AutoProcessor.from_pretrained(
 
 # Load the model on the available device(s)
 model = AutoModelForCausalLM.from_pretrained(
-    'jinaai/jina-vlm-v1', device_map='auto', trust_remote_code=True
+    'jinaai/jina-vlm-v1',
+    device_map='auto', 
+    trust_remote_code=True
 )
 
 # You can specify a different model dtype and/or attention implementation

blocks_jvlm.py

@@ -1,6 +1,7 @@
 # Copyright 2025 Jina AI. All rights reserved.
 
 from abc import ABCMeta, abstractmethod
+from contextlib import nullcontext
 from copy import deepcopy
 from functools import wraps
 from math import prod, sqrt
@@ -11,6 +12,7 @@ import torch
 import torch.backends.cuda
 import torch.nn as nn
 import torch.nn.functional as f
+from torch.nn.attention import SDPBackend, sdpa_kernel
 from transformers import PretrainedConfig
 from transformers.activations import ACT2FN
 from transformers.cache_utils import Cache
@@ -324,10 +326,11 @@ modeling_rope_utils.py
 
 
 def inv_freq_to_device(rope_forward):
-    """
-    Sometimes the inv_freq is calculated on the wrong device, or ends up in lower
-    precision than float32. This wrapper ensures that inv_freq is always on the right
-    device and in float32 precision.
+    """Sometimes the inv_freq is calculated on the wrong device, or ends up in lower
+    precision than float32.
+
+    This wrapper ensures that inv_freq is always on the right device and in float32
+    precision.
     """
 
     @wraps(rope_forward)
@@ -353,7 +356,6 @@ class RotaryEmbedding(nn.Module):
         theta: float,
         head_dim: int,
         hidden_size: int,
-        n_heads: int,
         partial_rotary_factor: float,
         device: Optional[torch.device] = None,
         scaling: Optional[Dict[str, Any]] = None,
@@ -366,7 +368,6 @@ class RotaryEmbedding(nn.Module):
         setattr(self.config, 'rope_theta', theta)
         setattr(self.config, 'partial_rotary_factor', partial_rotary_factor)
         setattr(self.config, 'head_dim', head_dim)
-        setattr(self.config, 'num_attention_heads', n_heads)
         setattr(self.config, 'hidden_size', hidden_size)
         setattr(self.config, 'rope_scaling', scaling or {})
 
@@ -377,9 +378,7 @@ class RotaryEmbedding(nn.Module):
         self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
         device = device or torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         seqlen = config.max_position_embeddings or config.max_sequence_length
-        invfreq, self.attention_scaling = self.rope_init_fn(
-            self.config, device, seqlen
-        )
+        invfreq, self.attention_scaling = self.rope_init_fn(self.config, device, seqlen)
         self.rope_init_device = device
         self.register_buffer('inv_freq', invfreq, persistent=False)
         self.original_inv_freq = self.inv_freq
@@ -617,11 +616,9 @@ def _create_causal_mask(seq_len: int, device: torch.device) -> torch.Tensor:
 def _ensure_finite(
     x: torch.Tensor, check_neg_inf: bool = True, check_pos_inf: bool = False
 ):
-    """
-    Modify ``x`` in place to replace ``float("-inf")`` with the minimum value of the
+    """Modify ``x`` in place to replace ``float("-inf")`` with the minimum value of the
     dtype when ``check_neg_inf`` is ``True`` and replace ``float("inf")`` with the
-    maximum value of the dtype when ``check_pos_inf`` is ``True``
-    """
+    maximum value of the dtype when ``check_pos_inf`` is ``True``"""
     if check_neg_inf:
         x.masked_fill_(x == float('-inf'), torch.finfo(x.dtype).min)
     if check_pos_inf:
@@ -641,14 +638,12 @@ def resolve_causal_mask(
         # shape: (batch_size, 1, 1, seq_len)
         if len(attention_mask.shape) == 2:
             attention_mask = attention_mask[:, : past_length + seq_len]
-            attention_mask = attention_mask.to(dtype=torch.float).view(
-                batch_size, -1
-            )[:, None, None, :]
+            attention_mask = attention_mask.to(dtype=torch.float).view(batch_size, -1)[
+                :, None, None, :
+            ]
         else:
             attention_mask = attention_mask.unsqueeze(1).to(dtype=torch.float)
-        attention_mask = (1.0 - attention_mask) * torch.finfo(
-            attention_mask.dtype
-        ).min
+        attention_mask = (1.0 - attention_mask) * torch.finfo(attention_mask.dtype).min
 
     # Merge attention mask with causal mask (attention bias)
     # NOTE: We need to initialize the attn bias in order for attn to
@@ -660,9 +655,7 @@ def resolve_causal_mask(
         or past_key_values is not None
     ):
         if causal_mask is None:
-            causal_mask = _create_causal_mask(
-                past_length + seq_len, device
-            )
+            causal_mask = _create_causal_mask(past_length + seq_len, device)
         elif causal_mask.dtype in (torch.int8, torch.bool):
             causal_mask = causal_mask.to(dtype=torch.float)
             causal_mask.masked_fill_(
@@ -719,6 +712,7 @@ def eager_attention_forward(
     dropout: float = 0.0,
     **_,
 ):
+    assert isinstance(module.num_key_value_groups, int)
     key_states = repeat_kv(key, module.num_key_value_groups)
     value_states = repeat_kv(value, module.num_key_value_groups)
 
@@ -745,7 +739,9 @@ def rotate_half(x: torch.Tensor):
 
 
 def apply_rotary_positional_embeddings(
-    x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor,
+    x: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
 ) -> torch.Tensor:
     return (x * cos + rotate_half(x) * sin).to(x.dtype)
 
@@ -890,7 +886,6 @@ class MHSDPA(nn.Module):
         attn_mask: Optional[torch.Tensor] = None,
         is_causal: Optional[bool] = None,
     ) -> Tuple[Callable, Optional[torch.Tensor], Optional[bool]]:
-
         if 'flash' in attn_implementation and self.fp32_attn:
             raise ValueError('Flash attention does not support fp32 attention')
         if self.sliding_window != -1 and 'flash' not in attn_implementation:
@@ -1071,9 +1066,7 @@ class FFN(nn.Module):
         if self.gated_activation:
             intermediate_size = 2 * self.intermediate_size
 
-        self.up = nn.Linear(
-            self.hidden_size, intermediate_size, bias=self.use_bias
-        )
+        self.up = nn.Linear(self.hidden_size, intermediate_size, bias=self.use_bias)
         self.down = nn.Linear(
             self.intermediate_size, self.output_size, bias=self.use_bias
         )
@@ -1245,6 +1238,8 @@ class VisionLanguageConnector(GradientCheckpointingLayer):
             assert config.attn_pooling_config is not None
             if config.pooling_type == ImagePooling2DType.attention_2wide:
                 pooling_input_size *= 2
+
+            attn_implementation, _ = self._resolve_attn_pooling(attn_implementation)
             self.pooling = MHSDPA(
                 config.attn_pooling_config,
                 hidden_size=pooling_input_size,
@@ -1285,11 +1280,29 @@ class VisionLanguageConnector(GradientCheckpointingLayer):
         self.projector_dropout = Dropout(config.projector_dropout)
         self.feature_dropout = Dropout(config.feature_dropout)
 
+    @staticmethod
+    def _resolve_attn_pooling(attn_implementation: Optional[str] = None):
+        """
+        Flash Attention can cause Inf grads in the attention pooling layer because of
+        very large batch sizes. Setting this to sdpa does not cost us much since
+        sequence lengths in the case of attention pooling are tiny
+        """
+        attn_runtime_ctx = nullcontext()
+        if (
+            attn_implementation is not None
+            and attn_implementation.startswith('flash')
+        ):
+            attn_implementation = 'sdpa'
+            attn_runtime_ctx = sdpa_kernel(backends=[SDPBackend.MATH])
+
+        return attn_implementation, attn_runtime_ctx
+
     def forward(
         self,
         image_features: torch.Tensor,
         image_masks: Optional[torch.Tensor] = None,
         attn_implementation: Optional[str] = None,
+        **kwargs: Unpack[FlashAttentionKwargs],
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
         # image_features:
         # (batch_size, num_crops(=num_image), num_patch, nximage_emb_dim)
@@ -1345,11 +1358,19 @@ class VisionLanguageConnector(GradientCheckpointingLayer):
                 dh=self.pooling_h,
                 dw=self.pooling_w,
             )
+            image_features = image_features.contiguous()
             if self.pooling_type == ImagePooling2DType.attention_meanq:
                 query = image_features.mean(-2, keepdim=True)
-                image_features, _ = self.pooling(
-                    xq=query, xk=image_features, attn_implementation=attn_implementation
+                attn_implementation, attn_runtime_ctx = self._resolve_attn_pooling(
+                    attn_implementation
                 )
+                with attn_runtime_ctx:
+                    image_features, _ = self.pooling(
+                        xq=query,
+                        xk=image_features,
+                        attn_implementation=attn_implementation,
+                        **kwargs,
+                    )
             elif self.pooling_type not in {
                 ImagePooling2DType.none,
                 ImagePooling2DType.stack,

config.json

@@ -4,6 +4,7 @@
   ],
   "auto_map": {
     "AutoConfig": "configuration_jvlm.JinaVLMConfig",
+    "AutoModel": "modeling_jvlm.JinaVLM",
     "AutoModelForCausalLM": "modeling_jvlm.JinaVLMForConditionalGeneration"
   },
   "bos_token_id": 151643,
@@ -214,4 +215,4 @@
       "spatial_merge_size": 2
     }
   }
-}
+}

configuration_jvlm.py

@@ -530,6 +530,11 @@ class JinaVLMTextConfig(PretrainedConfigWithDataclasses):
         self.rope_theta = rope_theta
         self.rope_scaling = rope_scaling
 
+    # Needed for vLLM
+    @property
+    def num_attention_heads(self) -> int:
+        return self.block_config.attn_config.n_heads
+
 
 class JinaVLMConfig(PretrainedConfig):
     """JinaVLM configuration.
@@ -545,7 +550,8 @@ class JinaVLMConfig(PretrainedConfig):
 
     model_type = 'jvlm'
     sub_configs = {
-        'vision_config': JinaVLMVisionConfig, 'text_config': JinaVLMTextConfig
+        'vision_config': JinaVLMVisionConfig,
+        'text_config': JinaVLMTextConfig,
     }
 
     def __init__(

image_processing_jvlm.py

@@ -437,6 +437,17 @@ class JinaVLMImageProcessor(BaseImageProcessor):
 
     """ Base cropping via resizing """
 
+    def base_get_n_image_patches(
+        self,
+        height: int,
+        width: int,
+        max_crops: int,
+    ) -> int:
+        raise NotImplementedError(
+            'Function `get_n_image_patches` is not implemented for cropping method '
+            f'{CroppingMethod.RESIZE}'
+        )
+
     def base_resize_cropping(self, image: np.ndarray):
         resized, mask = self.resize_image(image, list(self.base_input_size))
         resized = self.normalize_image(resized)
@@ -497,6 +508,117 @@ class JinaVLMImageProcessor(BaseImageProcessor):
 
         return candidate_tilings[ix]
 
+    @staticmethod
+    def _molmo_get_patches_from_tiling(
+        num_tiles,
+        pooling_size,
+        crop_patches,
+        crop_window_patches,
+        left_margin,
+        right_margin,
+    ) -> np.int32:
+        if num_tiles > 1:
+            left_crop_window_patches = (
+                (crop_window_patches + left_margin + pooling_size - 1)
+                // pooling_size
+                * pooling_size
+            )
+            middle_crop_window_patches = (
+                (crop_window_patches + pooling_size - 1) // pooling_size * pooling_size
+            )
+            right_crop_window_patches = (
+                (crop_window_patches + right_margin + pooling_size - 1)
+                // pooling_size
+                * pooling_size
+            )
+            return (
+                left_crop_window_patches
+                + (num_tiles - 2) * middle_crop_window_patches
+                + right_crop_window_patches
+            )
+        else:
+            single_crop_window_patches = (
+                (crop_patches + pooling_size - 1) // pooling_size * pooling_size
+            )
+            return single_crop_window_patches
+
+    def molmo_get_n_image_patches(
+        self,
+        height: int,
+        width: int,
+        max_crops: int,
+    ) -> int:
+        # Discard this many patches from the (left/top, right/bottom) of crops
+        left_margin, right_margin = self.overlap_margins
+        # Required for compatibility with image pooling
+        assert left_margin % self.pooling_w == 0 and right_margin % self.pooling_w == 0
+        assert left_margin % self.pooling_h == 0 and right_margin % self.pooling_h == 0
+        # pixels removed per dim
+        total_margin_pixels = self.patch_size * (right_margin + left_margin)
+        # patches per crop dim
+        crop_patches = self.base_input_size[0] // self.patch_size
+
+        # usable patches
+        crop_window_patches = crop_patches - (right_margin + left_margin)
+        crop_window_size = crop_window_patches * self.patch_size
+
+        # We assume hxw pooling, but can allow padding the right/bottom with extra
+        # patches if the number of patches per side is not divisible by h/w
+        assert (
+            crop_patches + self.pooling_h - 1
+        ) // self.pooling_h == self.token_length_h
+        assert (
+            crop_patches + self.pooling_w - 1
+        ) // self.pooling_w == self.token_length_w
+
+        # Decide how to tile the image, to account for the overlap margins we
+        # compute the tiling as if we had an image without the margins and were
+        # using a crop size without the margins
+        tiling = self._molmo_select_tiling(
+            height - total_margin_pixels,
+            width - total_margin_pixels,
+            crop_window_size,
+            max_crops,
+        )
+
+        # Now build the output tokens
+        h = self._molmo_get_patches_from_tiling(
+            tiling[0],
+            self.pooling_h,
+            crop_patches,
+            crop_window_patches,
+            left_margin,
+            right_margin,
+        )
+        w = self._molmo_get_patches_from_tiling(
+            tiling[1],
+            self.pooling_w,
+            crop_patches,
+            crop_window_patches,
+            left_margin,
+            right_margin,
+        )
+        # for each row of patches, add a patch token per patch
+        n_tokens = w.item() // self.pooling_w
+        if self.use_column_tokens:
+            # after each row, one column token is added
+            n_tokens += 1
+        # replicate each row of patch tokens by number of rows, i.e.
+        # proportional to image height
+        n_tokens *= h.item() // self.pooling_h
+        # add start and end image tokens
+        n_tokens += 2
+
+        # Global image goes first, so the order of patches in previous crops gets
+        # increased
+        n_thumbnail_tokens = self.token_length_w
+        if self.use_column_tokens:
+            n_thumbnail_tokens += 1
+        n_thumbnail_tokens *= self.token_length_h
+        n_thumbnail_tokens += 2
+
+        return n_tokens + n_thumbnail_tokens
+
     def molmo_overlap_and_resize_cropping(self, image: np.ndarray):
         # Discard this many patches from the (left/top, right/bottom) of crops
         left_margin, right_margin = self.overlap_margins
@@ -625,37 +747,23 @@ class JinaVLMImageProcessor(BaseImageProcessor):
         # new order into sparse structure of `patch_ordering` to fix it
         patch_ordering[valid] = patch_ordering_rh[patch_ordering_rh >= 0]
 
-        def get_num_patches(num_tiles, pooling_size) -> int:
-            if num_tiles > 1:
-                left_crop_window_patches = (
-                    (crop_window_patches + left_margin + pooling_size - 1)
-                    // pooling_size
-                    * pooling_size
-                )
-                middle_crop_window_patches = (
-                    (crop_window_patches + pooling_size - 1)
-                    // pooling_size
-                    * pooling_size
-                )
-                right_crop_window_patches = (
-                    (crop_window_patches + right_margin + pooling_size - 1)
-                    // pooling_size
-                    * pooling_size
-                )
-                return (
-                    left_crop_window_patches
-                    + (num_tiles - 2) * middle_crop_window_patches
-                    + right_crop_window_patches
-                )
-            else:
-                single_crop_window_patches = (
-                    (crop_patches + pooling_size - 1) // pooling_size * pooling_size
-                )
-                return single_crop_window_patches
-
         # Now build the output tokens
-        h = get_num_patches(tiling[0], self.pooling_h)
-        w = get_num_patches(tiling[1], self.pooling_w)
+        h = self._molmo_get_patches_from_tiling(
+            tiling[0],
+            self.pooling_h,
+            crop_patches,
+            crop_window_patches,
+            left_margin,
+            right_margin,
+        )
+        w = self._molmo_get_patches_from_tiling(
+            tiling[1],
+            self.pooling_w,
+            crop_patches,
+            crop_window_patches,
+            left_margin,
+            right_margin,
+        )
         # for each row of patches, add a patch token per patch
         per_row = np.full((w // self.pooling_w,), self.patch_token_id, dtype=np.int32)
         if self.use_column_tokens:
@@ -810,6 +918,14 @@ class JinaVLMImageProcessor(BaseImageProcessor):
 
         return slices, image_masks, patch_ordering_arr, best_grid
 
+    def minicpm_get_n_image_patches(
+        self, height: int, width: int, max_crops: int, with_thumbnail: bool = False
+    ) -> int:
+        raise NotImplementedError(
+            'Function `get_n_image_patches` is not implemented for cropping method '
+            f'{CroppingMethod.ADAPTIVE_SLICING}'
+        )
+
     def minicpm_adaptive_slicing(self, image: np.ndarray, with_thumbnail: bool = True):
         scale_resolution = self.base_input_size[0]
         refine_image, image_mask, best_grid = self._minicpm_refine_image_for_slicing(
@@ -946,23 +1062,12 @@ class JinaVLMImageProcessor(BaseImageProcessor):
         self.start_token_id = start_token_id
         self.end_token_id = end_token_id
 
-    def preprocess(
-        self,
-        images: ImageInput,
-        **kwargs: Unpack[JinaVLMImagesKwargs],
-    ) -> Dict[str, List[np.ndarray]]:
-        """Preprocess an image or batch of images."""
-        if images is None or len(images) == 0:
-            return {
-                'image_crops': [],
-                'image_tokens': [],
-                'image_input_idx': [],
-                'image_padding_mask': [],
-            }
-
+    def _resolve_images_kwargs(
+        self, **kwargs: Unpack[JinaVLMImagesKwargs]
+    ) -> JinaVLMImagesKwargs:
+        max_crops = self.max_crops
         if 'max_crops' in kwargs and kwargs['max_crops'] is not None:
             max_crops = kwargs['max_crops']
-            self.max_crops = max_crops
 
         min_pixels = self.min_pixels
         if 'min_pixels' in kwargs and kwargs['min_pixels'] is not None:
@@ -984,14 +1089,93 @@ class JinaVLMImageProcessor(BaseImageProcessor):
             size = {'shortest_edge': min_pixels, 'longest_edge': max_pixels}
         else:
             size = {**self.size}
-
+        min_pixels = size['shortest_edge']
+        max_pixels = size['longest_edge']
         do_resize = self.do_resize
         if 'do_resize' in kwargs and kwargs['do_resize'] is not None:
             do_resize = kwargs['do_resize']
-
         do_convert_rgb = self.do_convert_rgb
         if 'do_convert_rgb' in kwargs and kwargs['do_convert_rgb'] is not None:
             do_convert_rgb = kwargs['do_convert_rgb']
+        input_data_format = None
+        if 'input_data_format' in kwargs:
+            input_data_format = kwargs['input_data_format']
+
+        return JinaVLMImagesKwargs(
+            do_convert_rgb=do_convert_rgb,
+            do_resize=do_resize,
+            min_pixels=min_pixels,
+            max_pixels=max_pixels,
+            size=size,
+            max_crops=max_crops,
+            input_data_format=input_data_format,
+        )
+
+    def get_n_image_patches(
+        self,
+        height: int,
+        width: int,
+        **kwargs: Unpack[JinaVLMImagesKwargs],
+    ) -> int:
+        """A utility that returns number of image patches for a given image size.
+
+        Args:
+            height (`int`):
+                Height of the input image.
+            width (`int`):
+                Width of the input image.
+            **kwargs (`dict`, *optional*)
+                Any kwargs to override defaults of the image processor.
+        Returns:
+            `int`: Number of image patches
+        """
+        if self.cropping_method != CroppingMethod.OVERLAP_AND_RESIZE:
+            raise NotImplementedError(
+                'Function is only implemented for cropping method '
+                f'{CroppingMethod.OVERLAP_AND_RESIZE}'
+            )
+        kwargs = self._resolve_images_kwargs(**kwargs)
+        do_resize = kwargs['do_resize']
+        size = kwargs['size']
+        max_crops = kwargs['max_crops']
+        if do_resize:
+            height, width = smart_resize(
+                height,
+                width,
+                factor=self.patch_size,
+                min_pixels=size['shortest_edge'],
+                max_pixels=size['longest_edge'],
+            )
+
+        if self.cropping_method == CroppingMethod.RESIZE:
+            return self.base_get_n_image_patches(height, width, max_crops)
+        elif self.cropping_method == CroppingMethod.OVERLAP_AND_RESIZE:
+            return self.molmo_get_n_image_patches(height, width, max_crops)
+        elif self.cropping_method == CroppingMethod.ADAPTIVE_SLICING:
+            return self.minicpm_get_n_image_patches(height, width, max_crops)
+        return self.minicpm_get_n_image_patches(
+            height, width, max_crops, with_thumbnail=True
+        )
+
+    def preprocess(
+        self,
+        images: ImageInput,
+        **kwargs: Unpack[JinaVLMImagesKwargs],
+    ) -> Dict[str, List[np.ndarray]]:
+        """Preprocess an image or batch of images."""
+        if images is None or len(images) == 0:
+            return {
+                'image_crops': [],
+                'image_tokens': [],
+                'image_input_idx': [],
+                'image_padding_mask': [],
+            }
+        kwargs = self._resolve_images_kwargs(**kwargs)
+        do_convert_rgb = kwargs['do_convert_rgb']
+        do_resize = kwargs['do_resize']
+        input_data_format = kwargs['input_data_format']
+        size = kwargs['size']
+        self.max_crops = kwargs['max_crops']
 
         # noinspection PyTypeChecker
         images = self.fetch_images(images)
@@ -1001,16 +1185,11 @@ class JinaVLMImageProcessor(BaseImageProcessor):
                 'Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray '
                 'or torch.Tensor'
             )
-
         if do_convert_rgb:
             images = [convert_to_rgb(image) for image in images]
 
         # All transformations expect numpy arrays
         images = [to_numpy_array(image) for image in images]
-
-        input_data_format = None
-        if 'input_data_format' in kwargs:
-            input_data_format = kwargs['input_data_format']
         if input_data_format is None:
             # We assume that all images have the same channel dimension format.
             input_data_format = infer_channel_dimension_format(images[0])

modeling_jvlm.py

@@ -27,14 +27,13 @@ from .blocks_jvlm import (
     TransformerBlock,
     VisionLanguageConnector,
     build_layer_norm,
-    resolve_causal_mask
+    resolve_causal_mask,
 )
 from .configuration_jvlm import JinaVLMConfig, JinaVLMTextConfig, JinaVLMVisionConfig
 
 
 class JinaPreTrainedModel(PreTrainedModel):
     config: JinaVLMConfig
-    config_class = JinaVLMConfig
     base_model_prefix = 'model'
     supports_gradient_checkpointing = True
     _supports_flash_attn = True
@@ -51,8 +50,6 @@ class JinaPreTrainedModel(PreTrainedModel):
 
 class JinaVLMVisionModel(JinaPreTrainedModel):
     config: JinaVLMVisionConfig
-    config_class = JinaVLMVisionConfig
-    base_model_prefix = ''
 
     def __init__(self, config: JinaVLMVisionConfig, *args, **kwargs):
         super().__init__(config, *args, **kwargs)
@@ -186,7 +183,11 @@ class JinaVLMVisionModel(JinaPreTrainedModel):
             pos = pos_emb[None, :, :].to(x.dtype)
         return x + pos
 
-    def get_visual_features(self, images: torch.Tensor) -> BaseModelOutput:
+    def get_visual_features(
+        self,
+        images: torch.Tensor,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> BaseModelOutput:
         x, shape = self.patch_embed(images)
         if self.cls_embed is not None:
             cls = self.cls_embed.view(1, 1, -1).expand(x.shape[0], -1, -1).to(x.dtype)
@@ -201,7 +202,11 @@ class JinaVLMVisionModel(JinaPreTrainedModel):
         hidden_states = []
         attentions = []
         for layer in self.layers:
-            x, attn = layer(x, attn_implementation=self.config._attn_implementation)
+            x, attn = layer(
+                x,
+                attn_implementation=self.config._attn_implementation,
+                **kwargs,
+            )
             hidden_states.append(x)
             attentions.append(attn)
         x = self.post_lnorm(x)
@@ -214,12 +219,15 @@ class JinaVLMVisionModel(JinaPreTrainedModel):
         )
 
     def forward(
-        self, images: torch.Tensor, image_masks: torch.Tensor
+        self,
+        images: torch.Tensor,
+        image_masks: torch.Tensor,
+        **kwargs: Unpack[FlashAttentionKwargs],
     ) -> BaseModelOutput:
         b, t, n, d = images.shape
         mask = ~torch.all(images.view(b * t, n, d) == -1, dim=(1, 2), keepdim=True)
         images = images.view(b * t, n, d)
-        out = self.get_visual_features(images)
+        out = self.get_visual_features(images, **kwargs)
         image_features = out.hidden_states
 
         features = []
@@ -230,14 +238,13 @@ class JinaVLMVisionModel(JinaPreTrainedModel):
             features.append(feats)
         image_features = torch.cat(features, dim=-1)
         image_features = image_features * mask
-        image_features = image_features.view(b, t, n, -1)
-
+        image_features = image_features.view(b, t, n, -1).contiguous()
         image_features = self.vl_connector(
             image_features,
             image_masks,
             attn_implementation=self.config._attn_implementation,
+            **kwargs,
         )
-
         return BaseModelOutput(
             last_hidden_state=image_features,
             hidden_states=out.hidden_states,
@@ -246,11 +253,7 @@ class JinaVLMVisionModel(JinaPreTrainedModel):
 
 
 class JinaVLMTextModel(JinaPreTrainedModel):
-    """Decoder-only language model."""
-
     config: JinaVLMTextConfig
-    config_class = JinaVLMTextConfig
-    base_model_prefix = ''
 
     def __init__(self, config: JinaVLMTextConfig, *args, **kwargs):
         super().__init__(config, *args, **kwargs)
@@ -297,7 +300,6 @@ class JinaVLMTextModel(JinaPreTrainedModel):
                 theta=self.config.rope_theta,
                 head_dim=self.config.block_config.attn_config.head_dim,
                 hidden_size=self.config.hidden_size,
-                n_heads=self.config.block_config.attn_config.n_heads,
                 partial_rotary_factor=self.config.partial_rotary_factor,
                 scaling=self.config.rope_scaling,
             )
@@ -444,7 +446,7 @@ class JinaVLMTextModel(JinaPreTrainedModel):
 
 
 class JinaVLM(JinaPreTrainedModel):
-    base_model_prefix = ''
+    config: JinaVLMConfig
 
     def __init__(self, config: JinaVLMConfig):
         super().__init__(config)
@@ -493,7 +495,7 @@ class JinaVLM(JinaPreTrainedModel):
     ) -> BaseModelOutputWithPast:
         image_features = None
         if images is not None and images.shape[1] > 0:
-            image_out = self.vision_model(images, image_masks)
+            image_out = self.vision_model(images, image_masks, **kwargs)
             image_features = image_out.last_hidden_state
         return self.language_model(
             input_ids=input_ids,
@@ -512,10 +514,10 @@ class JinaVLM(JinaPreTrainedModel):
 
 
 class JinaVLMForConditionalGeneration(JinaPreTrainedModel, GenerationMixin):
-    _checkpoint_conversion_mapping = {}
-    _tied_weights_keys = ['lm_head.weight']
+    _tied_weights_keys = {
+        'lm_head.weight': 'model.language_model.embedding.embedding.weight'
+    }
     accepts_loss_kwargs = False
-    base_model_prefix = 'model'
     config: JinaVLMConfig
 
     def __init__(self, config: JinaVLMConfig):

processing_jvlm.py

@@ -10,11 +10,14 @@ from transformers.image_utils import ImageInput
 from transformers.processing_utils import (
     AllKwargsForChatTemplate,
     CommonKwargs,
+    MultiModalData,
     ProcessorMixin,
     Unpack,
 )
 from transformers.tokenization_utils_base import (
-    PaddingStrategy, PreTokenizedInput, TextInput,
+    PaddingStrategy,
+    PreTokenizedInput,
+    TextInput,
 )
 
 from .image_processing_jvlm import JinaVLMImageProcessor, JinaVLMImagesKwargs
@@ -38,7 +41,7 @@ class JinaVLMTextKwargs(TypedDict, total=False):
     is_split_into_words: Optional[bool]
 
 
-class JinaVLProcessingKwargs(JinaVLMTextKwargs, JinaVLMImagesKwargs, CommonKwargs):
+class JinaVLMProcessingKwargs(JinaVLMTextKwargs, JinaVLMImagesKwargs, CommonKwargs):
     return_labels: Optional[bool]
 
 
@@ -171,8 +174,8 @@ class JinaVLMProcessor(ProcessorMixin):
     def _collate(
         self,
         batch: Dict[str, List[Optional[np.ndarray]]],
-        max_sequence_length: Optional[int] = None,
-        max_crops: Optional[int] = None,
+        text_max_sequence_length: Optional[int] = None,
+        image_max_sequence_length: Optional[int] = None,
         padding: Union[
             PaddingStrategy.MAX_LENGTH, PaddingStrategy.LONGEST
         ] = PaddingStrategy.MAX_LENGTH,
@@ -185,10 +188,10 @@ class JinaVLMProcessor(ProcessorMixin):
             _padding_side = 'right'
             if key in self.TEXT_KEYS:
                 _padding_side = padding_side
-                max_len = max_sequence_length
+                max_len = text_max_sequence_length
                 dtype = np.int64
             elif key in self.IMAGE_KEYS:
-                max_len = max_crops
+                max_len = image_max_sequence_length
                 dtype = np.int64
                 if key == 'images':
                     dtype = np.float32
@@ -214,22 +217,22 @@ class JinaVLMProcessor(ProcessorMixin):
             shift = input_ids_padlens[:, np.newaxis, np.newaxis]
             shift = np.repeat(shift, n_image_tokens, axis=2)
             shift = np.repeat(shift, n_crops, axis=1)
-            image_input_idx[image_input_idx < 0] = -max_sequence_length
+            image_input_idx[image_input_idx < 0] = -text_max_sequence_length
             image_input_idx = image_input_idx + shift
             out['image_input_idx'] = image_input_idx
 
-        if max_sequence_length is not None:
+        if text_max_sequence_length is not None:
             image_input_idx = out.get('image_input_idx', [])
             n = len(image_input_idx)
             for i in range(n):
                 arr = image_input_idx[i]
                 if arr.ndim > 0 and arr.size > 0:
                     n_image_tokens = arr.max()
-                    if n_image_tokens > max_sequence_length - 3:
+                    if n_image_tokens > text_max_sequence_length - 3:
                         raise RuntimeError(
                             'Image tokens truncation at sequence boundary. Max '
-                            f'sequence length ({max_sequence_length}) is too small '
-                            'to fit the generated image tokens '
+                            f'sequence length ({text_max_sequence_length}) is too '
+                            'small to fit the generated image tokens '
                             f'({n_image_tokens}). Consider increasing the max '
                             'sequence length or tweaking the image processing '
                             'parameters (`max_crops`, `max_pixels`) to reduce the '
@@ -386,7 +389,7 @@ class JinaVLMProcessor(ProcessorMixin):
         text: Union[
             None, TextInput, PreTokenizedInput, list[TextInput], list[PreTokenizedInput]
         ] = None,
-        **kwargs: Unpack[JinaVLProcessingKwargs],
+        **kwargs: Unpack[JinaVLMProcessingKwargs],
     ) -> BatchFeature:
         """Main method to prepare for the model one or several sequences(s) and
         image(s). This method forwards the `text` and `kwargs` arguments to  the
@@ -489,9 +492,11 @@ class JinaVLMProcessor(ProcessorMixin):
                 )
 
         outputs = defaultdict(list)
+        n_images = []
         for idx in range(batch_size):
             _token_ids = token_ids[idx]
             _images = images[idx]
+            n_images.append(len(_images))
             image_inputs = self.image_processor(_images, **images_kwargs)
             image_crops = image_inputs['image_crops']
             image_tokens = image_inputs['image_tokens']
@@ -510,14 +515,42 @@ class JinaVLMProcessor(ProcessorMixin):
                 outputs[k].append(v)
 
         if padding != PaddingStrategy.DO_NOT_PAD:
+            text_max_sequence_length = max_length or self.max_sequence_length
+            max_crops = max_crops or self.max_crops
+            max_n_images = max(n_images)
+            image_max_sequence_length = (max_crops + 1) * max_n_images
             outputs = self._collate(
                 outputs,
-                max_sequence_length=max_length or self.max_sequence_length,
-                max_crops=max_crops or self.max_crops,
+                text_max_sequence_length=text_max_sequence_length,
+                image_max_sequence_length=image_max_sequence_length,
                 padding=padding,
                 padding_side=padding_side,
             )
         return BatchFeature(data=outputs, tensor_type=return_tensors)
 
+    def _get_num_multimodal_tokens(
+        self,
+        image_sizes: Optional[List[List[int]]] = None,
+        **kwargs: Unpack[JinaVLMImagesKwargs],
+    ) -> MultiModalData:
+        """Computes the number of placeholder tokens needed for multimodal inputs with
+        the given sizes.
+
+        Args:
+            image_sizes (`list[list[int]]`, *optional*):
+                The input sizes formatted as (height, width) per each image.
+        Returns:
+            `MultiModalData`: A `MultiModalData` object holding number of tokens per
+            each of the provided input modalities, along with other useful data.
+        """
+        data = {}
+        if image_sizes is not None:
+            n_patches = [
+                self.image_processor.get_n_image_patches(h, w, **kwargs)
+                for h, w in image_sizes
+            ]
+            data.update({'num_image_tokens': n_patches, 'num_image_patches': n_patches})
+        return MultiModalData(**data)
+
 
 JinaVLMProcessor.register_for_auto_class()

test_jvlm.py

@@ -11,7 +11,10 @@ os.environ['TOKENIZERS_PARALLELISM'] = 'false'
 
 import torch
 from transformers import (
-    AutoModelForCausalLM, AutoProcessor, GenerationConfig, TextStreamer
+    AutoModelForCausalLM,
+    AutoProcessor,
+    GenerationConfig,
+    TextStreamer,
 )
 from transformers.utils import is_flash_attn_2_available
 
@@ -60,7 +63,8 @@ def _build_conversations(
         try:
             result = urlparse(_path)
             return result.scheme in ('http', 'https')
-        except:
+        except Exception as e:
+            _ = str(e)
             return False
 
     images = images or []
@@ -83,8 +87,9 @@ def _build_conversations(
             images = [TEST_IMAGE]
             n_images = len(images)
         prompts = (
-            ['Describe the image in 100 words'] if n_images == 1 or map_mode else
-            ['Describe the images in 100 words']
+            ['Describe the image in 100 words']
+            if n_images == 1 or map_mode
+            else ['Describe the images in 100 words']
         )
     n_prompts = len(prompts)
 
@@ -119,8 +124,16 @@ def _build_conversations(
     allimages = []
     allprompts = []
     ordinals = [
-        'first', 'second', 'third', 'fourth', 'fifth',
-        'sixth', 'seventh', 'eighth', 'ninth', 'tenth',
+        'first',
+        'second',
+        'third',
+        'fourth',
+        'fifth',
+        'sixth',
+        'seventh',
+        'eighth',
+        'ninth',
+        'tenth',
     ]
     for images, prompt in examples:
         content = []
@@ -130,15 +143,17 @@ def _build_conversations(
             content.append({'type': 'text', 'text': prompt})
         if len(images) > 1 and image_labels:
             for idx, img in enumerate(images):
-                ordinal = ordinals[idx] if idx < len(ordinals) else f'{idx+1}th'
+                ordinal = ordinals[idx] if idx < len(ordinals) else f'{idx + 1}th'
                 image = images[idx]
                 descriptor = f'url: {image}'
                 if os.path.isfile(image):
                     descriptor = f'filename: {os.path.basename(image)}'
-                content.append({
-                    'type': 'text',
-                    'text': f'(this is the {ordinal} image, {descriptor})',
-                })
+                content.append(
+                    {
+                        'type': 'text',
+                        'text': f'(this is the {ordinal} image, {descriptor})',
+                    }
+                )
                 content.append({'type': 'image', 'image': img})
         else:
             content.extend([{'type': 'image', 'image': image} for image in images])
@@ -189,9 +204,7 @@ def _token_usage_report(
     tokens_per_image_list = []
 
     # Find all img_start and img_end positions in input_ids
-    start_positions = (input_ids == image_start_id).nonzero(
-        as_tuple=True
-    )[0].tolist()
+    start_positions = (input_ids == image_start_id).nonzero(as_tuple=True)[0].tolist()
     end_positions = (input_ids == image_end_id).nonzero(as_tuple=True)[0].tolist()
 
     if len(start_positions) > 0 and len(end_positions) > 0:
@@ -211,9 +224,8 @@ def _token_usage_report(
                 # Get the start and end indices for this image
                 start_idx_begin = idx * n_starts_per_image
                 end_idx_end = (idx + 1) * n_starts_per_image
-                if (
-                    start_idx_begin < len(start_positions) and
-                    end_idx_end <= len(end_positions)
+                if start_idx_begin < len(start_positions) and end_idx_end <= len(
+                    end_positions
                 ):
                     # First start position and last end position define the image span
                     first_start = start_positions[start_idx_begin]
@@ -233,10 +245,10 @@ def _token_usage_report(
 
     for idx in range(n_images):
         n_tokens = tokens_per_image_list[idx] if idx < len(tokens_per_image_list) else 0
-        pct = (n_tokens / max_sequence_length * 100)
+        pct = n_tokens / max_sequence_length * 100
         report.append(f'├── Image {idx + 1} → {n_tokens} tokens ({pct:.1f}%)')
 
-    text_pct = (text_token_count / max_sequence_length * 100)
+    text_pct = text_token_count / max_sequence_length * 100
     report.append(f'└── Text: {text_token_count} tokens ({text_pct:.1f}%)')
 
     return '\n'.join(report)
@@ -253,7 +265,7 @@ def test_jvlm():
         help=(
             'Model path (default: `"."`). Set this to `"jinaai/jina-vlm-v1"` if you '
             'are running this script outside this repo.'
-        )
+        ),
     )
     parser.add_argument(
         '-i',
@@ -339,7 +351,9 @@ def test_jvlm():
     print(f'Using dtype: {dtype}')
     print('Model path: ', args.model)
     processor = AutoProcessor.from_pretrained(
-        args.model, trust_remote_code=True, use_fast=False,
+        args.model,
+        trust_remote_code=True,
+        use_fast=False,
     )
     model = AutoModelForCausalLM.from_pretrained(
         args.model,
@@ -356,13 +370,13 @@ def test_jvlm():
     print('Done ✅')
     print()
 
-    print('--- Let\'s create some conversations ...')
+    print("--- Let's create some conversations ...")
     conversations, images, prompts = _build_conversations(
         args.image,
         args.prompt,
         map_mode=args.map,
         prompt_first=args.prompt_first,
-        image_labels=args.image_labels
+        image_labels=args.image_labels,
     )
     n_conversations = len(conversations)
     print(f'Built {n_conversations} conversations 🚀')
@@ -434,25 +448,28 @@ def test_jvlm():
             print(f'├── 🖼️Images: {images[idx]}')
             print(f'├── 📜Prompt: {prompts[idx]}')
             print(f'├── 💬Chat:{texts[idx]}')
-            print(f'└── 🧠Response:', end='')
+            print('└── 🧠Response:', end='')
             ith_inputs = {k: v[idx].unsqueeze(0) for k, v in device_inputs.items()}
             with (
                 timer,
                 torch.no_grad(),
-                torch.autocast(device.type, enabled=(device.type != 'mps'), dtype=dtype)
+                torch.autocast(
+                    device.type, enabled=(device.type != 'mps'), dtype=dtype
+                ),
             ):
                 output = model.generate(
                     **ith_inputs,
                     streamer=streamer,
                     generation_config=GenerationConfig(
-                        max_new_tokens=args.max_tokens, do_sample=False,
+                        max_new_tokens=args.max_tokens,
+                        do_sample=False,
                     ),
                     return_dict_in_generate=True,
                     use_model_defaults=True,
                 )
             generation_time += timer.time
 
-            out = output.sequences[0][len(input_prompts[idx].tolist()):]
+            out = output.sequences[0][len(input_prompts[idx].tolist()) :]
             generated_tokens += len(out)
             print('Token usage report:')
             print(token_usage_reports[idx])
@@ -470,7 +487,8 @@ def test_jvlm():
             output = model.generate(
                 **device_inputs,
                 generation_config=GenerationConfig(
-                    max_new_tokens=args.max_tokens, do_sample=False,
+                    max_new_tokens=args.max_tokens,
+                    do_sample=False,
                 ),
                 return_dict_in_generate=True,
                 use_model_defaults=True,
@@ -478,7 +496,7 @@ def test_jvlm():
         generation_time = timer.time
 
         for idx in range(n_conversations):
-            out = output.sequences[idx][len(input_prompts[idx].tolist()):]
+            out = output.sequences[idx][len(input_prompts[idx].tolist()) :]
             generated_tokens += len(out)
             response = processor.tokenizer.decode(out, skip_special_tokens=True)
             print(f'* Conversation {idx + 1}/{n_conversations}')