""" NaFlex Vision Transformer An improved version of the Vision Transformer with: 1. Encapsulated embedding and position encoding in a single module 2. Support for linear patch embedding on pre-patchified inputs 3. Support for NaFlex variable aspect, variable resolution 4. Support for FlexiViT variable patch size 5. Support for NaViT fractional/factorized position embedding Based on ideas from: - Original Vision Transformer: https://arxiv.org/abs/2010.11929 - FlexiViT: https://arxiv.org/abs/2212.08013 - NaViT: https://arxiv.org/abs/2307.06304 - NaFlex (SigLip-2): https://arxiv.org/abs/2502.14786 Hacked together by / Copyright 2025, Ross Wightman, Hugging Face """ import logging import math from dataclasses import dataclass, fields, replace from functools import partial from typing import Callable, Dict, List, Optional, Set, Tuple, Type, Union, Any import torch import torch.nn as nn import torch.nn.functional as F from torch.nn.utils.rnn import pad_sequence from timm.data import IMAGENET_INCEPTION_MEAN, IMAGENET_INCEPTION_STD from timm.layers import ( AttentionPoolLatent, Mlp, to_2tuple, get_act_layer, get_norm_layer, LayerNorm, _assert, ) from timm.models._builder import build_model_with_cfg from timm.models._features import feature_take_indices from timm.models._features_fx import register_notrace_function, register_notrace_module from timm.models._registry import register_model, generate_default_cfgs from timm.models._manipulate import checkpoint, checkpoint_seq, named_apply from .vision_transformer import Block, global_pool_nlc __all__ = ['NaFlexVitCfg', 'NaFlexVit'] _logger = logging.getLogger(__name__) @dataclass class NaFlexVitCfg: """Configuration for FlexVit model. This dataclass contains the bulk of model configuration parameters, with core parameters (img_size, in_chans, num_classes, etc.) remaining as direct constructor arguments for API compatibility. """ # Architecture parameters patch_size: Union[int, Tuple[int, int]] = 16 embed_dim: int = 768 depth: int = 12 num_heads: int = 12 mlp_ratio: float = 4.0 # Attention parameters qkv_bias: bool = True qk_norm: bool = False proj_bias: bool = True attn_drop_rate: float = 0.0 # Regularization init_values: Optional[float] = None # Layer-scale init values (layer-scale enabled if not None) drop_rate: float = 0.0 # Dropout rate for classifier pos_drop_rate: float = 0.0 # Dropout rate for position embeddings patch_drop_rate: float = 0.0 # Dropout rate for patch tokens proj_drop_rate: float = 0.0 # Dropout rate for linear projections drop_path_rate: float = 0.0 # Stochastic depth drop rate # Prefix token configuration class_token: bool = False # Use class token reg_tokens: int = 0 # Number of register tokens # Position embedding configuration pos_embed: str = 'learned' # Type of position embedding ('learned', 'factorized', 'rope', 'none') pos_embed_grid_size: Optional[Tuple[int, int]] = (16, 16) # Grid size for position embedding initialization pos_embed_interp_mode: str = 'bicubic' # Interpolation mode for position embedding resizing pos_embed_ar_preserving: bool = False # Whether to preserve aspect ratio during position embedding interpolation pos_embed_use_grid_sample: bool = False # Whether to use grid_sample for naflex position embedding interpolation # Image processing dynamic_img_pad: bool = False # Whether to enable dynamic padding for variable resolution # Architecture choices pre_norm: bool = False # Whether to apply normalization before attention/MLP layers (start of blocks) final_norm: bool = True # Whether to apply final normalization before pooling and classifier (end of blocks) fc_norm: Optional[bool] = None # Whether to normalize features before final classifier (after pooling) global_pool: str = 'map' # Type of global pooling for final sequence pool_include_prefix: bool = False # Whether to include class/register prefix tokens in global pooling # Weight initialization weight_init: str = '' # Weight initialization scheme fix_init: bool = True # Apply weight initialization fix (scaling w/ layer index) # Embedding configuration embed_proj_type: str = 'linear' # Type of embedding layer ('conv' or 'linear') input_norm_layer: Optional[str] = None # Normalization layer for embeddings input (before input projection) embed_norm_layer: Optional[str] = None # Normalization layer for embeddings (after input projection) # Layer implementations norm_layer: Optional[str] = None # Normalization layer for transformer blocks act_layer: Optional[str] = None # Activation layer for MLP blocks block_fn: Optional[str] = None # Transformer block implementation class name mlp_layer: Optional[str] = None # MLP implementation class name # Variable patch size support enable_patch_interpolator: bool = False # Enable dynamic patch size support def _overlay_kwargs(cfg: NaFlexVitCfg, **kwargs) -> NaFlexVitCfg: """Overlay kwargs onto config, replacing config values with provided kwargs.""" # Only update fields that exist in the config config_fields = set(cfg.__dataclass_fields__.keys()) config_kwargs = {k: v for k, v in kwargs.items() if k in config_fields} if config_kwargs: cfg = replace(cfg, **config_kwargs) return cfg def batch_patchify( x: torch.Tensor, patch_size: Tuple[int, int], pad: bool = True, ) -> Tuple[torch.Tensor, Tuple[int, int]]: """Patchify a batch of images. Args: x: Input tensor of shape [B, C, H, W]. patch_size: Patch dimensions (patch_h, patch_w). pad: Whether to pad images to be divisible by patch size. Returns: Tuple of (patches, grid_size) where patches has shape [B, N, P*P*C] and grid_size is (num_patches_h, num_patches_w). """ B, C, H, W = x.shape ph, pw = patch_size # Ensure the image is divisible by patch size if pad and (H % ph != 0 or W % pw != 0): pad_h = (ph - H % ph) % ph pad_w = (pw - W % pw) % pw x = F.pad(x, (0, pad_w, 0, pad_h)) nh, nw = H // ph, W // pw patches = x.view(B, C, nh, ph, nw, pw).permute(0, 2, 4, 3, 5, 1).reshape(B, nh * nw, ph * pw * C) # FIXME confirm we want 'channels last' in the patch channel layout, egg ph, ph, C instead of C, ph, hw return patches, (nh, nw) def calculate_naflex_grid_sizes(_coord: torch.Tensor): # Calculate the appropriate grid size from coords max_y = _coord[:, :, 0].amax(dim=1) + 1 max_x = _coord[:, :, 1].amax(dim=1) + 1 return [(int(h.item()), int(w.item())) for h, w in zip(max_y, max_x)] @register_notrace_module class NaFlexEmbeds(nn.Module): """NaFlex Embedding module for Vision Transformers. This module encapsulates the complete embedding process for Vision Transformers, supporting both standard and NaFlex (NaViT + FlexiViT) functionality: 1. Patch embedding (via Conv2d or Linear) 2. Class and register token preparation 3. Position embedding addition with interpolation support 4. Pre-normalization (if requested) 5. Dropout application NaFlex capabilities include: - Variable aspect ratio and resolution via patch coordinates - Patch type indicators for handling padding tokens in attention - Flexible position embedding interpolation for arbitrary grid sizes - Support for factorized position embeddings The patch embedding can be one of two types: - Conv2d-based (default): For standard image inputs [B, C, H, W] - Linear-based: For pre-patchified inputs [B, N, P*P*C] Args: patch_size: Size of patches for patch embedding in_chans: Number of input image channels embed_dim: Dimensionality of patch embedding proj_type: Type of embedding projection layer ('conv' or 'linear') input_norm_layer: Normalization layer applied to input (linear mode only) proj_norm_layer: Normalization layer applied after projection pos_embed: Type of position embedding ('learned', 'factorized', 'rope', 'none') pos_drop_rate: Dropout rate for position embeddings patch_drop_rate: Dropout rate for patch tokens class_token: Whether to include a class token reg_tokens: Number of register tokens to include bias: Whether to use bias in projection layers dynamic_img_pad: Whether to enable dynamic padding for variable resolution pos_embed_grid_size: Grid size for position embedding initialization pos_embed_interp_mode: Interpolation mode for position embedding resizing pos_embed_ar_preserving: Whether to preserve aspect ratio during position embedding interpolation default_img_size: Default image size for position embedding grid calculation """ def __init__( self, patch_size: Union[int, Tuple[int, int]] = 16, in_chans: int = 3, embed_dim: int = 768, proj_type: Optional[str] = None, proj_bias: bool = True, class_token: bool = True, reg_tokens: int = 0, dynamic_img_pad: bool = False, default_img_size: Optional[Union[int, Tuple[int, int]]] = None, pos_embed: str = 'learned', pos_embed_grid_size: Optional[Tuple[int, int]] = (14, 14), pos_embed_interp_mode: str = 'bicubic', pos_embed_ar_preserving: bool = False, pos_embed_use_grid_sample: bool = False, input_norm_layer: Optional[Type[nn.Module]] = None, proj_norm_layer: Union[bool, Optional[Type[nn.Module]]] = None, norm_layer: Optional[Type[nn.Module]] = None, pos_drop_rate: float = 0., patch_drop_rate: float = 0., enable_patch_interpolator: bool = False, ) -> None: """Initialize NaFlexEmbeds module. Args: patch_size: Size of patches for patch embedding. in_chans: Number of input image channels. embed_dim: Dimensionality of patch embedding. proj_type: Type of embedding projection layer ('conv' or 'linear'). proj_bias: Whether to use bias in projection layers. class_token: Whether to include a class token. reg_tokens: Number of register tokens to include. dynamic_img_pad: Whether to enable dynamic padding for variable resolution. default_img_size: Default image size for position embedding grid calculation. pos_embed: Type of position embedding ('learned', 'factorized', 'rope', 'none'). pos_embed_grid_size: Grid size for position embedding initialization. pos_embed_interp_mode: Interpolation mode for position embedding resizing. pos_embed_ar_preserving: Whether to preserve aspect ratio during interpolation. input_norm_layer: Normalization layer applied to input (linear mode only). proj_norm_layer: Normalization layer applied after projection. norm_layer: Default normalization layer. pos_drop_rate: Dropout rate for position embeddings. patch_drop_rate: Dropout rate for patch tokens. enable_patch_interpolator: Enable dynamic patch size support. """ super().__init__() self.has_class_token = class_token self.num_reg_tokens = reg_tokens self.pos_embed_interp_mode = pos_embed_interp_mode self.pos_embed_ar_preserving = pos_embed_ar_preserving self.pos_embed_use_grid_sample = pos_embed_use_grid_sample self.patch_size = to_2tuple(patch_size) self.in_chans = in_chans self.embed_dim = embed_dim self.dynamic_img_pad = dynamic_img_pad self.enable_patch_interpolator = enable_patch_interpolator # Calculate number of prefix tokens self.num_prefix_tokens = 1 if class_token else 0 self.num_prefix_tokens += reg_tokens # Create class and register tokens self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) if class_token else None self.reg_token = nn.Parameter(torch.zeros(1, reg_tokens, embed_dim)) if reg_tokens else None # Calculate grid size and number of patches self.default_img_size: Optional[Tuple[int, int]] = None self.pos_embed_grid_size: Optional[Tuple[int, int]] = None # Grid size used for learned pos embed init if pos_embed_grid_size is not None: # Highest priority, use provided pos_embed_grid_size self.pos_embed_grid_size = pos_embed_grid_size elif default_img_size is not None: # Fallback to calculating grid size from img_size + patch_size if img size provided. self.default_img_size = to_2tuple(default_img_size) self.pos_embed_grid_size = tuple([s // p for s, p in zip(self.default_img_size, self.patch_size)]) # Determine patch embedding type (linear or conv2d) if proj_type == 'linear': # Create linear projection for pre-patchified inputs # Input dimension is patch_size^2 * in_chans patch_dim = self.patch_size[0] * self.patch_size[1] * in_chans assert not (input_norm_layer is True and norm_layer is None), \ "`norm_layer` must be given when input_norm_layer=True" input_norm_layer = norm_layer if input_norm_layer is True else (input_norm_layer or None) self.norm_input = input_norm_layer(patch_dim) if input_norm_layer else None self.proj = nn.Linear(patch_dim, embed_dim, bias=proj_bias) self.flatten = False self.is_linear = True else: # Default to convolutional patch embedding for image inputs assert not input_norm_layer self.norm_input = None self.proj = nn.Conv2d( in_chans, embed_dim, kernel_size=patch_size, stride=patch_size, bias=proj_bias ) self.flatten = True self.is_linear = False # Create patch embedding interpolator if enabled if self.enable_patch_interpolator: from timm.layers import PatchEmbedInterpolator self.patch_interpolator = PatchEmbedInterpolator( base_patch_size=self.patch_size, in_chans=in_chans, embed_dim=embed_dim, interpolation=pos_embed_interp_mode, antialias=True, ) else: self.patch_interpolator = None # Create normalization layer after the projection assert not (proj_norm_layer is True and norm_layer is None), \ "`norm_layer` must be given when proj_norm_layer=True" proj_norm_layer = norm_layer if proj_norm_layer is True else (proj_norm_layer or None) self.norm = proj_norm_layer(embed_dim) if proj_norm_layer else nn.Identity() # Create position embedding if needed - only for patches, never for prefix tokens if pos_embed in ('factorized', 'learned') and self.pos_embed_grid_size is None: raise ValueError( "Cannot initialize position embeddings without grid_size." "Please provide img_size or pos_embed_grid_size.") self.pos_embed: Optional[torch.Tensor] = None self.pos_embed_y: Optional[torch.Tensor] = None self.pos_embed_x: Optional[torch.Tensor] = None if not pos_embed or pos_embed == 'none': self.pos_embed_type = 'none' elif pos_embed == 'rope': self.pos_embed_type = 'rope' # Rotary embeddings will be computed on-the-fly in the forward pass elif pos_embed == 'factorized': assert self.pos_embed_grid_size is not None h, w = self.pos_embed_grid_size self.pos_embed_type = 'factorized' self.pos_embed_y = nn.Parameter(torch.randn(1, h, embed_dim) * .02) self.pos_embed_x = nn.Parameter(torch.randn(1, w, embed_dim) * .02) else: assert self.pos_embed_grid_size is not None h, w = self.pos_embed_grid_size self.pos_embed = nn.Parameter(torch.randn(1, h, w, embed_dim) * .02) self.pos_embed_type = 'learned' # Dropout layers self.pos_drop = nn.Dropout(p=pos_drop_rate) if patch_drop_rate > 0: from timm.layers.patch_dropout import PatchDropout self.patch_drop = PatchDropout( patch_drop_rate, num_prefix_tokens=self.num_prefix_tokens, ) else: self.patch_drop = nn.Identity() def feature_info(self, location) -> Dict[str, Any]: """Get feature information for feature extraction. Args: location: Feature extraction location identifier Returns: Dictionary containing feature channel count and reduction factor """ return dict(num_chs=self.embed_dim, reduction=self.patch_size) def feat_ratio(self, as_scalar: bool = True) -> Union[int, Tuple[int, int]]: """Get the feature reduction ratio (stride) of the patch embedding. Args: as_scalar: Whether to return the maximum dimension as a scalar Returns: Feature reduction ratio as scalar or tuple """ if as_scalar: return max(self.patch_size) else: return self.patch_size def dynamic_feat_size(self, img_size: Tuple[int, int]) -> Tuple[int, int]: """Calculate grid (feature) size for given image size. Takes into account dynamic padding when enabled. Args: img_size: Input image size as (height, width) Returns: Grid size as (grid_height, grid_width) """ if self.dynamic_img_pad: return math.ceil(img_size[0] / self.patch_size[0]), math.ceil(img_size[1] / self.patch_size[1]) else: return img_size[0] // self.patch_size[0], img_size[1] // self.patch_size[1] #@torch.compiler.disable() def _apply_learned_naflex_pos_embed( self, x: torch.Tensor, patch_coord: torch.Tensor, ) -> None: """Apply learned position embeddings to NaFlex batch in-place. Interpolates learned 2D position embeddings for each sample in the batch based on their individual grid sizes. Args: x: Input tensor to add position embeddings to [B, N, C] patch_coord: Patch coordinates [B, N, 2] with (y, x) values """ # Calculate grid sizes from patch coordinates naflex_grid_sizes = calculate_naflex_grid_sizes(patch_coord) orig_h, orig_w = self.pos_embed.shape[1:3] pos_embed_nchw = self.pos_embed.permute(0, 3, 1, 2).float() # B,C,H,W def _interp2d(size): """ Return a flattened positional-embedding grid at an arbitrary spatial resolution. Converts the learned 2-D table stored in NCHW format (pos_embed_nchw) into a (1, H*W, C) sequence that matches the requested size. """ if (size[0] == orig_h) and (size[1] == orig_w): pos_embed_flat = self.pos_embed.reshape(1, orig_h * orig_w, -1) else: _interp_size = to_2tuple(max(size)) if self.pos_embed_ar_preserving else size pos_embed_flat = F.interpolate( pos_embed_nchw, size=_interp_size, mode=self.pos_embed_interp_mode, align_corners=False, antialias=True, )[:, :, :size[0], :size[1]].flatten(2).transpose(1, 2) return pos_embed_flat.to(dtype=x.dtype) # Determine unique grid sizes to avoid duplicate interpolation size_to_indices: Dict[Tuple[int, int], List[int]] = {} for bi, k in enumerate(naflex_grid_sizes): # k = h << 16 | w # FIXME can get jit compat with this size_to_indices.setdefault(k, []).append(bi) for k, batch_indices in size_to_indices.items(): # h, w = k >> 16, k & 0xFFFF # FIXME can get jit compat with this # Interpolate only once for this (h, w) pos_embed_flat = _interp2d(k) seq_len = min(x.shape[1], pos_embed_flat.shape[1]) x[:, :seq_len].index_add_( 0, torch.as_tensor(batch_indices, device=x.device), pos_embed_flat[:, :seq_len].expand(len(batch_indices), -1, -1) ) def _apply_learned_naflex_pos_embed_grid_sample( self, x: torch.Tensor, patch_coord: torch.Tensor, ) -> None: """Apply learned position embeddings to NaFlex batch using grid_sample. Uses F.grid_sample for efficient interpolation of learned 2D position embeddings based on patch coordinates. Based on proposal by https://github.com/stas-sl Args: x: Input tensor to add position embeddings to [B, N, C] patch_coord: Patch coordinates [B, N, 2] with (y, x) values """ device = x.device B, N, C = x.shape shapes = patch_coord.max(dim=1).values + 1 # (B, 2) containing [h_i, w_i] if self.pos_embed_ar_preserving: L_i = shapes.amax(dim=1) # (B,) max(h_i, w_i) L_global = L_i.amax() grid_size_y = grid_size_x = L_global scale_x = scale_y = L_global / L_i # uniform zoom (B,) else: grid_size_y, grid_size_x = shapes.amax(dim=0) # (2,) scale_y = grid_size_y / shapes[:, 0] # vertical zoom (B,) scale_x = grid_size_x / shapes[:, 1] # horizontal zoom (B,) theta = torch.zeros(B, 2, 3, device=device, dtype=torch.float32) theta[:, 0, 0] = scale_x theta[:, 1, 1] = scale_y theta[:, 0, 2] = scale_x - 1 # translate x theta[:, 1, 2] = scale_y - 1 # translate y grid = F.affine_grid(theta, (B, C, grid_size_y, grid_size_x), align_corners=False) pos_embed = F.grid_sample( self.pos_embed.permute(0, 3, 1, 2).expand(B, -1, -1, -1).float(), grid, mode=self.pos_embed_interp_mode, align_corners=False, padding_mode='border', ).to(dtype=x.dtype) # (B, C, H_out, W_out) bi = torch.arange(B, device=device).unsqueeze(1) x += pos_embed[bi, :, patch_coord[..., 0], patch_coord[..., 1]] # NOTE leave as '+=' def _apply_learned_pos_embed( self, x: torch.Tensor, grid_size: List[int], ) -> None: """Apply learned position embeddings to standard 2D batch in-place. Interpolates learned 2D position embeddings to match the specified grid size. Args: x: Input tensor to add position embeddings to [B, H*W, C] grid_size: Target grid size as [height, width] """ orig_h, orig_w = self.pos_embed.shape[1:3] if grid_size[0] == orig_h and grid_size[1] == orig_w: # No resize needed, just flatten pos_embed_flat = self.pos_embed.reshape(1, orig_h * orig_w, -1) else: # Resize if needed - directly using F.interpolate if self.pos_embed_ar_preserving: L = max(grid_size) _interp_size = L, L else: _interp_size = grid_size pos_embed_flat = F.interpolate( self.pos_embed.permute(0, 3, 1, 2).float(), # B,C,H,W size=_interp_size, mode=self.pos_embed_interp_mode, align_corners=False, antialias=True, )[:, :, :grid_size[0], :grid_size[1]].flatten(2).transpose(1, 2) pos_embed_flat = pos_embed_flat.to(dtype=x.dtype) x.add_(pos_embed_flat) def _apply_factorized_naflex_pos_embed( self, x: torch.Tensor, patch_coord: torch.Tensor, ) -> None: """Apply factorized position embeddings to NaFlex batch in-place. Uses separate Y and X position embedding tables that are interpolated and combined for each sample's grid size. Args: x: Input tensor to add position embeddings to [B, N, C] patch_coord: Patch coordinates [B, N, 2] with (y, x) values """ # Calculate grid sizes from patch coordinates naflex_grid_sizes = calculate_naflex_grid_sizes(patch_coord) assert len(naflex_grid_sizes) == x.size(0) # one (H,W) per sample # Handle each batch element separately with its own grid size orig_h, orig_w = self.pos_embed_y.shape[1], self.pos_embed_x.shape[1] # bucket samples that share the same (H, W) so we build each grid once size_to_indices: Dict[Tuple[int, int], List[int]] = {} for bi, k in enumerate(naflex_grid_sizes): size_to_indices.setdefault(k, []).append(bi) def _interp1d(table: torch.Tensor, new_length: int, orig_length: int) -> torch.Tensor: """ Resample a 1-D positional-embedding table to specified length and return it in (1, L, C) layout, dtype matching x. """ if new_length == orig_length: return table.to(dtype=x.dtype) return F.interpolate( table.permute(0, 2, 1).float(), # (1,C,L) → (1,C,L_out) size=new_length, mode='linear', align_corners=False, ).permute(0, 2, 1).to(dtype=x.dtype) # → (1,L_out,C) for k, batch_indices in size_to_indices.items(): target_h, target_w = k if self.pos_embed_ar_preserving: len_y = len_x = max(target_h, target_w) else: len_y, len_x = target_h, target_w pe_y = _interp1d(self.pos_embed_y, len_y, orig_h)[:, :target_h] # (1,H,C) pe_x = _interp1d(self.pos_embed_x, len_x, orig_w)[:, :target_w] # (1,W,C) # Broadcast, add and flatten to sequence layout (row major) pos = pe_y.unsqueeze(2) + pe_x.unsqueeze(1) # (1,H,W,C) pos = pos.flatten(1, 2) seq_len = min(x.shape[1], pos.shape[1]) x[:, :seq_len].index_add_( 0, torch.as_tensor(batch_indices, device=x.device), pos[:, :seq_len].expand(len(batch_indices), -1, -1) ) def _apply_factorized_naflex_pos_embed_grid_sample( self, x: torch.Tensor, patch_coord: torch.Tensor, ) -> None: """Apply factorized position embeddings to NaFlex batch using grid_sample. Uses F.grid_sample for efficient interpolation of separate Y and X position embedding tables based on patch coordinates. Based on proposal by https://github.com/stas-sl Args: x: Input tensor to add position embeddings to [B, N, C] patch_coord: Patch coordinates [B, N, 2] with (y, x) values """ device = x.device B, _, C = x.shape shapes = patch_coord.amax(dim=1) + 1 if self.pos_embed_ar_preserving: # Aspect ratio preserving mode: use square grid with uniform scaling L_i = shapes.amax(dim=1) # (B,) max(h_i, w_i) L_global = L_i.amax() grid_size_y = grid_size_x = L_global scale_x = scale_y = L_global / L_i # uniform zoom (B,) else: # Standard mode: different scaling for x and y grid_size_y, grid_size_x = shapes.amax(0) scale_x = grid_size_x / shapes[:, 1] # horizontal zoom (B,) scale_y = grid_size_y / shapes[:, 0] # vertical zoom (B,) def _interp1d(table: torch.Tensor, scale: torch.Tensor, out_length: torch.Tensor) -> torch.Tensor: pe = table.permute(0, 2, 1).unsqueeze(2).expand(B, -1, -1, -1).float() # (1, L, C) -> (B, C, 1, L) theta = torch.zeros(B, 2, 3, device=x.device) theta[:, 0, 0] = scale theta[:, 0, 2] = scale - 1 theta[:, 1, 1] = 1 grid = F.affine_grid(theta, (B, C, 1, out_length), align_corners=False) pe = F.grid_sample(pe, grid, mode='bilinear', align_corners=False, padding_mode='border') return pe.to(x.dtype) # Interpolate along each axis pe_x = _interp1d(self.pos_embed_x, scale=scale_x, out_length=grid_size_x) pe_y = _interp1d(self.pos_embed_y, scale=scale_y, out_length=grid_size_y) bi = torch.arange(B, device=device).unsqueeze(1) x += pe_x[bi, :, 0, patch_coord[..., 1]] + pe_y[bi, :, 0, patch_coord[..., 0]] def _apply_factorized_pos_embed( self, x: torch.Tensor, grid_size: List[int], ) -> None: """Apply factorized position embeddings to standard 2D batch in-place. Uses separate Y and X position embedding tables that are interpolated and combined for the specified grid size. Args: x: Input tensor to add position embeddings to [B, H*W, C] grid_size: Target grid size as [height, width] """ orig_h, orig_w = self.pos_embed_y.shape[1], self.pos_embed_x.shape[1] target_h, target_w = grid_size if self.pos_embed_ar_preserving: len_y = len_x = max(target_h, target_w) else: len_y, len_x = target_h, target_w def _interp1d(table: torch.Tensor, new_length: int, orig_length: int) -> torch.Tensor: if new_length == orig_length: return table.to(dtype=x.dtype) return F.interpolate( table.permute(0, 2, 1).float(), # (1,L,C) -> (1,C,L) size=new_length, mode='linear', align_corners=False, ).permute(0, 2, 1).to(dtype=x.dtype) # (1,L,C) # Interpolate embeddings pe_y = _interp1d(self.pos_embed_y, len_y, orig_h)[:, :target_h] # (1,H,C) pe_x = _interp1d(self.pos_embed_x, len_x, orig_w)[:, :target_w] # (1,W,C) # Broadcast, add and flatten to sequence layout (row major) pos_embed = pe_y.unsqueeze(2) + pe_x.unsqueeze(1) # (1, H, W, C) pos_embed_flat = pos_embed.flatten(1, 2) # (1, H*W, C) x.add_(pos_embed_flat) def forward( self, x: torch.Tensor, patch_coord: Optional[torch.Tensor] = None, ) -> torch.Tensor: """Forward pass for patch embedding with position encoding. Args: x: Input tensor. Supported formats: - [B, C, H, W] for conv mode - [B, N, P*P*C] for pre-patchified linear mode (normal) - [B, N, Ph, Pw, C] for pre-patchified linear mode (variable patch size) patch_coord: Optional patch coordinates [B, N, 2] for NaFlex mode. Returns: Embedded tensor with position encoding and class/register tokens. Shape: [B, num_prefix_tokens + N, embed_dim] """ grid_size: Optional[List[int]] = None B = x.shape[0] if self.is_linear: # Linear embedding path, works with NaFlex mode or standard 2D mode if patch_coord is None: # Standard 2D (B, C, H, W) mode _assert(x.ndim == 4, 'Expecting 2D image input with input ndim == 4') x, grid_size = batch_patchify(x, self.patch_size, pad=self.dynamic_img_pad) else: # Pre-patchified NaFlex mode # Variable patch size mode: [B, N, Ph, Pw, C], normal mode: [B, N, P*P*C] _assert(x.ndim == 5 or x.ndim == 3, 'Expecting patchified input with ndim == 3 or 5.') # Handle variable patch size projection if self.enable_patch_interpolator and x.ndim == 5: _assert(self.norm_input is None, 'input norm not supported with patch resizing') # Apply projection with interpolation x = self.patch_interpolator( x, self.proj.weight, self.proj.bias, patch_size=tuple(x.shape[2:4]), # patch size from [B, N, Ph, Pw, C] shape is_linear=True, ) else: # Standard projection x = x.flatten(2) # ensure [B, N, P*P*C], flatten Ph*Pw*C if separate if self.norm_input is not None: x = self.norm_input(x) x = self.proj(x) else: _assert(x.ndim == 4, 'Convolutional input must be 4D') if self.dynamic_img_pad: H, W = x.shape[-2:] pad_h = (self.patch_size[0] - H % self.patch_size[0]) % self.patch_size[0] pad_w = (self.patch_size[1] - W % self.patch_size[1]) % self.patch_size[1] x = F.pad(x, (0, pad_w, 0, pad_h)) x = self.proj(x) grid_size = x.shape[-2:] if self.flatten: x = x.flatten(2).transpose(1, 2) # NCHW -> NLC # Apply normalization after flattening x = self.norm(x) if self.pos_embed_type == 'learned': if grid_size is not None: # Standard 2D mode self._apply_learned_pos_embed(x, grid_size=grid_size) else: # NaFlex mode if self.pos_embed_use_grid_sample: self._apply_learned_naflex_pos_embed_grid_sample(x, patch_coord=patch_coord) else: self._apply_learned_naflex_pos_embed(x, patch_coord=patch_coord) elif self.pos_embed_type == 'factorized': if grid_size is not None: # Standard 2D mode self._apply_factorized_pos_embed(x, grid_size=grid_size) else: # NaFlex mode if self.pos_embed_use_grid_sample: self._apply_factorized_naflex_pos_embed_grid_sample(x, patch_coord=patch_coord) else: self._apply_factorized_naflex_pos_embed(x, patch_coord=patch_coord) elif self.pos_embed_type == 'rope': assert False, "ROPE not yet implemented" # Prepare and add class and register tokens to_cat = [] if self.cls_token is not None: to_cat.append(self.cls_token.expand(B, -1, -1)) if self.reg_token is not None: to_cat.append(self.reg_token.expand(B, -1, -1)) # Add tokens to the beginning if to_cat: x = torch.cat(to_cat + [x], dim=1) # Apply dropouts x = self.pos_drop(x) x = self.patch_drop(x) return x @register_notrace_function def create_attention_mask( patch_valid: torch.Tensor, num_prefix_tokens: int = 0, symmetric: bool = True, q_len: Optional[int] = None, dtype: torch.dtype = torch.float32, ) -> Optional[torch.Tensor]: """Creates an attention mask from patch validity information. Supports two modes controlled by `symmetric`: 1. `symmetric=True` (default): Creates a symmetric mask of shape [B, 1, seq_len, seq_len]. An attention pair (i, j) is allowed only if both token i and token j are valid. Suitable for standard self-attention. 2. `symmetric=False`: Creates a potentially non-square mask of shape [B, 1, q_len, kv_len]. An attention pair (q, k) is allowed only if the key/value token k is valid. Query token validity is not checked in the mask itself. Useful for cross-attention or specific self-attention implementations `q_len` can be specified. Used for NaFlex mode to handle variable token counts and padding tokens. Args: patch_valid: Tensor of shape [B, N] with True for valid patches, False for padding. num_prefix_tokens: Number of prefix tokens (class token, register tokens) to prepend, which are always considered valid. symmetric: If True, create a symmetric mask. If False, create an expanded mask based only on key/value validity. q_len: Query sequence length override. Only used when `symmetric` is False. Defaults to the key/value sequence length (`kv_len`) if None. dtype: Dtype of the output attention mask (e.g., torch.float32). Returns: Attention mask tensor. Additive mask (-inf for masked, 0 for unmasked). Shape is [B, 1, seq_len, seq_len] if symmetric=True, or [B, 1, q_len, kv_len] if symmetric=False. """ if patch_valid is None: return None patch_valid = patch_valid.bool() # Ensure boolean type B, N = patch_valid.shape kv_len = N # Initial key/value length is the number of patches # Prepend prefix tokens if any if num_prefix_tokens > 0: # Create prefix validity tensor on the same device/dtype base as patch_valid prefix_valid = patch_valid.new_ones((B, num_prefix_tokens), dtype=torch.bool) # Concatenate prefix and patch validity. Shape becomes [B, num_prefix_tokens + N] patch_valid = torch.cat([prefix_valid, patch_valid], dim=1) kv_len += num_prefix_tokens # Update total key/value sequence length if symmetric: # Symmetric mask is True where BOTH query and key are valid mask_bool = patch_valid.unsqueeze(-1) & patch_valid.unsqueeze(1) mask_bool = mask_bool.unsqueeze(1) # Add head dimension: [B, 1, seq_len, seq_len] else: # Expanded mask q_len = q_len or kv_len mask_bool = patch_valid[:, None, None, :].expand(B, 1, q_len, kv_len) # Create the float mask and apply masking using additive mask convention mask_float = torch.zeros_like(mask_bool, dtype=dtype) # Fill with negative infinity where mask_bool is False (masked positions) mask_float.masked_fill_(~mask_bool, torch.finfo(dtype).min) return mask_float @register_notrace_function def global_pool_naflex( x: torch.Tensor, patch_valid: Optional[torch.Tensor] = None, pool_type: str = 'token', num_prefix_tokens: int = 1, reduce_include_prefix: bool = False, ) -> torch.Tensor: """Global pooling with NaFlex support for masked tokens. Applies global pooling while respecting patch validity masks to exclude padding tokens from pooling operations. Args: x: Input tensor with shape [B, N, C] patch_valid: Optional validity mask for patches [B, N-num_prefix_tokens] pool_type: Type of pooling ('token', 'avg', 'avgmax', 'max') num_prefix_tokens: Number of prefix tokens (class/register) reduce_include_prefix: Whether to include prefix tokens in pooling reduction Returns: Pooled tensor with shape [B, C] """ if patch_valid is None or pool_type not in ('avg', 'avgmax', 'max'): # Fall back to standard pooling x = global_pool_nlc( x, pool_type=pool_type, num_prefix_tokens=num_prefix_tokens, reduce_include_prefix=reduce_include_prefix, ) return x # For NaFlex mode, we need to apply masked pooling to exclude padding tokens if num_prefix_tokens > 0: if reduce_include_prefix: # Include prefix tokens in pooling - they are always considered valid # patch_valid only covers patch tokens, so create combined validity mask prefix_valid = patch_valid.new_ones(x.shape[0], num_prefix_tokens) patch_valid = torch.cat([prefix_valid, patch_valid], dim=1) else: # Exclude prefix tokens from pooling (default behavior) x = x[:, num_prefix_tokens:] patch_valid_float = patch_valid.to(x.dtype) if pool_type == 'avg': # Compute masked average pooling, sum valid tokens and divide by count of valid tokens masked_sums = (x * patch_valid_float.unsqueeze(-1)).sum(dim=1) valid_counts = patch_valid_float.sum(dim=1, keepdim=True).clamp(min=1) pooled = masked_sums / valid_counts return pooled elif pool_type == 'avgmax': # For avgmax, compute masked average and masked max masked_sums = (x * patch_valid_float.unsqueeze(-1)).sum(dim=1) valid_counts = patch_valid_float.sum(dim=1, keepdim=True).clamp(min=1) masked_avg = masked_sums / valid_counts # For max pooling we set masked positions to large negative value masked_x = x.clone() masked_x[~patch_valid] = torch.finfo(masked_x.dtype).min masked_max = masked_x.amax(dim=1) # Combine average and max return 0.5 * (masked_avg + masked_max) elif pool_type == 'max': # For max pooling we set masked positions to large negative value masked_x = x.clone() masked_x[~patch_valid] = torch.finfo(masked_x.dtype).min return masked_x.amax(dim=1) else: assert False class NaFlexVit(nn.Module): """NaFlexVit: Vision Transformer with NaFlex support for flexible input handling. A flexible implementation of Vision Transformer that supports: - Standard image classification with various pooling strategies - NaFlex functionality for variable aspect ratios and resolutions - Linear patch embedding for pre-patchified inputs - Multiple position embedding strategies (learned, factorized, rope) - Comprehensive attention masking for efficient batch processing - Encapsulated embedding and position encoding in FlexEmbeds module - Compatible with standard ViT checkpoints through checkpoint filtering """ def __init__( self, cfg: Optional[NaFlexVitCfg] = None, in_chans: int = 3, num_classes: int = 1000, img_size: Optional[Union[int, Tuple[int, int]]] = None, **kwargs, ) -> None: """Initialize NaFlexVit model. Args: cfg: Model configuration. If None, uses default NaFlexVitCfg. in_chans: Number of input image channels. num_classes: Number of classification classes. img_size: Input image size (for backwards compatibility with classic vit). **kwargs: Additional config parameters to override cfg values. """ super().__init__() # Initialize config cfg = cfg or NaFlexVitCfg() if kwargs: cfg = _overlay_kwargs(cfg, **kwargs) # Validate configuration assert cfg.global_pool in ('', 'avg', 'avgmax', 'max', 'token', 'map') assert cfg.class_token or cfg.global_pool != 'token' assert cfg.pos_embed in ('', 'none', 'learned', 'factorized') # Resolve layer implementations norm_layer = get_norm_layer(cfg.norm_layer) or LayerNorm embed_norm_layer = get_norm_layer(cfg.embed_norm_layer) act_layer = get_act_layer(cfg.act_layer) or nn.GELU block_fn = cfg.block_fn or Block # TODO: Support configurable block_fn via string lookup mlp_layer = cfg.mlp_layer or Mlp # TODO: Support configurable mlp_layer via string lookup # Store instance variables self.num_classes = num_classes self.global_pool = cfg.global_pool self.num_features = self.head_hidden_size = self.embed_dim = cfg.embed_dim # for consistency with other models self.num_prefix_tokens = 1 if cfg.class_token else 0 self.num_prefix_tokens += cfg.reg_tokens self.num_reg_tokens = cfg.reg_tokens self.has_class_token = cfg.class_token self.pool_include_prefix = cfg.pool_include_prefix self.grad_checkpointing = False # Initialize embedding module (includes patch, position embedding, and class/reg tokens) # FlexEmbeds is always used - handles both linear and conv embedding self.embeds = NaFlexEmbeds( patch_size=cfg.patch_size, in_chans=in_chans, embed_dim=cfg.embed_dim, proj_type=cfg.embed_proj_type, proj_bias=not cfg.pre_norm, # disable bias if pre-norm is used (e.g. CLIP) class_token=cfg.class_token, reg_tokens=cfg.reg_tokens, default_img_size=img_size, dynamic_img_pad=cfg.dynamic_img_pad, pos_embed=cfg.pos_embed, pos_embed_grid_size=cfg.pos_embed_grid_size, pos_embed_interp_mode=cfg.pos_embed_interp_mode, pos_embed_ar_preserving=cfg.pos_embed_ar_preserving, pos_embed_use_grid_sample=cfg.pos_embed_use_grid_sample, proj_norm_layer=embed_norm_layer, pos_drop_rate=cfg.pos_drop_rate, patch_drop_rate=cfg.patch_drop_rate, enable_patch_interpolator=getattr(cfg, 'enable_patch_interpolator', False), ) self.norm_pre = norm_layer(cfg.embed_dim) if cfg.pre_norm else nn.Identity() # Transformer blocks dpr = [x.item() for x in torch.linspace(0, cfg.drop_path_rate, cfg.depth)] # stochastic depth decay rule self.blocks = nn.Sequential(*[ block_fn( dim=cfg.embed_dim, num_heads=cfg.num_heads, mlp_ratio=cfg.mlp_ratio, qkv_bias=cfg.qkv_bias, qk_norm=cfg.qk_norm, proj_bias=cfg.proj_bias, init_values=cfg.init_values, proj_drop=cfg.proj_drop_rate, attn_drop=cfg.attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer, act_layer=act_layer, mlp_layer=mlp_layer, ) for i in range(cfg.depth)]) # Feature info for downstream tasks patch_reduction = self.embeds.feat_ratio(as_scalar=True) self.feature_info = [ dict(module=f'blocks.{i}', num_chs=cfg.embed_dim, reduction=patch_reduction) for i in range(cfg.depth) ] self.norm = norm_layer(cfg.embed_dim) if cfg.final_norm and not cfg.fc_norm else nn.Identity() # Classifier Head if cfg.global_pool == 'map': self.attn_pool = AttentionPoolLatent( self.embed_dim, num_heads=cfg.num_heads, mlp_ratio=cfg.mlp_ratio, norm_layer=norm_layer, act_layer=act_layer, ) else: self.attn_pool = None # Handle fc_norm default value fc_norm = cfg.fc_norm if fc_norm is None: fc_norm = cfg.global_pool == 'avg' self.fc_norm = norm_layer(cfg.embed_dim) if cfg.final_norm and fc_norm else nn.Identity() self.head_drop = nn.Dropout(cfg.drop_rate) self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity() if cfg.weight_init != 'skip': self.init_weights(cfg.weight_init) if cfg.fix_init: self.fix_init_weight() def fix_init_weight(self) -> None: """Apply initialization weight fix with layer-wise scaling.""" def rescale(param: torch.Tensor, _layer_id: int) -> None: param.div_(math.sqrt(2.0 * _layer_id)) for layer_id, layer in enumerate(self.blocks): rescale(layer.attn.proj.weight.data, layer_id + 1) rescale(layer.mlp.fc2.weight.data, layer_id + 1) def init_weights(self, mode: str = '') -> None: """Initialize model weights according to specified scheme. Args: mode: Initialization mode ('jax', 'jax_nlhb', 'moco', or '') """ assert mode in ('jax', 'jax_nlhb', 'moco', '') head_bias = -math.log(self.num_classes) if 'nlhb' in mode else 0. named_apply(get_init_weights_vit(mode, head_bias), self) @torch.jit.ignore() def load_pretrained(self, checkpoint_path: str, prefix: str = '') -> None: # Custom loading for the new model structure from .vision_transformer import _load_weights as _orig_load_weights def _load_weights_adapter(model, checkpoint_path, prefix=''): """Adapter function to handle the different model structure""" state_dict = torch.load(checkpoint_path, map_location='cpu') if isinstance(state_dict, dict) and 'state_dict' in state_dict: state_dict = state_dict['state_dict'] # Map original keys to new structure for k in list(state_dict.keys()): if k.startswith('cls_token'): state_dict['embeds.' + k] = state_dict.pop(k) elif k.startswith('reg_token'): state_dict['embeds.' + k] = state_dict.pop(k) elif k.startswith('pos_embed'): state_dict['embeds.' + k] = state_dict.pop(k) elif k.startswith('patch_embed'): state_dict['embeds.' + k[12:]] = state_dict.pop(k) return _orig_load_weights(model, state_dict, prefix) _load_weights_adapter(self, checkpoint_path, prefix) @torch.jit.ignore def no_weight_decay(self) -> Set: """Get set of parameter names that should not have weight decay applied. Returns: Set of parameter names to skip during weight decay """ skip_list = {'embeds.pos_embed', 'embeds.cls_token', 'embeds.reg_token'} return skip_list @torch.jit.ignore def group_matcher(self, coarse: bool = False) -> Dict: """Get parameter group matcher for optimizer parameter grouping. Args: coarse: Whether to use coarse-grained grouping Returns: Dictionary mapping group names to regex patterns """ return dict( stem=r'^embeds', # stem and embed blocks=[(r'^blocks\.(\d+)', None), (r'^norm', (99999,))] ) @torch.jit.ignore def set_grad_checkpointing(self, enable: bool = True) -> None: """Enable or disable gradient checkpointing for memory efficiency. Args: enable: Whether to enable gradient checkpointing """ self.grad_checkpointing = enable if hasattr(self.embeds, 'patch_embed') and hasattr(self.embeds.patch_embed, 'set_grad_checkpointing'): self.embeds.patch_embed.set_grad_checkpointing(enable) @torch.jit.ignore def get_classifier(self) -> nn.Module: """Get the classification head module. Returns: Classification head module """ return self.head def reset_classifier(self, num_classes: int, global_pool: Optional[str] = None) -> None: """Reset the classification head with new number of classes and pooling. Args: num_classes: Number of classes for new classification head global_pool: Optional new global pooling type """ self.num_classes = num_classes if global_pool is not None: assert global_pool in ('', 'avg', 'avgmax', 'max', 'token', 'map') if global_pool == 'map' and self.attn_pool is None: assert False, "Cannot currently add attention pooling in reset_classifier()." elif global_pool != 'map' and self.attn_pool is not None: self.attn_pool = None # remove attention pooling self.global_pool = global_pool self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity() def forward_intermediates( self, x: Union[torch.Tensor, Dict[str, torch.Tensor]], indices: Optional[Union[int, List[int]]] = None, return_prefix_tokens: bool = False, norm: bool = False, stop_early: bool = False, output_fmt: str = 'NCHW', intermediates_only: bool = False, output_dict: bool = False, patch_coord: Optional[torch.Tensor] = None, patch_valid: Optional[torch.Tensor] = None, attn_mask: Optional[torch.Tensor] = None, ) -> Union[List[torch.Tensor], Tuple[torch.Tensor, List[torch.Tensor]], Dict[str, Any]]: """ Forward features that returns intermediates. Args: x: Input image tensor indices: Take last n blocks if int, all if None, select matching indices if sequence return_prefix_tokens: Return both prefix and spatial intermediate tokens norm: Apply norm layer to all intermediates stop_early: Stop iterating over blocks when last desired intermediate hit output_fmt: Shape of intermediate feature outputs intermediates_only: Only return intermediate features output_dict: Return outputs as a dictionary with 'image_features' and 'image_intermediates' keys patch_coord: Optional patch coordinates [B, N, 2] for NaFlex mode patch_valid: Optional patch type indicators (1=patch, 0=padding) for NaFlex attn_mask: Optional attention mask for masked attention Returns: A tuple with (final_features, intermediates), a list of intermediate features, or a dictionary containing 'image_features' and 'image_intermediates' (and optionally 'image_intermediates_prefix') """ # FIXME unfinished / untested assert output_fmt in ('NCHW', 'NLC'), 'Output format must be one of NCHW or NLC.' reshape = output_fmt == 'NCHW' intermediates = [] take_indices, max_index = feature_take_indices(len(self.blocks), indices) if isinstance(x, Dict): # Handle dictionary input from NaFlex collator patch_coord = x['patch_coord'] patch_valid = x['patch_valid'] patches = x['patches'] assert False, 'WIP, patch mode needs more work' else: patches = x height, width = x.shape[-2:] H, W = self.embeds.dynamic_feat_size((height, width)) # Create attention mask if patch_type is provided and mask is not if attn_mask is None and patch_valid is not None: attn_mask = create_attention_mask(patch_valid, self.num_prefix_tokens, patches.dtype) # Forward pass through embedding x = self.embeds(patches, patch_coord=patch_coord) x = self.norm_pre(x) # Forward pass through blocks if torch.jit.is_scripting() or not stop_early: # can't slice blocks in torchscript blocks = self.blocks else: blocks = self.blocks[:max_index + 1] for i, blk in enumerate(blocks): if attn_mask is not None: x = blk(x, attn_mask=attn_mask) elif self.grad_checkpointing and not torch.jit.is_scripting(): x = checkpoint(blk, x) else: x = blk(x) if i in take_indices: # normalize intermediates with final norm layer if enabled intermediates.append(self.norm(x) if norm else x) # Process intermediates if self.num_prefix_tokens: # split prefix (e.g. class, distill) and spatial feature tokens prefix_tokens = [y[:, 0:self.num_prefix_tokens] for y in intermediates] intermediates = [y[:, self.num_prefix_tokens:] for y in intermediates] else: prefix_tokens = None if reshape: # reshape to BCHW output format intermediates = [ y.reshape(y.shape[0], H, W, -1).permute(0, 3, 1, 2).contiguous() for y in intermediates ] # For dictionary output if output_dict: result_dict = {} # Intermediates are always included result_dict['image_intermediates'] = intermediates if prefix_tokens is not None and return_prefix_tokens: result_dict['image_intermediates_prefix'] = prefix_tokens # Only include features if not intermediates_only if not intermediates_only: x_final = self.norm(x) result_dict['image_features'] = x_final return result_dict # For non-dictionary output, maintain the original behavior if not torch.jit.is_scripting() and return_prefix_tokens and prefix_tokens is not None: # return_prefix not support in torchscript due to poor type handling intermediates = list(zip(intermediates, prefix_tokens)) if intermediates_only: return intermediates x = self.norm(x) return x, intermediates def forward_features( self, x: torch.Tensor, patch_coord: Optional[torch.Tensor] = None, patch_valid: Optional[torch.Tensor] = None, attn_mask: Optional[torch.Tensor] = None, ) -> torch.Tensor: if attn_mask is None: attn_mask = create_attention_mask( patch_valid, num_prefix_tokens=self.num_prefix_tokens, dtype=x.dtype ) # Pass through embedding module with patch coordinate/type support x = self.embeds(x, patch_coord=patch_coord) x = self.norm_pre(x) # Apply transformer blocks with masked attention if mask provided if attn_mask is not None: # We need to apply blocks one by one with mask for blk in self.blocks: x = blk(x, attn_mask=attn_mask) elif self.grad_checkpointing and not torch.jit.is_scripting(): x = checkpoint_seq(self.blocks, x) else: x = self.blocks(x) x = self.norm(x) return x def _pool( self, x: torch.Tensor, pool_type: Optional[str] = None, patch_valid: Optional[torch.Tensor] = None, ) -> torch.Tensor: if self.attn_pool is not None: attn_mask = create_attention_mask( patch_valid, num_prefix_tokens=self.num_prefix_tokens if self.pool_include_prefix else 0, symmetric=False, q_len=1, dtype=x.dtype, ) if not self.pool_include_prefix: x = x[:, self.num_prefix_tokens:] x = self.attn_pool(x, attn_mask=attn_mask) return x pool_type = self.global_pool if pool_type is None else pool_type x = global_pool_naflex( x, patch_valid, pool_type=pool_type, num_prefix_tokens=self.num_prefix_tokens, reduce_include_prefix=self.pool_include_prefix, ) return x def forward_head( self, x: torch.Tensor, pre_logits: bool = False, patch_valid: Optional[torch.Tensor] = None, ) -> torch.Tensor: x = self._pool(x, patch_valid=patch_valid) x = self.fc_norm(x) x = self.head_drop(x) return x if pre_logits else self.head(x) def forward( self, x: Union[torch.Tensor, Dict[str, torch.Tensor]], patch_coord: Optional[torch.Tensor] = None, patch_valid: Optional[torch.Tensor] = None, ) -> torch.Tensor: """Forward pass with optional NaFlex support. Args: x: Input tensor. Supported formats: - [B, C, H, W] standard image input - [B, N, P*P*C] pre-patchified tensor (flattened patches) - [B, N, Ph, Pw, C] pre-patchified tensor (variable patch size) - Dict from NaFlex collator patch_coord: Optional patch coordinates [B, N, 2] for NaFlex mode. patch_valid: Optional patch validity indicators for NaFlex. Returns: Model output tensor. """ if isinstance(x, Dict): # Handle dictionary input from NaFlex collator patch_coord = x['patch_coord'] patch_valid = x['patch_valid'] patches = x['patches'] # DEBUG, reconstruct patches # for i in range(len(patches)): # patch = patches[i][patch_valid[i]] # h = (patch_coord[i, :, 0].max() + 1).item() # w = (patch_coord[i, :, 1].max() + 1).item() # patch = patch.reshape(h, w, 16, 16, 3).permute(4, 0, 2, 1, 3) # patch = patch.reshape(3, h*16, w*16) # from torchvision.utils import save_image # save_image(patch, f'patch_{i}.jpg', normalize=True) else: patches = x # Create attention mask if patch_type is provided attn_mask = create_attention_mask( patch_valid, num_prefix_tokens=self.num_prefix_tokens, dtype=patches.dtype, ) # Forward features with mask x = self.forward_features( patches, patch_coord=patch_coord, patch_valid=patch_valid, attn_mask=attn_mask, ) # Pass mask to forward_head for masked pooling x = self.forward_head( x, patch_valid=patch_valid, ) return x def get_init_weights_vit(mode: str = 'jax', head_bias: float = 0.0) -> Callable: """Function imported from vision_transformer.py to maintain compatibility""" from .vision_transformer import init_weights_vit_jax, init_weights_vit_moco, init_weights_vit_timm if 'jax' in mode: return partial(init_weights_vit_jax, head_bias=head_bias) elif 'moco' in mode: return init_weights_vit_moco else: return init_weights_vit_timm def checkpoint_filter_fn(state_dict: Dict[str, Any], model: NaFlexVit) -> Dict[str, Any]: """Handle state dict conversion from original ViT to the new version with combined embedding.""" from .vision_transformer import checkpoint_filter_fn as orig_filter_fn # Handle CombinedEmbed module pattern out_dict = {} for k, v in state_dict.items(): # Convert tokens and embeddings to combined_embed structure if k == 'pos_embed': # Handle position embedding format conversion - from (1, N, C) to (1, H, W, C) if hasattr(model.embeds, 'pos_embed') and v.ndim == 3: num_cls_token = 0 num_reg_token = 0 if 'reg_token' in state_dict: num_reg_token = state_dict['reg_token'].shape[1] if 'cls_token' in state_dict: num_cls_token = state_dict['cls_token'].shape[1] num_prefix_tokens = num_cls_token + num_reg_token # Original format is (1, N, C), need to reshape to (1, H, W, C) num_patches = v.shape[1] num_patches_no_prefix = num_patches - num_prefix_tokens grid_size_no_prefix = math.sqrt(num_patches_no_prefix) grid_size = math.sqrt(num_patches) if (grid_size_no_prefix != grid_size and (grid_size_no_prefix.is_integer() and not grid_size.is_integer()) ): # make a decision, did the pos_embed of the original include the prefix tokens? num_patches = num_patches_no_prefix cls_token_emb = v[:, 0:num_cls_token] if cls_token_emb.numel(): state_dict['cls_token'] += cls_token_emb reg_token_emb = v[:, num_cls_token:num_reg_token] if reg_token_emb.numel(): state_dict['reg_token'] += reg_token_emb v = v[:, num_prefix_tokens:] grid_size = grid_size_no_prefix grid_size = int(grid_size) # Check if it's a perfect square for a standard grid if grid_size * grid_size == num_patches: # Reshape from (1, N, C) to (1, H, W, C) v = v.reshape(1, grid_size, grid_size, v.shape[2]) else: # Not a square grid, we need to get the actual dimensions if hasattr(model.embeds.patch_embed, 'grid_size'): h, w = model.embeds.patch_embed.grid_size if h * w == num_patches: # We have the right dimensions v = v.reshape(1, h, w, v.shape[2]) else: # Dimensions don't match, use interpolation _logger.warning( f"Position embedding size mismatch: checkpoint={num_patches}, model={(h * w)}. " f"Using default initialization and will resize in forward pass." ) # Keep v as is, the forward pass will handle resizing out_dict['embeds.pos_embed'] = v elif k == 'cls_token': out_dict['embeds.cls_token'] = v elif k == 'reg_token': out_dict['embeds.reg_token'] = v # Convert patch_embed.X to embeds.patch_embed.X elif k.startswith('patch_embed.'): suffix = k[12:] if suffix == 'proj.weight': v = v.permute(0, 2, 3, 1).flatten(1) new_key = 'embeds.' + suffix out_dict[new_key] = v else: out_dict[k] = v return out_dict def _cfg(url: str = '', **kwargs) -> Dict[str, Any]: return { 'url': url, 'num_classes': 1000, 'input_size': (3, 384, 384), 'pool_size': None, 'crop_pct': 1.0, 'interpolation': 'bicubic', 'mean': IMAGENET_INCEPTION_MEAN, 'std': IMAGENET_INCEPTION_STD, 'first_conv': 'embeds.proj', 'classifier': 'head', 'license': 'apache-2.0', **kwargs, } default_cfgs = generate_default_cfgs({ 'naflexvit_base_patch16_gap.e300_s576_in1k': _cfg( hf_hub_id='timm/', ), 'naflexvit_base_patch16_par_gap.e300_s576_in1k': _cfg( hf_hub_id='timm/', ), 'naflexvit_base_patch16_parfac_gap.e300_s576_in1k': _cfg( hf_hub_id='timm/', ), 'naflexvit_base_patch16_map.untrained': _cfg(), 'naflexvit_base_patch16_siglip.untrained': _cfg(), 'naflexvit_so400m_patch16_siglip.untrained': _cfg(), }) def _create_naflexvit(variant: str, pretrained: bool = False, **kwargs) -> NaFlexVit: out_indices = kwargs.pop('out_indices', 3) cfg = kwargs.pop('cfg', NaFlexVitCfg()) cfg_field_names = {f.name for f in fields(NaFlexVitCfg)} # pop in-place so the original kwargs is emptied of cfg-specific keys cfg_updates = {k: kwargs.pop(k) for k in list(kwargs) if k in cfg_field_names} if cfg_updates: cfg = _overlay_kwargs(cfg, **cfg_updates) model = build_model_with_cfg( NaFlexVit, variant, pretrained, pretrained_filter_fn=checkpoint_filter_fn, cfg=cfg, feature_cfg=dict(out_indices=out_indices, feature_cls='getter'), **kwargs, ) return model def _create_naflexvit_from_classic( variant: str, pretrained: bool = False, **kwargs, ) -> NaFlexVit: """Create FlexVit model from classic VisionTransformer configuration. This function handles the parameter mapping and configuration logic needed to create FlexVit models that are compatible with classic VisionTransformer configurations and pretrained weights. Args: variant: Model variant name pretrained: Whether to load pretrained weights **kwargs: Classic VisionTransformer parameters Returns: FlexVit model instance """ # Remove VisionTransformer-specific parameters that don't apply to FlexVit kwargs.pop('no_embed_class', None) kwargs.pop('dynamic_img_size', None) # Handle global pooling and fc_norm defaults that differ between ViT and FlexVit gp = kwargs.pop('global_pool', 'token') # Original ViTs default to cls token pooling fc_norm = kwargs.pop('fc_norm', None) # Original ViTs used fc_norm when not set and avg pooling used if fc_norm is None and gp == 'avg': fc_norm = True # Set FlexVit-specific defaults that differ from VisionTransformer flex_kwargs = { 'pos_embed_grid_size': None, # rely on img_size (// patch_size) that will be passed through 'class_token': kwargs.get('class_token', True), 'global_pool': gp, 'fc_norm': fc_norm, **kwargs # User overrides take precedence } return _create_naflexvit(variant, pretrained, **flex_kwargs) @register_model def naflexvit_base_patch16_gap(pretrained: bool = False, **kwargs) -> NaFlexVit: """ViT-Base with NaFlex functionality and global average pooling. """ cfg = NaFlexVitCfg( patch_size=16, embed_dim=768, depth=12, num_heads=12, init_values=1e-5, global_pool='avg', reg_tokens=4, fc_norm=True, ) model = _create_naflexvit('naflexvit_base_patch16_gap', pretrained=pretrained, cfg=cfg, **kwargs) return model @register_model def naflexvit_base_patch16_par_gap(pretrained: bool = False, **kwargs) -> NaFlexVit: """ViT-Base with NaFlex functionality, aspect preserving pos embed, global average pooling. """ cfg = NaFlexVitCfg( patch_size=16, embed_dim=768, depth=12, num_heads=12, init_values=1e-5, pos_embed_ar_preserving=True, global_pool='avg', reg_tokens=4, fc_norm=True, ) model = _create_naflexvit('naflexvit_base_patch16_par_gap', pretrained=pretrained, cfg=cfg, **kwargs) return model @register_model def naflexvit_base_patch16_parfac_gap(pretrained: bool = False, **kwargs) -> NaFlexVit: """ViT-Base with NaFlex functionality, aspect preserving & factorized pos embed, global average pooling. """ cfg = NaFlexVitCfg( patch_size=16, embed_dim=768, depth=12, num_heads=12, init_values=1e-5, pos_embed_ar_preserving=True, pos_embed='factorized', global_pool='avg', reg_tokens=4, fc_norm=True, ) model = _create_naflexvit('naflexvit_base_patch16_parfac_gap', pretrained=pretrained, cfg=cfg, **kwargs) return model @register_model def naflexvit_base_patch16_map(pretrained: bool = False, **kwargs) -> NaFlexVit: """ViT-Base with NaFlex functionality and MAP attention pooling. """ cfg = NaFlexVitCfg( patch_size=16, embed_dim=768, depth=12, num_heads=12, init_values=1e-5, global_pool='map', reg_tokens=1, ) model = _create_naflexvit('naflexvit_base_patch16_map', pretrained=pretrained, cfg=cfg, **kwargs) return model @register_model def naflexvit_so150m2_patch16_reg1_gap(pretrained: bool = False, **kwargs) -> NaFlexVit: """ViT-SO150M2 with NaFlex functionality for variable aspect ratios and resolutions. This model supports: 1. Variable aspect ratios and resolutions via patch coordinates 2. Position embedding interpolation for arbitrary grid sizes 3. Explicit patch coordinates and valid token masking """ cfg = NaFlexVitCfg( patch_size=16, embed_dim=832, depth=21, num_heads=13, mlp_ratio=34/13, init_values=1e-5, qkv_bias=False, reg_tokens=1, global_pool='avg', fc_norm=True, ) model = _create_naflexvit('naflexvit_so150m2_patch16_reg1_gap', pretrained=pretrained, cfg=cfg, **kwargs) return model @register_model def naflexvit_so150m2_patch16_reg1_map(pretrained: bool = False, **kwargs) -> NaFlexVit: """ViT-SO150M2 with NaFlex functionality for variable aspect ratios and resolutions. This model supports: 1. Variable aspect ratios and resolutions via patch coordinates 2. Position embedding interpolation for arbitrary grid sizes 3. Explicit patch coordinates and valid token masking """ cfg = NaFlexVitCfg( patch_size=16, embed_dim=832, depth=21, num_heads=13, mlp_ratio=34/13, init_values=1e-5, qkv_bias=False, reg_tokens=1, global_pool='map', ) model = _create_naflexvit('naflexvit_so150m2_patch16_reg1_map', pretrained=pretrained, cfg=cfg, **kwargs) return model @register_model def naflexvit_base_patch16_siglip(pretrained: bool = False, **kwargs) -> NaFlexVit: """ViT-Base with NaFlex functionality and SigLIP-style configuration. """ cfg = NaFlexVitCfg( patch_size=16, embed_dim=768, depth=12, num_heads=12, act_layer='gelu_tanh', global_pool='map', ) model = _create_naflexvit('naflexvit_base_patch16_siglip', pretrained=pretrained, cfg=cfg, **kwargs) return model @register_model def naflexvit_so400m_patch16_siglip(pretrained: bool = False, **kwargs) -> NaFlexVit: """ViT-SO400M with NaFlex functionality for variable aspect ratios and resolutions. """ cfg = NaFlexVitCfg( patch_size=16, embed_dim=1152, depth=27, num_heads=16, mlp_ratio=3.7362, act_layer='gelu_tanh', global_pool='map', ) model = _create_naflexvit('naflexvit_so400m_patch16_siglip', pretrained=pretrained, cfg=cfg, **kwargs) return model