# coding=utf-8
# PaGE: Patch-level Gaze Estimation with cross-attention scene/head interaction.
#
# This file is self-contained. Dependencies: torch, torchvision, timm, transformers (>=4.56, which ships DINOv3 built-in).
# The DINOv3 backbones are constructed from config only (no external checkpoint download); their weights live in this
# model's safetensors alongside the gaze decoder.
#
# Auto-map entry points: PaGEConfig, PaGEModel, PaGEImageProcessor.

from __future__ import annotations

import math
from typing import Optional, Tuple, Type, Union, List

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision

from timm.models.vision_transformer import Block
from timm.layers.mlp import SwiGLU, Mlp
from timm.layers import DropPath, LayerNorm, LayerScale, use_fused_attn

from transformers import PretrainedConfig, PreTrainedModel
from transformers.image_processing_utils import BaseImageProcessor

# --------------------------------------------------------------------------- #
# Robust DINOv3 import (built into transformers >= 4.56).                     #
# --------------------------------------------------------------------------- #
try:
    from transformers.models.dinov3_vit import DINOv3ViTModel, DINOv3ViTConfig
except Exception as e:  # pragma: no cover
    raise ImportError(
        "PaGE requires transformers>=4.56 with built-in DINOv3 support (`transformers.models.dinov3_vit`). "
        f"Import failed: {e!r}"
    )


# =========================================================================== #
#  Vendored utilities (from gazelle.utils)                                     #
# =========================================================================== #
def repeat_tensors(tensor, repeat_counts):
    repeated_tensors = [
        tensor[i:i + 1].repeat(repeat, *[1] * (tensor.ndim - 1))
        for i, repeat in enumerate(repeat_counts)
    ]
    return torch.cat(repeated_tensors, dim=0)


def split_tensors(tensor, split_counts):
    indices = torch.cumsum(torch.tensor([0] + split_counts), dim=0)
    return [tensor[indices[i]:indices[i + 1]] for i in range(len(split_counts))]


class TransposeLayerNorm(nn.Module):
    """Transpose 2D feature maps for layer norm, then transpose back."""

    def __init__(self, dim):
        super().__init__()
        self.ln = nn.LayerNorm(dim)

    def forward(self, x: torch.Tensor):
        x = x.permute(0, 2, 3, 1).contiguous()
        x = self.ln(x)
        x = x.permute(0, 3, 1, 2).contiguous()
        return x


def positionalencoding2d(d_model, height, width):
    if d_model % 4 != 0:
        raise ValueError(
            "Cannot use sin/cos positional encoding with odd dimension (got dim={:d})".format(d_model))
    pe = torch.zeros(d_model, height, width)
    d_model = int(d_model / 2)
    div_term = torch.exp(torch.arange(0., d_model, 2)
                         * -(math.log(10000.0) / d_model))
    pos_w = torch.arange(0., width).unsqueeze(1)
    pos_h = torch.arange(0., height).unsqueeze(1)
    pe[0:d_model:2, :, :] = torch.sin(pos_w * div_term).transpose(0, 1).unsqueeze(1).repeat(1, height, 1)
    pe[1:d_model:2, :, :] = torch.cos(pos_w * div_term).transpose(0, 1).unsqueeze(1).repeat(1, height, 1)
    pe[d_model::2, :, :] = torch.sin(pos_h * div_term).transpose(0, 1).unsqueeze(2).repeat(1, 1, width)
    pe[d_model + 1::2, :, :] = torch.cos(pos_h * div_term).transpose(0, 1).unsqueeze(2).repeat(1, 1, width)
    return pe


# =========================================================================== #
#  Vendored Axial 2D RoPE self-attention (from gazelle.rope_self_attention)    #
# =========================================================================== #
GridSize = Optional[Tuple[int, int]]
Rect = Union[Tuple[float, float, float, float], torch.Tensor]


class Axial2dRotaryEmbedding(nn.Module):
    def __init__(self, dim: int, base: float = 100.0) -> None:
        super().__init__()
        if dim <= 0 or dim % 4 != 0:
            raise ValueError(f"`dim` must be a positive multiple of 4, got {dim}.")
        self.dim = dim
        self.axis_dim = dim // 2
        self.base = base
        self.register_buffer("inv_freq", self._compute_inv_freq(), persistent=False)

    def _compute_inv_freq(self):
        return 1.0 / (self.base ** (torch.arange(0, self.axis_dim, 2, dtype=torch.float32) / self.axis_dim))

    def reset_inv_freq(self):
        """Recompute inv_freq (a non-persistent buffer that meta-init in from_pretrained can corrupt)."""
        self.inv_freq = self._compute_inv_freq()

    def _axis_cos_sin(self, coords, *, device, dtype):
        inv_freq = self.inv_freq.to(device=device, dtype=torch.float32)
        freqs = coords.to(device=device, dtype=torch.float32)[:, None] * inv_freq[None, :]
        return freqs.cos().to(dtype=dtype), freqs.sin().to(dtype=dtype)

    @staticmethod
    def _rotate_axis(x, cos, sin):
        x = x.reshape(*x.shape[:-1], -1, 2)
        x_even, x_odd = x.unbind(dim=-1)
        cos = cos[None, None, :, :]
        sin = sin[None, None, :, :]
        x_rot = torch.stack((x_even * cos - x_odd * sin, x_even * sin + x_odd * cos), dim=-1)
        return x_rot.flatten(-2)

    def forward(self, q, k, grid_size, num_front_tokens=0):
        _, _, n, head_dim = q.shape
        gh, gw = grid_size
        num_patch_tokens = gh * gw
        expected = num_front_tokens + num_patch_tokens
        if n != expected:
            raise ValueError(f"Token count mismatch: got N={n}, expected {num_front_tokens} front + {gh}*{gw} = {expected}.")
        if self.dim > head_dim:
            raise ValueError(f"RoPE dim {self.dim} exceeds head_dim {head_dim}.")
        q_front, q_patch = q[:, :, :num_front_tokens], q[:, :, num_front_tokens:]
        k_front, k_patch = k[:, :, :num_front_tokens], k[:, :, num_front_tokens:]
        yy, xx = torch.meshgrid(torch.arange(gh, device=q.device), torch.arange(gw, device=q.device), indexing="ij")
        yy = yy.reshape(-1); xx = xx.reshape(-1)
        cos_y, sin_y = self._axis_cos_sin(yy, device=q.device, dtype=q.dtype)
        cos_x, sin_x = self._axis_cos_sin(xx, device=q.device, dtype=q.dtype)

        def apply_rope(t):
            t_rope, t_pass = t[..., :self.dim], t[..., self.dim:]
            t_y, t_x = t_rope.split(self.axis_dim, dim=-1)
            t_y = self._rotate_axis(t_y, cos_y, sin_y)
            t_x = self._rotate_axis(t_x, cos_x, sin_x)
            return torch.cat((t_y, t_x, t_pass), dim=-1)

        q_patch = apply_rope(q_patch)
        k_patch = apply_rope(k_patch)
        q = torch.cat((q_front, q_patch), dim=2)
        k = torch.cat((k_front, k_patch), dim=2)
        return q, k


