optimum-rbln 0.9.3.post1__py3-none-any.whl

optimum/rbln/transformers/models/clip/modeling_clip.py

@@ -0,0 +1,384 @@
+# Copyright 2025 Rebellions Inc. All rights reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import TYPE_CHECKING, Optional, Tuple, Union
+
+import torch
+from transformers import CLIPTextConfig, CLIPTextModel, CLIPVisionConfig, CLIPVisionModel
+from transformers.modeling_outputs import BaseModelOutputWithPooling
+from transformers.models.clip.modeling_clip import CLIPTextModelOutput, CLIPVisionModelOutput
+
+from ....configuration_utils import RBLNCompileConfig
+from ....modeling import RBLNModel
+from ....utils.logging import get_logger
+from .configuration_clip import RBLNCLIPTextModelConfig, RBLNCLIPVisionModelConfig
+
+
+logger = get_logger(__name__)
+
+if TYPE_CHECKING:
+    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, CLIPTextModel, PreTrainedModel
+
+    from ....diffusers.modeling_diffusers import RBLNDiffusionMixin, RBLNDiffusionMixinConfig
+
+
+class _TextEncoder(torch.nn.Module):
+    def __init__(self, enc: "CLIPTextModel"):
+        super().__init__()
+        self.enc = enc
+
+    def forward(self, inp):
+        enc_out = self.enc(inp, output_hidden_states=True, return_dict=False)
+        return enc_out
+
+
+class RBLNCLIPTextModel(RBLNModel):
+    """
+    RBLN optimized CLIP text encoder model.
+
+    This class provides hardware-accelerated inference for CLIP text encoders
+    on RBLN devices, supporting text encoding for multimodal tasks.
+    """
+
+    _tp_support = False
+
+    @classmethod
+    def _wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNCLIPTextModelConfig) -> torch.nn.Module:
+        return _TextEncoder(model).eval()
+
+    @classmethod
+    def update_rbln_config_using_pipe(
+        cls, pipe: "RBLNDiffusionMixin", rbln_config: "RBLNDiffusionMixinConfig", submodule_name: str
+    ) -> "RBLNDiffusionMixinConfig":
+        return rbln_config
+
+    @classmethod
+    def _update_rbln_config(
+        cls,
+        preprocessors: Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"],
+        model: Optional["PreTrainedModel"] = None,
+        model_config: "CLIPTextConfig" = None,
+        rbln_config: Optional[RBLNCLIPTextModelConfig] = None,
+    ) -> RBLNCLIPTextModelConfig:
+        input_info = [
+            (
+                "input_ids",
+                [
+                    rbln_config.batch_size,
+                    model_config.max_position_embeddings,
+                ],
+                "int64",
+            ),
+        ]
+
+        rbln_config.set_compile_cfgs([RBLNCompileConfig(input_info=input_info)])
+        return rbln_config
+
+    def forward(self, input_ids: torch.LongTensor, return_dict: Optional[bool] = None, **kwargs) -> torch.FloatTensor:
+        """
+        Forward pass for the RBLN-optimized CLIP text encoder model.
+
+        Args:
+            input_ids (torch.LongTensor): The input ids to the model.
+            return_dict (Optional[bool]): Whether to return a dictionary of outputs.
+
+        Returns:
+            The model outputs. If return_dict=False is passed, returns a tuple of tensors. Otherwise, returns a CLIPTextModelOutput object.
+        """
+
+        # To ignore using attention_mask, we override forward method.
+        output = super().forward(input_ids, return_dict=return_dict)
+        return output
+
+    def _prepare_output(self, output, return_dict):
+        # Prepare model output based on return_dict flag.
+        # This method can be overridden by subclasses to provide task-specific output handling.
+
+        if not return_dict:
+            return (output,) if not isinstance(output, (tuple, list)) else output
+        else:
+            return CLIPTextModelOutput(
+                text_embeds=output[0],
+                last_hidden_state=output[1],
+                hidden_states=output[2:],
+            )
+
+
+class RBLNCLIPTextModelWithProjection(RBLNCLIPTextModel):
+    """
+    RBLN optimized CLIP text encoder model with projection layer.
+
+    This class extends RBLNCLIPTextModel with a projection layer for
+    multimodal embedding alignment tasks.
+    """
+
+
+class _VisionEncoder(torch.nn.Module):
+    def __init__(
+        self,
+        enc: CLIPVisionModel,
+        interpolate_pos_encoding: bool,
+        output_hidden_states: bool,
+        output_attentions: bool,
+    ):
+        super().__init__()
+        self.enc = enc
+        self.interpolate_pos_encoding = interpolate_pos_encoding
+        self.output_hidden_states = output_hidden_states
+        self.output_attentions = output_attentions
+
+    def forward(self, inp):
+        enc_out = self.enc(
+            inp,
+            output_hidden_states=self.output_hidden_states,
+            interpolate_pos_encoding=self.interpolate_pos_encoding,
+            output_attentions=self.output_attentions,
+            return_dict=False,
+        )
+        return enc_out
+
+
+class RBLNCLIPVisionModel(RBLNModel):
+    """
+    RBLN optimized CLIP vision encoder model.
+
+    This class provides hardware-accelerated inference for CLIP vision encoders
+    on RBLN devices, supporting image encoding for multimodal tasks.
+    """
+
+    _tp_support = False
+
+    @classmethod
+    def _wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNCLIPVisionModelConfig) -> torch.nn.Module:
+        wrapper_cfg = {
+            "interpolate_pos_encoding": rbln_config.interpolate_pos_encoding,
+            "output_hidden_states": rbln_config.output_hidden_states,
+            "output_attentions": rbln_config.output_attentions,
+        }
+        return _VisionEncoder(model, **wrapper_cfg).eval()
+
+    @classmethod
+    def update_rbln_config_using_pipe(
+        cls, pipe: "RBLNDiffusionMixin", rbln_config: "RBLNDiffusionMixinConfig", submodule_name: str
+    ) -> "RBLNDiffusionMixinConfig":
+        return rbln_config
+
+    @classmethod
+    def _update_rbln_config(
+        cls,
+        preprocessors: Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"],
+        model: Optional["PreTrainedModel"] = None,
+        model_config: "CLIPVisionConfig" = None,
+        rbln_config: Optional[RBLNCLIPVisionModelConfig] = None,
+    ) -> RBLNCLIPVisionModelConfig:
+        if rbln_config.image_size is None:
+            rbln_config.image_size = getattr(model_config, "image_size", None)
+
+        if isinstance(rbln_config.image_size, int):
+            rbln_config.image_size = (rbln_config.image_size, rbln_config.image_size)
+
+        if rbln_config.image_size is None:
+            raise ValueError("`rbln_image_size` should be specified!")
+
+        if rbln_config.output_attentions is None:
+            rbln_config.output_attentions = getattr(model_config, "output_attentions", False)
+
+        if rbln_config.output_hidden_states is None:
+            rbln_config.output_hidden_states = getattr(model_config, "output_hidden_states", False)
+
+        rbln_compile_config = RBLNCompileConfig(
+            input_info=[
+                (
+                    "pixel_values",
+                    [
+                        rbln_config.batch_size,
+                        3,
+                        rbln_config.image_height,
+                        rbln_config.image_width,
+                    ],
+                    "float32",
+                )
+            ]
+        )
+
+        rbln_config.set_compile_cfgs([rbln_compile_config])
+        return rbln_config
+
+    def forward(
+        self,
+        pixel_values: torch.FloatTensor,
+        return_dict: bool = True,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        interpolate_pos_encoding: bool = False,
+        **kwargs,
+    ) -> Union[Tuple, BaseModelOutputWithPooling]:
+        """
+        Forward pass for the RBLN-optimized CLIP vision encoder model.
+
+        Args:
+            pixel_values (torch.Tensor): The pixel values to the model.
+            return_dict (bool): Whether to return a dictionary of outputs.
+            output_attentions (Optional[bool]): Whether to return attentions.
+            output_hidden_states (Optional[bool]): Whether to return hidden states.
+            interpolate_pos_encoding (bool): Whether to interpolate position encoding.
+
+        Returns:
+            The model outputs. If return_dict=False is passed, returns a tuple of tensors. Otherwise, returns a BaseModelOutputWithPooling object.
+        """
+
+        if len(kwargs) > 0 and any(value is not None for value in kwargs.values()):
+            logger.warning(
+                f"Currently, optimum-rbln does not support kwargs {kwargs.keys()} for {self.__class__.__name__}."
+            )
+
+        output_attentions = output_attentions if output_attentions is not None else self.rbln_config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.rbln_config.output_hidden_states
+        )
+
+        if output_attentions != self.rbln_config.output_attentions:
+            raise ValueError(
+                f"Variable output_attentions {output_attentions} is not equal to rbln_config.output_attentions {self.rbln_config.output_attentions} "
+                f"Please compile again with the correct argument."
+            )
+
+        if output_hidden_states != self.rbln_config.output_hidden_states:
+            raise ValueError(
+                f"Variable output_hidden_states {output_hidden_states} is not equal to rbln_config.output_hidden_states {self.rbln_config.output_hidden_states} "
+                f"Please compile again with the correct argument."
+            )
+
+        if interpolate_pos_encoding != self.rbln_config.interpolate_pos_encoding:
+            raise ValueError(
+                f"Variable interpolate_pos_encoding {interpolate_pos_encoding} is not equal to rbln_config.interpolate_pos_encoding {self.rbln_config.interpolate_pos_encoding} "
+                f"Please compile again with the correct argument."
+            )
+
+        output = super().forward(pixel_values, return_dict=return_dict)
+        return output
+
+    def _prepare_output(self, output, return_dict):
+        # Prepare model output based on return_dict flag.
+        # This method can be overridden by subclasses to provide task-specific output handling.
+        last_hidden_state = output.pop(0)
+        pooler_output = output.pop(0)
+        vision_config = self.config.vision_config if hasattr(self.config, "vision_config") else self.config
+
+        if self.rbln_config.output_hidden_states:
+            hidden_states = ()
+            num_hidden_layers = vision_config.num_hidden_layers
+            for _ in range(num_hidden_layers + 1):
+                hidden_states += (output.pop(0),)
+        else:
+            hidden_states = None
+
+        if self.rbln_config.output_attentions:
+            attentions = ()
+            num_hidden_layers = vision_config.num_hidden_layers
+            for _ in range(num_hidden_layers):
+                attentions += (output.pop(0),)
+        else:
+            attentions = None
+
+        if not return_dict:
+            return tuple(
+                item for item in (last_hidden_state, pooler_output, hidden_states, attentions) if item is not None
+            )
+        else:
+            return BaseModelOutputWithPooling(
+                last_hidden_state=last_hidden_state,
+                pooler_output=pooler_output,
+                hidden_states=hidden_states,
+                attentions=attentions,
+            )
+
+
+class RBLNCLIPVisionModelWithProjection(RBLNCLIPVisionModel):
+    """
+    RBLN optimized CLIP vision encoder model with projection layer.
+
+    This class extends RBLNCLIPVisionModel with a projection layer for
+    multimodal embedding alignment tasks.
+    """
+
+    def forward(
+        self,
+        pixel_values: torch.FloatTensor,
+        return_dict: bool = True,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        interpolate_pos_encoding: bool = False,
+        **kwargs,
+    ) -> Union[Tuple, CLIPVisionModelOutput]:
+        """
+        Forward pass for the RBLN-optimized CLIP vision encoder model with projection.
+
+        Args:
+            pixel_values (torch.Tensor): The pixel values to the model.
+            return_dict (bool): Whether to return a dictionary of outputs.
+            output_attentions (Optional[bool]): Whether to return attentions.
+            output_hidden_states (Optional[bool]): Whether to return hidden states.
+            interpolate_pos_encoding (bool): Whether to interpolate position encoding.
+
+        Returns:
+            The model outputs. If return_dict=False is passed, returns a tuple of tensors. Otherwise, returns a CLIPVisionModelOutput object.
+        """
+
+        return super().forward(
+            pixel_values=pixel_values,
+            return_dict=return_dict,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            interpolate_pos_encoding=interpolate_pos_encoding,
+            **kwargs,
+        )
+
+    def _prepare_output(self, output, return_dict):
+        # Prepare model output based on return_dict flag.
+        # This method can be overridden by subclasses to provide task-specific output handling.
+
+        image_embeds = output.pop(0) if isinstance(output, (tuple, list)) else output
+        last_hidden_state = output.pop(0)
+
+        vision_config = self.config.vision_config if hasattr(self.config, "vision_config") else self.config
+
+        if self.rbln_config.output_hidden_states:
+            hidden_states = ()
+            num_hidden_layers = vision_config.num_hidden_layers
+            for _ in range(num_hidden_layers + 1):
+                hidden_states += (output.pop(0),)
+        else:
+            hidden_states = None
+
+        if self.rbln_config.output_attentions:
+            attentions = ()
+            num_hidden_layers = vision_config.num_hidden_layers
+            for _ in range(num_hidden_layers):
+                attentions += (output.pop(0),)
+        else:
+            attentions = None
+
+        if not return_dict:
+            return tuple(
+                item for item in (image_embeds, last_hidden_state, hidden_states, attentions) if item is not None
+            )
+
+        else:
+            return CLIPVisionModelOutput(
+                image_embeds=image_embeds,
+                last_hidden_state=last_hidden_state,
+                hidden_states=hidden_states,
+                attentions=attentions,
+            )

