ByteDance/Sa2VA-1B

1# --------------------------------------------------------2# InternVL3# Copyright (c) 2024 OpenGVLab4# Licensed under The MIT License [see LICENSE for details]5# --------------------------------------------------------6 7from typing import Optional, Tuple, Union8 9import torch10import torch.nn.functional as F11import torch.utils.checkpoint12from einops import rearrange13from timm.models.layers import DropPath14from torch import nn15from transformers.activations import ACT2FN16from transformers.modeling_outputs import (BaseModelOutput,17                                           BaseModelOutputWithPooling)18from transformers.modeling_utils import PreTrainedModel19from transformers.utils import logging20 21from .configuration_intern_vit import InternVisionConfig22 23try:24    from .flash_attention import FlashAttention25    has_flash_attn = True26except:27    print('FlashAttention is not installed.')28    has_flash_attn = False29 30logger = logging.get_logger(__name__)31 32 33class InternRMSNorm(nn.Module):34    def __init__(self, hidden_size, eps=1e-6):35        super().__init__()36        self.weight = nn.Parameter(torch.ones(hidden_size))37        self.variance_epsilon = eps38 39    def forward(self, hidden_states):40        input_dtype = hidden_states.dtype41        hidden_states = hidden_states.to(torch.float32)42        variance = hidden_states.pow(2).mean(-1, keepdim=True)43        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)44        return self.weight * hidden_states.to(input_dtype)45 46 47try:48    from apex.normalization import FusedRMSNorm49 50    InternRMSNorm = FusedRMSNorm  # noqa51 52    logger.info('Discovered apex.normalization.FusedRMSNorm - will use it instead of InternRMSNorm')53except ImportError:54    # using the normal InternRMSNorm55    pass56except Exception:57    logger.warning('discovered apex but it failed to load, falling back to InternRMSNorm')58    pass59 60 61NORM2FN = {62    'rms_norm': InternRMSNorm,63    'layer_norm': nn.LayerNorm,64}65 66 67class InternVisionEmbeddings(nn.Module):68    def __init__(self, config: InternVisionConfig):69        super().__init__()70        self.config = config71        self.embed_dim = config.hidden_size72        self.image_size = config.image_size73        self.patch_size = config.patch_size74 75        self.class_embedding = nn.Parameter(76            torch.randn(1, 1, self.embed_dim),77        )78 79        self.patch_embedding = nn.Conv2d(80            in_channels=3, out_channels=self.embed_dim, kernel_size=self.patch_size, stride=self.patch_size81        )82 83        self.num_patches = (self.image_size // self.patch_size) ** 284        self.num_positions = self.num_patches + 185 86        self.position_embedding = nn.Parameter(torch.randn(1, self.num_positions, self.embed_dim))87 88    def _get_pos_embed(self, pos_embed, H, W):89        target_dtype = pos_embed.dtype90        pos_embed = pos_embed.float().reshape(91            1, self.image_size // self.patch_size, self.image_size // self.patch_size, -1).permute(0, 3, 1, 2)92        pos_embed = F.interpolate(pos_embed, size=(H, W), mode='bicubic', align_corners=False). \93            reshape(1, -1, H * W).permute(0, 2, 1).to(target_dtype)94        return pos_embed95 96    def forward(self, pixel_values: torch.FloatTensor) -> torch.Tensor:97        target_dtype = self.patch_embedding.weight.dtype98        patch_embeds = self.patch_embedding(pixel_values)  # shape = [*, channel, width, height]99        batch_size, _, height, width = patch_embeds.shape100        patch_embeds = patch_embeds.flatten(2).transpose(1, 2)101        class_embeds = self.class_embedding.expand(batch_size, 1, -1).to(target_dtype)102        embeddings = torch.cat([class_embeds, patch_embeds], dim=1)103        position_embedding = torch.cat([104            self.position_embedding[:, :1, :],105            self._get_pos_embed(self.position_embedding[:, 1:, :], height, width)106        ], dim=1)107        embeddings = embeddings + position_embedding.to(target_dtype)108        return embeddings109 110 111class InternAttention(nn.Module):112    """Multi-headed attention from 'Attention Is All You Need' paper"""113 114    def __init__(self, config: InternVisionConfig):115        super().__init__()116        self.config = config117        self.embed_dim = config.hidden_size118        self.num_heads = config.num_attention_heads119        self.use_flash_attn = config.use_flash_attn and has_flash_attn120        if config.use_flash_attn and not has_flash_attn:121            print('Warning: Flash Attention is not available, use_flash_attn is set to False.')122        self.head_dim = self.embed_dim // self.num_heads123        if self.head_dim * self.num_heads != self.embed_dim:124            raise ValueError(125                f'embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:'126                f' {self.num_heads}).'127            )128 129        self.scale = self.head_dim ** -0.5130        self.qkv = nn.Linear(self.embed_dim, 3 * self.embed_dim, bias=config.qkv_bias)131        self.attn_drop = nn.Dropout(config.attention_dropout)132        self.proj_drop = nn.Dropout(config.dropout)133 134        self.qk_normalization = config.qk_normalization135 136        if self.qk_normalization:137            self.q_norm = InternRMSNorm(self.embed_dim, eps=config.layer_norm_eps)138            self.k_norm = InternRMSNorm(self.embed_dim, eps=config.layer_norm_eps)139 140        if self.use_flash_attn:141            self.inner_attn = FlashAttention(attention_dropout=config.attention_dropout)142        self.proj = nn.Linear(self.embed_dim, self.embed_dim)143 144    def _naive_attn(self, x):145        B, N, C = x.shape146        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)147        q, k, v = qkv.unbind(0)  # make torchscript happy (cannot use tensor as tuple)148 149        if self.qk_normalization:150            B_, H_, N_, D_ = q.shape151            q = self.q_norm(q.transpose(1, 2).flatten(-2, -1)).view(B_, N_, H_, D_).transpose(1, 2)152            k = self.k_norm(k.transpose(1, 2).flatten(-2, -1)).view(B_, N_, H_, D_).transpose(1, 2)153 154        attn = ((q * self.scale) @ k.transpose(-2, -1))155        attn = attn.softmax(dim=-1)156        attn = self.attn_drop(attn)157 158        x = (attn @ v).transpose(1, 2).reshape(B, N, C)159        x = self.proj(x)160        x = self.proj_drop(x)161        return x162 163    def _flash_attn(self, x, key_padding_mask=None, need_weights=False):164        qkv = self.qkv(x)165        qkv = rearrange(qkv, 'b s (three h d) -> b s three h d', three=3, h=self.num_heads)166 167        if self.qk_normalization:168            q, k, v = qkv.unbind(2)169            q = self.q_norm(q.flatten(-2, -1)).view(q.shape)170            k = self.k_norm(k.flatten(-2, -1)).view(k.shape)171            qkv = torch.stack([q, k, v], dim=2)172 173        context, _ = self.inner_attn(174            qkv, key_padding_mask=key_padding_mask, need_weights=need_weights, causal=False175        )176        outs = self.proj(rearrange(context, 'b s h d -> b s (h d)'))177        outs = self.proj_drop(outs)178        return outs179 180    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:181        x = self._naive_attn(hidden_states) if not self.use_flash_attn else self._flash_attn(hidden_states)182        return x183 184 185class InternMLP(nn.Module):186    def __init__(self, config: InternVisionConfig):187        super().__init__()188        self.config = config189        self.act = ACT2FN[config.hidden_act]190        self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)191        self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)192 193    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:194        hidden_states = self.fc1(hidden_states)195        hidden_states = self.act(hidden_states)196        hidden_states = self.fc2(hidden_states)197        return hidden_states198 199 200class InternVisionEncoderLayer(nn.Module):201    def __init__(self, config: InternVisionConfig, drop_path_rate: float):202        super().__init__()203        self.embed_dim = config.hidden_size204        self.intermediate_size = config.intermediate_size205        self.norm_type = config.norm_type206 207        self.attn = InternAttention(config)208        self.mlp = InternMLP(config)209        self.norm1 = NORM2FN[self.norm_type](self.embed_dim, eps=config.layer_norm_eps)210        self.norm2 = NORM2FN[self.norm_type](self.embed_dim, eps=config.layer_norm_eps)211 212        self.ls1 = nn.Parameter(config.initializer_factor * torch.ones(self.embed_dim))213        self.ls2 = nn.Parameter(config.initializer_factor * torch.ones(self.embed_dim))214        self.drop_path1 = DropPath(drop_path_rate) if drop_path_rate > 0. else nn.Identity()215        self.drop_path2 = DropPath(drop_path_rate) if drop_path_rate > 0. else nn.Identity()216 217    def forward(218            self,219            hidden_states: torch.Tensor,220    ) -> Tuple[torch.FloatTensor, Optional[torch.FloatTensor], Optional[Tuple[torch.FloatTensor]]]:221        """222        Args:223            hidden_states (`Tuple[torch.FloatTensor, Optional[torch.FloatTensor]]`): input to the layer of shape `(batch, seq_len, embed_dim)`224        """225        hidden_states = hidden_states + self.drop_path1(self.attn(self.norm1(hidden_states)) * self.ls1)226 227        hidden_states = hidden_states + self.drop_path2(self.mlp(self.norm2(hidden_states)) * self.ls2)228 229        return hidden_states230 231 232class InternVisionEncoder(nn.Module):233    """234    Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a235    [`InternEncoderLayer`].236 237    Args:238        config (`InternConfig`):239            The corresponding vision configuration for the `InternEncoder`.240    """241 242    def __init__(self, config: InternVisionConfig):243        super().__init__()244        self.config = config245        # stochastic depth decay rule246        dpr = [x.item() for x in torch.linspace(0, config.drop_path_rate, config.num_hidden_layers)]247        self.layers = nn.ModuleList([248            InternVisionEncoderLayer(config, dpr[idx]) for idx in range(config.num_hidden_layers)])249        self.gradient_checkpointing = True250 251    def forward(252            self,253            inputs_embeds,254            output_hidden_states: Optional[bool] = None,255            return_dict: Optional[bool] = None,256    ) -> Union[Tuple, BaseModelOutput]:257        r"""258        Args:259            inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):260                Embedded representation of the inputs. Should be float, not int tokens.261            output_hidden_states (`bool`, *optional*):262                Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors263                for more detail.264            return_dict (`bool`, *optional*):265                Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.266        """267        output_hidden_states = (268            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states269        )270        return_dict = return_dict if return_dict is not None else self.config.use_return_dict271 272        encoder_states = () if output_hidden_states else None273        hidden_states = inputs_embeds274 275        for idx, encoder_layer in enumerate(self.layers):276            if output_hidden_states:277                encoder_states = encoder_states + (hidden_states,)278            if self.gradient_checkpointing and self.training:279                layer_outputs = torch.utils.checkpoint.checkpoint(280                    encoder_layer,281                    hidden_states)282            else:283                layer_outputs = encoder_layer(284                    hidden_states,285                )286            hidden_states = layer_outputs287 288        if output_hidden_states:289            encoder_states = encoder_states + (hidden_states,)290 291        if not return_dict:292            return tuple(v for v in [hidden_states, encoder_states] if v is not None)293        return BaseModelOutput(294            last_hidden_state=hidden_states, hidden_states=encoder_states295        )296 297 298class InternVisionModel(PreTrainedModel):299    main_input_name = 'pixel_values'300    _supports_flash_attn_2 = True301    config_class = InternVisionConfig302    _no_split_modules = ['InternVisionEncoderLayer']303 304    def __init__(self, config: InternVisionConfig):305        super().__init__(config)306        self.config = config307 308        self.embeddings = InternVisionEmbeddings(config)309        self.encoder = InternVisionEncoder(config)310 311    def resize_pos_embeddings(self, old_size, new_size, patch_size):312        pos_emb = self.embeddings.position_embedding313        _, num_positions, embed_dim = pos_emb.shape314        cls_emb = pos_emb[:, :1, :]315        pos_emb = pos_emb[:, 1:, :].reshape(1, old_size // patch_size, old_size // patch_size, -1).permute(0, 3, 1, 2)316        pos_emb = F.interpolate(pos_emb.float(), size=new_size // patch_size, mode='bicubic', align_corners=False)317        pos_emb = pos_emb.to(cls_emb.dtype).reshape(1, embed_dim, -1).permute(0, 2, 1)318        pos_emb = torch.cat([cls_emb, pos_emb], dim=1)319        self.embeddings.position_embedding = nn.Parameter(pos_emb)320        self.embeddings.image_size = new_size321        logger.info('Resized position embeddings from {} to {}'.format(old_size, new_size))322 323    def get_input_embeddings(self):324        return self.embeddings325 326    def forward(327            self,328            pixel_values: Optional[torch.FloatTensor] = None,329            output_hidden_states: Optional[bool] = None,330            return_dict: Optional[bool] = None,331            pixel_embeds: Optional[torch.FloatTensor] = None,332    ) -> Union[Tuple, BaseModelOutputWithPooling]:333        output_hidden_states = (334            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states335        )336        return_dict = return_dict if return_dict is not None else self.config.use_return_dict337 338        if pixel_values is None and pixel_embeds is None:339            raise ValueError('You have to specify pixel_values or pixel_embeds')340 341        if pixel_embeds is not None:342            hidden_states = pixel_embeds343        else:344            if len(pixel_values.shape) == 4:345                hidden_states = self.embeddings(pixel_values)346            else:347                raise ValueError(f'wrong pixel_values size: {pixel_values.shape}')348        encoder_outputs = self.encoder(349            inputs_embeds=hidden_states,350            output_hidden_states=output_hidden_states,351            return_dict=return_dict,352        )353        last_hidden_state = encoder_outputs.last_hidden_state354        pooled_output = last_hidden_state[:, 0, :]355 356        if not return_dict:357            return (last_hidden_state, pooled_output) + encoder_outputs[1:]358 359        return BaseModelOutputWithPooling(360            last_hidden_state=last_hidden_state,361            pooler_output=pooled_output,362            hidden_states=encoder_outputs.hidden_states,363            attentions=encoder_outputs.attentions,364        )365