class AxialRoPEAttention(nn.Module):
    fused_attn: bool

    def __init__(self, dim, num_heads=8, attn_head_dim=None, dim_out=None, qkv_bias=False,
                 qk_norm=False, scale_norm=False, proj_bias=True, attn_drop=0.0, proj_drop=0.0,
                 norm_layer=None, grid_size=None, num_front_tokens=0, rope_base=100.0,
                 rope_dim=None, device=None, dtype=None):
        super().__init__()
        dd = {"device": device, "dtype": dtype}
        dim_out = dim_out or dim
        head_dim = attn_head_dim or dim // num_heads
        if attn_head_dim is None:
            assert dim % num_heads == 0
        if qk_norm or scale_norm:
            assert norm_layer is not None
        rope_dim = head_dim if rope_dim is None else rope_dim
        if rope_dim > head_dim:
            raise ValueError(f"`rope_dim`={rope_dim} exceeds head_dim={head_dim}.")
        if rope_dim % 4 != 0:
            raise ValueError("For axial 2D RoPE, `rope_dim` must be divisible by 4.")
        if num_front_tokens < 0:
            raise ValueError("`num_front_tokens` must be non-negative.")
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.attn_dim = num_heads * head_dim
        self.scale = head_dim ** -0.5
        self.fused_attn = use_fused_attn()
        self.grid_size = grid_size
        self.num_front_tokens = num_front_tokens
        self.qkv = nn.Linear(dim, self.attn_dim * 3, bias=qkv_bias, **dd)
        self.q_norm = norm_layer(head_dim, **dd) if qk_norm else nn.Identity()
        self.k_norm = norm_layer(head_dim, **dd) if qk_norm else nn.Identity()
        self.rope = Axial2dRotaryEmbedding(rope_dim, base=rope_base)
        self.attn_drop = nn.Dropout(attn_drop)
        self.norm = norm_layer(self.attn_dim, **dd) if scale_norm else nn.Identity()
        self.proj = nn.Linear(self.attn_dim, dim_out, bias=proj_bias, **dd)
        self.proj_drop = nn.Dropout(proj_drop)

    def set_grid_size(self, grid_size):
        self.grid_size = grid_size

    def _infer_grid_size(self, num_patch_tokens):
        if self.grid_size is not None:
            gh, gw = self.grid_size
            if gh * gw != num_patch_tokens:
                raise ValueError(f"`grid_size={self.grid_size}` implies {gh * gw} patches, got {num_patch_tokens}.")
            return gh, gw
        side = math.isqrt(num_patch_tokens)
        if side * side != num_patch_tokens:
            raise ValueError("Cannot infer a non-square patch grid from the token sequence.")
        return side, side

    def forward(self, x, attn_mask=None, is_causal=False):
        b, n, _ = x.shape
        num_patch_tokens = n - self.num_front_tokens
        if num_patch_tokens <= 0:
            raise ValueError(f"Expected patch tokens after {self.num_front_tokens} front tokens, got N={n}.")
        qkv = self.qkv(x).reshape(b, n, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
        q, k, v = qkv.unbind(0)
        q = self.q_norm(q); k = self.k_norm(k)
        grid_size = self._infer_grid_size(num_patch_tokens)
        q, k = self.rope(q, k, grid_size=grid_size, num_front_tokens=self.num_front_tokens)
        if self.fused_attn:
            x = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask,
                                               dropout_p=self.attn_drop.p if self.training else 0.0, is_causal=is_causal)
        else:
            q = q * self.scale
            attn = q @ k.transpose(-2, -1)
            attn = attn.softmax(dim=-1)
            attn = self.attn_drop(attn)
            x = attn @ v
        x = x.transpose(1, 2).reshape(b, n, self.attn_dim)
        x = self.norm(x)
        x = self.proj(x)
        x = self.proj_drop(x)
        return x


class AxialRoPEBlock(nn.Module):
    def __init__(self, dim, num_heads, mlp_ratio=4.0, qkv_bias=False, qk_norm=False,
                 scale_attn_norm=False, scale_mlp_norm=False, proj_bias=True, proj_drop=0.0,
                 attn_drop=0.0, init_values=None, drop_path=0.0, act_layer=nn.GELU,
                 norm_layer=LayerNorm, mlp_layer=Mlp, attn_layer=None, depth=0,
                 grid_size=None, num_front_tokens=0, rope_base=100.0, rope_dim=None, device=None, dtype=None):
        super().__init__()
        dd = {"device": device, "dtype": dtype}
        self.norm1 = norm_layer(dim, **dd)
        self.attn = AxialRoPEAttention(dim=dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_norm=qk_norm,
                                       scale_norm=scale_attn_norm, proj_bias=proj_bias, attn_drop=attn_drop,
                                       proj_drop=proj_drop, norm_layer=norm_layer, grid_size=grid_size,
                                       num_front_tokens=num_front_tokens, rope_base=rope_base, rope_dim=rope_dim, **dd)
        self.ls1 = LayerScale(dim, init_values=init_values, **dd) if init_values is not None else nn.Identity()
        self.drop_path1 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
        self.norm2 = norm_layer(dim, **dd)
        self.mlp = mlp_layer(in_features=dim, hidden_features=int(dim * mlp_ratio), act_layer=act_layer,
                             norm_layer=norm_layer if scale_mlp_norm else None, bias=proj_bias, drop=proj_drop, **dd)
        self.ls2 = LayerScale(dim, init_values=init_values, **dd) if init_values is not None else nn.Identity()
        self.drop_path2 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()

    def set_grid_size(self, grid_size):
        self.attn.set_grid_size(grid_size)

    def forward(self, x, attn_mask=None, is_causal=False):
        x = x + self.drop_path1(self.ls1(self.attn(self.norm1(x), attn_mask=attn_mask, is_causal=is_causal)))
        x = x + self.drop_path2(self.ls2(self.mlp(self.norm2(x))))
        return x


# =========================================================================== #
#  Vendored cross-attention (no RoPE) (from gazelle.cross_attention)           #
# =========================================================================== #
def _cross_attn_mask(attn_mask, dtype):
    if attn_mask is None:
        return None
    if attn_mask.dtype == torch.bool:
        bias = torch.zeros_like(attn_mask, dtype=dtype)
        bias.masked_fill_(~attn_mask, float("-inf"))
        return bias
    return attn_mask


class CrossAttention(nn.Module):
    fused_attn: bool

    def __init__(self, dim, num_heads=8, attn_head_dim=None, dim_out=None, qkv_bias=False,
                 qk_norm=False, scale_norm=False, proj_bias=True, attn_drop=0.0, proj_drop=0.0, norm_layer=None):
        super().__init__()
        dim_out = dim_out or dim
        if attn_head_dim is None:
            assert dim % num_heads == 0
            head_dim = dim // num_heads
        else:
            head_dim = attn_head_dim
        if qk_norm or scale_norm:
            assert norm_layer is not None
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.attn_dim = num_heads * head_dim
        self.scale = head_dim ** -0.5
        self.fused_attn = use_fused_attn()
        self.q = nn.Linear(dim, self.attn_dim, bias=qkv_bias)
        self.k = nn.Linear(dim, self.attn_dim, bias=qkv_bias)
        self.v = nn.Linear(dim, self.attn_dim, bias=qkv_bias)
        self.q_norm = norm_layer(head_dim) if qk_norm else nn.Identity()
        self.k_norm = norm_layer(head_dim) if qk_norm else nn.Identity()
        self.attn_drop = nn.Dropout(attn_drop)
        self.norm = norm_layer(self.attn_dim) if scale_norm else nn.Identity()
        self.proj = nn.Linear(self.attn_dim, dim_out, bias=proj_bias)
        self.proj_drop = nn.Dropout(proj_drop)

    def forward(self, x_q, x_kv, attn_mask=None):
        B, Nq, _ = x_q.shape
        Bkv, Nk, _ = x_kv.shape
        assert B == Bkv
        q = self.q(x_q).reshape(B, Nq, self.num_heads, self.head_dim).transpose(1, 2)
        k = self.k(x_kv).reshape(B, Nk, self.num_heads, self.head_dim).transpose(1, 2)
        v = self.v(x_kv).reshape(B, Nk, self.num_heads, self.head_dim).transpose(1, 2)
        q = self.q_norm(q); k = self.k_norm(k)
        if self.fused_attn:
            x = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask,
                                               dropout_p=self.attn_drop.p if self.training else 0.0)
        else:
            q = q * self.scale
            attn = q @ k.transpose(-2, -1)
            attn_bias = _cross_attn_mask(attn_mask, attn.dtype)
            if attn_bias is not None:
                attn = attn + attn_bias
            attn = attn.softmax(dim=-1)
            attn = self.attn_drop(attn)
            x = attn @ v
        x = x.transpose(1, 2).reshape(B, Nq, self.attn_dim)
        x = self.norm(x)
        x = self.proj(x)
        x = self.proj_drop(x)
        return x


