transformers 5.0.0rc0__py3-none-any.whl → 5.0.0rc2__py3-none-any.whl

transformers/models/lfm2_vl/processing_lfm2_vl.py

@@ -165,63 +165,103 @@ class Lfm2VlProcessor(ProcessorMixin):
         image_sizes: list[list[int]],
         use_image_special_tokens: bool,
         **images_kwargs,
-    ):
-        prompt_strings = []
+    ) -> list[str]:
+        use_thumbnail = images_kwargs.get("use_thumbnail", self.image_processor.use_thumbnail)
+        image_data = iter(zip(image_rows, image_cols, image_sizes))
 
-        image_data = iter(zip(*[image_rows, image_cols, image_sizes]))
+        prompt_strings = []
         for sample_text, sample_images in zip(text, images):
-            split_sample = sample_text.split(self.image_token)
-            sample_text_with_image_tokens = ""
-            for i, image in enumerate(sample_images):
-                sample_text_with_image_tokens += split_sample[i]
-                if use_image_special_tokens:
-                    sample_text_with_image_tokens += self.image_start_token
+            text_parts = sample_text.split(self.image_token)
+            result_parts = []
+
+            for i, _ in enumerate(sample_images):
+                result_parts.append(text_parts[i])
 
                 rows, cols, image_size = next(image_data)
-                num_thumbnail_tokens, num_tokens_per_tile = self._get_image_num_tokens(image_size, **images_kwargs)
-
-                if rows > 1 or cols > 1:
-                    for row in range(rows):
-                        for col in range(cols):
-                            if use_image_special_tokens:
-                                sample_text_with_image_tokens += f"<|img_row_{row + 1}_col_{col + 1}|>"
-                            sample_text_with_image_tokens += self.image_token * num_tokens_per_tile
-
-                    if num_thumbnail_tokens > 0:
-                        if use_image_special_tokens:
-                            sample_text_with_image_tokens += self.image_thumbnail_token
-                        sample_text_with_image_tokens += self.image_token * num_thumbnail_tokens
-                else:
-                    sample_text_with_image_tokens += self.image_token * num_thumbnail_tokens
+                tokens_per_tile, tokens_for_image = self._get_image_num_tokens(image_size, **images_kwargs)
+                image_tokens = self._build_image_tokens(
+                    rows,
+                    cols,
+                    tokens_per_tile,
+                    tokens_for_image,
+                    use_thumbnail,
+                    use_image_special_tokens,
+                )
+                result_parts.append(image_tokens)
 
-                if use_image_special_tokens:
-                    sample_text_with_image_tokens += self.image_end_token
+            # Add remaining text after the last image
+            if len(sample_images) < len(text_parts):
+                result_parts.append(text_parts[-1])
 
-                sample_text_with_image_tokens += split_sample[i + 1]
-            prompt_strings.append(sample_text_with_image_tokens)
+            prompt_strings.append("".join(result_parts))
 
         return prompt_strings
 
