from __future__ import annotations from dataclasses import dataclass from typing import Any import torch from torch import nn from torch.nn import functional as F from transformers import PreTrainedModel from transformers.utils import ModelOutput from .configuration_efficient_rae import EfficientRAEConfig @dataclass class DepthAttentionRecord: source_layers: tuple[int, ...] weights: torch.Tensor @dataclass class EfficientRAEOutput(ModelOutput): sample: torch.Tensor | None = None latents: torch.Tensor | None = None reconstructed_features: torch.Tensor | None = None teacher_features: torch.Tensor | None = None pooler_depth_attentions: tuple[DepthAttentionRecord, ...] | None = None unpooler_depth_attentions: tuple[DepthAttentionRecord, ...] | None = None @dataclass class CompressorOutput(ModelOutput): compact: torch.Tensor | None = None corrupted_compact: torch.Tensor | None = None restored: torch.Tensor | None = None semantic_prediction: torch.Tensor | None = None noise_levels: torch.Tensor | None = None mask_ratios: torch.Tensor | None = None pooler_depth_attentions: tuple[DepthAttentionRecord, ...] | None = None unpooler_depth_attentions: tuple[DepthAttentionRecord, ...] | None = None def _rotate_pairs( x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor ) -> torch.Tensor: x_even = x[..., 0::2] x_odd = x[..., 1::2] rotated = torch.stack((-x_odd, x_even), dim=-1).flatten(-2) return x * cos + rotated * sin def apply_2d_rope( query: torch.Tensor, key: torch.Tensor, height: int, width: int, ) -> tuple[torch.Tensor, torch.Tensor]: """Apply parameter-free 2D RoPE to tensors shaped [B, heads, tokens, dim].""" if query.shape[-2] != height * width: raise ValueError("token count does not match the supplied spatial grid") half = query.shape[-1] // 2 quarter = half // 2 inv_freq = 1.0 / ( 10_000 ** ( torch.arange(quarter, device=query.device, dtype=torch.float32) / max(quarter, 1) ) ) rows = torch.arange(height, device=query.device, dtype=torch.float32) cols = torch.arange(width, device=query.device, dtype=torch.float32) row_grid, col_grid = torch.meshgrid(rows, cols, indexing="ij") def frequencies(positions: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: angles = positions.flatten()[:, None] * inv_freq[None, :] angles = angles.repeat_interleave(2, dim=-1) return ( angles.cos().to(dtype=query.dtype)[None, None], angles.sin().to(dtype=query.dtype)[None, None], ) row_cos, row_sin = frequencies(row_grid) col_cos, col_sin = frequencies(col_grid) def rotate(tensor: torch.Tensor) -> torch.Tensor: row_part, col_part = tensor.split(half, dim=-1) row_part = _rotate_pairs(row_part, row_cos, row_sin) col_part = _rotate_pairs(col_part, col_cos, col_sin) return torch.cat((row_part, col_part), dim=-1) return rotate(query), rotate(key) class FeedForward(nn.Module): def __init__( self, hidden_size: int, intermediate_size: int, dropout: float, output_size: int | None = None, ): super().__init__() self.up_proj = nn.Linear(hidden_size, intermediate_size) self.activation = nn.GELU() self.dropout = nn.Dropout(dropout) self.down_proj = nn.Linear( intermediate_size, output_size if output_size is not None else hidden_size ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states = self.up_proj(hidden_states) hidden_states = self.activation(hidden_states) hidden_states = self.dropout(hidden_states) return self.down_proj(hidden_states) class SpatialAttention(nn.Module): def __init__( self, hidden_size: int, num_heads: int, dropout: float, attention_dropout: float, ): super().__init__() self.num_heads = num_heads self.head_size = hidden_size // num_heads self.attention_dropout = attention_dropout self.query = nn.Linear(hidden_size, hidden_size) self.key = nn.Linear(hidden_size, hidden_size) self.value = nn.Linear(hidden_size, hidden_size) self.output = nn.Linear(hidden_size, hidden_size) self.output_dropout = nn.Dropout(dropout) def _split_heads(self, tensor: torch.Tensor) -> torch.Tensor: batch, tokens, _ = tensor.shape return tensor.reshape(batch, tokens, self.num_heads, self.head_size).transpose( 1, 2 ) def forward( self, hidden_states: torch.Tensor, grid_size: tuple[int, int], source_states: list[tuple[int, torch.Tensor, torch.Tensor]], current_layer: int, depth_attention: bool, depth_attention_stride: int, output_depth_attentions: bool, ) -> tuple[ torch.Tensor, tuple[int, torch.Tensor, torch.Tensor], DepthAttentionRecord | None, ]: query = self._split_heads(self.query(hidden_states)) key = self._split_heads(self.key(hidden_states)) value = self._split_heads(self.value(hidden_states)) mixed_value = value record = None if depth_attention: selected = [ source for source in source_states if source[0] % depth_attention_stride == 0 ] source_layers = tuple(source[0] for source in selected) + (current_layer,) source_keys = [source[1] for source in selected] + [key] source_values = [source[2] for source in selected] + [value] stacked_keys = torch.stack(source_keys, dim=-2) stacked_values = torch.stack(source_values, dim=-2) scale = self.head_size**-0.5 scores = torch.einsum("bhnd,bhnsd->bhns", query, stacked_keys) * scale weights = torch.softmax(scores, dim=-1, dtype=torch.float32).to(query.dtype) mixed_value = torch.einsum("bhns,bhnsd->bhnd", weights, stacked_values) if output_depth_attentions: record = DepthAttentionRecord(source_layers, weights) rotary_query, rotary_key = apply_2d_rope( query, key, height=grid_size[0], width=grid_size[1] ) context = F.scaled_dot_product_attention( rotary_query, rotary_key, mixed_value, dropout_p=self.attention_dropout if self.training else 0.0, ) context = context.transpose(1, 2).contiguous().flatten(2) context = self.output_dropout(self.output(context)) return context, (current_layer, key, mixed_value), record class DepthTransformerBlock(nn.Module): def __init__(self, config: EfficientRAEConfig): super().__init__() self.attention_norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.attention = SpatialAttention( config.hidden_size, config.num_attention_heads, config.dropout, config.attention_dropout, ) self.mlp_norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.mlp = FeedForward( config.hidden_size, config.intermediate_size, config.dropout ) self.residual_dropout = nn.Dropout(config.dropout) def forward( self, hidden_states: torch.Tensor, grid_size: tuple[int, int], source_states: list[tuple[int, torch.Tensor, torch.Tensor]], current_layer: int, depth_attention: bool, depth_attention_stride: int, output_depth_attentions: bool = False, ) -> tuple[ torch.Tensor, tuple[int, torch.Tensor, torch.Tensor], DepthAttentionRecord | None, ]: attention_output, source, record = self.attention( self.attention_norm(hidden_states), grid_size, source_states, current_layer, depth_attention, depth_attention_stride, output_depth_attentions, ) hidden_states = hidden_states + attention_output hidden_states = hidden_states + self.residual_dropout( self.mlp(self.mlp_norm(hidden_states)) ) return hidden_states, source, record class SemanticAttentionPooler(nn.Module): """Pool compact spatial tokens into the frozen backbone token sequence.""" def __init__(self, config: EfficientRAEConfig): super().__init__() self.num_heads = config.num_attention_heads self.head_size = config.encoder_hidden_size // self.num_heads self.attention_dropout = config.attention_dropout self.queries = nn.Parameter( torch.empty(1, config.semantic_num_queries, config.encoder_hidden_size) ) self.key = nn.Linear(config.hidden_size, config.encoder_hidden_size) self.value = nn.Linear(config.hidden_size, config.encoder_hidden_size) self.output = nn.Linear(config.encoder_hidden_size, config.encoder_hidden_size) self.output_dropout = nn.Dropout(config.dropout) def _split_heads(self, tensor: torch.Tensor) -> torch.Tensor: batch, tokens, _ = tensor.shape return tensor.reshape(batch, tokens, self.num_heads, self.head_size).transpose( 1, 2 ) def forward(self, compact: torch.Tensor) -> torch.Tensor: batch = compact.shape[0] tokens = compact.flatten(2).transpose(1, 2) query = self._split_heads(self.queries.expand(batch, -1, -1)) key = self._split_heads(self.key(tokens)) value = self._split_heads(self.value(tokens)) context = F.scaled_dot_product_attention( query, key, value, dropout_p=self.attention_dropout if self.training else 0.0, ) context = context.transpose(1, 2).contiguous().flatten(2) return self.output_dropout(self.output(context)) class SpatialPooler(nn.Module): def __init__(self, config: EfficientRAEConfig): super().__init__() self.config = config self.input_projection = ( nn.Identity() if config.encoder_hidden_size == config.hidden_size else nn.Linear(config.encoder_hidden_size, config.hidden_size) ) self.layers = nn.ModuleList( [ DepthTransformerBlock(config) for _ in range(config.pooler_num_hidden_layers) ] ) shuffled_size = config.hidden_size * config.pool_size**2 self.reduction_norm = nn.RMSNorm(shuffled_size, eps=config.rms_norm_eps) self.reduction = FeedForward( shuffled_size, config.intermediate_size, config.dropout, output_size=config.hidden_size, ) self.out_norm = nn.LayerNorm( config.hidden_size, eps=config.layer_norm_eps, elementwise_affine=False ) self.semantic_pooler = SemanticAttentionPooler(config) def forward( self, features: torch.Tensor, output_depth_attentions: bool = False, ) -> tuple[torch.Tensor, tuple[DepthAttentionRecord, ...] | None]: if features.ndim != 4: raise ValueError( "features must have shape [batch, channels, height, width]" ) batch, channels, height, width = features.shape if channels != self.config.encoder_hidden_size: raise ValueError( f"expected {self.config.encoder_hidden_size} channels, got {channels}" ) if height % self.config.pool_size or width % self.config.pool_size: raise ValueError("feature grid must be divisible by pool_size") hidden_states = features.flatten(2).transpose(1, 2) hidden_states = self.input_projection(hidden_states) sources: list[tuple[int, torch.Tensor, torch.Tensor]] = [] records: list[DepthAttentionRecord] = [] for layer_index, layer in enumerate(self.layers): hidden_states, source, record = layer( hidden_states, (height, width), sources, layer_index, self.config.pooler_depth_attention, self.config.depth_attention_stride, output_depth_attentions, ) sources.append(source) if record is not None: records.append(record) hidden_states = hidden_states.transpose(1, 2).reshape( batch, self.config.hidden_size, height, width ) hidden_states = F.pixel_unshuffle(hidden_states, self.config.pool_size) compact_height, compact_width = hidden_states.shape[-2:] hidden_states = hidden_states.flatten(2).transpose(1, 2) hidden_states = self.reduction(self.reduction_norm(hidden_states)) hidden_states = self.out_norm(hidden_states) compact = hidden_states.transpose(1, 2).reshape( batch, self.config.hidden_size, compact_height, compact_width, ) return compact, tuple(records) if output_depth_attentions else None class SpatialUnpooler(nn.Module): def __init__(self, config: EfficientRAEConfig): super().__init__() self.config = config self.input_projection = nn.Linear(config.hidden_size, config.hidden_size) self.expansion_norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.expansion = nn.Linear( config.hidden_size, config.hidden_size * config.pool_size**2 ) self.expansion_dropout = nn.Dropout(config.dropout) self.layers = nn.ModuleList( [ DepthTransformerBlock(config) for _ in range(config.unpooler_num_hidden_layers) ] ) self.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.head = nn.Linear(config.hidden_size, config.encoder_hidden_size) def forward( self, compact: torch.Tensor, output_depth_attentions: bool = False, ) -> tuple[torch.Tensor, tuple[DepthAttentionRecord, ...] | None]: if compact.ndim != 4: raise ValueError("compact must have shape [batch, channels, height, width]") batch, channels, height, width = compact.shape if channels != self.config.hidden_size: raise ValueError( f"expected {self.config.hidden_size} channels, got {channels}" ) hidden_states = compact.flatten(2).transpose(1, 2) hidden_states = self.input_projection(hidden_states) hidden_states = self.expansion_dropout( self.expansion(self.expansion_norm(hidden_states)) ) hidden_states = hidden_states.transpose(1, 2).reshape( batch, self.config.hidden_size * self.config.pool_size**2, height, width, ) hidden_states = F.pixel_shuffle(hidden_states, self.config.pool_size) full_height, full_width = hidden_states.shape[-2:] hidden_states = hidden_states.flatten(2).transpose(1, 2) sources: list[tuple[int, torch.Tensor, torch.Tensor]] = [] records: list[DepthAttentionRecord] = [] for layer_index, layer in enumerate(self.layers): hidden_states, source, record = layer( hidden_states, (full_height, full_width), sources, layer_index, self.config.unpooler_depth_attention, self.config.depth_attention_stride, output_depth_attentions, ) sources.append(source) if record is not None: records.append(record) restored = self.head(self.norm(hidden_states)) restored = restored.transpose(1, 2).reshape( batch, self.config.encoder_hidden_size, full_height, full_width, ) return restored, tuple(records) if output_depth_attentions else None class EfficientRAECompressor(nn.Module): def __init__(self, config: EfficientRAEConfig): super().__init__() self.config = config self.pooler = SpatialPooler(config) self.unpooler = SpatialUnpooler(config) self.mask_token = nn.Parameter(torch.empty(1, config.hidden_size, 1, 1)) self.apply(_initialize_module) nn.init.normal_(self.pooler.semantic_pooler.queries, std=0.02) nn.init.normal_(self.mask_token, std=0.02) def corrupt( self, compact: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: batch, channels, height, width = compact.shape noise_levels = torch.empty( batch, device=compact.device, dtype=torch.float32 ).uniform_(0.0, self.config.noise_t_max) interpolation = noise_levels.to(compact.dtype).view(batch, 1, 1, 1) corrupted = (1.0 - interpolation) * compact + interpolation * torch.randn_like( compact ) sampled_ratios = torch.empty( batch, device=compact.device, dtype=torch.float32 ).uniform_(self.config.mask_ratio_min, self.config.mask_ratio_max) sampled_ratios.clamp_min_(0.0) token_count = height * width mask_counts = torch.floor(sampled_ratios * token_count).to(torch.long) ordering = torch.rand(batch, token_count, device=compact.device).argsort(dim=1) selected = torch.arange(token_count, device=compact.device).expand(batch, -1) selected = selected < mask_counts[:, None] mask = torch.zeros( batch, token_count, device=compact.device, dtype=torch.bool ).scatter_(1, ordering, selected) tokens = corrupted.flatten(2).transpose(1, 2) mask_token = self.mask_token.flatten(2).transpose(1, 2) tokens = torch.where(mask[:, :, None], mask_token, tokens) corrupted = tokens.transpose(1, 2).reshape(batch, channels, height, width) realized_ratios = mask_counts.to(torch.float32) / token_count return corrupted, noise_levels, realized_ratios def forward( self, features: torch.Tensor, output_depth_attentions: bool = False, output_semantic: bool = False, ) -> CompressorOutput: compact, pooler_attentions = self.pooler( features, output_depth_attentions=output_depth_attentions ) if self.training: corrupted, noise_levels, mask_ratios = self.corrupt(compact) else: corrupted = compact noise_levels = compact.new_zeros(compact.shape[0], dtype=torch.float32) mask_ratios = compact.new_zeros(compact.shape[0], dtype=torch.float32) restored, unpooler_attentions = self.unpooler( corrupted, output_depth_attentions=output_depth_attentions ) semantic_prediction = ( self.pooler.semantic_pooler(corrupted) if output_semantic else None ) return CompressorOutput( compact=compact, corrupted_compact=corrupted, restored=restored, semantic_prediction=semantic_prediction, noise_levels=noise_levels, mask_ratios=mask_ratios, pooler_depth_attentions=pooler_attentions, unpooler_depth_attentions=unpooler_attentions, ) def _initialize_module(module: nn.Module) -> None: if isinstance(module, nn.Linear): nn.init.xavier_uniform_(module.weight) if module.bias is not None: nn.init.zeros_(module.bias) elif isinstance(module, nn.RMSNorm): nn.init.ones_(module.weight) class EfficientRAEModel(PreTrainedModel): config_class = EfficientRAEConfig # The serialized checkpoint is an adapter, not the nested RAE base model. # Leaving this empty prevents Transformers from treating adapter keys as an # unprefixed base-model checkpoint during from_pretrained(). base_model_prefix = "" main_input_name = "pixel_values" supports_gradient_checkpointing = False def __init__( self, config: EfficientRAEConfig, rae: nn.Module | None = None, **kwargs: Any, ): super().__init__(config) if rae is None: raise ValueError( "rae must be loaded independently and passed as rae=... when " "constructing or loading EfficientRAEModel" ) self.rae = rae self.compressor = EfficientRAECompressor(config) self._freeze_base_model() @classmethod def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs): requested_loading_info = bool(kwargs.get("output_loading_info", False)) kwargs["output_loading_info"] = True model, loading_info = super().from_pretrained( pretrained_model_name_or_path, *model_args, **kwargs ) required = { "compressor.mask_token", "compressor.pooler.semantic_pooler.queries", } missing = sorted(required.intersection(loading_info["missing_keys"])) if missing: raise ValueError( "checkpoint predates compressor format v2 and cannot be loaded; " f"missing: {', '.join(missing)}" ) if requested_loading_info: return model, loading_info return model @property def encoder(self) -> nn.Module: return self.rae.encoder @property def decoder(self) -> nn.Module: return self.rae.decoder @property def pooler(self) -> SpatialPooler: return self.compressor.pooler @property def unpooler(self) -> SpatialUnpooler: return self.compressor.unpooler def _freeze_base_model(self) -> None: self.rae.requires_grad_(False) self.rae.eval() def train(self, mode: bool = True) -> "EfficientRAEModel": super().train(mode) self.rae.eval() return self def state_dict(self, *args: Any, **kwargs: Any) -> dict[str, torch.Tensor]: state = super().state_dict(*args, **kwargs) return { key: value for key, value in state.items() if key.startswith("compressor.") } def full_state_dict(self) -> dict[str, torch.Tensor]: return super().state_dict() def compress( self, teacher_features: torch.Tensor, output_depth_attentions: bool = False, ) -> tuple[torch.Tensor, tuple[DepthAttentionRecord, ...] | None]: return self.pooler( teacher_features, output_depth_attentions=output_depth_attentions ) def uncompress( self, compact_latents: torch.Tensor, output_depth_attentions: bool = False, ) -> tuple[torch.Tensor, tuple[DepthAttentionRecord, ...] | None]: return self.unpooler( compact_latents, output_depth_attentions=output_depth_attentions ) def encode(self, pixel_values: torch.Tensor) -> torch.Tensor: teacher_features = self.rae.encode(pixel_values) compact, _ = self.compress(teacher_features) return compact def decode(self, compact_latents: torch.Tensor) -> torch.Tensor: restored, _ = self.uncompress(compact_latents) return self.rae.decode(restored) def forward( self, pixel_values: torch.Tensor, return_features: bool = False, output_depth_attentions: bool = False, return_dict: bool = True, ) -> EfficientRAEOutput | tuple[torch.Tensor, ...]: rae = self.rae teacher = rae.encode(pixel_values) compressor_output = self.compressor( teacher, output_depth_attentions=output_depth_attentions ) sample = rae.decode(compressor_output.restored) output = EfficientRAEOutput( sample=sample, latents=compressor_output.compact if return_features else None, reconstructed_features=( compressor_output.restored if return_features else None ), teacher_features=teacher if return_features else None, pooler_depth_attentions=compressor_output.pooler_depth_attentions, unpooler_depth_attentions=compressor_output.unpooler_depth_attentions, ) if return_dict: return output return output.to_tuple() EfficientRAEModel.register_for_auto_class("AutoModel")