optimum-rbln 0.1.12__py3-none-any.whl → 0.1.13__py3-none-any.whl

optimum/rbln/transformers/models/phi/modeling_phi.py

@@ -21,28 +21,18 @@
 # copied, modified, or distributed without prior written permission
 # from Rebellions Inc.
 
-import inspect
-import logging
-from typing import TYPE_CHECKING, Any, Callable
-
-from transformers import PhiForCausalLM
-
-from ..decoderonly import RBLNDecoderOnlyModelForCausalLM
+from ....utils import logging
+from ...models.decoderonly import RBLNDecoderOnlyModelForCausalLM
 from .phi_architecture import PhiWrapper
 
 
-if TYPE_CHECKING:
-    from transformers import PreTrainedModel
-
-    from ....modeling_config import RBLNConfig
-
-logger = logging.getLogger(__name__)
+logger = logging.get_logger(__name__)
 
 
 class RBLNPhiForCausalLM(RBLNDecoderOnlyModelForCausalLM):
     """
     The Phi Model transformer with a language modeling head (linear layer) on top.
-    This model inherits from [`RBLNMultiModel`]. Check the superclass documentation for the generic methods the library implements for all its models.
+    This model inherits from [`RBLNDecoderOnlyModelForCausalLM`]. 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 PhiForCausalLM model on RBLN devices.
     It implements the methods to convert a pre-trained transformers PhiForCausalLM model into a RBLN transformer model by:
@@ -50,20 +40,4 @@ class RBLNPhiForCausalLM(RBLNDecoderOnlyModelForCausalLM):
     - compiling the resulting graph using the RBLN compiler.
     """
 
-    @classmethod
-    def wrap_model_if_needed(self, model: "PreTrainedModel", rbln_config: "RBLNConfig"):
-        rbln_max_seq_len = rbln_config.model_cfg["max_seq_len"]
-        return PhiWrapper(model, rbln_max_seq_len).eval()
-
-    def __getattr__(self, __name: str) -> Any:
-        def redirect(func):
-            return lambda *pargs, **kwargs: func(self, *pargs, **kwargs)
-
-        val = getattr(PhiForCausalLM, __name)
-
-        if isinstance(val, Callable) and "self" in set(
-            inspect.signature(val).parameters
-        ):
-            return redirect(val)
-
-        return val
+    _decoder_wrapper_cls = PhiWrapper

optimum/rbln/transformers/models/phi/phi_architecture.py

@@ -33,46 +33,12 @@ from transformers.modeling_outputs import (
 from ...cache_utils import RebelDynamicCache
 from ..decoderonly import (
     DecoderOnlyWrapper,
-    DynamicNTKScalingRotaryEmbedding,
-    LinearScalingRotaryEmbedding,
-    RotaryEmbedding,
     apply_rotary_pos_emb,
     slice_and_unsqueeze_cos_sin,
 )
 
 
 class PhiWrapper(DecoderOnlyWrapper):
-    def _init_rope(self):
-        if self.rope_scaling is None:
-            rotary_emb = RotaryEmbedding(
-                int(self.config.partial_rotary_factor * self.head_dim),
-                max_position_embeddings=self.max_position_embeddings,
-                base=self.config.rope_theta,
-            )
-        else:
-            scaling_type = self.rope_scaling["type"]
-            scaling_factor = self.rope_scaling["factor"]
-            if scaling_type == "linear":
-                rotary_emb = LinearScalingRotaryEmbedding(
-                    int(self.config.partial_rotary_factor * self.head_dim),
-                    max_position_embeddings=self.max_position_embeddings,
-                    scaling_factor=scaling_factor,
-                    base=self.config.rope_theta,
-                    max_seq_len=self.max_seq_len,
-                )
-            elif scaling_type == "dynamic":
-                rotary_emb = DynamicNTKScalingRotaryEmbedding(
-                    int(self.config.partial_rotary_factor * self.head_dim),
-                    max_position_embeddings=self.max_position_embeddings,
-                    scaling_factor=scaling_factor,
-                    base=self.config.rope_theta,
-                    max_seq_len=self.max_seq_len,
-                )
-            else:
-                raise ValueError(f"Unknown RoPE scaling type {scaling_type}")
-
-        return rotary_emb
-
     def get_forward_dict(self):
         forward_dict = {}
         forward_dict.update(
@@ -86,6 +52,41 @@ class PhiWrapper(DecoderOnlyWrapper):
 
 
 class PhiAttention:
+    def _attn(self, query_state, key_state, value_state, attn_mask, past_key_value, batch_idx=0, is_prefill=False):
+        # reshape for removing repeat_kv (batch=1 , num_head, 1, q_len=1, head_dim)
+        key_state = key_state.unsqueeze(2)
+        value_state = value_state.unsqueeze(2)
+        attn_mask = attn_mask.unsqueeze(2)
+
+        query_state = query_state.view(
+            1,
+            self.num_key_value_heads,
+            self.num_heads // self.num_key_value_heads,
+            -1,
+            self.head_dim,
+        )
+
+        key_state, value_state = past_key_value.update(key_state, value_state, self.layer_idx, batch_idx, is_prefill)
+
+        # Queries and keys upcast to fp32 is required by Phi-2 to avoid overflow
+        attn_weights = torch.matmul(
+            query_state.to(torch.float32),
+            key_state.to(torch.float32).transpose(3, 4),
+        ) / math.sqrt(self.head_dim)
+        attn_weights = attn_weights + attn_mask
+
+        # upcast attention to fp32
+        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_state.dtype)
+        attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)
+        attn_output = torch.matmul(attn_weights, value_state)
+
+        # reshape for removing repeat_kv
+        attn_output = attn_output.view(1, self.num_heads, -1, self.head_dim)
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.reshape(1, -1, self.num_heads * self.head_dim)
+
+        return attn_output, key_state, value_state
+
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -95,7 +96,6 @@ class PhiAttention:
         output_attentions: bool = False,
         cos: Optional[torch.Tensor] = None,
         sin: Optional[torch.Tensor] = None,