+    def _build_image_tokens(
+        self,
+        rows: int,
+        cols: int,
+        tokens_per_tile: int,
+        tokens_for_image: int,
+        use_thumbnail: bool,
+        use_image_special_tokens: bool,
+    ) -> str:
+        """Build the expanded token string for a single image."""
+        parts = []
+
+        if use_image_special_tokens:
+            parts.append(self.image_start_token)
+
+        is_multi_tile = rows > 1 or cols > 1
+        if is_multi_tile:
+            for row in range(rows):
+                for col in range(cols):
+                    if use_image_special_tokens:
+                        parts.append(f"<|img_row_{row + 1}_col_{col + 1}|>")
+                    parts.append(self.image_token * tokens_per_tile)
+
+            if use_thumbnail:
+                if use_image_special_tokens:
+                    parts.append(self.image_thumbnail_token)
+                parts.append(self.image_token * tokens_for_image)
+        else:
+            parts.append(self.image_token * tokens_for_image)
+
+        if use_image_special_tokens:
+            parts.append(self.image_end_token)
+
+        return "".join(parts)
+
+    def _compute_tokens_per_tile(self, tile_size: int, encoder_patch_size: int, downsample_factor: int) -> int:
+        """Compute the number of tokens for a single tile."""
+        num_patches = tile_size // encoder_patch_size
+        downsampled_patches = math.ceil(num_patches / downsample_factor)
+        return downsampled_patches * downsampled_patches
+
+    def _compute_tokens_for_image(self, image_size: list[int], encoder_patch_size: int, downsample_factor: int) -> int:
+        """Compute the number of tokens for a resized image (used for single-tile or thumbnail)."""
+        image_height, image_width = image_size
+        patches_h = math.ceil((image_height // encoder_patch_size) / downsample_factor)
+        patches_w = math.ceil((image_width // encoder_patch_size) / downsample_factor)
+        return patches_h * patches_w
+
     def _get_image_num_tokens(self, image_size: list[int], **images_kwargs) -> tuple[int, int]:
+        """
+        Compute token counts for image processing.
+
+        Returns:
+            tuple[int, int]: (tokens_per_tile, tokens_for_image)
+                - tokens_per_tile: tokens for each tile in multi-tile mode
+                - tokens_for_image: tokens for the resized image (single-tile) or thumbnail (multi-tile)
+        """
         tile_size = images_kwargs.get("tile_size", self.image_processor.tile_size)
         downsample_factor = images_kwargs.get("downsample_factor", self.image_processor.downsample_factor)
         encoder_patch_size = images_kwargs.get("encoder_patch_size", self.image_processor.encoder_patch_size)
-        use_thumbnail = images_kwargs.get("use_thumbnail", self.image_processor.use_thumbnail)
-
-        thumbnail_tokens = 0
-        if use_thumbnail:
-            image_height, image_width = image_size
-            num_patches_height = image_height // encoder_patch_size
-            num_patches_width = image_width // encoder_patch_size
-            dwn_num_patches_height = math.ceil(num_patches_height / downsample_factor)
-            dwn_num_patches_width = math.ceil(num_patches_width / downsample_factor)
-            thumbnail_tokens = dwn_num_patches_height * dwn_num_patches_width
 
-        num_patches_tile = tile_size // encoder_patch_size
-        dwn_num_patches_tile = math.ceil(num_patches_tile / downsample_factor)
-        tile_tokens = dwn_num_patches_tile * dwn_num_patches_tile
+        tokens_per_tile = self._compute_tokens_per_tile(tile_size, encoder_patch_size, downsample_factor)
+        tokens_for_image = self._compute_tokens_for_image(image_size, encoder_patch_size, downsample_factor)
 
-        return thumbnail_tokens, tile_tokens
+        return tokens_per_tile, tokens_for_image
 
     def batch_decode(self, *args, **kwargs):
         """

transformers/models/lightglue/image_processing_lightglue_fast.py

@@ -174,9 +174,8 @@ class LightGlueImageProcessorFast(BaseImageProcessorFast):
         stacked_pairs = [torch.stack(pair, dim=0) for pair in image_pairs]
 
         # Return in same format as slow processor
-        image_pairs = torch.stack(stacked_pairs, dim=0) if return_tensors else stacked_pairs
 
-        return BatchFeature(data={"pixel_values": image_pairs})
+        return BatchFeature(data={"pixel_values": stacked_pairs}, tensor_type=return_tensors)
 
     def post_process_keypoint_matching(
         self,

transformers/models/lightglue/modeling_lightglue.py

@@ -27,6 +27,7 @@ from torch import nn
 from torch.nn.utils.rnn import pad_sequence
 
 from ...activations import ACT2FN
+from ...integrations import use_kernelized_func
 from ...modeling_flash_attention_utils import FlashAttentionKwargs
 from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
 from ...processing_utils import Unpack
@@ -174,6 +175,7 @@ def eager_attention_forward(
     return attn_output, attn_weights
 
 
+@use_kernelized_func(apply_rotary_pos_emb)
 class LightGlueAttention(nn.Module):
     """Multi-headed attention from 'Attention Is All You Need' paper"""
 
@@ -199,7 +201,6 @@ class LightGlueAttention(nn.Module):
         self.o_proj = nn.Linear(
             config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
         )
-        self.rotary_fn = apply_rotary_pos_emb
 
     def forward(
         self,
@@ -870,6 +871,7 @@ class LightGlueForKeypointMatching(LightGluePreTrainedModel):
         labels: Optional[torch.LongTensor] = None,
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
+        **kwargs,
     ) -> Union[tuple, "LightGlueKeypointMatchingOutput"]:
         loss = None
         if labels is not None:

transformers/models/lightglue/modular_lightglue.py

@@ -927,6 +927,7 @@ class LightGlueForKeypointMatching(LightGluePreTrainedModel):
         labels: Optional[torch.LongTensor] = None,
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
+        **kwargs,
     ) -> Union[tuple, "LightGlueKeypointMatchingOutput"]:
         loss = None
         if labels is not None:

transformers/models/lilt/modeling_lilt.py

@@ -21,6 +21,7 @@ import torch
 from torch import nn
 from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
 
+from ... import initialization as init
 from ...activations import ACT2FN
 from ...modeling_layers import GradientCheckpointingLayer
 from ...modeling_outputs import (
@@ -279,11 +280,9 @@ class LiltSelfAttention(nn.Module):
         new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
         context_layer = context_layer.view(*new_context_layer_shape)
 
-        outputs = (
-            ((context_layer, layout_context_layer), attention_probs)
-            if output_attentions
-            else ((context_layer, layout_context_layer),)
-        )
+        outputs = (context_layer, layout_context_layer)
+        if output_attentions:
+            outputs = outputs + (attention_probs,)
 
         return outputs
 
@@ -327,9 +326,9 @@ class LiltAttention(nn.Module):
             attention_mask,
             output_attentions,
         )
-        attention_output = self.output(self_outputs[0][0], hidden_states)
-        layout_attention_output = self.layout_output(self_outputs[0][1], layout_inputs)
-        outputs = ((attention_output, layout_attention_output),) + self_outputs[1:]  # add attentions if we output them
+        attention_output = self.output(self_outputs[0], hidden_states)
+        layout_attention_output = self.layout_output(self_outputs[1], layout_inputs)
+        outputs = (attention_output, layout_attention_output) + self_outputs[2:]  # add attentions if we output them
         return outputs
 
 
@@ -395,10 +394,10 @@ class LiltLayer(GradientCheckpointingLayer):
             attention_mask,
             output_attentions=output_attentions,
         )
-        attention_output = self_attention_outputs[0][0]
-        layout_attention_output = self_attention_outputs[0][1]
+        attention_output = self_attention_outputs[0]
+        layout_attention_output = self_attention_outputs[1]
 
-        outputs = self_attention_outputs[1:]  # add self attentions if we output attention weights
+        outputs = self_attention_outputs[2:]  # add self attentions if we output attention weights
 
         layer_output = apply_chunking_to_forward(
             self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output
@@ -406,7 +405,7 @@ class LiltLayer(GradientCheckpointingLayer):
         layout_layer_output = apply_chunking_to_forward(
             self.layout_feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, layout_attention_output
         )
-        outputs = ((layer_output, layout_layer_output),) + outputs
+        outputs = (layer_output, layout_layer_output) + outputs
 
         return outputs
 
@@ -451,11 +450,11 @@ class LiltEncoder(nn.Module):
                 output_attentions,
             )
 
-            hidden_states = layer_outputs[0][0]
-            layout_inputs = layer_outputs[0][1]
+            hidden_states = layer_outputs[0]
+            layout_inputs = layer_outputs[1]
 
             if output_attentions:
-                all_self_attentions = all_self_attentions + (layer_outputs[1],)
+                all_self_attentions = all_self_attentions + (layer_outputs[2],)
 
         if output_hidden_states:
             all_hidden_states = all_hidden_states + (hidden_states,)
@@ -500,6 +499,11 @@ class LiltPreTrainedModel(PreTrainedModel):
     supports_gradient_checkpointing = True
     _no_split_modules = []
 
+    def _init_weights(self, module):
+        super()._init_weights(module)
+        if isinstance(module, LiltTextEmbeddings):
+            init.copy_(module.position_ids, torch.arange(module.position_ids.shape[-1]).expand((1, -1)))
+
 
 @auto_docstring
 class LiltModel(LiltPreTrainedModel):
@@ -538,6 +542,7 @@ class LiltModel(LiltPreTrainedModel):
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
+        **kwargs,
     ) -> Union[tuple[torch.Tensor], BaseModelOutputWithPooling]:
         r"""
         bbox (`torch.LongTensor` of shape `(batch_size, sequence_length, 4)`, *optional*):
@@ -665,6 +670,7 @@ class LiltForSequenceClassification(LiltPreTrainedModel):
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
+        **kwargs,
     ) -> Union[tuple[torch.Tensor], SequenceClassifierOutput]:
         r"""
         bbox (`torch.LongTensor` of shape `(batch_size, sequence_length, 4)`, *optional*):
@@ -780,6 +786,7 @@ class LiltForTokenClassification(LiltPreTrainedModel):
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
+        **kwargs,
     ) -> Union[tuple[torch.Tensor], TokenClassifierOutput]:
         r"""
         bbox (`torch.LongTensor` of shape `(batch_size, sequence_length, 4)`, *optional*):
@@ -897,6 +904,7 @@ class LiltForQuestionAnswering(LiltPreTrainedModel):
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
+        **kwargs,
     ) -> Union[tuple[torch.Tensor], QuestionAnsweringModelOutput]:
         r"""
         bbox (`torch.LongTensor` of shape `(batch_size, sequence_length, 4)`, *optional*):

transformers/models/llama/modeling_llama.py

@@ -26,7 +26,7 @@ from torch import nn
 from ...activations import ACT2FN
 from ...cache_utils import Cache, DynamicCache
 from ...generation import GenerationMixin
-from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub
+from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func
 from ...masking_utils import create_causal_mask
 from ...modeling_layers import (
     GenericForQuestionAnswering,
@@ -42,7 +42,7 @@ from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
 from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
 from ...processing_utils import Unpack
 from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging
-from ...utils.generic import check_model_inputs
+from ...utils.generic import check_model_inputs, maybe_autocast
 from .configuration_llama import LlamaConfig
 
 
@@ -87,7 +87,7 @@ class LlamaRotaryEmbedding(nn.Module):
         inv_freq, self.attention_scaling = rope_init_fn(self.config, device)
 
         self.register_buffer("inv_freq", inv_freq, persistent=False)
-        self.original_inv_freq = inv_freq
+        self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False)
 
     @staticmethod
     def compute_default_rope_parameters(
@@ -126,7 +126,7 @@ class LlamaRotaryEmbedding(nn.Module):
         position_ids_expanded = position_ids[:, None, :].float()
 
         device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
-        with torch.autocast(device_type=device_type, enabled=False):  # Force float32
+        with maybe_autocast(device_type=device_type, enabled=False):  # Force float32
             freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
             emb = torch.cat((freqs, freqs), dim=-1)
             cos = emb.cos() * self.attention_scaling
@@ -224,6 +224,7 @@ def eager_attention_forward(
     return attn_output, attn_weights
 
 
+@use_kernelized_func(apply_rotary_pos_emb)
 class LlamaAttention(nn.Module):
     """Multi-headed attention from 'Attention Is All You Need' paper"""
 
@@ -249,7 +250,6 @@ class LlamaAttention(nn.Module):
         self.o_proj = nn.Linear(
             config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
         )
-        self.rotary_fn = apply_rotary_pos_emb
 
     def forward(
         self,

transformers/models/llama/tokenization_llama.py

@@ -12,11 +12,12 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from typing import Optional, Union
 
-from tokenizers import AddedToken, Tokenizer, decoders, pre_tokenizers
+from tokenizers import Tokenizer, decoders, pre_tokenizers
 from tokenizers.models import BPE
 
-from ...tokenization_utils_base import _get_prepend_scheme, generate_merges
+from ...tokenization_utils_base import _get_prepend_scheme
 from ...tokenization_utils_tokenizers import TokenizersBackend
 from ...utils import logging
 
@@ -61,6 +62,10 @@ class LlamaTokenizer(TokenizersBackend):
     refer to this superclass for more information regarding those methods.
 
     Args:
+        vocab (`str`, `dict` or `list`, *optional*):
+            Path to the vocabulary file, a dictionary or a list of tokens.
+        merges (`str` or `list`, *optional*):
+            Path to the merges file or a list of merges.
         clean_up_tokenization_spaces (`bool`, *optional*, defaults to `False`):
             Whether or not to cleanup spaces after decoding, cleanup consists in removing potential artifacts like
             extra spaces.
@@ -84,42 +89,32 @@ class LlamaTokenizer(TokenizersBackend):
     vocab_files_names = VOCAB_FILES_NAMES
     padding_side = "left"
     model_input_names = ["input_ids", "attention_mask"]
+    model = BPE
 
     def __init__(
         self,
+        vocab: Optional[Union[str, dict, list]] = None,
+        merges: Optional[Union[str, list]] = None,
         clean_up_tokenization_spaces=False,
         unk_token="<unk>",
         bos_token="<s>",
         eos_token="</s>",
-        add_bos_token=True,
-        add_eos_token=False,
         use_default_system_prompt=False,
         legacy=False,
         add_prefix_space=None,
-        vocab=None,
-        merges=None,
         **kwargs,
     ):
         self.add_prefix_space = add_prefix_space if add_prefix_space is not None else True
-
-        if vocab is not None:
-            self._vocab = (
-                {token: idx for idx, (token, _score) in enumerate(vocab)} if isinstance(vocab, list) else vocab
-            )
-        else:
+        self.legacy = legacy
+        self._vocab = vocab
+        if vocab is None:
             self._vocab = {
                 str(unk_token): 0,
                 str(bos_token): 1,
                 str(eos_token): 2,
             }
 
-        special_tokens = {str(eos_token), str(bos_token), str(unk_token)}
-
-        filtered_vocab = {t: i for t, i in self._vocab.items() if t not in special_tokens}
-        if merges is not None:
-            self._merges = [tuple(merge) if isinstance(merge, list) else merge for merge in merges]
-        else:
-            self._merges = generate_merges(filtered_vocab)
+        self._merges = merges or []
         self._tokenizer = Tokenizer(
             BPE(vocab=self._vocab, merges=self._merges, fuse_unk=True, byte_fallback=True, dropout=None)
         )
@@ -138,40 +133,17 @@ class LlamaTokenizer(TokenizersBackend):
             sequence += [decoders.Strip(content=" ", left=1)]
 
         self._tokenizer.decoder = decoders.Sequence(sequence)
-        tokenizer_object = self._tokenizer
-
+        self.use_default_system_prompt = use_default_system_prompt
         super().__init__(
-            tokenizer_object=tokenizer_object,
             clean_up_tokenization_spaces=clean_up_tokenization_spaces,
             unk_token=unk_token,
             bos_token=bos_token,
             eos_token=eos_token,
-            add_bos_token=add_bos_token,
-            add_eos_token=add_eos_token,
             use_default_system_prompt=use_default_system_prompt,
             add_prefix_space=add_prefix_space,
             **kwargs,
         )
 
-        self._add_bos_token = add_bos_token
-        self._add_eos_token = add_eos_token
-        self.use_default_system_prompt = use_default_system_prompt
-
-        self._post_init()
-
-    def _post_init(self):
-        """Post-initialization setup that needs to run after _tokenizer is set."""
-        # Only set pre_tokenizer/normalizer for Llama-3 style tokenizers (use Sequence)
-        pre_tok = self._tokenizer.pre_tokenizer
-        if pre_tok is None or type(pre_tok).__name__ != "Sequence":
-            self._tokenizer.pre_tokenizer = pre_tokenizers.Metaspace(
-                replacement="▁", prepend_scheme="first", split=False
-            )
-            self._tokenizer.normalizer = None
-            self.add_tokens([AddedToken(token, special=True) for token in self.all_special_tokens])
-        super()._post_init()
-        self.update_post_processor()
-
 
 __all__ = ["LlamaTokenizer", "LlamaTokenizerFast"]
 

transformers/models/llama4/image_processing_llama4_fast.py

@@ -419,10 +419,9 @@ class Llama4ImageProcessorFast(BaseImageProcessorFast):
                 )
                 grouped_processed_images[shape] = torch.cat([processed_images, global_tiles.unsqueeze(1)], dim=1)
         processed_images = reorder_images(grouped_processed_images, grouped_images_index)
-        aspect_ratios_list = reorder_images(grouped_aspect_ratios, grouped_images_index)
+        aspect_ratios = reorder_images(grouped_aspect_ratios, grouped_images_index)
 
         processed_images = torch.cat(processed_images, dim=0) if return_tensors else processed_images
-        aspect_ratios = torch.stack(aspect_ratios_list, dim=0) if return_tensors else aspect_ratios_list
         return BatchFeature(
             data={"pixel_values": processed_images, "aspect_ratios": aspect_ratios}, tensor_type=return_tensors
         )

transformers/models/llama4/modeling_llama4.py

@@ -40,7 +40,7 @@ from ...modeling_rope_utils import (
 from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
 from ...processing_utils import Unpack
 from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging
-from ...utils.generic import check_model_inputs
+from ...utils.generic import check_model_inputs, maybe_autocast
 from .configuration_llama4 import Llama4Config, Llama4TextConfig
 
 
@@ -188,7 +188,7 @@ class Llama4TextRotaryEmbedding(nn.Module):
         inv_freq, self.attention_scaling = rope_init_fn(self.config, device)
 
         self.register_buffer("inv_freq", inv_freq, persistent=False)
-        self.original_inv_freq = inv_freq
+        self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False)
 
     @staticmethod
     def compute_default_rope_parameters(
@@ -228,7 +228,7 @@ class Llama4TextRotaryEmbedding(nn.Module):
         position_ids_expanded = position_ids[:, None, :].float()
 
         device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
-        with torch.autocast(device_type=device_type, enabled=False):  # Force float32
+        with maybe_autocast(device_type=device_type, enabled=False):  # Force float32
             freqs = (inv_freq_expanded.to(x.device) @ position_ids_expanded).transpose(1, 2)
             freqs_cis = torch.polar(torch.ones_like(freqs), freqs)  # Convert to complex representation
             freqs_cis = freqs_cis * self.attention_scaling
@@ -1072,6 +1072,7 @@ class Llama4VisionModel(Llama4PreTrainedModel):
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
+        **kwargs,
     ) -> Union[BaseModelOutput, tuple[torch.Tensor, ...]]:
         r"""
 
@@ -1386,6 +1387,7 @@ class Llama4ForConditionalGeneration(Llama4PreTrainedModel, GenerationMixin):
         attention_mask=None,
         cache_position=None,
         logits_to_keep=None,
+        is_first_iteration=False,
         **kwargs,
     ):
         # Overwritten -- in specific circumstances we don't want to forward image inputs to the model
@@ -1397,12 +1399,15 @@ class Llama4ForConditionalGeneration(Llama4PreTrainedModel, GenerationMixin):
             attention_mask=attention_mask,
             cache_position=cache_position,
             logits_to_keep=logits_to_keep,
+            is_first_iteration=is_first_iteration,
             **kwargs,
         )
 