optimum/rbln/transformers/models/colpali/__init__.py

@@ -0,0 +1,2 @@
+from .configuration_colpali import RBLNColPaliForRetrievalConfig
+from .modeling_colpali import RBLNColPaliForRetrieval

optimum/rbln/transformers/models/colpali/colpali_architecture.py

@@ -0,0 +1,218 @@
+from typing import List, Optional, Tuple, Union
+
+import torch
+from torch import nn
+from transformers import GemmaForCausalLM, GemmaModel
+
+from ..decoderonly.decoderonly_architecture import RotaryEmbedding, apply_rotary_pos_emb
+
+
+def slice_and_unsqueeze_cos_sin(cos, sin, position_ids):
+    """Slice cos[cache_position], sin[cache_position] vector for the query."""
+    cos = cos[position_ids[0]][None, None, None, :, :]
+    sin = sin[position_ids[0]][None, None, None, :, :]
+
+    return cos, sin
+
+
+class RBLNColPaliForRetrievalWrapper(nn.Module):
+    def __init__(
+        self,
+        causal_lm: GemmaForCausalLM,
+        embedding_proj_layer: nn.Module,
+        max_seq_len: int,
+        output_hidden_states: bool = False,
+    ):
+        super().__init__()
+        self.text_config = causal_lm.config.text_config
+        self.rotary_emb = self.get_rotary_emb(max_seq_len=max_seq_len)
+
+        self.output_hidden_states = output_hidden_states
+        self.language_model = self.convert_to_rbln_language_model(causal_lm.model.language_model, max_seq_len)
+
+        self.num_hidden_layers = getattr(self.text_config, "num_hidden_layers", None)
+        self.embedding_proj_layer = embedding_proj_layer
+
+    def get_rotary_emb(self, max_seq_len):
+        return RotaryEmbedding(config=self.text_config, max_seq_len_cached=max_seq_len)
+
+    def convert_to_rbln_language_model(self, gemma_model: GemmaModel, max_seq_len: int):
+        new_layers = []
+        for layer in gemma_model.layers:
+            new_self_attn = ColPaliAttention(
+                layer.self_attn,
+            )
+            new_layer = ColPaliLayer(layer, new_self_attn)
+            new_layers.append(new_layer)
+
+        new_model = ColPaliModel(
+            gemma_model,
+            new_layers,
+            output_hidden_states=self.output_hidden_states,
+            max_seq_len=max_seq_len,
+        )
+
+        return new_model
+
+    def forward(self, inputs_embeds: torch.Tensor, attention_mask: torch.Tensor, position_ids: torch.Tensor):
+        attention_mask = (1.0 - attention_mask) * torch.finfo(torch.float32).min
+        attention_mask = attention_mask[:, None, None, None, :]
+
+        hidden_states, all_hidden_states = self.language_model(
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            rotary_emb=self.rotary_emb,
+            position_ids=position_ids,
+        )
+        embeddings = self.embedding_proj_layer(hidden_states)
+
+        if self.output_hidden_states:
+            return embeddings, all_hidden_states
+        else:
+            return embeddings
+
+
+class ColPaliModel(nn.Module):
+    def __init__(
+        self, model, layers: List["ColPaliLayer"], output_hidden_states: bool = False, max_seq_len: int = 2048
+    ):
+        super().__init__()
+        self._original_mod = model
+        self.layers = nn.ModuleList(layers)
+        self.output_hidden_states = output_hidden_states
+        self.norm = self._original_mod.norm
+        self.hidden_size = self._original_mod.config.hidden_size
+        self.max_seq_len = max_seq_len
+
+    def forward(
+        self,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        attention_mask: torch.Tensor = None,
+        rotary_emb: Optional[Union[nn.Module, torch.Tensor]] = None,
+        position_ids: Optional[torch.Tensor] = None,
+    ):
+        hidden_states = inputs_embeds * self.hidden_size**0.5
+
+        cos, sin = rotary_emb(hidden_states, self.max_seq_len)  # dtype carrier, max_seq_len
+        cos, sin = slice_and_unsqueeze_cos_sin(cos, sin, position_ids)
+
+        all_hidden_states = () if self.output_hidden_states else None
+        for layer in self.layers:
+            if self.output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            hidden_states = layer(
+                hidden_states=hidden_states,
+                attention_mask=attention_mask,
+                cos=cos,
+                sin=sin,
+            )
+        hidden_states = self.norm(hidden_states)
+
+        if self.output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        return hidden_states, all_hidden_states
+
+
+class ColPaliLayer(nn.Module):
+    def __init__(self, layer, self_attn: "ColPaliAttention"):
+        super().__init__()
+        self._original_mod = layer
+        self.self_attn = self_attn
+        self.mlp = layer.mlp
+        self.input_layernorm = layer.input_layernorm
+        self.post_attention_layernorm = layer.post_attention_layernorm
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        cos: Optional[torch.Tensor] = None,
+        sin: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.FloatTensor]:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+
+        # Self Attention
+        hidden_states = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            cos=cos,
+            sin=sin,
+        )
+        hidden_states = residual + hidden_states
+
+        # Fully Connected
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+
+        return hidden_states
+
+
+class ColPaliAttention(nn.Module):
+    def __init__(self, self_attn):
+        super().__init__()
+        self._original_mod = self_attn
+        self.num_heads = getattr(self._original_mod, "num_heads", None) or getattr(
+            self._original_mod.config, "num_attention_heads"
+        )
+        self.head_dim = self._original_mod.head_dim
+        self.scaling = self.head_dim**-0.5
+
+        if hasattr(self._original_mod, "num_key_value_heads"):
+            self.num_key_value_heads = self._original_mod.num_key_value_heads
+        elif hasattr(self._original_mod, "config") and hasattr(self._original_mod.config, "num_key_value_heads"):
+            self.num_key_value_heads = self._original_mod.config.num_key_value_heads
+        else:
+            self.num_key_value_heads = self.num_heads
+
+        self.__post_init__()
+
+    def __post_init__(self):
+        self.q_proj = self._original_mod.q_proj
+        self.k_proj = self._original_mod.k_proj
+        self.v_proj = self._original_mod.v_proj
+        self.o_proj = self._original_mod.o_proj
+
+    def projection(self, hidden_states) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+
+        return query_states, key_states, value_states
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+        cos: Optional[torch.Tensor] = None,
+        sin: Optional[torch.Tensor] = None,
+    ):
+        batch_size, query_length, _ = hidden_states.size()
+
+        query_states, key_states, value_states = self.projection(hidden_states=hidden_states)
+
+        query_states = query_states.view(batch_size, query_length, 1, self.num_heads, self.head_dim).transpose(1, 3)
+        key_states = key_states.view(batch_size, query_length, 1, self.num_key_value_heads, self.head_dim).transpose(
+            1, 3
+        )
+        value_states = value_states.view(
+            batch_size, query_length, 1, self.num_key_value_heads, self.head_dim
+        ).transpose(1, 3)
+
+        if cos is not None and sin is not None:
+            query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+        attn_weights = torch.matmul(query_states, key_states.transpose(3, 4)) * self.scaling
+        attn_weights = attn_weights + attention_mask
+        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32)
+        attn_output = torch.matmul(attn_weights, value_states)
+        attn_output = attn_output.transpose(1, 3)
+
+        attn_output = attn_output.reshape(batch_size, query_length, -1)
+        attn_output = self.o_proj(attn_output)
+
+        return attn_output

