Upload folder using huggingface_hub

README.md

@@ -0,0 +1,6 @@
+---
+license: apache-2.0
+---
+
+This is an intermediate stage of pretraining for the OneVision-Encoder model, where only images have been trained so far. More detailed documentation will be added later.
+

config.json

@@ -0,0 +1,27 @@
+{
+  "architectures": [
+    "OneVisionEncoderModel"
+  ],
+  "attention_dropout": 0.0,
+  "dtype": "bfloat16",
+  "hidden_act": "gelu",
+  "hidden_size": 1024,
+  "image_size": 448,
+  "initializer_range": 0.02,
+  "intermediate_size": 4096,
+  "layer_norm_eps": 1e-06,
+  "layer_norm_type": "layer_norm",
+  "model_type": "onevision_encoder",
+  "num_attention_heads": 16,
+  "num_channels": 3,
+  "num_hidden_layers": 24,
+  "patch_size": 14,
+  "rope_theta": 10000.0,
+  "rope_temporal_size": 64,
+  "transformers_version": "4.57.1",
+  "use_head": true,
+  "auto_map": {
+    "AutoConfig": "configuration_onevision_encoder.OneVisionEncoderConfig",
+    "AutoModel": "modeling_onevision_encoder.OneVisionEncoderModel"
+  }
+}

configuration_onevision_encoder.py