-        if cache_position[0] == 0:
-            # If we're in cached decoding stage, pixel values should be None because input ids do not contain special image token anymore
-            # Otherwise we need pixel values to be passed to model
+        if is_first_iteration or not kwargs.get("use_cache", True):
+            # Pixel values are used only in the first iteration if available
+            # In subsquent iterations, they are already merged with text and cached
+            # NOTE: first iteration doesn't have to be prefill, it can be the first
+            # iteration with a question and cached system prompt (continue generate from cache)
             model_inputs["pixel_values"] = pixel_values
 
         return model_inputs

transformers/models/llava/image_processing_llava_fast.py

@@ -149,7 +149,6 @@ class LlavaImageProcessorFast(BaseImageProcessorFast):
             processed_images_grouped[shape] = stacked_images
 
         processed_images = reorder_images(processed_images_grouped, grouped_images_index)
-        processed_images = torch.stack(processed_images, dim=0) if return_tensors else processed_images
 
         return BatchFeature(data={"pixel_values": processed_images}, tensor_type=return_tensors)
 

transformers/models/llava/modeling_llava.py

@@ -202,10 +202,11 @@ class LlavaModel(LlavaPreTrainedModel):
         image_features = self.multi_modal_projector(selected_image_feature)
 
         if "image_sizes" in kwargs:
