modeling_monkeyocrv2_vitae_vision.py

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.checkpoint
from transformers.modeling_utils import PreTrainedModel

flash_attn_available = True

try:
    from flash_attn import flash_attn_varlen_func
except ImportError:
    flash_attn_available = False

from .configuration_monkeyocrv2_vitae import MonkeyOCRv2ViTAEEncoderConfig

# ---------------------------------------------------------------------------
# 2D-RoPE helpers
# ---------------------------------------------------------------------------

def rotate_half(x):
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2:]
    return torch.cat((-x2, x1), dim=-1)


def apply_rotary_pos_emb_vision(tensor: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
    orig_dtype = tensor.dtype
    tensor = tensor.float()
    cos = freqs.cos()
    sin = freqs.sin()
    cos = cos.unsqueeze(1).repeat(1, 1, 2).unsqueeze(0).float()
    sin = sin.unsqueeze(1).repeat(1, 1, 2).unsqueeze(0).float()
    output = (tensor * cos) + (rotate_half(tensor) * sin)
    return output.to(orig_dtype)


class VisionRotaryEmbedding(nn.Module):
    def __init__(self, dim: int, theta: float = 10000.0) -> None:
        super().__init__()
        if dim % 2 != 0:
            raise ValueError("2D-RoPE requires an even per-axis rotary dimension.")
        inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float) / dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)

    def forward(self, seqlen: int) -> torch.Tensor:
        seq = torch.arange(seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype)
        freqs = torch.outer(seq, self.inv_freq)
        return freqs


# ---------------------------------------------------------------------------
# Norms
# ---------------------------------------------------------------------------

class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self._norm(x.float()).type_as(x) * self.weight

    def _norm(self, x: torch.Tensor) -> torch.Tensor:
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)


class RMSNorm2d(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))
        self.bias = nn.Parameter(torch.zeros(dim))
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        orig_dtype = x.dtype
        x = x.float()
        x = x * torch.rsqrt(x.pow(2).mean(1, keepdim=True) + self.eps)
        x = x * self.weight[:, None, None] + self.bias[:, None, None]
        return x.to(orig_dtype)


# ---------------------------------------------------------------------------
# Attention
# ---------------------------------------------------------------------------

class VisionAttentionSdpa(nn.Module):
    def __init__(self, dim: int, num_heads: int, bias: bool = False):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.qkv = nn.Linear(dim, dim * 3, bias=bias)
        self.proj = nn.Linear(dim, dim, bias=bias)

    def forward(self, x: torch.Tensor, cu_seqlens: torch.Tensor,
                rotary_pos_emb: torch.Tensor = None) -> torch.Tensor:
        S = x.shape[0]
        q, k, v = self.qkv(x).reshape(S, 3, self.num_heads, self.head_dim).permute(1, 0, 2, 3).unbind(0)
        if rotary_pos_emb is not None:
            q = apply_rotary_pos_emb_vision(q.unsqueeze(0), rotary_pos_emb).squeeze(0)
            k = apply_rotary_pos_emb_vision(k.unsqueeze(0), rotary_pos_emb).squeeze(0)
        mask = torch.zeros([1, S, S], device=q.device, dtype=torch.bool)
        for i in range(1, len(cu_seqlens)):
            mask[..., cu_seqlens[i - 1]:cu_seqlens[i], cu_seqlens[i - 1]:cu_seqlens[i]] = True
        q = q.transpose(0, 1).unsqueeze(0)
        k = k.transpose(0, 1).unsqueeze(0)
        v = v.transpose(0, 1).unsqueeze(0)
        if mask.stride(-1) != 1:
            mask = torch.empty_like(mask, memory_format=torch.contiguous_format).copy_(mask)
        out = F.scaled_dot_product_attention(q, k, v, mask, dropout_p=0.0)
        out = out.squeeze(0).transpose(0, 1).reshape(S, -1)
        return self.proj(out)


class VisionFlashAttention2(nn.Module):
    def __init__(self, dim: int, num_heads: int, is_causal: bool = False, bias: bool = False):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.is_causal = is_causal
        self.qkv = nn.Linear(dim, dim * 3, bias=bias)
        self.proj = nn.Linear(dim, dim, bias=bias)

    def forward(self, x: torch.Tensor, cu_seqlens: torch.Tensor,
                rotary_pos_emb: torch.Tensor = None) -> torch.Tensor:
        S = x.shape[0]
        q, k, v = self.qkv(x).reshape(S, 3, self.num_heads, self.head_dim).permute(1, 0, 2, 3).unbind(0)
        if rotary_pos_emb is not None:
            q = apply_rotary_pos_emb_vision(q.unsqueeze(0), rotary_pos_emb).squeeze(0)
            k = apply_rotary_pos_emb_vision(k.unsqueeze(0), rotary_pos_emb).squeeze(0)
        max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
        out = flash_attn_varlen_func(q, k, v, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen,
                                     causal=self.is_causal)
        out = out.reshape(S, -1)
        return self.proj(out)


class VisionAttentionEager(nn.Module):
    def __init__(self, dim: int, num_heads: int, bias: bool = False):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.qkv = nn.Linear(dim, dim * 3, bias=bias)
        self.proj = nn.Linear(dim, dim, bias=bias)

    def forward(self, x: torch.Tensor, cu_seqlens: torch.Tensor,
                rotary_pos_emb: torch.Tensor = None) -> torch.Tensor:
        S = x.shape[0]
        q, k, v = self.qkv(x).reshape(S, 3, self.num_heads, self.head_dim).permute(1, 0, 2, 3).unbind(0)
        if rotary_pos_emb is not None:
            q = apply_rotary_pos_emb_vision(q.unsqueeze(0), rotary_pos_emb).squeeze(0)
            k = apply_rotary_pos_emb_vision(k.unsqueeze(0), rotary_pos_emb).squeeze(0)
        mask = torch.full([1, S, S], torch.finfo(q.dtype).min, device=q.device, dtype=q.dtype)
        for i in range(1, len(cu_seqlens)):
            mask[..., cu_seqlens[i - 1]:cu_seqlens[i], cu_seqlens[i - 1]:cu_seqlens[i]] = 0
        q = q.transpose(0, 1); k = k.transpose(0, 1); v = v.transpose(0, 1)
        w = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim) + mask
        w = F.softmax(w, dim=-1, dtype=torch.float32).to(q.dtype)
        out = torch.matmul(w, v).transpose(0, 1).reshape(S, -1)
        return self.proj(out)


