optimum-rbln 0.7.3a5__py3-none-any.whl → 0.7.3.post1__py3-none-any.whl

optimum/rbln/__init__.py

@@ -78,9 +78,13 @@ _import_structure = {
         "RBLNAutoencoderKL",
         "RBLNControlNetModel",
         "RBLNPriorTransformer",
+        "RBLNKandinskyV22CombinedPipeline",
+        "RBLNKandinskyV22Img2ImgCombinedPipeline",
         "RBLNKandinskyV22InpaintCombinedPipeline",
         "RBLNKandinskyV22InpaintPipeline",
+        "RBLNKandinskyV22Img2ImgPipeline",
         "RBLNKandinskyV22PriorPipeline",
+        "RBLNKandinskyV22Pipeline",
         "RBLNStableDiffusionPipeline",
         "RBLNStableDiffusionXLPipeline",
         "RBLNUNet2DConditionModel",
@@ -107,8 +111,12 @@ if TYPE_CHECKING:
         RBLNAutoencoderKL,
         RBLNControlNetModel,
         RBLNDiffusionMixin,
+        RBLNKandinskyV22CombinedPipeline,
+        RBLNKandinskyV22Img2ImgCombinedPipeline,
+        RBLNKandinskyV22Img2ImgPipeline,
         RBLNKandinskyV22InpaintCombinedPipeline,
         RBLNKandinskyV22InpaintPipeline,
+        RBLNKandinskyV22Pipeline,
         RBLNKandinskyV22PriorPipeline,
         RBLNMultiControlNetModel,
         RBLNPriorTransformer,

optimum/rbln/__version__.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.7.3a5'
-__version_tuple__ = version_tuple = (0, 7, 3, 'a5')
+__version__ = version = '0.7.3.post1'
+__version_tuple__ = version_tuple = (0, 7, 3)

optimum/rbln/diffusers/__init__.py

@@ -24,9 +24,13 @@ ALL_IMPORTABLE_CLASSES.update(LOADABLE_CLASSES["optimum.rbln"])
 
 _import_structure = {
     "pipelines": [
+        "RBLNKandinskyV22CombinedPipeline",
+        "RBLNKandinskyV22Img2ImgCombinedPipeline",
         "RBLNKandinskyV22InpaintCombinedPipeline",
         "RBLNKandinskyV22InpaintPipeline",
+        "RBLNKandinskyV22Img2ImgPipeline",
         "RBLNKandinskyV22PriorPipeline",
+        "RBLNKandinskyV22Pipeline",
         "RBLNStableDiffusionPipeline",
         "RBLNStableDiffusionXLPipeline",
         "RBLNStableDiffusionImg2ImgPipeline",
@@ -66,8 +70,12 @@ if TYPE_CHECKING:
         RBLNVQModel,
     )
     from .pipelines import (
+        RBLNKandinskyV22CombinedPipeline,
+        RBLNKandinskyV22Img2ImgCombinedPipeline,
+        RBLNKandinskyV22Img2ImgPipeline,
         RBLNKandinskyV22InpaintCombinedPipeline,
         RBLNKandinskyV22InpaintPipeline,
+        RBLNKandinskyV22Pipeline,
         RBLNKandinskyV22PriorPipeline,
         RBLNMultiControlNetModel,
         RBLNStableDiffusion3Img2ImgPipeline,

optimum/rbln/diffusers/modeling_diffusers.py

@@ -23,7 +23,6 @@ from ..modeling import RBLNModel
 from ..modeling_config import RUNTIME_KEYWORDS, ContextRblnConfig, use_rbln_config
 from ..utils.decorator_utils import remove_compile_time_kwargs
 from ..utils.logging import get_logger
-from . import pipelines
 
 
 logger = get_logger(__name__)
@@ -67,6 +66,7 @@ class RBLNDiffusionMixin:
           as keys in rbln_config
     """
 
+    _connected_classes = {}
     _submodules = []
     _prefix = {}
 
@@ -103,37 +103,6 @@ class RBLNDiffusionMixin:
                 }
             )
             submodule_config = submodule_cls.update_rbln_config_using_pipe(model, submodule_config)
-        elif hasattr(pipelines, submodule_class_name):
-            submodule_config = rbln_config.get(submodule_name, {})
-            submodule_config = copy.deepcopy(submodule_config)
-
-            submodule_cls: RBLNModel = getattr(importlib.import_module("optimum.rbln"), f"{submodule_class_name}")
-            prefix = cls._prefix.get(submodule_name, "")
-            connected_submodules = cls._connected_classes.get(submodule_name)._submodules
-            pipe_global_config = {k: v for k, v in submodule_config.items() if k not in connected_submodules}
-            submodule_config = {k: v for k, v in submodule_config.items() if k in connected_submodules}
-            for key in submodule_config.keys():
-                submodule_config[key].update(pipe_global_config)
-
-            for connected_submodule_name in connected_submodules:
-                connected_submodule_config = rbln_config.pop(prefix + connected_submodule_name, {})
-                if connected_submodule_name in submodule_config:
-                    submodule_config[connected_submodule_name].update(connected_submodule_config)
-                else:
-                    submodule_config[connected_submodule_name] = connected_submodule_config
-
-            pipe_global_config = {
-                k: v for k, v in rbln_config.items() if k != submodule_class_name and not isinstance(v, dict)
-            }
-
-            for connected_submodule_name in connected_submodules:
-                for k, v in pipe_global_config.items():
-                    if "guidance_scale" in k:
-                        if prefix + "guidance_scale" == k:
-                            submodule_config[connected_submodule_name]["guidance_scale"] = v
-                    else:
-                        submodule_config[connected_submodule_name][k] = v
-            rbln_config[submodule_name] = submodule_config
         else:
             raise ValueError(f"submodule {submodule_name} isn't supported")
         return submodule_config
@@ -199,25 +168,8 @@ class RBLNDiffusionMixin:
         else:
             # raise error if any of submodules are torch module.
             model_index_config = cls.load_config(pretrained_model_name_or_path=model_id)
-            if cls._load_connected_pipes:
-                submodules = []
-                for submodule in cls._submodules:
-                    submodule_config = rbln_config.pop(submodule, {})
-                    prefix = cls._prefix.get(submodule, "")
-                    connected_submodules = cls._connected_classes.get(submodule)._submodules
-                    for connected_submodule_name in connected_submodules:
-                        connected_submodule_config = submodule_config.pop(connected_submodule_name, {})
-                        if connected_submodule_config:
-                            rbln_config[prefix + connected_submodule_name] = connected_submodule_config
-                        submodules.append(prefix + connected_submodule_name)
-                pipe_global_config = {k: v for k, v in rbln_config.items() if k not in submodules}
-                for submodule in submodules:
-                    if submodule in rbln_config:
-                        rbln_config[submodule].update(pipe_global_config)
-            else:
-                submodules = cls._submodules
-
-            for submodule_name in submodules:
+            rbln_config = cls._flatten_rbln_config(rbln_config)
+            for submodule_name in cls._submodules:
                 if isinstance(kwargs.get(submodule_name), torch.nn.Module):
                     raise AssertionError(
                         f"{submodule_name} is not compiled torch module. If you want to compile, set `export=True`."
@@ -266,9 +218,89 @@ class RBLNDiffusionMixin:
             lora_scales=lora_scales,
         )
 
-        compiled_submodules = cls._compile_submodules(model, passed_submodules, model_save_dir, rbln_config)
+        if cls._load_connected_pipes:
+            compiled_submodules = cls._compile_pipelines(model, passed_submodules, model_save_dir, rbln_config)
+        else:
+            compiled_submodules = cls._compile_submodules(model, passed_submodules, model_save_dir, rbln_config)
         return cls._construct_pipe(model, compiled_submodules, model_save_dir, rbln_config)
 
+    @classmethod
+    def _prepare_rbln_config(
+        cls,
+        rbln_config,
+    ) -> Dict[str, Any]:
+        prepared_config = {}
+        for connected_pipe_name, connected_pipe_cls in cls._connected_classes.items():
+            connected_pipe_config = rbln_config.pop(connected_pipe_name, {})
+            prefix = cls._prefix.get(connected_pipe_name, "")
+            guidance_scale = rbln_config.pop(f"{prefix}guidance_scale", None)
+            if "guidance_scale" not in connected_pipe_config and guidance_scale is not None:
+                connected_pipe_config["guidance_scale"] = guidance_scale
+            for submodule_name in connected_pipe_cls._submodules:
+                submodule_config = rbln_config.pop(prefix + submodule_name, {})
+                if submodule_name not in connected_pipe_config:
+                    connected_pipe_config[submodule_name] = {}
+                connected_pipe_config[submodule_name].update(
+                    {k: v for k, v in submodule_config.items() if k not in connected_pipe_config[submodule_name]}
+                )
+            prepared_config[connected_pipe_name] = connected_pipe_config
+        prepared_config.update(rbln_config)
+        return prepared_config
+
+    @classmethod
+    def _flatten_rbln_config(
+        cls,
+        rbln_config,
+    ) -> Dict[str, Any]:
+        prepared_config = cls._prepare_rbln_config(rbln_config)
+        flattened_config = {}
+        pipe_global_config = {k: v for k, v in prepared_config.items() if k not in cls._connected_classes.keys()}
+        for connected_pipe_name, connected_pipe_cls in cls._connected_classes.items():
+            connected_pipe_config = prepared_config.pop(connected_pipe_name)
+            prefix = cls._prefix.get(connected_pipe_name, "")
+            connected_pipe_global_config = {
+                k: v for k, v in connected_pipe_config.items() if k not in connected_pipe_cls._submodules
+            }
+            for submodule_name in connected_pipe_cls._submodules:
+                flattened_config[prefix + submodule_name] = connected_pipe_config[submodule_name]
+                flattened_config[prefix + submodule_name].update(
+                    {
+                        k: v
+                        for k, v in connected_pipe_global_config.items()
+                        if k not in flattened_config[prefix + submodule_name]
+                    }
+                )
+        flattened_config.update(pipe_global_config)
+        return flattened_config
+
+    @classmethod
+    def _compile_pipelines(
+        cls,
+        model: torch.nn.Module,
+        passed_submodules: Dict[str, RBLNModel],
+        model_save_dir: Optional[PathLike],
+        rbln_config: Dict[str, Any],
+    ) -> Dict[str, RBLNModel]:
+        compiled_submodules = {}
+
+        rbln_config = cls._prepare_rbln_config(rbln_config)
+        pipe_global_config = {k: v for k, v in rbln_config.items() if k not in cls._connected_classes.keys()}
+        for connected_pipe_name, connected_pipe_cls in cls._connected_classes.items():
+            connected_pipe_submodules = {}
+            prefix = cls._prefix.get(connected_pipe_name, "")
+            for submodule_name in connected_pipe_cls._submodules:
+                connected_pipe_submodules[submodule_name] = passed_submodules.get(prefix + submodule_name, None)
+            connected_pipe = getattr(model, connected_pipe_name)
+            connected_pipe_config = {}
+            connected_pipe_config.update(pipe_global_config)
+            connected_pipe_config.update(rbln_config[connected_pipe_name])
+            connected_pipe_compiled_submodules = connected_pipe_cls._compile_submodules(
+                connected_pipe, connected_pipe_submodules, model_save_dir, connected_pipe_config, prefix
+            )
+            for submodule_name, compiled_submodule in connected_pipe_compiled_submodules.items():
+                compiled_submodules[prefix + submodule_name] = compiled_submodule
+        return compiled_submodules
+
     @classmethod
     def _compile_submodules(
         cls,
@@ -307,41 +339,6 @@ class RBLNDiffusionMixin:
                     model_save_dir=model_save_dir,
                     rbln_config=submodule_rbln_config,
                 )
-            elif hasattr(pipelines, submodule.__class__.__name__):
-                connected_pipe = submodule
-                connected_pipe_model_save_dir = model_save_dir
-                connected_pipe_rbln_config = submodule_rbln_config
-                connected_pipe_cls: RBLNDiffusionMixin = getattr(
-                    importlib.import_module("optimum.rbln"), connected_pipe.__class__.__name__
-                )
-                submodule_dict = {}
-                for name in connected_pipe.config.keys():
-                    if hasattr(connected_pipe, name):
-                        submodule_dict[name] = getattr(connected_pipe, name)
-                connected_pipe = connected_pipe_cls(**submodule_dict)
-                connected_pipe_submodules = {}
-                prefix = cls._prefix.get(submodule_name, "")
-                for name in connected_pipe_cls._submodules:
-                    if prefix + name in passed_submodules:
-                        connected_pipe_submodules[name] = passed_submodules.get(prefix + name)
-
-                connected_pipe_compiled_submodules = connected_pipe_cls._compile_submodules(
-                    model=connected_pipe,
-                    passed_submodules=connected_pipe_submodules,
-                    model_save_dir=model_save_dir,
-                    rbln_config=connected_pipe_rbln_config,
-                    prefix=prefix,
-                )
-                connected_pipe = connected_pipe_cls._construct_pipe(
-                    connected_pipe,
-                    connected_pipe_compiled_submodules,
-                    connected_pipe_model_save_dir,
-                    connected_pipe_rbln_config,
-                )
-
-                for name in connected_pipe_cls._submodules:
-                    compiled_submodules[prefix + name] = getattr(connected_pipe, name)
-                submodule = connected_pipe
             else:
                 raise ValueError(f"Unknown class of submodule({submodule_name}) : {submodule.__class__.__name__} ")
 
@@ -374,23 +371,16 @@ class RBLNDiffusionMixin:
     @classmethod
     def _construct_pipe(cls, model, submodules, model_save_dir, rbln_config):
         # Construct finalize pipe setup with compiled submodules and configurations
-        submodule_names = []
-        for submodule_name in cls._submodules:
-            submodule = getattr(model, submodule_name)
-            if hasattr(pipelines, submodule.__class__.__name__):
-                prefix = cls._prefix.get(submodule_name, "")
-                connected_pipe_submodules = submodules[submodule_name].__class__._submodules
-                connected_pipe_submodules = [prefix + name for name in connected_pipe_submodules]
-                submodule_names += connected_pipe_submodules
-                setattr(model, submodule_name, submodules[submodule_name])
-            else:
-                submodule_names.append(submodule_name)
-
         if model_save_dir is not None:
             # To skip saving original pytorch modules
-            for submodule_name in submodule_names:
+            for submodule_name in cls._submodules:
                 delattr(model, submodule_name)
 
+            if cls._load_connected_pipes:
+                for connected_pipe_name, connected_pipe_cls in cls._connected_classes.items():
+                    for submodule_name in connected_pipe_cls._submodules:
+                        delattr(getattr(model, connected_pipe_name), submodule_name)
+
             # Direct calling of `save_pretrained` causes config.unet = (None, None).
             # So config must be saved again, later.
             model.save_pretrained(model_save_dir)
@@ -398,10 +388,15 @@ class RBLNDiffusionMixin:
             # Causing warning messeages.
 
         update_dict = {}
-        for submodule_name in submodule_names:
+        for submodule_name in cls._submodules:
             # replace submodule
             setattr(model, submodule_name, submodules[submodule_name])
             update_dict[submodule_name] = ("optimum.rbln", submodules[submodule_name].__class__.__name__)
+        if cls._load_connected_pipes:
+            for connected_pipe_name, connected_pipe_cls in cls._connected_classes.items():
+                prefix = cls._prefix.get(connected_pipe_name, "")
+                for submodule_name in connected_pipe_cls._submodules:
+                    setattr(getattr(model, connected_pipe_name), submodule_name, submodules[prefix + submodule_name])
 
         # Update config to be able to load from model directory.
         #
@@ -420,16 +415,9 @@ class RBLNDiffusionMixin:
         if rbln_config.get("optimize_host_memory") is False:
             # Keep compiled_model objs to further analysis. -> TODO: remove soon...
             model.compiled_models = []
-            if model._load_connected_pipes:
-                for name in cls._submodules:
-                    connected_pipe = getattr(model, name)
-                    for submodule_name in connected_pipe.__class__._submodules:
-                        submodule = getattr(connected_pipe, submodule_name)
-                        model.compiled_models.extend(submodule.compiled_models)
-            else:
-                for name in cls._submodules:
-                    submodule = getattr(model, name)
-                    model.compiled_models.extend(submodule.compiled_models)
+            for name in cls._submodules:
+                submodule = getattr(model, name)
+                model.compiled_models.extend(submodule.compiled_models)
 
         return model
 

optimum/rbln/diffusers/models/autoencoders/vq_model.py

@@ -90,9 +90,17 @@ class RBLNVQModel(RBLNModel):
         model_config: "PretrainedConfig",
         rbln_kwargs: Dict[str, Any] = {},
     ) -> RBLNConfig:
-        batch_size = rbln_kwargs.get("batch_size") or 1
-        height = rbln_kwargs.get("img_height") or 512
-        width = rbln_kwargs.get("img_width") or 512
+        batch_size = rbln_kwargs.get("batch_size")
+        if batch_size is None:
+            batch_size = 1
+
+        height = rbln_kwargs.get("img_height")
+        if height is None:
+            height = 512
+
+        width = rbln_kwargs.get("img_width")
+        if width is None:
+            width = 512
 
         if hasattr(model_config, "block_out_channels"):
             scale_factor = 2 ** (len(model_config.block_out_channels) - 1)

optimum/rbln/diffusers/models/unets/unet_2d_condition.py

@@ -176,15 +176,22 @@ class RBLNUNet2DConditionModel(RBLNModel):
             raise ValueError("Both image height and image width must be given or not given")
         elif image_size[0] is None and image_size[1] is None:
             if rbln_config["img2img_pipeline"]:
-                # In case of img2img, sample size of unet is determined by vae encoder.
-                vae_sample_size = pipe.vae.config.sample_size
-                if isinstance(vae_sample_size, int):
-                    sample_size = vae_sample_size // scale_factor
-                else:
-                    sample_size = (
-                        vae_sample_size[0] // scale_factor,
-                        vae_sample_size[1] // scale_factor,
+                if hasattr(pipe, "vae"):
+                    # In case of img2img, sample size of unet is determined by vae encoder.
+                    vae_sample_size = pipe.vae.config.sample_size
+                    if isinstance(vae_sample_size, int):
+                        sample_size = vae_sample_size // scale_factor
+                    else:
+                        sample_size = (
+                            vae_sample_size[0] // scale_factor,
+                            vae_sample_size[1] // scale_factor,
+                        )
+                elif hasattr(pipe, "movq"):
+                    logger.warning(
+                        "RBLN config 'img_height' and 'img_width' should have been provided for this pipeline. "
+                        "Both variable will be set 512 by default."
                     )
+                    sample_size = (512 // scale_factor, 512 // scale_factor)
             else:
                 sample_size = pipe.unet.config.sample_size
         else:

optimum/rbln/diffusers/pipelines/__init__.py

@@ -26,9 +26,13 @@ _import_structure = {
         "RBLNStableDiffusionXLControlNetPipeline",
     ],
     "kandinsky2_2": [
+        "RBLNKandinskyV22CombinedPipeline",
+        "RBLNKandinskyV22Img2ImgCombinedPipeline",
         "RBLNKandinskyV22InpaintCombinedPipeline",
         "RBLNKandinskyV22InpaintPipeline",
+        "RBLNKandinskyV22Img2ImgPipeline",
         "RBLNKandinskyV22PriorPipeline",
+        "RBLNKandinskyV22Pipeline",
     ],
     "stable_diffusion": [
         "RBLNStableDiffusionImg2ImgPipeline",
@@ -55,8 +59,12 @@ if TYPE_CHECKING:
         RBLNStableDiffusionXLControlNetPipeline,
     )
     from .kandinsky2_2 import (
+        RBLNKandinskyV22CombinedPipeline,
+        RBLNKandinskyV22Img2ImgCombinedPipeline,
+        RBLNKandinskyV22Img2ImgPipeline,
         RBLNKandinskyV22InpaintCombinedPipeline,
         RBLNKandinskyV22InpaintPipeline,
+        RBLNKandinskyV22Pipeline,
         RBLNKandinskyV22PriorPipeline,
     )
     from .stable_diffusion import (

optimum/rbln/diffusers/pipelines/kandinsky2_2/__init__.py

@@ -12,6 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .pipeline_kandinsky2_2_combined import RBLNKandinskyV22InpaintCombinedPipeline
+from .pipeline_kandinsky2_2 import RBLNKandinskyV22Pipeline
+from .pipeline_kandinsky2_2_combined import (
+    RBLNKandinskyV22CombinedPipeline,
+    RBLNKandinskyV22Img2ImgCombinedPipeline,
+    RBLNKandinskyV22InpaintCombinedPipeline,
+)
+from .pipeline_kandinsky2_2_img2img import RBLNKandinskyV22Img2ImgPipeline
 from .pipeline_kandinsky2_2_inpaint import RBLNKandinskyV22InpaintPipeline
 from .pipeline_kandinsky2_2_prior import RBLNKandinskyV22PriorPipeline

optimum/rbln/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2.py

@@ -0,0 +1,25 @@
+# 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.
+
+from diffusers import KandinskyV22Pipeline
+
+from ...modeling_diffusers import RBLNDiffusionMixin
+
+
+class RBLNKandinskyV22Pipeline(RBLNDiffusionMixin, KandinskyV22Pipeline):
+    original_class = KandinskyV22Pipeline
+    _submodules = ["unet", "movq"]
+
+    def get_compiled_image_size(self):
+        return self.movq.image_size

optimum/rbln/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_combined.py

@@ -14,6 +14,8 @@
 
 from diffusers import (
     DDPMScheduler,
+    KandinskyV22CombinedPipeline,
+    KandinskyV22Img2ImgCombinedPipeline,
     KandinskyV22InpaintCombinedPipeline,
     PriorTransformer,
     UnCLIPScheduler,
@@ -28,14 +30,118 @@ from transformers import (
 )
 
 from ...modeling_diffusers import RBLNDiffusionMixin
+from .pipeline_kandinsky2_2 import RBLNKandinskyV22Pipeline
+from .pipeline_kandinsky2_2_img2img import RBLNKandinskyV22Img2ImgPipeline
 from .pipeline_kandinsky2_2_inpaint import RBLNKandinskyV22InpaintPipeline
 from .pipeline_kandinsky2_2_prior import RBLNKandinskyV22PriorPipeline
 
 
+class RBLNKandinskyV22CombinedPipeline(RBLNDiffusionMixin, KandinskyV22CombinedPipeline):
+    original_class = KandinskyV22CombinedPipeline
+    _connected_classes = {"prior_pipe": RBLNKandinskyV22PriorPipeline, "decoder_pipe": RBLNKandinskyV22Pipeline}
+    _submodules = ["prior_image_encoder", "prior_text_encoder", "prior_prior", "unet", "movq"]
+    _prefix = {"prior_pipe": "prior_"}
+
+    def __init__(
+        self,
+        unet: "UNet2DConditionModel",
+        scheduler: "DDPMScheduler",
+        movq: "VQModel",
+        prior_prior: "PriorTransformer",
+        prior_image_encoder: "CLIPVisionModelWithProjection",
+        prior_text_encoder: "CLIPTextModelWithProjection",
+        prior_tokenizer: "CLIPTokenizer",
+        prior_scheduler: "UnCLIPScheduler",
+        prior_image_processor: "CLIPImageProcessor",
+    ):
+        RBLNDiffusionMixin.__init__(self)
+        super(KandinskyV22CombinedPipeline, self).__init__()
+
+        self.register_modules(
+            unet=unet,
+            scheduler=scheduler,
+            movq=movq,
+            prior_prior=prior_prior,
+            prior_image_encoder=prior_image_encoder,
+            prior_text_encoder=prior_text_encoder,
+            prior_tokenizer=prior_tokenizer,
+            prior_scheduler=prior_scheduler,
+            prior_image_processor=prior_image_processor,
+        )
+
+        self.prior_pipe = RBLNKandinskyV22PriorPipeline(
+            prior=prior_prior,
+            image_encoder=prior_image_encoder,
+            text_encoder=prior_text_encoder,
+            tokenizer=prior_tokenizer,
+            scheduler=prior_scheduler,
+            image_processor=prior_image_processor,
+        )
+        self.decoder_pipe = RBLNKandinskyV22Pipeline(
+            unet=unet,
+            scheduler=scheduler,
+            movq=movq,
+        )
+
+    def get_compiled_image_size(self):
+        return self.movq.image_size
+
+
+class RBLNKandinskyV22Img2ImgCombinedPipeline(RBLNDiffusionMixin, KandinskyV22Img2ImgCombinedPipeline):
+    original_class = KandinskyV22Img2ImgCombinedPipeline
+    _connected_classes = {"prior_pipe": RBLNKandinskyV22PriorPipeline, "decoder_pipe": RBLNKandinskyV22Img2ImgPipeline}
+    _submodules = ["prior_image_encoder", "prior_text_encoder", "prior_prior", "unet", "movq"]
+    _prefix = {"prior_pipe": "prior_"}
+
+    def __init__(
+        self,
+        unet: "UNet2DConditionModel",
+        scheduler: "DDPMScheduler",
+        movq: "VQModel",
+        prior_prior: "PriorTransformer",
+        prior_image_encoder: "CLIPVisionModelWithProjection",
+        prior_text_encoder: "CLIPTextModelWithProjection",
+        prior_tokenizer: "CLIPTokenizer",
+        prior_scheduler: "UnCLIPScheduler",
+        prior_image_processor: "CLIPImageProcessor",
+    ):
+        RBLNDiffusionMixin.__init__(self)
+        super(KandinskyV22Img2ImgCombinedPipeline, self).__init__()
+
+        self.register_modules(
+            unet=unet,
+            scheduler=scheduler,
+            movq=movq,
+            prior_prior=prior_prior,
+            prior_image_encoder=prior_image_encoder,
+            prior_text_encoder=prior_text_encoder,
+            prior_tokenizer=prior_tokenizer,
+            prior_scheduler=prior_scheduler,
+            prior_image_processor=prior_image_processor,
+        )
+
+        self.prior_pipe = RBLNKandinskyV22PriorPipeline(
+            prior=prior_prior,
+            image_encoder=prior_image_encoder,
+            text_encoder=prior_text_encoder,
+            tokenizer=prior_tokenizer,
+            scheduler=prior_scheduler,
+            image_processor=prior_image_processor,
+        )
+        self.decoder_pipe = RBLNKandinskyV22Img2ImgPipeline(
+            unet=unet,
+            scheduler=scheduler,
+            movq=movq,
+        )
+
+    def get_compiled_image_size(self):
+        return self.movq.image_size
+
+
 class RBLNKandinskyV22InpaintCombinedPipeline(RBLNDiffusionMixin, KandinskyV22InpaintCombinedPipeline):
     original_class = KandinskyV22InpaintCombinedPipeline
     _connected_classes = {"prior_pipe": RBLNKandinskyV22PriorPipeline, "decoder_pipe": RBLNKandinskyV22InpaintPipeline}
-    _submodules = ["prior_pipe", "decoder_pipe"]
+    _submodules = ["prior_image_encoder", "prior_text_encoder", "prior_prior", "unet", "movq"]
     _prefix = {"prior_pipe": "prior_"}
 
     def __init__(

optimum/rbln/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_img2img.py

@@ -0,0 +1,25 @@
+# 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.
+
+from diffusers import KandinskyV22Img2ImgPipeline
+
+from ...modeling_diffusers import RBLNDiffusionMixin
+
+
+class RBLNKandinskyV22Img2ImgPipeline(RBLNDiffusionMixin, KandinskyV22Img2ImgPipeline):
+    original_class = KandinskyV22Img2ImgPipeline
+    _submodules = ["unet", "movq"]
+
+    def get_compiled_image_size(self):
+        return self.movq.image_size

optimum/rbln/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_inpaint.py

@@ -20,3 +20,6 @@ from ...modeling_diffusers import RBLNDiffusionMixin
 class RBLNKandinskyV22InpaintPipeline(RBLNDiffusionMixin, KandinskyV22InpaintPipeline):
     original_class = KandinskyV22InpaintPipeline
     _submodules = ["unet", "movq"]
+
+    def get_compiled_image_size(self):
+        return self.movq.image_size

optimum/rbln/modeling_base.py

@@ -282,15 +282,6 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
             **kwargs,
         )
 
-    @classmethod
-    def _check_compiled_models(
-        cls, compiled_models: Dict[str, rebel.RBLNCompiledModel], rbln_config: RBLNConfig, config: "PretrainedConfig"
-    ):
-        # check compiled model can create runtimes.
-        # this logic currently only works in LLM
-        # fail when LLM model using Paged Attention can't guarantee max sequence length
-        pass
-
     @classmethod
     def _from_compiled_models(
         cls,
@@ -305,8 +296,6 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
         if isinstance(model_save_dir, str):
             model_save_dir = Path(model_save_dir)
 
-        cls._check_compiled_models(compiled_models=rbln_compiled_models, rbln_config=rbln_config, config=config)
-
         # FIXME:: Should we convert it?
         compiled_model_names = [cfg.compiled_model_name for cfg in rbln_config.compile_cfgs]
         rbln_compiled_models = [rbln_compiled_models[cm_name] for cm_name in compiled_model_names]

optimum/rbln/transformers/models/bart/modeling_bart.py

@@ -108,6 +108,8 @@ class RBLNBartModel(RBLNModel):
 
 
 class RBLNBartForConditionalGeneration(RBLNModelForSeq2SeqLM):
+    support_paged_causal_attn = True
+
     @classmethod
     def wrap_model_if_needed(self, model: "PreTrainedModel", rbln_config: "RBLNConfig"):
         enc_max_seq_len = (

optimum/rbln/transformers/models/decoderonly/decoderonly_architecture.py

@@ -98,6 +98,18 @@ def validate_attention_method(
                 "this requirement, or consider switching `rbln_attn_impl` to 'eager' for shorter lengths."
             )
 
+    if rbln_kvcache_block_size is not None:
+        if rbln_attn_impl == "flash_attn" and rbln_kvcache_partition_len != rbln_kvcache_block_size:
+            raise ValueError(
+                f" When using 'flash attention', the `rbln_kvcache_block_size` ({rbln_kvcache_block_size})  "
+                f"must always be set equal to the `rbln_kvcache_partition_len` {rbln_kvcache_partition_len}."
+            )
+        elif rbln_attn_impl == "eager" and rbln_kvcache_block_size != rbln_max_seq_len:
+            raise ValueError(
+                f" When using 'eager attention', the `rbln_kvcache_block_size` ({rbln_kvcache_block_size})  "
+                f"must always be set equal to the `rbln_max_seq_len` {rbln_max_seq_len}."
+            )
+
     return rbln_attn_impl, rbln_kvcache_partition_len, rbln_kvcache_block_size
 
 

optimum/rbln/transformers/models/decoderonly/modeling_decoderonly.py

@@ -98,9 +98,9 @@ class RBLNRuntimeModel(RBLNPytorchRuntime):
         """
 
         NO_BLOCKS_ERROR = (
-            "No memory blocks are available for allocation."
-            "The generate() API cannot complete this inference task because Paged Attention is not fully supported by optimum-rbln."
-            "This is supported by vllm-rbln (see: https://docs.rbln.ai/software/model_serving/vllm_support/vllm-rbln.html)."
+            "No memory blocks are available for allocation. "
+            "The generate() API cannot complete this inference task because Paged Attention is not fully supported by optimum-rbln. "
+            "This is supported by vllm-rbln (see: https://docs.rbln.ai/software/model_serving/vllm_support/vllm-rbln.html). "
             "Using vllm-rbln should fix this issue and enhance inference performance."
         )
 
@@ -575,59 +575,58 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
         nbits_per_param: int,
         n_model_params: int,
     ) -> int:
-        num_attention_heads = getattr(config, "n_head", None) or getattr(config, "num_attention_heads")
-        num_layers = getattr(config, "n_layer", None) or getattr(config, "num_hidden_layers")
-        head_dim = getattr(config, "head_dim", None) or config.hidden_size // num_attention_heads
-        vocab_size = config.vocab_size
-        hidden_size = getattr(config, "n_embd", None) or getattr(config, "hidden_size")
-        num_key_value_heads = getattr(config, "num_key_value_heads", None) or num_attention_heads
-
-        TARGET_DRAM_LIMIT = int(tensor_parallel_size * 15.7 * 2**30)  # 16GB # TODO(jongho): 더 정확한 값
-
         def align(x: int, nbytes: int) -> int:
             return int(math.ceil(x / nbytes) * nbytes)
 
         def align_2MB(x: int) -> int:
             return align(x, 2 * 1024 * 1024)
 
-        def get_kernel_size() -> int:
-            # TODO: Implement
-            lm_heads_params = align(vocab_size, 64) * hidden_size
-            lm_heads_nbytes = (
-                align_2MB(lm_heads_params * nbits_per_param // 8 / tensor_parallel_size) * tensor_parallel_size
-            )
+        num_attention_heads = getattr(config, "n_head", None) or getattr(config, "num_attention_heads")
+        num_layers = getattr(config, "n_layer", None) or getattr(config, "num_hidden_layers")
+        head_dim = getattr(config, "head_dim", None) or config.hidden_size // num_attention_heads
+        vocab_size = config.vocab_size
+        hidden_size = getattr(config, "n_embd", None) or getattr(config, "hidden_size")
+        num_key_value_heads = getattr(config, "num_key_value_heads", None) or num_attention_heads
 
-            params = n_model_params - lm_heads_params
-            layer_nbytes = (
-                align_2MB(params * nbits_per_param // 8 / num_layers / tensor_parallel_size)
-                * num_layers
-                * tensor_parallel_size
-            )
+        # TODO(jongho): Update if target npu is REBEL.
+        ATOM_DRAM_NBYTES = 16 * 2**30
+        ATOM_SYS_DRAM_NBYTES = 288 * 2**20
+        available_dram = tensor_parallel_size * (ATOM_DRAM_NBYTES - ATOM_SYS_DRAM_NBYTES)
 
-            return layer_nbytes + lm_heads_nbytes
+        # Get estimated kernel size (approximated)
+        lm_heads_params = align(vocab_size, 64) * hidden_size
+        lm_heads_nbytes = (
+            align_2MB(lm_heads_params * nbits_per_param // 8 / tensor_parallel_size) * tensor_parallel_size
+        )
+        params = n_model_params - lm_heads_params
+        layer_nbytes = (
+            align_2MB(params * nbits_per_param // 8 / num_layers / tensor_parallel_size)
+            * num_layers
+            * tensor_parallel_size
+        )
+        kernel_size = layer_nbytes + lm_heads_nbytes
 
-        available_dram = TARGET_DRAM_LIMIT - get_kernel_size()
+        available_dram -= kernel_size
 
-        buffer = 2**30  # 1GB
-        if tensor_parallel_size <= 2:
+        # TODO: Accurate buffer estimation
+        buffer = 2**30  # 1GB Buffer
+        if tensor_parallel_size <= 4:
             buffer /= 4
 
         available_dram -= buffer
 
-        def get_nbytes_per_block() -> int:
-            return (
-                align_2MB(
-                    kvcache_block_size
-                    * head_dim
-                    * math.ceil(num_key_value_heads / tensor_parallel_size)  # Shard
-                    * 2  # (fp16)
-                )
-                * num_layers
-                * 2  # (k, v)
-                * tensor_parallel_size
+        # Estimate nbytes per a single kvcache block
+        nbytes_per_block = (
+            align_2MB(
+                kvcache_block_size
+                * head_dim
+                * math.ceil(num_key_value_heads / tensor_parallel_size)  # Shard
+                * 2  # (fp16)
             )
-
-        nbytes_per_block = get_nbytes_per_block()
+            * num_layers
+            * 2  # (k, v)
+            * tensor_parallel_size
+        )
         n_blocks = available_dram // nbytes_per_block
 
         return n_blocks, nbytes_per_block

optimum/rbln/transformers/models/seq2seq/modeling_seq2seq.py

@@ -50,6 +50,7 @@ class RBLNRuntimeDecoder(RBLNPytorchRuntime):
         runtime: rebel.Runtime,
         batch_size: int,
         dec_max_seq_len: int,
+        support_paged_causal_attn: Optional[bool] = None,
         use_attention_mask: Optional[bool] = None,
         **kwargs: Any,
     ) -> None:
@@ -57,7 +58,10 @@ class RBLNRuntimeDecoder(RBLNPytorchRuntime):
         self.batch_size = batch_size
         self.dec_max_seq_len = dec_max_seq_len
         self.use_attention_mask = use_attention_mask
-        self.default_block_tables = torch.arange(0, self.batch_size, dtype=torch.int16).view(self.batch_size, 1)
+        if support_paged_causal_attn:
+            self.default_block_tables = torch.arange(0, self.batch_size, dtype=torch.int16).view(self.batch_size, 1)
+        else:
+            self.default_block_tables = None
 
     def forward(
         self,
@@ -94,7 +98,7 @@ class RBLNRuntimeDecoder(RBLNPytorchRuntime):
             decoder_attention_mask if self.use_attention_mask else None,
             attention_mask,
             cache_position,
-            block_tables,
+            block_tables=block_tables,
         )
 
         return Seq2SeqLMOutput(logits=lm_logits)
@@ -115,6 +119,7 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
 
     main_input_name = "input_ids"
     auto_model_class = AutoModelForSeq2SeqLM
+    support_paged_causal_attn = None
 
     def __post_init__(self, **kwargs):
         batch_size = self.rbln_config.model_cfg["batch_size"]
@@ -130,6 +135,7 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
             main_input_name="input_ids",
             batch_size=batch_size,
             dec_max_seq_len=dec_max_seq_len,
+            support_paged_causal_attn=self.support_paged_causal_attn,
             use_attention_mask=self.use_attention_mask,
         )
 
@@ -186,13 +192,16 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
         rbln_dec_max_seq_len = rbln_kwargs.get("dec_max_seq_len", None)
         rbln_batch_size = rbln_kwargs.get("batch_size", None)
         rbln_batch_size = 1 if rbln_batch_size is None else rbln_batch_size
-        rbln_use_attention_mask = rbln_kwargs.get("use_attention_mask", None)
 
-        if rbln_use_attention_mask is None:
-            rbln_use_attention_mask = False
-            rbln_npu = rbln_kwargs.get("npu", None) or rebel.get_npu_name()
-            if rbln_npu == "RBLN-CA02":
-                rbln_use_attention_mask = True
+        if cls.support_paged_causal_attn:
+            rbln_use_attention_mask = rbln_kwargs.get("use_attention_mask", None)
+            if rbln_use_attention_mask is None:
+                rbln_use_attention_mask = False
+                rbln_npu = rbln_kwargs.get("npu", None) or rebel.get_npu_name()
+                if rbln_npu == "RBLN-CA02":
+                    rbln_use_attention_mask = True
+        else:
+            rbln_use_attention_mask = True
 
         n_layer = getattr(model_config, "decoder_layers", None) or getattr(model_config, "num_layers")
         n_head = getattr(model_config, "decoder_attention_heads", None) or getattr(model_config, "num_heads")
@@ -265,11 +274,6 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
                 [rbln_batch_size, 1],
                 "int32",
             ),
-            (
-                "block_tables",
-                [rbln_batch_size, 1],
-                "int16",
-            ),
         ]
         dec_input_info.extend(
             [
@@ -302,6 +306,8 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
             ]
         )
 
+        if cls.support_paged_causal_attn:
+            dec_input_info.insert(3, ("block_tables", [rbln_batch_size, 1], "int16"))
         if rbln_use_attention_mask:
             dec_input_info.insert(1, ("attention_mask", [rbln_batch_size, rbln_dec_max_seq_len], "float32"))
 

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

@@ -13,9 +13,8 @@
 # limitations under the License.
 
 import inspect
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Tuple, Union
 
-import rebel
 import torch
 from transformers import (
     AutoModelForTextEncoding,
@@ -23,7 +22,7 @@ from transformers import (
     T5EncoderModel,
     T5ForConditionalGeneration,
 )
-from transformers.modeling_outputs import BaseModelOutput, Seq2SeqLMOutput
+from transformers.modeling_outputs import BaseModelOutput
 
 from ....diffusers.modeling_diffusers import RBLNDiffusionMixin
 from ....modeling import RBLNModel
@@ -58,63 +57,6 @@ class RBLNRuntimeModel(RBLNPytorchRuntime):
         )
 
 
-class RBLNRuntimeEncoder(RBLNPytorchRuntime):
-    mandatory_members = ["main_input_name"]
-
-    def forward(self, *args: List[torch.Tensor], **kwargs: Dict[str, torch.Tensor]):
-        _ = super().forward(*args, **kwargs)
-        return BaseModelOutput(last_hidden_state=torch.tensor([1.0]))
-
-
-class RBLNRuntimeDecoder(RBLNPytorchRuntime):
-    mandatory_members = ["main_input_name"]
-
-    def __init__(
-        self,
-        runtime: rebel.Runtime,
-        batch_size: int,
-        dec_max_seq_len: int,
-        **kwargs: Any,
-    ) -> None:
-        super().__init__(runtime, **kwargs)
-        self.batch_size = batch_size
-        self.dec_max_seq_len = dec_max_seq_len
-
-    def forward(
-        self,
-        decoder_input_ids: Optional[torch.LongTensor] = None,
-        attention_mask: Optional[torch.FloatTensor] = None,
-        decoder_attention_mask: Optional[torch.BoolTensor] = None,
-        cache_position: Optional[torch.Tensor] = None,
-        **kwargs,
-    ) -> Tuple[torch.FloatTensor]:
-        batch_size = decoder_input_ids.shape[0]
-        if batch_size != self.batch_size:
-            raise RuntimeError(
-                f"Batch size mismatch: got {batch_size}, expected {self.batch_size} (compiled batch size)."
-            )
-
-        if batch_size != cache_position.shape[0]:
-            raise RuntimeError(f"Cache position size mismatch: got {cache_position.shape[0]}, expected {batch_size}.")
-
-        for b_idx in range(self.batch_size):
-            decoding_step = cache_position[b_idx].item()
-            if not (0 <= decoding_step < self.dec_max_seq_len):
-                raise ValueError(
-                    f"Decoding step {decoding_step} out of bounds for attention mask with shape {self.dec_attn_mask.shape}."
-                )
-            decoder_attention_mask[b_idx, : decoding_step + 1] = 1
-
-        lm_logits = super().forward(
-            decoder_input_ids,
-            decoder_attention_mask,
-            attention_mask,
-            cache_position,
-        )
-
-        return Seq2SeqLMOutput(logits=lm_logits)
-
-
 class T5EncoderWrapper(torch.nn.Module):
     def __init__(self, model: "T5EncoderModel") -> None:
         super().__init__()
@@ -247,20 +189,7 @@ class RBLNT5EncoderModel(RBLNModel):
 
 
 class RBLNT5ForConditionalGeneration(RBLNModelForSeq2SeqLM):
-    def __post_init__(self, **kwargs):
-        batch_size = self.rbln_config.model_cfg["batch_size"]
-        dec_max_seq_len = self.rbln_config.model_cfg["dec_max_seq_len"]
-
-        self.encoder = RBLNRuntimeEncoder(
-            runtime=self.model[0],
-            main_input_name="input_ids",
-        )
-        self.decoder = RBLNRuntimeDecoder(
-            runtime=self.model[1],
-            main_input_name="input_ids",
-            batch_size=batch_size,
-            dec_max_seq_len=dec_max_seq_len,
-        )
+    support_causal_paged_attn = False
 
     @classmethod
     def wrap_model_if_needed(self, model: "PreTrainedModel", rbln_config: "RBLNConfig"):
@@ -279,139 +208,3 @@ class RBLNT5ForConditionalGeneration(RBLNModelForSeq2SeqLM):
             return redirect(val)
 
         return val
-
-    @classmethod
-    def _get_rbln_config(
-        cls,
-        preprocessors: Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"],
-        model_config: "PretrainedConfig",
-        rbln_kwargs: Dict[str, Any] = {},
-    ) -> RBLNConfig:
-        rbln_enc_max_seq_len = rbln_kwargs.get("enc_max_seq_len", None)
-        rbln_dec_max_seq_len = rbln_kwargs.get("dec_max_seq_len", None)
-        rbln_batch_size = rbln_kwargs.get("batch_size", None)
-        rbln_batch_size = 1 if rbln_batch_size is None else rbln_batch_size
-
-        n_layer = getattr(model_config, "decoder_layers", None) or getattr(model_config, "num_layers")
-        n_head = getattr(model_config, "decoder_attention_heads", None) or getattr(model_config, "num_heads")
-        d_kv = (
-            model_config.d_kv
-            if hasattr(model_config, "d_kv")
-            else model_config.d_model // model_config.encoder_attention_heads
-        )
-
-        max_position_embeddings = getattr(model_config, "n_positions", None) or getattr(
-            model_config, "max_position_embeddings", None
-        )
-
-        rbln_pad_token_id = getattr(model_config, "pad_token_id", None)
-        if rbln_pad_token_id is None:
-            rbln_pad_token_id = getattr(model_config, "bos_token_id", None)
-            if rbln_pad_token_id is None:
-                rbln_pad_token_id = getattr(model_config, "eos_token_id", None)
-                if rbln_pad_token_id is None:
-                    rbln_pad_token_id = -1
-
-        if rbln_enc_max_seq_len is None:
-            rbln_enc_max_seq_len = max_position_embeddings
-            if rbln_enc_max_seq_len is None:
-                for tokenizer in preprocessors:
-                    if hasattr(tokenizer, "model_max_length"):
-                        rbln_enc_max_seq_len = tokenizer.model_max_length
-                        break
-                if rbln_enc_max_seq_len is None:
-                    raise ValueError("`rbln_enc_max_seq_len` should be specified!")
-        if max_position_embeddings is not None and rbln_enc_max_seq_len > max_position_embeddings:
-            raise ValueError("`rbln_enc_max_seq_len` should be less or equal than max_position_embeddings!")
-
-        if rbln_dec_max_seq_len is None:
-            rbln_dec_max_seq_len = max_position_embeddings
-            if rbln_dec_max_seq_len is None:
-                for tokenizer in preprocessors:
-                    if hasattr(tokenizer, "model_max_length"):
-                        rbln_dec_max_seq_len = tokenizer.model_max_length
-                        break
-                if rbln_dec_max_seq_len is None:
-                    raise ValueError("`rbln_dec_max_seq_len` should be specified!")
-
-        if max_position_embeddings is not None and rbln_dec_max_seq_len > max_position_embeddings:
-            raise ValueError("`rbln_dec_max_seq_len` should be less or equal than max_position_embeddings!")
-
-        # model input info
-        enc_input_info = [
-            ("input_ids", [1, rbln_enc_max_seq_len], "int64"),
-            ("attention_mask", [1, rbln_enc_max_seq_len], "float32"),
-            (
-                "cross_key_value_states",
-                [
-                    n_layer * 2,
-                    rbln_batch_size,
-                    n_head,
-                    rbln_enc_max_seq_len,
-                    d_kv,
-                ],
-                "float32",
-            ),
-            ("block_tables", [1], "int16"),
-        ]
-
-        dec_input_info = [
-            ("input_ids", [rbln_batch_size, 1], "int64"),
-            ("attention_mask", [rbln_batch_size, rbln_dec_max_seq_len], "float32"),
-            ("encoder_attention_mask", [rbln_batch_size, rbln_enc_max_seq_len], "float32"),
-            (
-                "cache_position",
-                [rbln_batch_size, 1],
-                "int32",
-            ),
-        ]
-        dec_input_info.extend(
-            [
-                (
-                    "cross_key_value_states",
-                    [
-                        n_layer * 2,
-                        rbln_batch_size,
-                        n_head,
-                        rbln_enc_max_seq_len,
-                        d_kv,
-                    ],
-                    "float32",
-                )
-            ]
-        )
-        dec_input_info.extend(
-            [
-                (
-                    f"self_key_value_states_{i}",
-                    [
-                        rbln_batch_size,
-                        n_head,
-                        rbln_dec_max_seq_len,
-                        d_kv,
-                    ],
-                    "float32",
-                )
-                for i in range(n_layer * 2)
-            ]
-        )
-
-        enc_compile_config = RBLNCompileConfig(compiled_model_name="encoder", input_info=enc_input_info)
-        dec_compile_config = RBLNCompileConfig(compiled_model_name="decoder", input_info=dec_input_info)
-
-        rbln_config = RBLNConfig(
-            rbln_cls=cls.__name__,
-            compile_cfgs=[enc_compile_config, dec_compile_config],
-            rbln_kwargs=rbln_kwargs,
-        )
-
-        rbln_config.model_cfg.update(
-            {
-                "enc_max_seq_len": rbln_enc_max_seq_len,
-                "dec_max_seq_len": rbln_dec_max_seq_len,
-                "batch_size": rbln_batch_size,
-                "pad_token_id": rbln_pad_token_id,
-            }
-        )
-
-        return rbln_config

optimum/rbln/utils/import_utils.py

@@ -28,6 +28,13 @@ class VersionCompat:
 
 
 RBLN_VERSION_COMPATS = {
+    "0.7.3": [
+        VersionCompat(
+            package_name="rebel-compiler",
+            min_version="0.7.3",
+            max_version="0.7.4",
+        ),
+    ],
     "0.7.2": [
         VersionCompat(
             package_name="rebel-compiler",

{optimum_rbln-0.7.3a5.dist-info → optimum_rbln-0.7.3.post1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: optimum-rbln
-Version: 0.7.3a5
+Version: 0.7.3.post1
 Summary: Optimum RBLN is the interface between the Hugging Face 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

{optimum_rbln-0.7.3a5.dist-info → optimum_rbln-0.7.3.post1.dist-info}/RECORD

@@ -1,31 +1,33 @@
-optimum/rbln/__init__.py,sha256=eHi15YM3989AcX52jka9rUmgAtlp1PHqMNwBEdOfuu8,6554
-optimum/rbln/__version__.py,sha256=MC3yJ2-M633KXIQTqHjv3l9eWLzkJkdlmhuQkRiV278,519
+optimum/rbln/__init__.py,sha256=ZDzXcl-oAcYJhKjJMpotjbTih9awo7HzUb6T3MUEP6Q,6894
+optimum/rbln/__version__.py,sha256=aegWGVZeZJ9bIegWWNAgPL2y9SAs5kPTsXCQi0EZ9go,517
 optimum/rbln/modeling.py,sha256=nJsAs5zs--VVOYGFjYNpqfxYIemJIK4Lr0WEzlDLdP0,8390
-optimum/rbln/modeling_base.py,sha256=Ow73GVJF1N5cDFO8_rgirtGj1wC-cXBDyqXHW5PCybA,22270
+optimum/rbln/modeling_base.py,sha256=dNCL-BhrWCpuOVkZaj8-MW567Tf4lLo3p3Z3ldjWJfU,21779
 optimum/rbln/modeling_config.py,sha256=7104bxmrvKW4Q6XTruQayiIGl8GHDFmPkJ3cknMIInE,11335
-optimum/rbln/diffusers/__init__.py,sha256=pOyoXv3-JRzTBSwPKbgLS9H6F2K9dJdReEmpGhcLQYU,3283
-optimum/rbln/diffusers/modeling_diffusers.py,sha256=zqVNgH9oeOx2iNE7VsW_FinVf4s6G5Idyh4TKz7XJJg,21116
+optimum/rbln/diffusers/__init__.py,sha256=Hq87CbtiCy85YmK2SB-OmUyfv77oe3j4bsTenTRnu6w,3623
+optimum/rbln/diffusers/modeling_diffusers.py,sha256=IS6Mlgexofap7f9Lefk5cKFP7ejSG_oWN3v2PX9_IDQ,20118
 optimum/rbln/diffusers/models/__init__.py,sha256=mkCvJyH1KcwrsUvYSq_bVC79oOfyqtBSFDyPS1_48wA,1478
 optimum/rbln/diffusers/models/controlnet.py,sha256=EM_HlzCdaZdnnK0oGpY2fQeigPqHhlwh4NHCzlmoumI,10512
 optimum/rbln/diffusers/models/autoencoders/__init__.py,sha256=dg17ZTUsiqTcbIaEE4fqew9uRbao0diQ21PXvRKIqKg,679
 optimum/rbln/diffusers/models/autoencoders/autoencoder_kl.py,sha256=rCbC32bJnfXtsLdVvNVVHpRAkCYy6jeCSwIZ-JSReWk,9220
 optimum/rbln/diffusers/models/autoencoders/vae.py,sha256=gB9HR7Bf7wpIXLv-Js4Pc3oyWRlqEe4cms4sI2AJicY,4380
-optimum/rbln/diffusers/models/autoencoders/vq_model.py,sha256=GunIau02_-lodYZBzd0ktJSNRT5axEFIZxSAfj2Mlyo,5974
+optimum/rbln/diffusers/models/autoencoders/vq_model.py,sha256=b36QqPbayjApKivceQVVyQxHyR1ZOZ1ffuGgdALEPTQ,6117
 optimum/rbln/diffusers/models/transformers/__init__.py,sha256=V8rSR7WzHs-i8Cwb_MNxhY2NFbwPgxu24vGtkwl-6tk,706
 optimum/rbln/diffusers/models/transformers/prior_transformer.py,sha256=VG9cQo-_eppDvQSW1q1euAGBt1socUHetN_fIN2u1iU,6169
 optimum/rbln/diffusers/models/transformers/transformer_sd3.py,sha256=n_krmMgiRxWrG--567PNpk58EG_X7x7H4gidIkRvwjo,7308
 optimum/rbln/diffusers/models/unets/__init__.py,sha256=MaICuK9CWjgzejXy8y2NDrphuEq1rkzanF8u45k6O5I,655
-optimum/rbln/diffusers/models/unets/unet_2d_condition.py,sha256=xHnBzFrm7aNaolxrsotbjo9GkbNiNdTleXQoeqGLlhg,15540
-optimum/rbln/diffusers/pipelines/__init__.py,sha256=DAsM4eNks3hEY-bsUKSxRKmgwUWDGDlw82gfplSOdO8,2800
+optimum/rbln/diffusers/models/unets/unet_2d_condition.py,sha256=QIjVWQQf8KBn5rU7lvipdm3gNBxZl7l6HCAj7p5FjLU,15977
+optimum/rbln/diffusers/pipelines/__init__.py,sha256=5KLZ5LrpMzBya2e_3_PvEoPwG24U8JMexfw_ygZREKc,3140
 optimum/rbln/diffusers/pipelines/controlnet/__init__.py,sha256=n1Ef22TSeax-kENi_d8K6wGGHSNEo9QkUeygELHgcao,983
 optimum/rbln/diffusers/pipelines/controlnet/multicontrolnet.py,sha256=JWKtnZYBIfgmbAo0SLFIvHBQCv2BPSFNvpcdjG4GUOY,4113
 optimum/rbln/diffusers/pipelines/controlnet/pipeline_controlnet.py,sha256=dGdw5cwJLS4CLv6IHskk5ZCcPgS7UDuHKbfOZ8ojNUs,35187
 optimum/rbln/diffusers/pipelines/controlnet/pipeline_controlnet_img2img.py,sha256=7xCiXrH4ToCTHohVGFXqO7_f9G8HShYaHgZxoMZARkQ,33664
 optimum/rbln/diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl.py,sha256=Gzt2wg4dgFg0TV3Bu0cs8Xru3wVrxWUxxgciwZ-QKLE,44755
 optimum/rbln/diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl_img2img.py,sha256=RfwxNX_zQWFtvvFQJ5bt3qtHbdYdQV_3XLHm9WYCKOs,46084
-optimum/rbln/diffusers/pipelines/kandinsky2_2/__init__.py,sha256=YFqA76_XiMNxPwqotbHug2kd7jCbOXOu5NlxG2hbaVs,808
-optimum/rbln/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_combined.py,sha256=9szfe1NvOr1mgDnSPZvBGq1b65RElUrqLVhuErY3Dmw,2962
-optimum/rbln/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_inpaint.py,sha256=WxBbHAZSAKDSWhFerPvUlIhhWEsejW4NmhwmWX-_b54,856
+optimum/rbln/diffusers/pipelines/kandinsky2_2/__init__.py,sha256=I4YQq2HfA3xONbWsdJ870IEJPyLWeCDDG-UCJsu9YO8,1035
+optimum/rbln/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2.py,sha256=aNFGOjth8tDvPrjYLbRWrkHr6p-8AFgcQx1Qay1fw70,904
+optimum/rbln/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_combined.py,sha256=BVXOpdrezWVTCibpuAMu9KkD5oEQUY00cSqm6dFbTnk,7020
+optimum/rbln/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_img2img.py,sha256=fEs-WgJqWs5zvuCkKb7MuZokH9Mi6q-0DOEKxzfWxzo,932
+optimum/rbln/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_inpaint.py,sha256=Ad2ZYCXaMiYpB0mz-8X1CGhILxrVbt7rRIXt6IPwYBM,932
 optimum/rbln/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior.py,sha256=Mf7tzrXetwCgt7LuXfkX-CX1hltLgNZdwF9bHxAbDJM,874
 optimum/rbln/diffusers/pipelines/stable_diffusion/__init__.py,sha256=gz6CbP4T6w8XH3PIGRIJXTmKFsChJIkwcAEAsiR5Ydg,830
 optimum/rbln/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py,sha256=DgRLzO9HxtgE1jICmHoHaqeVXM4Ih-5uo2JqNMAPMcc,876
@@ -53,14 +55,14 @@ optimum/rbln/transformers/models/auto/auto_factory.py,sha256=IK9jFrJ3EEzYQa9_aKp
 optimum/rbln/transformers/models/auto/modeling_auto.py,sha256=Un9qoqdy3dO8JBza_bTJF_6_fRVNM9QisihSgTRFI-o,3933
 optimum/rbln/transformers/models/bart/__init__.py,sha256=32HPe0_GIO0hp9U464Iv6Jd7M-1nop9g8hA1UZMHhyw,674
 optimum/rbln/transformers/models/bart/bart_architecture.py,sha256=Oo-Cdne7igKEex8wwP-gztKJHgs5GLHQjK1oc3IZIDE,5801
-optimum/rbln/transformers/models/bart/modeling_bart.py,sha256=iI3ubPOVvHmhLt0wEz_vkOfMyNTHVNjmnkLtbpOX760,5797
+optimum/rbln/transformers/models/bart/modeling_bart.py,sha256=6IpWXlBCd02v66KF77oEWfrv8-FnPBYjjjL_8KZL3Ow,5835
 optimum/rbln/transformers/models/bert/__init__.py,sha256=YVV7k_laU6yJBawZrgjIWjRmIF-Y4oQQHqyf8lsraQs,691
 optimum/rbln/transformers/models/bert/modeling_bert.py,sha256=p3utRqf3dv9_RkHwaMCa1EfXttNJkqCJUIZo3CeZ9YY,4674
 optimum/rbln/transformers/models/clip/__init__.py,sha256=H9vuBwrmFO0-CqZhXUrKF-uQL6igCqMlqrT1X_ELaAI,754
 optimum/rbln/transformers/models/clip/modeling_clip.py,sha256=NiSm7bHs4SReHDUr53BBWSX0Y8bkKOeUSpsBDrp8YDw,6628
 optimum/rbln/transformers/models/decoderonly/__init__.py,sha256=pDogsdpJKKB5rqnVFrRjwfhUvOSV-jZ3oARMsqSvOOQ,665
-optimum/rbln/transformers/models/decoderonly/decoderonly_architecture.py,sha256=7OIKteJLKNxOLOg0w3lLOM7TxZovQn4jkglI9wRkrtQ,40609
-optimum/rbln/transformers/models/decoderonly/modeling_decoderonly.py,sha256=gcIkpRSsJycLtkFyVxU8PblzNhKlsIy5fDSDUlwhflM,38884
+optimum/rbln/transformers/models/decoderonly/decoderonly_architecture.py,sha256=m93-qKN7NMw3i0XDmFmttmRIRK4np_fWtLFlBb2RFgU,41351
+optimum/rbln/transformers/models/decoderonly/modeling_decoderonly.py,sha256=uGdPGcFrWm2gAwFLjfBiALwFsl49VGCReVi4NUfOPxM,38898
 optimum/rbln/transformers/models/dpt/__init__.py,sha256=gP1tkR3XMNlHq1GT87ugIVvb2o_1eAUg1JaniXjy1Lw,651
 optimum/rbln/transformers/models/dpt/modeling_dpt.py,sha256=ZsS2SOiqcA4azULB-WFEMQZbgIoOyVUKqVKqrw_tWzA,3430
 optimum/rbln/transformers/models/exaone/__init__.py,sha256=zYH_5tVa8-juEdsOIky7I33WSC3Zuhoq1upI0OHYeVw,859
@@ -90,10 +92,10 @@ optimum/rbln/transformers/models/qwen2/__init__.py,sha256=RAMWc21W_2I6DH9xBjeNxP
 optimum/rbln/transformers/models/qwen2/modeling_qwen2.py,sha256=9-aFDvjMzPNUyGOz0qo33RE18bUFGYZ3Wt_68zb5uJY,1530
 optimum/rbln/transformers/models/qwen2/qwen2_architecture.py,sha256=XlNAMYAcDLohnSAhIFGKOPuCB5XLgzYs5ABWdeQSaZs,720
 optimum/rbln/transformers/models/seq2seq/__init__.py,sha256=EmEMV4rOYqKyruX85d0fR73-b8N6BSD6CPcbpYdBuVk,651
-optimum/rbln/transformers/models/seq2seq/modeling_seq2seq.py,sha256=NPfJf9Uk_bYOae7hXGHwteGiWH0va63Z-D93RmAMENg,17611
+optimum/rbln/transformers/models/seq2seq/modeling_seq2seq.py,sha256=9Pf9Y86ABDfhwIenlZqYfgqjbyFmtKBiPnbCD7zxw4M,18017
 optimum/rbln/transformers/models/seq2seq/seq2seq_architecture.py,sha256=tvzacIZam1sIr_1BvvZ_fDr8u5dXAiYiynFdX9tArtY,18877
 optimum/rbln/transformers/models/t5/__init__.py,sha256=1skR1RmnG62WTAP3-F5P1x-V_ReFhMyirH3u56vWwvc,675
-optimum/rbln/transformers/models/t5/modeling_t5.py,sha256=nKRR3eH1EAu1YkKvhlqGyTrJXDRd-IWB5LOeG9jrcb4,16021
+optimum/rbln/transformers/models/t5/modeling_t5.py,sha256=8PAhPlYT1dmpcWM7hUMmZV9lPd4d75CuMuFen1pzr3Q,8088
 optimum/rbln/transformers/models/t5/t5_architecture.py,sha256=AArCQhZRETVM583wlIRzMFOSYq7t2nzxaAeyhZxyxKk,9508
 optimum/rbln/transformers/models/wav2vec2/__init__.py,sha256=YpgA0K-vyg9veh0eL_jxauosbRpb_kpGKHvvQLBspKM,649
 optimum/rbln/transformers/models/wav2vec2/modeling_wav2vec2.py,sha256=JYJmV52j6cBwim4RanVJryfKnV80V96ol0A-oR6o7cg,3856
@@ -108,13 +110,13 @@ optimum/rbln/transformers/utils/rbln_quantization.py,sha256=gwBVHf97sQgPNmGa0wq8
 optimum/rbln/utils/__init__.py,sha256=ieDBT2VFTt2E0M4v_POLBpuGW9LxSydpb_DuPd6PQqc,712
 optimum/rbln/utils/decorator_utils.py,sha256=xu-TrsNi33SRC2a7DBsyoo6-pEQxWKZPZSmM9QlDe2Y,3745
 optimum/rbln/utils/hub.py,sha256=bNmOJGEO9Jfux4Cg8Xli-898I4mxk20KuwQOhP0Zs1U,4198
-optimum/rbln/utils/import_utils.py,sha256=n4HcvZPzFW2ytl45qJ4ZQYlrRSoOb0-nnqhyT2_JA8M,4224
+optimum/rbln/utils/import_utils.py,sha256=uMldLJmDVMj5uHvxBfb96uV29bfGEDvlksLY26GOHAs,4389
 optimum/rbln/utils/logging.py,sha256=VKKBmlQSdg6iZCGmAXaWYiW67K84jyp1QJhLQSSjPPE,3453
 optimum/rbln/utils/model_utils.py,sha256=DfD_Z2qvZHqcddXqnzTM1AN8khanj3-DXK2lJvVxDvs,1278
 optimum/rbln/utils/runtime_utils.py,sha256=5-DYniyP59nx-mrrbi7AqA77L85b4Cm5oLpaxidSyss,3699
 optimum/rbln/utils/save_utils.py,sha256=hG5uOtYmecSXZuGTvCXsTM-SiyZpr5q3InUGCCq_jzQ,3619
 optimum/rbln/utils/submodule.py,sha256=oZoGrItB8WqY4i-K9WJPlLlcLohc1YGB9OHB8_XZw3A,4071
-optimum_rbln-0.7.3a5.dist-info/METADATA,sha256=XgkOm4f_xhli40HbenyjQYm66ZGna1Pv1prBnpF5N5E,5300
-optimum_rbln-0.7.3a5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-optimum_rbln-0.7.3a5.dist-info/licenses/LICENSE,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
-optimum_rbln-0.7.3a5.dist-info/RECORD,,
+optimum_rbln-0.7.3.post1.dist-info/METADATA,sha256=dKER74SsqGQwVQgTXVM854y97xzhfRl5LKaGedd4IIw,5304
+optimum_rbln-0.7.3.post1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+optimum_rbln-0.7.3.post1.dist-info/licenses/LICENSE,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
+optimum_rbln-0.7.3.post1.dist-info/RECORD,,