class CrossAttentionBlock(nn.Module):
    def __init__(self, dim, num_heads, qkv_bias=False, qk_norm=False, scale_attn_norm=False,
                 proj_bias=True, proj_drop=0.0, attn_drop=0.0, init_values=None, drop_path=0.0,
                 norm_layer=LayerNorm):
        super().__init__()
        self.norm_q = norm_layer(dim)
        self.norm_kv = norm_layer(dim)
        self.attn = CrossAttention(dim=dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_norm=qk_norm,
                                   scale_norm=scale_attn_norm, proj_bias=proj_bias, attn_drop=attn_drop,
                                   proj_drop=proj_drop, norm_layer=norm_layer)
        self.ls1 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()

    def forward(self, x_q, x_kv, attn_mask=None):
        x_q = x_q + self.drop_path(self.ls1(self.attn(self.norm_q(x_q), self.norm_kv(x_kv), attn_mask=attn_mask)))
        return x_q


# =========================================================================== #
#  Vendored Axial 2D RoPE cross-attention (from gazelle.rope_cross_attention)  #
# =========================================================================== #
def _infer_square_grid_size(num_patch_tokens, *, name):
    side = math.isqrt(num_patch_tokens)
    if side * side != num_patch_tokens:
        raise ValueError(f"Cannot infer square grid for {name} from {num_patch_tokens} patch tokens.")
    return side, side


def _as_batched_rect(rect, *, batch_size, device, dtype, name):
    rect = torch.as_tensor(rect, device=device, dtype=dtype)
    if rect.ndim == 1:
        if rect.shape[0] != 4:
            raise ValueError(f"{name}_rect must have shape [4] or [B, 4].")
        rect = rect[None, :].expand(batch_size, 4)
    elif rect.ndim == 2:
        if rect.shape[1] != 4:
            raise ValueError(f"{name}_rect must have shape [4] or [B, 4].")
        if rect.shape[0] == 1:
            rect = rect.expand(batch_size, 4)
        elif rect.shape[0] != batch_size:
            raise ValueError(f"{name}_rect has batch size {rect.shape[0]}, expected {batch_size}.")
    else:
        raise ValueError(f"{name}_rect must have shape [4] or [B, 4].")
    return rect


def _native_grid_coords(grid_size, *, batch_size, device, dtype):
    gh, gw = grid_size
    yy, xx = torch.meshgrid(torch.arange(gh, device=device, dtype=dtype),
                            torch.arange(gw, device=device, dtype=dtype), indexing="ij")
    coords = torch.stack((yy.reshape(-1), xx.reshape(-1)), dim=-1)
    return coords[None, :, :].expand(batch_size, -1, -1)


def _rect_grid_coords(grid_size, rect, *, align_corners):
    b = rect.shape[0]; gh, gw = grid_size
    device = rect.device; dtype = rect.dtype
    y0, x0, y1, x1 = rect.unbind(dim=-1)
    if align_corners:
        if gh == 1:
            ys = ((y0 + y1 - 1.0) * 0.5)[:, None]
        else:
            iy = torch.linspace(0.0, 1.0, gh, device=device, dtype=dtype)
            ys = y0[:, None] + iy[None, :] * ((y1 - 1.0) - y0)[:, None]
        if gw == 1:
            xs = ((x0 + x1 - 1.0) * 0.5)[:, None]
        else:
            ix = torch.linspace(0.0, 1.0, gw, device=device, dtype=dtype)
            xs = x0[:, None] + ix[None, :] * ((x1 - 1.0) - x0)[:, None]
    else:
        iy = torch.arange(gh, device=device, dtype=dtype) + 0.5
        ix = torch.arange(gw, device=device, dtype=dtype) + 0.5
        ys = y0[:, None] + iy[None, :] * ((y1 - y0) / gh)[:, None] - 0.5
        xs = x0[:, None] + ix[None, :] * ((x1 - x0) / gw)[:, None] - 0.5
    yy = ys[:, :, None].expand(b, gh, gw)
    xx = xs[:, None, :].expand(b, gh, gw)
    return torch.stack((yy.reshape(b, -1), xx.reshape(b, -1)), dim=-1)


def _as_batched_patch_coords(patch_coords, *, batch_size, num_patch_tokens, device, dtype, name):
    patch_coords = torch.as_tensor(patch_coords, device=device, dtype=dtype)
    if patch_coords.ndim == 2:
        if patch_coords.shape != (num_patch_tokens, 2):
            raise ValueError(f"{name}_patch_coords must have shape [{num_patch_tokens}, 2] or [B, {num_patch_tokens}, 2], got {tuple(patch_coords.shape)}.")
        patch_coords = patch_coords[None, :, :].expand(batch_size, -1, -1)
    elif patch_coords.ndim == 3:
        if patch_coords.shape[1:] != (num_patch_tokens, 2):
            raise ValueError(f"{name}_patch_coords must have shape [B, {num_patch_tokens}, 2], got {tuple(patch_coords.shape)}.")
        if patch_coords.shape[0] == 1:
            patch_coords = patch_coords.expand(batch_size, -1, -1)
        elif patch_coords.shape[0] != batch_size:
            raise ValueError(f"{name}_patch_coords has batch size {patch_coords.shape[0]}, expected {batch_size}.")
    else:
        raise ValueError(f"{name}_patch_coords must have shape [N_patch, 2] or [B, N_patch, 2].")
    return patch_coords


def make_stream_patch_coords(*, batch_size, num_patch_tokens, grid_size, rect, patch_coords,
                             device, dtype=torch.float32, align_corners=False, name):
    if grid_size is None:
        grid_size = _infer_square_grid_size(num_patch_tokens, name=name)
    gh, gw = grid_size
    expected = gh * gw
    if expected != num_patch_tokens:
        raise ValueError(f"{name}_grid_size={grid_size} implies {expected} patch tokens, but {name} stream has {num_patch_tokens}.")
    if patch_coords is not None and rect is not None:
        raise ValueError(f"Provide either {name}_patch_coords or {name}_rect, not both.")
    if patch_coords is not None:
        return _as_batched_patch_coords(patch_coords, batch_size=batch_size, num_patch_tokens=num_patch_tokens,
                                        device=device, dtype=dtype, name=name)
    if rect is not None:
        rect = _as_batched_rect(rect, batch_size=batch_size, device=device, dtype=dtype, name=name)
        return _rect_grid_coords(grid_size, rect, align_corners=align_corners)
    return _native_grid_coords(grid_size, batch_size=batch_size, device=device, dtype=dtype)


class Axial2dCrossRotaryEmbedding(nn.Module):
    def __init__(self, dim, base=100.0):
        super().__init__()
        if dim <= 0 or dim % 4 != 0:
            raise ValueError(f"dim must be a positive multiple of 4, got {dim}.")
        self.dim = dim
        self.axis_dim = dim // 2
        self.base = base
        self.register_buffer("inv_freq", self._compute_inv_freq(), persistent=False)

    def _compute_inv_freq(self):
        return 1.0 / (self.base ** (torch.arange(0, self.axis_dim, 2, dtype=torch.float32) / self.axis_dim))

    def reset_inv_freq(self):
        self.inv_freq = self._compute_inv_freq()

    def _axis_cos_sin(self, coords, *, out_dtype):
        if coords.ndim != 2:
            raise ValueError(f"coords must have shape [B, N], got {tuple(coords.shape)}.")
        coords = coords.to(dtype=torch.float32)
        inv_freq = self.inv_freq.to(device=coords.device, dtype=torch.float32)
        freqs = coords[..., None] * inv_freq[None, None, :]
        return freqs.cos().to(dtype=out_dtype), freqs.sin().to(dtype=out_dtype)

    @staticmethod
    def _rotate_axis(x, cos, sin):
        x = x.reshape(*x.shape[:-1], -1, 2)
        x_even, x_odd = x.unbind(dim=-1)
        cos = cos[:, None, :, :]
        sin = sin[:, None, :, :]
        x_rot = torch.stack((x_even * cos - x_odd * sin, x_even * sin + x_odd * cos), dim=-1)
        return x_rot.flatten(-2)

    def rotate_one(self, x, coords_yx, *, num_front_tokens, stream_name):
        b, _, n_total, head_dim = x.shape
        if self.dim > head_dim:
            raise ValueError(f"RoPE dim {self.dim} exceeds head_dim {head_dim} for {stream_name}.")
        if num_front_tokens < 0:
            raise ValueError(f"{stream_name}_num_front_tokens must be non-negative.")
        n_patch = n_total - num_front_tokens
        if n_patch <= 0:
            raise ValueError(f"{stream_name} has no patch tokens after {num_front_tokens} front tokens.")
        if coords_yx.shape != (b, n_patch, 2):
            raise ValueError(f"{stream_name}_coords_yx must have shape [{b}, {n_patch}, 2], got {tuple(coords_yx.shape)}.")
        coords_yx = coords_yx.to(device=x.device)
        x_front = x[:, :, :num_front_tokens, :]
        x_patch = x[:, :, num_front_tokens:, :]
        y = coords_yx[..., 0]
        x_coord = coords_yx[..., 1]
        cos_y, sin_y = self._axis_cos_sin(y, out_dtype=x.dtype)
        cos_x, sin_x = self._axis_cos_sin(x_coord, out_dtype=x.dtype)
        x_rope = x_patch[..., :self.dim]
        x_pass = x_patch[..., self.dim:]
        x_y, x_x = x_rope.split(self.axis_dim, dim=-1)
        x_y = self._rotate_axis(x_y, cos_y, sin_y)
        x_x = self._rotate_axis(x_x, cos_x, sin_x)
        x_patch = torch.cat((x_y, x_x, x_pass), dim=-1)
        if num_front_tokens == 0:
            return x_patch
        return torch.cat((x_front, x_patch), dim=2)

    def forward(self, q, k, *, q_coords_yx, kv_coords_yx, q_num_front_tokens=0, kv_num_front_tokens=0):
        q = self.rotate_one(q, q_coords_yx, num_front_tokens=q_num_front_tokens, stream_name="q")
        k = self.rotate_one(k, kv_coords_yx, num_front_tokens=kv_num_front_tokens, stream_name="kv")
        return q, k