def _build_attn(dim: int, num_heads: int, attn_impl: str, is_causal: bool, bias: bool) -> nn.Module:
    if attn_impl == "flash_attention_2":
        if flash_attn_available:
            return VisionFlashAttention2(dim, num_heads, is_causal=is_causal, bias=bias)
        attn_impl = "sdpa"
    if attn_impl == "sdpa":
        return VisionAttentionSdpa(dim, num_heads, bias=bias)
    return VisionAttentionEager(dim, num_heads, bias=bias)


# ---------------------------------------------------------------------------
# MLP  (standard 2-layer GELU, matching official ViTAEv2)
# ---------------------------------------------------------------------------

class Mlp(nn.Module):
    def __init__(self, dim: int, hidden_dim: int, bias: bool = False):
        super().__init__()
        self.fc1 = nn.Linear(dim, hidden_dim, bias=bias)
        self.fc2 = nn.Linear(hidden_dim, dim, bias=bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.fc2(F.gelu(self.fc1(x)))


class SwiGLU(nn.Module):
    def __init__(self, dim: int, hidden_dim: int, bias: bool = False):
        super().__init__()
        self.gate_proj = nn.Linear(dim, hidden_dim, bias=bias)
        self.up_proj = nn.Linear(dim, hidden_dim, bias=bias)
        self.down_proj = nn.Linear(hidden_dim, dim, bias=bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))


# ---------------------------------------------------------------------------
# Window-attention helpers
# ---------------------------------------------------------------------------

def window_partition(x: torch.Tensor, window_size: int):
    """x: [B, H, W, C] → [nW*B, ws, ws, C]"""
    B, H, W, C = x.shape
    x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
    return x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)


def window_reverse(windows: torch.Tensor, window_size: int, H: int, W: int) -> torch.Tensor:
    """windows: [nW*B, ws, ws, C] → [B, H, W, C]"""
    B = int(windows.shape[0] / (H * W / window_size / window_size))
    x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)
    return x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)


# ---------------------------------------------------------------------------
# NormalCell with 2D-RoPE
# ---------------------------------------------------------------------------