-            split_sizes = [
-                (height // self.vision_tower.patch_size) * (width // self.vision_tower.patch_size)
-                for height, width in kwargs["image_sizes"]
-            ]
+            split_sizes = (
+                (torch.as_tensor(kwargs["image_sizes"], device=image_features.device) // self.vision_tower.patch_size)
+                .prod(dim=-1)
+                .tolist()
+            )
             image_features = torch.split(image_features.squeeze(0), split_sizes)
         else:
             image_features = list(image_features)
@@ -437,6 +438,7 @@ class LlavaForConditionalGeneration(LlavaPreTrainedModel, GenerationMixin):
         attention_mask=None,
         cache_position=None,
         logits_to_keep=None,
+        is_first_iteration=False,
         **kwargs,
     ):
         # Overwritten -- in specific circumstances we don't want to forward image inputs to the model
@@ -448,12 +450,15 @@ class LlavaForConditionalGeneration(LlavaPreTrainedModel, GenerationMixin):
             attention_mask=attention_mask,
             cache_position=cache_position,
             logits_to_keep=logits_to_keep,
+            is_first_iteration=is_first_iteration,
             **kwargs,
         )
 
-        if cache_position[0] == 0:
-            # If we're in cached decoding stage, pixel values should be None because input ids do not contain special image token anymore
-            # Otherwise we need pixel values to be passed to model
+        if is_first_iteration or not kwargs.get("use_cache", True):
+            # Pixel values are used only in the first iteration if available
+            # In subsquent iterations, they are already merged with text and cached
+            # NOTE: first iteration doesn't have to be prefill, it can be the first
+            # iteration with a question and cached system prompt (continue generate from cache)
             model_inputs["pixel_values"] = pixel_values
 
         return model_inputs