@@ -0,0 +1,99 @@
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+
+class OneVisionEncoderConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`OneVisionEncoderModel`]. It is used to instantiate a
+    OneVision Encoder model according to the specified arguments, defining the model architecture. Instantiating a configuration
+    with the defaults will yield a similar configuration to that of the OneVision Encoder architecture.
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+    Args:
+        hidden_size (`int`, *optional*, defaults to 1024):
+            Dimensionality of the encoder layers and the pooler layer.
+        intermediate_size (`int`, *optional*, defaults to 4096):
+            Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
+        num_hidden_layers (`int`, *optional*, defaults to 24):
+            Number of hidden layers in the Transformer encoder.
+        num_attention_heads (`int`, *optional*, defaults to 16):
+            Number of attention heads for each attention layer in the Transformer encoder.
+        num_channels (`int`, *optional*, defaults to 3):
+            The number of input channels.
+        image_size (`int`, *optional*, defaults to 224):
+            The size (resolution) of each image.
+        patch_size (`int`, *optional*, defaults to 16):
+            The size (resolution) of each patch.
+        hidden_act (`str` or `function`, *optional*, defaults to `"gelu"`):
+            The non-linear activation function (function or string) in the encoder and pooler.
+        layer_norm_eps (`float`, *optional*, defaults to 1e-6):
+            The epsilon used by the layer normalization layers.
+        layer_norm_type (`str`, *optional*, defaults to `"layer_norm"`):
+            The type of layer normalization to use. Supported values: `"layer_norm"`, `"rms_norm"`.
+        attention_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the attention probabilities.
+        initializer_range (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        rope_theta (`float`, *optional*, defaults to 10000.0):
+            The base period of the RoPE embeddings.
+        use_head (`bool`, *optional*, defaults to `True`):
+            Whether to use the pooling head.
+
+    Example:
+
+    ```python
+    >>> from configuration_onevision_encoder import OneVisionEncoderConfig
+    >>> from modeling_onevision_encoder import OneVisionEncoderModel
+
+    >>> # Initializing a OneVisionEncoder configuration
+    >>> configuration = OneVisionEncoderConfig()
+
+    >>> # Initializing a model (with random weights) from the configuration
+    >>> model = OneVisionEncoderModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```
+    """
+
+    model_type = "onevision_encoder"
+
+    def __init__(
+        self,
+        hidden_size=1024,
+        intermediate_size=4096,
+        num_hidden_layers=24,
+        num_attention_heads=16,
+        num_channels=3,
+        image_size=448,
+        patch_size=16,
+        hidden_act="gelu",
+        layer_norm_eps=1e-6,
+        layer_norm_type="layer_norm",
+        attention_dropout=0.0,
+        initializer_range=0.02,
+        rope_theta=10000.0,
+        rope_temporal_size=64,
+        use_head=True,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_channels = num_channels
+        self.image_size = image_size
+        self.patch_size = patch_size
+        self.hidden_act = hidden_act
+        self.layer_norm_eps = layer_norm_eps
+        self.layer_norm_type = layer_norm_type
+        self.attention_dropout = attention_dropout
+        self.initializer_range = initializer_range
+        self.rope_theta = rope_theta
+        self.rope_temporal_size = rope_temporal_size  # None=use actual frames, int=fixed size (legacy: 64)
+        self.use_head = use_head

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:844d3db0ab48ae305cc86e69c9ba7971b1a4d22ef8b37450824e39f1d99a5156
+size 631063168

modeling_onevision_encoder.py

@@ -0,0 +1,620 @@
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+from transformers.modeling_outputs import (BaseModelOutput,
+                                           BaseModelOutputWithPooling)
+from transformers.modeling_utils import PreTrainedModel
+from transformers.models.siglip.modeling_siglip import SiglipMLP
+from transformers.utils import (add_start_docstrings,
+                                add_start_docstrings_to_model_forward, logging,
+                                replace_return_docstrings)
+
+from .configuration_onevision_encoder import OneVisionEncoderConfig
+
+try:
+    from flash_attn import flash_attn_func
+    _flash_attn_available = True
+except ImportError:
+    _flash_attn_available = False
+
+logger = logging.get_logger(__name__)
+
+
+# ---------------------------------------------------------------------------
+# Model Docstrings
+# ---------------------------------------------------------------------------
+
+ONEVISION_ENCODER_START_DOCSTRING = r"""
+    This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
+    library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
+    etc.)
+
+    This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
+    Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
+    and behavior.
+
+    Parameters:
+        config ([`OneVisionEncoderConfig`]): Model configuration class with all the parameters of the model.
+            Initializing with a config file does not load the weights associated with the model, only the
+            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
+"""
+
+ONEVISION_ENCODER_INPUTS_DOCSTRING = r"""
+    Args:
+        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` or `(batch_size, num_channels, num_frames, height, width)`):
+            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`].
+        visible_indices (`torch.Tensor`, *optional*):
+            Indices of visible patches for masking. Used in MAE-style pretraining or inference.
+        output_attentions (`bool`, *optional*):
+            Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
+            tensors for more detail.
+        output_hidden_states (`bool`, *optional*):
+            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
+            more detail.
+        return_dict (`bool`, *optional*):
+            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+"""
+
+
+# ---------------------------------------------------------------------------
+# Helper Functions & Layers
+# ---------------------------------------------------------------------------
+
+def get_norm_layer(config):
+    if config.layer_norm_type == "rms_norm":
+        return nn.RMSNorm(config.hidden_size, eps=config.layer_norm_eps)
+    else:
+        return nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+
+def rotate_half(x):
+    """
+    Interleaved rotation to match Source model's implementation.
+    (x1, x2, x3, x4) -> (-x2, x1, -x4, x3)
+    """
+    x_even = x[..., ::2]
+    x_odd = x[..., 1::2]
+    return torch.stack((-x_odd, x_even), dim=-1).flatten(-2)
+
+
+def apply_rotary_pos_emb(q, k, freqs):
+    # q, k: (B, H, L, D)
+    # freqs: (B, L, D)
+
+    # We need to broadcast freqs to match heads
+    # (B, L, D) -> (B, 1, L, D)
+
+    # !!! CRITICAL FIX: Cast cos/sin to q.dtype (bf16/fp16) immediately
+    # freqs are typically float32, so cos() returns float32.
+    # Without this cast, (q * cos) upcasts q to float32, causing FlashAttention to fail.
+    cos = freqs.cos().unsqueeze(1).to(q.dtype)
+    sin = freqs.sin().unsqueeze(1).to(q.dtype)
+
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+class VideoRotaryEmbeddingSplit466(nn.Module):
+    """
+    3D (T,H,W) Rotary frequency constructor with 4:6:6 split.
+    """
+    def __init__(self, config: OneVisionEncoderConfig):
+        super().__init__()
+        head_dim = config.hidden_size // config.num_attention_heads
+        base = config.rope_theta
+
+        assert head_dim % 2 == 0, "head_dim must be even for rotary."
+        assert head_dim % 16 == 0, "head_dim must be divisible by 16."
+        half = head_dim // 2
+        assert half % 16 == 0, "head_dim//2 must also be divisible by 16 to split into 4:6:6."
+
+        self.head_dim = head_dim
+        self.half = half
+
+        unit = half // 16
+        self.t_size = 4 * unit
+        self.h_size = 6 * unit
+        self.w_size = 6 * unit
+
+        self.register_buffer("inv_freq_t", 1.0 / (base ** (torch.arange(self.t_size, dtype=torch.float32) / self.t_size)), persistent=False)
+        self.register_buffer("inv_freq_h", 1.0 / (base ** (torch.arange(self.h_size, dtype=torch.float32) / self.h_size)), persistent=False)
+        self.register_buffer("inv_freq_w", 1.0 / (base ** (torch.arange(self.w_size, dtype=torch.float32) / self.w_size)), persistent=False)
+
+    def forward(self, t: int, h: int, w: int, device=None):
+        if device is None: device = self.inv_freq_t.device
+
+        inv_t = self.inv_freq_t.to(device=device)
+        inv_h = self.inv_freq_h.to(device=device)
+        inv_w = self.inv_freq_w.to(device=device)
+
+        ft = torch.outer(torch.arange(t, device=device, dtype=torch.float32), inv_t)
+        fh = torch.outer(torch.arange(h, device=device, dtype=torch.float32), inv_h)
+        fw = torch.outer(torch.arange(w, device=device, dtype=torch.float32), inv_w)
+
+        t_ids = torch.arange(t, device=device).repeat_interleave(h * w)
+        h_ids = torch.arange(h, device=device).repeat_interleave(w).repeat(t)
+        w_ids = torch.arange(w, device=device).repeat(h).repeat(t)
+
+        freqs = torch.cat([ft[t_ids], fh[h_ids], fw[w_ids]], dim=-1)
+        return freqs
+
+
+class Siglip2MultiheadAttentionPoolingHead(nn.Module):
+    """
+    Multi-Head Attention Pooling with a learned probe (PMA-style).
+    """
+    def __init__(self, config: OneVisionEncoderConfig):
+        super().__init__()
+        self.embed_dim = config.hidden_size
+        self.probe = nn.Parameter(torch.randn(1, 1, config.hidden_size))
+        self.attention = nn.MultiheadAttention(config.hidden_size, config.num_attention_heads, batch_first=True)
+        self.norm = nn.RMSNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.mlp = SiglipMLP(config)
+
+    def forward(self, hidden_states):
+        batch_size = hidden_states.shape[0]
+        probe = self.probe.repeat(batch_size, 1, 1)
+
+        attn_output, _ = self.attention(probe, hidden_states, hidden_states)
+
+        residual = attn_output
+        attn_output = self.norm(attn_output)
+        attn_output = residual + self.mlp(attn_output)
+
+        return attn_output[:, 0]
+
+
+# ---------------------------------------------------------------------------
+# Modeling Components
+# ---------------------------------------------------------------------------
+
+class OneVisionEncoderEmbeddings(nn.Module):
+    def __init__(self, config: OneVisionEncoderConfig):
+        super().__init__()
+        self.config = config
+        self.embed_dim = config.hidden_size
+        self.image_size = config.image_size
+        self.patch_size = config.patch_size
+
+        self.patch_embedding = nn.Conv2d(
+            in_channels=config.num_channels,
+            out_channels=self.embed_dim,
+            kernel_size=self.patch_size,
+            stride=self.patch_size,
+            bias=False,
+        )
+
+    def forward(self, pixel_values: torch.FloatTensor) -> torch.Tensor:
+        # Handle 4D (B, C, H, W) or 5D (B, C, T, H, W) inputs
+        if pixel_values.dim() == 4:
+             pixel_values = pixel_values.unsqueeze(2) # (B, C, 1, H, W)
+
+        batch_size, channels, t_frames, height, width = pixel_values.shape
+
+        # Merge time into batch for Conv2d
+        x_2d = pixel_values.permute(0, 2, 1, 3, 4).reshape(batch_size * t_frames, channels, height, width)
+
+        # Patch Embed
+        embeddings = self.patch_embedding(x_2d)  # (B*T, C, Hp, Wp)
+        embeddings = embeddings.flatten(2).transpose(1, 2) # (B*T, L_frame, C)
+
+        # Flatten all patches
+        total_patches = t_frames * (height // self.patch_size) * (width // self.patch_size)
+        embeddings = embeddings.reshape(batch_size, total_patches, self.embed_dim)
+
+        return embeddings
+
+
+class OneVisionEncoderAttention(nn.Module):
+    """Multi-headed attention with RoPE support"""
+    def __init__(self, config: OneVisionEncoderConfig):
+        super().__init__()
+        self.config = config
+        self.embed_dim = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = self.embed_dim // self.num_heads
+        if self.head_dim * self.num_heads != self.embed_dim:
+             raise ValueError(
+                f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`: {self.num_heads})."
+            )
+
+        self.scale = self.head_dim**-0.5
+        self.dropout = config.attention_dropout
+
+        self.k_proj = nn.Linear(self.embed_dim, self.embed_dim)
+        self.v_proj = nn.Linear(self.embed_dim, self.embed_dim)
+        self.q_proj = nn.Linear(self.embed_dim, self.embed_dim)
+        self.out_proj = nn.Linear(self.embed_dim, self.embed_dim)
+
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        rotary_pos_emb: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+
+        batch_size, q_len, _ = hidden_states.size()
+
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+
+        # (B, L, H, D) -> Transpose to (B, H, L, D)
+        query_states = query_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = key_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+
+        if rotary_pos_emb is not None:
+            query_states, key_states = apply_rotary_pos_emb(query_states, key_states, rotary_pos_emb)
+
+        # Calculate attention scores
+        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) * self.scale
+
+        if attention_mask is not None:
+            if attention_mask.size() != (batch_size, 1, q_len, q_len):
+                if attention_mask.dim() == 3:
+                     attention_mask = attention_mask.unsqueeze(1)
+            attn_weights = attn_weights + attention_mask
+
+        # FIX: Remove dtype=torch.float32 to stay in original dtype (bf16/fp16)
+        attn_weights = nn.functional.softmax(attn_weights, dim=-1)
+        attn_weights = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)
+
+        attn_output = torch.matmul(attn_weights, value_states)
+
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.reshape(batch_size, q_len, self.embed_dim)
+
+        attn_output = self.out_proj(attn_output)
+
+        return attn_output, attn_weights if output_attentions else None
+
+
+class OneVisionEncoderFlashAttention2(nn.Module):
+    """
+    Multi-headed attention with RoPE support using Flash Attention 2.
+    This module implements the same attention mechanism as OneVisionEncoderAttention but uses
+    Flash Attention for improved performance and memory efficiency.
+    """
+    def __init__(self, config: OneVisionEncoderConfig):
+        super().__init__()
+        self.config = config
+        self.embed_dim = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = self.embed_dim // self.num_heads
+        if self.head_dim * self.num_heads != self.embed_dim:
+            raise ValueError(
+                f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`: {self.num_heads})."
+            )
+
+        self.scale = self.head_dim**-0.5
+        self.dropout = config.attention_dropout
+
+        self.k_proj = nn.Linear(self.embed_dim, self.embed_dim)
+        self.v_proj = nn.Linear(self.embed_dim, self.embed_dim)
+        self.q_proj = nn.Linear(self.embed_dim, self.embed_dim)
+        self.out_proj = nn.Linear(self.embed_dim, self.embed_dim)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        rotary_pos_emb: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        """
+        Forward pass using Flash Attention 2.
+        """
+        batch_size, q_len, _ = hidden_states.size()
+
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+
+        # Flash Attention requires (B, L, H, D) format
+        query_states = query_states.view(batch_size, q_len, self.num_heads, self.head_dim)
+        key_states = key_states.view(batch_size, q_len, self.num_heads, self.head_dim)
+        value_states = value_states.view(batch_size, q_len, self.num_heads, self.head_dim)
+
+        # Apply RoPE if provided
+        if rotary_pos_emb is not None:
+            # Transpose for RoPE application: (B, L, H, D) -> (B, H, L, D)
+            query_states = query_states.transpose(1, 2)
+            key_states = key_states.transpose(1, 2)
+            # NOTE: apply_rotary_pos_emb now ensures NO float32 cast happens
+            query_states, key_states = apply_rotary_pos_emb(query_states, key_states, rotary_pos_emb)
+            # Transpose back: (B, H, L, D) -> (B, L, H, D)
+            query_states = query_states.transpose(1, 2)
+            key_states = key_states.transpose(1, 2)
+
+        # Flash Attention forward pass
+        if not _flash_attn_available:
+            raise ImportError("flash_attn is not installed. Please install it to use OneVisionEncoderFlashAttention2.")
+
+        attn_output = flash_attn_func(
+            query_states,
+            key_states,
+            value_states,
+            dropout_p=self.dropout if self.training else 0.0,
+            softmax_scale=self.scale,
+            causal=False,
+        )
+
+        # Reshape to (B, L, embed_dim)
+        attn_output = attn_output.reshape(batch_size, q_len, self.embed_dim)
+
+        # No extra casting here.
+        attn_output = self.out_proj(attn_output)
+
+        return attn_output, None
+
+
+ONEVISION_ENCODER_ATTENTION_CLASSES = {
+    "eager": OneVisionEncoderAttention,
+    "flash_attention_2": OneVisionEncoderFlashAttention2,
+}
+
+
+class OneVisionEncoderEncoderLayer(nn.Module):
+    def __init__(self, config: OneVisionEncoderConfig):
+        super().__init__()
+        self.embed_dim = config.hidden_size
+        # Get attention implementation from config, default to "flash_attention_2"
+        attn_implementation = getattr(config, "_attn_implementation", "flash_attention_2")
+        if attn_implementation not in ONEVISION_ENCODER_ATTENTION_CLASSES:
+            # Fallback to eager if flash_attention_2 is not available
+            if not _flash_attn_available and attn_implementation == "flash_attention_2":
+                attn_implementation = "eager"
+            else:
+                raise ValueError(
+                    f"Unknown attention implementation: {attn_implementation}. "
+                    f"Available implementations: {list(ONEVISION_ENCODER_ATTENTION_CLASSES.keys())}"
+                )
+        self.self_attn = ONEVISION_ENCODER_ATTENTION_CLASSES[attn_implementation](config)
+        self.layer_norm1 = get_norm_layer(config)
+        self.mlp = SiglipMLP(config)
+        self.layer_norm2 = get_norm_layer(config)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        rotary_pos_emb: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+
+        residual = hidden_states
+        hidden_states = self.layer_norm1(hidden_states)
+
+        hidden_states, attn_weights = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            rotary_pos_emb=rotary_pos_emb,
+            output_attentions=output_attentions,
+        )
+        hidden_states = residual + hidden_states
+
+        residual = hidden_states
+        hidden_states = self.layer_norm2(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+
+        outputs = (hidden_states, attn_weights) if output_attentions else (hidden_states,)
+        return outputs
+
+
+class OneVisionEncoderEncoder(nn.Module):
+    def __init__(self, config: OneVisionEncoderConfig):
+        super().__init__()
+        self.config = config
+        self.layers = nn.ModuleList([OneVisionEncoderEncoderLayer(config) for _ in range(config.num_hidden_layers)])
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        rotary_pos_emb: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+        output_hidden_states: bool = False,
+        return_dict: bool = True,
+    ) -> Union[Tuple, BaseModelOutput]:
+
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attentions = () if output_attentions else None
+
+        for layer in self.layers:
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
+
+            layer_outputs = layer(
+                hidden_states,
+                attention_mask=attention_mask,
+                rotary_pos_emb=rotary_pos_emb,
+                output_attentions=output_attentions,
+            )
+
+            hidden_states = layer_outputs[0]
+
+            if output_attentions:
+                all_self_attentions = all_self_attentions + (layer_outputs[1],)
+
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + (hidden_states,)
+
+        if not return_dict:
+            return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None)
+
+        return BaseModelOutput(
+            last_hidden_state=hidden_states,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attentions,
+        )
+
+
+# ---------------------------------------------------------------------------
+# Main Models
+# ---------------------------------------------------------------------------
+
+@add_start_docstrings(
+    "The bare OneVision Encoder Model outputting raw hidden-states without any specific head on top.",
+    ONEVISION_ENCODER_START_DOCSTRING,
+)
+class OneVisionEncoderPreTrainedModel(PreTrainedModel):
+    config_class = OneVisionEncoderConfig
+    base_model_prefix = "onevision_encoder"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["OneVisionEncoderEncoderLayer"]
+    _supports_flash_attn_2 = True
+
+    def _init_weights(self, module):
+        """Initialize the weights"""
+        std = self.config.initializer_range
+        if isinstance(module, (nn.Linear, nn.Conv2d)):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+        elif isinstance(module, (nn.LayerNorm, nn.RMSNorm)):
+            # Fix: RMSNorm doesn't have bias, must check hasattr first
+            module.weight.data.fill_(1.0)
+            if hasattr(module, 'bias') and module.bias is not None:
+                module.bias.data.zero_()
+
+
+@add_start_docstrings(
+    "OneVision Encoder Model with a vision transformer encoder.",
+    ONEVISION_ENCODER_START_DOCSTRING,
+)
+class OneVisionEncoderModel(OneVisionEncoderPreTrainedModel):
+    def __init__(self, config: OneVisionEncoderConfig):
+        super().__init__(config)
+        self.config = config
+
+        self.embeddings = OneVisionEncoderEmbeddings(config)
+        self.layernorm_pre = get_norm_layer(config)
+        self.encoder = OneVisionEncoderEncoder(config)
+        self.video_rope = VideoRotaryEmbeddingSplit466(config)
+
+        if config.use_head:
+             self.layernorm_post = get_norm_layer(config)
+             self.head = Siglip2MultiheadAttentionPoolingHead(config)
+        else:
+             self.layernorm_post = None
+             self.head = None
+
+        self.post_init()
+
+
+    @add_start_docstrings_to_model_forward(ONEVISION_ENCODER_INPUTS_DOCSTRING)
+    @replace_return_docstrings(output_type=BaseModelOutputWithPooling, config_class=OneVisionEncoderConfig)
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        visible_indices: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPooling]:
+        r"""
+        Returns:
+
+        Examples:
+
+        ```python
+        >>> from transformers import AutoModel, AutoImageProcessor
+        >>> from PIL import Image
+
+        >>> model = AutoModel.from_pretrained("lmms-lab-encoder/onevision-encoder-large", trust_remote_code=True)
+        >>> preprocessor = AutoImageProcessor.from_pretrained("lmms-lab-encoder/onevision-encoder-large", trust_remote_code=True)
+        >>> image = Image.open("path/to/your/image.jpg")  # Replace with your image path
+        >>> pixel_values = preprocessor(images=image, return_tensors="pt")["pixel_values"]
+        >>> outputs = model(pixel_values)
+        >>> last_hidden_states = outputs.last_hidden_state
+        >>> pooled_output = outputs.pooler_output
+        ```
+        """
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # Determine video dimensions for RoPE
+        # Note: pixel_values passed to embeddings can be 4D or 5D
+        if pixel_values.dim() == 5:
+             # Use config.rope_temporal_size if set, otherwise use actual frame count
+             t_frames = self.config.rope_temporal_size if self.config.rope_temporal_size is not None else pixel_values.shape[2]
+             height = pixel_values.shape[3]
+             width = pixel_values.shape[4]
+        else:
+             t_frames = 1
+             height = pixel_values.shape[2]
+             width = pixel_values.shape[3]
+
+        # 1. Embeddings
+        hidden_states = self.embeddings(pixel_values)
+        batch_size, total_patches, _ = hidden_states.shape
+
+        # 2. Visible Indices Handling
+        if visible_indices is None:
+             visible_indices = torch.arange(total_patches, device=pixel_values.device).unsqueeze(0).expand(batch_size, -1)
+
+        # 3. RoPE Construction
+        freqs_full = self.video_rope(
+            t=t_frames,
+            h=height // self.config.patch_size,
+            w=width // self.config.patch_size,
+            device=pixel_values.device
+        )
+        freqs_visible = freqs_full[visible_indices]
+
+        # Concatenate D/2 + D/2 -> D for applying rope
+        freqs_visible = torch.cat([freqs_visible, freqs_visible], dim=-1)
+
+        # 4. Pre-Norm & Encoder
+        hidden_states = self.layernorm_pre(hidden_states)
+
+        # fix: gather hidden_states to match freqs_visible when using sparse visible_indices
+        num_visible = visible_indices.shape[1]
+        if num_visible != total_patches:
+            # sparse mode: select only visible patches
+            hidden_states = hidden_states.gather(
+                1, visible_indices.unsqueeze(-1).expand(-1, -1, hidden_states.shape[-1])
+            )
+
+        encoder_outputs = self.encoder(
+            hidden_states,
+            attention_mask=None,
+            rotary_pos_emb=freqs_visible,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        sequence_output = encoder_outputs[0]
+
+        # Apply post-norm if configured
+        if self.layernorm_post is not None:
+            sequence_output = self.layernorm_post(sequence_output)
+
+        # 5. Pooling Head
+        pooled_output = None
+        if self.head is not None:
+            pooled_output = self.head(sequence_output)
+
+        if not return_dict:
+            return (sequence_output, pooled_output) + encoder_outputs[1:]
+
+        return BaseModelOutputWithPooling(
+            last_hidden_state=sequence_output,
+            pooler_output=pooled_output,
+            hidden_states=encoder_outputs.hidden_states,
+            attentions=encoder_outputs.attentions,
+        )

ov_encoder_ocr_128gpus/ov_encoder_ocr_000.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3348cba862ed1f50c720b6d15834e96644adc28023dfc6328279de3062d0ddb6
+size 11691298

ov_encoder_ocr_128gpus/ov_encoder_ocr_001.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bf5d7f57d9c399c3447316cb986e8cf92bf462fc09a402211bb9b69feb43d6b7
+size 11691298

ov_encoder_ocr_128gpus/ov_encoder_ocr_002.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8f2032f7c23a78da0d6873b008d1e5abcc4ef008b19315a73fa091f807e34540
+size 11691298

ov_encoder_ocr_128gpus/ov_encoder_ocr_003.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1bb370040a9b6e567e69ae03e03db9337021e06f0deb487646336303008ce59b
+size 11687202

ov_encoder_ocr_128gpus/ov_encoder_ocr_004.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:431d12749d2f0807bf654197e3dcc40ed06561763717c696ae2b9ca38c2c90d8
+size 11687202

ov_encoder_ocr_128gpus/ov_encoder_ocr_005.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dea1c86154812f00dce2564cf9e87a2c04c727e1bf25965a5751e27110c7314f
+size 11687202

ov_encoder_ocr_128gpus/ov_encoder_ocr_006.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2d812476e61b20e0de90d5becb69d9d25fb3aa9e7d1e6c513f48b778fc74a097
+size 11687202

ov_encoder_ocr_128gpus/ov_encoder_ocr_007.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5a5bff2e065bf477d7b33ca18bdf917fe238020f3de1dc6fb708be30eaeaa835
+size 11687202

ov_encoder_ocr_128gpus/ov_encoder_ocr_008.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cd5eec98b18936ce8400007131747a51e3de9cc84349593ab08d9129c73fa094
+size 11687202

ov_encoder_ocr_128gpus/ov_encoder_ocr_009.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aa5b0851a6de31d5480422578c97de4b00cadddebad37068c668a60347f9b4c6
+size 11687202

ov_encoder_ocr_128gpus/ov_encoder_ocr_010.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0599eddc26e172ecb6970da89f8effbbe107917fdd72e56d415781fd48c52039
+size 11687202

ov_encoder_ocr_128gpus/ov_encoder_ocr_011.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:12ab27abd42bcd4d3af28a1a69aabceb9576fb4998e79552a6c61492193f9b6b
+size 11687202

ov_encoder_ocr_128gpus/ov_encoder_ocr_012.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6a8af7c0db84441e191aeda00bfd2fc9b23dbce7f5ba0eb35d95ab01a6947fad
+size 11687202

ov_encoder_ocr_128gpus/ov_encoder_ocr_013.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e1e5b374455edbc99d336856bf05395cc007a5fe92e6393d531624fb8183ad0a
+size 11687202

ov_encoder_ocr_128gpus/ov_encoder_ocr_014.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6ae41f6d8ac8c68c3aaacdca96d773c6481c3892ea112c974a3278a10ee385a1
+size 11687202

ov_encoder_ocr_128gpus/ov_encoder_ocr_015.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4bad6df8014630a9c62064fb07a4a20b8e5e31fe9f0df8b7cdd5f567ab02b5d4
+size 11687202

ov_encoder_ocr_128gpus/ov_encoder_ocr_016.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:98924ecae09cb64ed09a768f753f4d0f5a19eb7d4bf900908d2ebc1cf4ea0056
+size 11687202

ov_encoder_ocr_128gpus/ov_encoder_ocr_017.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3845e51a7092bd8411b0cc3a12927c1d8ca2ce55ca199ea14d42fe85bccb828a
+size 11687202

ov_encoder_ocr_128gpus/ov_encoder_ocr_018.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:59c36c4676f5c72d1e791aa66980d42a7036201dc7f3379db783801fd292f6ba
+size 11687202

ov_encoder_ocr_128gpus/ov_encoder_ocr_019.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:252e2afad8fd5d39bd7c6a1ecc125d5281b4ff1d47d342d9e6a6c45212040355
+size 11687202

ov_encoder_ocr_128gpus/ov_encoder_ocr_020.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:61eb05f4cccde6f3beed5873a0dbb0642b0f0ed59f4680a333d482a4bcb5248f
+size 11687202

ov_encoder_ocr_128gpus/ov_encoder_ocr_021.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4dad3fda872db9e9b27345222225e50fc79147b80f61baa48965cbeeb44d1117
+size 11687202

ov_encoder_ocr_128gpus/ov_encoder_ocr_022.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d2976fb756e5359c7a987cbbae8ce42d834290696cee28f00c590c80d1878dd8
+size 11687202

ov_encoder_ocr_128gpus/ov_encoder_ocr_023.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:685e06b275fb5c7cbef71cd2da37a1d995800e9cca7e40b6c143d0a732fb8e2e
+size 11687202

ov_encoder_ocr_128gpus/ov_encoder_ocr_024.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8d4fb09f4fff5a806b88235a82e8c068b8212d285a362af492e1674122248bd4
+size 11687202

ov_encoder_ocr_128gpus/ov_encoder_ocr_025.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:996adaebc64490b69559a7a705df309f6958544433bb887594686008b2206126
+size 11687202

ov_encoder_ocr_128gpus/ov_encoder_ocr_026.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b4d2e247f9a30812b26dd993615ad2297f1d0cfa09a4547c3bfcc2528b4204a4
+size 11687202

ov_encoder_ocr_128gpus/ov_encoder_ocr_027.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:383c80c3465365ea3448e1830780a4afc1e028edcece4b06f706122105aae6fd
+size 11687202

ov_encoder_ocr_128gpus/ov_encoder_ocr_028.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a3f49ca0675460a7b0384e4de87ee7cf79dfc5781fb871e0ffe669095bef8a5a
+size 11687202

ov_encoder_ocr_128gpus/ov_encoder_ocr_029.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e2f6dfeb0ebf305fe4437c82e12820004fd06b4cb75f36b6c4b5c53e77d77b3c
+size 11687202

ov_encoder_ocr_128gpus/ov_encoder_ocr_030.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e59dde5aa9bf708eedb3d4ab376857a4b77f3528a8c017aee4a1102797f4c006
+size 11687202

ov_encoder_ocr_128gpus/ov_encoder_ocr_031.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b758a0cd71c3e7fa472de827ad226d87eed23d6dbc80d227c13499a1c6b46c6e
+size 11687202

ov_encoder_ocr_128gpus/ov_encoder_ocr_032.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:916850d752e95a7a3ba2f08ddd9948ab4be6db5826ad8fb53bb3c2a4fe0d1187
+size 11687202

ov_encoder_ocr_128gpus/ov_encoder_ocr_033.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:261ab1a151ee8032845f25c588bc3531bcdb35304186b0506d9399c07414baf5
+size 11687202

ov_encoder_ocr_128gpus/ov_encoder_ocr_034.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:282b5abfd232860a637aca7d68d04779fe82c7e78cd4b6283096be473b46625f
+size 11687202

ov_encoder_ocr_128gpus/ov_encoder_ocr_035.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:18f9cb3db7f2f00a49062e0f607b9a3c68378bdbe0cdfd4045e0afd70a9b21fb
+size 11687202

ov_encoder_ocr_128gpus/ov_encoder_ocr_036.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:efac1e9fe1d813c58df675664bb15f2017f2ae7aa1c071a7a323bc25b8ba6b30
+size 11687202

ov_encoder_ocr_128gpus/ov_encoder_ocr_037.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ee7fb4ae1f6941a90efd76a789ef8bc1569ad92e58f7dd0d32d9d0a6f2680978
+size 11687202

ov_encoder_ocr_128gpus/ov_encoder_ocr_038.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:468533ee9c2ec1b26ec4daeccde6fde827282f9e01bf6b78547872f8518dbeef
+size 11687202

ov_encoder_ocr_128gpus/ov_encoder_ocr_039.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:058a37013c5e6d7b6541e6ea6c186ecc2e59bb94a0ccf83d5d14db70519e66b9
+size 11687202

ov_encoder_ocr_128gpus/ov_encoder_ocr_040.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a412be9045b2c50b478c281f8f09f658b45edbccf08ebf347a1808bc39a8c6d2
+size 11687202

ov_encoder_ocr_128gpus/ov_encoder_ocr_041.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ce50dea63053025448b069ed2267d83b3606b3f94f4202636ab43f1bbe18a523
+size 11687202

ov_encoder_ocr_128gpus/ov_encoder_ocr_042.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1b907f2b05a60d2a8113e247bd7dd9ce9371c9ad2afe5f63329f25e92c39eada
+size 11687202

ov_encoder_ocr_128gpus/ov_encoder_ocr_043.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b050812ef8b405592f662db18e95a11a0a256bbbc857972464bbb02a23c0afdb
+size 11687202

ov_encoder_ocr_128gpus/ov_encoder_ocr_044.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:357daf3e5ba5ba12179a0be299627a5e17075fdd5300dfd931a5b4cd212ef28c
+size 11687202