class AxialRoPECrossAttention(nn.Module):
    fused_attn: bool

    def __init__(self, dim, num_heads=8, attn_head_dim=None, dim_out=None, qkv_bias=False,
                 qk_norm=False, scale_norm=False, proj_bias=True, attn_drop=0.0, proj_drop=0.0,
                 norm_layer=None, q_num_front_tokens=0, kv_num_front_tokens=0, rope_base=100.0,
                 rope_dim=None, align_corners=False, device=None, dtype=None):
        super().__init__()
        dd = {"device": device, "dtype": dtype}
        dim_out = dim_out or dim
        if attn_head_dim is None:
            assert dim % num_heads == 0
            head_dim = dim // num_heads
        else:
            head_dim = attn_head_dim
        if qk_norm or scale_norm:
            assert norm_layer is not None
        rope_dim = head_dim if rope_dim is None else rope_dim
        if rope_dim > head_dim:
            raise ValueError(f"rope_dim={rope_dim} exceeds head_dim={head_dim}.")
        if rope_dim % 4 != 0:
            raise ValueError("For axial 2D RoPE, `rope_dim` must be divisible by 4.")
        if q_num_front_tokens < 0:
            raise ValueError("q_num_front_tokens must be non-negative.")
        if kv_num_front_tokens < 0:
            raise ValueError("kv_num_front_tokens must be non-negative.")
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.attn_dim = num_heads * head_dim
        self.scale = head_dim ** -0.5
        self.fused_attn = use_fused_attn()
        self.q_num_front_tokens = q_num_front_tokens
        self.kv_num_front_tokens = kv_num_front_tokens
        self.align_corners = align_corners
        self.q = nn.Linear(dim, self.attn_dim, bias=qkv_bias, **dd)
        self.k = nn.Linear(dim, self.attn_dim, bias=qkv_bias, **dd)
        self.v = nn.Linear(dim, self.attn_dim, bias=qkv_bias, **dd)
        self.q_norm = norm_layer(head_dim, **dd) if qk_norm else nn.Identity()
        self.k_norm = norm_layer(head_dim, **dd) if qk_norm else nn.Identity()
        self.rope = Axial2dCrossRotaryEmbedding(dim=rope_dim, base=rope_base)
        self.attn_drop = nn.Dropout(attn_drop)
        self.norm = norm_layer(self.attn_dim, **dd) if scale_norm else nn.Identity()
        self.proj = nn.Linear(self.attn_dim, dim_out, bias=proj_bias, **dd)
        self.proj_drop = nn.Dropout(proj_drop)

    def forward(self, x_q, x_kv, attn_mask=None, *, q_grid_size=None, kv_grid_size=None,
                q_rect=None, kv_rect=None, q_patch_coords=None, kv_patch_coords=None,
                q_num_front_tokens=None, kv_num_front_tokens=None):
        b, nq, _ = x_q.shape
        b_kv, nk, _ = x_kv.shape
        if b != b_kv:
            raise ValueError(f"x_q and x_kv must have the same batch size, got {b} and {b_kv}.")
        q_num_front_tokens = self.q_num_front_tokens if q_num_front_tokens is None else q_num_front_tokens
        kv_num_front_tokens = self.kv_num_front_tokens if kv_num_front_tokens is None else kv_num_front_tokens
        q_num_patch_tokens = nq - q_num_front_tokens
        kv_num_patch_tokens = nk - kv_num_front_tokens
        if q_num_patch_tokens <= 0:
            raise ValueError(f"x_q has no patch tokens after {q_num_front_tokens} front tokens.")
        if kv_num_patch_tokens <= 0:
            raise ValueError(f"x_kv has no patch tokens after {kv_num_front_tokens} front tokens.")
        coord_dtype = torch.float32
        q_coords_yx = make_stream_patch_coords(batch_size=b, num_patch_tokens=q_num_patch_tokens, grid_size=q_grid_size,
                                               rect=q_rect, patch_coords=q_patch_coords, device=x_q.device,
                                               dtype=coord_dtype, align_corners=self.align_corners, name="q")
        kv_coords_yx = make_stream_patch_coords(batch_size=b, num_patch_tokens=kv_num_patch_tokens, grid_size=kv_grid_size,
                                                rect=kv_rect, patch_coords=kv_patch_coords, device=x_kv.device,
                                                dtype=coord_dtype, align_corners=self.align_corners, name="kv")
        q = self.q(x_q).reshape(b, nq, self.num_heads, self.head_dim).transpose(1, 2)
        k = self.k(x_kv).reshape(b, nk, self.num_heads, self.head_dim).transpose(1, 2)
        v = self.v(x_kv).reshape(b, nk, self.num_heads, self.head_dim).transpose(1, 2)
        q = self.q_norm(q); k = self.k_norm(k)
        q, k = self.rope(q, k, q_coords_yx=q_coords_yx, kv_coords_yx=kv_coords_yx,
                         q_num_front_tokens=q_num_front_tokens, kv_num_front_tokens=kv_num_front_tokens)
        if self.fused_attn:
            x = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask,
                                               dropout_p=self.attn_drop.p if self.training else 0.0)
        else:
            q = q * self.scale
            attn = q @ k.transpose(-2, -1)
            attn_bias = _cross_attn_mask(attn_mask, attn.dtype)
            if attn_bias is not None:
                attn = attn + attn_bias
            attn = attn.softmax(dim=-1)
            attn = self.attn_drop(attn)
            x = attn @ v
        x = x.transpose(1, 2).reshape(b, nq, self.attn_dim)
        x = self.norm(x)
        x = self.proj(x)
        x = self.proj_drop(x)
        return x


