remove some unuse code

modeling_step_vl.py

@@ -290,10 +290,6 @@ class StepRoboticsModel(StepRoboticsPreTrainedModel, GenerationMixin):
     def _get_vision_model_output(self,
                                  input_tensor: torch.Tensor) -> torch.Tensor:
         return self.vision_model(input_tensor)
-    
-    def _get_pooled_vision_model_output(
-            self, input_tensor: torch.Tensor) -> torch.Tensor:
-        return self.vision_model.pool(input_tensor)
 
     def _process_image_input(
             self, image_input: StepVLImageInputs) -> tuple[torch.Tensor, ...]:

vision_encoder.py

@@ -53,7 +53,6 @@ class EncoderRope2D(nn.Module):
         max_grid_height: int,
         max_grid_width: int,
         use_cls_token: bool = False,
-        freqs_for: Literal["lang", "pixel", "constant"] = "lang",
         theta: Union[int, float] = 10000,
         max_freq: int = 10,
         num_freqs: int = 1,
@@ -65,7 +64,6 @@ class EncoderRope2D(nn.Module):
         self.max_grid_width = max_grid_width
         self.use_cls_token = use_cls_token
         self.theta = theta * theta_rescale_factor**(dim / (dim - 2))
-        self.freqs_for = freqs_for
         self.max_freq = max_freq
         self.num_freqs = num_freqs
         cache = self._compute_2d_freqs()
@@ -73,15 +71,9 @@ class EncoderRope2D(nn.Module):
 
     def _compute_inv_freq(self, base: Union[int, float],
                           dim: int) -> torch.Tensor:
-        if self.freqs_for == "lang":
-            freqs = 1.0 / (base**(
-                torch.arange(0, dim, 2)[:(dim // 2)].float() / dim))
-        elif self.freqs_for == "pixel":
-            freqs = torch.linspace(1.0, self.max_freq / 2, dim // 2) * torch.pi
-        elif self.freqs_for == "constant":
-            freqs = torch.ones(self.num_freqs).float()
-        else:
-            raise ValueError(f"Unsupported freqs_for value: {self.freqs_for}")
+
+        freqs = 1.0 / (base**(
+            torch.arange(0, dim, 2)[:(dim // 2)].float() / dim))
         return freqs
 
     def _compute_freqs(self, t: torch.Tensor, inv_freq: torch.Tensor):
@@ -309,14 +301,9 @@ class EncoderVisionTransformer(nn.Module):
 
     def forward(self,
                 hidden_states: torch.Tensor,
-                grid_hw: tuple[int, int],
-                layer_idx: int = -1) -> torch.Tensor:
-        
-        stop_idx = (self.layers + layer_idx) % self.layers
-        for idx, block in enumerate(self.resblocks):
+                grid_hw: tuple[int, int]) -> torch.Tensor:
+        for block in self.resblocks:
             hidden_states = block(hidden_states, grid_hw=grid_hw)
-            if idx == stop_idx:
-                break
         return hidden_states
 
 
@@ -432,10 +419,7 @@ class StepRoboticsVisionEncoder(nn.Module):
 
         return pos_embed[None, ...]
 
-    def forward(self,
-                pixel_values: torch.Tensor,
-                layer_idx: int = -1,
-                strip_cls_token: bool = False) -> torch.Tensor:
+    def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
         """
         Args:
             pixel_values: Image tensor of shape (B, C, H, W).
@@ -457,12 +441,12 @@ class StepRoboticsVisionEncoder(nn.Module):
             pos_emb = self.sample_abs_posemb(grid_h, grid_w)
             hidden_state = hidden_state + pos_emb
         hidden_state = self.ln_pre(hidden_state)
-        hidden_state = self.transformer(hidden_state, grid_hw=(grid_h, grid_w), layer_idx=layer_idx)
+        hidden_state = self.transformer(hidden_state, grid_hw=(grid_h, grid_w))
 
         if self.use_ln_post:
             hidden_state = self.ln_post(hidden_state)
 
-        if strip_cls_token and self.use_cls_token:
+        if self.use_cls_token:
             hidden_state = hidden_state[:, 1:, :]
 
         return hidden_state