-        rotary_pos_emb=None,
         **kwargs,
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
         bsz, q_len, _ = hidden_states.size()
@@ -108,24 +108,18 @@ class PhiAttention:
             query_states = self.q_layernorm(query_states)
             key_states = self.k_layernorm(key_states)
 
-        query_states = query_states.view(
-            bsz, q_len, self.num_heads, self.head_dim
-        ).transpose(1, 2)
-        key_states = key_states.view(
-            bsz, q_len, self.num_key_value_heads, self.head_dim
-        ).transpose(1, 2)
-        value_states = value_states.view(
-            bsz, q_len, self.num_key_value_heads, self.head_dim
-        ).transpose(1, 2)
+        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
 
         # Partial rotary embedding
         query_rot, query_pass = (
-            query_states[..., : rotary_pos_emb.dim],
-            query_states[..., rotary_pos_emb.dim :],
+            query_states[..., : self.rotary_ndims],
+            query_states[..., self.rotary_ndims :],
         )
         key_rot, key_pass = (
-            key_states[..., : rotary_pos_emb.dim],
-            key_states[..., rotary_pos_emb.dim :],
+            key_states[..., : self.rotary_ndims],
+            key_states[..., self.rotary_ndims :],
         )
 
         # [batch_size, seq_length, num_heads, head_dim // config.partial_rotary_factor]
@@ -135,113 +129,38 @@ class PhiAttention:
         query_states = torch.cat((query_rot, query_pass), dim=-1)
         key_states = torch.cat((key_rot, key_pass), dim=-1)
 
-        # Decoder
-        if (batch_index is None or batch_index == -1) and bsz > 1:
-            all_key_states = []
-            all_value_states = []
-            all_attn_output = []
-
+        # Decoder (bsz > 1)
+        if bsz > 1:
+            iterate_results = {"key_states": [], "value_states": [], "attn_output": []}
             for b in range(bsz):
