lmms-lab-encoder/onevision-encoder-large-lang
826
1from typing import Optional, Tuple, Union2 3import torch4import torch.nn as nn5 6from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling7from transformers.modeling_utils import PreTrainedModel8from transformers.models.siglip.modeling_siglip import SiglipMLP9from transformers.utils import (10 add_start_docstrings,11 add_start_docstrings_to_model_forward,12 logging,13 replace_return_docstrings,14)15 16from .configuration_onevision_encoder import OneVisionEncoderConfig17 18 19try:20 from flash_attn import flash_attn_func21 22 _flash_attn_available = True23except ImportError:24 _flash_attn_available = False25 26logger = logging.get_logger(__name__)27 28 29# ---------------------------------------------------------------------------30# Model Docstrings31# ---------------------------------------------------------------------------32 33ONEVISION_ENCODER_START_DOCSTRING = r"""34 This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the35 library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads36 etc.)37 38 This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.39 Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage40 and behavior.41 42 Parameters:43 config ([`OneVisionEncoderConfig`]): Model configuration class with all the parameters of the model.44 Initializing with a config file does not load the weights associated with the model, only the45 configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.46"""47 48ONEVISION_ENCODER_INPUTS_DOCSTRING = r"""49 Args:50 pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` or `(batch_size, num_channels, num_frames, height, width)`):51 Pixel values. Pixel values can be obtained using [`AutoImageProcessor`].52 visible_indices (`torch.Tensor`, *optional*):53 Indices of visible patches for masking. Used in MAE-style pretraining or inference.54 output_attentions (`bool`, *optional*):55 Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned56 tensors for more detail.57 output_hidden_states (`bool`, *optional*):58 Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for59 more detail.60 return_dict (`bool`, *optional*):61 Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.62"""63 64 65# ---------------------------------------------------------------------------66# Helper Functions & Layers67# ---------------------------------------------------------------------------68 69 70def get_norm_layer(config):71 if config.layer_norm_type == "rms_norm":72 return nn.RMSNorm(config.hidden_size, eps=config.layer_norm_eps)73 else:74 return nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)75 76 77def rotate_half(x):78 """79 Interleaved rotation to match Source model's implementation.80 (x1, x2, x3, x4) -> (-x2, x1, -x4, x3)81 """82 x_even = x[..., ::2]83 x_odd = x[..., 1::2]84 return torch.stack((-x_odd, x_even), dim=-1).flatten(-2)85 86 87def apply_rotary_pos_emb(q, k, freqs):88 # q, k: (B, H, L, D)89 # freqs: (B, L, D)90 91 # We need to broadcast freqs to match heads92 # (B, L, D) -> (B, 1, L, D)93 94 # !!! CRITICAL FIX: Cast cos/sin to q.dtype (bf16/fp16) immediately95 # freqs are typically float32, so cos() returns float32.96 # Without this cast, (q * cos) upcasts q to float32, causing FlashAttention to fail.97 cos = freqs.cos().unsqueeze(1).to(q.dtype)98 sin = freqs.sin().unsqueeze(1).to(q.dtype)99 100 q_embed = (q * cos) + (rotate_half(q) * sin)101 k_embed = (k * cos) + (rotate_half(k) * sin)102 return q_embed, k_embed103 104 105class VideoRotaryEmbeddingSplit466(nn.Module):106 """107 3D (T,H,W) Rotary frequency constructor with 4:6:6 split.108 """109 110 def __init__(self, config: OneVisionEncoderConfig):111 super().__init__()112 head_dim = config.hidden_size // config.num_attention_heads113 base = config.rope_theta114 115 assert head_dim % 2 == 0, "head_dim must be even for rotary."116 assert head_dim % 16 == 0, "head_dim must be divisible by 16."117 half = head_dim // 2118 assert half % 16 == 0, "head_dim//2 must also be divisible by 16 to split into 4:6:6."119 120 self.head_dim = head_dim121 self.half = half122 123 unit = half // 16124 self.t_size = 4 * unit125 self.h_size = 6 * unit126 self.w_size = 6 * unit127 128 self.register_buffer(129 "inv_freq_t",130 1.0 / (base ** (torch.arange(self.t_size, dtype=torch.float32) / self.t_size)),131 persistent=False,132 )133 self.register_buffer(134 "inv_freq_h",135 1.0 / (base ** (torch.arange(self.h_size, dtype=torch.float32) / self.h_size)),136 persistent=False,137 )138 self.register_buffer(139 "inv_freq_w",140 1.0 / (base ** (torch.arange(self.w_size, dtype=torch.float32) / self.w_size)),141 persistent=False,142 )143 144 def forward(self, t: int, h: int, w: int, device=None):145 if device is None:146 device = self.inv_freq_t.device147 148 inv_t = self.inv_freq_t.to(device=device)149 inv_h = self.inv_freq_h.to(device=device)150 inv_w = self.inv_freq_w.to(device=device)151 152 ft = torch.outer(torch.arange(t, device=device, dtype=torch.float32), inv_t)153 fh = torch.outer(torch.arange(h, device=device, dtype=torch.float32), inv_h)154 fw = torch.outer(torch.arange(w, device=device, dtype=torch.float32), inv_w)155 156 t_ids = torch.arange(t, device=device).repeat_interleave(h * w)157 h_ids = torch.arange(h, device=device).repeat_interleave(w).repeat(t)158 w_ids = torch.arange(w, device=device).repeat(h).repeat(t)159 160 freqs = torch.cat([ft[t_ids], fh[h_ids], fw[w_ids]], dim=-1)161 return freqs162 163 def forward_from_positions(self, patch_positions: torch.Tensor) -> torch.Tensor:164 """165 Compute rotary position embeddings from explicit patch positions.166 167 Args:168 patch_positions: [batch_size, seq_len, 3] tensor with [t, h, w] positions for each patch169 170 Returns:171 freqs: [batch_size, seq_len, half] tensor of position frequencies172 """173 device = patch_positions.device174 inv_t = self.inv_freq_t.to(device=device)175 inv_h = self.inv_freq_h.to(device=device)176 inv_w = self.inv_freq_w.to(device=device)177 178 t_pos = patch_positions[..., 0].float() # [batch_size, seq_len]179 h_pos = patch_positions[..., 1].float() # [batch_size, seq_len]180 w_pos = patch_positions[..., 2].float() # [batch_size, seq_len]181 182 # Use einsum for batched outer product: [batch_size, seq_len] x [dim] -> [batch_size, seq_len, dim]183 ft = torch.einsum("bs,d->bsd", t_pos, inv_t)184 fh = torch.einsum("bs,d->bsd", h_pos, inv_h)185 fw = torch.einsum("bs,d->bsd", w_pos, inv_w)186 187 return torch.cat([ft, fh, fw], dim=-1)188 189 190class Siglip2MultiheadAttentionPoolingHead(nn.Module):191 """192 Multi-Head Attention Pooling with a learned probe (PMA-style).193 """194 195 def __init__(self, config: OneVisionEncoderConfig):196 super().__init__()197 self.embed_dim = config.hidden_size198 self.probe = nn.Parameter(torch.randn(1, 1, config.hidden_size))199 self.attention = nn.MultiheadAttention(config.hidden_size, config.num_attention_heads, batch_first=True)200 self.norm = nn.RMSNorm(config.hidden_size, eps=config.layer_norm_eps)201 self.mlp = SiglipMLP(config)202 203 def forward(self, hidden_states):204 batch_size = hidden_states.shape[0]205 probe = self.probe.repeat(batch_size, 1, 1)206 207 attn_output, _ = self.attention(probe, hidden_states, hidden_states)208 209 residual = attn_output210 attn_output = self.norm(attn_output)211 attn_output = residual + self.mlp(attn_output)212 213 return attn_output[:, 0]214 215 216# ---------------------------------------------------------------------------217# Modeling Components218# ---------------------------------------------------------------------------219 220 221class OneVisionEncoderEmbeddings(nn.Module):222 def __init__(self, config: OneVisionEncoderConfig):223 super().__init__()224 self.config = config225 self.embed_dim = config.hidden_size226 self.image_size = config.image_size227 self.patch_size = config.patch_size228 229 self.patch_embedding = nn.Conv2d(230 in_channels=config.num_channels,231 out_channels=self.embed_dim,232 kernel_size=self.patch_size,233 stride=self.patch_size,234 bias=False,235 )236 237 def forward(self, pixel_values: torch.FloatTensor) -> torch.Tensor:238 # Handle 4D (B, C, H, W) or 5D (B, C, T, H, W) inputs239 if pixel_values.dim() == 4:240 pixel_values = pixel_values.unsqueeze(2) # (B, C, 1, H, W)241 242 batch_size, channels, t_frames, height, width = pixel_values.shape243 244 # Merge time into batch for Conv2d245 x_2d = pixel_values.permute(0, 2, 1, 3, 4).reshape(batch_size * t_frames, channels, height, width)246 247 # Patch Embed248 embeddings = self.patch_embedding(x_2d) # (B*T, C, Hp, Wp)249 embeddings = embeddings.flatten(2).transpose(1, 2) # (B*T, L_frame, C)250 251 # Flatten all patches252 total_patches = t_frames * (height // self.patch_size) * (width // self.patch_size)253 embeddings = embeddings.reshape(batch_size, total_patches, self.embed_dim)254 255 return embeddings256 257 258class OneVisionEncoderAttention(nn.Module):259 """Multi-headed attention with RoPE support"""260 261 def __init__(self, config: OneVisionEncoderConfig):262 super().__init__()263 self.config = config264 self.embed_dim = config.hidden_size265 self.num_heads = config.num_attention_heads266 self.head_dim = self.embed_dim // self.num_heads267 if self.head_dim * self.num_heads != self.embed_dim:268 raise ValueError(269 f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`: {self.num_heads})."270 )271 272 self.scale = self.head_dim**-0.5273 self.dropout = config.attention_dropout274 275 self.k_proj = nn.Linear(self.embed_dim, self.embed_dim)276 self.v_proj = nn.Linear(self.embed_dim, self.embed_dim)277 self.q_proj = nn.Linear(self.embed_dim, self.embed_dim)278 self.out_proj = nn.Linear(self.embed_dim, self.embed_dim)279 280 def forward(281 self,282 hidden_states: torch.Tensor,283 attention_mask: Optional[torch.Tensor] = None,284 rotary_pos_emb: Optional[torch.Tensor] = None,285 output_attentions: bool = False,286 ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:287 batch_size, q_len, _ = hidden_states.size()288 289 query_states = self.q_proj(hidden_states)290 key_states = self.k_proj(hidden_states)291 value_states = self.v_proj(hidden_states)292 293 # (B, L, H, D) -> Transpose to (B, H, L, D)294 query_states = query_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)295 key_states = key_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)296 value_states = value_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)297 298 if rotary_pos_emb is not None:299 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, rotary_pos_emb)300 301 # Calculate attention scores302 attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) * self.scale303 304 if attention_mask is not None:305 if attention_mask.size() != (batch_size, 1, q_len, q_len):306 if attention_mask.dim() == 3:307 attention_mask = attention_mask.unsqueeze(1)308 attn_weights = attn_weights + attention_mask309 310 # FIX: Remove dtype=torch.float32 to stay in original dtype (bf16/fp16)311 attn_weights = nn.functional.softmax(attn_weights, dim=-1)312 attn_weights = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)313 314 attn_output = torch.matmul(attn_weights, value_states)315 316 attn_output = attn_output.transpose(1, 2).contiguous()317 attn_output = attn_output.reshape(batch_size, q_len, self.embed_dim)318 319 attn_output = self.out_proj(attn_output)320 321 return attn_output, attn_weights if output_attentions else None322 323 324class OneVisionEncoderFlashAttention2(nn.Module):325 """326 Multi-headed attention with RoPE support using Flash Attention 2.327 This module implements the same attention mechanism as OneVisionEncoderAttention but uses328 Flash Attention for improved performance and memory efficiency.329 """330 331 def __init__(self, config: OneVisionEncoderConfig):332 super().__init__()333 self.config = config334 self.embed_dim = config.hidden_size335 self.num_heads = config.num_attention_heads336 self.head_dim = self.embed_dim // self.num_heads337 if self.head_dim * self.num_heads != self.embed_dim:338 raise ValueError(339 f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`: {self.num_heads})."340 )341 342 self.scale = self.head_dim**-0.5343 self.dropout = config.attention_dropout344 345 self.k_proj = nn.Linear(self.embed_dim, self.embed_dim)346 self.v_proj = nn.Linear(self.embed_dim, self.embed_dim)347 self.q_proj = nn.Linear(self.embed_dim, self.embed_dim)348 self.out_proj = nn.Linear(self.embed_dim, self.embed_dim)349 350 def forward(351 self,352 hidden_states: torch.Tensor,353 attention_mask: Optional[torch.Tensor] = None,354 rotary_pos_emb: Optional[torch.Tensor] = None,355 output_attentions: bool = False,356 ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:357 """358 Forward pass using Flash Attention 2.359 """360 batch_size, q_len, _ = hidden_states.size()361 362 query_states = self.q_proj(hidden_states)363 key_states = self.k_proj(hidden_states)364 value_states = self.v_proj(hidden_states)365 366 # Flash Attention requires (B, L, H, D) format367 query_states = query_states.view(batch_size, q_len, self.num_heads, self.head_dim)368 key_states = key_states.view(batch_size, q_len, self.num_heads, self.head_dim)369 value_states = value_states.view(batch_size, q_len, self.num_heads, self.head_dim)370 371 # Apply RoPE if provided372 if rotary_pos_emb is not None:373 # Transpose for RoPE application: (B, L, H, D) -> (B, H, L, D)374 query_states = query_states.transpose(1, 2)375 key_states = key_states.transpose(1, 2)376 # NOTE: apply_rotary_pos_emb now ensures NO float32 cast happens377 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, rotary_pos_emb)378 # Transpose back: (B, H, L, D) -> (B, L, H, D)379 query_states = query_states.transpose(1, 2)380 key_states = key_states.transpose(1, 2)381 382 # Flash Attention forward pass383 if not _flash_attn_available:384 raise ImportError("flash_attn is not installed. Please install it to use OneVisionEncoderFlashAttention2.")385 386 attn_output = flash_attn_func(387 query_states,388 key_states,389 value_states,390 dropout_p=self.dropout if self.training else 0.0,391 softmax_scale=self.scale,392 causal=False,393 )394 395 # Reshape to (B, L, embed_dim)396 attn_output = attn_output.reshape(batch_size, q_len, self.embed_dim)397 398 # No extra casting here.399 attn_output = self.out_proj(attn_output)400 401 return attn_output, None402 403 404ONEVISION_ENCODER_ATTENTION_CLASSES = {405 "eager": OneVisionEncoderAttention,406 "flash_attention_2": OneVisionEncoderFlashAttention2,407}408 409 410class OneVisionEncoderEncoderLayer(nn.Module):411 def __init__(self, config: OneVisionEncoderConfig):412 super().__init__()413 self.embed_dim = config.hidden_size414 # Get attention implementation from config, default to "flash_attention_2"415 attn_implementation = getattr(config, "_attn_implementation", "flash_attention_2")416 if attn_implementation not in ONEVISION_ENCODER_ATTENTION_CLASSES:417 # Fallback to eager if flash_attention_2 is not available418 if not _flash_attn_available and attn_implementation == "flash_attention_2":419 attn_implementation = "eager"420 else:421 raise ValueError(422 f"Unknown attention implementation: {attn_implementation}. "423 f"Available implementations: {list(ONEVISION_ENCODER_ATTENTION_CLASSES.keys())}"424 )425 self.self_attn = ONEVISION_ENCODER_ATTENTION_CLASSES[attn_implementation](config)426 self.layer_norm1 = get_norm_layer(config)427 self.mlp = SiglipMLP(config)428 self.layer_norm2 = get_norm_layer(config)429 430 def forward(431 self,432 hidden_states: torch.Tensor,433 attention_mask: Optional[torch.Tensor] = None,434 rotary_pos_emb: Optional[torch.Tensor] = None,435 output_attentions: bool = False,436 ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:437 residual = hidden_states438 hidden_states = self.layer_norm1(hidden_states)439 440 hidden_states, attn_weights = self.self_attn(441 hidden_states=hidden_states,442 attention_mask=attention_mask,443 rotary_pos_emb=rotary_pos_emb,444 output_attentions=output_attentions,445 )446 hidden_states = residual + hidden_states447 448 residual = hidden_states449 hidden_states = self.layer_norm2(hidden_states)450 hidden_states = self.mlp(hidden_states)451 hidden_states = residual + hidden_states452 453 outputs = (hidden_states, attn_weights) if output_attentions else (hidden_states,)454 return outputs455 456 457class OneVisionEncoderEncoder(nn.Module):458 def __init__(self, config: OneVisionEncoderConfig):459 super().__init__()460 self.config = config461 self.layers = nn.ModuleList([OneVisionEncoderEncoderLayer(config) for _ in range(config.num_hidden_layers)])462 463 def forward(464 self,465 hidden_states: torch.Tensor,466 attention_mask: Optional[torch.Tensor] = None,467 rotary_pos_emb: Optional[torch.Tensor] = None,468 output_attentions: bool = False,469 output_hidden_states: bool = False,470 return_dict: bool = True,471 ) -> Union[tuple, BaseModelOutput]:472 all_hidden_states = () if output_hidden_states else None473 all_self_attentions = () if output_attentions else None474 475 for layer in self.layers:476 if output_hidden_states:477 all_hidden_states = all_hidden_states + (hidden_states,)478 479 layer_outputs = layer(480 hidden_states,481 attention_mask=attention_mask,482 rotary_pos_emb=rotary_pos_emb,483 output_attentions=output_attentions,484 )485 486 hidden_states = layer_outputs[0]487 488 if output_attentions:489 all_self_attentions = all_self_attentions + (layer_outputs[1],)490 491 if output_hidden_states:492 all_hidden_states = all_hidden_states + (hidden_states,)493 494 if not return_dict:495 return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None)496 497 return BaseModelOutput(498 last_hidden_state=hidden_states,499 hidden_states=all_hidden_states,500 attentions=all_self_attentions,501 )502 503 504# ---------------------------------------------------------------------------505# Main Models506# ---------------------------------------------------------------------------507 508 509@add_start_docstrings(510 "The bare OneVision Encoder Model outputting raw hidden-states without any specific head on top.",511 ONEVISION_ENCODER_START_DOCSTRING,512)513class OneVisionEncoderPreTrainedModel(PreTrainedModel):514 config_class = OneVisionEncoderConfig515 base_model_prefix = "onevision_encoder"516 supports_gradient_checkpointing = True517 _no_split_modules = ["OneVisionEncoderEncoderLayer"]518 _supports_flash_attn_2 = True519 520 def _init_weights(self, module):521 """Initialize the weights"""522 std = self.config.initializer_range523 if isinstance(module, (nn.Linear, nn.Conv2d)):524 module.weight.data.normal_(mean=0.0, std=std)525 if module.bias is not None:526 module.bias.data.zero_()527 elif isinstance(module, nn.Embedding):528 module.weight.data.normal_(mean=0.0, std=std)529 if module.padding_idx is not None:530 module.weight.data[module.padding_idx].zero_()531 elif isinstance(module, (nn.LayerNorm, nn.RMSNorm)):532 # Fix: RMSNorm doesn't have bias, must check hasattr first533 module.weight.data.fill_(1.0)534 if hasattr(module, "bias") and module.bias is not None:535 module.bias.data.zero_()536 537 538@add_start_docstrings(539 "OneVision Encoder Model with a vision transformer encoder.",540 ONEVISION_ENCODER_START_DOCSTRING,541)542class OneVisionEncoderModel(OneVisionEncoderPreTrainedModel):543 def __init__(self, config: OneVisionEncoderConfig):544 super().__init__(config)545 self.config = config546 547 self.embeddings = OneVisionEncoderEmbeddings(config)548 self.layernorm_pre = get_norm_layer(config)549 self.encoder = OneVisionEncoderEncoder(config)550 self.video_rope = VideoRotaryEmbeddingSplit466(config)551 552 if config.use_head:553 self.layernorm_post = get_norm_layer(config)554 self.head = Siglip2MultiheadAttentionPoolingHead(config)555 else:556 self.layernorm_post = None557 self.head = None558 559 self.post_init()560 561 @add_start_docstrings_to_model_forward(ONEVISION_ENCODER_INPUTS_DOCSTRING)562 @replace_return_docstrings(output_type=BaseModelOutputWithPooling, config_class=OneVisionEncoderConfig)563 def forward(564 self,565 pixel_values: torch.Tensor,566 visible_indices: Optional[torch.Tensor] = None,567 patch_positions: Optional[torch.Tensor] = None,568 output_attentions: Optional[bool] = None,569 output_hidden_states: Optional[bool] = None,570 return_dict: Optional[bool] = None,571 ) -> Union[tuple, BaseModelOutputWithPooling]:572 r"""573 Returns:574 575 Examples:576 577 ```python578 >>> from transformers import AutoModel, AutoImageProcessor579 >>> from PIL import Image580 581 >>> model = AutoModel.from_pretrained("lmms-lab-encoder/onevision-encoder-large", trust_remote_code=True)582 >>> preprocessor = AutoImageProcessor.from_pretrained("lmms-lab-encoder/onevision-encoder-large", trust_remote_code=True)583 >>> image = Image.open("path/to/your/image.jpg") # Replace with your image path584 >>> pixel_values = preprocessor(images=image, return_tensors="pt")["pixel_values"]585 >>> outputs = model(pixel_values)586 >>> last_hidden_states = outputs.last_hidden_state587 >>> pooled_output = outputs.pooler_output588 ```589 """590 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions591 output_hidden_states = (592 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states593 )594 return_dict = return_dict if return_dict is not None else self.config.use_return_dict595 596 # Determine video dimensions for RoPE597 # Note: pixel_values passed to embeddings can be 4D or 5D598 if pixel_values.dim() == 5:599 # Use config.rope_temporal_size if set, otherwise use actual frame count600 t_frames = (601 self.config.rope_temporal_size if self.config.rope_temporal_size is not None else pixel_values.shape[2]602 )603 height = pixel_values.shape[3]604 width = pixel_values.shape[4]605 else:606 t_frames = 1607 height = pixel_values.shape[2]608 width = pixel_values.shape[3]609 610 # 1. Embeddings611 hidden_states = self.embeddings(pixel_values)612 batch_size, total_patches, _ = hidden_states.shape613 614 # 2. Visible Indices Handling615 if visible_indices is None:616 visible_indices = (617 torch.arange(total_patches, device=pixel_values.device).unsqueeze(0).expand(batch_size, -1)618 )619 620 # 3. RoPE Construction621 if patch_positions is not None:622 freqs_visible = self.video_rope.forward_from_positions(patch_positions)623 else:624 freqs_full = self.video_rope(625 t=t_frames,626 h=height // self.config.patch_size,627 w=width // self.config.patch_size,628 device=pixel_values.device,629 )630 freqs_visible = freqs_full[visible_indices]631 632 # Concatenate D/2 + D/2 -> D for applying rope633 freqs_visible = torch.cat([freqs_visible, freqs_visible], dim=-1)634 635 # 4. Pre-Norm & Encoder636 hidden_states = self.layernorm_pre(hidden_states)637 638 # fix: gather hidden_states to match freqs_visible when using sparse visible_indices639 num_visible = visible_indices.shape[1]640 if num_visible != total_patches:641 # sparse mode: select only visible patches642 hidden_states = hidden_states.gather(643 1, visible_indices.unsqueeze(-1).expand(-1, -1, hidden_states.shape[-1])644 )645 646 encoder_outputs = self.encoder(647 hidden_states,648 attention_mask=None,649 rotary_pos_emb=freqs_visible,650 output_attentions=output_attentions,651 output_hidden_states=output_hidden_states,652 return_dict=return_dict,653 )654 655 sequence_output = encoder_outputs[0]656 657 # Apply post-norm if configured658 if self.layernorm_post is not None:659 sequence_output = self.layernorm_post(sequence_output)660 661 # 5. Pooling Head662 pooled_output = None663 if self.head is not None:664 pooled_output = self.head(sequence_output)665 666 if not return_dict:667 return (sequence_output, pooled_output) + encoder_outputs[1:]668 669 return BaseModelOutputWithPooling(670 last_hidden_state=sequence_output,671 pooler_output=pooled_output,672 hidden_states=encoder_outputs.hidden_states,673 attentions=encoder_outputs.attentions,674 )675