Aluode/PerceptionLabPortable

1# coding=utf-82# Copyright 2021 The OpenAI Team Authors and The HuggingFace 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 IdeficsVision model: a copy of CLIPVisionModel using a simpler config object"""16 17import math18from dataclasses import dataclass19from typing import Callable, Optional, Union20 21import torch22from torch import nn23 24from ...activations import ACT2FN25from ...modeling_layers import GradientCheckpointingLayer26from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling27from ...modeling_utils import ALL_ATTENTION_FUNCTIONS28from ...utils import (29    ModelOutput,30    can_return_tuple,31    logging,32)33from .configuration_idefics import IdeficsVisionConfig34 35 36logger = logging.get_logger(__name__)37 38 39@dataclass40class IdeficsVisionModelOutput(ModelOutput):41    """42    Base class for vision model's outputs that also contains image embeddings of the pooling of the last hidden states.43 44    Args:45        image_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim)` *optional* returned when model is initialized with `with_projection=True`):46            The image embeddings obtained by applying the projection layer to the pooler_output.47        last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):48            Sequence of hidden-states at the output of the last layer of the model.49        hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):50            Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +51            one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.52 53            Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.54        attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):55            Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,56            sequence_length)`.57 58            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention59            heads.60    """61 62    image_embeds: Optional[torch.FloatTensor] = None63    last_hidden_state: Optional[torch.FloatTensor] = None64    hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None65    attentions: Optional[tuple[torch.FloatTensor, ...]] = None66 67 68# Adapted from transformers.models.clip.modeling_clip.CLIPVisionEmbeddings69class IdeficsVisionEmbeddings(nn.Module):70    def __init__(self, config: IdeficsVisionConfig):71        super().__init__()72        self.config = config73        self.embed_dim = config.hidden_size74        self.image_size = config.image_size75        self.patch_size = config.patch_size76 77        self.class_embedding = nn.Parameter(torch.randn(self.embed_dim))78 79        self.patch_embedding = nn.Conv2d(80            in_channels=config.num_channels,81            out_channels=self.embed_dim,82            kernel_size=self.patch_size,83            stride=self.patch_size,84            bias=False,85        )86 87        self.num_patches = (self.image_size // self.patch_size) ** 288        self.num_positions = self.num_patches + 189        self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim)90        self.register_buffer("position_ids", torch.arange(self.num_positions).expand((1, -1)), persistent=False)91 92    # Heavily inspired from https://github.com/huggingface/transformers/blob/v4.33.0/src/transformers/models/vit/modeling_vit.py#L8293    def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: int) -> torch.Tensor:94        """95        This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher96        resolution images.97 98        Source:99        https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174100        """101 102        num_patches = embeddings.shape[1] - 1103        pos_embed = self.position_embedding(self.position_ids)104        num_positions = pos_embed.shape[1] - 1105        if num_patches == num_positions and height == width:106            return pos_embed107        class_pos_embed = pos_embed[:, 0]108        patch_pos_embed = pos_embed[:, 1:]109 110        embed_dim = embeddings.shape[-1]111        num_h_patches = height // self.config.patch_size112        num_w_patches = width // self.config.patch_size113        # we add a small number to avoid floating point error in the interpolation114        # see discussion at https://github.com/facebookresearch/dino/issues/8115        num_h_patches, num_w_patches = num_h_patches + 0.1, num_w_patches + 0.1116        sqrt_num_positions = math.sqrt(num_positions)117        patch_pos_embed = patch_pos_embed.reshape(1, int(sqrt_num_positions), int(sqrt_num_positions), embed_dim)118        patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2)119        fp32_upcasting = patch_pos_embed.dtype == torch.bfloat16120        if fp32_upcasting:121            logger.warning_once(122                "Upcasting patch_pos_embed to fp32 for interpolation since `upsample_bicubic2d_out_frame` in nn.functional.interpolate "123                "is not implemented for 'torch.bfloat16' dtype. This will result in a slight overhead."124            )125            patch_pos_embed = patch_pos_embed.to(torch.float)126        patch_pos_embed = nn.functional.interpolate(127            patch_pos_embed,128            scale_factor=(num_h_patches / sqrt_num_positions, num_w_patches / sqrt_num_positions),129            mode="bicubic",130            align_corners=False,131        )132        if fp32_upcasting:133            patch_pos_embed = patch_pos_embed.to(torch.bfloat16)134        if int(num_h_patches) != patch_pos_embed.shape[-2] or int(num_w_patches) != patch_pos_embed.shape[-1]:135            raise ValueError(136                f"Number of patches for images ({int(num_h_patches), int(num_w_patches)}) don't match the "137                f"shape of position embedding ({patch_pos_embed.shape[-2], patch_pos_embed.shape[-1]})"138            )139        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, embed_dim)140        return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1)141 142    def forward(self, pixel_values: torch.FloatTensor, interpolate_pos_encoding: bool = False) -> torch.Tensor:143        batch_size, num_channels, height, width = pixel_values.shape144        if not interpolate_pos_encoding:145            if height != self.image_size or width != self.image_size:146                raise ValueError(147                    f"Input image size ({height}*{width}) doesn't match model"148                    f" ({self.image_size}*{self.image_size}). You should try to set `interpolate_pos_encoding=True`"149                )150 151        target_dtype = self.patch_embedding.weight.dtype152        patch_embeds = self.patch_embedding(pixel_values.to(dtype=target_dtype))  # shape = [*, width, grid, grid]153 154        patch_embeds = patch_embeds.flatten(2).transpose(1, 2)155 156        class_embeds = self.class_embedding.expand(batch_size, 1, -1)157        embeddings = torch.cat([class_embeds, patch_embeds], dim=1)158 159        # add positional encoding to each token160        if interpolate_pos_encoding:161            embeddings = embeddings + self.interpolate_pos_encoding(embeddings, height, width)162        else:163            embeddings = embeddings + self.position_embedding(self.position_ids)164 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 192class IdeficsVisionAttention(nn.Module):193    """Multi-headed attention from 'Attention Is All You Need' paper"""194 195    def __init__(self, config: IdeficsVisionConfig):196        super().__init__()197        self.config = config198        self.embed_dim = config.hidden_size199        self.num_heads = config.num_attention_heads200        self.head_dim = self.embed_dim // self.num_heads201        if self.head_dim * self.num_heads != self.embed_dim:202            raise ValueError(203                f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"204                f" {self.num_heads})."205            )206        self.scale = self.head_dim**-0.5207        self.dropout = config.attention_dropout208        self.is_causal = False209 210        self.k_proj = nn.Linear(self.embed_dim, self.embed_dim)211        self.v_proj = nn.Linear(self.embed_dim, self.embed_dim)212        self.q_proj = nn.Linear(self.embed_dim, self.embed_dim)213        self.out_proj = nn.Linear(self.embed_dim, self.embed_dim)214 215    def forward(216        self,217        hidden_states: torch.Tensor,218        attention_mask: Optional[torch.Tensor] = None,219        causal_attention_mask: Optional[torch.Tensor] = None,220        output_attentions: Optional[bool] = False,221    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:222        """Input shape: Batch x Time x Channel"""223 224        batch_size, seq_length, embed_dim = hidden_states.shape225 226        queries = self.q_proj(hidden_states)227        keys = self.k_proj(hidden_states)228        values = self.v_proj(hidden_states)229 230        queries = queries.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2)231        keys = keys.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2)232        values = values.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2)233        # CLIP text model uses both `causal_attention_mask` and `attention_mask`234        # in case FA2 kernel is called, `is_causal` should be inferred from `causal_attention_mask`235        if self.config._attn_implementation != "flash_attention_2":236            if attention_mask is not None and causal_attention_mask is not None:237                attention_mask = attention_mask + causal_attention_mask238            elif causal_attention_mask is not None:239                attention_mask = causal_attention_mask240        else:241            self.is_causal = causal_attention_mask is not None242 243        attention_interface: Callable = eager_attention_forward244        if self.config._attn_implementation != "eager":245            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]246 247        attn_output, attn_weights = attention_interface(248            self,249            queries,250            keys,251            values,252            attention_mask,253            is_causal=self.is_causal,254            scaling=self.scale,255            dropout=0.0 if not self.training else self.dropout,256        )257 258        attn_output = attn_output.reshape(batch_size, seq_length, embed_dim).contiguous()259        attn_output = self.out_proj(attn_output)260        if not output_attentions:261            attn_weights = None262        return attn_output, attn_weights263 264 265# Copied from transformers.models.clip.modeling_clip.CLIPMLP with CLIP->IdeficsVision266class IdeficsVisionMLP(nn.Module):267    def __init__(self, config):268        super().__init__()269        self.config = config270        self.activation_fn = ACT2FN[config.hidden_act]271        self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)272        self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)273 274    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:275        hidden_states = self.fc1(hidden_states)276        hidden_states = self.activation_fn(hidden_states)277        hidden_states = self.fc2(hidden_states)278        return hidden_states279 280 281# Copied from transformers.models.altclip.modeling_altclip.AltCLIPEncoderLayer with AltCLIP->IdeficsVision282class IdeficsVisionEncoderLayer(GradientCheckpointingLayer):283    def __init__(self, config: IdeficsVisionConfig):284        super().__init__()285        self.embed_dim = config.hidden_size286        self.self_attn = IdeficsVisionAttention(config)287        self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)288        self.mlp = IdeficsVisionMLP(config)289        self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)290 291    def forward(292        self,293        hidden_states: torch.Tensor,294        attention_mask: torch.Tensor,295        causal_attention_mask: torch.Tensor,296        output_attentions: Optional[bool] = False,297    ) -> tuple[torch.FloatTensor]:298        """299        Args:300            hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`301            attention_mask (`torch.FloatTensor`): attention mask of size302                `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.303                `(config.encoder_attention_heads,)`.304            output_attentions (`bool`, *optional*):305                Whether or not to return the attentions tensors of all attention layers. See `attentions` under306                returned tensors for more detail.307        """308        residual = hidden_states309 310        hidden_states = self.layer_norm1(hidden_states)311        hidden_states, attn_weights = self.self_attn(312            hidden_states=hidden_states,313            attention_mask=attention_mask,314            causal_attention_mask=causal_attention_mask,315            output_attentions=output_attentions,316        )317        hidden_states = residual + hidden_states318 319        residual = hidden_states320        hidden_states = self.layer_norm2(hidden_states)321        hidden_states = self.mlp(hidden_states)322        hidden_states = residual + hidden_states323 324        outputs = (hidden_states,)325 326        if output_attentions:327            outputs += (attn_weights,)328 329        return outputs330 331 332# Copied from transformers.models.altclip.modeling_altclip.AltCLIPEncoder with AltCLIP->IdeficsVision333class IdeficsVisionEncoder(nn.Module):334    """335    Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a336    [`IdeficsVisionEncoderLayer`].337 338    Args:339        config: IdeficsVisionConfig340    """341 342    def __init__(self, config: IdeficsVisionConfig):343        super().__init__()344        self.config = config345        self.layers = nn.ModuleList([IdeficsVisionEncoderLayer(config) for _ in range(config.num_hidden_layers)])346        self.gradient_checkpointing = False347 348    @can_return_tuple349    def forward(350        self,351        inputs_embeds,352        attention_mask: Optional[torch.Tensor] = None,353        causal_attention_mask: Optional[torch.Tensor] = None,354        output_attentions: Optional[bool] = None,355        output_hidden_states: Optional[bool] = None,356        return_dict: Optional[bool] = None,357    ) -> Union[tuple, BaseModelOutput]:358        r"""359        Args:360            inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):361                Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation.362                This is useful if you want more control over how to convert `input_ids` indices into associated vectors363                than the model's internal embedding lookup matrix.364            attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):365                Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:366 367                - 1 for tokens that are **not masked**,368                - 0 for tokens that are **masked**.369 370                [What are attention masks?](../glossary#attention-mask)371            causal_attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):372                Causal mask for the text model. Mask values selected in `[0, 1]`:373 374                - 1 for tokens that are **not masked**,375                - 0 for tokens that are **masked**.376 377                [What are attention masks?](../glossary#attention-mask)378            output_attentions (`bool`, *optional*):379                Whether or not to return the attentions tensors of all attention layers. See `attentions` under380                returned tensors for more detail.381            output_hidden_states (`bool`, *optional*):382                Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors383                for more detail.384            return_dict (`bool`, *optional*):385                Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.386        """387        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions388        output_hidden_states = (389            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states390        )391        return_dict = return_dict if return_dict is not None else self.config.use_return_dict392 393        encoder_states = () if output_hidden_states else None394        all_attentions = () if output_attentions else None395 396        hidden_states = inputs_embeds397        for idx, encoder_layer in enumerate(self.layers):398            if output_hidden_states:399                encoder_states = encoder_states + (hidden_states,)400            layer_outputs = encoder_layer(401                hidden_states,402                attention_mask,403                causal_attention_mask,404                output_attentions=output_attentions,405            )406 407            hidden_states = layer_outputs[0]408 409            if output_attentions:410                all_attentions = all_attentions + (layer_outputs[1],)411 412        if output_hidden_states:413            encoder_states = encoder_states + (hidden_states,)414 415        return BaseModelOutput(416            last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions417        )418 419 420# Adapted from transformers.models.clip.modeling_clip.CLIPVisionTransformer421class IdeficsVisionTransformer(nn.Module):422    def __init__(self, config: IdeficsVisionConfig):423        super().__init__()424        self.config = config425        embed_dim = config.hidden_size426 427        self.embeddings = IdeficsVisionEmbeddings(config)428        self.pre_layrnorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)429        self.encoder = IdeficsVisionEncoder(config)430        self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)431 432    # Adapted from transformers.models.clip.modeling_clip.CLIPVisionTransformer.forward433    def forward(434        self,435        pixel_values: Optional[torch.FloatTensor] = None,436        output_attentions: Optional[bool] = None,437        output_hidden_states: Optional[bool] = None,438        interpolate_pos_encoding: Optional[bool] = False,439        return_dict: Optional[bool] = None,440    ) -> Union[tuple, BaseModelOutputWithPooling]:441        r"""442        Returns:443 444        """445        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions446        output_hidden_states = (447            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states448        )449        return_dict = return_dict if return_dict is not None else self.config.use_return_dict450 451        if pixel_values is None:452            raise ValueError("You have to specify pixel_values")453 454        hidden_states = self.embeddings(pixel_values, interpolate_pos_encoding=interpolate_pos_encoding)455        hidden_states = self.pre_layrnorm(hidden_states)456 457        encoder_outputs = self.encoder(458            inputs_embeds=hidden_states,459            output_attentions=output_attentions,460            output_hidden_states=output_hidden_states,461            return_dict=return_dict,462        )463 464        last_hidden_state = encoder_outputs[0]465        pooled_output = last_hidden_state[:, 0, :]466        pooled_output = self.post_layernorm(pooled_output)467 468        if not return_dict:469            return (last_hidden_state, pooled_output) + encoder_outputs[1:]470 471        return BaseModelOutputWithPooling(472            last_hidden_state=last_hidden_state,473            pooler_output=pooled_output,474            hidden_states=encoder_outputs.hidden_states,475            attentions=encoder_outputs.attentions,476        )477