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

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/transformers/models/t5/modeling_t5.py

@@ -17,6 +17,7 @@ from typing import TYPE_CHECKING, Any, Callable
 
 import torch
 from transformers import AutoModelForTextEncoding, T5EncoderModel, T5ForConditionalGeneration
+from transformers.modeling_outputs import BaseModelOutputWithPastAndCrossAttentions
 
 from ...modeling_generic import RBLNTransformerEncoderForFeatureExtraction
 from ...models.seq2seq import RBLNModelForSeq2SeqLM
@@ -64,7 +65,7 @@ class RBLNT5EncoderModel(RBLNTransformerEncoderForFeatureExtraction):
     """
 
     auto_model_class = AutoModelForTextEncoding
-    rbln_model_input_names = ["input_ids", "attention_mask"]
+    output_class = BaseModelOutputWithPastAndCrossAttentions
 
     @classmethod
     def wrap_model_if_needed(self, model: "PreTrainedModel", rbln_config: RBLNT5EncoderModelConfig):
@@ -74,11 +75,16 @@ class RBLNT5EncoderModel(RBLNTransformerEncoderForFeatureExtraction):
     def update_rbln_config_using_pipe(
         cls, pipe: "RBLNDiffusionMixin", rbln_config: "RBLNDiffusionMixinConfig", submodule_name: str
     ) -> "RBLNDiffusionMixinConfig":
-        submodule_config = getattr(rbln_config, submodule_name)
-        submodule_config.max_seq_len = rbln_config.max_seq_len or 256
-        submodule_config.model_input_names = ["input_ids"]
         return rbln_config
 
+    def forward(self, input_ids=None, attention_mask=None, **kwargs):
+        input_dict = {"input_ids": input_ids.long()}
+        if attention_mask is not None:
+            input_dict["attention_mask"] = attention_mask.long()
+
+        output = super().forward(**input_dict, **kwargs)
+        return output
+
 
 class RBLNT5ForConditionalGeneration(RBLNModelForSeq2SeqLM):
     """

optimum/rbln/utils/runtime_utils.py

@@ -43,6 +43,9 @@ class RBLNPytorchRuntime:
     def __repr__(self) -> str:
         return repr(self.runtime)
 