class NormalCell2DRoPE(nn.Module):
    """ViTAEv2 NormalCell with windowed or global attention + 2D-RoPE + PCM."""

    def __init__(self, dim: int, num_heads: int, mlp_ratio: float = 4.0,
                 window_size: int = 7, window_attn: bool = True,
                 bias: bool = False, norm_eps: float = 1e-6,
                 attn_impl: str = "sdpa", is_causal: bool = False,
                 pcm_group: int = None):
        super().__init__()
        self.dim = dim
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        if self.head_dim % 4 != 0:
            raise ValueError("2D-RoPE requires attention head_dim divisible by 4.")
        self.window_attn = window_attn
        self.window_size = window_size

        self.norm1 = RMSNorm(dim, eps=norm_eps)
        self.norm2 = RMSNorm(dim, eps=norm_eps)
        self.qkv = nn.Linear(dim, dim * 3, bias=bias)
        self.proj = nn.Linear(dim, dim, bias=bias)

        hidden_dim = int(dim * mlp_ratio)
        swiglu_hidden_dim = math.ceil((2 * hidden_dim) / 3)
        self.mlp = SwiGLU(dim, swiglu_hidden_dim, bias=bias)

        g = pcm_group if pcm_group is not None else max(1, dim // 64)
        self.pcm = nn.Sequential(
            nn.Conv2d(dim, hidden_dim, 3, 1, 1, groups=g),
            RMSNorm2d(hidden_dim, eps=norm_eps),
            nn.SiLU(inplace=True),
            nn.Conv2d(hidden_dim, dim, 3, 1, 1, groups=g),
            RMSNorm2d(dim, eps=norm_eps),
            nn.SiLU(inplace=True),
            nn.Conv2d(dim, dim, 3, 1, 1, groups=g),
        )
        self._attn_impl = attn_impl
        self._is_causal = is_causal

        half_d = self.head_dim // 2
        ws = window_size
        h_ids = torch.arange(ws).unsqueeze(1).expand(ws, ws).reshape(-1)
        w_ids = torch.arange(ws).unsqueeze(0).expand(ws, ws).reshape(-1)
        theta = 1.0 / (10000.0 ** (torch.arange(0, half_d, 2, dtype=torch.float32) / half_d))
        freq_h = torch.outer(h_ids.float(), theta)  # [ws*ws, half_d//2]
        freq_w = torch.outer(w_ids.float(), theta)
        freqs = torch.stack([freq_h, freq_w], dim=2).reshape(ws * ws, half_d)  # [ws*ws, half_d]
        cos_emb = freqs.cos().repeat(1, 2)[None, None]  # [1, 1, ws*ws, head_dim]
        sin_emb = freqs.sin().repeat(1, 2)[None, None]
        self.register_buffer("_local_rope_cos", cos_emb, persistent=False)
        self.register_buffer("_local_rope_sin", sin_emb, persistent=False)

    def _apply_local_rope_window(self, q: torch.Tensor, ws: int) -> torch.Tensor:
        q_f = q.float()
        q_out = q_f * self._local_rope_cos + rotate_half(q_f) * self._local_rope_sin
        return q_out.to(q.dtype)

    def _window_attn(self, x: torch.Tensor, H: int, W: int) -> torch.Tensor:
        ws = self.window_size
        pad_h = (ws - H % ws) % ws
        pad_w = (ws - W % ws) % ws
        x2d = x.view(H, W, self.dim).unsqueeze(0)
        if pad_h + pad_w > 0:
            x2d = x2d.permute(0, 3, 1, 2)
            x2d = F.pad(x2d, (0, pad_w, 0, pad_h))
            x2d = x2d.permute(0, 2, 3, 1)
        Hp, Wp = H + pad_h, W + pad_w
        wins = window_partition(x2d, ws)               # [nW, ws, ws, dim]
        nWB = wins.shape[0]
        wins = wins.view(nWB, ws * ws, self.dim)
        q, k, v = self.qkv(wins).reshape(nWB, ws * ws, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4).unbind(0)
        q = self._apply_local_rope_window(q, ws)
        k = self._apply_local_rope_window(k, ws)
        attn = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
        attn = F.softmax(attn, dim=-1, dtype=torch.float32).to(q.dtype)
        out = torch.matmul(attn, v).transpose(1, 2).reshape(nWB, ws * ws, self.dim)
        out = self.proj(out)
        out = out.view(nWB, ws, ws, self.dim)
        out = window_reverse(out, ws, Hp, Wp)
        if pad_h + pad_w > 0:
            out = out[:, :H, :W, :].contiguous()
        return out.view(-1, self.dim)

    def _global_attn(self, x: torch.Tensor, cu_seqlens: torch.Tensor,
                     rotary_pos_emb: torch.Tensor) -> torch.Tensor:
        S = x.shape[0]
        q, k, v = self.qkv(x).reshape(S, 3, self.num_heads, self.head_dim).permute(1, 0, 2, 3).unbind(0)
        if rotary_pos_emb is not None:
            q = apply_rotary_pos_emb_vision(q.unsqueeze(0), rotary_pos_emb).squeeze(0)
            k = apply_rotary_pos_emb_vision(k.unsqueeze(0), rotary_pos_emb).squeeze(0)
        if flash_attn_available and self._attn_impl == "flash_attention_2":
            max_sl = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
            out = flash_attn_varlen_func(q, k, v, cu_seqlens, cu_seqlens, max_sl, max_sl, causal=False)
        else:
            mask = torch.zeros([1, S, S], device=q.device, dtype=torch.bool)
            for i in range(1, len(cu_seqlens)):
                mask[..., cu_seqlens[i - 1]:cu_seqlens[i], cu_seqlens[i - 1]:cu_seqlens[i]] = True
            qt = q.transpose(0, 1).unsqueeze(0)
            kt = k.transpose(0, 1).unsqueeze(0)
            vt = v.transpose(0, 1).unsqueeze(0)
            if mask.stride(-1) != 1:
                mask = torch.empty_like(mask, memory_format=torch.contiguous_format).copy_(mask)
            out = F.scaled_dot_product_attention(qt, kt, vt, mask, dropout_p=0.0)
            out = out.squeeze(0).transpose(0, 1)
        out = out.reshape(S, -1)
        return self.proj(out)

    def _pcm_single(self, x_flat: torch.Tensor, H: int, W: int) -> torch.Tensor:
        x2d = x_flat.view(1, H, W, self.dim).permute(0, 3, 1, 2)  # [1, dim, H, W]
        return self.pcm(x2d).permute(0, 2, 3, 1).reshape(H * W, self.dim)

    def _pcm(self, x: torch.Tensor, cu_seqlens: torch.Tensor, grid_hw: torch.Tensor) -> torch.Tensor:
        seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
        outs, token_idx = [], 0
        for img_idx, sl in enumerate(seqlens):
            H, W = int(grid_hw[img_idx, 0]), int(grid_hw[img_idx, 1])
            outs.append(self._pcm_single(x[token_idx: token_idx + sl], H, W))
            token_idx += sl
        return torch.cat(outs, dim=0)

    def forward(self, hidden_states: torch.Tensor, cu_seqlens: torch.Tensor,
                rotary_pos_emb: torch.Tensor, grid_hw: torch.Tensor = None) -> torch.Tensor:
        shortcut = hidden_states
        x_norm = self.norm1(hidden_states)

        if self.window_attn and grid_hw is not None:
            seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
            attn_outs, pcm_outs = [], []
            token_idx = 0
            for img_idx, sl in enumerate(seqlens):
                H, W = int(grid_hw[img_idx, 0]), int(grid_hw[img_idx, 1])
                attn_outs.append(self._window_attn(x_norm[token_idx: token_idx + sl], H, W))
                pcm_outs.append(self._pcm_single(shortcut[token_idx: token_idx + sl], H, W))
                token_idx += sl
            attn_out = torch.cat(attn_outs, dim=0)
            pcm_out  = torch.cat(pcm_outs, dim=0)
        else:
            attn_out = self._global_attn(x_norm, cu_seqlens, rotary_pos_emb)
            pcm_out  = self._pcm(shortcut, cu_seqlens, grid_hw) if grid_hw is not None else 0.0

        hidden_states = shortcut + attn_out + pcm_out
        hidden_states = hidden_states + self.mlp(self.norm2(hidden_states))
        return hidden_states


# ---------------------------------------------------------------------------
# ReductionCell
# ---------------------------------------------------------------------------

class ReductionCell(nn.Module):
    """
    ViTAEv2 ReductionCell: downsamples tokens via PRM (4 dilation branches) +
    single Transformer attention block + 3-conv PCM shortcut.

    RC_tokens_type in ViTAEv2-S: RC0&RC1 -> window attention (W), RC2&RC3 -> global attention (F).
    RC_group in ViTAEv2-S: [1, 16, 32, 64] for RC0, RC1, RC2, RC3.
    """

    def __init__(self, in_dim: int, out_dim: int,
                 num_heads: int = 1,
                 window_attn: bool = True,
                 window_size: int = 7,
                 pcm_group: int = 1,
                 prm_embed_dim: int = 64,
                 downsample_ratio: int = 2,
                 kernel_size: int = 3,
                 norm_eps: float = 1e-6,
                 bias: bool = False,
                 attn_impl: str = "sdpa"):
        super().__init__()
        self.in_dim = in_dim
        self.out_dim = out_dim
        self.num_heads = num_heads
        self.head_dim = out_dim // num_heads
        self.window_attn = window_attn
        self.window_size = window_size
        self.prm_embed_dim = prm_embed_dim
        self.downsample_ratio = downsample_ratio
        self.kernel_size = kernel_size
        self._attn_impl = attn_impl

        dilations = [1, 2, 3, 4]
        self.prm_convs = nn.ModuleList()
        for d in dilations:
            pad = math.ceil(((kernel_size - 1) * d + 1 - downsample_ratio) / 2)
            self.prm_convs.append(nn.Sequential(
                nn.Conv2d(in_dim, prm_embed_dim, kernel_size=kernel_size, stride=downsample_ratio, padding=pad, dilation=d),
                nn.GELU(),
            ))
        self.prm_proj = nn.Linear(prm_embed_dim * 4, out_dim, bias=bias)

        pcm_strides = []
        residual = downsample_ratio // 2
        for _ in range(3):
            pcm_strides.append((residual > 0) + 1)
            residual //= 2
        assert residual == 0

        g = pcm_group
        self.pcm = nn.Sequential(
            nn.Conv2d(in_dim, prm_embed_dim, 3, stride=pcm_strides[0], padding=1, groups=g),
            RMSNorm2d(prm_embed_dim, eps=norm_eps),
            nn.SiLU(inplace=True),
            nn.Conv2d(prm_embed_dim, prm_embed_dim, 3, stride=pcm_strides[1], padding=1, groups=g),
            RMSNorm2d(prm_embed_dim, eps=norm_eps),
            nn.SiLU(inplace=True),
            nn.Conv2d(prm_embed_dim, out_dim, 3, stride=pcm_strides[2], padding=1, groups=g),
        )

        self.norm_attn = RMSNorm(out_dim, eps=norm_eps)
        self.qkv = nn.Linear(out_dim, out_dim * 3, bias=bias)
        self.attn_proj = nn.Linear(out_dim, out_dim, bias=bias)
        self.norm_ffn = RMSNorm(out_dim, eps=norm_eps)
        self.mlp = Mlp(out_dim, out_dim, bias=bias)

    def _window_attn_rc(self, y: torch.Tensor, H: int, W: int) -> torch.Tensor:
        """Window self-attention on flat tokens [H*W, out_dim]. Returns [H*W, out_dim]."""
        ws = self.window_size
        y2d = y.view(1, H, W, self.out_dim)
        pad_h = (ws - H % ws) % ws
        pad_w = (ws - W % ws) % ws
        if pad_h + pad_w > 0:
            y2d = F.pad(y2d.permute(0, 3, 1, 2), (0, pad_w, 0, pad_h)).permute(0, 2, 3, 1)
        Hp, Wp = H + pad_h, W + pad_w
        wins = window_partition(y2d, ws)                          # [nW, ws, ws, out_dim]
        nW = wins.shape[0]
        wins_flat = wins.view(nW, ws * ws, self.out_dim)
        q, k, v = self.qkv(wins_flat).reshape(nW, ws * ws, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4).unbind(0)
        attn = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
        attn = F.softmax(attn, dim=-1, dtype=torch.float32).to(q.dtype)
        out = torch.matmul(attn, v).transpose(1, 2).reshape(nW, ws * ws, self.out_dim)
        out = self.attn_proj(out)
        out = window_reverse(out.view(nW, ws, ws, self.out_dim), ws, Hp, Wp)
        if pad_h + pad_w > 0:
            out = out[:, :H, :W, :].contiguous()
        return out.view(H * W, self.out_dim)

    def _global_attn_rc(self, y: torch.Tensor, cu_seqlens: torch.Tensor,
                         attn_impl: str = "sdpa") -> torch.Tensor:
        """Global attention over packed tokens using cu_seqlens. Returns same shape."""
        S = y.shape[0]
        q, k, v = self.qkv(y).reshape(S, 3, self.num_heads, self.head_dim).permute(1, 0, 2, 3).unbind(0)
        if flash_attn_available and attn_impl == "flash_attention_2":
            max_sl = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
            out = flash_attn_varlen_func(q, k, v, cu_seqlens, cu_seqlens, max_sl, max_sl, causal=False)
        else:
            mask = torch.zeros([1, S, S], device=q.device, dtype=torch.bool)
            for i in range(1, len(cu_seqlens)):
                mask[..., cu_seqlens[i - 1]:cu_seqlens[i], cu_seqlens[i - 1]:cu_seqlens[i]] = True
            qt = q.transpose(0, 1).unsqueeze(0)
            kt = k.transpose(0, 1).unsqueeze(0)
            vt = v.transpose(0, 1).unsqueeze(0)
            if mask.stride(-1) != 1:
                mask = torch.empty_like(mask, memory_format=torch.contiguous_format).copy_(mask)
            out = F.scaled_dot_product_attention(qt, kt, vt, mask, dropout_p=0.0)
            out = out.squeeze(0).transpose(0, 1)
        return self.attn_proj(out.reshape(S, self.out_dim))

    def forward(self, x: torch.Tensor, cu_seqlens: torch.Tensor, grid_hw: torch.Tensor):
        """
        x: [total_tokens, in_dim],  cu_seqlens: [N+1],  grid_hw: [N, 2]
        Returns: (x_out [total_new, out_dim], new_cu_seqlens [N+1], new_grid_hw [N, 2])
        """
        seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
        prm_list, pcm_list, new_hw_list = [], [], []
        token_idx = 0

        for img_idx, sl in enumerate(seqlens):
            H, W = int(grid_hw[img_idx, 0]), int(grid_hw[img_idx, 1])
            tok = x[token_idx: token_idx + sl]
            x2d = tok.view(1, H, W, self.in_dim).permute(0, 3, 1, 2)  # [1, in_dim, H, W]

            ys = [conv(x2d) for conv in self.prm_convs]        # each [1, prm_embed_dim, H//2, W//2]
            Hn, Wn = ys[0].shape[2], ys[0].shape[3]
            y = torch.cat(ys, dim=1)                            # [1, prm_embed_dim*4, Hn, Wn]
            y = y.permute(0, 2, 3, 1).reshape(Hn * Wn, self.prm_embed_dim * 4)
            prm_list.append(self.prm_proj(y))                   # [Hn*Wn, out_dim]

            pcm_list.append(self.pcm(x2d).permute(0, 2, 3, 1).reshape(Hn * Wn, self.out_dim))
            new_hw_list.append((Hn, Wn))
            token_idx += sl

        prm_out = torch.cat(prm_list, dim=0)
        pcm_out = torch.cat(pcm_list, dim=0)
        new_lens = torch.tensor([h * w for h, w in new_hw_list], dtype=torch.long, device=x.device)
        new_cu_seqlens = F.pad(new_lens.cumsum(0, dtype=torch.int32), (1, 0), value=0)
        new_grid_hw = torch.tensor(new_hw_list, dtype=torch.long, device=x.device)

        y = prm_out + pcm_out
        y_norm = self.norm_attn(y)
        if self.window_attn:
            attn_outs = []
            idx = 0
            for Hn, Wn in new_hw_list:
                sl_n = Hn * Wn
                attn_outs.append(self._window_attn_rc(y_norm[idx: idx + sl_n], Hn, Wn))
                idx += sl_n
            attn_out = torch.cat(attn_outs, dim=0)
        else:
            attn_out = self._global_attn_rc(y_norm, new_cu_seqlens, self._attn_impl)
        y = y + attn_out
        y = y + self.mlp(self.norm_ffn(y))
        return y, new_cu_seqlens, new_grid_hw


class ConvStem(nn.Module):
    def __init__(self, num_channels: int = 3, out_dim: int = 64,
                 patch_size: int = 32, temporal_patch_size: int = 1):
        super().__init__()
        self.num_channels = num_channels
        self.patch_size = patch_size
        self.temporal_patch_size = temporal_patch_size
        inter = out_dim // 2
        self.conv1 = nn.Sequential(
            nn.Conv2d(num_channels, inter, 3, stride=2, padding=1, bias=False),
            RMSNorm2d(inter),
            nn.GELU(),
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(inter, out_dim, 3, stride=2, padding=1, bias=False),
            RMSNorm2d(out_dim),
            nn.GELU(),
        )
        self.proj = nn.Conv2d(out_dim, out_dim, 3, stride=1, padding=1)
        self.norm = RMSNorm(out_dim)

    def forward(self, pixel_values: torch.Tensor, grid_thw: torch.Tensor):
        """
        pixel_values: [total_tiles, C * temporal * ph * pw]  — flat from Qwen2VLImageProcessor
                   OR [total_tiles, C, ph, pw]               — already spatial
        grid_thw: [N, 3]  (t, h_tiles, w_tiles)

        Returns:
            x: [total_tokens, out_dim]  packed
            cu_seqlens: [N+1]
            grid_hw: [N, 2]  (H_feat, W_feat) after 4x downsampling
        """
        # Reshape flat pixel_values [N, C*T*ph*pw] → [N, C, ph, pw]
        if pixel_values.ndim == 2:
            C = self.num_channels
            T = self.temporal_patch_size
            hw = int(math.sqrt(pixel_values.shape[1] / (C * T)))
            pixel_values = pixel_values.view(-1, C, T, hw, hw)[:, :, 0]  # [N, C, ph, pw]

        N = grid_thw.shape[0]
        outputs = []
        hw_list = []
        token_idx = 0

        for img_idx in range(N):
            t, h, w = int(grid_thw[img_idx, 0]), int(grid_thw[img_idx, 1]), int(grid_thw[img_idx, 2])
            num_tiles = t * h * w
            tiles = pixel_values[token_idx: token_idx + num_tiles]  # [h*w, C, ph, pw]
            token_idx += num_tiles
            ph = tiles.shape[-2]
            pw = tiles.shape[-1]
            img = tiles.view(h, w, -1, ph, pw).permute(2, 0, 3, 1, 4).reshape(-1, h * ph, w * pw).unsqueeze(0)
            out = self.conv1(img)                                   # [1, inter, H/2, W/2]
            out = self.conv2(out)                                   # [1, dim, H/4, W/4]
            out = self.proj(out)                                    # [1, dim, H/4, W/4]
            Hn, Wn = out.shape[2], out.shape[3]
            out = out.squeeze(0).permute(1, 2, 0).reshape(Hn * Wn, -1)  # [Hn*Wn, dim]
            out = self.norm(out)
            outputs.append(out)
            hw_list.append((Hn, Wn))

        x = torch.cat(outputs, dim=0)
        grid_hw = torch.tensor(hw_list, dtype=torch.long, device=x.device)
        lens = torch.tensor([h * w for h, w in hw_list], dtype=torch.long, device=x.device)
        cu_seqlens = F.pad(lens.cumsum(0, dtype=torch.int32), (1, 0), value=0)
        return x, cu_seqlens, grid_hw


class StemReductionCell(nn.Module):
    """
    First-stage reduction cell replacing ConvStem: reads raw pixel tiles and applies
    4x PRM (kernel=7, dilations=[1,2,3,4], stride=4) + PCM +
    window-attention block. Matches the official ViTAEv2-S RC0 (downsample_ratio=4).

    Forward interface is identical to the old ConvStem:
      forward(pixel_values, grid_thw) -> (packed_tokens, cu_seqlens, grid_hw)
    """

    def __init__(self, num_channels: int = 3, out_dim: int = 64, num_heads: int = 1,
                 window_size: int = 7, prm_embed_dim: int = 64,
                 downsample_ratio: int = 4, kernel_size: int = 7, pcm_group: int = 1,
                 patch_size: int = 32, temporal_patch_size: int = 1,
                 norm_eps: float = 1e-6, bias: bool = False, attn_impl: str = "sdpa"):
        super().__init__()
        self.num_channels = num_channels
        self.patch_size = patch_size
        self.temporal_patch_size = temporal_patch_size
        self.out_dim = out_dim
        self.num_heads = num_heads
        self.head_dim = out_dim // num_heads
        self.window_size = window_size
        self.prm_embed_dim = prm_embed_dim
        self.downsample_ratio = downsample_ratio
        self.kernel_size = kernel_size
        self._attn_impl = attn_impl

        in_dim = num_channels
        dilations = [1, 2, 3, 4]
        self.prm_convs = nn.ModuleList()
        for d in dilations:
            pad = math.ceil(((kernel_size - 1) * d + 1 - downsample_ratio) / 2)
            self.prm_convs.append(nn.Sequential(
                nn.Conv2d(in_dim, prm_embed_dim, kernel_size=kernel_size,
                          stride=downsample_ratio, padding=pad, dilation=d),
                nn.GELU(),
            ))
        self.prm_proj = nn.Linear(prm_embed_dim * 4, out_dim, bias=bias)

        pcm_strides = []
        residual = downsample_ratio // 2
        for _ in range(3):
            pcm_strides.append((residual > 0) + 1)
            residual //= 2
        assert residual == 0

        inter = prm_embed_dim
        self.pcm = nn.Sequential(
            nn.Conv2d(in_dim, inter, 3, stride=pcm_strides[0], padding=1, groups=pcm_group),
            RMSNorm2d(inter, eps=norm_eps),
            nn.SiLU(inplace=True),
            nn.Conv2d(inter, inter, 3, stride=pcm_strides[1], padding=1, groups=pcm_group),
            RMSNorm2d(inter, eps=norm_eps),
            nn.SiLU(inplace=True),
            nn.Conv2d(inter, out_dim, 3, stride=pcm_strides[2], padding=1, groups=pcm_group),
        )

        self.norm_attn = RMSNorm(out_dim, eps=norm_eps)
        self.qkv = nn.Linear(out_dim, out_dim * 3, bias=bias)
        self.attn_proj = nn.Linear(out_dim, out_dim, bias=bias)
        self.norm_ffn = RMSNorm(out_dim, eps=norm_eps)
        self.mlp = Mlp(out_dim, out_dim, bias=bias)

    def _window_attn(self, y: torch.Tensor, H: int, W: int) -> torch.Tensor:
        """Window self-attention on flat tokens [H*W, out_dim]."""
        ws = self.window_size
        y2d = y.view(1, H, W, self.out_dim)
        pad_h = (ws - H % ws) % ws
        pad_w = (ws - W % ws) % ws
        if pad_h + pad_w > 0:
            y2d = F.pad(y2d.permute(0, 3, 1, 2), (0, pad_w, 0, pad_h)).permute(0, 2, 3, 1)
        Hp, Wp = H + pad_h, W + pad_w
        wins = window_partition(y2d, ws)                    # [nW, ws, ws, out_dim]
        nW = wins.shape[0]
        wins_flat = wins.view(nW, ws * ws, self.out_dim)
        q, k, v = (self.qkv(wins_flat)
                       .reshape(nW, ws * ws, 3, self.num_heads, self.head_dim)
                       .permute(2, 0, 3, 1, 4).unbind(0))
        attn = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
        attn = F.softmax(attn, dim=-1, dtype=torch.float32).to(q.dtype)
        out  = torch.matmul(attn, v).transpose(1, 2).reshape(nW, ws * ws, self.out_dim)
        out  = self.attn_proj(out)
        out  = window_reverse(out.view(nW, ws, ws, self.out_dim), ws, Hp, Wp)
        if pad_h + pad_w > 0:
            out = out[:, :H, :W, :].contiguous()
        return out.view(H * W, self.out_dim)

    def forward(self, pixel_values: torch.Tensor, grid_thw: torch.Tensor):
        """
        pixel_values : [total_tiles, C, ph, pw]  or  [total_tiles, C*T*ph*pw]
        grid_thw     : [N, 3]  (t, h_tiles, w_tiles)
        Returns      : (x [total_tokens, out_dim], cu_seqlens [N+1], grid_hw [N, 2])
        """
        # Reshape flat tiles if preprocessor returned [N, C*T*ph*pw]
        if pixel_values.ndim == 2:
            C = self.num_channels
            T = self.temporal_patch_size
            hw_px = int(math.sqrt(pixel_values.shape[1] / (C * T)))
            pixel_values = pixel_values.view(-1, C, T, hw_px, hw_px)[:, :, 0]

        N = grid_thw.shape[0]
        prm_list, pcm_list, hw_list = [], [], []
        tile_idx = 0

        for img_idx in range(N):
            t, h, w = (int(grid_thw[img_idx, 0]),
                       int(grid_thw[img_idx, 1]),
                       int(grid_thw[img_idx, 2]))
            num_tiles = t * h * w
            tiles = pixel_values[tile_idx: tile_idx + num_tiles]   # [h*w, C, ph, pw]
            tile_idx += num_tiles
            ph, pw = tiles.shape[-2], tiles.shape[-1]
            img = (tiles.view(h, w, -1, ph, pw)
                        .permute(2, 0, 3, 1, 4)
                        .reshape(-1, h * ph, w * pw)
                        .unsqueeze(0))

            ys = [conv(img) for conv in self.prm_convs]
            Hn, Wn = ys[0].shape[2], ys[0].shape[3]
            y_cat  = torch.cat(ys, dim=1)                              # [1, emb*4, Hn, Wn]
            y_flat = y_cat.permute(0, 2, 3, 1).reshape(Hn * Wn, self.prm_embed_dim * 4)
            prm_list.append(self.prm_proj(y_flat))                     # [Hn*Wn, out_dim]

            pcm_list.append(
                self.pcm(img).permute(0, 2, 3, 1).reshape(Hn * Wn, self.out_dim)
            )
            hw_list.append((Hn, Wn))

        prm_out = torch.cat(prm_list, dim=0)     # [total_tokens, out_dim]
        pcm_out = torch.cat(pcm_list, dim=0)

        new_lens   = torch.tensor([h * w for h, w in hw_list],
                                   dtype=torch.long, device=prm_out.device)
        cu_seqlens = F.pad(new_lens.cumsum(0, dtype=torch.int32), (1, 0), value=0)
        grid_hw    = torch.tensor(hw_list, dtype=torch.long, device=prm_out.device)

        x = prm_out + pcm_out

        x_norm = self.norm_attn(x)
        attn_outs, idx = [], 0
        for Hn, Wn in hw_list:
            sl = Hn * Wn
            attn_outs.append(self._window_attn(x_norm[idx: idx + sl], Hn, Wn))
            idx += sl
        x = x + torch.cat(attn_outs, dim=0)
        x = x + self.mlp(self.norm_ffn(x))
        return x, cu_seqlens, grid_hw



# ---------------------------------------------------------------------------
# 2D-RoPE for global attention
# ---------------------------------------------------------------------------

def compute_2d_rope(grid_hw: torch.Tensor, rotary_emb: VisionRotaryEmbedding):
    emb_device = rotary_emb.inv_freq.device
    pos_ids_list = []
    for h, w in grid_hw.tolist():
        h, w = int(h), int(w)
        hpos = torch.arange(h, device=emb_device).unsqueeze(1).expand(h, w).flatten()
        wpos = torch.arange(w, device=emb_device).unsqueeze(0).expand(h, w).flatten()
        pos_ids_list.append(torch.stack([hpos, wpos], dim=-1))
    pos_ids = torch.cat(pos_ids_list, dim=0)
    max_size = grid_hw.max().item()
    full_emb = rotary_emb(int(max_size))
    emb = full_emb[pos_ids]
    return torch.stack([emb[:, 0], emb[:, 1]], dim=2).reshape(emb.shape[0], -1)


# ---------------------------------------------------------------------------
# PatchMerger
# ---------------------------------------------------------------------------

class PatchMerger(nn.Module):
    def __init__(self, dim: int, context_dim: int, init_std: float = 0.02):
        super().__init__()
        self.norm = RMSNorm(context_dim)
        self.mlp = nn.Sequential(
            nn.Linear(context_dim, context_dim),
            nn.GELU(),
            nn.Linear(context_dim, dim),
        )
        nn.init.normal_(self.mlp[0].weight, std=init_std)
        nn.init.zeros_(self.mlp[0].bias)
        nn.init.normal_(self.mlp[2].weight, std=init_std)
        nn.init.zeros_(self.mlp[2].bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.mlp(self.norm(x))

# ---------------------------------------------------------------------------
# Main ViTAEv2 Vision Transformer
# ---------------------------------------------------------------------------

class MonkeyOCRv2ViTAEVisionTransformer(PreTrainedModel):
    _supports_flash_attn = True
    _supports_sdpa = True
    _no_split_modules = ["StemReductionCell", "ReductionCell", "NormalCell2DRoPE"]
    config_class = MonkeyOCRv2ViTAEEncoderConfig
    def __init__(self, config: MonkeyOCRv2ViTAEEncoderConfig) -> None:
        super().__init__(config)
        self.config = config

        sd = config.stage_dims      # [64, 128, 256, 512]  (ViTAEv2-S)
        sdp = config.stage_depths    # [2, 2, 8, 2]
        sh = config.stage_heads      # [1, 2, 4, 8]  (head_dim=64 at all stages)
        dsr = getattr(config, 'downsample_ratios', [4, 2, 2, 2])
        ksz = getattr(config, 'kernel_sizes', [7, 3, 3, 3])
        rct = getattr(config, 'rc_tokens_type', ["window", "window", "transformer", "transformer"])
        nct = getattr(config, 'nc_tokens_type', ["window", "window", "transformer", "transformer"])
        ncg = getattr(config, 'nc_groups', [1, 32, 64, 128])
        rcg = list(getattr(config, 'rc_groups', [1, 16, 32, 64]))
        rcpem = list(getattr(config, 'rc_embed_dims', [64, 64, 128, 256]))
        rch = list(getattr(config, 'rc_heads', [1, 1, 2, 4]))
        if len(rcg) == 3:
            rcg = [1] + rcg
        if len(rch) == 3:
            rch = [1] + rch
        if len(rcpem) == 3:
            rcpem = [getattr(config, 'prm_embed_dim', 64)] + rcpem
        ws = config.window_size
        mr = config.mlp_ratio
        bias = config.use_bias
        eps = config.rms_norm_eps
        attn_impl = config.attn_implementation
        is_causal = config.is_causal

        self.stem = StemReductionCell(
            num_channels=config.num_channels,
            out_dim=sd[0],
            num_heads=rch[0],
            window_size=ws,
            prm_embed_dim=rcpem[0],
            downsample_ratio=dsr[0],
            kernel_size=ksz[0],
            pcm_group=rcg[0],
            patch_size=config.patch_size,
            temporal_patch_size=config.temporal_patch_size,
            norm_eps=eps,
            bias=bias,
            attn_impl=attn_impl,
        )

        self.stage1 = nn.ModuleList([
            NormalCell2DRoPE(sd[0], sh[0], mr, ws, nct[0] == "window", bias, eps, attn_impl, is_causal, ncg[0])
            for _ in range(sdp[0])])

        self.rc1 = ReductionCell(sd[0], sd[1], num_heads=rch[1], window_attn=rct[1] == "window", window_size=ws, pcm_group=rcg[1], prm_embed_dim=rcpem[1], downsample_ratio=dsr[1], kernel_size=ksz[1], norm_eps=eps, bias=bias, attn_impl=attn_impl)

        self.stage2 = nn.ModuleList([
            NormalCell2DRoPE(sd[1], sh[1], mr, ws, nct[1] == "window", bias, eps, attn_impl, is_causal, ncg[1])
            for _ in range(sdp[1])])

        self.rc2 = ReductionCell(sd[1], sd[2], num_heads=rch[2], window_attn=rct[2] == "window", window_size=ws, pcm_group=rcg[2], prm_embed_dim=rcpem[2], downsample_ratio=dsr[2], kernel_size=ksz[2], norm_eps=eps, bias=bias, attn_impl=attn_impl)

        head_dim_s3 = sd[2] // sh[2]
        self.rotary_pos_emb_s3 = VisionRotaryEmbedding(head_dim_s3 // 2)
        self.stage3 = nn.ModuleList([
            NormalCell2DRoPE(sd[2], sh[2], mr, ws, nct[2] == "window", bias, eps, attn_impl, is_causal, ncg[2])
            for _ in range(sdp[2])])

        self.rc3 = ReductionCell(sd[2], sd[3], num_heads=rch[3], window_attn=rct[3] == "window", window_size=ws, pcm_group=rcg[3], prm_embed_dim=rcpem[3], downsample_ratio=dsr[3], kernel_size=ksz[3], norm_eps=eps, bias=bias, attn_impl=attn_impl)

        head_dim_s4 = sd[3] // sh[3]
        self.rotary_pos_emb = VisionRotaryEmbedding(head_dim_s4 // 2)
        self.stage4 = nn.ModuleList([
            NormalCell2DRoPE(sd[3], sh[3], mr, ws, nct[3] == "window", bias, eps, attn_impl, is_causal, ncg[3])
            for _ in range(sdp[3])])

        if config.post_norm:
            self.post_trunk_norm = RMSNorm(sd[3], eps=eps)

        self.gradient_checkpointing = False
        self._gradient_checkpointing_func = torch.utils.checkpoint.checkpoint
        self.apply(self._init_weights)

    def _init_weights(self, module):
        std = self.config.initializer_range
        if isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=std)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Conv2d):
            module.weight.data.normal_(mean=0.0, std=std)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, (RMSNorm, RMSNorm2d)):
            module.weight.data.fill_(1.0)
            if hasattr(module, "bias") and module.bias is not None:
                module.bias.data.zero_()

    @property
    def dtype(self):
        return self.stem.prm_proj.weight.dtype

    @property
    def device(self):
        return self.stem.prm_proj.weight.device

    def _run_stage(self, cells, x, cu_seqlens, rotary_pos_emb, grid_hw):
        for cell in cells:
            if self.gradient_checkpointing and self.training:
                x = self._gradient_checkpointing_func(cell.__call__, x, cu_seqlens, rotary_pos_emb, grid_hw)
            else:
                x = cell(x, cu_seqlens, rotary_pos_emb, grid_hw)
        return x

    def forward(self, pixel_values: torch.Tensor, grid_thw: torch.Tensor) -> tuple:
        pixel_values = pixel_values.to(dtype=self.dtype)

        all_hidden_states =[]
        all_grid_hw = []
        # --- Stage 1 (1/4) ---
        x, cu_seqlens, grid_hw = self.stem(pixel_values, grid_thw)
        x = self._run_stage(self.stage1, x, cu_seqlens, None, grid_hw)
        all_hidden_states.append(x.clone())
        all_grid_hw.append(grid_hw.clone())

        # --- Stage 2 (1/8) ---
        x, cu_seqlens, grid_hw = self.rc1(x, cu_seqlens, grid_hw)
        x = self._run_stage(self.stage2, x, cu_seqlens, None, grid_hw)
        all_hidden_states.append(x.clone())
        all_grid_hw.append(grid_hw.clone())

        # --- Stage 3 (1/16) → global attention + 2D-RoPE ---
        x, cu_seqlens, grid_hw = self.rc2(x, cu_seqlens, grid_hw)
        rotary_pos_emb_s3 = compute_2d_rope(grid_hw, self.rotary_pos_emb_s3)
        x = self._run_stage(self.stage3, x, cu_seqlens, rotary_pos_emb_s3, grid_hw)
        all_hidden_states.append(x.clone())
        all_grid_hw.append(grid_hw.clone())

        # --- Stage 4 (1/32) → global attention ---
        x, cu_seqlens, grid_hw = self.rc3(x, cu_seqlens, grid_hw)
        rotary_pos_emb = compute_2d_rope(grid_hw, self.rotary_pos_emb)
        x = self._run_stage(self.stage4, x, cu_seqlens, rotary_pos_emb, grid_hw)

        if self.config.post_norm:
            x = self.post_trunk_norm(x)
        
        all_hidden_states.append(x)
        all_grid_hw.append(grid_hw)

        return all_hidden_states, all_grid_hw