optimum-rbln 0.9.1__tar.gz → 0.9.2a1__tar.gz

{optimum_rbln-0.9.1 → optimum_rbln-0.9.2a1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: optimum-rbln
-Version: 0.9.1
+Version: 0.9.2a1
 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
@@ -29,7 +29,7 @@ Requires-Dist: packaging>=24.1
 Requires-Dist: torch==2.7.0
 Requires-Dist: torchaudio<=2.7.0
 Requires-Dist: torchvision<=0.22.0
-Requires-Dist: transformers==4.51.3
+Requires-Dist: transformers==4.53.1
 Description-Content-Type: text/markdown
 
 

{optimum_rbln-0.9.1 → optimum_rbln-0.9.2a1}/pyproject.toml

@@ -33,7 +33,7 @@ dependencies = [
     "torchaudio<=2.7.0",
     "torchvision<=0.22.0",
     "accelerate>=1.0.1",
-    "transformers==4.51.3",
+    "transformers==4.53.1",
     "diffusers==0.35.1",
     "packaging>=24.1",
 ]
@@ -57,6 +57,7 @@ tests = [
     "soundfile>=0.13.1",
     "librosa>=0.11.0",
     "simphile>=1.0.2",
+    "torchcodec>=0.5.0"
 ]
 quality = [
     "ruff>=0.3.3",

{optimum_rbln-0.9.1 → optimum_rbln-0.9.2a1}/src/optimum/rbln/__version__.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.9.1'
-__version_tuple__ = version_tuple = (0, 9, 1)
+__version__ = version = '0.9.2a1'
+__version_tuple__ = version_tuple = (0, 9, 2, 'a1')
 
 __commit_id__ = commit_id = None

{optimum_rbln-0.9.1 → optimum_rbln-0.9.2a1}/src/optimum/rbln/configuration_utils.py

@@ -534,18 +534,18 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
     submodules: List[str] = []
     subclass_non_save_attributes = []
 
-    def init_submodule_config(
+    def initialize_submodule_config(
         self,
-        submodule_config_cls: Type["RBLNModelConfig"],
         submodule_config: Optional[Union[Dict[str, Any], "RBLNModelConfig"]] = None,
+        force_kwargs: bool = False,
         **kwargs: Any,
     ) -> "RBLNModelConfig":
-        # Initialize a submodule config from a dict or a RBLNModelConfig.
-        # kwargs is specified from the predecessor config.
-
         if submodule_config is None:
             submodule_config = {}
 
+        if isinstance(submodule_config, RBLNModelConfig):
+            return submodule_config
+
         if isinstance(submodule_config, dict):
             from_predecessor = self._runtime_options.copy()
             from_predecessor.update(
@@ -559,13 +559,60 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
 
             init_kwargs = from_predecessor
             init_kwargs.update(submodule_config)
-            submodule_config = submodule_config_cls(**init_kwargs)
 
-        if not isinstance(submodule_config, submodule_config_cls):
+            if force_kwargs:
+                for key, value in kwargs.items():
+                    if key in init_kwargs:
+                        if init_kwargs[key] != value:
+                            raise ValueError(
+                                f"Parameter conflict for '{key}': submodule_config has {init_kwargs[key]}, "
+                                f"but kwargs has {value}. Using kwargs value: {value}"
+                            )
+                        init_kwargs[key] = value
+
+            if "cls_name" in init_kwargs:
+                config_cls = get_rbln_config_class(init_kwargs["cls_name"])
+            else:
+                return init_kwargs
+
+            submodule_config = config_cls(**init_kwargs)
+
+        if not isinstance(submodule_config, RBLNModelConfig):
             raise TypeError(f"Invalid submodule config type: {type(submodule_config)}")
 
         return submodule_config
 
+    def filter_parameters(self, config_cls: Type["RBLNModelConfig"], parameters: Dict[str, Any]) -> Dict[str, Any]:
+        import importlib
+
+        model_cls_name = config_cls.__name__.replace("Config", "")
+        modeling_module_name = config_cls.__module__.replace("configuration_", "modeling_")
+
+        model_cls = None
+        try:
+            modeling_module = importlib.import_module(modeling_module_name)
+            if hasattr(modeling_module, model_cls_name):
+                model_cls = getattr(modeling_module, model_cls_name)
+        except ImportError:
+            logger.debug(f"Could not import modeling module: {modeling_module_name}")
+
+        filtered_out_params = set()
+
+        if model_cls is not None:
+            if not getattr(model_cls, "_tp_support", False):
+                filtered_out_params.add("tensor_parallel_size")
+
+        filtered_params = {}
+        for key, value in parameters.items():
+            if key in filtered_out_params:
+                logger.debug(
+                    f"Parameter '{key}' filtered out for {config_cls.__name__} (not supported by model flags)."
+                )
+            else:
+                filtered_params[key] = value
+
+        return filtered_params
+
     def __setattr__(self, key, value):
         if (
             key != "_attributes_map"

{optimum_rbln-0.9.1 → optimum_rbln-0.9.2a1}/src/optimum/rbln/diffusers/configurations/pipelines/configuration_controlnet.py

@@ -93,20 +93,27 @@ class RBLNStableDiffusionControlNetPipelineBaseConfig(RBLNModelConfig):
         elif (img_height is not None and img_width is None) or (img_height is None and img_width is not None):
             raise ValueError("Both img_height and img_width must be provided together if used")
 
-        self.text_encoder = self.init_submodule_config(RBLNCLIPTextModelConfig, text_encoder, batch_size=batch_size)
-        self.unet = self.init_submodule_config(
-            RBLNUNet2DConditionModelConfig,
+        self.text_encoder = self.initialize_submodule_config(
+            text_encoder,
+            cls_name="RBLNCLIPTextModelConfig",
+            batch_size=batch_size,
+        )
+        self.unet = self.initialize_submodule_config(
             unet,
+            cls_name="RBLNUNet2DConditionModelConfig",
             sample_size=sample_size,
         )
-        self.vae = self.init_submodule_config(
-            RBLNAutoencoderKLConfig,
+        self.vae = self.initialize_submodule_config(
             vae,
+            cls_name="RBLNAutoencoderKLConfig",
             batch_size=batch_size,
             uses_encoder=self.__class__._vae_uses_encoder,
             sample_size=image_size,  # image size is equal to sample size in vae
         )
-        self.controlnet = self.init_submodule_config(RBLNControlNetModelConfig, controlnet)
+        self.controlnet = self.initialize_submodule_config(
+            controlnet,
+            cls_name="RBLNControlNetModelConfig",
+        )
 
         # Get default guidance scale from original class to set UNet and ControlNet batch size
         if guidance_scale is None:
@@ -235,23 +242,32 @@ class RBLNStableDiffusionXLControlNetPipelineBaseConfig(RBLNModelConfig):
         elif (img_height is not None and img_width is None) or (img_height is None and img_width is not None):
             raise ValueError("Both img_height and img_width must be provided together if used")
 
-        self.text_encoder = self.init_submodule_config(RBLNCLIPTextModelConfig, text_encoder, batch_size=batch_size)
-        self.text_encoder_2 = self.init_submodule_config(
-            RBLNCLIPTextModelWithProjectionConfig, text_encoder_2, batch_size=batch_size
+        self.text_encoder = self.initialize_submodule_config(
+            text_encoder,
+            cls_name="RBLNCLIPTextModelConfig",
+            batch_size=batch_size,
         )
-        self.unet = self.init_submodule_config(
-            RBLNUNet2DConditionModelConfig,
+        self.text_encoder_2 = self.initialize_submodule_config(
+            text_encoder_2,
+            cls_name="RBLNCLIPTextModelWithProjectionConfig",
+            batch_size=batch_size,
+        )
+        self.unet = self.initialize_submodule_config(
             unet,
+            cls_name="RBLNUNet2DConditionModelConfig",
             sample_size=sample_size,
         )
-        self.vae = self.init_submodule_config(
-            RBLNAutoencoderKLConfig,
+        self.vae = self.initialize_submodule_config(
             vae,
+            cls_name="RBLNAutoencoderKLConfig",
             batch_size=batch_size,
             uses_encoder=self.__class__._vae_uses_encoder,
             sample_size=image_size,  # image size is equal to sample size in vae
         )
-        self.controlnet = self.init_submodule_config(RBLNControlNetModelConfig, controlnet)
+        self.controlnet = self.initialize_submodule_config(
+            controlnet,
+            cls_name="RBLNControlNetModelConfig",
+        )
 
         # Get default guidance scale from original class to set UNet and ControlNet batch size
         guidance_scale = (

{optimum_rbln-0.9.1 → optimum_rbln-0.9.2a1}/src/optimum/rbln/diffusers/configurations/pipelines/configuration_cosmos.py

@@ -63,12 +63,15 @@ class RBLNCosmosPipelineBaseConfig(RBLNModelConfig):
         """
         super().__init__(**kwargs)
 
-        self.text_encoder = self.init_submodule_config(
-            RBLNT5EncoderModelConfig, text_encoder, batch_size=batch_size, max_seq_len=max_seq_len
+        self.text_encoder = self.initialize_submodule_config(
+            text_encoder,
+            cls_name="RBLNT5EncoderModelConfig",
+            batch_size=batch_size,
+            max_seq_len=max_seq_len,
         )
-        self.transformer = self.init_submodule_config(
-            RBLNCosmosTransformer3DModelConfig,
+        self.transformer = self.initialize_submodule_config(
             transformer,
+            cls_name="RBLNCosmosTransformer3DModelConfig",
             batch_size=batch_size,
             max_seq_len=max_seq_len,
             height=height,
@@ -76,18 +79,18 @@ class RBLNCosmosPipelineBaseConfig(RBLNModelConfig):
             num_frames=num_frames,
             fps=fps,
         )
-        self.vae = self.init_submodule_config(
-            RBLNAutoencoderKLCosmosConfig,
+        self.vae = self.initialize_submodule_config(
             vae,
+            cls_name="RBLNAutoencoderKLCosmosConfig",
             batch_size=batch_size,
             uses_encoder=self.__class__._vae_uses_encoder,
             height=height,
             width=width,
             num_frames=num_frames,
         )
-        self.safety_checker = self.init_submodule_config(
-            RBLNCosmosSafetyCheckerConfig,
+        self.safety_checker = self.initialize_submodule_config(
             safety_checker,
+            cls_name="RBLNCosmosSafetyCheckerConfig",
             batch_size=batch_size,
             height=height,
             width=width,

{optimum_rbln-0.9.1 → optimum_rbln-0.9.2a1}/src/optimum/rbln/diffusers/configurations/pipelines/configuration_kandinsky2_2.py

@@ -88,10 +88,14 @@ class RBLNKandinskyV22PipelineBaseConfig(RBLNModelConfig):
         elif (img_height is not None and img_width is None) or (img_height is None and img_width is not None):
             raise ValueError("Both img_height and img_width must be provided together if used")
 
-        self.unet = self.init_submodule_config(RBLNUNet2DConditionModelConfig, unet, sample_size=sample_size)
-        self.movq = self.init_submodule_config(
-            RBLNVQModelConfig,
+        self.unet = self.initialize_submodule_config(
+            unet,
+            cls_name="RBLNUNet2DConditionModelConfig",
+            sample_size=sample_size,
+        )
+        self.movq = self.initialize_submodule_config(
             movq,
+            cls_name="RBLNVQModelConfig",
             batch_size=batch_size,
             sample_size=image_size,  # image size is equal to sample size in vae
             uses_encoder=self._movq_uses_encoder,
@@ -173,14 +177,20 @@ class RBLNKandinskyV22PriorPipelineConfig(RBLNModelConfig):
             accommodate classifier-free guidance.
         """
         super().__init__(**kwargs)
-        self.text_encoder = self.init_submodule_config(
-            RBLNCLIPTextModelWithProjectionConfig, text_encoder, batch_size=batch_size
+        self.text_encoder = self.initialize_submodule_config(
+            text_encoder,
+            cls_name="RBLNCLIPTextModelWithProjectionConfig",
+            batch_size=batch_size,
         )
-        self.image_encoder = self.init_submodule_config(
-            RBLNCLIPVisionModelWithProjectionConfig, image_encoder, batch_size=batch_size
+        self.image_encoder = self.initialize_submodule_config(
+            image_encoder,
+            cls_name="RBLNCLIPVisionModelWithProjectionConfig",
+            batch_size=batch_size,
+        )
+        self.prior = self.initialize_submodule_config(
+            prior,
+            cls_name="RBLNPriorTransformerConfig",
         )
-
-        self.prior = self.init_submodule_config(RBLNPriorTransformerConfig, prior)
 
         # Get default guidance scale from original class to set UNet batch size
         if guidance_scale is None:
@@ -286,18 +296,18 @@ class RBLNKandinskyV22CombinedPipelineBaseConfig(RBLNModelConfig):
         elif (img_height is not None and img_width is None) or (img_height is None and img_width is not None):
             raise ValueError("Both img_height and img_width must be provided together if used")
 
-        self.prior_pipe = self.init_submodule_config(
-            RBLNKandinskyV22PriorPipelineConfig,
+        self.prior_pipe = self.initialize_submodule_config(
             prior_pipe,
+            cls_name="RBLNKandinskyV22PriorPipelineConfig",
             prior=prior_prior,
             image_encoder=prior_image_encoder,
             text_encoder=prior_text_encoder,
             batch_size=batch_size,
             guidance_scale=guidance_scale,
         )
-        self.decoder_pipe = self.init_submodule_config(
-            self._decoder_pipe_cls,
+        self.decoder_pipe = self.initialize_submodule_config(
             decoder_pipe,
+            cls_name=self._decoder_pipe_cls.__name__,
             unet=unet,
             movq=movq,
             batch_size=batch_size,

{optimum_rbln-0.9.1 → optimum_rbln-0.9.2a1}/src/optimum/rbln/diffusers/configurations/pipelines/configuration_stable_diffusion.py

@@ -90,18 +90,22 @@ class RBLNStableDiffusionPipelineBaseConfig(RBLNModelConfig):
         elif (img_height is not None and img_width is None) or (img_height is None and img_width is not None):
             raise ValueError("Both img_height and img_width must be provided together if used")
 
-        self.text_encoder = self.init_submodule_config(RBLNCLIPTextModelConfig, text_encoder, batch_size=batch_size)
-        self.unet = self.init_submodule_config(
-            RBLNUNet2DConditionModelConfig,
+        self.text_encoder = self.initialize_submodule_config(
+            text_encoder,
+            cls_name="RBLNCLIPTextModelConfig",
+            batch_size=batch_size,
+        )
+        self.unet = self.initialize_submodule_config(
             unet,
+            cls_name="RBLNUNet2DConditionModelConfig",
             sample_size=sample_size,
         )
-        self.vae = self.init_submodule_config(
-            RBLNAutoencoderKLConfig,
+        self.vae = self.initialize_submodule_config(
             vae,
+            cls_name="RBLNAutoencoderKLConfig",
             batch_size=batch_size,
             uses_encoder=self.__class__._vae_uses_encoder,
-            sample_size=image_size,  # image size is equal to sample size in vae
+            sample_size=image_size,
         )
 
         # Get default guidance scale from original class to set UNet batch size

{optimum_rbln-0.9.1 → optimum_rbln-0.9.2a1}/src/optimum/rbln/diffusers/configurations/pipelines/configuration_stable_diffusion_3.py

@@ -100,27 +100,31 @@ class RBLNStableDiffusion3PipelineBaseConfig(RBLNModelConfig):
 
         max_seq_len = max_seq_len or 256
 
-        self.text_encoder = self.init_submodule_config(
-            RBLNCLIPTextModelWithProjectionConfig, text_encoder, batch_size=batch_size
+        self.text_encoder = self.initialize_submodule_config(
+            text_encoder,
+            cls_name="RBLNCLIPTextModelWithProjectionConfig",
+            batch_size=batch_size,
         )
-        self.text_encoder_2 = self.init_submodule_config(
-            RBLNCLIPTextModelWithProjectionConfig, text_encoder_2, batch_size=batch_size
+        self.text_encoder_2 = self.initialize_submodule_config(
+            text_encoder_2,
+            cls_name="RBLNCLIPTextModelWithProjectionConfig",
+            batch_size=batch_size,
         )
-        self.text_encoder_3 = self.init_submodule_config(
-            RBLNT5EncoderModelConfig,
+        self.text_encoder_3 = self.initialize_submodule_config(
             text_encoder_3,
+            cls_name="RBLNT5EncoderModelConfig",
             batch_size=batch_size,
             max_seq_len=max_seq_len,
             model_input_names=["input_ids"],
         )
-        self.transformer = self.init_submodule_config(
-            RBLNSD3Transformer2DModelConfig,
+        self.transformer = self.initialize_submodule_config(
             transformer,
+            cls_name="RBLNSD3Transformer2DModelConfig",
             sample_size=sample_size,
         )
-        self.vae = self.init_submodule_config(
-            RBLNAutoencoderKLConfig,
+        self.vae = self.initialize_submodule_config(
             vae,
+            cls_name="RBLNAutoencoderKLConfig",
             batch_size=batch_size,
             uses_encoder=self.__class__._vae_uses_encoder,
             sample_size=image_size,

{optimum_rbln-0.9.1 → optimum_rbln-0.9.2a1}/src/optimum/rbln/diffusers/configurations/pipelines/configuration_stable_diffusion_xl.py

@@ -93,18 +93,25 @@ class RBLNStableDiffusionXLPipelineBaseConfig(RBLNModelConfig):
         elif (img_height is not None and img_width is None) or (img_height is None and img_width is not None):
             raise ValueError("Both img_height and img_width must be provided together if used")
 
-        self.text_encoder = self.init_submodule_config(RBLNCLIPTextModelConfig, text_encoder, batch_size=batch_size)
-        self.text_encoder_2 = self.init_submodule_config(
-            RBLNCLIPTextModelWithProjectionConfig, text_encoder_2, batch_size=batch_size
+        self.text_encoder = self.initialize_submodule_config(
+            text_encoder,
+            cls_name="RBLNCLIPTextModelConfig",
+            batch_size=batch_size,
+        )
+        self.text_encoder_2 = self.initialize_submodule_config(
+            text_encoder_2,
+            cls_name="RBLNCLIPTextModelWithProjectionConfig",
+            batch_size=batch_size,
         )
-        self.unet = self.init_submodule_config(
-            RBLNUNet2DConditionModelConfig,
+
+        self.unet = self.initialize_submodule_config(
             unet,
+            cls_name="RBLNUNet2DConditionModelConfig",
             sample_size=sample_size,
         )
-        self.vae = self.init_submodule_config(
-            RBLNAutoencoderKLConfig,
+        self.vae = self.initialize_submodule_config(
             vae,
+            cls_name="RBLNAutoencoderKLConfig",
             batch_size=batch_size,
             uses_encoder=self.__class__._vae_uses_encoder,
             sample_size=image_size,  # image size is equal to sample size in vae

{optimum_rbln-0.9.1 → optimum_rbln-0.9.2a1}/src/optimum/rbln/diffusers/pipelines/cosmos/configuration_cosmos_guardrail.py

@@ -15,7 +15,7 @@
 from typing import Any, Optional, Tuple
 
 from ....configuration_utils import RBLNAutoConfig, RBLNModelConfig
-from ....transformers import RBLNLlamaForCausalLMConfig, RBLNSiglipVisionModelConfig
+from ....transformers import RBLNSiglipVisionModelConfig
 
 
 class RBLNVideoSafetyModelConfig(RBLNModelConfig):
@@ -81,30 +81,28 @@ class RBLNCosmosSafetyCheckerConfig(RBLNModelConfig):
 
         tensor_parallel_size = kwargs.get("tensor_parallel_size")
 
-        self.llamaguard3 = self.init_submodule_config(
-            RBLNLlamaForCausalLMConfig,
+        self.llamaguard3 = self.initialize_submodule_config(
             llamaguard3,
+            cls_name="RBLNLlamaForCausalLMConfig",
             batch_size=batch_size,
             tensor_parallel_size=tensor_parallel_size,
             max_seq_len=max_seq_len,
         )
-
-        self.siglip_encoder = self.init_submodule_config(
-            RBLNSiglipVisionModelConfig,
+        self.siglip_encoder = self.initialize_submodule_config(
             siglip_encoder,
+            cls_name="RBLNSiglipVisionModelConfig",
             batch_size=batch_size,
             image_size=(384, 384),
         )
-
-        self.video_safety_model = self.init_submodule_config(
-            RBLNVideoSafetyModelConfig,
+        self.video_safety_model = self.initialize_submodule_config(
             video_safety_model,
+            cls_name="RBLNVideoSafetyModelConfig",
             batch_size=batch_size,
             input_size=1152,
         )
-        self.face_blur_filter = self.init_submodule_config(
-            RBLNRetinaFaceFilterConfig,
+        self.face_blur_filter = self.initialize_submodule_config(
             face_blur_filter,
+            cls_name="RBLNRetinaFaceFilterConfig",
             batch_size=batch_size,
             image_size=image_size,
         )

{optimum_rbln-0.9.1 → optimum_rbln-0.9.2a1}/src/optimum/rbln/transformers/models/blip_2/configuration_blip_2.py

@@ -15,6 +15,10 @@
 from typing import Any, Optional
 
 from ....configuration_utils import RBLNModelConfig
+from ....utils.logging import get_logger
+
+
+logger = get_logger(__name__)
 
 
 class RBLNBlip2VisionModelConfig(RBLNModelConfig):
@@ -25,6 +29,16 @@ class RBLNBlip2VisionModelConfig(RBLNModelConfig):
     RBLN-optimized BLIP-2 vision encoder models for multimodal tasks.
     """
 
+    def __init__(
+        self,
+        batch_size: Optional[int] = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.batch_size = batch_size or 1
+        if not isinstance(self.batch_size, int) or self.batch_size < 0:
+            raise ValueError(f"batch_size must be a positive integer, got {self.batch_size}")
+
 
 class RBLNBlip2QFormerModelConfig(RBLNModelConfig):
     """
@@ -36,6 +50,7 @@ class RBLNBlip2QFormerModelConfig(RBLNModelConfig):
 
     def __init__(
         self,
+        batch_size: Optional[int] = None,
         num_query_tokens: Optional[int] = None,
         image_text_hidden_size: Optional[int] = None,
         **kwargs,
@@ -47,11 +62,22 @@ class RBLNBlip2QFormerModelConfig(RBLNModelConfig):
             kwargs: Additional arguments passed to the parent RBLNModelConfig.
         """
         super().__init__(**kwargs)
+        self.batch_size = batch_size or 1
+        if not isinstance(self.batch_size, int) or self.batch_size < 0:
+            raise ValueError(f"batch_size must be a positive integer, got {self.batch_size}")
+
         self.num_query_tokens = num_query_tokens
         self.image_text_hidden_size = image_text_hidden_size
 
 
 class RBLNBlip2ForConditionalGenerationConfig(RBLNModelConfig):
+    """
+    Configuration class for RBLNBlip2ForConditionalGeneration.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized BLIP-2 models for conditional generation tasks that involve both image and text inputs.
+    """
+
     submodules = ["vision_model", "qformer", "language_model"]
 
     def __init__(
@@ -78,6 +104,12 @@ class RBLNBlip2ForConditionalGenerationConfig(RBLNModelConfig):
         if not isinstance(self.batch_size, int) or self.batch_size < 0:
             raise ValueError(f"batch_size must be a positive integer, got {self.batch_size}")
 
-        self.vision_model = self.init_submodule_config(RBLNBlip2VisionModelConfig, vision_model)
-        self.language_model = language_model
-        self.qformer = self.init_submodule_config(RBLNBlip2QFormerModelConfig, qformer)
+        if self.batch_size != 1:
+            logger.warning("Ignore batch_size for Blip2 vision model. It will be set to 1.")
+            logger.warning("Ignore batch_size for Blip2 qformer. It will be set to 1.")
+
+        self.vision_model = self.initialize_submodule_config(
+            submodule_config=vision_model, batch_size=1, force_kwargs=True
+        )
+        self.qformer = self.initialize_submodule_config(submodule_config=qformer, batch_size=1, force_kwargs=True)
+        self.language_model = self.initialize_submodule_config(submodule_config=language_model)

{optimum_rbln-0.9.1 → optimum_rbln-0.9.2a1}/src/optimum/rbln/transformers/models/blip_2/modeling_blip_2.py

@@ -30,34 +30,30 @@ from transformers.utils import logging
 
 from ....configuration_utils import RBLNCompileConfig, RBLNModelConfig
 from ....modeling import RBLNModel
+from ...utils.rbln_runtime_wrapper import LoopProcessor
 
 
 logger = logging.get_logger(__name__)
 
 if TYPE_CHECKING:
+    import rebel
     from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer
 
 
-class LoopProjector:
-    def __init__(self, language_projection) -> None:
-        self.language_projection = language_projection
+class LoopProjector(LoopProcessor):
+    def __init__(self, language_projection: Union[RBLNModel, "rebel.Runtime"]):
+        super().__init__(model=language_projection)
 
-    def forward(self, *args, **kwargs):
-        query_output = args[0]
+    def _get_batch_size(self, query_output, **kwargs):
+        return query_output.shape[0]
 
-        batch_size = query_output.shape[0]
-        outputs = []
-        for i in range(batch_size):
-            outputs.append(self.language_projection(query_output[i : i + 1]))
-
-        outputs = torch.cat(outputs, dim=0)
-        return outputs
+    def _prepare_inputs_for_iteration(self, index, common_inputs, query_output, **kwargs):
+        query_output_item = query_output[index : index + 1]
+        return ([query_output_item], {})
 
-    def __call__(self, *args: Any, **kwds: Any) -> Any:
-        return self.forward(*args, **kwds)
-
-    def __repr__(self) -> str:
-        return repr(self.language_projection)
+    def _process_outputs(self, outputs: list, **kwargs):
+        output = torch.cat(outputs, dim=0)
+        return output
 
 
 class RBLNBlip2VisionModel(RBLNModel):
@@ -68,6 +64,8 @@ class RBLNBlip2VisionModel(RBLNModel):
     on RBLN devices, supporting image encoding for multimodal vision-language tasks.
     """
 
+    _tp_support = False
+
     def get_input_embeddings(self):
         return self.embeddings
 
@@ -96,8 +94,7 @@ class RBLNBlip2VisionModel(RBLNModel):
             (
                 "pixel_values",
                 [
-                    # support for vllm CB (prefill)
-                    1,
+                    rbln_config.batch_size,
                     model_config.num_channels,
                     model_config.image_size,
                     model_config.image_size,
@@ -147,6 +144,8 @@ class RBLNBlip2QFormerModel(RBLNModel):
     mechanisms for multimodal understanding tasks.
     """
 
+    _tp_support = False
+
     def get_input_embeddings(self):
         return self.embeddings.word_embeddings
 
@@ -200,7 +199,7 @@ class RBLNBlip2QFormerModel(RBLNModel):
             (
                 "query_embeds",
                 [
-                    1,
+                    rbln_config.batch_size,
                     rbln_config.num_query_tokens,
                     model_config.hidden_size,
                 ],
@@ -209,7 +208,7 @@ class RBLNBlip2QFormerModel(RBLNModel):
             (
                 "encoder_hidden_states",
                 [
-                    1,
+                    rbln_config.batch_size,
                     # image_text_hidden_size + cls token
                     rbln_config.image_text_hidden_size + 1,
                     model_config.encoder_hidden_size,
@@ -219,7 +218,7 @@ class RBLNBlip2QFormerModel(RBLNModel):
             (
                 "encoder_attention_mask",
                 # image_text_hidden_size + cls token
-                [1, rbln_config.image_text_hidden_size + 1],
+                [rbln_config.batch_size, rbln_config.image_text_hidden_size + 1],
                 "int64",
             ),
         ]