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