modeling_siglino.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import einops as E
from typing import Optional, Dict, Union, Tuple
from transformers import PreTrainedModel
from transformers.modeling_outputs import BaseModelOutput

# Relative imports from your local files
from .configuration_siglino import SigLinoConfig
from .attention import Attention, create_attention_mask
from .moe import MoE, FeedForward
from .rope import (
    precompute_freqs_cis,
    precompute_golden_freqs_cis,
    apply_golden_freqs_cis_to_visual_pos,
)

class PytorchGELUTanh(nn.Module):
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return F.gelu(x, approximate="tanh")

class Siglip2MLP(nn.Module):
    def __init__(self, hidden_size: int, intermediate_size: int):
        super().__init__()
        self.activation_fn = PytorchGELUTanh()
        self.fc1 = nn.Linear(hidden_size, intermediate_size)
        self.fc2 = nn.Linear(intermediate_size, hidden_size)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        hidden_states = self.fc1(hidden_states)
        hidden_states = self.activation_fn(hidden_states)
        hidden_states = self.fc2(hidden_states)
        return hidden_states

class Siglip2MultiheadAttentionPoolingHead(nn.Module):
    def __init__(self, hidden_size: int, num_attention_heads: int, output_dim: int):
        super().__init__()
        self.probe = nn.Parameter(torch.randn(1, 1, hidden_size))
        self.attention = nn.MultiheadAttention(hidden_size, num_attention_heads, batch_first=True)
        self.layernorm = nn.LayerNorm(hidden_size, eps=1e-5)
        self.mlp = Siglip2MLP(hidden_size, 4304)
        self.num_heads = num_attention_heads

    def forward(self, hidden_state: torch.Tensor, attention_mask: torch.Tensor | None = None) -> torch.Tensor:
        batch_size = hidden_state.shape[0]
        probe = self.probe.repeat(batch_size, 1, 1)

        if attention_mask is not None:
            # Mask expansion logic kept from your original model.py
            # Note: This uses einops and specific expansion for MHA
            def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: int | None = None):
                bsz, src_len = mask.size()
                tgt_len = tgt_len if tgt_len is not None else src_len
                expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)
                inverted_mask = torch.tensor(1.0, dtype=dtype, device=mask.device) - expanded_mask
                return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min)

            attention_mask = E.rearrange(attention_mask, "(b s) -> b s", b=batch_size)
            target_len, source_len = probe.shape[1], hidden_state.shape[1]
            attention_mask = _expand_mask(attention_mask, hidden_state.dtype, target_len)
            attention_mask = attention_mask.repeat(1, self.num_heads, target_len, 1)
            attention_mask = attention_mask.reshape(-1, target_len, source_len)

        hidden_state = self.attention(probe, hidden_state, hidden_state, attn_mask=attention_mask)[0]
        residual = hidden_state
        hidden_state = self.layernorm(hidden_state)
        hidden_state = residual + self.mlp(hidden_state)
        return hidden_state[:, 0]

