optimum-rbln 0.2.1a5__py3-none-any.whl → 0.7.2rc0__py3-none-any.whl

optimum/rbln/__init__.py

@@ -48,6 +48,7 @@ _import_structure = {
         "RBLNCLIPTextModel",
         "RBLNCLIPTextModelWithProjection",
         "RBLNCLIPVisionModel",
+        "RBLNCLIPVisionModelWithProjection",
         "RBLNDPTForDepthEstimation",
         "RBLNExaoneForCausalLM",
         "RBLNGemmaForCausalLM",
@@ -74,11 +75,15 @@ _import_structure = {
         "RBLNBertForMaskedLM",
     ],
     "diffusers": [
+        "RBLNAutoencoderKL",
+        "RBLNControlNetModel",
+        "RBLNPriorTransformer",
+        "RBLNKandinskyV22InpaintCombinedPipeline",
+        "RBLNKandinskyV22InpaintPipeline",
+        "RBLNKandinskyV22PriorPipeline",
         "RBLNStableDiffusionPipeline",
         "RBLNStableDiffusionXLPipeline",
-        "RBLNAutoencoderKL",
         "RBLNUNet2DConditionModel",
-        "RBLNControlNetModel",
         "RBLNStableDiffusionImg2ImgPipeline",
         "RBLNStableDiffusionInpaintPipeline",
         "RBLNStableDiffusionControlNetImg2ImgPipeline",
@@ -88,6 +93,7 @@ _import_structure = {
         "RBLNStableDiffusionControlNetPipeline",
         "RBLNStableDiffusionXLControlNetPipeline",
         "RBLNStableDiffusionXLControlNetImg2ImgPipeline",
+        "RBLNVQModel",
         "RBLNSD3Transformer2DModel",
         "RBLNStableDiffusion3Img2ImgPipeline",
         "RBLNStableDiffusion3InpaintPipeline",
@@ -101,7 +107,11 @@ if TYPE_CHECKING:
         RBLNAutoencoderKL,
         RBLNControlNetModel,
         RBLNDiffusionMixin,
+        RBLNKandinskyV22InpaintCombinedPipeline,
+        RBLNKandinskyV22InpaintPipeline,
+        RBLNKandinskyV22PriorPipeline,
         RBLNMultiControlNetModel,
+        RBLNPriorTransformer,
         RBLNSD3Transformer2DModel,
         RBLNStableDiffusion3Img2ImgPipeline,
         RBLNStableDiffusion3InpaintPipeline,
@@ -117,6 +127,7 @@ if TYPE_CHECKING:
         RBLNStableDiffusionXLInpaintPipeline,
         RBLNStableDiffusionXLPipeline,
         RBLNUNet2DConditionModel,
+        RBLNVQModel,
     )
     from .modeling import (
         RBLNBaseModel,
@@ -148,6 +159,7 @@ if TYPE_CHECKING:
         RBLNCLIPTextModel,
         RBLNCLIPTextModelWithProjection,
         RBLNCLIPVisionModel,
+        RBLNCLIPVisionModelWithProjection,
         RBLNDistilBertForQuestionAnswering,
         RBLNDPTForDepthEstimation,
         RBLNExaoneForCausalLM,

optimum/rbln/__version__.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.2.1a5'
-__version_tuple__ = version_tuple = (0, 2, 1)
+__version__ = version = '0.7.2rc0'
+__version_tuple__ = version_tuple = (0, 7, 2)

optimum/rbln/diffusers/__init__.py

@@ -24,6 +24,9 @@ ALL_IMPORTABLE_CLASSES.update(LOADABLE_CLASSES["optimum.rbln"])
 
 _import_structure = {
     "pipelines": [
+        "RBLNKandinskyV22InpaintCombinedPipeline",
+        "RBLNKandinskyV22InpaintPipeline",
+        "RBLNKandinskyV22PriorPipeline",
         "RBLNStableDiffusionPipeline",
         "RBLNStableDiffusionXLPipeline",
         "RBLNStableDiffusionImg2ImgPipeline",
@@ -44,6 +47,8 @@ _import_structure = {
         "RBLNUNet2DConditionModel",
         "RBLNControlNetModel",
         "RBLNSD3Transformer2DModel",
+        "RBLNPriorTransformer",
+        "RBLNVQModel",
     ],
     "modeling_diffusers": [
         "RBLNDiffusionMixin",
@@ -55,10 +60,15 @@ if TYPE_CHECKING:
     from .models import (
         RBLNAutoencoderKL,
         RBLNControlNetModel,
+        RBLNPriorTransformer,
         RBLNSD3Transformer2DModel,
         RBLNUNet2DConditionModel,
+        RBLNVQModel,
     )
     from .pipelines import (
+        RBLNKandinskyV22InpaintCombinedPipeline,
+        RBLNKandinskyV22InpaintPipeline,
+        RBLNKandinskyV22PriorPipeline,
         RBLNMultiControlNetModel,
         RBLNStableDiffusion3Img2ImgPipeline,
         RBLNStableDiffusion3InpaintPipeline,

optimum/rbln/diffusers/modeling_diffusers.py

@@ -23,6 +23,7 @@ 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 +68,7 @@ class RBLNDiffusionMixin:
     """
 
     _submodules = []
+    _prefix = {}
 
     @classmethod
     @property
@@ -84,25 +86,50 @@ class RBLNDiffusionMixin:
     ) -> Dict[str, Any]:
         submodule = getattr(model, submodule_name)
         submodule_class_name = submodule.__class__.__name__
+        if isinstance(submodule, torch.nn.Module):
+            if submodule_class_name == "MultiControlNetModel":
+                submodule_class_name = "ControlNetModel"
 
-        if submodule_class_name == "MultiControlNetModel":
-            submodule_class_name = "ControlNetModel"
+            submodule_cls: RBLNModel = getattr(importlib.import_module("optimum.rbln"), f"RBLN{submodule_class_name}")
 
-        submodule_cls: RBLNModel = getattr(importlib.import_module("optimum.rbln"), f"RBLN{submodule_class_name}")
+            submodule_config = rbln_config.get(submodule_name, {})
+            submodule_config = copy.deepcopy(submodule_config)
 
-        submodule_config = rbln_config.get(submodule_name, {})
-        submodule_config = copy.deepcopy(submodule_config)
+            pipe_global_config = {k: v for k, v in rbln_config.items() if k not in cls._submodules}
 
-        pipe_global_config = {k: v for k, v in rbln_config.items() if k not in cls._submodules}
+            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,
+                }
+            )
+            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
+            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
 
-        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,
-            }
-        )
-        submodule_config = submodule_cls.update_rbln_config_using_pipe(model, submodule_config)
+            submodules = copy.deepcopy(cls._submodules)
+            submodules += [prefix + connected_submodule_name for connected_submodule_name in connected_submodules]
+
+            pipe_global_config = {k: v for k, v in rbln_config.items() if k not in submodules}
+            for connected_submodule_name in connected_submodules:
+                submodule_config[connected_submodule_name].update(
+                    {k: v for k, v in pipe_global_config.items() if k not in submodule_config}
+                )
+            rbln_config[submodule_name] = submodule_config
+        else:
+            raise ValueError(f"submodule {submodule_name} isn't supported")
         return submodule_config
 
     @staticmethod
@@ -165,8 +192,26 @@ class RBLNDiffusionMixin:
 
         else:
             # raise error if any of submodules are torch module.
-            model_index_config = None
-            for submodule_name in cls._submodules:
+            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:
                 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`."
@@ -181,9 +226,6 @@ class RBLNDiffusionMixin:
                 if not any(kwd in submodule_config for kwd in RUNTIME_KEYWORDS):
                     continue
 
-                if model_index_config is None:
-                    model_index_config = cls.load_config(pretrained_model_name_or_path=model_id)
-
                 module_name, class_name = model_index_config[submodule_name]
                 if module_name != "optimum.rbln":
                     raise ValueError(
@@ -228,6 +270,7 @@ class RBLNDiffusionMixin:
         passed_submodules: Dict[str, RBLNModel],
         model_save_dir: Optional[PathLike],
         rbln_config: Dict[str, Any],
+        prefix: Optional[str] = "",
     ) -> Dict[str, RBLNModel]:
         compiled_submodules = {}
 
@@ -245,17 +288,54 @@ class RBLNDiffusionMixin:
                     controlnets=submodule,
                     model_save_dir=model_save_dir,
                     controlnet_rbln_config=submodule_rbln_config,
+                    prefix=prefix,
                 )
             elif isinstance(submodule, torch.nn.Module):
                 submodule_cls: RBLNModel = getattr(
                     importlib.import_module("optimum.rbln"), f"RBLN{submodule.__class__.__name__}"
                 )
+                subfolder = prefix + submodule_name
                 submodule = submodule_cls.from_model(
                     model=submodule,
-                    subfolder=submodule_name,
+                    subfolder=subfolder,
                     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__} ")
 
@@ -268,6 +348,7 @@ class RBLNDiffusionMixin:
         controlnets: "MultiControlNetModel",
         model_save_dir: Optional[PathLike],
         controlnet_rbln_config: Dict[str, Any],
+        prefix: Optional[str] = "",
     ):
         # Compile multiple ControlNet models for a MultiControlNet setup
         from .models.controlnet import RBLNControlNetModel
@@ -276,7 +357,7 @@ class RBLNDiffusionMixin:
         compiled_controlnets = [
             RBLNControlNetModel.from_model(
                 model=controlnet,
-                subfolder="controlnet" if i == 0 else f"controlnet_{i}",
+                subfolder=f"{prefix}controlnet" if i == 0 else f"{prefix}controlnet_{i}",
                 model_save_dir=model_save_dir,
                 rbln_config=controlnet_rbln_config,
             )
@@ -287,10 +368,21 @@ 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 cls._submodules:
+            for submodule_name in submodule_names:
                 delattr(model, submodule_name)
 
             # Direct calling of `save_pretrained` causes config.unet = (None, None).
@@ -300,7 +392,7 @@ class RBLNDiffusionMixin:
             # Causing warning messeages.
 
         update_dict = {}
-        for submodule_name in cls._submodules:
+        for submodule_name in submodule_names:
             # replace submodule
             setattr(model, submodule_name, submodules[submodule_name])
             update_dict[submodule_name] = ("optimum.rbln", submodules[submodule_name].__class__.__name__)

optimum/rbln/diffusers/models/__init__.py

@@ -20,20 +20,26 @@ from transformers.utils import _LazyModule
 _import_structure = {
     "autoencoders": [
         "RBLNAutoencoderKL",
+        "RBLNVQModel",
     ],
     "unets": [
         "RBLNUNet2DConditionModel",
     ],
     "controlnet": ["RBLNControlNetModel"],
-    "transformers": ["RBLNSD3Transformer2DModel"],
+    "transformers": [
+        "RBLNPriorTransformer",
+        "RBLNSD3Transformer2DModel",
+    ],
 }
 
 if TYPE_CHECKING:
     from .autoencoders import (
         RBLNAutoencoderKL,
+        RBLNVQModel,
     )
     from .controlnet import RBLNControlNetModel
     from .transformers import (
+        RBLNPriorTransformer,
         RBLNSD3Transformer2DModel,
     )
     from .unets import (

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

@@ -13,3 +13,4 @@
 # limitations under the License.
 
 from .autoencoder_kl import RBLNAutoencoderKL
+from .vq_model import RBLNVQModel

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

@@ -12,11 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, List
 
 import torch  # noqa: I001
-from diffusers import AutoencoderKL
+from diffusers import AutoencoderKL, VQModel
 from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution
+from diffusers.models.autoencoders.vq_model import VQEncoderOutput
 from diffusers.models.modeling_outputs import AutoencoderKLOutput
 
 from ....utils.logging import get_logger
@@ -72,3 +73,52 @@ class _VAEEncoder(torch.nn.Module):
     def forward(self, x):
         vae_out = _VAEEncoder.encode(self.vae, x, return_dict=False)
         return vae_out
+
+
+class RBLNRuntimeVQEncoder(RBLNPytorchRuntime):
+    def encode(self, x: torch.FloatTensor, **kwargs) -> torch.FloatTensor:
+        h = self.forward(x.contiguous())
+        return VQEncoderOutput(latents=h)
+
+
+class RBLNRuntimeVQDecoder(RBLNPytorchRuntime):
+    def decode(self, h: torch.Tensor, force_not_quantize: bool = False, shape=None, **kwargs) -> List[torch.Tensor]:
+        if not (force_not_quantize and not self.lookup_from_codebook):
+            raise ValueError(
+                "Currently, the `decode` method of the class `RBLNVQModel` is executed successfully only if `force_not_quantize` is True and `config.lookup_from_codebook` is False"
+            )
+        commit_loss = torch.zeros((h.shape[0])).to(h.device, dtype=h.dtype)
+        dec = self.forward(h.contiguous())
+        return dec, commit_loss
+
+
+class _VQEncoder(torch.nn.Module):
+    def __init__(self, vq_model: VQModel):
+        super().__init__()
+        self.vq_model = vq_model
+
+    def encode(self, x: torch.Tensor, return_dict: bool = True):
+        h = self.vq_model.encoder(x)
+        h = self.vq_model.quant_conv(h)
+        return h
+
+    def forward(self, x: torch.Tensor):
+        vq_out = self.encode(x)
+        return vq_out
+
+
+class _VQDecoder(torch.nn.Module):
+    def __init__(self, vq_model: VQModel):
+        super().__init__()
+        self.vq_model = vq_model
+
+    def decode(self, h: torch.Tensor, force_not_quantize: bool = False, return_dict: bool = True, shape=None):
+        quant = h
+        quant2 = self.vq_model.post_quant_conv(quant)
+        quant = quant if self.vq_model.config.norm_type == "spatial" else None
+        dec = self.vq_model.decoder(quant2, quant)
+        return dec
+
+    def forward(self, h: torch.Tensor):
+        vq_out = self.decode(h)
+        return vq_out

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

@@ -0,0 +1,159 @@
+# 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 typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+
+import rebel
+import torch
+from diffusers import VQModel
+from diffusers.models.autoencoders.vae import DecoderOutput
+from diffusers.models.autoencoders.vq_model import VQEncoderOutput
+from transformers import PretrainedConfig
+
+from ....modeling import RBLNModel
+from ....modeling_config import DEFAULT_COMPILED_MODEL_NAME, RBLNCompileConfig, RBLNConfig
+from ....utils.logging import get_logger
+from ...modeling_diffusers import RBLNDiffusionMixin
+from .vae import RBLNRuntimeVQDecoder, RBLNRuntimeVQEncoder, _VQDecoder, _VQEncoder
+
+
+if TYPE_CHECKING:
+    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer
+
+logger = get_logger(__name__)
+
+
+class RBLNVQModel(RBLNModel):
+    auto_model_class = VQModel
+    config_name = "config.json"
+    hf_library_name = "diffusers"
+
+    def __post_init__(self, **kwargs):
+        super().__post_init__(**kwargs)
+
+        self.encoder = RBLNRuntimeVQEncoder(runtime=self.model[0], main_input_name="x")
+        self.decoder = RBLNRuntimeVQDecoder(runtime=self.model[1], main_input_name="z")
+        self.decoder.lookup_from_codebook = self.config.lookup_from_codebook
+        height = self.rbln_config.model_cfg.get("img_height", 512)
+        width = self.rbln_config.model_cfg.get("img_width", 512)
+        self.image_size = [height, width]
+
+    @classmethod
+    def get_compiled_model(cls, model, rbln_config: RBLNConfig):
+        encoder_model = _VQEncoder(model)
+        decoder_model = _VQDecoder(model)
+        encoder_model.eval()
+        decoder_model.eval()
+
+        enc_compiled_model = cls.compile(encoder_model, rbln_compile_config=rbln_config.compile_cfgs[0])
+        dec_compiled_model = cls.compile(decoder_model, rbln_compile_config=rbln_config.compile_cfgs[1])
+
+        return {"encoder": enc_compiled_model, "decoder": dec_compiled_model}
+
+    @classmethod
+    def update_rbln_config_using_pipe(cls, pipe: RBLNDiffusionMixin, rbln_config: Dict[str, Any]) -> Dict[str, Any]:
+        batch_size = rbln_config.get("batch_size")
+        if batch_size is None:
+            batch_size = 1
+        img_height = rbln_config.get("img_height")
+        if img_height is None:
+            img_height = 512
+        img_width = rbln_config.get("img_width")
+        if img_width is None:
+            img_width = 512
+
+        rbln_config.update(
+            {
+                "batch_size": batch_size,
+                "img_height": img_height,
+                "img_width": img_width,
+            }
+        )
+
+        return rbln_config
+
+    @classmethod
+    def _get_rbln_config(
+        cls,
+        preprocessors: Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"],
+        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
+
+        if hasattr(model_config, "block_out_channels"):
+            scale_factor = 2 ** (len(model_config.block_out_channels) - 1)
+        else:
+            # image processor default value 8 (int)
+            scale_factor = 8
+
+        enc_shape = (height, width)
+        dec_shape = (height // scale_factor, width // scale_factor)
+
+        enc_input_info = [
+            (
+                "x",
+                [batch_size, model_config.in_channels, enc_shape[0], enc_shape[1]],
+                "float32",
+            )
+        ]
+        dec_input_info = [
+            (
+                "h",
+                [batch_size, model_config.latent_channels, dec_shape[0], dec_shape[1]],
+                "float32",
+            )
+        ]
+
+        enc_rbln_compile_config = RBLNCompileConfig(compiled_model_name="encoder", input_info=enc_input_info)
+        dec_rbln_compile_config = RBLNCompileConfig(compiled_model_name="decoder", input_info=dec_input_info)
+
+        compile_cfgs = [enc_rbln_compile_config, dec_rbln_compile_config]
+        rbln_config = RBLNConfig(
+            rbln_cls=cls.__name__,
+            compile_cfgs=compile_cfgs,
+            rbln_kwargs=rbln_kwargs,
+        )
+        return rbln_config
+
+    @classmethod
+    def _create_runtimes(
+        cls,
+        compiled_models: List[rebel.RBLNCompiledModel],
+        rbln_device_map: Dict[str, int],
+        activate_profiler: Optional[bool] = None,
+    ) -> List[rebel.Runtime]:
+        if len(compiled_models) == 1:
+            device_val = rbln_device_map[DEFAULT_COMPILED_MODEL_NAME]
+            return [
+                compiled_models[0].create_runtime(
+                    tensor_type="pt", device=device_val, activate_profiler=activate_profiler
+                )
+            ]
+
+        device_vals = [rbln_device_map["encoder"], rbln_device_map["decoder"]]
+        return [
+            compiled_model.create_runtime(tensor_type="pt", device=device_val, activate_profiler=activate_profiler)
+            for compiled_model, device_val in zip(compiled_models, device_vals)
+        ]
+
+    def encode(self, x: torch.FloatTensor, **kwargs) -> torch.FloatTensor:
+        posterior = self.encoder.encode(x)
+        return VQEncoderOutput(latents=posterior)
+
+    def decode(self, h: torch.FloatTensor, **kwargs) -> torch.FloatTensor:
+        dec, commit_loss = self.decoder.decode(h, **kwargs)
+        return DecoderOutput(sample=dec, commit_loss=commit_loss)

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

@@ -12,4 +12,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from .prior_transformer import RBLNPriorTransformer
 from .transformer_sd3 import RBLNSD3Transformer2DModel

optimum/rbln/diffusers/models/transformers/prior_transformer.py

@@ -0,0 +1,174 @@
+# 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 pathlib import Path
+from typing import Any, Dict, Optional, Union
+
+import torch
+from diffusers.models.transformers.prior_transformer import PriorTransformer, PriorTransformerOutput
+from transformers import PretrainedConfig, PreTrainedModel
+
+from ....modeling import RBLNModel
+from ....modeling_config import RBLNCompileConfig, RBLNConfig
+from ....utils.logging import get_logger
+from ....utils.runtime_utils import RBLNPytorchRuntime
+from ...modeling_diffusers import RBLNDiffusionMixin
+
+
+logger = get_logger(__name__)
+
+
+class RBLNRuntimePriorTransformer(RBLNPytorchRuntime):
+    def forward(
+        self, hidden_states, timestep, proj_embedding, encoder_hidden_states, attention_mask, return_dict: bool = True
+    ):
+        predicted_image_embedding = super().forward(
+            hidden_states,
+            timestep,
+            proj_embedding,
+            encoder_hidden_states,
+            attention_mask,
+        )
+        if return_dict:
+            return PriorTransformerOutput(predicted_image_embedding=predicted_image_embedding)
+        else:
+            return (predicted_image_embedding,)
+
+
+class _PriorTransformer(torch.nn.Module):
+    def __init__(self, prior: PriorTransformer):
+        super().__init__()
+        self._prior = prior
+
+    def forward(
+        self,
+        hidden_states,
+        timestep,
+        proj_embedding,
+        encoder_hidden_states,
+        attention_mask,
+        return_dict=True,
+    ):
+        return self._prior.forward(
+            hidden_states,
+            timestep,
+            proj_embedding,
+            encoder_hidden_states,
+            attention_mask,
+            return_dict=False,
+        )
+
+
+class RBLNPriorTransformer(RBLNModel):
+    hf_library_name = "diffusers"
+    auto_model_class = PriorTransformer
+
+    def __post_init__(self, **kwargs):
+        super().__post_init__(**kwargs)
+        self.runtime = RBLNRuntimePriorTransformer(runtime=self.model[0])
+        artifacts = torch.load(self.model_save_dir / self.subfolder / "torch_artifacts.pth", weights_only=False)
+        self.clip_mean = artifacts["clip_mean"]
+        self.clip_std = artifacts["clip_std"]
+
+    @classmethod
+    def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNConfig) -> torch.nn.Module:
+        return _PriorTransformer(model).eval()
+
+    @classmethod
+    def update_rbln_config_using_pipe(cls, pipe: RBLNDiffusionMixin, rbln_config: Dict[str, Any]) -> Dict[str, Any]:
+        batch_size = rbln_config.get("batch_size")
+        if not batch_size:
+            do_classifier_free_guidance = rbln_config.get("guidance_scale", 5.0) > 1.0
+            batch_size = 2 if do_classifier_free_guidance else 1
+        else:
+            if rbln_config.get("guidance_scale"):
+                logger.warning(
+                    "guidance_scale is ignored because batch size is explicitly specified. "
+                    "To ensure consistent behavior, consider removing the guidance scale or "
+                    "adjusting the batch size configuration as needed."
+                )
+        embedding_dim = rbln_config.get("embedding_dim", pipe.prior.config.embedding_dim)
+        num_embeddings = rbln_config.get("num_embeddings", pipe.prior.config.num_embeddings)
+
+        rbln_config.update(
+            {
+                "batch_size": batch_size,
+                "embedding_dim": embedding_dim,
+                "num_embeddings": num_embeddings,
+            }
+        )
+
+        return rbln_config
+
+    @classmethod
+    def save_torch_artifacts(
+        cls,
+        model: "PreTrainedModel",
+        save_dir_path: Path,
+        subfolder: str,
+        rbln_config: RBLNConfig,
+    ):
+        save_dict = {}
+        save_dict["clip_mean"] = model.clip_mean
+        save_dict["clip_std"] = model.clip_std
+        torch.save(save_dict, save_dir_path / subfolder / "torch_artifacts.pth")
+
+    @classmethod
+    def _get_rbln_config(
+        cls,
+        preprocessors,
+        model_config: PretrainedConfig,
+        rbln_kwargs,
+    ) -> RBLNConfig:
+        batch_size = rbln_kwargs.get("batch_size") or 1
+        embedding_dim = rbln_kwargs.get("embedding_dim") or model_config.embedding_dim
+        num_embeddings = rbln_kwargs.get("num_embeddings") or model_config.num_embeddings
+
+        input_info = [
+            ("hidden_states", [batch_size, embedding_dim], "float32"),
+            ("timestep", [], "float32"),
+            ("proj_embedding", [batch_size, embedding_dim], "float32"),
+            ("encoder_hidden_states", [batch_size, num_embeddings, embedding_dim], "float32"),
+            ("attention_mask", [batch_size, num_embeddings], "float32"),
+        ]
+
+        rbln_compile_config = RBLNCompileConfig(input_info=input_info)
+        rbln_config = RBLNConfig(
+            rbln_cls=cls.__name__,
+            compile_cfgs=[rbln_compile_config],
+            rbln_kwargs=rbln_kwargs,
+        )
+        return rbln_config
+
+    def forward(
+        self,
+        hidden_states,
+        timestep: Union[torch.Tensor, float, int],
+        proj_embedding: torch.Tensor,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.BoolTensor] = None,
+        return_dict: bool = True,
+    ):
+        return self.runtime.forward(
+            hidden_states.contiguous(),
+            timestep.float(),
+            proj_embedding,
+            encoder_hidden_states,
+            attention_mask.float(),
+            return_dict,
+        )
+
+    def post_process_latents(self, prior_latents):
+        prior_latents = (prior_latents * self.clip_std) + self.clip_mean
+        return prior_latents

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

@@ -115,6 +115,29 @@ class _UNet_SDXL(torch.nn.Module):
         return unet_out
 
 
+class _UNet_Kandinsky(torch.nn.Module):
+    def __init__(self, unet: "UNet2DConditionModel"):
+        super().__init__()
+        self.unet = unet
+
+    def forward(
+        self,
+        sample: torch.Tensor,
+        timestep: Union[torch.Tensor, float, int],
+        image_embeds: torch.Tensor,
+    ) -> torch.Tensor:
+        added_cond_kwargs = {"image_embeds": image_embeds}
+
+        unet_out = self.unet(
+            sample=sample,
+            timestep=timestep,
+            encoder_hidden_states=None,
+            added_cond_kwargs=added_cond_kwargs,
+            return_dict=False,
+        )
+        return unet_out
+
+
 class RBLNUNet2DConditionModel(RBLNModel):
     hf_library_name = "diffusers"
     auto_model_class = UNet2DConditionModel
@@ -138,6 +161,8 @@ class RBLNUNet2DConditionModel(RBLNModel):
     def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNConfig) -> torch.nn.Module:
         if model.config.addition_embed_type == "text_time":
             return _UNet_SDXL(model).eval()
+        elif model.config.addition_embed_type == "image":
+            return _UNet_Kandinsky(model).eval()
         else:
             return _UNet_SD(model).eval()
 
@@ -146,6 +171,7 @@ class RBLNUNet2DConditionModel(RBLNModel):
         cls, pipe: RBLNDiffusionMixin, rbln_config: Dict[str, Any]
     ) -> Union[int, Tuple[int, int]]:
         image_size = (rbln_config.get("img_height"), rbln_config.get("img_width"))
+        scale_factor = pipe.movq_scale_factor if hasattr(pipe, "movq_scale_factor") else pipe.vae_scale_factor
         if (image_size[0] is None) != (image_size[1] is None):
             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:
@@ -153,22 +179,23 @@ class RBLNUNet2DConditionModel(RBLNModel):
                 # 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 // pipe.vae_scale_factor
+                    sample_size = vae_sample_size // scale_factor
                 else:
                     sample_size = (
-                        vae_sample_size[0] // pipe.vae_scale_factor,
-                        vae_sample_size[1] // pipe.vae_scale_factor,
+                        vae_sample_size[0] // scale_factor,
+                        vae_sample_size[1] // scale_factor,
                     )
             else:
                 sample_size = pipe.unet.config.sample_size
         else:
-            sample_size = (image_size[0] // pipe.vae_scale_factor, image_size[1] // pipe.vae_scale_factor)
+            sample_size = (image_size[0] // scale_factor, image_size[1] // scale_factor)
 
         return sample_size
 
     @classmethod
     def update_rbln_config_using_pipe(cls, pipe: RBLNDiffusionMixin, rbln_config: Dict[str, Any]) -> Dict[str, Any]:
         text_model_hidden_size = pipe.text_encoder_2.config.hidden_size if hasattr(pipe, "text_encoder_2") else None
+        image_model_hidden_size = pipe.unet.config.encoder_hid_dim if hasattr(pipe, "unet") else None
 
         batch_size = rbln_config.get("batch_size")
         if not batch_size:
@@ -184,10 +211,12 @@ class RBLNUNet2DConditionModel(RBLNModel):
                     "adjusting the batch size configuration as needed."
                 )
 
+        max_seq_len = pipe.text_encoder.config.max_position_embeddings if hasattr(pipe, "text_encoder") else None
         rbln_config.update(
             {
-                "max_seq_len": pipe.text_encoder.config.max_position_embeddings,
+                "max_seq_len": max_seq_len,
                 "text_model_hidden_size": text_model_hidden_size,
+                "image_model_hidden_size": image_model_hidden_size,
                 "sample_size": cls.get_unet_sample_size(pipe, rbln_config),
                 "batch_size": batch_size,
                 "is_controlnet": "controlnet" in pipe.config.keys(),
@@ -218,15 +247,16 @@ class RBLNUNet2DConditionModel(RBLNModel):
         if isinstance(sample_size, int):
             sample_size = (sample_size, sample_size)
 
-        if max_seq_len is None:
-            raise ValueError("`rbln_max_seq_len` (ex. text_encoder's max_position_embeddings) must be specified.")
-
         input_info = [
             ("sample", [batch_size, model_config.in_channels, sample_size[0], sample_size[1]], "float32"),
             ("timestep", [], "float32"),
-            ("encoder_hidden_states", [batch_size, max_seq_len, model_config.cross_attention_dim], "float32"),
         ]
 
+        if max_seq_len is not None:
+            input_info.append(
+                ("encoder_hidden_states", [batch_size, max_seq_len, model_config.cross_attention_dim], "float32"),
+            )
+
         if is_controlnet:
             # down block addtional residuals
             first_shape = [batch_size, model_config.block_out_channels[0], sample_size[0], sample_size[1]]
@@ -256,11 +286,15 @@ class RBLNUNet2DConditionModel(RBLNModel):
             ]
             input_info.append(("mid_block_additional_residual", shape, "float32"))
 
-        if hasattr(model_config, "addition_embed_type") and model_config.addition_embed_type == "text_time":
-            rbln_text_model_hidden_size = rbln_kwargs["text_model_hidden_size"]
-            rbln_in_features = model_config.projection_class_embeddings_input_dim
-            input_info.append(("text_embeds", [batch_size, rbln_text_model_hidden_size], "float32"))
-            input_info.append(("time_ids", [batch_size, 6], "float32"))
+        if hasattr(model_config, "addition_embed_type"):
+            if model_config.addition_embed_type == "text_time":
+                rbln_text_model_hidden_size = rbln_kwargs["text_model_hidden_size"]
+                rbln_in_features = model_config.projection_class_embeddings_input_dim
+                input_info.append(("text_embeds", [batch_size, rbln_text_model_hidden_size], "float32"))
+                input_info.append(("time_ids", [batch_size, 6], "float32"))
+            elif model_config.addition_embed_type == "image":
+                rbln_image_model_hidden_size = rbln_kwargs["image_model_hidden_size"]
+                input_info.append(("image_embeds", [batch_size, rbln_image_model_hidden_size], "float32"))
 
         rbln_compile_config = RBLNCompileConfig(input_info=input_info)
 
@@ -323,6 +357,15 @@ class RBLNUNet2DConditionModel(RBLNModel):
                 ),
             )
 
+        if "image_embeds" in added_cond_kwargs:
+            return (
+                super().forward(
+                    sample.contiguous(),
+                    timestep.float(),
+                    **added_cond_kwargs,
+                ),
+            )
+
         return (
             super().forward(
                 sample.contiguous(),

optimum/rbln/diffusers/pipelines/__init__.py

@@ -25,6 +25,11 @@ _import_structure = {
         "RBLNStableDiffusionXLControlNetImg2ImgPipeline",
         "RBLNStableDiffusionXLControlNetPipeline",
     ],
+    "kandinsky2_2": [
+        "RBLNKandinskyV22InpaintCombinedPipeline",
+        "RBLNKandinskyV22InpaintPipeline",
+        "RBLNKandinskyV22PriorPipeline",
+    ],
     "stable_diffusion": [
         "RBLNStableDiffusionImg2ImgPipeline",
         "RBLNStableDiffusionPipeline",
@@ -49,6 +54,11 @@ if TYPE_CHECKING:
         RBLNStableDiffusionXLControlNetImg2ImgPipeline,
         RBLNStableDiffusionXLControlNetPipeline,
     )
+    from .kandinsky2_2 import (
+        RBLNKandinskyV22InpaintCombinedPipeline,
+        RBLNKandinskyV22InpaintPipeline,
+        RBLNKandinskyV22PriorPipeline,
+    )
     from .stable_diffusion import (
         RBLNStableDiffusionImg2ImgPipeline,
         RBLNStableDiffusionInpaintPipeline,

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

@@ -0,0 +1,17 @@
+# 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 .pipeline_kandinsky2_2_combined import RBLNKandinskyV22InpaintCombinedPipeline
+from .pipeline_kandinsky2_2_inpaint import RBLNKandinskyV22InpaintPipeline
+from .pipeline_kandinsky2_2_prior import RBLNKandinskyV22PriorPipeline

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

@@ -0,0 +1,83 @@
+# 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 (
+    DDPMScheduler,
+    KandinskyV22InpaintCombinedPipeline,
+    PriorTransformer,
+    UnCLIPScheduler,
+    UNet2DConditionModel,
+    VQModel,
+)
+from transformers import (
+    CLIPImageProcessor,
+    CLIPTextModelWithProjection,
+    CLIPTokenizer,
+    CLIPVisionModelWithProjection,
+)
+
+from ...modeling_diffusers import RBLNDiffusionMixin
+from .pipeline_kandinsky2_2_inpaint import RBLNKandinskyV22InpaintPipeline
+from .pipeline_kandinsky2_2_prior import RBLNKandinskyV22PriorPipeline
+
+
+class RBLNKandinskyV22InpaintCombinedPipeline(RBLNDiffusionMixin, KandinskyV22InpaintCombinedPipeline):
+    original_class = KandinskyV22InpaintCombinedPipeline
+    _connected_classes = {"prior_pipe": RBLNKandinskyV22PriorPipeline, "decoder_pipe": RBLNKandinskyV22InpaintPipeline}
+    _submodules = ["prior_pipe", "decoder_pipe"]
+    _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(KandinskyV22InpaintCombinedPipeline, 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 = RBLNKandinskyV22InpaintPipeline(
+            unet=unet,
+            scheduler=scheduler,
+            movq=movq,
+        )
+
+    def get_compiled_image_size(self):
+        return self.movq.image_size

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

@@ -0,0 +1,22 @@
+# 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 KandinskyV22InpaintPipeline
+
+from ...modeling_diffusers import RBLNDiffusionMixin
+
+
+class RBLNKandinskyV22InpaintPipeline(RBLNDiffusionMixin, KandinskyV22InpaintPipeline):
+    original_class = KandinskyV22InpaintPipeline
+    _submodules = ["unet", "movq"]

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

@@ -0,0 +1,22 @@
+# 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 KandinskyV22PriorPipeline
+
+from ...modeling_diffusers import RBLNDiffusionMixin
+
+
+class RBLNKandinskyV22PriorPipeline(RBLNDiffusionMixin, KandinskyV22PriorPipeline):
+    original_class = KandinskyV22PriorPipeline
+    _submodules = ["text_encoder", "image_encoder", "prior"]

optimum/rbln/transformers/__init__.py

@@ -40,6 +40,7 @@ _import_structure = {
         "RBLNCLIPTextModel",
         "RBLNCLIPTextModelWithProjection",
         "RBLNCLIPVisionModel",
+        "RBLNCLIPVisionModelWithProjection",
         "RBLNDPTForDepthEstimation",
         "RBLNExaoneForCausalLM",
         "RBLNGemmaForCausalLM",
@@ -99,6 +100,7 @@ if TYPE_CHECKING:
         RBLNCLIPTextModel,
         RBLNCLIPTextModelWithProjection,
         RBLNCLIPVisionModel,
+        RBLNCLIPVisionModelWithProjection,
         RBLNDPTForDepthEstimation,
         RBLNExaoneForCausalLM,
         RBLNGemmaForCausalLM,

optimum/rbln/transformers/models/__init__.py

@@ -34,7 +34,12 @@ _import_structure = {
     ],
     "bart": ["RBLNBartForConditionalGeneration", "RBLNBartModel"],
     "bert": ["RBLNBertModel", "RBLNBertForQuestionAnswering", "RBLNBertForMaskedLM"],
-    "clip": ["RBLNCLIPTextModel", "RBLNCLIPTextModelWithProjection", "RBLNCLIPVisionModel"],
+    "clip": [
+        "RBLNCLIPTextModel",
+        "RBLNCLIPTextModelWithProjection",
+        "RBLNCLIPVisionModel",
+        "RBLNCLIPVisionModelWithProjection",
+    ],
     "dpt": ["RBLNDPTForDepthEstimation"],
     "exaone": ["RBLNExaoneForCausalLM"],
     "gemma": ["RBLNGemmaForCausalLM"],
@@ -68,7 +73,12 @@ if TYPE_CHECKING:
     )
     from .bart import RBLNBartForConditionalGeneration, RBLNBartModel
     from .bert import RBLNBertForMaskedLM, RBLNBertForQuestionAnswering, RBLNBertModel
-    from .clip import RBLNCLIPTextModel, RBLNCLIPTextModelWithProjection, RBLNCLIPVisionModel
+    from .clip import (
+        RBLNCLIPTextModel,
+        RBLNCLIPTextModelWithProjection,
+        RBLNCLIPVisionModel,
+        RBLNCLIPVisionModelWithProjection,
+    )
     from .dpt import RBLNDPTForDepthEstimation
     from .exaone import RBLNExaoneForCausalLM
     from .gemma import RBLNGemmaForCausalLM

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

@@ -12,4 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .modeling_clip import RBLNCLIPTextModel, RBLNCLIPTextModelWithProjection, RBLNCLIPVisionModel
+from .modeling_clip import (
+    RBLNCLIPTextModel,
+    RBLNCLIPTextModelWithProjection,
+    RBLNCLIPVisionModel,
+    RBLNCLIPVisionModelWithProjection,
+)

optimum/rbln/transformers/models/clip/modeling_clip.py

@@ -22,7 +22,7 @@ from transformers import (
     CLIPVisionModel,
 )
 from transformers.modeling_outputs import BaseModelOutputWithPooling
-from transformers.models.clip.modeling_clip import CLIPTextModelOutput
+from transformers.models.clip.modeling_clip import CLIPTextModelOutput, CLIPVisionModelOutput
 
 from ....diffusers.modeling_diffusers import RBLNDiffusionMixin
 from ....modeling import RBLNModel
@@ -116,6 +116,10 @@ class RBLNCLIPVisionModel(RBLNModel):
     def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNConfig) -> torch.nn.Module:
         return _VisionEncoder(model).eval()
 
+    @classmethod
+    def update_rbln_config_using_pipe(cls, pipe: RBLNDiffusionMixin, rbln_config: Dict[str, Any]) -> Dict[str, Any]:
+        return rbln_config
+
     @classmethod
     def _get_rbln_config(
         cls,
@@ -179,3 +183,24 @@ class RBLNCLIPVisionModel(RBLNModel):
             pooler_output=output[1],
             hidden_states=output[2:],
         )
+
+
+class RBLNCLIPVisionModelWithProjection(RBLNCLIPVisionModel):
+    def forward(
+        self,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        **kwargs,
+    ) -> Union[Tuple, CLIPVisionModelOutput]:
+        if len(kwargs) > 0 and any(kwargs.values()):
+            logger.warning(f"Currently, optimum-rbln does not support kwargs {kwargs.keys()} for {self.__class__}.")
+
+        output = super().forward(pixel_values)
+        image_embeds = output[0]
+        last_hidden_state = output[1]
+        hidden_states = output[2:]
+
+        return CLIPVisionModelOutput(
+            image_embeds=image_embeds,
+            last_hidden_state=last_hidden_state,
+            hidden_states=hidden_states,
+        )

{optimum_rbln-0.2.1a5.dist-info → optimum_rbln-0.7.2rc0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: optimum-rbln
-Version: 0.2.1a5
+Version: 0.7.2rc0
 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.2.1a5.dist-info → optimum_rbln-0.7.2rc0.dist-info}/RECORD

@@ -1,26 +1,32 @@
-optimum/rbln/__init__.py,sha256=sLCjJu_MLZEKDOwHIlJP4u4GzGZx-1kqHTYGw5B4xDg,6096
-optimum/rbln/__version__.py,sha256=J4Eyn4HLzB0UpyosVo-P3LCDkB5knEOS6Nu24mnl5NA,413
+optimum/rbln/__init__.py,sha256=eHi15YM3989AcX52jka9rUmgAtlp1PHqMNwBEdOfuu8,6554
+optimum/rbln/__version__.py,sha256=OIstx-UVfreSh_0960pBWlg0VhVw0Sy7VpRV6nHTpqA,414
 optimum/rbln/modeling.py,sha256=REImAAKO82CqSNABR-9E1jJEsWch9amSOwOOQhFEYLY,8283
 optimum/rbln/modeling_base.py,sha256=fQ0bI1Bb6GJquRXftmSSN9K-TXLhFltZJ6C-2w43xMg,21193
 optimum/rbln/modeling_config.py,sha256=7104bxmrvKW4Q6XTruQayiIGl8GHDFmPkJ3cknMIInE,11335
-optimum/rbln/diffusers/__init__.py,sha256=68FTAMpbbMflm8qiSqfM5J2_gFb3iU3fng6AL0TG47A,2913
-optimum/rbln/diffusers/modeling_diffusers.py,sha256=E1x-iOKEJCUB6ml0RgtFEVPPk6J6pqEF-JTEyOZzOyc,14928
-optimum/rbln/diffusers/models/__init__.py,sha256=aSL5_yd-y8Q6DxNvfQ-yl-BUNyMzI1P6AikjQMKZzpI,1357
+optimum/rbln/diffusers/__init__.py,sha256=pOyoXv3-JRzTBSwPKbgLS9H6F2K9dJdReEmpGhcLQYU,3283
+optimum/rbln/diffusers/modeling_diffusers.py,sha256=t7kwC-NEw-Y01EK8Xi4bqlsUOfabkNrLkmEHfRCHKG0,20320
+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=nMfnwEwuOje-qKofAw-uOWUWcYV_YvnaN68IGfDdqHg,645
+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=A-F2TRJ2vL4gNXiMT_hRGeanIFKWxJ1QaKmYVp41rwI,2513
-optimum/rbln/diffusers/models/transformers/__init__.py,sha256=TEhARgQJx_NUZzI6M8gt3aWbdzmLHnM6FMSQd9M9zCk,654
+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/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=Z0-eAZw1Gah24y6uOO5m9-GRruBppCSdV2NQZLNtBaI,14021
-optimum/rbln/diffusers/pipelines/__init__.py,sha256=i8AQJSoV9clLTill7wP5ECci6E7lC2gBaNuqfhYklZk,2469
+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/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/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
 optimum/rbln/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py,sha256=qU7wN2gFUjFImuB6CGDY2SC4aZw4VhaRKu92eA_Fa08,904
@@ -37,11 +43,11 @@ optimum/rbln/ops/__init__.py,sha256=MbAHevg59fXQOFgrnsFFZ97s8-YrgvYCYML_sgKEEfM,
 optimum/rbln/ops/attn.py,sha256=QYvSMg4ps_PenHwpLVhuYRoOAFvHIo19nY0ZEdj4nTE,9700
 optimum/rbln/ops/flash_attn.py,sha256=Zn5nkouY3kk6MBivQpPjgGh4oepjpi8F3tnTrmrNfpg,2304
 optimum/rbln/ops/kv_cache_update.py,sha256=9W4WCO1Dtfy0u5i978JJRa7uLbqrfR2lHuoPynb07fw,3143
-optimum/rbln/transformers/__init__.py,sha256=SdOjpa4Iufo6aOJPvjQwD_vz28dmmNV9AtF5Cz9ajLM,4167
+optimum/rbln/transformers/__init__.py,sha256=AGo3BqVIZrsOzYsQAnnQ25HCstTPBclrXbvvUxVMlqE,4255
 optimum/rbln/transformers/modeling_alias.py,sha256=yx7FnZQWAnrWzivaO5hI7T6i-fyLzt2tMIXG2oDNbPo,1657
 optimum/rbln/transformers/modeling_generic.py,sha256=SD7XjpjnCn-ejNAUWgkaaHV6Fv2Y6K-hbXEXXb9W_H4,18177
 optimum/rbln/transformers/modeling_rope_utils.py,sha256=3zwkhYUyTZhxCJUSmwCc88iiY1TppRWEY9ShwUqNB2k,14293
-optimum/rbln/transformers/models/__init__.py,sha256=wucrA1ybpDfNcrySwdVeK5PZEYl-3ONXJvGpHGTvteo,3683
+optimum/rbln/transformers/models/__init__.py,sha256=zGnYODR-_T65tv6jFjtC8l01LC4vjfm41bM4doCXRvY,3835
 optimum/rbln/transformers/models/auto/__init__.py,sha256=GvGbb3ZpMv-h6euXeZ42jSizoOfrL2O1uvpAnfKxYEo,1034
 optimum/rbln/transformers/models/auto/auto_factory.py,sha256=IK9jFrJ3EEzYQa9_aKpcp2TO68M5YGkA-HcfBVpA2QU,7027
 optimum/rbln/transformers/models/auto/modeling_auto.py,sha256=Un9qoqdy3dO8JBza_bTJF_6_fRVNM9QisihSgTRFI-o,3933
@@ -50,8 +56,8 @@ optimum/rbln/transformers/models/bart/bart_architecture.py,sha256=dTkgMpNkyh4vT_
 optimum/rbln/transformers/models/bart/modeling_bart.py,sha256=ADRbE-5N3xJ60AzzjJ4BZs_THmB71qs4XTr9iFqsEqE,5667
 optimum/rbln/transformers/models/bert/__init__.py,sha256=YVV7k_laU6yJBawZrgjIWjRmIF-Y4oQQHqyf8lsraQs,691
 optimum/rbln/transformers/models/bert/modeling_bert.py,sha256=-nv-sgmHkyHQIoQvF8-lXOJiL4eaa1pq8MpdN4uRi9M,4668
-optimum/rbln/transformers/models/clip/__init__.py,sha256=ssJqlEt318ti2QaEakGh_tO3Ap1VSPCVF-ymUuvjAJs,698
-optimum/rbln/transformers/models/clip/modeling_clip.py,sha256=E1QfVNq1sTCp7uvuha1ZPfXMwvMTkGV9L4oFdmy1w4g,5724
+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=eT1fbKDL92yGBXtUKA_JibD4kiRPdf3tAFJHP5nlfH4,36646
 optimum/rbln/transformers/models/decoderonly/modeling_decoderonly.py,sha256=2OO8MEgFgcl1VPrQXxqkvmRJJEuFdexwu8XqbHDbR6Y,27609
@@ -108,7 +114,7 @@ optimum/rbln/utils/model_utils.py,sha256=DfD_Z2qvZHqcddXqnzTM1AN8khanj3-DXK2lJvV
 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.2.1a5.dist-info/METADATA,sha256=WSMoEbo3z3TMFB1lqbdJsu4ZeVI9AtewXktRjMk6WQw,5300
-optimum_rbln-0.2.1a5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-optimum_rbln-0.2.1a5.dist-info/licenses/LICENSE,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
-optimum_rbln-0.2.1a5.dist-info/RECORD,,
+optimum_rbln-0.7.2rc0.dist-info/METADATA,sha256=4lKtPSN8dCRMWCpiwizg0MxQtvlVfQnyWVIFhCfAMLs,5301
+optimum_rbln-0.7.2rc0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+optimum_rbln-0.7.2rc0.dist-info/licenses/LICENSE,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
+optimum_rbln-0.7.2rc0.dist-info/RECORD,,