optimum/rbln/transformers/models/colpali/configuration_colpali.py

@@ -0,0 +1,84 @@
+# Copyright 2025 Rebellions Inc. All rights reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import Any, List, Optional, Union
+
+from ....configuration_utils import RBLNModelConfig
+from ....utils.logging import get_logger
+
+
+logger = get_logger(__name__)
+
+
+class RBLNColPaliForRetrievalConfig(RBLNModelConfig):
+    """
+    Configuration class for RBLN ColPali models for document retrieval.
+
+    This class extends RBLNModelConfig with specific configurations for ColPali models,
+    including vision tower settings and multi-sequence length support.
+
+    Example usage:
+        ```python
+        from optimum.rbln import RBLNColPaliForRetrieval, RBLNColPaliForRetrievalConfig
+
+        # Create a configuration object
+        config = RBLNColPaliForRetrievalConfig(
+            max_seq_lens=1152,
+            output_hidden_states=False,
+            tensor_parallel_size=4
+        )
+
+        # Use the configuration with from_pretrained
+        model = RBLNColPaliForRetrieval.from_pretrained(
+            "vidore/colpali-v1.3-hf",
+            export=True,
+            rbln_config=config
+        )
+        ```
+    """
+
+    submodules = ["vision_tower"]
+
+    def __init__(
+        self,
+        batch_size: Optional[int] = None,
+        max_seq_lens: Union[int, List[int]] = None,
+        output_hidden_states: Optional[bool] = None,
+        vision_tower: Optional[RBLNModelConfig] = None,
+        **kwargs: Any,
+    ):
+        """
+        Args:
+            batch_size (Optional[int]): The batch size for the model.
+            vision_tower (Optional[RBLNModelConfig]): Configuration for the vision encoder component.
+            max_seq_lens (Union[int, List[int]]): The maximum sequence lengths for the language model.
+                This can be multiple values, and the model will be compiled for each max_seq_len, allowing selection of the most appropriate max_seq_len at inference time.
+            output_hidden_states (Optional[bool]): Whether to output the hidden states of the language model.
+            vision_tower (Optional[RBLNModelConfig]): Configuration for the vision encoder component.
+            kwargs: Additional arguments passed to the parent RBLNModelConfig.
+        Raises:
+            ValueError: If batch_size is not a positive integer.
+        """
+        super().__init__(**kwargs)
+        self.batch_size = batch_size or 1
+        if not isinstance(self.batch_size, int) or self.batch_size < 0:
+            raise ValueError(f"batch_size must be a positive integer, got {self.batch_size}")
+
+        if self.batch_size != 1:
+            logger.warning("Ignore batch_size for ColPali vision tower. It will be set to 1.")
+
+        self.vision_tower = self.initialize_submodule_config(
+            submodule_config=vision_tower, batch_size=1, force_kwargs=True
+        )
+        self.max_seq_lens = max_seq_lens
+        self.output_hidden_states = output_hidden_states