from __future__ import annotations from transformers.configuration_utils import PretrainedConfig from transformers.utils import ( ModelOutput, ) from transformers.activations import ACT2FN class ViTMAEConfig(PretrainedConfig): r""" This is the configuration class to store the configuration of a [`ViTMAEModel`]. It is used to instantiate an ViT MAE model according to the specified arguments, defining the model architecture. Instantiating a configuration with the defaults will yield a similar configuration to that of the ViT [facebook/vit-mae-base](https://huggingface.co/facebook/vit-mae-base) architecture. Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the documentation from [`PretrainedConfig`] for more information. Args: hidden_size (`int`, *optional*, defaults to 768): Dimensionality of the encoder layers and the pooler layer. num_hidden_layers (`int`, *optional*, defaults to 12): Number of hidden layers in the Transformer encoder. num_attention_heads (`int`, *optional*, defaults to 12): Number of attention heads for each attention layer in the Transformer encoder. intermediate_size (`int`, *optional*, defaults to 3072): Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder. hidden_act (`str` or `function`, *optional*, defaults to `"gelu"`): The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`, `"relu"`, `"selu"` and `"gelu_new"` are supported. hidden_dropout_prob (`float`, *optional*, defaults to 0.0): The dropout probabilitiy for all fully connected layers in the embeddings, encoder, and pooler. attention_probs_dropout_prob (`float`, *optional*, defaults to 0.0): The dropout ratio for the attention probabilities. initializer_range (`float`, *optional*, defaults to 0.02): The standard deviation of the truncated_normal_initializer for initializing all weight matrices. layer_norm_eps (`float`, *optional*, defaults to 1e-12): The epsilon used by the layer normalization layers. image_size (`int`, *optional*, defaults to 224): The size (resolution) of each image. patch_size (`int`, *optional*, defaults to 16): The size (resolution) of each patch. num_channels (`int`, *optional*, defaults to 3): The number of input channels. qkv_bias (`bool`, *optional*, defaults to `True`): Whether to add a bias to the queries, keys and values. decoder_num_attention_heads (`int`, *optional*, defaults to 16): Number of attention heads for each attention layer in the decoder. decoder_hidden_size (`int`, *optional*, defaults to 512): Dimensionality of the decoder. decoder_num_hidden_layers (`int`, *optional*, defaults to 8): Number of hidden layers in the decoder. decoder_intermediate_size (`int`, *optional*, defaults to 2048): Dimensionality of the "intermediate" (i.e., feed-forward) layer in the decoder. mask_ratio (`float`, *optional*, defaults to 0.75): The ratio of the number of masked tokens in the input sequence. norm_pix_loss (`bool`, *optional*, defaults to `False`): Whether or not to train with normalized pixels (see Table 3 in the paper). Using normalized pixels improved representation quality in the experiments of the authors. Example: ```python >>> from transformers import ViTMAEConfig, ViTMAEModel >>> # Initializing a ViT MAE vit-mae-base style configuration >>> configuration = ViTMAEConfig() >>> # Initializing a model (with random weights) from the vit-mae-base style configuration >>> model = ViTMAEModel(configuration) >>> # Accessing the model configuration >>> configuration = model.config ```""" model_type = "vit_mae" def __init__( self, hidden_size=768, num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072, hidden_act="gelu", hidden_dropout_prob=0.0, attention_probs_dropout_prob=0.0, initializer_range=0.02, layer_norm_eps=1e-12, image_size=224, patch_size=16, num_channels=3, qkv_bias=True, decoder_num_attention_heads=16, decoder_hidden_size=512, decoder_num_hidden_layers=8, decoder_intermediate_size=2048, mask_ratio=0.75, norm_pix_loss=False, **kwargs, ): super().__init__(**kwargs) self.hidden_size = hidden_size self.num_hidden_layers = num_hidden_layers self.num_attention_heads = num_attention_heads self.intermediate_size = intermediate_size self.hidden_act = hidden_act self.hidden_dropout_prob = hidden_dropout_prob self.attention_probs_dropout_prob = attention_probs_dropout_prob self.initializer_range = initializer_range self.layer_norm_eps = layer_norm_eps self.image_size = image_size self.patch_size = patch_size self.num_channels = num_channels self.qkv_bias = qkv_bias self.decoder_num_attention_heads = decoder_num_attention_heads self.decoder_hidden_size = decoder_hidden_size self.decoder_num_hidden_layers = decoder_num_hidden_layers self.decoder_intermediate_size = decoder_intermediate_size self.mask_ratio = mask_ratio self.norm_pix_loss = norm_pix_loss # coding=utf-8 # Copyright 2022 Facebook AI and The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # 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. """PyTorch ViT MAE (masked autoencoder) model.""" import collections.abc import math from copy import deepcopy from dataclasses import dataclass from typing import Optional, Set, Tuple, Union import numpy as np import torch from torch import nn # correct the above import to the following from transformers.modeling_outputs import BaseModelOutput try: from flash_attn import flash_attn_func HAS_FLASH_ATTN = True # if set to False, FlashAttention is unavailable except: print('FlashAttention is not installed.') HAS_FLASH_ATTN = False @dataclass class ViTMAEModelOutput(ModelOutput): """ Class for ViTMAEModel's outputs, with potential hidden states and attentions. Args: last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): Sequence of hidden-states at the output of the last layer of the model. mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`): Tensor indicating which patches are masked (1) and which are not (0). ids_restore (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Tensor containing the original index of the (shuffled) masked patches. hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the model at the output of each layer plus the initial embedding outputs. attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads. """ last_hidden_state: torch.FloatTensor = None mask: torch.LongTensor = None ids_restore: torch.LongTensor = None hidden_states: Optional[Tuple[torch.FloatTensor]] = None attentions: Optional[Tuple[torch.FloatTensor]] = None @dataclass class ViTMAEDecoderOutput(ModelOutput): """ Class for ViTMAEDecoder's outputs, with potential hidden states and attentions. Args: logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, patch_size ** 2 * num_channels)`): Pixel reconstruction logits. hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the model at the output of each layer plus the initial embedding outputs. attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads. """ logits: torch.FloatTensor = None hidden_states: Optional[Tuple[torch.FloatTensor]] = None attentions: Optional[Tuple[torch.FloatTensor]] = None @dataclass class ViTMAEForPreTrainingOutput(ModelOutput): """ Class for ViTMAEForPreTraining's outputs, with potential hidden states and attentions. Args: loss (`torch.FloatTensor` of shape `(1,)`): Pixel reconstruction loss. logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, patch_size ** 2 * num_channels)`): Pixel reconstruction logits. mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`): Tensor indicating which patches are masked (1) and which are not (0). ids_restore (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Tensor containing the original index of the (shuffled) masked patches. hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the model at the output of each layer plus the initial embedding outputs. attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads. """ loss: Optional[torch.FloatTensor] = None logits: torch.FloatTensor = None mask: torch.LongTensor = None ids_restore: torch.LongTensor = None hidden_states: Optional[Tuple[torch.FloatTensor]] = None attentions: Optional[Tuple[torch.FloatTensor]] = None def get_2d_sincos_pos_embed(embed_dim, grid_size, add_cls_token=False): """ Create 2D sin/cos positional embeddings. Args: embed_dim (`int`): Embedding dimension. grid_size (`int`): The grid height and width. add_cls_token (`bool`, *optional*, defaults to `False`): Whether or not to add a classification (CLS) token. Returns: (`torch.FloatTensor` of shape (grid_size*grid_size, embed_dim) or (1+grid_size*grid_size, embed_dim): the position embeddings (with or without classification token) """ grid_h = np.arange(grid_size, dtype=np.float32) grid_w = np.arange(grid_size, dtype=np.float32) grid = np.meshgrid(grid_w, grid_h) # here w goes first grid = np.stack(grid, axis=0) grid = grid.reshape([2, 1, grid_size, grid_size]) pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid) if add_cls_token: pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0) return pos_embed def get_2d_sincos_pos_embed_from_grid(embed_dim, grid): if embed_dim % 2 != 0: raise ValueError("embed_dim must be even") # use half of dimensions to encode grid_h emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0]) # (H*W, D/2) emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1]) # (H*W, D/2) emb = np.concatenate([emb_h, emb_w], axis=1) # (H*W, D) return emb def get_1d_sincos_pos_embed_from_grid(embed_dim, pos): """ embed_dim: output dimension for each position pos: a list of positions to be encoded: size (M,) out: (M, D) """ if embed_dim % 2 != 0: raise ValueError("embed_dim must be even") omega = np.arange(embed_dim // 2, dtype=float) omega /= embed_dim / 2.0 omega = 1.0 / 10000**omega # (D/2,) pos = pos.reshape(-1) # (M,) out = np.einsum("m,d->md", pos, omega) # (M, D/2), outer product emb_sin = np.sin(out) # (M, D/2) emb_cos = np.cos(out) # (M, D/2) emb = np.concatenate([emb_sin, emb_cos], axis=1) # (M, D) return emb class ViTMAEEmbeddings(nn.Module): """ Construct the CLS token, position and patch embeddings. """ def __init__(self, config): super().__init__() self.cls_token = nn.Parameter(torch.zeros(1, 1, config.hidden_size)) self.patch_embeddings = ViTMAEPatchEmbeddings(config) self.num_patches = self.patch_embeddings.num_patches # fixed sin-cos embedding self.position_embeddings = nn.Parameter( torch.zeros(1, self.num_patches + 1, config.hidden_size), requires_grad=False ) self.config = config self.initialize_weights() def initialize_weights(self): # initialize (and freeze) position embeddings by sin-cos embedding pos_embed = get_2d_sincos_pos_embed( self.position_embeddings.shape[-1], int(self.patch_embeddings.num_patches**0.5), add_cls_token=True ) self.position_embeddings.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0)) # initialize patch_embeddings like nn.Linear (instead of nn.Conv2d) w = self.patch_embeddings.projection.weight.data torch.nn.init.xavier_uniform_(w.view([w.shape[0], -1])) # timm's trunc_normal_(std=.02) is effectively normal_(std=0.02) as cutoff is too big (2.) torch.nn.init.normal_(self.cls_token, std=self.config.initializer_range) def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: int) -> torch.Tensor: """ This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher resolution images. Source: https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174 """ num_patches = embeddings.shape[1] - 1 num_positions = self.position_embeddings.shape[1] - 1 if num_patches == num_positions and height == width: return self.position_embeddings class_pos_embed = self.position_embeddings[:, 0, :] patch_pos_embed = self.position_embeddings[:, 1:, :] dim = embeddings.shape[-1] h0 = height // self.config.patch_size w0 = width // self.config.patch_size # we add a small number to avoid floating point error in the interpolation # see discussion at https://github.com/facebookresearch/dino/issues/8 h0, w0 = h0 + 0.1, w0 + 0.1 patch_pos_embed = patch_pos_embed.reshape(1, int(math.sqrt(num_positions)), int(math.sqrt(num_positions)), dim) patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2) patch_pos_embed = nn.functional.interpolate( patch_pos_embed, scale_factor=(h0 / math.sqrt(num_positions), w0 / math.sqrt(num_positions)), mode="bicubic", align_corners=False, ) if int(h0) != patch_pos_embed.shape[-2] or int(w0) != patch_pos_embed.shape[-1]: raise ValueError("Width or height does not match with the interpolated position embeddings") patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim) return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1) def random_masking(self, sequence, noise=None): """ Perform per-sample random masking by per-sample shuffling. Per-sample shuffling is done by argsort random noise. Args: sequence (`torch.LongTensor` of shape `(batch_size, sequence_length, dim)`) noise (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*) which is mainly used for testing purposes to control randomness and maintain the reproducibility """ batch_size, seq_length, dim = sequence.shape len_keep = int(seq_length * (1 - self.config.mask_ratio)) if noise is None: noise = torch.rand(batch_size, seq_length, device=sequence.device) # noise in [0, 1] # sort noise for each sample ids_shuffle = torch.argsort(noise, dim=1).to(sequence.device) # ascend: small is keep, large is remove ids_restore = torch.argsort(ids_shuffle, dim=1).to(sequence.device) # keep the first subset ids_keep = ids_shuffle[:, :len_keep] sequence_unmasked = torch.gather(sequence, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, dim)) # generate the binary mask: 0 is keep, 1 is remove mask = torch.ones([batch_size, seq_length], device=sequence.device) mask[:, :len_keep] = 0 # unshuffle to get the binary mask mask = torch.gather(mask, dim=1, index=ids_restore) return sequence_unmasked, mask, ids_restore def forward(self, pixel_values, noise=None, interpolate_pos_encoding: bool = False): batch_size, num_channels, height, width = pixel_values.shape embeddings = self.patch_embeddings(pixel_values, interpolate_pos_encoding=interpolate_pos_encoding) if interpolate_pos_encoding: position_embeddings = self.interpolate_pos_encoding(embeddings, height, width) else: position_embeddings = self.position_embeddings # add position embeddings w/o cls token embeddings = embeddings + position_embeddings[:, 1:, :] # masking: length -> length * config.mask_ratio embeddings, mask, ids_restore = self.random_masking(embeddings, noise) # append cls token cls_token = self.cls_token + position_embeddings[:, :1, :] cls_tokens = cls_token.expand(embeddings.shape[0], -1, -1) embeddings = torch.cat((cls_tokens, embeddings), dim=1) return embeddings, mask, ids_restore class ViTMAEPatchEmbeddings(nn.Module): """ This class turns `pixel_values` of shape `(batch_size, num_channels, height, width)` into the initial `hidden_states` (patch embeddings) of shape `(batch_size, seq_length, hidden_size)` to be consumed by a Transformer. """ def __init__(self, config): super().__init__() image_size, patch_size = config.image_size, config.patch_size num_channels, hidden_size = config.num_channels, config.hidden_size image_size = image_size if isinstance(image_size, collections.abc.Iterable) else (image_size, image_size) patch_size = patch_size if isinstance(patch_size, collections.abc.Iterable) else (patch_size, patch_size) num_patches = (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0]) self.image_size = image_size self.patch_size = patch_size self.num_channels = num_channels self.num_patches = num_patches self.projection = nn.Conv2d(num_channels, hidden_size, kernel_size=patch_size, stride=patch_size) def forward(self, pixel_values, interpolate_pos_encoding: bool = False): batch_size, num_channels, height, width = pixel_values.shape if num_channels != self.num_channels: raise ValueError( "Make sure that the channel dimension of the pixel values match with the one set in the configuration." ) if not interpolate_pos_encoding and (height != self.image_size[0] or width != self.image_size[1]): raise ValueError( f"Input image size ({height}*{width}) doesn't match model ({self.image_size[0]}*{self.image_size[1]})." ) x = self.projection(pixel_values).flatten(2).transpose(1, 2) return x # Copied from transformers.models.vit.modeling_vit.ViTSelfAttention ViT->ViTMAE class ViTMAESelfAttention(nn.Module): def __init__(self, config: ViTMAEConfig) -> None: super().__init__() if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"): raise ValueError( f"The hidden size {config.hidden_size,} is not a multiple of the number of attention " f"heads {config.num_attention_heads}." ) self.has_flash_attn = HAS_FLASH_ATTN self.num_attention_heads = config.num_attention_heads self.attention_head_size = int(config.hidden_size / config.num_attention_heads) self.all_head_size = self.num_attention_heads * self.attention_head_size self.query = nn.Linear(config.hidden_size, self.all_head_size, bias=config.qkv_bias) self.key = nn.Linear(config.hidden_size, self.all_head_size, bias=config.qkv_bias) self.value = nn.Linear(config.hidden_size, self.all_head_size, bias=config.qkv_bias) self.dropout = nn.Dropout(config.attention_probs_dropout_prob) self.attention_probs_dropout_prob = config.attention_probs_dropout_prob def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor: new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size) x = x.view(new_x_shape) return x.permute(0, 2, 1, 3) def reshape_for_flash(self, x: torch.Tensor) -> torch.Tensor: """Reshape to (B, N, num_heads, head_dim) for flash_attn_func.""" new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size) return x.view(new_x_shape) def forward( self, hidden_states, head_mask: Optional[torch.Tensor] = None, output_attentions: bool = False ) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor]]: # Use FlashAttention when available and no special outputs are needed if self.has_flash_attn and not output_attentions and head_mask is None: query_layer = self.reshape_for_flash(self.query(hidden_states)) key_layer = self.reshape_for_flash(self.key(hidden_states)) value_layer = self.reshape_for_flash(self.value(hidden_states)) context_layer = flash_attn_func( query_layer, key_layer, value_layer, dropout_p=self.attention_probs_dropout_prob if self.training else 0.0, causal=False, deterministic=True ) context_layer = context_layer.reshape(hidden_states.size(0), -1, self.all_head_size) return (context_layer,) # Fallback: manual attention computation mixed_query_layer = self.query(hidden_states) key_layer = self.transpose_for_scores(self.key(hidden_states)) value_layer = self.transpose_for_scores(self.value(hidden_states)) query_layer = self.transpose_for_scores(mixed_query_layer) attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2)) attention_scores = attention_scores / math.sqrt(self.attention_head_size) attention_probs = nn.functional.softmax(attention_scores, dim=-1) attention_probs = self.dropout(attention_probs) if head_mask is not None: attention_probs = attention_probs * head_mask context_layer = torch.matmul(attention_probs, value_layer) context_layer = context_layer.permute(0, 2, 1, 3).contiguous() new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,) context_layer = context_layer.view(new_context_layer_shape) outputs = (context_layer, attention_probs) if output_attentions else (context_layer,) return outputs # Copied from transformers.models.vit.modeling_vit.ViTSdpaSelfAttention ViT->ViTMAE class ViTMAESdpaSelfAttention(ViTMAESelfAttention): def __init__(self, config: ViTMAEConfig) -> None: super().__init__(config) self.attention_probs_dropout_prob = config.attention_probs_dropout_prob def forward( self, hidden_states, head_mask: Optional[torch.Tensor] = None, output_attentions: bool = False ) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor]]: # Use FlashAttention when available and no special outputs are needed if self.has_flash_attn and not output_attentions and head_mask is None: query_layer = self.reshape_for_flash(self.query(hidden_states)) key_layer = self.reshape_for_flash(self.key(hidden_states)) value_layer = self.reshape_for_flash(self.value(hidden_states)) context_layer = flash_attn_func( query_layer, key_layer, value_layer, dropout_p=self.attention_probs_dropout_prob if self.training else 0.0, causal=False, deterministic=True ) context_layer = context_layer.reshape(hidden_states.size(0), -1, self.all_head_size) return context_layer, None # Fallback: PyTorch SDPA mixed_query_layer = self.query(hidden_states) key_layer = self.transpose_for_scores(self.key(hidden_states)) value_layer = self.transpose_for_scores(self.value(hidden_states)) query_layer = self.transpose_for_scores(mixed_query_layer) context_layer = torch.nn.functional.scaled_dot_product_attention( query_layer, key_layer, value_layer, head_mask, self.attention_probs_dropout_prob if self.training else 0.0, is_causal=False, scale=None, ) context_layer = context_layer.permute(0, 2, 1, 3).contiguous() new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,) context_layer = context_layer.view(new_context_layer_shape) return context_layer, None # Copied from transformers.models.vit.modeling_vit.ViTSelfOutput with ViT->ViTMAE class ViTMAESelfOutput(nn.Module): """ The residual connection is defined in ViTMAELayer instead of here (as is the case with other models), due to the layernorm applied before each block. """ def __init__(self, config: ViTMAEConfig) -> None: super().__init__() self.dense = nn.Linear(config.hidden_size, config.hidden_size) self.dropout = nn.Dropout(config.hidden_dropout_prob) def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor: hidden_states = self.dense(hidden_states) hidden_states = self.dropout(hidden_states) return hidden_states # Copied from transformers.models.vit.modeling_vit.ViTAttention with ViT->ViTMAE class ViTMAEAttention(nn.Module): def __init__(self, config: ViTMAEConfig) -> None: super().__init__() self.attention = ViTMAESelfAttention(config) self.output = ViTMAESelfOutput(config) def forward( self, hidden_states: torch.Tensor, head_mask: Optional[torch.Tensor] = None, output_attentions: bool = False, ) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor]]: self_outputs = self.attention(hidden_states, head_mask, output_attentions) attention_output = self.output(self_outputs[0], hidden_states) outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them return outputs # Copied from transformers.models.vit.modeling_vit.ViTIntermediate ViT->ViTMAE class ViTMAEIntermediate(nn.Module): def __init__(self, config: ViTMAEConfig) -> None: super().__init__() self.dense = nn.Linear(config.hidden_size, config.intermediate_size) if isinstance(config.hidden_act, str): self.intermediate_act_fn = ACT2FN[config.hidden_act] else: self.intermediate_act_fn = config.hidden_act def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states = self.dense(hidden_states) hidden_states = self.intermediate_act_fn(hidden_states) return hidden_states # Copied from transformers.models.vit.modeling_vit.ViTOutput ViT->ViTMAE class ViTMAEOutput(nn.Module): def __init__(self, config: ViTMAEConfig) -> None: super().__init__() self.dense = nn.Linear(config.intermediate_size, config.hidden_size) self.dropout = nn.Dropout(config.hidden_dropout_prob) def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor: hidden_states = self.dense(hidden_states) hidden_states = self.dropout(hidden_states) hidden_states = hidden_states + input_tensor return hidden_states # Copied from transformers.models.vit.modeling_vit.ViTLayer with ViT->ViTMAE,VIT->VITMAE class ViTMAELayer(nn.Module): """This corresponds to the Block class in the timm implementation.""" def __init__(self, config: ViTMAEConfig) -> None: super().__init__() self.chunk_size_feed_forward = config.chunk_size_feed_forward self.seq_len_dim = 1 self.attention = ViTMAEAttention(config) # no SPDA by default self.intermediate = ViTMAEIntermediate(config) self.output = ViTMAEOutput(config) self.layernorm_before = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) self.layernorm_after = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) def forward( self, hidden_states: torch.Tensor, head_mask: Optional[torch.Tensor] = None, output_attentions: bool = False, ) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor]]: self_attention_outputs = self.attention( self.layernorm_before(hidden_states), # in ViTMAE, layernorm is applied before self-attention head_mask, output_attentions=output_attentions, ) attention_output = self_attention_outputs[0] outputs = self_attention_outputs[1:] # add self attentions if we output attention weights # first residual connection hidden_states = attention_output + hidden_states # in ViTMAE, layernorm is also applied after self-attention layer_output = self.layernorm_after(hidden_states) layer_output = self.intermediate(layer_output) # second residual connection is done here layer_output = self.output(layer_output, hidden_states) outputs = (layer_output,) + outputs return outputs class GeneralDecoder(nn.Module): def __init__(self, config, num_patches): super().__init__() self.decoder_embed = nn.Linear(config.hidden_size, config.decoder_hidden_size, bias=True) self.decoder_pos_embed = nn.Parameter( torch.zeros(1, num_patches + 1, config.decoder_hidden_size), requires_grad=False ) # fixed sin-cos embedding decoder_config = deepcopy(config) decoder_config.hidden_size = config.decoder_hidden_size decoder_config.num_hidden_layers = config.decoder_num_hidden_layers decoder_config.num_attention_heads = config.decoder_num_attention_heads decoder_config.intermediate_size = config.decoder_intermediate_size self.decoder_layers = nn.ModuleList( [ViTMAELayer(decoder_config) for _ in range(config.decoder_num_hidden_layers)] ) self.decoder_norm = nn.LayerNorm(config.decoder_hidden_size, eps=config.layer_norm_eps) self.decoder_pred = nn.Linear( config.decoder_hidden_size, config.patch_size**2 * config.num_channels, bias=True ) # encoder to decoder self.gradient_checkpointing = False self.config = config self.num_patches = num_patches self.initialize_weights(num_patches) self.decoder_config = decoder_config self.set_trainable_cls_token() def set_trainable_cls_token(self, tensor: Optional[torch.Tensor] = None): # register a trainable CLS token tensor = torch.zeros(1, 1, self.decoder_config.hidden_size) if tensor is None else tensor self.trainable_cls_token = nn.Parameter(tensor) def interpolate_pos_encoding(self, embeddings: torch.Tensor) -> torch.Tensor: """ This method is a modified version of the interpolation function for ViT-mae model at the deocder, that allows to interpolate the pre-trained decoder position encodings, to be able to use the model on higher resolution images. Source: https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174 """ # -1 removes the class dimension since we later append it without interpolation embeddings_positions = embeddings.shape[1] - 1 num_positions = self.decoder_pos_embed.shape[1] - 1 # Separation of class token and patch tokens class_pos_embed = self.decoder_pos_embed[:, 0, :] patch_pos_embed = self.decoder_pos_embed[:, 1:, :] # To retain the final 3d tensor with the required dimensions dim = self.decoder_pos_embed.shape[-1] # Increasing a dimension to enable bicubic interpolation patch_pos_embed = patch_pos_embed.reshape(1, 1, -1, dim) # permute to bring the dimension to be interpolated, to the last patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2) # Interpolating the decoder position embeddings shape wrt embeddings shape i.e (x). # 1 keeps the other dimension constant patch_pos_embed = nn.functional.interpolate( patch_pos_embed, scale_factor=(1, embeddings_positions / num_positions), mode="bicubic", align_corners=False, ) # Converting back to the original shape patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim) # Adding the class token back return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1) def interpolate_latent(self, x: torch.Tensor) -> torch.Tensor: b, l, c = x.shape if l == self.num_patches: return x # interpolate the latent #print(f"interpolating latent from {l} to {self.num_patches}, x.shape = {x.shape}") h, w = int(l**0.5), int(l**0.5) x = x.reshape(b, h, w, c) x = x.permute(0, 3, 1, 2) target_size = (int(self.num_patches**0.5), int(self.num_patches**0.5)) x = nn.functional.interpolate(x, size=target_size, mode="bilinear", align_corners=False) x = x.permute(0, 2, 3, 1).contiguous().view(b, self.num_patches, c) return x def initialize_weights(self, num_patches): # initialize (and freeze) position embeddings by sin-cos embedding decoder_pos_embed = get_2d_sincos_pos_embed( self.decoder_pos_embed.shape[-1], int(num_patches**0.5), add_cls_token=True ) self.decoder_pos_embed.data.copy_(torch.from_numpy(decoder_pos_embed).float().unsqueeze(0)) # timm's trunc_normal_(std=.02) is effectively normal_(std=0.02) as cutoff is too big (2.) # torch.nn.init.normal_(self.mask_token, std=self.config.initializer_range) def unpatchify(self, patchified_pixel_values, original_image_size: Optional[Tuple[int, int]] = None): """ Args: patchified_pixel_values (`torch.FloatTensor` of shape `(batch_size, num_patches, patch_size**2 * num_channels)`: Patchified pixel values. original_image_size (`Tuple[int, int]`, *optional*): Original image size. Returns: `torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`: Pixel values. """ patch_size, num_channels = self.config.patch_size, self.config.num_channels original_image_size = ( original_image_size if original_image_size is not None else (self.config.image_size, self.config.image_size) ) original_height, original_width = original_image_size num_patches_h = original_height // patch_size num_patches_w = original_width // patch_size # sanity check if num_patches_h * num_patches_w != patchified_pixel_values.shape[1]: raise ValueError( f"The number of patches in the patchified pixel values {patchified_pixel_values.shape[1]}, does not match the number of patches on original image {num_patches_h}*{num_patches_w}" ) # unpatchify batch_size = patchified_pixel_values.shape[0] patchified_pixel_values = patchified_pixel_values.reshape( batch_size, num_patches_h, num_patches_w, patch_size, patch_size, num_channels, ) patchified_pixel_values = torch.einsum("nhwpqc->nchpwq", patchified_pixel_values) pixel_values = patchified_pixel_values.reshape( batch_size, num_channels, num_patches_h * patch_size, num_patches_w * patch_size, ) return pixel_values def forward( self, hidden_states, output_attentions=False, output_hidden_states=False, return_dict=True, interpolate_pos_encoding: bool = False, drop_cls_token: bool = False, ): # embed tokens x = self.decoder_embed(hidden_states) #print(f"x.shape = {x.shape}") if drop_cls_token: x_ = x[:, 1:, :] # no cls token x_ = self.interpolate_latent(x_) else: x_ = self.interpolate_latent(x) # interpolate the whole latent cls_token = self.trainable_cls_token.expand(x_.shape[0], -1, -1) x = torch.cat([cls_token, x_], dim=1) # add pos embed if interpolate_pos_encoding: assert drop_cls_token, "interpolate_pos_encoding only works with drop_cls_token=True" decoder_pos_embed = self.interpolate_pos_encoding(x) else: decoder_pos_embed = self.decoder_pos_embed hidden_states = x + decoder_pos_embed #print(f"hidden_states.shape = {hidden_states.shape}") # apply Transformer layers (blocks) all_hidden_states = () if output_hidden_states else None all_self_attentions = () if output_attentions else None for i, layer_module in enumerate(self.decoder_layers): if output_hidden_states: all_hidden_states = all_hidden_states + (hidden_states,) if self.gradient_checkpointing and self.training: layer_outputs = self._gradient_checkpointing_func( layer_module.__call__, hidden_states, None, output_attentions, ) else: layer_outputs = layer_module(hidden_states, head_mask=None, output_attentions=output_attentions) hidden_states = layer_outputs[0] if output_attentions: all_self_attentions = all_self_attentions + (layer_outputs[1],) if output_hidden_states: all_hidden_states = all_hidden_states + (hidden_states,) hidden_states = self.decoder_norm(hidden_states) # predictor projection logits = self.decoder_pred(hidden_states) # remove cls token logits = logits[:, 1:, :] if not return_dict: return tuple(v for v in [logits, all_hidden_states, all_self_attentions] if v is not None) return ViTMAEDecoderOutput( logits=logits, hidden_states=all_hidden_states, attentions=all_self_attentions, ) import json from math import sqrt from pathlib import Path from typing import Any, Dict, Optional import torch import torch.nn as nn from timm.models.vision_transformer import Attention from transformers import AutoConfig, AutoImageProcessor try: from diffusers.configuration_utils import ConfigMixin, register_to_config from diffusers.models.modeling_utils import ModelMixin except Exception: # pragma: no cover class ConfigMixin: config_name = "config.json" class ModelMixin(nn.Module): pass def register_to_config(init): return init _PIXEL_DECODER_DEFAULTS: Dict[str, Any] = { "decoder_hidden_size": 1024, "decoder_intermediate_size": 4096, "decoder_num_attention_heads": 16, "decoder_num_hidden_layers": 24, "layer_norm_eps": 1e-12, "num_channels": 3, "hidden_act": "gelu", "qkv_bias": True, "hidden_dropout_prob": 0.0, "attention_probs_dropout_prob": 0.0, "initializer_range": 0.02, } def _resolve_local_asset_path(path: str, model_root: Optional[Path] = None) -> str: candidate = Path(path) if candidate.is_file(): return str(candidate) search_roots = [] if model_root is not None: search_roots.extend([model_root, model_root.parent]) for root in search_roots: for probe in (root / candidate, root / candidate.name, root / "decoder_config" / "config.json"): if probe.is_file(): return str(probe) return path def _build_pixel_decoder_config( *, encoder_hidden_size: int, decoder_patch_size: int, base_patches: int, decoder_config_path: Optional[str] = None, overrides: Optional[Dict[str, Any]] = None, ) -> ViTMAEConfig: pixel_decoder = dict(_PIXEL_DECODER_DEFAULTS) if decoder_config_path is not None: resolved = _resolve_local_asset_path(decoder_config_path) if Path(resolved).is_file(): with open(resolved, encoding="utf-8") as handle: pixel_decoder.update(json.load(handle)) if overrides: pixel_decoder.update({key: value for key, value in overrides.items() if value is not None}) return ViTMAEConfig( hidden_size=encoder_hidden_size, patch_size=decoder_patch_size, image_size=int(decoder_patch_size * sqrt(base_patches)), **pixel_decoder, ) def _load_image_normalization(encoder_config_path: str) -> tuple[list[float], list[float]]: try: proc = AutoImageProcessor.from_pretrained(encoder_config_path, use_fast=True, local_files_only=True) return proc.image_mean, proc.image_std except Exception: pass try: proc = AutoImageProcessor.from_pretrained(encoder_config_path, use_fast=True, local_files_only=False) return proc.image_mean, proc.image_std except Exception: if "siglip" in encoder_config_path.lower(): return [0.5, 0.5, 0.5], [0.5, 0.5, 0.5] return [0.485, 0.456, 0.406], [0.229, 0.224, 0.225] class PAEDecoder(ModelMixin, ConfigMixin): """PAE decoder: latent decompressor + ViT-MAE pixel decoder.""" @register_to_config def __init__( self, encoder_hidden_size: int = 1024, encoder_num_heads: int = 16, encoder_input_size: int = 224, encoder_patch_size: int = 14, encoder_config_path: str = "facebook/dinov2-with-registers-large", latent_dim: int = 32, decoder_patch_size: int = 16, decoder_config_path: Optional[str] = None, decoder_hidden_size: int = 1024, decoder_intermediate_size: int = 4096, decoder_num_attention_heads: int = 16, decoder_num_hidden_layers: int = 24, layer_norm_eps: float = 1e-12, num_channels: int = 3, hidden_act: str = "gelu", qkv_bias: bool = True, hidden_dropout_prob: float = 0.0, attention_probs_dropout_prob: float = 0.0, initializer_range: float = 0.02, image_mean: Optional[list[float]] = None, image_std: Optional[list[float]] = None, latent_mean: Optional[list[float]] = None, latent_std: Optional[list[float]] = None, latent_multiplier: float = 1.0, **kwargs, ) -> None: super().__init__() if image_mean is not None and image_std is not None: norm_mean, norm_std = image_mean, image_std else: norm_mean, norm_std = _load_image_normalization(encoder_config_path) self.register_buffer("encoder_mean", torch.tensor(norm_mean).view(1, 3, 1, 1), persistent=True) self.register_buffer("encoder_std", torch.tensor(norm_std).view(1, 3, 1, 1), persistent=True) self.encoder_input_size = encoder_input_size self.encoder_patch_size = encoder_patch_size self.base_patches = (self.encoder_input_size // self.encoder_patch_size) ** 2 self.latent_dim = latent_dim self.latent_decompressor = nn.ModuleList( [ nn.Conv2d(latent_dim, encoder_hidden_size, kernel_size=3, padding=1), Attention(dim=encoder_hidden_size, num_heads=encoder_num_heads), ] ) pixel_decoder_config = _build_pixel_decoder_config( encoder_hidden_size=encoder_hidden_size, decoder_patch_size=decoder_patch_size, base_patches=self.base_patches, decoder_config_path=decoder_config_path, overrides={ "decoder_hidden_size": decoder_hidden_size, "decoder_intermediate_size": decoder_intermediate_size, "decoder_num_attention_heads": decoder_num_attention_heads, "decoder_num_hidden_layers": decoder_num_hidden_layers, "layer_norm_eps": layer_norm_eps, "num_channels": num_channels, "hidden_act": hidden_act, "qkv_bias": qkv_bias, "hidden_dropout_prob": hidden_dropout_prob, "attention_probs_dropout_prob": attention_probs_dropout_prob, "initializer_range": initializer_range, }, ) self.decoder = GeneralDecoder(pixel_decoder_config, num_patches=self.base_patches) @classmethod def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs): model_root = Path(pretrained_model_name_or_path) config_file = model_root / cls.config_name if config_file.is_file(): with open(config_file, encoding="utf-8") as handle: config = json.load(handle) decoder_config_path = config.get("decoder_config_path") if decoder_config_path is not None: resolved = _resolve_local_asset_path(decoder_config_path, model_root=model_root) if resolved != decoder_config_path: kwargs.setdefault("decoder_config_path", resolved) return super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs) def decode(self, latents: torch.Tensor) -> torch.Tensor: hidden_states = self.latent_decompressor[0](latents) batch_size, channels, height, width = hidden_states.shape num_tokens = height * width hidden_states = hidden_states.view(batch_size, channels, num_tokens).transpose(1, 2) hidden_states = self.latent_decompressor[1](hidden_states) output = self.decoder(hidden_states, drop_cls_token=False).logits x_rec = self.decoder.unpatchify(output) x_rec = x_rec * self.encoder_std.to(x_rec.device) + self.encoder_mean.to(x_rec.device) return x_rec def forward(self, latents: torch.Tensor) -> torch.Tensor: return self.decode(latents)