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 AutoModel, 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


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 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)
            ]
        )
        self.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        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,
        )

    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 = self.norm(hidden_states)
        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))
        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.encoder_hidden_size, config.encoder_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.pooler = SpatialPooler(config)
        self.unpooler = SpatialUnpooler(config)
        self.apply(_initialize_module)

    def forward(
        self,
        features: torch.Tensor,
        output_depth_attentions: bool = False,
    ) -> tuple[
        torch.Tensor,
        torch.Tensor,
        tuple[DepthAttentionRecord, ...] | None,
        tuple[DepthAttentionRecord, ...] | None,
    ]:
        compact, pooler_attentions = self.pooler(
            features, output_depth_attentions=output_depth_attentions
        )
        restored, unpooler_attentions = self.unpooler(
            compact, output_depth_attentions=output_depth_attentions
        )
        return compact, restored, pooler_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,
        **base_model_kwargs: Any,
    ):
        super().__init__(config)
        if rae is None:
            rae = AutoModel.from_pretrained(
                config.base_model_name_or_path,
                revision=config.base_model_revision,
                trust_remote_code=True,
                **base_model_kwargs,
            )
        self.rae = rae
        self.compressor = EfficientRAECompressor(config)
        self._freeze_base_model()

    @classmethod
    def from_base_model(
        cls,
        base_model_name_or_path: str,
        revision: str | None = None,
        **config_overrides: Any,
    ) -> "EfficientRAEModel":
        rae = AutoModel.from_pretrained(
            base_model_name_or_path,
            revision=revision,
            trust_remote_code=True,
        )
        base_config = rae.config
        config_values = dict(
            base_model_name_or_path=base_model_name_or_path,
            base_model_revision=revision or getattr(base_config, "_commit_hash", None),
            image_size=base_config.image_size,
            patch_size=base_config.patch_size,
            encoder_hidden_size=base_config.hidden_size,
            hidden_size=base_config.hidden_size,
            num_attention_heads=getattr(base_config, "encoder_num_attention_heads", 16),
        )
        config_values.update(config_overrides)
        config = EfficientRAEConfig(**config_values)
        return cls(config, rae=rae)

    @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, ...]:
        teacher = self.rae.encode(pixel_values)
        compact, restored, pooler_attentions, unpooler_attentions = self.compressor(
            teacher, output_depth_attentions=output_depth_attentions
        )
        sample = self.rae.decode(restored)
        output = EfficientRAEOutput(
            sample=sample,
            latents=compact if return_features else None,
            reconstructed_features=restored if return_features else None,
            teacher_features=teacher if return_features else None,
            pooler_depth_attentions=pooler_attentions,
            unpooler_depth_attentions=unpooler_attentions,
        )
        if return_dict:
            return output
        return output.to_tuple()


EfficientRAEModel.register_for_auto_class("AutoModel")