modeling_efficient_rae.py

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")