class Adapter(nn.Module):
    def __init__(self, in_dim: int, out_dim: int, bias: bool = True):
        super().__init__()
        self.fc1 = nn.Linear(in_dim, out_dim)
        self.norm = nn.LayerNorm(out_dim)
        self.act = nn.GELU()
        self.fc2 = nn.Linear(out_dim, out_dim, bias=bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.fc1(x)
        x = self.norm(x)
        x = self.act(x)
        x = self.fc2(x)
        return x

class TransformerBlock(nn.Module):
    def __init__(self, layer_id: int, config: SigLinoConfig):
        super().__init__()
        self.dim = config.dim
        self.parameterized_norm = getattr(config, 'parameterized_norm', True)
        if self.parameterized_norm:
            self.attention_norm = nn.RMSNorm(config.dim, eps=config.norm_eps)
            self.ffn_norm = nn.RMSNorm(config.dim, eps=config.norm_eps)

        self.attention = Attention(
            dim=config.dim,
            n_heads=config.n_heads,
            n_kv_heads=config.n_kv_heads,
            head_dim=config.head_dim,
            use_qk_norm=config.use_qk_norm,
            enable_3d_rope=config.enable_3d_rope,
            use_flex_attn=config.use_flex_attn,
            use_sink_attn=True,
        )

        # Handle MoE initialization from config dict
        moe_args = config.moe_args
        if isinstance(moe_args, dict):
            from .moe import MoEArgs
            moe_args = MoEArgs(**moe_args)

        first_n_dense = getattr(config, 'first_n_layers_dense', 0)
        use_dense = layer_id < first_n_dense
        if use_dense:
            ffn_hidden = getattr(config, 'ffn_dim', None) or config.moe_dim
            activation = getattr(config, 'activation', 'silu')
            self.feed_forward = FeedForward(config.dim, ffn_hidden, activation=activation)
            self.moe_enabled = False
        elif moe_args and moe_args.num_experts > 0:
            self.moe = MoE(moe_args, dim=config.dim, hidden_dim=config.moe_dim)
            self.moe_enabled = True
        else:
            self.feed_forward = FeedForward(config.dim, config.moe_dim)
            self.moe_enabled = False

        self.weight_init_std = 0.02 / (2 * (layer_id + 1)) ** 0.5

    def forward(self, x, freqs_cis, freqs_cis_2d=None, pos_thw=None, attention_masks=None, compile=False):
        if self.parameterized_norm:
            x_norm = self.attention_norm(x)
        else:
            x_norm = F.rms_norm(x, (x.size(-1),))
        h = x + self.attention(
            x_norm,
            freqs_cis,
            freqs_cis_2d,
            pos_thw,
            attention_masks=attention_masks,
            compile=compile,
        )
        h_norm = self.ffn_norm(h) if self.parameterized_norm else F.rms_norm(h, (h.size(-1),))
        out = h + self.moe(h_norm) if self.moe_enabled else h + self.feed_forward(h_norm)
        return out

class SigLinoPreTrainedModel(PreTrainedModel):
    config_class = SigLinoConfig
    base_model_prefix = "siglino"
    main_input_name = "pixel_values"
    _no_split_modules = ["TransformerBlock"]

    def _init_weights(self, module):
        # Weight initialization is handled by the internal init_weights call in __init__
        pass

    def _apply(self, fn):
        # Prevent casting complex RoPE buffers (freqs_cis) to real dtypes on model.to(bf16/fp16)
        complex_buffers = {}
        for name, buf in list(self.named_buffers(recurse=False)):
            if buf is not None and buf.is_complex():
                complex_buffers[name] = buf
                del self._buffers[name]

        ret = super()._apply(fn)

        for name, buf in complex_buffers.items():
            dummy = torch.tensor([0.0], device=buf.device)
            res = fn(dummy)

            if not res.is_complex():
                new_buf = buf.to(device=res.device)
            else:
                new_buf = fn(buf)

            persistent = name not in self._non_persistent_buffers_set
            self.register_buffer(name, new_buf, persistent=persistent)

        return ret


class SigLinoModel(SigLinoPreTrainedModel):
    def __init__(self, config: SigLinoConfig):
        super().__init__(config)
        self.config = config
        self.n_layers = config.n_layers
        self.patch_size = config.spatial_patch_size
        self.n_storage_tokens = config.n_storage_tokens

        # Patch embedding
        self.n_pixels_per_patch = config.temporal_patch_size * config.spatial_patch_size ** 2
        self.img_projector = nn.Linear(
            self.n_pixels_per_patch * config.channel_size,
            config.dim,
            bias=False,
        )

        self.cls_token = nn.Parameter(torch.empty(1, 1, config.dim))
        if self.n_storage_tokens > 0:
            self.storage_tokens = nn.Parameter(torch.empty(1, self.n_storage_tokens, config.dim))

        # RoPE
        head_dim = config.head_dim or config.dim // config.n_heads
        d = head_dim // 2
        self.register_buffer("freqs_cis_golden", self._precompute_golden_freqs_cis(d, config))
        self.register_buffer("freqs_cis", self._precompute_freqs_cis(d, config), persistent=False)

        self.layers = nn.ModuleList([TransformerBlock(i, config) for i in range(config.n_layers)])
        self.norm = nn.RMSNorm(config.dim, eps=config.norm_eps)

        # Teacher adapters
        teachers_dict = dict(zip(config.teachers, config.teachers_dim))
        dinov3_dim = teachers_dict.get("dinov3", 1280)
        siglip2_dim = teachers_dict.get("siglip2", 1152)

        self.dinov3_adapter = Adapter(config.dim, dinov3_dim, bias=False)
        self.siglip2_adapter = Adapter(config.dim, siglip2_dim, bias=False)
        self.layer_norm_dinov3 = nn.LayerNorm(dinov3_dim)
        self.siglip2_multihead_attention_pooling_head = Siglip2MultiheadAttentionPoolingHead(
            siglip2_dim, 16, siglip2_dim
        )

        self.post_init()

    def _precompute_freqs_cis(self, head_dim: int, config: SigLinoConfig) -> torch.Tensor:
        return precompute_freqs_cis(head_dim, config.max_seq_len, config.rope_theta)

    def _precompute_golden_freqs_cis(self, head_dim: int, config: SigLinoConfig) -> torch.Tensor:
        return precompute_golden_freqs_cis(
            config.n_heads, head_dim, config.rope_min_freqs, config.rope_max_freqs
        )

    def _get_thw_pos(self, batch_size, num_patches, spatial_shapes, device):
        N = batch_size
        R = 1 + self.n_storage_tokens
        S = R + num_patches
        tpos = torch.zeros((N, S), dtype=torch.float32, device=device)
        hpos = torch.zeros((N, S), dtype=torch.float32, device=device)
        wpos = torch.zeros((N, S), dtype=torch.float32, device=device)

        for n in range(N):
            H, W = spatial_shapes[n].tolist()
            h_coords = torch.arange(H, device=device).float()
            w_coords = torch.arange(W, device=device).float()
            xlim, ylim = (W / H) ** 0.5, (H / W) ** 0.5
            h_norm = -ylim + 2 * ylim * h_coords / max(H - 1, 1)
            w_norm = -xlim + 2 * xlim * w_coords / max(W - 1, 1)
            
            # Vectorized fill for patches
            h_grid, w_grid = torch.meshgrid(h_norm, w_norm, indexing='ij')
            hpos[n, R:R+H*W] = h_grid.reshape(-1)
            wpos[n, R:R+H*W] = w_grid.reshape(-1)
            
            hpos[n, :R], wpos[n, :R] = float('nan'), float('nan')

        return torch.stack([tpos, hpos, wpos], dim=0)

    def forward(
        self,
        pixel_values: torch.Tensor,
        padding_mask: Optional[torch.Tensor] = None,
        spatial_shapes: Optional[torch.Tensor] = None,
        output_hidden_states: bool = False,
        return_dict: bool = True,
        compile: bool = True,
    ) -> Union[Dict, Tuple]:
        N, L, _ = pixel_values.shape
        device = pixel_values.device
        R = 1 + self.n_storage_tokens

        if padding_mask is None:
            padding_mask = torch.ones((N, L), dtype=pixel_values.dtype, device=device)
        
        h_NLD = self.img_projector(pixel_values)
        cls_expanded = self.cls_token.expand(N, -1, -1)
        if self.n_storage_tokens > 0:
            reg_expanded = self.storage_tokens.expand(N, -1, -1)
            h_NSD = torch.cat([cls_expanded, reg_expanded, h_NLD], dim=1)
        else:
            h_NSD = torch.cat([cls_expanded, h_NLD], dim=1)

        S = h_NSD.shape[1]
        cls_reg_mask = torch.ones((N, R), dtype=padding_mask.dtype, device=device)
        full_mask = torch.cat([cls_reg_mask, padding_mask], dim=1)
        
        # FlexAttention Mask
        def mask_mod(b, h, q_idx, kv_idx):
            return full_mask.bool()[b, q_idx] & full_mask.bool()[b, kv_idx]
        
        block_mask = create_attention_mask(mask_mod, N, None, S, S)

        # RoPE
        thw_pos = self._get_thw_pos(N, L, spatial_shapes, device)
        pos_thw = E.rearrange(thw_pos, "p n s -> n s p").to(dtype=torch.float32)
        patch_mask_2d = torch.zeros((N, S), dtype=torch.bool, device=device)
        patch_mask_2d[:, R:] = padding_mask.bool()
        pos_thw[:, :, 1:] = pos_thw[:, :, 1:].masked_fill(~patch_mask_2d.unsqueeze(-1), float("nan"))
        
        freqs_cis_golden = apply_golden_freqs_cis_to_visual_pos(
            self.freqs_cis_golden.to(dtype=pos_thw.dtype), pos_thw[:, :, 1:]
        )

        all_hidden_states = () if output_hidden_states else None
        for layer in self.layers:
            if output_hidden_states:
                all_hidden_states += (h_NSD,)
            h_NSD = layer(h_NSD, self.freqs_cis, freqs_cis_2d=freqs_cis_golden, 
                          pos_thw=pos_thw, attention_masks=block_mask, compile=compile)

        h_NSD = self.norm(h_NSD)
        
        # Feature Extraction & Adapters
        cls_feats = h_NSD[:, 0]
        patch_feats = h_NSD[:, R:]
        
        student_patch_dinov3 = self.dinov3_adapter(patch_feats)
        student_patch_siglip = self.siglip2_adapter(patch_feats)
        student_cls_dinov3 = self.dinov3_adapter(cls_feats)

        h_sig = self.siglip2_adapter(h_NSD)
        siglip_attn_mask = full_mask.reshape(-1)
        student_summary_siglip = self.siglip2_multihead_attention_pooling_head(h_sig, siglip_attn_mask)

        output = {
            "last_hidden_state": h_NSD,
            "patch_features": {
                "dinov3": student_patch_dinov3,
                "siglip2": student_patch_siglip,
                "siglino": patch_feats,
            },
            "summary_features": {
                "dinov3": student_cls_dinov3,
                "siglip2": student_summary_siglip,
                "siglino": cls_feats,
            },
            "hidden_states": all_hidden_states,
        }

        if not return_dict:
            return tuple(v for v in output.values() if v is not None)
        return output