optimum-rbln 0.7.2rc2__py3-none-any.whl → 0.7.3__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.2rc2'
-__version_tuple__ = version_tuple = (0, 7, 2)
+__version__ = version = '0.7.3'
+__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,17 +66,16 @@ class RBLNDiffusionMixin:
           as keys in rbln_config
     """
 
+    _connected_classes = {}
     _submodules = []
     _prefix = {}
 
     @classmethod
-    @property
-    def img2img_pipeline(cls):
+    def is_img2img_pipeline(cls):
         return "Img2Img" in cls.__name__
 
     @classmethod
-    @property
-    def inpaint_pipeline(cls):
+    def is_inpaint_pipeline(cls):
         return "Inpaint" in cls.__name__
 
     @classmethod
@@ -100,42 +98,11 @@ class RBLNDiffusionMixin:
             submodule_config.update({k: v for k, v in pipe_global_config.items() if k not in submodule_config})
             submodule_config.update(
                 {
-                    "img2img_pipeline": cls.img2img_pipeline,
-                    "inpaint_pipeline": cls.inpaint_pipeline,
+                    "img2img_pipeline": cls.is_img2img_pipeline(),
+                    "inpaint_pipeline": cls.is_inpaint_pipeline(),
                 }
             )
             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
@@ -201,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`."
@@ -268,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,
@@ -309,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__} ")
 
@@ -376,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)
@@ -400,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.
         #
@@ -422,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.py

@@ -134,6 +134,9 @@ class RBLNModel(RBLNBaseModel):
         for preprocessor in preprocessors:
             preprocessor.save_pretrained(save_dir_path / subfolder)
 
+        # ad-hoc
+        rbln_kwargs["n_model_params"] = sum(p.numel() for p in model.parameters())
+
         # Get compilation arguments (e.g. input_info)
         rbln_config: RBLNConfig = cls.get_rbln_config(
             preprocessors=preprocessors, model_config=config, rbln_kwargs=rbln_kwargs
@@ -196,7 +199,7 @@ class RBLNModel(RBLNBaseModel):
         **kwargs,
     ) -> "PreTrainedModel":
         kwargs = cls.update_kwargs(kwargs)
-        return cls.hf_class.from_pretrained(
+        return cls.get_hf_class().from_pretrained(
             model_id,
             subfolder=subfolder,
             revision=revision,