optimum-rbln 0.8.1a4__py3-none-any.whl → 0.8.1a5__py3-none-any.whl

optimum/rbln/__init__.py

@@ -70,6 +70,8 @@ _import_structure = {
         "RBLNCLIPVisionModelConfig",
         "RBLNCLIPVisionModelWithProjection",
         "RBLNCLIPVisionModelWithProjectionConfig",
+        "RBLNColPaliForRetrieval",
+        "RBLNColPaliForRetrievalConfig",
         "RBLNDecoderOnlyModelForCausalLM",
         "RBLNDecoderOnlyModelForCausalLMConfig",
         "RBLNDistilBertForQuestionAnswering",
@@ -297,6 +299,8 @@ if TYPE_CHECKING:
         RBLNCLIPVisionModelConfig,
         RBLNCLIPVisionModelWithProjection,
         RBLNCLIPVisionModelWithProjectionConfig,
+        RBLNColPaliForRetrieval,
+        RBLNColPaliForRetrievalConfig,
         RBLNDecoderOnlyModelForCausalLM,
         RBLNDecoderOnlyModelForCausalLMConfig,
         RBLNDistilBertForQuestionAnswering,

optimum/rbln/__version__.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.8.1a4'
-__version_tuple__ = version_tuple = (0, 8, 1, 'a4')
+__version__ = version = '0.8.1a5'
+__version_tuple__ = version_tuple = (0, 8, 1, 'a5')

optimum/rbln/diffusers/models/controlnet.py

@@ -16,7 +16,7 @@ from typing import TYPE_CHECKING, Dict, Optional, Union
 
 import torch
 from diffusers import ControlNetModel
-from diffusers.models.controlnet import ControlNetOutput
+from diffusers.models.controlnets.controlnet import ControlNetOutput
 from transformers import PretrainedConfig
 
 from ...configuration_utils import RBLNCompileConfig, RBLNModelConfig

optimum/rbln/transformers/__init__.py

@@ -50,6 +50,8 @@ _import_structure = {
         "RBLNBlip2QFormerModelConfig",
         "RBLNBlip2VisionModel",
         "RBLNBlip2VisionModelConfig",
+        "RBLNColPaliForRetrieval",
+        "RBLNColPaliForRetrievalConfig",
         "RBLNCLIPTextModel",
         "RBLNCLIPTextModelConfig",
         "RBLNCLIPTextModelWithProjection",

optimum/rbln/transformers/models/__init__.py

@@ -69,6 +69,10 @@ _import_structure = {
         "RBLNCLIPVisionModelWithProjection",
         "RBLNCLIPVisionModelWithProjectionConfig",
     ],
+    "colpali": [
+        "RBLNColPaliForRetrieval",
+        "RBLNColPaliForRetrievalConfig",
+    ],
     "distilbert": [
         "RBLNDistilBertForQuestionAnswering",
         "RBLNDistilBertForQuestionAnsweringConfig",
@@ -193,6 +197,10 @@ if TYPE_CHECKING:
         RBLNCLIPVisionModelWithProjection,
         RBLNCLIPVisionModelWithProjectionConfig,
     )
+    from .colpali import (
+        RBLNColPaliForRetrieval,
+        RBLNColPaliForRetrievalConfig,
+    )
     from .decoderonly import (
         RBLNDecoderOnlyModelForCausalLM,
         RBLNDecoderOnlyModelForCausalLMConfig,

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,221 @@
+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
+        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, 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,68 @@
+# 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 List, Optional, Union
+
+from ....configuration_utils import RBLNModelConfig
+
+
+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,
+        max_seq_lens: Union[int, List[int]] = None,
+        output_hidden_states: Optional[bool] = None,
+        vision_tower: Optional[RBLNModelConfig] = None,
+        **kwargs,
+    ):
+        """
+        Args:
+            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.
+            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+        Raises:
+            ValueError: If batch_size is not a positive integer.
+        """
+        super().__init__(**kwargs)
+        self.vision_tower = vision_tower
+        self.max_seq_lens = max_seq_lens
+        self.output_hidden_states = output_hidden_states

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

@@ -0,0 +1,383 @@
+# 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.
+
+import bisect
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Optional, Union
+
+import torch
+from transformers import (
+    PretrainedConfig,
+    PreTrainedModel,
+)
+from transformers.modeling_outputs import BaseModelOutputWithPooling
+from transformers.modeling_utils import no_init_weights
+from transformers.models.colpali.modeling_colpali import ColPaliForRetrievalOutput
+from transformers.models.paligemma.modeling_paligemma import PaliGemmaMultiModalProjector
+
+from ....configuration_utils import RBLNCompileConfig, RBLNModelConfig
+from ....modeling import RBLNModel
+from .colpali_architecture import RBLNColPaliForRetrievalWrapper
+
+
+if TYPE_CHECKING:
+    from transformers import (
+        AutoFeatureExtractor,
+        AutoProcessor,
+        AutoTokenizer,
+        PretrainedConfig,
+    )
+
+
+class LoopVisionTower:
+    def __init__(self, vision_tower: RBLNModel) -> None:
+        self.vision_tower = vision_tower
+
+    def forward(self, pixel_values, **kwargs):
+        batch_size = pixel_values.shape[0]
+        outputs = []
+        for i in range(batch_size):
+            outputs.append(self.vision_tower(pixel_values[i : i + 1]))
+
+        last_hidden_states = [output.last_hidden_state for output in outputs]
+        last_hidden_states = torch.cat(last_hidden_states, dim=0)
+
+        return BaseModelOutputWithPooling(
+            last_hidden_state=last_hidden_states,
+        )
+
+    def __call__(self, *args: Any, **kwds: Any) -> Any:
+        return self.forward(*args, **kwds)
+
+    def __repr__(self) -> str:
+        return repr(self.vision_tower)
+
+
+class LoopLanguageModel:
+    def __init__(self, language_model: RBLNModel, rbln_config: RBLNModelConfig) -> None:
+        self.language_model = language_model
+        self.rbln_config = rbln_config
+
+    def prepare_inputs(self, inputs_embeds: torch.Tensor, attention_mask: torch.Tensor):
+        input_len = inputs_embeds.shape[1]
+        idx = bisect.bisect_left(self.rbln_config.max_seq_lens, input_len)
+        if idx == len(self.rbln_config.max_seq_lens):
+            raise ValueError(
+                f"Required seq_len({input_len}) is larger than available max_seq_lens({self.rbln_config.max_seq_lens})."
+            )
+        else:
+            max_seq_len = self.rbln_config.max_seq_lens[idx]
+
+        inputs_embed = torch.nn.functional.pad(inputs_embeds, (0, 0, 0, max_seq_len - input_len))
+        attn_mask = torch.nn.functional.pad(attention_mask, (0, max_seq_len - input_len)).to(torch.float32)
+        position_ids = torch.arange(max_seq_len, dtype=torch.int32).view(1, -1)
+
+        return inputs_embed, attn_mask, position_ids
+
+    def forward(self, inputs_embeds: torch.Tensor, attention_mask: torch.Tensor, **kwargs):
+        padded_inputs_embed, padded_attn_mask, padded_position_ids = self.prepare_inputs(inputs_embeds, attention_mask)
+        input_batch_size = inputs_embeds.shape[0]
+        input_seq_len = inputs_embeds.shape[1]
+
+        all_embeddings = []
+        all_hidden_states = []
+        for i in range(input_batch_size):
+            outputs = self.language_model(
+                inputs_embeds=padded_inputs_embed[i : i + 1],
+                attention_mask=padded_attn_mask[i : i + 1],
+                position_ids=padded_position_ids,
+            )
+
+            if self.rbln_config.output_hidden_states:
+                embedding = outputs[0]
+                hidden_states = outputs[1:]
+            else:
+                embedding = outputs
+                hidden_states = None
+
+            all_embeddings.append(embedding)
+            all_hidden_states.append(hidden_states)
+
+        embeddings = torch.cat(all_embeddings, dim=0)[:, :input_seq_len]
+        if self.rbln_config.output_hidden_states:
+            hidden_states = [
+                torch.cat(
+                    [batch_hidden_states[layer_idx][:, :input_seq_len] for batch_hidden_states in all_hidden_states],
+                    dim=0,
+                )
+                for layer_idx in range(len(all_hidden_states[0]))
+            ]
+            return embeddings, tuple(hidden_states)
+        else:
+            return embeddings
+
+    def __call__(self, *args: Any, **kwds: Any) -> Any:
+        return self.forward(*args, **kwds)
+
+    def __repr__(self) -> str:
+        return repr(self.language_model)
+
+
+class RBLNColPaliForRetrieval(RBLNModel):
+    """
+    The ColPali Model transformer for document retrieval using vision-language models.
+    This model inherits from [`RBLNModel`]. Check the superclass documentation for the generic methods the library implements for all its models.
+
+    A class to convert and run pre-trained transformers based ColPaliForRetrieval model on RBLN devices.
+    It implements the methods to convert a pre-trained transformers ColPaliForRetrieval model into a RBLN transformer model by:
+
+    - transferring the checkpoint weights of the original into an optimized RBLN graph,
+    - compiling the resulting graph using the RBLN compiler.
+
+    **Configuration:**
+    This model uses [`RBLNColPaliForRetrievalConfig`] for configuration. When calling methods like `from_pretrained` or `from_model`,
+    the `rbln_config` parameter should be an instance of [`RBLNColPaliForRetrievalConfig`] or a dictionary conforming to its structure.
+
+    See the [`RBLNColPaliForRetrievalConfig`] class for all available configuration options.
+
+    Examples:
+        ```python
+        from optimum.rbln import RBLNColPaliForRetrieval
+
+        # Simple usage using rbln_* arguments
+        # `max_seq_lens` is automatically inferred from the model config
+        model = RBLNColPaliForRetrieval.from_pretrained(
+            "vidore/colpali-v1.3-hf",
+            export=True,
+            rbln_max_seq_lens=1152,
+        )
+
+        # Using a config dictionary
+        rbln_config = {
+            "max_seq_lens": 1152,
+            "output_hidden_states": False,
+        }
+        model = RBLNColPaliForRetrieval.from_pretrained(
+            "vidore/colpali-v1.3-hf",
+            export=True,
+            rbln_config=rbln_config
+        )
+
+        # Using a RBLNColPaliForRetrievalConfig instance (recommended for type checking)
+        from optimum.rbln import RBLNColPaliForRetrievalConfig
+
+        config = RBLNColPaliForRetrievalConfig(
+            max_seq_lens=1152,
+            output_hidden_states=False,
+            tensor_parallel_size=4
+        )
+        model = RBLNColPaliForRetrieval.from_pretrained(
+            "vidore/colpali-v1.3-hf",
+            export=True,
+            rbln_config=config
+        )
+        ```
+    """
+
+    auto_model_class = None
+    _rbln_submodules = [
+        {"name": "vision_tower"},
+    ]
+
+    def __post_init__(self, **kwargs):
+        self.vision_tower = LoopVisionTower(self.rbln_submodules[0])
+        self.language_model = LoopLanguageModel(self.model[0], self.rbln_config)
+
+        artifacts = torch.load(self.model_save_dir / self.subfolder / "torch_artifacts.pth", weights_only=False)
+        self.embed_tokens = self._create_embedding_layer()
+        self.embed_tokens.load_state_dict(artifacts["embed_tokens"])
+        self.multi_modal_projector = self._create_multi_modal_projector()
+        self.multi_modal_projector.load_state_dict(artifacts["multi_modal_projector"])
+
+        return super().__post_init__(**kwargs)
+
+    def _create_embedding_layer(self):
+        with no_init_weights():
+            embed_tokens = torch.nn.Embedding(
+                self.config.text_config.vocab_size,
+                self.config.text_config.hidden_size,
+                self.config.text_config.pad_token_id,
+            )
+        return embed_tokens
+
+    def _create_multi_modal_projector(self):
+        with no_init_weights():
+            multi_modal_projector = PaliGemmaMultiModalProjector(self.config.vlm_config)
+        return multi_modal_projector
+
+    @classmethod
+    def wrap_model_if_needed(cls, model: "PreTrainedModel", rbln_config: RBLNModelConfig):
+        return RBLNColPaliForRetrievalWrapper(
+            causal_lm=model.vlm.language_model,
+            embedding_proj_layer=model.embedding_proj_layer,
+            max_seq_len=max(rbln_config.max_seq_lens),
+            output_hidden_states=rbln_config.output_hidden_states,
+        )
+
+    @classmethod
+    def save_torch_artifacts(
+        cls,
+        model: "PreTrainedModel",
+        save_dir_path: Path,
+        subfolder: str,
+        rbln_config: RBLNModelConfig,
+    ):
+        save_dict = {}
+        save_dict["embed_tokens"] = model.vlm.get_input_embeddings().state_dict()
+        save_dict["multi_modal_projector"] = model.vlm.multi_modal_projector.state_dict()
+        torch.save(save_dict, save_dir_path / subfolder / "torch_artifacts.pth")
+
+    @classmethod
+    def _update_rbln_config(
+        cls,
+        preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]],
+        model: Optional["PreTrainedModel"] = None,
+        model_config: Optional["PretrainedConfig"] = None,
+        rbln_config: Optional[RBLNModelConfig] = None,
+    ) -> RBLNModelConfig:
+        hidden_size = model_config.vlm_config.text_config.hidden_size
+        if rbln_config.max_seq_lens is None:
+            rbln_config.max_seq_lens = [model_config.vlm_config.text_config.max_position_embeddings]
+        if isinstance(rbln_config.max_seq_lens, int):
+            rbln_config.max_seq_lens = [rbln_config.max_seq_lens]
+        rbln_config.max_seq_lens = sorted(set(rbln_config.max_seq_lens))
+
+        if rbln_config.output_hidden_states is None:
+            rbln_config.output_hidden_states = model_config.vlm_config.text_config.output_hidden_states
+
+        input_infos = []
+        for max_seq_len in rbln_config.max_seq_lens:
+            input_info = [
+                ("inputs_embeds", [1, max_seq_len, hidden_size], "float32"),
+                ("attention_mask", [1, max_seq_len], "float32"),
+                ("position_ids", [1, max_seq_len], "int32"),
+            ]
+            input_infos.append(input_info)
+
+        rbln_compile_config = RBLNCompileConfig(input_info=input_infos)
+        rbln_config.set_compile_cfgs([rbln_compile_config])
+
+        return rbln_config
+
+    @classmethod
+    def from_model(cls, model: "PreTrainedModel", *args, **kwargs):
+        if not hasattr(model, "vision_tower"):
+            model.vision_tower = model.vlm.vision_tower
+            del model.vlm.vision_tower
+        model = super().from_model(model, *args, **kwargs)
+        return model
+
+    @classmethod
+    def get_pytorch_model(cls, *args, **kwargs):
+        model = super().get_pytorch_model(*args, **kwargs)
+        model.vision_tower = model.vlm.vision_tower
+        del model.vlm.vision_tower
+
+        return model
+
+    def get_image_features(self, pixel_values: torch.Tensor):
+        # Projects the last hidden state from the vision model into language model space.
+        # Args:
+        #     pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`)
+        #        The tensors corresponding to the input images.
+        # Returns:
+        #     image_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`).
+
+        vision_outputs = self.vision_tower(pixel_values).last_hidden_state
+        image_features = self.multi_modal_projector(vision_outputs)
+        image_features = image_features / (self.config.text_config.hidden_size**0.5)
+        return image_features
+
+    def _preprocess_inputs(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        **kwargs,
+    ):
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        # Replace image id woth PAD if the image token if OOV, to avoid index-errors
+        if input_ids is not None and self.config.vlm_config.image_token_index >= self.config.text_config.vocab_size:
+            special_image_mask = input_ids == self.config.vlm_config.image_token_index
+            llm_input_ids = input_ids.clone()
+            llm_input_ids[special_image_mask] = 0
+        else:
+            llm_input_ids = input_ids
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(llm_input_ids)
+
+        # Merge text and images
+        image_features = None
+        if pixel_values is not None:
+            image_features = self.get_image_features(pixel_values)
+            special_image_mask = (input_ids == self.config.vlm_config.image_token_index).unsqueeze(-1)
+            special_image_mask = special_image_mask.expand_as(inputs_embeds).to(inputs_embeds.device)
+
+            image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype)
+            inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features)
+
+        return inputs_embeds, image_features
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs,
+    ) -> ColPaliForRetrievalOutput:
+        if pixel_values is not None:
+            pixel_values = pixel_values.to(dtype=self.dtype)
+
+        if output_attentions:
+            raise ValueError("output_attentions is not supported for RBLNColPaliForRetrieval")
+
+        if output_hidden_states is not None and 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."
+            )
+
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        inputs_embeds, image_features = self._preprocess_inputs(
+            input_ids=input_ids, inputs_embeds=inputs_embeds, pixel_values=pixel_values
+        )
+
+        # Embedding_proj_layer is fused on the bottom of the language model.
+        outputs = self.language_model(inputs_embeds=inputs_embeds, attention_mask=attention_mask)
+
+        embeddings = outputs if not self.rbln_config.output_hidden_states else outputs[0]
+        hidden_states = None if not self.rbln_config.output_hidden_states else outputs[1]
+
+        # L2 normalization
+        embeddings = embeddings / embeddings.norm(dim=-1, keepdim=True)  # (batch_size, sequence_length, dim)
+
+        if attention_mask is not None:
+            embeddings = embeddings * attention_mask.unsqueeze(-1)  # (batch_size, sequence_length, dim)
+
+        if not return_dict:
+            return (embeddings, hidden_states, image_features)
+        else:
+            return ColPaliForRetrievalOutput(
+                embeddings=embeddings,
+                hidden_states=hidden_states,
+                image_hidden_states=image_features,
+            )