class AxialRoPECrossAttentionBlock(nn.Module):
    def __init__(self, dim, num_heads, qkv_bias=False, qk_norm=False, scale_attn_norm=False,
                 proj_bias=True, proj_drop=0.0, attn_drop=0.0, init_values=None, drop_path=0.0,
                 norm_layer=LayerNorm, q_num_front_tokens=0, kv_num_front_tokens=0, rope_base=100.0,
                 rope_dim=None, align_corners=False, device=None, dtype=None):
        super().__init__()
        dd = {"device": device, "dtype": dtype}
        self.norm_q = norm_layer(dim, **dd)
        self.norm_kv = norm_layer(dim, **dd)
        self.attn = AxialRoPECrossAttention(dim=dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_norm=qk_norm,
                                            scale_norm=scale_attn_norm, proj_bias=proj_bias, attn_drop=attn_drop,
                                            proj_drop=proj_drop, norm_layer=norm_layer, q_num_front_tokens=q_num_front_tokens,
                                            kv_num_front_tokens=kv_num_front_tokens, rope_base=rope_base, rope_dim=rope_dim,
                                            align_corners=align_corners, **dd)
        self.ls1 = LayerScale(dim, init_values=init_values, **dd) if init_values is not None else nn.Identity()
        self.drop_path1 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()

    def forward(self, x_q, x_kv, attn_mask=None, *, q_grid_size=None, kv_grid_size=None,
                q_rect=None, kv_rect=None, q_patch_coords=None, kv_patch_coords=None,
                q_num_front_tokens=None, kv_num_front_tokens=None):
        x_q = x_q + self.drop_path1(self.ls1(self.attn(
            self.norm_q(x_q), self.norm_kv(x_kv), attn_mask=attn_mask, q_grid_size=q_grid_size,
            kv_grid_size=kv_grid_size, q_rect=q_rect, kv_rect=kv_rect, q_patch_coords=q_patch_coords,
            kv_patch_coords=kv_patch_coords, q_num_front_tokens=q_num_front_tokens,
            kv_num_front_tokens=kv_num_front_tokens)))
        return x_q


# =========================================================================== #
#  Block factories + SceneHeadInteraction (from gazelle.cross_attention_model) #
# =========================================================================== #
def get_vit_block(dim=256, num_heads=8, mlp_ratio=4, mlp_layer=SwiGLU, drop_path=0.1,
                  act_layer=nn.GELU, pos_encoding="rope", num_front_tokens=4, rope_base=100.0):
    if pos_encoding not in {"rope", "sinusoidal", "ape"}:
        raise ValueError(f"pos_encoding must be one of: rope, sinusoidal, ape, got {pos_encoding}")
    if pos_encoding == "rope":
        return AxialRoPEBlock(dim=dim, num_heads=num_heads, mlp_ratio=mlp_ratio, mlp_layer=mlp_layer,
                              drop_path=drop_path, act_layer=act_layer, num_front_tokens=num_front_tokens, rope_base=rope_base)
    else:
        return Block(dim=dim, num_heads=num_heads, mlp_ratio=mlp_ratio, mlp_layer=mlp_layer,
                     drop_path=drop_path, act_layer=act_layer)


def get_cross_attn_block(dim=256, num_heads=8, drop_path=0.1, pos_encoding="rope",
                         q_num_front_tokens=0, kv_num_front_tokens=0, rope_base=100.0):
    if pos_encoding not in {"rope", "sinusoidal", "ape"}:
        raise ValueError(f"pos_encoding must be one of: rope, sinusoidal, ape, got {pos_encoding}")
    if pos_encoding == "rope":
        return AxialRoPECrossAttentionBlock(dim=dim, num_heads=num_heads, drop_path=drop_path,
                                            q_num_front_tokens=q_num_front_tokens, kv_num_front_tokens=kv_num_front_tokens,
                                            rope_base=rope_base)
    else:
        return CrossAttentionBlock(dim=dim, num_heads=num_heads, drop_path=drop_path)


class SceneHeadInteraction(nn.Module):
    """Variant A2: synchronous and symmetric feature interaction."""

    def __init__(self, dim, num_heads=8, mlp_ratio=4, mlp_layer=SwiGLU, act_layer=nn.GELU,
                 drop_path=0.0, num_front_tokens=0, pos_encoding="rope"):
        super().__init__()
        if pos_encoding not in {"rope", "sinusoidal", "ape"}:
            raise ValueError(f"pos_encoding must be one of: rope, sinusoidal, ape, got {pos_encoding}")
        self.pos_encoding = pos_encoding
        self.cross_attn_scene = get_cross_attn_block(dim=dim, num_heads=num_heads, pos_encoding=pos_encoding,
                                                      q_num_front_tokens=num_front_tokens, kv_num_front_tokens=num_front_tokens, drop_path=drop_path)
        self.vit_block_scene = get_vit_block(dim=dim, num_heads=num_heads, mlp_ratio=mlp_ratio, mlp_layer=mlp_layer,
                                             drop_path=drop_path, act_layer=act_layer, num_front_tokens=num_front_tokens, pos_encoding=pos_encoding)
        self.cross_attn_head = get_cross_attn_block(dim=dim, num_heads=num_heads, pos_encoding=pos_encoding,
                                                    q_num_front_tokens=num_front_tokens, kv_num_front_tokens=num_front_tokens, drop_path=drop_path)
        self.vit_block_head = get_vit_block(dim=dim, num_heads=num_heads, mlp_ratio=mlp_ratio, mlp_layer=mlp_layer,
                                            drop_path=drop_path, act_layer=act_layer, num_front_tokens=num_front_tokens, pos_encoding=pos_encoding)

    def forward(self, tokens):
        scene_tokens = tokens["scene_tokens"]
        head_tokens = tokens["head_tokens"]
        head_rects = tokens["head_rects"]
        if self.pos_encoding == "rope":
            out_scene_tokens = self.cross_attn_scene(scene_tokens, head_tokens, kv_rect=head_rects)
            out_head_tokens = self.cross_attn_head(head_tokens, scene_tokens, q_rect=head_rects)
        else:
            out_scene_tokens = self.cross_attn_scene(scene_tokens, head_tokens)
            out_head_tokens = self.cross_attn_head(head_tokens, scene_tokens)
        out_scene_tokens = self.vit_block_scene(out_scene_tokens)
        out_head_tokens = self.vit_block_head(out_head_tokens)
        return {"scene_tokens": out_scene_tokens, "head_tokens": out_head_tokens, "head_rects": head_rects}