transformers/models/llava_next/image_processing_llava_next_fast.py

@@ -260,7 +260,6 @@ class LlavaNextImageProcessorFast(BaseImageProcessorFast):
 
         if do_pad:
             processed_images = self._pad_for_batching(processed_images)
-        processed_images = torch.stack(processed_images, dim=0) if return_tensors else processed_images
         return BatchFeature(
             data={"pixel_values": processed_images, "image_sizes": image_sizes}, tensor_type=return_tensors
         )

transformers/models/llava_next/modeling_llava_next.py

@@ -692,6 +692,7 @@ class LlavaNextForConditionalGeneration(LlavaNextPreTrainedModel, GenerationMixi
         attention_mask=None,
         cache_position=None,
         logits_to_keep=None,
+        is_first_iteration=False,
         **kwargs,
     ):
         # Overwritten -- in specific circumstances we don't want to forward image inputs to the model
@@ -703,12 +704,15 @@ class LlavaNextForConditionalGeneration(LlavaNextPreTrainedModel, GenerationMixi
             attention_mask=attention_mask,
             cache_position=cache_position,
             logits_to_keep=logits_to_keep,
+            is_first_iteration=is_first_iteration,
             **kwargs,
         )
 
-        # If we're in cached decoding stage, pixel values should be None because input ids do not contain special image token anymore
-        # Otherwise we need pixel values to be passed to model
-        if cache_position[0] == 0:
+        # Pixel values are used only in the first iteration if available
+        # In subsquent iterations, they are already merged with text and cached
+        # NOTE: first iteration doesn't have to be prefill, it can be the first
+        # iteration with a question and cached system prompt (continue generate from cache)
+        if is_first_iteration or not kwargs.get("use_cache", True):
             model_inputs["pixel_values"] = pixel_values
             model_inputs["image_sizes"] = image_sizes