+    def parameters(self):
+        yield torch.tensor([1.0], dtype=torch.float32, device=torch.device("cpu"))
+
 
 class UnavailableRuntime:
     """

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: optimum-rbln
-Version: 0.8.1a4
+Version: 0.8.1a6
 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.1a6.dist-info}/RECORD

@@ -1,42 +1,52 @@
-optimum/rbln/__init__.py,sha256=qJJTumXhoFnawXGpeGJbAm4J4A9FFwD1SQ2MqcKDXoM,14436
-optimum/rbln/__version__.py,sha256=hdBV0MOKkAsGp6FVqyauDmHCC6gC0y_cyykn1_s49sg,519
+optimum/rbln/__init__.py,sha256=MZCYmY4Y_Zfk0TGo3xK52osHDLZHz4cSdduXZt6RfSI,15316
+optimum/rbln/__version__.py,sha256=eSvwBarNUZxC-xtDA3baJ04Swhx3b8yoQ_vvkfSuClA,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
-optimum/rbln/diffusers/__init__.py,sha256=XL6oKPHbPCV6IVCw3fu0-M9mD2KO_x6unx5kJdAtpVY,6180
-optimum/rbln/diffusers/modeling_diffusers.py,sha256=imodzf_GBCTIdY8bqb9ylzhyV1nt_t8j8-f4F1tKQqo,18777
-optimum/rbln/diffusers/configurations/__init__.py,sha256=Sk_sQVTuTl01RVgYViWknQSLmulxKaISS0w-oPdNoBQ,1164
-optimum/rbln/diffusers/configurations/models/__init__.py,sha256=P3vif5I4wYeol50jzHCZ1ttujuEFZSYJPzUdSF6_jsU,407
+optimum/rbln/diffusers/__init__.py,sha256=_3FoA0uxCdFd6YK9PMaptFmR9XvdB_PcvYR1MkbGpN8,6957
+optimum/rbln/diffusers/modeling_diffusers.py,sha256=oWNgeh9S1W-GqAaiDK6ZgD1ys2MkZbsehdIxPbdZkmQ,19842
+optimum/rbln/diffusers/configurations/__init__.py,sha256=vMRnPY4s-Uju43xP038D2EA18X_mhy2YfsZVpSU-VoA,1322
+optimum/rbln/diffusers/configurations/models/__init__.py,sha256=-4oanQrr9IVznY-ly9ivEnR_zFlfMxzLV6ixLH6TT3w,567
 optimum/rbln/diffusers/configurations/models/configuration_autoencoder_kl.py,sha256=ySetuNq6koleFIZ542zZLTzEEyl_CTul9l12ufWlQ_Y,3218
+optimum/rbln/diffusers/configurations/models/configuration_autoencoder_kl_cosmos.py,sha256=brjtidoogeCc4fYuJ7GxWkJylSu_QGPwJT8l_FQldBI,4060
 optimum/rbln/diffusers/configurations/models/configuration_controlnet.py,sha256=VDO_YFS_QhcHhuRIXQL53JZXEO27yoKHtecq5hd2la8,2637
+optimum/rbln/diffusers/configurations/models/configuration_cosmos_transformer.py,sha256=MnRjn6yI7ST6QFU2MsZqdE2I_FE6vJOOWw4FTTi3KHo,3033
 optimum/rbln/diffusers/configurations/models/configuration_prior_transformer.py,sha256=vE8RsXc27Z4-9k0KEM_vP7AWd5UUYvDgfX1g6nUrPp4,2224
 optimum/rbln/diffusers/configurations/models/configuration_transformer_sd3.py,sha256=TAwHUyVy_9HSEZdXIuFCtrBfNIuYIedklJaCut5wEys,2412
 optimum/rbln/diffusers/configurations/models/configuration_unet_2d_condition.py,sha256=mxcbrOqLMnPpP-jnjSeRWPj2zwPMsgeQSq6LzhG2btA,3630
 optimum/rbln/diffusers/configurations/models/configuration_vq_model.py,sha256=dslGcfCZL_hNeVyjV-4FnCT1POmXuiaLbr6NcQSKgHg,3259
-optimum/rbln/diffusers/configurations/pipelines/__init__.py,sha256=HJlu5lRZJWy4nYjBw3-ed93Pfb5QQmUbCJZKDW1bGH4,1160
+optimum/rbln/diffusers/configurations/pipelines/__init__.py,sha256=RfJXQiYvgGc3Rp7JYk5s0AQd0XB5JCAb37_riGWQAYg,1268
 optimum/rbln/diffusers/configurations/pipelines/configuration_controlnet.py,sha256=odA7UGzAb4bZcey3PeUB0-Z6bLkgCIv_i_eshMUa76A,14298
+optimum/rbln/diffusers/configurations/pipelines/configuration_cosmos.py,sha256=iECUFpM453j3QX54v-idh-_wi1_GfMvtiaawfUOUxl4,4485
 optimum/rbln/diffusers/configurations/pipelines/configuration_kandinsky2_2.py,sha256=PxF5hELuJ4vERzFX-OLT5fPV0MFaD6Ocw7TApbI6dYE,16460
 optimum/rbln/diffusers/configurations/pipelines/configuration_stable_diffusion.py,sha256=IB5Qm7XoTsJ70st3VjuaUqI1OIaA5PGP73JcV7wDznc,6663
-optimum/rbln/diffusers/configurations/pipelines/configuration_stable_diffusion_3.py,sha256=5pDsxol2tm9hYs8u6_6713VwHxCo-iNhAK5G4JVwNwU,7952
+optimum/rbln/diffusers/configurations/pipelines/configuration_stable_diffusion_3.py,sha256=7L-x8nivIifFDi81TDaojAijYVME68gfe32r5-9qdDY,7997
 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/autoencoders/__init__.py,sha256=dg17ZTUsiqTcbIaEE4fqew9uRbao0diQ21PXvRKIqKg,679
+optimum/rbln/diffusers/models/__init__.py,sha256=TgrKGOPDJyEtDLiF22L1q1M8tCal4D9wBcrAtE5Zs5Y,1624
+optimum/rbln/diffusers/models/controlnet.py,sha256=6owledPe9BXhbZOG8lbuuYvpBU0UrQV7zmat6SoMXOM,10585
+optimum/rbln/diffusers/models/autoencoders/__init__.py,sha256=kpINW4bWwy-Q2doPME4nZ8gXRmkK2DRv2kDdbZuQ3m8,738
 optimum/rbln/diffusers/models/autoencoders/autoencoder_kl.py,sha256=UwaYFHXKRJTzJDmfYMC7-xvaWTh7JUDQYD3bRiQs4ZA,8367
-optimum/rbln/diffusers/models/autoencoders/vae.py,sha256=ja9yLhPYGmg1d3Kec6fS-6XgfS0yVJXuVsNDD0X3yHM,4048
+optimum/rbln/diffusers/models/autoencoders/autoencoder_kl_cosmos.py,sha256=jazjvHD2ApEZajH0WNMSzAv3ebtmDPKFVh3qR50vaeo,8316
+optimum/rbln/diffusers/models/autoencoders/vae.py,sha256=ZX6XH9eZxexSSN8sSKFDcvEK9mMuEQNyoalSpOOqQrE,5419
 optimum/rbln/diffusers/models/autoencoders/vq_model.py,sha256=QYbRbnVy4id_P70nIW9wIFeHCP7ITCA3rqNzV1DpdV0,6572
-optimum/rbln/diffusers/models/transformers/__init__.py,sha256=V8rSR7WzHs-i8Cwb_MNxhY2NFbwPgxu24vGtkwl-6tk,706
+optimum/rbln/diffusers/models/transformers/__init__.py,sha256=3oTqAOok-dUR2KealC41CKt36dpKq3IT4kabmHkrCpg,767
 optimum/rbln/diffusers/models/transformers/prior_transformer.py,sha256=SWoeVK--BRMwuXVABNVtonmzJDusx0iz4Q3EAvJ9uN8,5395
+optimum/rbln/diffusers/models/transformers/transformer_cosmos.py,sha256=_YwjOqIGluIt-Nvax8_YJb7BtyZcSf9OAJoYSZsVB2I,12704
 optimum/rbln/diffusers/models/transformers/transformer_sd3.py,sha256=yF7sS0QvawowpV9hR5GeT8DaE8CCp3mj1njHHd9cKTc,6630
 optimum/rbln/diffusers/models/unets/__init__.py,sha256=MaICuK9CWjgzejXy8y2NDrphuEq1rkzanF8u45k6O5I,655
 optimum/rbln/diffusers/models/unets/unet_2d_condition.py,sha256=v3WS9EGKROE_QClXrxC7rmRko1BspAvAbeIfh83LK88,15832
-optimum/rbln/diffusers/pipelines/__init__.py,sha256=5KLZ5LrpMzBya2e_3_PvEoPwG24U8JMexfw_ygZREKc,3140
+optimum/rbln/diffusers/pipelines/__init__.py,sha256=n41EZwEw23tVpe_cgDnuflST1lPiX1y8XQ-3ktTqBIo,3425
 optimum/rbln/diffusers/pipelines/controlnet/__init__.py,sha256=n1Ef22TSeax-kENi_d8K6wGGHSNEo9QkUeygELHgcao,983
 optimum/rbln/diffusers/pipelines/controlnet/multicontrolnet.py,sha256=3S9dogIHW8Bqg5kIlCudhCQG-4g3FcdOPEWhBOf7CJA,4059
 optimum/rbln/diffusers/pipelines/controlnet/pipeline_controlnet.py,sha256=G96bh4D9Cu-w4F9gZBQF6wNzhJQv9kvI34ZFsuEDjSw,35714
 optimum/rbln/diffusers/pipelines/controlnet/pipeline_controlnet_img2img.py,sha256=deGtaqgNumcvCKzKoHZrS-3UZxxWBP0ESizdfvCJlBE,34186
 optimum/rbln/diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl.py,sha256=2w6dmGQuBWqVoocn27z2yMkG7fL7_MVDBcQNJPJsRXU,45300
 optimum/rbln/diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl_img2img.py,sha256=HX56itORMqXLjZcwv25C-_z3JyZn3v6BpfIjsrDO3mE,46640
+optimum/rbln/diffusers/pipelines/cosmos/__init__.py,sha256=h2j6S8IJPVHeNU8qmW9vyXMgHBw0d7kQcuMAA5YoHPU,795
+optimum/rbln/diffusers/pipelines/cosmos/configuration_cosmos_guardrail.py,sha256=kDVnUBBGdumpDj2DaOpo5MSsFvlFIGY6BU1LZaFVqao,3327
+optimum/rbln/diffusers/pipelines/cosmos/cosmos_guardrail.py,sha256=aFKYqaQJtFMULMRQXRMmKnSW3v6gUXKz6cJDbf0kNFs,15100
+optimum/rbln/diffusers/pipelines/cosmos/pipeline_cosmos_text2world.py,sha256=m1RQr5C1G5GUI1bXI7BvZN_WWl75LvYFyC2ApoEaUAI,3910
+optimum/rbln/diffusers/pipelines/cosmos/pipeline_cosmos_video2world.py,sha256=icJXQbzCZygDuDs_PDJtq7qQr9uX3q-ljJOT6sXB8JQ,3917
 optimum/rbln/diffusers/pipelines/kandinsky2_2/__init__.py,sha256=I4YQq2HfA3xONbWsdJ870IEJPyLWeCDDG-UCJsu9YO8,1035
 optimum/rbln/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2.py,sha256=2vgZZt0JrZlbHbUlBHrT0zKvYhuX1a4vwy3mxTPHisM,1335
 optimum/rbln/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_combined.py,sha256=cuN7w5LsnkSy0WUaCBPxVmYOJiBMihdh-KM7iV0tHL0,8149
@@ -61,11 +71,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 +95,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 +158,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
@@ -160,7 +174,7 @@ optimum/rbln/transformers/models/siglip/configuration_siglip.py,sha256=fXL4g985o
 optimum/rbln/transformers/models/siglip/modeling_siglip.py,sha256=se74bZTSTjqcV5UoBowSHVUTrE6CapoEcqgzqMs7Q9Y,7827
 optimum/rbln/transformers/models/t5/__init__.py,sha256=R1Q8Z1vaIdx4rDjeCmm_ZMSgewWaqaI0l93AHwewtew,818
 optimum/rbln/transformers/models/t5/configuration_t5.py,sha256=cOa2NAAocNSSECTtLpOsVVLfn4-7l7K7jz3nf6CSNCs,912
-optimum/rbln/transformers/models/t5/modeling_t5.py,sha256=SasIteIdpM-7TG6iPiRVlchd_3jaG7phFvdQJ_3VXjo,4924
+optimum/rbln/transformers/models/t5/modeling_t5.py,sha256=pdAWBLVknTzbma0Ij-VQ2Qve-frPjxL-AwMyU-zouPY,5123
 optimum/rbln/transformers/models/t5/t5_architecture.py,sha256=X_9X4QRhkiiMrwFHv3mzER3yGmF9oQ2U-HdH6jbwVmw,9824
 optimum/rbln/transformers/models/time_series_transformer/__init__.py,sha256=xJaFWQawlwtv4H5tVFcY1pxLYzjHtMAlLq6nXysdkN8,1243
 optimum/rbln/transformers/models/time_series_transformer/configuration_time_series_transformer.py,sha256=FCzESVMNtpJ4Me-Vxr4i-naUfiR3S7YXWH_szLuoYeA,1390
@@ -188,10 +202,10 @@ optimum/rbln/utils/hub.py,sha256=Z_R9Ic9VAew8bUmlaAlxZf5JGMDBivHvvFRI557pILY,419
 optimum/rbln/utils/import_utils.py,sha256=fpOERIIxXm-cDYGn1NN6c7aWDPQYVitPQW2MiyZ9NEY,5471
 optimum/rbln/utils/logging.py,sha256=VKKBmlQSdg6iZCGmAXaWYiW67K84jyp1QJhLQSSjPPE,3453
 optimum/rbln/utils/model_utils.py,sha256=4k5879Kh75m3x_vS4-qOGfqsOiAvc2kdNFFfvsFvz3k,1748
-optimum/rbln/utils/runtime_utils.py,sha256=LoKNK3AQNV_BSScstIZWjICkJf265MnUgy360BOocVI,5454
+optimum/rbln/utils/runtime_utils.py,sha256=kTcDSNGOMFifRclVrNtMXqEUJxcoUmvEBYFvl1fORWo,5564
 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.1a6.dist-info/METADATA,sha256=d5MR8fiDCLMklaKLRimUqEnqA520uVvxqKxA3fsrWIU,5299
+optimum_rbln-0.8.1a6.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+optimum_rbln-0.8.1a6.dist-info/licenses/LICENSE,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
+optimum_rbln-0.8.1a6.dist-info/RECORD,,