# =========================================================================== #
#  DINOv3 backbone wrapper (config-only construction, weights from safetensors) #
# =========================================================================== #
class PaGEBackbone(nn.Module):
    """
    Wraps a transformers built-in DINOv3ViTModel. Output: patch tokens -> [B, C, H', W'].
    The DINOv3 model is built from config only (no external download).
    """

    def __init__(self, dinov3_config: DINOv3ViTConfig, in_size=(512, 512)):
        super().__init__()
        self.in_size = in_size
        self.model = DINOv3ViTModel(dinov3_config)
        self.patch_size = int(dinov3_config.patch_size)
        self.embed_dim = int(dinov3_config.hidden_size)
        # CLS(1) + num_register_tokens
        self._num_front = 1 + int(getattr(dinov3_config, "num_register_tokens", 0) or 0)
        # DINOv3ViTModel's internal naming differs across transformers versions:
        #   4.56.x  -> layer stack flattened:  model.layer.N.*        (model.<dinov3 top-level>)
        #   5.6.x   -> layer stack nested:     model.model.layer.N.*  (extra `.model` wrapper)
        # A single safetensors file must load under both, so remap incoming keys at load time.
        self._register_load_state_dict_pre_hook(self._remap_dinov3_keys)

    @staticmethod
    def _dinov3_has_nested_layer(dinov3_module) -> bool:
        """True if this transformers version nests the layer stack under an inner `.model`
        (transformers >= 5.x). In 4.56.x the layers are flattened onto the DINOv3ViTModel itself."""
        inner = getattr(dinov3_module, "model", None)
        if not isinstance(inner, nn.Module):
            return False
        # inner's own keys are relative to it: "layer.0.*" when nested, never "embeddings.*".
        return any(k.startswith("layer.") for k in inner.state_dict().keys())

    def _remap_dinov3_keys(self, state_dict, prefix, *args, **kwargs):
        """Normalize DINOv3 backbone keys (embeddings / layer / norm / rope_embeddings)
        from whichever convention the checkpoint uses into the one this transformers
        version expects."""
        nested = self._dinov3_has_nested_layer(self.model)
        model_pref = prefix + "model."
        new = {}
        for k in list(state_dict.keys()):
            if not k.startswith(model_pref):
                continue
            rest = k[len(model_pref):]  # after "<prefix>model."
            # rest is one of: "embeddings...", "norm...", "rope_embeddings...", "layer...",
            #                  or "model.layer..." (nested-conv checkpoint under a flat version, etc.)
            if rest.startswith("model.layer."):
                core = rest[len("model."):]  # -> "layer..."
            else:
                core = rest  # "layer..." / "embeddings..." / "norm..." / "rope_embeddings..."
            if nested and core.startswith("layer."):
                target = model_pref + "model." + core
            else:
                target = model_pref + core
            if target != k:
                new[target] = state_dict.pop(k)
        state_dict.update(new)

    def _get_patch_tokens(self, x: torch.Tensor) -> torch.Tensor:
        out = self.model(pixel_values=x, return_dict=True)
        tokens = getattr(out, "last_hidden_state", None)
        if tokens is None:
            tokens = out[0]
        if not torch.is_tensor(tokens) or tokens.dim() != 3:
            raise RuntimeError("Unexpected DINOv3 output format.")
        tokens = tokens[:, self._num_front:, :]  # drop CLS + register tokens
        return tokens

    def forward(self, x) -> torch.Tensor:
        if isinstance(x, (list, tuple)):
            # head stream comes in as a 1-element list (one backbone branch); unwrap it
            assert len(x) == 1
            x = x[0]
        b, c, h, w = x.shape
        out_h, out_w = self.get_out_size((h, w))
        patch_tokens = self._get_patch_tokens(x)
        if patch_tokens.shape[1] != out_h * out_w:
            raise RuntimeError(
                f"[PaGEBackbone] token count mismatch: {patch_tokens.shape[1]} vs {out_h * out_w}. "
                f"patch_size={self.patch_size}, input={(h, w)}")
        feat = patch_tokens.view(b, out_h, out_w, -1).permute(0, 3, 1, 2).contiguous()
        return feat

    def get_dimension(self):
        return self.embed_dim

    def get_out_size(self, in_size):
        h, w = in_size
        return (h // self.patch_size, w // self.patch_size)


# =========================================================================== #
#  PaGE config                                                                #
# =========================================================================== #
class PaGEConfig(PretrainedConfig):
    model_type = "page"

    def __init__(
        self,
        # gaze decoder
        dim: int = 256,
        num_heads: int = 8,
        mlp_ratio: float = 4.0,
        mlp_layer: str = "geglu",
        pos_encoding: str = "rope",
        n_scene_self_attn_layers: int = 1,
        n_head_self_attn_layers: int = 1,
        n_scene_head_interaction_layers: int = 5,
        n_reg_tokens: int = 4,
        heatmap_out_size: Tuple[int, int] = (64, 64),
        dino_feature_dropout: float = 0.1,
        drop_path: float = 0.1,
        use_head_prompt: bool = False,
        inout: bool = True,
        # input sizes
        scene_in_size: Tuple[int, int] = (512, 512),
        head_in_size: Tuple[int, int] = (256, 256),
        # image preprocessing
        image_mean: Tuple[float, float, float] = (0.485, 0.456, 0.406),
        image_std: Tuple[float, float, float] = (0.229, 0.224, 0.225),
        # DINOv3 backbone config (shared by scene & head branches)
        dinov3_hidden_size: int = 768,
        dinov3_num_hidden_layers: int = 12,
        dinov3_num_attention_heads: int = 12,
        dinov3_intermediate_size: int = 3072,
        dinov3_num_register_tokens: int = 4,
        dinov3_patch_size: int = 16,
        dinov3_use_gated_mlp: bool = False,
        dinov3_layerscale_value: float = 1.0,
        dinov3_drop_path_rate: float = 0.0,
        dinov3_layer_norm_eps: float = 1e-5,
        **kwargs,
    ):
        self.dim = dim
        self.num_heads = num_heads
        self.mlp_ratio = mlp_ratio
        self.mlp_layer = mlp_layer
        self.pos_encoding = pos_encoding
        self.n_scene_self_attn_layers = n_scene_self_attn_layers
        self.n_head_self_attn_layers = n_head_self_attn_layers
        self.n_scene_head_interaction_layers = n_scene_head_interaction_layers
        self.n_reg_tokens = n_reg_tokens
        self.heatmap_out_size = tuple(heatmap_out_size)
        self.dino_feature_dropout = dino_feature_dropout
        self.drop_path = drop_path
        self.use_head_prompt = use_head_prompt
        self.inout = inout
        self.scene_in_size = tuple(scene_in_size)
        self.head_in_size = tuple(head_in_size)
        self.image_mean = tuple(image_mean)
        self.image_std = tuple(image_std)
        self.dinov3_hidden_size = dinov3_hidden_size
        self.dinov3_num_hidden_layers = dinov3_num_hidden_layers
        self.dinov3_num_attention_heads = dinov3_num_attention_heads
        self.dinov3_intermediate_size = dinov3_intermediate_size
        self.dinov3_num_register_tokens = dinov3_num_register_tokens
        self.dinov3_patch_size = dinov3_patch_size
        self.dinov3_use_gated_mlp = dinov3_use_gated_mlp
        self.dinov3_layerscale_value = dinov3_layerscale_value
        self.dinov3_drop_path_rate = dinov3_drop_path_rate
        self.dinov3_layer_norm_eps = dinov3_layer_norm_eps
        super().__init__(**kwargs)

    def to_dinov3_config(self) -> DINOv3ViTConfig:
        return DINOv3ViTConfig(
            hidden_size=self.dinov3_hidden_size,
            num_hidden_layers=self.dinov3_num_hidden_layers,
            num_attention_heads=self.dinov3_num_attention_heads,
            intermediate_size=self.dinov3_intermediate_size,
            num_register_tokens=self.dinov3_num_register_tokens,
            patch_size=self.dinov3_patch_size,
            use_gated_mlp=self.dinov3_use_gated_mlp,
            layerscale_value=self.dinov3_layerscale_value,
            drop_path_rate=self.dinov3_drop_path_rate,
            layer_norm_eps=self.dinov3_layer_norm_eps,
            image_size=self.scene_in_size[0],
        )


# =========================================================================== #
#  PaGE model (CrossGaze architecture)                                        #
# =========================================================================== #
class PaGEPreTrainedModel(PreTrainedModel):
    config_class = PaGEConfig
    base_model_prefix = "page"
    supports_gradient_checkpointing = False

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            nn.init.trunc_normal_(module.weight, std=0.02)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Conv2d):
            nn.init.trunc_normal_(module.weight, std=0.02)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.trunc_normal_(module.weight, std=0.02)
        elif isinstance(module, nn.LayerNorm):
            nn.init.zeros_(module.bias)
            nn.init.ones_(module.weight)

    # ------------------------------------------------------------------ #
    #  Version-safe loading                                              #
    # ------------------------------------------------------------------ #
    # The DINOv3 backbones are built from transformers' built-in DINOv3ViTModel, whose
    # internal parameter naming changed between transformers 4.56.x (layers flattened:
    # `model.layer.N`) and 5.x (layers nested: `model.model.layer.N`). The checkpoints store
    # one convention; loading under the other leaves the backbone randomly initialized.
    # `from_pretrained` in transformers >= 5 bypasses `nn.Module._load_state_dict_pre_hook`,
    # so the remap hook on PaGEBackbone is not invoked by it. We therefore reload the backbone
    # weights ourselves through `nn.Module.load_state_dict` (which *does* fire the hook).
    @staticmethod
    def _collect_safetensors(path_or_repo, **kwargs):
        """Return the full state_dict from a local dir or a HF repo id."""
        import os as _os
        import glob as _glob
        from safetensors.torch import load_file as _load_file
        state = {}
        if _os.path.isdir(path_or_repo):
            index = _os.path.join(path_or_repo, "model.safetensors.index.json")
            if _os.path.isfile(index):
                import json as _json
                wm = _json.load(open(index))["weight_map"]
                files = sorted(set(wm.values()))
            else:
                files = ["model.safetensors"]
            for f in files:
                state.update(_load_file(_os.path.join(path_or_repo, f)))
        else:
            from huggingface_hub import hf_hub_download
            import json as _json
            repo_id = path_or_repo
            try:
                idx_path = hf_hub_download(repo_id=repo_id, filename="model.safetensors.index.json")
                wm = _json.load(open(idx_path))["weight_map"]
                files = sorted(set(wm.values()))
            except Exception:
                files = ["model.safetensors"]
            for f in files:
                p = hf_hub_download(repo_id=repo_id, filename=f)
                state.update(_load_file(p))
        return state

    @classmethod
    def from_pretrained(cls, *args, **kwargs):
        model = super().from_pretrained(*args, **kwargs)
        # Reload DINOv3 backbone weights version-safely (the remap hook fires here).
        try:
            path_or_repo = args[0] if args else kwargs.get("pretrained_model_name_or_path")
            full_state = cls._collect_safetensors(path_or_repo)
            for branch in ("scene_branch_backbone", "head_branch_backbone"):
                if not hasattr(model, branch):
                    continue
                bb = getattr(model, branch)
                bb_state = {k[len(branch) + 1:]: v for k, v in full_state.items()
                            if k.startswith(branch + ".")}
                if bb_state:
                    bb.load_state_dict(bb_state, strict=False)
            # Recompute RoPE inv_freq buffers: they are non-persistent and transformers' meta-init
            # during from_pretrained leaves them as garbage, which would corrupt axial RoPE.
            for m in model.modules():
                if hasattr(m, "reset_inv_freq") and callable(m.reset_inv_freq):
                    m.reset_inv_freq()
        except Exception as e:  # pragma: no cover
            import warnings
            warnings.warn(f"PaGE: version-safe backbone reload skipped ({e!r}). "
                          "Backbone weights may be random if your transformers version mismatches the checkpoint.")
        return model