-                query_state = query_states[b].unsqueeze(0)
-                attn_mask = attention_mask[b].unsqueeze(0)
-                key_state = key_states[b].unsqueeze(0)
-                value_state = value_states[b].unsqueeze(0)
-
-                # reshape for removing repeat_kv (batch=1 , num_head, 1, q_len=1, head_dim)
-                key_state = key_state.unsqueeze(2)
-                value_state = value_state.unsqueeze(2)
-                attn_mask = attn_mask.unsqueeze(2)
-
-                query_state = query_state.view(
-                    1,
-                    self.num_key_value_heads,
-                    self.num_heads // self.num_key_value_heads,
-                    q_len,
-                    self.head_dim,
+                attn_output, key_state, value_state = PhiAttention._attn(
+                    self,
+                    query_states[b].unsqueeze(0),
+                    key_states[b].unsqueeze(0),
+                    value_states[b].unsqueeze(0),
+                    attention_mask[b].unsqueeze(0),
+                    past_key_value,
+                    batch_idx=b,
+                    is_prefill=False,
                 )
-
-                key_state, value_state = past_key_value.update(
-                    key_state,
-                    value_state,
-                    self.layer_idx,
-                    b,
-                )
-
-                # Queries and keys upcast to fp32 is required by Phi-2 to avoid overflow
-                attn_weights = torch.matmul(
-                    query_state.to(torch.float32),
-                    key_state.to(torch.float32).transpose(3, 4),
-                ) / math.sqrt(self.head_dim)
-                attn_weights = attn_weights + attn_mask
-
-                # upcast attention to fp32
-                attn_weights = nn.functional.softmax(
-                    attn_weights, dim=-1, dtype=torch.float32
-                ).to(query_states.dtype)
-                attn_weights = nn.functional.dropout(
-                    attn_weights, p=self.attention_dropout, training=self.training
-                )
-                attn_output = torch.matmul(attn_weights, value_state)
-
-                # reshape for removing repeat_kv
-                attn_output = attn_output.view(1, self.num_heads, q_len, self.head_dim)
-                attn_output = attn_output.transpose(1, 2).contiguous()
-                attn_output = attn_output.reshape(
-                    1, q_len, self.num_heads * self.head_dim
-                )
-
-                all_key_states.append(key_state)
-                all_value_states.append(value_state)
-                all_attn_output.append(attn_output)
-
-            key_states = torch.cat(all_key_states, dim=0)
-            value_states = torch.cat(all_value_states, dim=0)
-            attn_output = torch.cat(all_attn_output, dim=0)
+                iterate_results["key_states"].append(key_state)
+                iterate_results["value_states"].append(value_state)
+                iterate_results["attn_output"].append(attn_output)
+
+            key_states = torch.cat(iterate_results["key_states"], dim=0)
+            value_states = torch.cat(iterate_results["value_states"], dim=0)
+            attn_output = torch.cat(iterate_results["attn_output"], dim=0)
+        # Prefill & Decoder (bsz == 1)
         else:
-            if batch_index is None or batch_index == -1:
-                batch_index = 0
-
-            # reshape for removing repeat_kv
-            key_states = key_states.unsqueeze(2)
-            value_states = value_states.unsqueeze(2)
-            attention_mask = attention_mask.unsqueeze(2)
-            query_states = query_states.view(
-                1,
-                self.num_key_value_heads,
-                self.num_heads // self.num_key_value_heads,
-                q_len,
-                self.head_dim,
-            )
-
-            key_states, value_states = past_key_value.update(
+            attn_output, key_states, value_states = PhiAttention._attn(
+                self,
+                query_states,
                 key_states,
                 value_states,
-                self.layer_idx,
-                batch_index,
-                read_first_step=True,
-            )
-
-            # Queries and keys upcast to fp32 is required by Phi-2 to avoid overflow
-            attn_weights = torch.matmul(
-                query_states.to(torch.float32),
-                key_states.to(torch.float32).transpose(3, 4),
-            ) / math.sqrt(self.head_dim)
-            attn_weights = attn_weights + attention_mask
-
-            # upcast attention to fp32
-            attn_weights = torch.nn.functional.softmax(
-                attn_weights, dim=-1, dtype=torch.float32
-            ).to(value_states.dtype)
-            attn_weights = torch.nn.functional.dropout(
-                attn_weights, p=self.attention_dropout, training=self.training
-            )
-            attn_output = torch.matmul(attn_weights, value_states)
-
-            # reshape for removing repeat_kv
-            attn_output = attn_output.view(1, self.num_heads, q_len, self.head_dim)
-            attn_output = attn_output.transpose(1, 2).contiguous()
-            attn_output = attn_output.reshape(
-                bsz, q_len, self.num_heads * self.head_dim
+                attention_mask,
+                past_key_value,
+                batch_idx=batch_index,
+                is_prefill=True,
             )
 
         attn_output = self.dense(attn_output)
@@ -265,12 +184,9 @@ class PhiDecoderLayer:
         batch_ids: Optional[torch.LongTensor] = None,
         cos: Optional[torch.Tensor] = None,
         sin: Optional[torch.Tensor] = None,
-        rotary_pos_emb=None,
         forward_dict: Optional[Dict[str, classmethod]] = None,
         **kwargs,
-    ) -> Tuple[
-        torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]
-    ]:
+    ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
         """
         Args:
             hidden_states (`torch.FloatTensor`):
@@ -294,9 +210,7 @@ class PhiDecoderLayer:
         hidden_states = self.input_layernorm(hidden_states)
 
         # Self Attention
-        attn_outputs, self_attn_weights, key_states, value_states = forward_dict[
-            "decoder_layer"
-        ](
+        attn_outputs, self_attn_weights, key_states, value_states = forward_dict["decoder_layer"](
             self.self_attn,
             hidden_states=hidden_states,
             attention_mask=attention_mask,
@@ -307,7 +221,6 @@ class PhiDecoderLayer:
             use_cache=use_cache,
             cos=cos,
             sin=sin,
-            rotary_pos_emb=rotary_pos_emb,
             **kwargs,
         )
         past_key_value.assign(key_states, value_states, layer_idx)
@@ -339,6 +252,8 @@ class PhiModel:
         use_cache: Optional[bool] = True,
         output_attentions: Optional[bool] = False,
         output_hidden_states: Optional[bool] = False,
+        cache_pos_for_partitions: Optional[torch.Tensor] = None,
+        kvcache_partition_size: Optional[torch.Tensor] = None,
         forward_dict: Optional[Dict[str, classmethod]] = None,
         rotary_pos_emb=None,
     ) -> BaseModelOutputWithPast:
@@ -378,7 +293,8 @@ class PhiModel:
                 batch_ids=batch_ids,
                 cos=cos,
                 sin=sin,
-                rotary_pos_emb=rotary_pos_emb,
+                cache_pos_for_partitions=cache_pos_for_partitions,
+                kvcache_partition_size=kvcache_partition_size,
                 forward_dict=forward_dict,
             )
 

optimum/rbln/transformers/models/qwen2/modeling_qwen2.py

@@ -21,28 +21,18 @@
 # copied, modified, or distributed without prior written permission
 # from Rebellions Inc.
 
-import inspect
-import logging
-from typing import TYPE_CHECKING, Any, Callable
-
-from transformers import Qwen2ForCausalLM
-
-from ..decoderonly import RBLNDecoderOnlyModelForCausalLM
+from ....utils import logging
+from ...models.decoderonly import RBLNDecoderOnlyModelForCausalLM
 from .qwen2_architecture import QWEN2Wrapper
 
 
-if TYPE_CHECKING:
-    from transformers import PreTrainedModel
-
-    from ....modeling_config import RBLNConfig
-
-logger = logging.getLogger(__name__)
+logger = logging.get_logger(__name__)
 
 
 class RBLNQwen2ForCausalLM(RBLNDecoderOnlyModelForCausalLM):
     """
     The Llama Model transformer with a language modeling head (linear layer) on top.
