Aluode/PerceptionLabPortable
0
1# coding=utf-82# Copyright 2024 the HuggingFace Inc. team. All rights reserved.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8# http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15"""PyTorch Idefics3 model."""16 17from dataclasses import dataclass18from typing import Callable, Optional, Union19 20import torch21from torch import nn22 23from ...activations import ACT2FN24from ...cache_utils import Cache, DynamicCache25from ...generation import GenerationMixin26from ...modeling_attn_mask_utils import _prepare_4d_attention_mask27from ...modeling_flash_attention_utils import FlashAttentionKwargs28from ...modeling_layers import GradientCheckpointingLayer29from ...modeling_outputs import BaseModelOutput, ModelOutput30from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel31from ...processing_utils import Unpack32from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging33from ...utils.generic import check_model_inputs34from ..auto import AutoModel35from .configuration_idefics3 import Idefics3Config, Idefics3VisionConfig36 37 38logger = logging.get_logger(__name__)39 40 41@dataclass42@auto_docstring(43 custom_intro="""44 Base class for Idefics3 model's outputs that may also contain a past key/values (to speed up sequential decoding).45 """46)47class Idefics3BaseModelOutputWithPast(ModelOutput):48 r"""49 last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):50 Sequence of hidden-states at the output of the last layer of the model.51 If `past_key_values` is used only the last hidden-state of the sequences of shape `(batch_size, 1,52 hidden_size)` is output.53 past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):54 It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).55 56 Contains pre-computed hidden-states (key and values in the self-attention blocks and optionally if57 `config.is_encoder_decoder=True` in the cross-attention blocks) that can be used (see `past_key_values`58 input) to speed up sequential decoding.59 image_hidden_states (`tuple(torch.FloatTensor)`, *optional*):60 Tuple of `torch.FloatTensor` (one for the output of the image embeddings, `(batch_size, num_images,61 sequence_length, hidden_size)`.62 image_hidden_states of the model produced by the vision encoder63 """64 65 last_hidden_state: Optional[torch.FloatTensor] = None66 past_key_values: Optional[Cache] = None67 hidden_states: Optional[tuple[torch.FloatTensor]] = None68 attentions: Optional[tuple[torch.FloatTensor]] = None69 image_hidden_states: Optional[tuple[torch.FloatTensor]] = None70 71 72@dataclass73@auto_docstring(74 custom_intro="""75 Base class for Idefics causal language model (or autoregressive) outputs.76 """77)78class Idefics3CausalLMOutputWithPast(ModelOutput):79 r"""80 loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):81 Language modeling loss (for next-token prediction).82 logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):83 Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).84 past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):85 It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).86 87 Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see88 `past_key_values` input) to speed up sequential decoding.89 image_hidden_states (`tuple(torch.FloatTensor)`, *optional*):90 Tuple of `torch.FloatTensor` (one for the output of the image embeddings, `(batch_size, num_images,91 sequence_length, hidden_size)`.92 image_hidden_states of the model produced by the vision encoder93 """94 95 loss: Optional[torch.FloatTensor] = None96 logits: Optional[torch.FloatTensor] = None97 past_key_values: Optional[Cache] = None98 hidden_states: Optional[tuple[torch.FloatTensor]] = None99 attentions: Optional[tuple[torch.FloatTensor]] = None100 image_hidden_states: Optional[tuple[torch.FloatTensor]] = None101 102 103# Copied from transformers.models.idefics2.modeling_idefics2.Idefics2VisionEmbeddings with Idefics2->Idefics3104class Idefics3VisionEmbeddings(nn.Module):105 """106 This is a modified version of `siglip.modelign_siglip.SiglipVisionEmbeddings` to enable images of variable107 resolution.108 109 The modifications are adapted from [Patch n' Pack: NaViT, a Vision Transformer for any Aspect Ratio and Resolution](https://huggingface.co/papers/2307.06304)110 which allows treating images in their native aspect ratio and without the need to resize them to the same111 fixed size. In particular, we start from the original pre-trained SigLIP model112 (which uses images of fixed-size square images) and adapt it by training on images of variable resolutions.113 """114 115 def __init__(self, config: Idefics3VisionConfig):116 super().__init__()117 self.embed_dim = config.hidden_size118 self.image_size = config.image_size119 self.patch_size = config.patch_size120 121 self.patch_embedding = nn.Conv2d(122 in_channels=config.num_channels,123 out_channels=self.embed_dim,124 kernel_size=self.patch_size,125 stride=self.patch_size,126 padding="valid",127 )128 129 self.num_patches_per_side = self.image_size // self.patch_size130 self.num_patches = self.num_patches_per_side**2131 self.num_positions = self.num_patches132 self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim)133 134 def forward(self, pixel_values: torch.FloatTensor, patch_attention_mask: torch.BoolTensor) -> torch.Tensor:135 batch_size, _, max_im_h, max_im_w = pixel_values.shape136 137 patch_embeds = self.patch_embedding(pixel_values)138 embeddings = patch_embeds.flatten(2).transpose(1, 2)139 140 max_nb_patches_h, max_nb_patches_w = max_im_h // self.patch_size, max_im_w // self.patch_size141 boundaries = torch.arange(142 1 / self.num_patches_per_side, 1.0, 1 / self.num_patches_per_side, device=pixel_values.device143 )144 position_ids = torch.full(145 size=(batch_size, max_nb_patches_h * max_nb_patches_w), fill_value=0, device=pixel_values.device146 )147 148 for batch_idx, p_attn_mask in enumerate(patch_attention_mask):149 nb_patches_h = p_attn_mask[:, 0].sum()150 nb_patches_w = p_attn_mask[0].sum()151 152 h_indices = torch.arange(nb_patches_h, device=position_ids.device, dtype=pixel_values.dtype)153 w_indices = torch.arange(nb_patches_w, device=position_ids.device, dtype=pixel_values.dtype)154 155 fractional_coords_h = h_indices / nb_patches_h * (1 - 1e-6)156 fractional_coords_w = w_indices / nb_patches_w * (1 - 1e-6)157 158 bucket_coords_h = torch.bucketize(fractional_coords_h, boundaries, right=True)159 bucket_coords_w = torch.bucketize(fractional_coords_w, boundaries, right=True)160 161 pos_ids = (bucket_coords_h[:, None] * self.num_patches_per_side + bucket_coords_w).flatten()162 position_ids[batch_idx][p_attn_mask.view(-1)] = pos_ids163 164 embeddings = embeddings + self.position_embedding(position_ids)165 return embeddings166 167 168# Copied from transformers.models.siglip.modeling_siglip.eager_attention_forward169def eager_attention_forward(170 module: nn.Module,171 query: torch.Tensor,172 key: torch.Tensor,173 value: torch.Tensor,174 attention_mask: Optional[torch.Tensor],175 scaling: float,176 dropout: float = 0.0,177 **kwargs,178):179 attn_weights = torch.matmul(query, key.transpose(-1, -2)) * scaling180 if attention_mask is not None:181 attn_weights = attn_weights + attention_mask182 183 attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)184 attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)185 186 attn_output = torch.matmul(attn_weights, value)187 attn_output = attn_output.transpose(1, 2).contiguous()188 189 return attn_output, attn_weights190 191 192# Copied from transformers.models.siglip.modeling_siglip.SiglipAttention with Siglip->Idefics3Vision193class Idefics3VisionAttention(nn.Module):194 """Multi-headed attention from 'Attention Is All You Need' paper"""195 196 # Copied from transformers.models.clip.modeling_clip.CLIPAttention.__init__197 def __init__(self, config):198 super().__init__()199 self.config = config200 self.embed_dim = config.hidden_size201 self.num_heads = config.num_attention_heads202 self.head_dim = self.embed_dim // self.num_heads203 if self.head_dim * self.num_heads != self.embed_dim:204 raise ValueError(205 f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"206 f" {self.num_heads})."207 )208 self.scale = self.head_dim**-0.5209 self.dropout = config.attention_dropout210 211 self.k_proj = nn.Linear(self.embed_dim, self.embed_dim)212 self.v_proj = nn.Linear(self.embed_dim, self.embed_dim)213 self.q_proj = nn.Linear(self.embed_dim, self.embed_dim)214 self.out_proj = nn.Linear(self.embed_dim, self.embed_dim)215 216 # Ignore copy217 self.is_causal = False218 219 def forward(220 self,221 hidden_states: torch.Tensor,222 attention_mask: Optional[torch.Tensor] = None,223 **kwargs,224 ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:225 """Input shape: Batch x Time x Channel"""226 227 batch_size, seq_length, embed_dim = hidden_states.shape228 229 queries = self.q_proj(hidden_states)230 keys = self.k_proj(hidden_states)231 values = self.v_proj(hidden_states)232 233 queries = queries.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2)234 keys = keys.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2)235 values = values.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2)236 237 attention_interface: Callable = eager_attention_forward238 if self.config._attn_implementation != "eager":239 attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]240 241 attn_output, attn_weights = attention_interface(242 self,243 queries,244 keys,245 values,246 attention_mask,247 is_causal=self.is_causal,248 scaling=self.scale,249 dropout=0.0 if not self.training else self.dropout,250 )251 252 attn_output = attn_output.reshape(batch_size, seq_length, embed_dim).contiguous()253 attn_output = self.out_proj(attn_output)254 255 return attn_output, attn_weights256 257 258# Copied from transformers.models.siglip.modeling_siglip.SiglipMLP with Siglip->Idefics3Vision259class Idefics3VisionMLP(nn.Module):260 def __init__(self, config):261 super().__init__()262 self.config = config263 self.activation_fn = ACT2FN[config.hidden_act]264 self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)265 self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)266 267 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:268 hidden_states = self.fc1(hidden_states)269 hidden_states = self.activation_fn(hidden_states)270 hidden_states = self.fc2(hidden_states)271 return hidden_states272 273 274class Idefics3SimpleMLP(nn.Module):275 def __init__(self, config):276 super().__init__()277 input_size = config.vision_config.hidden_size * (config.scale_factor**2)278 output_size = config.text_config.hidden_size279 self.proj = nn.Linear(input_size, output_size, bias=False)280 281 def forward(self, x):282 return self.proj(x)283 284 285# Copied from transformers.models.idefics2.modeling_idefics2.Idefics2EncoderLayer with Idefics2->Idefics3286class Idefics3EncoderLayer(GradientCheckpointingLayer):287 def __init__(self, config: Idefics3VisionConfig):288 super().__init__()289 self.embed_dim = config.hidden_size290 self.self_attn = Idefics3VisionAttention(config)291 self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)292 self.mlp = Idefics3VisionMLP(config)293 self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)294 295 @auto_docstring296 # Copied from transformers.models.siglip.modeling_siglip.SiglipEncoderLayer.forward297 def forward(298 self,299 hidden_states: torch.Tensor,300 attention_mask: torch.Tensor,301 **kwargs: Unpack[TransformersKwargs],302 ) -> torch.FloatTensor:303 residual = hidden_states304 305 hidden_states = self.layer_norm1(hidden_states)306 hidden_states, _ = self.self_attn(307 hidden_states=hidden_states,308 attention_mask=attention_mask,309 **kwargs,310 )311 hidden_states = residual + hidden_states312 313 residual = hidden_states314 hidden_states = self.layer_norm2(hidden_states)315 hidden_states = self.mlp(hidden_states)316 hidden_states = residual + hidden_states317 318 return hidden_states319 320 321# Copied from transformers.models.siglip.modeling_siglip.SiglipEncoder with Siglip->Idefics3322class Idefics3Encoder(nn.Module):323 """324 Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a325 [`Idefics3EncoderLayer`].326 327 Args:328 config: Idefics3Config329 """330 331 def __init__(self, config: Idefics3Config):332 super().__init__()333 self.config = config334 self.layers = nn.ModuleList([Idefics3EncoderLayer(config) for _ in range(config.num_hidden_layers)])335 self.gradient_checkpointing = False336 337 # Ignore copy338 @auto_docstring339 def forward(340 self,341 inputs_embeds,342 attention_mask: Optional[torch.Tensor] = None,343 ) -> Union[tuple, BaseModelOutput]:344 hidden_states = inputs_embeds345 for encoder_layer in self.layers:346 layer_outputs = encoder_layer(347 hidden_states,348 attention_mask,349 )350 351 hidden_states = layer_outputs352 353 return BaseModelOutput(last_hidden_state=hidden_states)354 355 356# Copied from transformers.models.llama.modeling_llama.repeat_kv357def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:358 """359 This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,360 num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)361 """362 batch, num_key_value_heads, slen, head_dim = hidden_states.shape363 if n_rep == 1:364 return hidden_states365 hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)366 return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)367 368 369# Copied from transformers.models.llama.modeling_llama.LlamaRMSNorm with Llama->Idefics3370class Idefics3RMSNorm(nn.Module):371 def __init__(self, hidden_size, eps=1e-6):372 """373 Idefics3RMSNorm is equivalent to T5LayerNorm374 """375 super().__init__()376 self.weight = nn.Parameter(torch.ones(hidden_size))377 self.variance_epsilon = eps378 379 def forward(self, hidden_states):380 input_dtype = hidden_states.dtype381 hidden_states = hidden_states.to(torch.float32)382 variance = hidden_states.pow(2).mean(-1, keepdim=True)383 hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)384 return self.weight * hidden_states.to(input_dtype)385 386 def extra_repr(self):387 return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"388 389 390class Idefics3Connector(nn.Module):391 def __init__(self, config):392 super().__init__()393 self.scale_factor = config.scale_factor394 self.modality_projection = Idefics3SimpleMLP(config)395 396 def pixel_shuffle(self, x, scale_factor=2):397 bsz, seq, embed_dim = x.size()398 height = width = int(seq**0.5)399 x = x.view(bsz, height, width, embed_dim)400 x = x.view(bsz, height, int(width / scale_factor), embed_dim * scale_factor)401 x = x.permute(0, 2, 1, 3)402 x = x.reshape(bsz, int(width / scale_factor), int(height / scale_factor), embed_dim * (scale_factor**2))403 x = x.permute(0, 2, 1, 3)404 x = x.reshape(bsz, int(seq / (scale_factor**2)), embed_dim * (scale_factor**2))405 return x406 407 def forward(self, image_hidden_states):408 image_hidden_states = self.pixel_shuffle(image_hidden_states, self.scale_factor)409 image_hidden_states = self.modality_projection(image_hidden_states)410 return image_hidden_states411 412 413@auto_docstring414class Idefics3PreTrainedModel(PreTrainedModel):415 config: Idefics3Config416 base_model_prefix = "model"417 supports_gradient_checkpointing = True418 _no_split_modules = ["Idefics3VisionAttention", "Idefics3DecoderLayer"]419 _skip_keys_device_placement = "past_key_values"420 _supports_flash_attn = True421 _supports_sdpa = True422 _supports_flex_attn = True423 424 _supports_attention_backend = True425 426 def _init_weights(self, module):427 std = getattr(self.config, "initializer_range", self.config.get_text_config().initializer_range)428 429 if isinstance(module, (nn.Linear, nn.Conv2d)):430 module.weight.data.normal_(mean=0.0, std=std)431 if module.bias is not None:432 module.bias.data.zero_()433 elif isinstance(module, nn.Embedding):434 module.weight.data.normal_(mean=0.0, std=std)435 if module.padding_idx is not None:436 module.weight.data[module.padding_idx].zero_()437 elif isinstance(module, nn.LayerNorm):438 module.weight.data.fill_(1.0)439 module.bias.data.zero_()440 elif isinstance(module, Idefics3RMSNorm):441 module.weight.data.fill_(1.0)442 443 444@auto_docstring(445 custom_intro="""446 The Idefics3 Vision Transformer Model outputting raw image embedding.447 """448)449class Idefics3VisionTransformer(Idefics3PreTrainedModel):450 config: Idefics3VisionConfig451 _supports_sdpa = True452 _supports_flash_attn = True453 _supports_flex_attn = True454 _can_record_outputs = {455 "hidden_states": Idefics3EncoderLayer,456 "attentions": Idefics3VisionAttention,457 }458 459 def __init__(self, config: Idefics3VisionConfig):460 super().__init__(config)461 embed_dim = config.hidden_size462 463 self.embeddings = Idefics3VisionEmbeddings(config)464 self.encoder = Idefics3Encoder(config)465 self.patch_size = config.patch_size466 self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)467 468 # Copied from transformers.models.idefics2.modeling_idefics2.Idefics2VisionTransformer.get_input_embeddings469 def get_input_embeddings(self):470 return self.embeddings471 472 # Copied from transformers.models.idefics2.modeling_idefics2.Idefics2VisionTransformer.set_input_embeddings473 def set_input_embeddings(self, value):474 self.embeddings = value475 476 @check_model_inputs(tie_last_hidden_states=False)477 def forward(478 self,479 pixel_values,480 patch_attention_mask: Optional[torch.BoolTensor] = None,481 **kwargs: Unpack[TransformersKwargs],482 ) -> Union[tuple, BaseModelOutput]:483 batch_size = pixel_values.size(0)484 if patch_attention_mask is None:485 patch_size = self.patch_size486 patch_attention_mask = torch.ones(487 (488 batch_size,489 pixel_values.size(2) // patch_size,490 pixel_values.size(3) // patch_size,491 )492 )493 patch_attention_mask = patch_attention_mask.to(dtype=torch.bool, device=pixel_values.device)494 495 hidden_states = self.embeddings(pixel_values=pixel_values, patch_attention_mask=patch_attention_mask)496 497 patch_attention_mask = patch_attention_mask.view(batch_size, -1)498 # The call to `_upad_input` in `_flash_attention_forward` is expensive499 # So when the `patch_attention_mask` is full of 1s (i.e. attending to the whole sequence),500 # avoiding passing the attention_mask, which is equivalent to attending to the full sequence501 if self.config._attn_implementation != "flash_attention_2":502 patch_attention_mask = _prepare_4d_attention_mask(patch_attention_mask, hidden_states.dtype)503 elif not torch.any(~patch_attention_mask):504 patch_attention_mask = None505 506 encoder_outputs: BaseModelOutput = self.encoder(507 inputs_embeds=hidden_states,508 attention_mask=patch_attention_mask,509 )510 511 last_hidden_state = encoder_outputs.last_hidden_state512 last_hidden_state = self.post_layernorm(last_hidden_state)513 514 return BaseModelOutput(515 last_hidden_state=last_hidden_state,516 )517 518 519@auto_docstring(520 custom_intro="""521 Idefics3 model consisting of a SIGLIP vision encoder and Llama3 language decoder522 """523)524class Idefics3Model(Idefics3PreTrainedModel):525 def __init__(self, config: Idefics3Config):526 super().__init__(config)527 self.padding_idx = self.config.text_config.pad_token_id528 self.vocab_size = self.config.text_config.vocab_size529 530 self.vision_model = Idefics3VisionTransformer._from_config(config.vision_config)531 self.connector = Idefics3Connector(config)532 self.text_model = AutoModel.from_config(config.text_config)533 534 self.image_seq_len = int(535 ((config.vision_config.image_size // config.vision_config.patch_size) ** 2) / (config.scale_factor**2)536 )537 self.image_token_id = self.config.image_token_id538 539 self.post_init()540 541 # Copied from transformers.models.idefics2.modeling_idefics2.Idefics2Model.enable_input_require_grads542 def enable_input_require_grads(self):543 """544 Enables the gradients for the input embeddings.545 546 This is useful for lora when using gradient checkpointing.547 c.f. https://github.com/huggingface/peft/issues/1402#issuecomment-1913675032548 549 Override to set output.requires_grad = True for both the decoder's and vision model's embeddings.550 """551 552 def get_lowest_module(module):553 if len(list(module.children())) == 0:554 # If the module has no children, it is a leaf module (e.g., Linear, Conv2d, etc.)555 return module556 else:557 # Recursively call the function on each child module558 return get_lowest_module(list(module.children())[0])559 560 def make_inputs_require_grads(module, input, output):561 output.requires_grad_(True)562 563 self._text_require_grads_hook = self.get_input_embeddings().register_forward_hook(make_inputs_require_grads)564 self._vision_require_grads_hook = get_lowest_module(self.vision_model).register_forward_hook(565 make_inputs_require_grads566 )567 568 # Copied from transformers.models.idefics2.modeling_idefics2.Idefics2Model.disable_input_require_grads569 def disable_input_require_grads(self):570 self._text_require_grads_hook.remove()571 self._vision_require_grads_hook.remove()572 573 # Copied from transformers.models.idefics2.modeling_idefics2.Idefics2Model.get_input_embeddings574 def get_input_embeddings(self):575 return self.text_model.get_input_embeddings()576 577 # Copied from transformers.models.idefics2.modeling_idefics2.Idefics2Model.set_input_embeddings578 def set_input_embeddings(self, value):579 self.text_model.set_input_embeddings(value)580 581 def inputs_merger(582 self,583 input_ids: torch.LongTensor,584 inputs_embeds: Optional[torch.Tensor],585 image_hidden_states: Optional[torch.Tensor],586 ):587 """588 This method aims at merging the token embeddings with the image hidden states into one single sequence of vectors that are fed to the transformer LM.589 The merging happens as follows:590 - The text token sequence is: `tok_1 tok_2 tok_3 <fake_token_around_image> <image> <image> ... <image> <fake_token_around_image> tok_4`.591 - We get the image hidden states for the image through the vision encoder and that hidden state, after a pixel shuffle operation, is then projected into the text embedding space.592 We thus have a sequence of image hidden states of size (1, image_seq_len, hidden_dim), where 1 is for batch_size of 1 image and hidden_dim is the hidden_dim of the LM transformer.593 - The merging happens so that we obtain the following sequence: `vector_tok_1 vector_tok_2 vector_tok_3 vector_fake_tok_around_image {sequence of image_seq_len image hidden states} vector_fake_toke_around_image vector_tok_4`. That sequence is fed to the LM.594 - To fit the format of that sequence, `input_ids`, `input_embeds`, `attention_mask` are all 3 adapted to insert the image hidden states.595 """596 if input_ids is None:597 special_image_mask = inputs_embeds == self.get_input_embeddings()(598 torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device)599 )600 special_image_mask = special_image_mask.all(-1)601 else:602 special_image_mask = input_ids == self.config.image_token_id603 604 special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)605 image_hidden_states = image_hidden_states.to(inputs_embeds.device, inputs_embeds.dtype)606 inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_hidden_states)607 return inputs_embeds608 609 def get_image_features(610 self, pixel_values: torch.FloatTensor, pixel_attention_mask: Optional[torch.LongTensor] = None611 ):612 """613 Encodes images into continuous embeddings that can be forwarded to the language model.614 615 Args:616 pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`):617 The tensors corresponding to the input images.618 pixel_attention_mask (`torch.LongTensor`, *optional*):619 The attention mask indicating padded regions in the image.620 """621 batch_size, num_images, num_channels, height, width = pixel_values.shape622 pixel_values = pixel_values.to(dtype=self.dtype) # fp16 compatibility623 pixel_values = pixel_values.view(batch_size * num_images, *pixel_values.shape[2:])624 625 # Remove padding images - padding images are full 0.626 nb_values_per_image = pixel_values.shape[1:].numel()627 real_images_inds = (pixel_values == 0.0).sum(dim=(-1, -2, -3)) != nb_values_per_image628 pixel_values = pixel_values[real_images_inds].contiguous()629 630 # Handle the vision attention mask631 if pixel_attention_mask is None:632 pixel_attention_mask = torch.ones(633 size=(pixel_values.size(0), pixel_values.size(2), pixel_values.size(3)),634 dtype=torch.bool,635 device=pixel_values.device,636 )637 else:638 # Remove padding images from the mask639 pixel_attention_mask = pixel_attention_mask.view(batch_size * num_images, *pixel_attention_mask.shape[2:])640 pixel_attention_mask = pixel_attention_mask[real_images_inds].contiguous()641 642 patch_size = self.config.vision_config.patch_size643 patches_subgrid = pixel_attention_mask.unfold(dimension=1, size=patch_size, step=patch_size)644 patches_subgrid = patches_subgrid.unfold(dimension=2, size=patch_size, step=patch_size)645 patch_attention_mask = (patches_subgrid.sum(dim=(-1, -2)) > 0).bool()646 647 # Get sequence from the vision encoder648 image_hidden_states = self.vision_model(pixel_values=pixel_values, patch_attention_mask=patch_attention_mask)649 image_hidden_states.last_hidden_state650 651 # Modality projection & resampling652 image_hidden_states = self.connector(image_hidden_states.last_hidden_state)653 return image_hidden_states654 655 @can_return_tuple656 @auto_docstring(657 custom_intro="""658 Inputs fed to the model can have an arbitrary number of images. To account for this, pixel_values fed to659 the model have image padding -> (batch_size, max_num_images, 3, max_heights, max_widths) where660 max_num_images is the maximum number of images among the batch_size samples in the batch.661 Padding images are not needed beyond padding the pixel_values at the entrance of the model.662 For efficiency, we only pass through the vision_model's forward the real images by663 discarding the padding images i.e. pixel_values of size (image_batch_size, 3, height, width) where664 image_batch_size would be 7 when num_images_per_sample=[1, 3, 1, 2] and max_num_images would be 3.665 """666 )667 def forward(668 self,669 input_ids: Optional[torch.LongTensor] = None,670 attention_mask: Optional[torch.Tensor] = None,671 position_ids: Optional[torch.LongTensor] = None,672 past_key_values: Optional[Cache] = None,673 inputs_embeds: Optional[torch.FloatTensor] = None,674 pixel_values: Optional[torch.FloatTensor] = None,675 pixel_attention_mask: Optional[torch.BoolTensor] = None,676 image_hidden_states: Optional[torch.FloatTensor] = None,677 use_cache: Optional[bool] = None,678 output_attentions: Optional[bool] = None,679 output_hidden_states: Optional[bool] = None,680 cache_position: Optional[torch.LongTensor] = None,681 return_dict: Optional[bool] = None,682 **kwargs: Unpack[FlashAttentionKwargs],683 ) -> Union[tuple, Idefics3BaseModelOutputWithPast]:684 r"""685 pixel_attention_mask (`torch.Tensor` of shape `(batch_size, image_size, image_size)`, *optional*):686 Mask to avoid performing attention on padding pixel indices.687 image_hidden_states (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`):688 The hidden states of the image encoder after modality projection.689 """690 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions691 output_hidden_states = (692 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states693 )694 use_cache = use_cache if use_cache is not None else self.config.use_cache695 return_dict = return_dict if return_dict is not None else self.config.use_return_dict696 697 if self.training and self.text_model.gradient_checkpointing and use_cache:698 logger.warning_once(699 "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."700 )701 use_cache = False702 703 # retrieve input_ids and inputs_embeds704 if input_ids is not None:705 batch_size, seq_length = input_ids.shape706 elif inputs_embeds is not None:707 batch_size, seq_length, _ = inputs_embeds.shape708 else:709 raise ValueError("You have to specify either input_ids or inputs_embeds")710 711 if use_cache and past_key_values is None:712 past_key_values = DynamicCache(config=self.config)713 714 if inputs_embeds is None:715 inputs_embeds = self.text_model.get_input_embeddings()(input_ids).to(self.device)716 717 # START VISUAL INPUTS INTEGRATION718 if pixel_values is not None and image_hidden_states is not None:719 raise ValueError("You cannot specify both pixel_values and image_hidden_states at the same time")720 elif pixel_values is not None:721 image_hidden_states = self.get_image_features(pixel_values, pixel_attention_mask)722 elif image_hidden_states is not None:723 image_hidden_states = image_hidden_states.to(dtype=self.dtype, device=input_ids.device)724 725 if image_hidden_states is not None:726 # When we generate, we don't want to replace the potential image_token_id that we generated by images727 # that simply don't exist728 inputs_embeds = self.inputs_merger(729 input_ids=input_ids,730 inputs_embeds=inputs_embeds,731 image_hidden_states=image_hidden_states,732 )733 734 outputs = self.text_model(735 inputs_embeds=inputs_embeds,736 attention_mask=attention_mask,737 position_ids=position_ids,738 past_key_values=past_key_values,739 use_cache=use_cache,740 output_attentions=output_attentions,741 output_hidden_states=output_hidden_states,742 cache_position=cache_position,743 return_dict=True,744 **kwargs,745 )746 747 return Idefics3BaseModelOutputWithPast(748 last_hidden_state=outputs.last_hidden_state,749 past_key_values=outputs.past_key_values,750 hidden_states=outputs.hidden_states,751 attentions=outputs.attentions,752 image_hidden_states=image_hidden_states,753 )754 755 756@auto_docstring(757 custom_intro="""758 The Idefics3 Model with a language modeling head. It is made up a SigLIP vision encoder, with a language modeling head on top.759 """760)761class Idefics3ForConditionalGeneration(Idefics3PreTrainedModel, GenerationMixin):762 _tied_weights_keys = ["lm_head.weight"]763 764 # Copied from transformers.models.idefics2.modeling_idefics2.Idefics2ForConditionalGeneration.__init__ with Idefics2->Idefics3765 def __init__(self, config):766 super().__init__(config)767 self.model = Idefics3Model(config)768 self.image_token_id = self.config.image_token_id769 770 self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False)771 self.vocab_size = config.text_config.vocab_size772 773 # Initialize weights and apply final processing774 self.post_init()775 776 # Copied from transformers.models.idefics2.modeling_idefics2.Idefics2ForConditionalGeneration.enable_input_require_grads777 def enable_input_require_grads(self):778 """779 Enables the gradients for the input embeddings. This is useful for fine-tuning adapter weights while keeping780 the model weights fixed.781 """782 783 def make_inputs_require_grads(module, input, output):784 output.requires_grad_(True)785 786 self._text_require_grads_hook = self.get_input_embeddings().register_forward_hook(make_inputs_require_grads)787 self._vision_require_grads_hook = self.model.vision_model.get_input_embeddings().register_forward_hook(788 make_inputs_require_grads789 )790 791 # Copied from transformers.models.idefics2.modeling_idefics2.Idefics2ForConditionalGeneration.disable_input_require_grads792 def disable_input_require_grads(self):793 self._text_require_grads_hook.remove()794 self._vision_require_grads_hook.remove()795 796 # Copied from transformers.models.idefics2.modeling_idefics2.Idefics2ForConditionalGeneration.get_input_embeddings797 def get_input_embeddings(self):798 return self.model.text_model.get_input_embeddings()799 800 # Copied from transformers.models.idefics2.modeling_idefics2.Idefics2ForConditionalGeneration.set_input_embeddings801 def set_input_embeddings(self, value):802 self.model.text_model.set_input_embeddings(value)803 804 def get_image_features(805 self, pixel_values: torch.FloatTensor, pixel_attention_mask: Optional[torch.LongTensor] = None806 ):807 return self.model.get_image_features(pixel_values=pixel_values, pixel_attention_mask=pixel_attention_mask)808 809 @can_return_tuple810 @auto_docstring811 def forward(812 self,813 input_ids: Optional[torch.LongTensor] = None,814 attention_mask: Optional[torch.Tensor] = None,815 position_ids: Optional[torch.LongTensor] = None,816 past_key_values: Optional[Cache] = None,817 inputs_embeds: Optional[torch.FloatTensor] = None,818 pixel_values: Optional[torch.FloatTensor] = None,819 pixel_attention_mask: Optional[torch.BoolTensor] = None,820 image_hidden_states: Optional[torch.FloatTensor] = None,821 labels: Optional[torch.LongTensor] = None,822 use_cache: Optional[bool] = None,823 output_attentions: Optional[bool] = None,824 output_hidden_states: Optional[bool] = None,825 cache_position: Optional[torch.LongTensor] = None,826 return_dict: Optional[bool] = None,827 logits_to_keep: Union[int, torch.Tensor] = 0,828 **kwargs: Unpack[TransformersKwargs],829 ) -> Union[tuple, Idefics3CausalLMOutputWithPast]:830 r"""831 pixel_attention_mask (`torch.Tensor` of shape `(batch_size, image_size, image_size)`, *optional*):832 Mask to avoid performing attention on padding pixel indices.833 image_hidden_states (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`):834 The hidden states of the image encoder after modality projection.835 labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):836 Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,837 config.vocab_size]` or `model.image_token_id` (where `model` is your instance of `Idefics3ForConditionalGeneration`).838 Tokens with indices set to `model.image_token_id` are ignored (masked), the loss is only839 computed for the tokens with labels in `[0, ..., config.vocab_size]`.840 841 Example:842 843 ```python844 >>> import requests845 >>> import torch846 >>> from PIL import Image847 >>> from io import BytesIO848 849 >>> from transformers import AutoProcessor, AutoModelForVision2Seq850 >>> from transformers.image_utils import load_image851 852 >>> # Note that passing the image urls (instead of the actual pil images) to the processor is also possible853 >>> image1 = load_image("https://cdn.britannica.com/61/93061-050-99147DCE/Statue-of-Liberty-Island-New-York-Bay.jpg")854 >>> image2 = load_image("https://cdn.britannica.com/59/94459-050-DBA42467/Skyline-Chicago.jpg")855 >>> image3 = load_image("https://cdn.britannica.com/68/170868-050-8DDE8263/Golden-Gate-Bridge-San-Francisco.jpg")856 857 >>> processor = AutoProcessor.from_pretrained("HuggingFaceM4/Idefics3-8B-Llama3")858 >>> model = AutoModelForVision2Seq.from_pretrained("HuggingFaceM4/Idefics3-8B-Llama3", dtype=torch.bfloat16, device_map="auto")859 860 >>> # Create inputs861 >>> messages = [862 ... {863 ... "role": "user",864 ... "content": [865 ... {"type": "image"},866 ... {"type": "text", "text": "In this image, we can see the city of New York, and more specifically the Statue of Liberty."},867 ... {"type": "image"},868 ... {"type": "text", "text": "What can we see in this image?"},869 ... ]870 ... },871 ... {872 ... "role": "user",873 ... "content": [874 ... {"type": "image"},875 ... {"type": "text", "text": "In which city is that bridge located?"},876 ... ]877 ... }878 ... ]879 880 >>> prompts = [processor.apply_chat_template([message], add_generation_prompt=True) for message in messages]881 >>> images = [[image1, image2], [image3]]882 >>> inputs = processor(text=prompts, images=images, padding=True, return_tensors="pt").to(model.device)883 884 >>> # Generate885 >>> generated_ids = model.generate(**inputs, max_new_tokens=256)886 >>> generated_texts = processor.batch_decode(generated_ids, skip_special_tokens=True)887 888 >>> print(generated_texts[0])889 Assistant: There are buildings, trees, lights, and water visible in this image.890 891 >>> print(generated_texts[1])892 Assistant: The bridge is in San Francisco.893 ```"""894 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions895 output_hidden_states = (896 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states897 )898 return_dict = return_dict if return_dict is not None else self.config.use_return_dict899 900 # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)901 outputs = self.model(902 input_ids=input_ids,903 attention_mask=attention_mask,904 position_ids=position_ids,905 past_key_values=past_key_values,906 inputs_embeds=inputs_embeds,907 pixel_values=pixel_values,908 pixel_attention_mask=pixel_attention_mask,909 image_hidden_states=image_hidden_states,910 use_cache=use_cache,911 output_attentions=output_attentions,912 output_hidden_states=output_hidden_states,913 cache_position=cache_position,914 return_dict=True,915 **kwargs,916 )917 918 hidden_states = outputs[0]919 # Only compute necessary logits, and do not upcast them to float if we are not computing the loss920 slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep921 logits = self.lm_head(hidden_states[:, slice_indices, :])922 923 loss = None924 if labels is not None:925 loss = self.loss_function(926 logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size, **kwargs927 )928 929 return Idefics3CausalLMOutputWithPast(930 loss=loss,931 logits=logits,932 past_key_values=outputs.past_key_values,933 hidden_states=outputs.hidden_states,934 attentions=outputs.attentions,935 image_hidden_states=outputs.image_hidden_states,936 )937 938 # Copied from transformers.models.idefics2.modeling_idefics2.Idefics2ForConditionalGeneration.prepare_inputs_for_generation939 def prepare_inputs_for_generation(940 self,941 input_ids,942 past_key_values=None,943 attention_mask=None,944 inputs_embeds=None,945 cache_position=None,946 pixel_values=None,947 pixel_attention_mask=None,948 image_hidden_states=None,949 logits_to_keep=None,950 **kwargs,951 ):952 # Overwritten -- there are mutually exclusive inputs (if the logic to make `image_hidden_states` take953 # precedence is moved to the model, we can remove this fn)954 955 model_inputs = super().prepare_inputs_for_generation(956 input_ids,957 past_key_values=past_key_values,958 attention_mask=attention_mask,959 inputs_embeds=inputs_embeds,960 cache_position=cache_position,961 pixel_values=pixel_values,962 pixel_attention_mask=pixel_attention_mask,963 image_hidden_states=image_hidden_states,964 logits_to_keep=logits_to_keep,965 **kwargs,966 )967 968 if image_hidden_states is not None or cache_position[0] != 0:969 model_inputs["pixel_values"] = None970 model_inputs["pixel_attention_mask"] = None971 972 return model_inputs973 974 975__all__ = ["Idefics3ForConditionalGeneration", "Idefics3PreTrainedModel", "Idefics3Model", "Idefics3VisionTransformer"]976