optimum-rbln 0.2.1a4__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.1a4'
-__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