class PaGEModel(PaGEPreTrainedModel):
    """
    PaGE gaze target estimation model with ViT-adapter-style cross attention between scene and head features.

    Inputs (dict):
        - "images":      scene image tensor [B, 3, H_scene, W_scene]
        - "head_images": list of head-crop tensors, one tensor [sum(Np), 3, H_head, W_head]
        - "bboxes":      list (len B) of lists of bboxes; each bbox is (xmin, ymin, xmax, ymax) in [0,1] image coords

    Outputs (dict):
        - "heatmap": list (len B) of [Np, H_out, W_out] heatmaps (sigmoid applied)
        - "inout":  list (len B) of [Np] in/out scores (sigmoid applied) if inout else None
    """

    def __init__(self, config: PaGEConfig):
        super().__init__(config)
        cfg = config
        dinov3_cfg = cfg.to_dinov3_config()

        self.scene_branch_backbone = PaGEBackbone(dinov3_cfg, in_size=cfg.scene_in_size)
        self.head_branch_backbone = PaGEBackbone(dinov3_cfg, in_size=cfg.head_in_size)

        self.dim = cfg.dim
        self.n_scene_self_attn_layers = cfg.n_scene_self_attn_layers
        self.n_head_self_attn_layers = cfg.n_head_self_attn_layers
        self.n_scene_head_interaction_layers = cfg.n_scene_head_interaction_layers
        self.scene_featmap_h, self.scene_featmap_w = self.scene_branch_backbone.get_out_size(cfg.scene_in_size)
        self.head_featmap_h, self.head_featmap_w = self.head_branch_backbone.get_out_size(cfg.head_in_size)
        self.n_reg_tokens = cfg.n_reg_tokens
        self.n_front_tokens = cfg.n_reg_tokens + 1 if cfg.inout else cfg.n_reg_tokens
        self.heatmap_out_size = tuple(cfg.heatmap_out_size)
        self.inout = cfg.inout
        self.pos_encoding = cfg.pos_encoding
        self.use_head_prompt = cfg.use_head_prompt

        self.scene_proj = nn.Sequential(
            nn.Dropout2d(cfg.dino_feature_dropout),
            nn.Conv2d(self.scene_branch_backbone.get_dimension(), self.dim, 1),
        )
        self.head_proj = nn.Sequential(
            nn.Dropout2d(cfg.dino_feature_dropout),
            nn.Conv2d(self.head_branch_backbone.get_dimension(), self.dim, 1),
        )
        if self.use_head_prompt:
            self.head_position_token = nn.Embedding(1, self.dim)

        if self.pos_encoding == "ape":
            self.scene_seq_len = self.n_reg_tokens + self.scene_featmap_h * self.scene_featmap_w
            self.head_seq_len = self.n_reg_tokens + self.head_featmap_h * self.head_featmap_w
            self.scene_ape = nn.Parameter(torch.zeros((1, self.scene_seq_len, self.dim)))
            self.head_ape = nn.Parameter(torch.zeros((1, self.head_seq_len, self.dim)))
        elif self.pos_encoding == "sinusoidal":
            self.register_buffer("scene_pos_embed", positionalencoding2d(self.dim, self.scene_featmap_h, self.scene_featmap_w).squeeze(0).squeeze(0))
            self.register_buffer("head_pos_embed", positionalencoding2d(self.dim, self.head_featmap_h, self.head_featmap_w).squeeze(0).squeeze(0))

        if self.inout:
            self.scene_inout_token = nn.Parameter(torch.zeros((1, 1, self.dim)))
            self.head_inout_token = nn.Parameter(torch.zeros((1, 1, self.dim)))

        if self.n_reg_tokens > 0:
            self.scene_register_tokens = nn.Parameter(torch.zeros((1, self.n_reg_tokens, self.dim)))
            self.head_register_tokens = nn.Parameter(torch.zeros((1, self.n_reg_tokens, self.dim)))

        if cfg.mlp_layer == "mlp":
            mlp_layer = Mlp; act_layer = nn.GELU
        elif cfg.mlp_layer == "geglu":
            mlp_layer = SwiGLU; act_layer = nn.GELU
        elif cfg.mlp_layer == "swiglu":
            mlp_layer = SwiGLU; act_layer = nn.SiLU
        else:
            raise ValueError(f"mlp_layer must be mlp/geglu/swiglu, got {cfg.mlp_layer}")

        self.scene_self_attn_layers = nn.Sequential(*[
            get_vit_block(dim=self.dim, num_heads=cfg.num_heads, mlp_ratio=cfg.mlp_ratio, mlp_layer=mlp_layer,
                          drop_path=cfg.drop_path, act_layer=act_layer, num_front_tokens=self.n_front_tokens, pos_encoding=cfg.pos_encoding)
            for _ in range(cfg.n_scene_self_attn_layers)
        ]) if cfg.n_scene_self_attn_layers > 0 else nn.Identity()
        self.head_self_attn_layers = nn.Sequential(*[
            get_vit_block(dim=self.dim, num_heads=cfg.num_heads, mlp_ratio=cfg.mlp_ratio, mlp_layer=mlp_layer,
                          drop_path=cfg.drop_path, act_layer=act_layer, num_front_tokens=self.n_front_tokens, pos_encoding=cfg.pos_encoding)
            for _ in range(cfg.n_head_self_attn_layers)
        ]) if cfg.n_head_self_attn_layers > 0 else nn.Identity()
        self.scene_head_interaction_layers = nn.Sequential(*[
            SceneHeadInteraction(dim=self.dim, num_heads=cfg.num_heads, mlp_ratio=cfg.mlp_ratio, mlp_layer=mlp_layer,
                                 drop_path=cfg.drop_path, act_layer=act_layer, pos_encoding=cfg.pos_encoding, num_front_tokens=self.n_front_tokens)
            for _ in range(cfg.n_scene_head_interaction_layers)
        ])

        self.heatmap_head = nn.Sequential(
            nn.ConvTranspose2d(self.dim, self.dim, kernel_size=2, stride=2),
            nn.Conv2d(self.dim, 1, kernel_size=1, bias=False),
        )
        if self.inout:
            self.inout_head = nn.Sequential(
                nn.Linear(self.dim * 2, 128),
                nn.GELU(),
                nn.Dropout(0.1),
                nn.Linear(128, 1),
            )

        self.post_init()

    # ------------------------------------------------------------------ #
    def get_input_head_maps(self, bboxes):
        head_maps = []
        head_rects = []
        for bbox_list in bboxes:
            img_head_maps = []
            img_head_rects = []
            for bbox in bbox_list:
                if bbox is None:
                    img_head_maps.append(torch.zeros(self.scene_featmap_h, self.scene_featmap_w))
                else:
                    xmin, ymin, xmax, ymax = bbox
                    width, height = self.scene_featmap_w, self.scene_featmap_h
                    xmin = round(xmin * width); ymin = round(ymin * height)
                    xmax = round(xmax * width); ymax = round(ymax * height)
                    head_map = torch.zeros((height, width))
                    head_map[ymin:ymax, xmin:xmax] = 1
                    img_head_maps.append(head_map)
                    img_head_rects.append(torch.Tensor([ymin, xmin, ymax, xmax]))
            head_maps.append(torch.stack(img_head_maps))
            head_rects.append(torch.stack(img_head_rects))
        return head_maps, head_rects

    def get_logits(self, input, return_tokens=False):
        num_ppl_per_img = [len(bbox_list) for bbox_list in input["bboxes"]]
        for head_stream_images in input["head_images"]:
            if sum(num_ppl_per_img) != len(head_stream_images):
                raise ValueError(f"bboxes and head crops mismatch: {sum(num_ppl_per_img)} bboxes vs {len(head_stream_images)} head crops.")

        scene_featmap = self.scene_branch_backbone(input["images"])
        scene_featmap = self.scene_proj(scene_featmap)
        scene_dino_tokens = scene_featmap.flatten(start_dim=2).permute(0, 2, 1)
        if self.pos_encoding == "sinusoidal":
            scene_featmap = scene_featmap + self.scene_pos_embed
        scene_featmap = repeat_tensors(scene_featmap, num_ppl_per_img)

        head_featmap = self.head_branch_backbone(input["head_images"])
        head_featmap = self.head_proj(head_featmap)
        head_dino_tokens = head_featmap.flatten(start_dim=2).permute(0, 2, 1)
        if self.pos_encoding == "sinusoidal":
            head_featmap = head_featmap + self.head_pos_embed

        head_maps, head_rects = self.get_input_head_maps(input["bboxes"])
        head_maps = torch.cat(head_maps, dim=0).to(scene_featmap.device)
        head_rects = torch.cat(head_rects, dim=0).to(scene_featmap.device)
        if self.use_head_prompt:
            head_map_embeddings = head_maps.unsqueeze(dim=1) * self.head_position_token.weight.unsqueeze(-1).unsqueeze(-1)
            scene_featmap = scene_featmap + head_map_embeddings

        scene_tokens = scene_featmap.flatten(start_dim=2).permute(0, 2, 1)
        head_tokens = head_featmap.flatten(start_dim=2).permute(0, 2, 1)

        if self.n_reg_tokens > 0:
            scene_tokens = torch.cat([self.scene_register_tokens.expand(sum(num_ppl_per_img), -1, -1), scene_tokens], dim=1)
            head_tokens = torch.cat([self.head_register_tokens.expand(sum(num_ppl_per_img), -1, -1), head_tokens], dim=1)
        if self.inout:
            scene_tokens = torch.cat([self.scene_inout_token.expand(sum(num_ppl_per_img), -1, -1), scene_tokens], dim=1)
            head_tokens = torch.cat([self.head_inout_token.expand(sum(num_ppl_per_img), -1, -1), head_tokens], dim=1)
        if self.pos_encoding == "ape":
            scene_tokens = scene_tokens + self.scene_ape.expand(sum(num_ppl_per_img), -1, -1)
            head_tokens = head_tokens + self.head_ape.expand(sum(num_ppl_per_img), -1, -1)

        scene_tokens = self.scene_self_attn_layers(scene_tokens)
        head_tokens = self.head_self_attn_layers(head_tokens)
        tokens = self.scene_head_interaction_layers({"scene_tokens": scene_tokens, "head_tokens": head_tokens, "head_rects": head_rects})
        scene_tokens = tokens["scene_tokens"][:, self.n_front_tokens:, :]
        scene_inout_token = tokens["scene_tokens"][:, 0, :]
        head_inout_token = tokens["head_tokens"][:, 0, :]

        if self.inout:
            inout_features = torch.cat((scene_inout_token, head_inout_token), dim=1)
            inout_preds = self.inout_head(inout_features).squeeze(dim=-1)
            inout_preds = split_tensors(inout_preds, num_ppl_per_img)

        scene_featmap = scene_tokens.reshape(scene_tokens.shape[0], self.scene_featmap_h, self.scene_featmap_w,
                                             scene_tokens.shape[2]).permute(0, 3, 1, 2)
        heatmap = self.heatmap_head(scene_featmap).squeeze(dim=1)
        heatmap = torchvision.transforms.functional.resize(heatmap, self.heatmap_out_size)
        heatmap_preds = split_tensors(heatmap, num_ppl_per_img)

        if return_tokens:
            return {"scene_tokens": tokens["scene_tokens"], "head_tokens": tokens["head_tokens"],
                    "scene_dino_tokens": scene_dino_tokens, "head_dino_tokens": head_dino_tokens,
                    "heatmap": heatmap_preds, "inout": inout_preds if self.inout else None}
        return {"heatmap": heatmap_preds, "inout": inout_preds if self.inout else None}

    def forward(self, input):
        """Inference forward (applies sigmoid). Do NOT use for training (numerical stability)."""
        logits = self.get_logits(input)
        heatmap_preds = [torch.sigmoid(h) for h in logits["heatmap"]]
        inout_preds = [torch.sigmoid(i) for i in logits["inout"]] if logits["inout"] is not None else None
        return {"heatmap": heatmap_preds, "inout": inout_preds if self.inout else None}