-    This model inherits from [`RBLNMultiModel`]. Check the superclass documentation for the generic methods the library implements for all its models.
+    This model inherits from [`RBLNDecoderOnlyModelForCausalLM`]. 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 LlamaForCausalLM model on RBLN devices.
     It implements the methods to convert a pre-trained transformers LlamaForCausalLM model into a RBLN transformer model by:
@@ -50,18 +40,4 @@ class RBLNQwen2ForCausalLM(RBLNDecoderOnlyModelForCausalLM):
     - compiling the resulting graph using the RBLN compiler.
     """
 
-    @classmethod
-    def wrap_model_if_needed(self, model: "PreTrainedModel", rbln_config: "RBLNConfig"):
-        rbln_max_seq_len = rbln_config.model_cfg["max_seq_len"]
-        return QWEN2Wrapper(model, rbln_max_seq_len).eval()
-
-    def __getattr__(self, __name: str) -> Any:
-        def redirect(func):
-            return lambda *pargs, **kwargs: func(self, *pargs, **kwargs)
-
-        val = getattr(Qwen2ForCausalLM, __name)
-
-        if isinstance(val, Callable) and "self" in set(inspect.signature(val).parameters):
-            return redirect(val)
-
-        return val
+    _decoder_wrapper_cls = QWEN2Wrapper

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

@@ -21,5 +21,5 @@
 # copied, modified, or distributed without prior written permission
 # from Rebellions Inc.
 
-from .modeling_t5 import RBLNT5ForConditionalGeneration
+from .modeling_t5 import RBLNT5EncoderModel, RBLNT5ForConditionalGeneration
 from .t5_architecture import T5DecoderWrapper, T5EncoderWrapper

optimum/rbln/transformers/models/t5/modeling_t5.py

@@ -22,11 +22,16 @@
 # from Rebellions Inc.
 
 import inspect
-from typing import TYPE_CHECKING, Any, Callable
+from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Union
 
-from transformers import T5ForConditionalGeneration
+from transformers import (
+    AutoModelForTextEncoding,
+    PretrainedConfig,
+    T5ForConditionalGeneration,
+)
 
-from ....modeling_config import RBLNConfig
+from ....modeling_base import RBLNModel
+from ....modeling_config import RBLNCompileConfig, RBLNConfig
 from ....utils.logging import get_logger
 from ...models.seq2seq import RBLNModelForSeq2SeqLM
 from .t5_architecture import T5Wrapper
@@ -35,7 +40,55 @@ from .t5_architecture import T5Wrapper
 logger = get_logger()
 
 if TYPE_CHECKING:
-    from transformers import PreTrainedModel
+    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, PreTrainedModel
+
+
+class RBLNT5EncoderModel(RBLNModel):
+    auto_model_class = AutoModelForTextEncoding
+
+    @classmethod
+    def _get_rbln_config(
+        cls,
+        preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]],
+        model_config: Optional["PretrainedConfig"] = None,
+        rbln_kwargs: Dict[str, Any] = {},
+    ) -> RBLNConfig:
+        rbln_max_seq_len = rbln_kwargs.get("max_seq_len", None)
+        rbln_batch_size = rbln_kwargs.get("batch_size", None)
+
+        max_position_embeddings = getattr(model_config, "n_positions", None)
+
+        if rbln_max_seq_len is None:
+            rbln_max_seq_len = max_position_embeddings
+            if rbln_max_seq_len is None:
+                for tokenizer in preprocessors:
+                    if hasattr(tokenizer, "model_max_length"):
+                        rbln_max_seq_len = tokenizer.model_max_length
+                        break
+                if rbln_max_seq_len is None:
+                    raise ValueError("`rbln_max_seq_len` should be specified!")
+
+        if max_position_embeddings is not None and rbln_max_seq_len > max_position_embeddings:
+            raise ValueError("`rbln_max_seq_len` should be less or equal than max_position_embeddings!")
+
+        if rbln_batch_size is None:
+            rbln_batch_size = 1
+
+        input_info = [
+            ("input_ids", [rbln_batch_size, rbln_max_seq_len], "int64"),
+            ("attention_mask", [rbln_batch_size, rbln_max_seq_len], "int64"),
+        ]
+
+        rbln_compile_config = RBLNCompileConfig(input_info=input_info)
+
+        rbln_config = RBLNConfig(
+            rbln_cls=cls.__name__,
+            compile_cfgs=[rbln_compile_config],
+            rbln_kwargs=rbln_kwargs,
+        )
+
+        rbln_config.model_cfg.update({"max_seq_len": rbln_max_seq_len})
+        return rbln_config
 
 
 class RBLNT5ForConditionalGeneration(RBLNModelForSeq2SeqLM):

optimum/rbln/transformers/models/whisper/modeling_whisper.py

@@ -102,7 +102,7 @@ class RBLNRuntimeDecoder(RBLNPytorchRuntime):
 class RBLNWhisperForConditionalGeneration(RBLNModel, RBLNWhisperGenerationMixin):
     """
     The Whisper Model with a language modeling head. Can be used for automatic speech recognition.