optimum/rbln/transformers/models/qwen2_5_vl/qwen2_5_vl_architecture.py

@@ -79,7 +79,7 @@ class Qwen2_5_VLVisionFullAttention(nn.Module):
         super().__init__()
         self._origin_model = model
         self.num_heads = model.num_heads
-        self.head_dim = model.head_dim
+        self.head_dim = getattr(model, "head_dim", model.proj.in_features // model.num_heads)
         self.qkv = model.qkv
         self.proj = model.proj
 
@@ -114,7 +114,7 @@ class Qwen2_5_VLVisionWindowAttention(nn.Module):
         super().__init__()
         self._origin_model = model
         self.num_heads = model.num_heads
-        self.head_dim = model.head_dim
+        self.head_dim = getattr(model, "head_dim", model.proj.in_features // model.num_heads)
         self.qkv = model.qkv
         self.proj = model.proj
         self.window_seq_len = window_seq_len

{optimum_rbln-0.8.1a4.dist-info → optimum_rbln-0.8.1a5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: optimum-rbln
-Version: 0.8.1a4
+Version: 0.8.1a5
 Summary: Optimum RBLN is the interface between the HuggingFace Transformers and Diffusers libraries and RBLN accelerators. It provides a set of tools enabling easy model loading and inference on single and multiple rbln device settings for different downstream tasks.
 Project-URL: Homepage, https://rebellions.ai
 Project-URL: Documentation, https://docs.rbln.ai
@@ -25,9 +25,9 @@ Requires-Python: <3.13,>=3.9
 Requires-Dist: accelerate>=1.0.1
 Requires-Dist: diffusers==0.34.0
 Requires-Dist: packaging>=24.1
-Requires-Dist: torch==2.6.0
-Requires-Dist: torchaudio<=2.6.0
-Requires-Dist: torchvision<=0.21.0
+Requires-Dist: torch==2.7.0
+Requires-Dist: torchaudio<=2.7.0
+Requires-Dist: torchvision<=0.22.0
 Requires-Dist: transformers==4.51.3
 Description-Content-Type: text/markdown
 

{optimum_rbln-0.8.1a4.dist-info → optimum_rbln-0.8.1a5.dist-info}/RECORD

@@ -1,5 +1,5 @@
-optimum/rbln/__init__.py,sha256=qJJTumXhoFnawXGpeGJbAm4J4A9FFwD1SQ2MqcKDXoM,14436
-optimum/rbln/__version__.py,sha256=hdBV0MOKkAsGp6FVqyauDmHCC6gC0y_cyykn1_s49sg,519
+optimum/rbln/__init__.py,sha256=Z5GM8hmc_cgNzhdfOAKbAQr-vFP24kC-IbiRaIOIxxE,14584
+optimum/rbln/__version__.py,sha256=Ln2yvKWXaraNsP7hCs26LOEd96BBrL7JNrmQ42n0dqA,519
 optimum/rbln/configuration_utils.py,sha256=o5oer7fBdE-MHLGNXoP35FjmuQbMmjEIDv0QE_k3kpo,32336
 optimum/rbln/modeling.py,sha256=ZlJ_tOCWiFjDIlwJ_B_HOCO0kBduWrBAbW9VSEVIAFg,12088
 optimum/rbln/modeling_base.py,sha256=5fUb1FaxfjApzJIkT8-SrPhuygGo_1Uc0i7UedawOeE,23393
@@ -20,7 +20,7 @@ optimum/rbln/diffusers/configurations/pipelines/configuration_stable_diffusion.p
 optimum/rbln/diffusers/configurations/pipelines/configuration_stable_diffusion_3.py,sha256=5pDsxol2tm9hYs8u6_6713VwHxCo-iNhAK5G4JVwNwU,7952
 optimum/rbln/diffusers/configurations/pipelines/configuration_stable_diffusion_xl.py,sha256=zg7aRyp8jYJuAnb_dTg6HdACCcAvhv1jX2FhEfRD6V0,7114
 optimum/rbln/diffusers/models/__init__.py,sha256=mkCvJyH1KcwrsUvYSq_bVC79oOfyqtBSFDyPS1_48wA,1478
-optimum/rbln/diffusers/models/controlnet.py,sha256=yKPQTO2jwb9VRMagiqzEXMAwJfcyAnfqMD7Lc8AOsr8,10573
+optimum/rbln/diffusers/models/controlnet.py,sha256=6owledPe9BXhbZOG8lbuuYvpBU0UrQV7zmat6SoMXOM,10585
 optimum/rbln/diffusers/models/autoencoders/__init__.py,sha256=dg17ZTUsiqTcbIaEE4fqew9uRbao0diQ21PXvRKIqKg,679
 optimum/rbln/diffusers/models/autoencoders/autoencoder_kl.py,sha256=UwaYFHXKRJTzJDmfYMC7-xvaWTh7JUDQYD3bRiQs4ZA,8367
 optimum/rbln/diffusers/models/autoencoders/vae.py,sha256=ja9yLhPYGmg1d3Kec6fS-6XgfS0yVJXuVsNDD0X3yHM,4048
@@ -61,11 +61,11 @@ optimum/rbln/ops/flash_attn.py,sha256=z39DJZSk94630ueoOCkiybxR5gzvNR-SRADHs0F6pz
 optimum/rbln/ops/kv_cache_update.py,sha256=HjnHBR-oFrJQibsVnkYb0P5_-wEma8jl0mkjkylwakU,1270
 optimum/rbln/ops/linear.py,sha256=1_7Hg-9wXxhu97fqPobotLQx17k7VPeSSL91_9Z7EDg,1018
 optimum/rbln/ops/sliding_window_attn.py,sha256=EQrV_yRGc5z6kvwEsAcLP028bJWkQg2UPI3xubt9skU,3487
-optimum/rbln/transformers/__init__.py,sha256=fE-kzDnWj0ueAG-xDrIKdBX59wCE__8m86uBMBOEb9g,9031
+optimum/rbln/transformers/__init__.py,sha256=MF7OaGf-KI9rz4EOzejxHTDYUB3RO2L02BquTe0PXmI,9107
 optimum/rbln/transformers/configuration_generic.py,sha256=kNhPWtzF0IovUnrsXfxXdXITqgpfCAAedjfB6jSAhEg,5131
 optimum/rbln/transformers/modeling_generic.py,sha256=u1JzjWcPsQgH_rqBzRVr582NARqOk7XVKgY4CdEfXe8,12228
 optimum/rbln/transformers/modeling_rope_utils.py,sha256=6Zg3r-TeUk4WQAlr95pqfhuoAD_RQ4njT1rbO9uPL0Q,14379
-optimum/rbln/transformers/models/__init__.py,sha256=-rc_00p4d58cdM2ylmgURxoAGKgIRF7X7r6z1w6h3mo,10061
+optimum/rbln/transformers/models/__init__.py,sha256=VVQJgpUUnN4MPAQlOsxsw63w7WPK05ggFfRkGYuZFJQ,10266
 optimum/rbln/transformers/models/audio_spectrogram_transformer/__init__.py,sha256=I2vL4lrzbT5p4eJcH-EKHzEfcPkj_XVsie7jb9q6yic,775
 optimum/rbln/transformers/models/audio_spectrogram_transformer/configuration_audio_spectrogram_transformer.py,sha256=z7LJiVJPmnlCM3mcyhPJP8AufSrxO_dsPeJ51onq-Nc,833
 optimum/rbln/transformers/models/audio_spectrogram_transformer/modeling_audio_spectrogram_transformer.py,sha256=FIKEVWpIt6-JQX9B_rAfCrAPqdUHtR2i8D_X2k7639E,1498
@@ -85,6 +85,10 @@ optimum/rbln/transformers/models/blip_2/modeling_blip_2.py,sha256=gx9pPXQfaIjDUN
 optimum/rbln/transformers/models/clip/__init__.py,sha256=TLeXDqcFK6M6v9x7Xr64kBbqGu3hFHM7p754dQ8UVQc,938
 optimum/rbln/transformers/models/clip/configuration_clip.py,sha256=mgtR_lS1_g5vAh_wWarff3-pwM_tzzRAWm7XkfhGwmo,3019
 optimum/rbln/transformers/models/clip/modeling_clip.py,sha256=0u1JTlO47qoH_-qxWGvXLc67whddLzcuLoMB5KaMh94,7285
+optimum/rbln/transformers/models/colpali/__init__.py,sha256=n3rueXT_oC0N8myoZiic0YkVK24CW5hZBPa-0L8so6Y,119
+optimum/rbln/transformers/models/colpali/colpali_architecture.py,sha256=bWG7TehWRZkTh2y6mGkpd85_onWAyiyKdaQC9TFsy3E,8065
+optimum/rbln/transformers/models/colpali/configuration_colpali.py,sha256=yPzLYON6qRJlBkzxFfIBzBWd2KjYWvdClO4iAqd_V7E,2609
+optimum/rbln/transformers/models/colpali/modeling_colpali.py,sha256=jzvJCBrrCXSpjfmJ3O-VvPNFGWGaNbpOV09JwLPAZWs,15757
 optimum/rbln/transformers/models/decoderonly/__init__.py,sha256=vQYZDDdoddwA7yKc5zzrq2Zs9sax-0p8rNF_aYfF4bk,1006
 optimum/rbln/transformers/models/decoderonly/configuration_decoderonly.py,sha256=cakn8RGo8gS3nmXdEqOfC2xUBOMGInROgLEbCOoLFR0,13398
 optimum/rbln/transformers/models/decoderonly/decoderonly_architecture.py,sha256=YAn8J_lIq4IS-HM_gbi5Qov8_osxhWtBr5z_28QRbGM,49667
@@ -144,7 +148,7 @@ optimum/rbln/transformers/models/qwen2/qwen2_architecture.py,sha256=XlNAMYAcDLoh
 optimum/rbln/transformers/models/qwen2_5_vl/__init__.py,sha256=rAW3DKQUzGL6EMwa5r1iLu94yhpiZpk6zfoD7TtYXrc,865
 optimum/rbln/transformers/models/qwen2_5_vl/configuration_qwen2_5_vl.py,sha256=U3ngIfkA58itqQZqTf-gbISMPoV7ipDttI7V2uwK_18,4155
 optimum/rbln/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py,sha256=Q4U-avMkby-CunNXEERqvRZx9duC5i-6UmfF1376ciU,26336
-optimum/rbln/transformers/models/qwen2_5_vl/qwen2_5_vl_architecture.py,sha256=PAQz__9o_f5phlozhhXAB8JErBlS1jc4FYZkZkSYJuI,7312
+optimum/rbln/transformers/models/qwen2_5_vl/qwen2_5_vl_architecture.py,sha256=oU4MyNeDHzqD3dl1DgwrMev07yvMFhl_hXvV6tRdXCo,7422
 optimum/rbln/transformers/models/resnet/__init__.py,sha256=0QqtEQF1IMYgEmmfXMGarCDS8kJB5tzODfwTEzDVZRg,837
 optimum/rbln/transformers/models/resnet/configuration_resnet.py,sha256=KQd887jgNOl_Am3b407P2OvKtzkkeBS1cEhCfiN0tJg,769
 optimum/rbln/transformers/models/resnet/modeling_resnet.py,sha256=E8vg3Rw_KsHt6vaOg0ungZD7sXe0T4OMP0X8NFG1EXI,816
@@ -191,7 +195,7 @@ optimum/rbln/utils/model_utils.py,sha256=4k5879Kh75m3x_vS4-qOGfqsOiAvc2kdNFFfvsF
 optimum/rbln/utils/runtime_utils.py,sha256=LoKNK3AQNV_BSScstIZWjICkJf265MnUgy360BOocVI,5454
 optimum/rbln/utils/save_utils.py,sha256=hG5uOtYmecSXZuGTvCXsTM-SiyZpr5q3InUGCCq_jzQ,3619
 optimum/rbln/utils/submodule.py,sha256=w5mgPgncI740gVKMu3S-69DGNdUSI0bTZxegQGcZ98Y,5011
-optimum_rbln-0.8.1a4.dist-info/METADATA,sha256=jo7yVVPhX8QJJK0WE1x2ReG_VbuNiyhAkAPj9Um90A8,5299
-optimum_rbln-0.8.1a4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-optimum_rbln-0.8.1a4.dist-info/licenses/LICENSE,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
-optimum_rbln-0.8.1a4.dist-info/RECORD,,
+optimum_rbln-0.8.1a5.dist-info/METADATA,sha256=yetswBiXM1Cce75lQOgrUw3pNMuaxt6XoaclWnDlGIE,5299
+optimum_rbln-0.8.1a5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+optimum_rbln-0.8.1a5.dist-info/licenses/LICENSE,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
+optimum_rbln-0.8.1a5.dist-info/RECORD,,