# =========================================================================== #
#  Image processor (dual-stream: scene + per-person head crops)               #
# =========================================================================== #
class PaGEImageProcessor(BaseImageProcessor):
    """
    Produces the input dict expected by PaGEModel.forward from a scene image + head crops.

    Convenience:
        proc = AutoImageProcessor.from_pretrained(repo, trust_remote_code=True)
        inputs = proc(scene_pil, head_crops=[pil0, pil1, ...], bboxes=[[(xmin,ymin,xmax,ymax), ...]])
        out = model(inputs)
    """

    model_input_names = ["pixel_values"]

    def __init__(self, scene_size=(512, 512), head_size=(256, 256),
                 image_mean=(0.485, 0.456, 0.406), image_std=(0.229, 0.224, 0.225),
                 resample=2, **kwargs):
        super().__init__(**kwargs)
        self.scene_size = tuple(scene_size)
        self.head_size = tuple(head_size)
        self.image_mean = tuple(image_mean)
        self.image_std = tuple(image_std)
        self.resample = resample

    @classmethod
    def from_dict(cls, image_processor_dict, **kwargs):
        return cls(**{**image_processor_dict, **kwargs})

    def _to_tensor(self, pil_img, size):
        import numpy as np
        from PIL import Image
        if not isinstance(pil_img, Image.Image):
            pil_img = to_pil_image(pil_img)
        pil_img = pil_img.convert("RGB").resize((size[1], size[0]), self.resample)  # PIL resize is (W, H)
        arr = np.asarray(pil_img, dtype=np.float32) / 255.0  # H, W, 3
        arr = (arr - np.array(self.image_mean, dtype=np.float32)) / np.array(self.image_std, dtype=np.float32)
        arr = np.transpose(arr, (2, 0, 1))  # 3, H, W
        return torch.from_numpy(arr)

    def preprocess(self, scene_image, head_crops=None, bboxes=None, **kwargs):
        """
        scene_image: PIL image (or tensor) of the full scene.
        head_crops:  list of PIL images, one per person (length == total bboxes across scene).
                     Pass [None] * Np if you only have bboxes (zero head maps); but a real crop is expected.
        bboxes:      list of bbox lists, one per scene image. Each bbox: (xmin, ymin, xmax, ymax) in [0,1].
        Returns: dict with "images", "head_images", "bboxes" ready for PaGEModel.forward.
        """
        if bboxes is None:
            raise ValueError("bboxes is required.")
        scene_tensor = self._to_tensor(scene_image, self.scene_size).unsqueeze(0)  # 1, 3, H, W
        head_tensors = [self._to_tensor(hc, self.head_size) for hc in (head_crops or [])]
        if head_tensors:
            head_batch = torch.stack(head_tensors, dim=0)
        else:
            head_batch = torch.zeros(0, 3, *self.head_size)
        return {"images": scene_tensor, "head_images": [head_batch], "bboxes": bboxes}