-    This model inherits from [`RBLNMultiModel`]. Check the superclass documentation for the generic methods the library implements for all its models.
+    This model inherits from [`RBLNDecoderOnlyModelForCausalLM`]. 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 LlamaForCausalLM model on RBLN devices.
     It implements the methods to convert a pre-trained transformers LlamaForCausalLM model into a RBLN transformer model by:

optimum/rbln/transformers/utils/rbln_quantization.py

@@ -31,8 +31,13 @@ from torch.nn import functional as F
 
 # Constants
 QUANTIZED_WEIGHTS = {
-    "q_proj", "k_proj", "v_proj", "o_proj",
-    "gate_proj", "up_proj", "down_proj",
+    "q_proj",
+    "k_proj",
+    "v_proj",
+    "o_proj",
+    "gate_proj",
+    "up_proj",
+    "down_proj",
 }
 
 
@@ -81,6 +86,7 @@ def create_qlinear(layer: Linear) -> Linear:
     """
     Converts a standard linear layer to a quantized linear (qlinear) layer with a custom forward pass.
     """
+
     def qlinear_forward(self, inputs: torch.Tensor) -> torch.Tensor:
         if inputs.dtype != self.scales.dtype:
             raise TypeError(f"Expected input dtype {self.scales.dtype}, but got {inputs.dtype}")

