Upload folder using huggingface_hub

configuration_onevision_encoder.py

@@ -1,6 +1,7 @@
 from transformers.configuration_utils import PretrainedConfig
 from transformers.utils import logging
 
+
 logger = logging.get_logger(__name__)
 
 
@@ -26,7 +27,7 @@ class OneVisionEncoderConfig(PretrainedConfig):
             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):
+        patch_size (`int`, *optional*, defaults to 14):
             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.
@@ -70,7 +71,7 @@ class OneVisionEncoderConfig(PretrainedConfig):
         num_attention_heads=16,
         num_channels=3,
         image_size=448,
-        patch_size=16,
+        patch_size=14,
         hidden_act="gelu",
         layer_norm_eps=1e-6,
         layer_norm_type="layer_norm",

modeling_onevision_encoder.py

@@ -2,18 +2,23 @@ from typing import Optional, Tuple, Union
 
 import torch
 import torch.nn as nn
-from transformers.modeling_outputs import (BaseModelOutput,
-                                           BaseModelOutputWithPooling)
+
+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 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
@@ -61,6 +66,7 @@ ONEVISION_ENCODER_INPUTS_DOCSTRING = r"""
 # 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)
@@ -100,6 +106,7 @@ 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
@@ -118,12 +125,25 @@ class VideoRotaryEmbeddingSplit466(nn.Module):
         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)
+        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
+        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)
@@ -140,11 +160,38 @@ class VideoRotaryEmbeddingSplit466(nn.Module):
         freqs = torch.cat([ft[t_ids], fh[h_ids], fw[w_ids]], dim=-1)
         return freqs
 
+    def forward_from_positions(self, patch_positions: torch.Tensor) -> torch.Tensor:
+        """
+        Compute rotary position embeddings from explicit patch positions.
+
+        Args:
+            patch_positions: [batch_size, seq_len, 3] tensor with [t, h, w] positions for each patch
+
+        Returns:
+            freqs: [batch_size, seq_len, half] tensor of position frequencies
+        """
+        device = patch_positions.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)
+
+        t_pos = patch_positions[..., 0].float()  # [batch_size, seq_len]
+        h_pos = patch_positions[..., 1].float()  # [batch_size, seq_len]
+        w_pos = patch_positions[..., 2].float()  # [batch_size, seq_len]
+
+        # Use einsum for batched outer product: [batch_size, seq_len] x [dim] -> [batch_size, seq_len, dim]
+        ft = torch.einsum("bs,d->bsd", t_pos, inv_t)
+        fh = torch.einsum("bs,d->bsd", h_pos, inv_h)
+        fw = torch.einsum("bs,d->bsd", w_pos, inv_w)
+
+        return torch.cat([ft, fh, fw], dim=-1)
+
 
 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
@@ -170,6 +217,7 @@ class Siglip2MultiheadAttentionPoolingHead(nn.Module):
 # Modeling Components
 # ---------------------------------------------------------------------------
 
+
 class OneVisionEncoderEmbeddings(nn.Module):
     def __init__(self, config: OneVisionEncoderConfig):
         super().__init__()
@@ -189,7 +237,7 @@ class OneVisionEncoderEmbeddings(nn.Module):
     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)
+            pixel_values = pixel_values.unsqueeze(2)  # (B, C, 1, H, W)
 
         batch_size, channels, t_frames, height, width = pixel_values.shape
 
@@ -198,7 +246,7 @@ class OneVisionEncoderEmbeddings(nn.Module):
 
         # 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)
+        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)
@@ -209,6 +257,7 @@ class OneVisionEncoderEmbeddings(nn.Module):
 
 class OneVisionEncoderAttention(nn.Module):
     """Multi-headed attention with RoPE support"""
+
     def __init__(self, config: OneVisionEncoderConfig):
         super().__init__()
         self.config = config
@@ -216,7 +265,7 @@ class OneVisionEncoderAttention(nn.Module):
         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(
+            raise ValueError(
                 f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`: {self.num_heads})."
             )
 
@@ -228,7 +277,6 @@ class OneVisionEncoderAttention(nn.Module):
         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,
@@ -236,7 +284,6 @@ class OneVisionEncoderAttention(nn.Module):
         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)
@@ -257,7 +304,7 @@ class OneVisionEncoderAttention(nn.Module):
         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)
+                    attention_mask = attention_mask.unsqueeze(1)
             attn_weights = attn_weights + attention_mask
 
         # FIX: Remove dtype=torch.float32 to stay in original dtype (bf16/fp16)
@@ -280,6 +327,7 @@ class OneVisionEncoderFlashAttention2(nn.Module):
     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
@@ -386,7 +434,6 @@ class OneVisionEncoderEncoderLayer(nn.Module):
         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)
 
@@ -421,8 +468,7 @@ class OneVisionEncoderEncoder(nn.Module):
         output_attentions: bool = False,
         output_hidden_states: bool = False,
         return_dict: bool = True,
-    ) -> Union[Tuple, BaseModelOutput]:
-
+    ) -> Union[tuple, BaseModelOutput]:
         all_hidden_states = () if output_hidden_states else None
         all_self_attentions = () if output_attentions else None
 
@@ -459,6 +505,7 @@ class OneVisionEncoderEncoder(nn.Module):
 # Main Models
 # ---------------------------------------------------------------------------
 
+
 @add_start_docstrings(
     "The bare OneVision Encoder Model outputting raw hidden-states without any specific head on top.",
     ONEVISION_ENCODER_START_DOCSTRING,
@@ -484,7 +531,7 @@ class OneVisionEncoderPreTrainedModel(PreTrainedModel):
         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:
+            if hasattr(module, "bias") and module.bias is not None:
                 module.bias.data.zero_()
 
 
@@ -503,25 +550,25 @@ class OneVisionEncoderModel(OneVisionEncoderPreTrainedModel):
         self.video_rope = VideoRotaryEmbeddingSplit466(config)
 
         if config.use_head:
-             self.layernorm_post = get_norm_layer(config)
-             self.head = Siglip2MultiheadAttentionPoolingHead(config)
+            self.layernorm_post = get_norm_layer(config)
+            self.head = Siglip2MultiheadAttentionPoolingHead(config)
         else:
-             self.layernorm_post = None
-             self.head = None
+            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,
+        patch_positions: Optional[torch.Tensor] = None,
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
-    ) -> Union[Tuple, BaseModelOutputWithPooling]:
+    ) -> Union[tuple, BaseModelOutputWithPooling]:
         r"""
         Returns:
 
@@ -549,14 +596,16 @@ class OneVisionEncoderModel(OneVisionEncoderPreTrainedModel):
         # 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]
+            # 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]
+            t_frames = 1
+            height = pixel_values.shape[2]
+            width = pixel_values.shape[3]
 
         # 1. Embeddings
         hidden_states = self.embeddings(pixel_values)
@@ -564,16 +613,21 @@ class OneVisionEncoderModel(OneVisionEncoderPreTrainedModel):
 
         # 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)
+            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]
+        if patch_positions is not None:
+            freqs_visible = self.video_rope.forward_from_positions(patch_positions)
+        else:
+            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)