optimum/rbln/utils/context.py

@@ -0,0 +1,58 @@
+# Copyright 2024 Rebellions Inc.
+
+# 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.
+
+# Portions of this software are licensed under the Apache License,
+# Version 2.0. See the NOTICE file distributed with this work for
+# additional information regarding copyright ownership.
+
+# All other portions of this software, including proprietary code,
+# are the intellectual property of Rebellions Inc. and may not be
+# copied, modified, or distributed without prior written permission
+# from Rebellions Inc.
+
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Union
+
+from optimum.exporters import TasksManager
+from transformers import AutoConfig, AutoModel
+
+
+@contextmanager
+def override_auto_classes(config_func=None, model_func=None, skip_taskmanager=True):
+    """Temporarily override Auto classes with original model classes"""
+    original_config = AutoConfig.from_pretrained
+    original_model = AutoModel.from_pretrained
+    original_get_model_from_task = TasksManager.get_model_from_task
+
+    def get_model_from_task(
+        task: str,
+        model_name_or_path: Union[str, Path],
+        **kwargs,
+    ):
+        return model_func(model_name_or_path, **kwargs)
+
+    def none_func(*args, **kwargs):
+        return None
+
+    try:
+        AutoConfig.from_pretrained = config_func or none_func
+        AutoModel.from_pretrained = model_func or none_func
+        if skip_taskmanager:
+            TasksManager.get_model_from_task = none_func if model_func is None else get_model_from_task
+        yield
+    finally:
+        AutoConfig.from_pretrained = original_config
+        AutoModel.from_pretrained = original_model
+        TasksManager.get_model_from_task = original_get_model_from_task

optimum/rbln/utils/decorator_utils.py

@@ -0,0 +1,55 @@
+from functools import wraps
+
+from .logging import get_logger
+
+
+logger = get_logger(__name__)
+
+
+def remove_compile_time_kwargs(func):
+    """
+    Decorator to handle compile-time parameters during inference.
+
+    For RBLN-optimized pipelines, several parameters must be determined during compilation
+    and cannot be modified during inference. This decorator:
+    1. Removes and warns about LoRA scale in cross_attention_kwargs
+    2. Removes and warns about image dimension parameters (height, width)
+
+    Args:
+        func: The pipeline's __call__ method to be wrapped
+    """
+
+    @wraps(func)
+    def wrapper(self, *args, **kwargs):
+        height_exists = "height" in kwargs and kwargs["height"] is not None
+        width_exists = "width" in kwargs and kwargs["width"] is not None
+        if height_exists or width_exists:
+            logger.warning(
+                "Image dimension parameters (`height`, `width`) will be ignored during inference. "
+                "Image dimensions must be specified during model compilation using from_pretrained()."
+            )
+            kwargs.pop("width", None)
+            kwargs.pop("height", None)
+
+        if "cross_attention_kwargs" in kwargs:
+            cross_attention_kwargs = kwargs.get("cross_attention_kwargs")
+            if not cross_attention_kwargs:
+                return func(self, *args, **kwargs)
+
+            has_scale = "scale" in cross_attention_kwargs
+            if has_scale:
+                logger.warning(
+                    "LoRA scale in cross_attention_kwargs will be ignored during inference. "
+                    "To adjust LoRA scale, specify it during model compilation using from_pretrained()."
+                )
+
+                # If scale is the only key, set to None
+                # Otherwise, remove scale and preserve other settings
+                if len(cross_attention_kwargs) == 1:
+                    kwargs["cross_attention_kwargs"] = None
+                else:
+                    kwargs["cross_attention_kwargs"].pop("scale")
+
+        return func(self, *args, **kwargs)
+
+    return wrapper

optimum/rbln/utils/import_utils.py

@@ -37,6 +37,13 @@ class VersionCompat:
 
 
 RBLN_VERSION_COMPATS = {
+    "0.1.13": [
+        VersionCompat(
+            package_name="rebel-compiler",
+            min_version="0.6.0",
+            max_version="0.6.1",
+        ),
+    ],
     "0.1.12": [
         VersionCompat(
             package_name="rebel-compiler",