optimum-rbln 0.7.4a4__py3-none-any.whl → 0.7.4a6__py3-none-any.whl

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

@@ -12,73 +12,72 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Dict, List, Tuple, Union
 
 import rebel
-import torch  # noqa: I001
+import torch
 from diffusers import AutoencoderKL
+from diffusers.models.autoencoders.vae import DecoderOutput
 from diffusers.models.modeling_outputs import AutoencoderKLOutput
 from transformers import PretrainedConfig
 
+from ....configuration_utils import RBLNCompileConfig
 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 ...configurations import RBLNAutoencoderKLConfig
 from .vae import RBLNRuntimeVAEDecoder, RBLNRuntimeVAEEncoder, _VAEDecoder, _VAEEncoder
 
 
 if TYPE_CHECKING:
     import torch
-    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, PretrainedConfig
+    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, PretrainedConfig, PreTrainedModel
+
+    from ...modeling_diffusers import RBLNDiffusionMixin, RBLNDiffusionMixinConfig
 
 logger = get_logger(__name__)
 
 
 class RBLNAutoencoderKL(RBLNModel):
     auto_model_class = AutoencoderKL
-    config_name = "config.json"
     hf_library_name = "diffusers"
+    _rbln_config_class = RBLNAutoencoderKLConfig
 
     def __post_init__(self, **kwargs):
         super().__post_init__(**kwargs)
 
-        if self.rbln_config.model_cfg.get("img2img_pipeline") or self.rbln_config.model_cfg.get("inpaint_pipeline"):
+        if self.rbln_config.uses_encoder:
             self.encoder = RBLNRuntimeVAEEncoder(runtime=self.model[0], main_input_name="x")
-            self.decoder = RBLNRuntimeVAEDecoder(runtime=self.model[1], main_input_name="z")
         else:
-            self.decoder = RBLNRuntimeVAEDecoder(runtime=self.model[0], main_input_name="z")
+            self.encoder = None
 
-        self.image_size = self.rbln_config.model_cfg["sample_size"]
+        self.decoder = RBLNRuntimeVAEDecoder(runtime=self.model[-1], main_input_name="z")
+        self.image_size = self.rbln_config.image_size
 
     @classmethod
-    def get_compiled_model(cls, model, rbln_config: RBLNConfig):
-        def compile_img2img():
-            encoder_model = _VAEEncoder(model)
-            decoder_model = _VAEDecoder(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])
+    def get_compiled_model(cls, model, rbln_config: RBLNAutoencoderKLConfig) -> Dict[str, rebel.RBLNCompiledModel]:
+        if rbln_config.uses_encoder:
+            expected_models = ["encoder", "decoder"]
+        else:
+            expected_models = ["decoder"]
 
-            return {"encoder": enc_compiled_model, "decoder": dec_compiled_model}
+        compiled_models = {}
+        for i, model_name in enumerate(expected_models):
+            if model_name == "encoder":
+                wrapped_model = _VAEEncoder(model)
+            else:
+                wrapped_model = _VAEDecoder(model)
 
-        def compile_text2img():
-            decoder_model = _VAEDecoder(model)
-            decoder_model.eval()
+            wrapped_model.eval()
 
-            dec_compiled_model = cls.compile(decoder_model, rbln_compile_config=rbln_config.compile_cfgs[0])
+            compiled_models[model_name] = cls.compile(wrapped_model, rbln_compile_config=rbln_config.compile_cfgs[i])
 
-            return dec_compiled_model
-
-        if rbln_config.model_cfg.get("img2img_pipeline") or rbln_config.model_cfg.get("inpaint_pipeline"):
-            return compile_img2img()
-        else:
-            return compile_text2img()
+        return compiled_models
 
     @classmethod
-    def get_vae_sample_size(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"))
+    def get_vae_sample_size(
+        cls, pipe: "RBLNDiffusionMixin", rbln_config: RBLNAutoencoderKLConfig, return_vae_scale_factor: bool = False
+    ) -> Tuple[int, int]:
+        sample_size = rbln_config.sample_size
         noise_module = getattr(pipe, "unet", None) or getattr(pipe, "transformer", None)
         vae_scale_factor = (
             pipe.vae_scale_factor
@@ -91,139 +90,121 @@ class RBLNAutoencoderKL(RBLNModel):
                 "Cannot find noise processing or predicting module attributes. ex. U-Net, Transformer, ..."
             )
 
-        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")
+        if sample_size is None:
+            sample_size = noise_module.config.sample_size
+            if isinstance(sample_size, int):
+                sample_size = (sample_size, sample_size)
+            sample_size = (sample_size[0] * vae_scale_factor, sample_size[1] * vae_scale_factor)
 
-        elif image_size[0] is None and image_size[1] is None:
-            if rbln_config["img2img_pipeline"]:
-                sample_size = noise_module.config.sample_size
-            elif rbln_config["inpaint_pipeline"]:
-                sample_size = noise_module.config.sample_size * vae_scale_factor
-            else:
-                # In case of text2img, sample size of vae decoder is determined by unet.
-                noise_module_sample_size = noise_module.config.sample_size
-                if isinstance(noise_module_sample_size, int):
-                    sample_size = noise_module_sample_size * vae_scale_factor
-                else:
-                    sample_size = (
-                        noise_module_sample_size[0] * vae_scale_factor,
-                        noise_module_sample_size[1] * vae_scale_factor,
-                    )
+        if return_vae_scale_factor:
+            return sample_size, vae_scale_factor
         else:
-            sample_size = (image_size[0], image_size[1])
-
-        return sample_size
+            return sample_size
 
     @classmethod
-    def update_rbln_config_using_pipe(cls, pipe: RBLNDiffusionMixin, rbln_config: Dict[str, Any]) -> Dict[str, Any]:
-        rbln_config.update({"sample_size": cls.get_vae_sample_size(pipe, rbln_config)})
+    def update_rbln_config_using_pipe(
+        cls, pipe: "RBLNDiffusionMixin", rbln_config: "RBLNDiffusionMixinConfig", submodule_name: str
+    ) -> "RBLNDiffusionMixinConfig":
+        rbln_config.vae.sample_size, rbln_config.vae.vae_scale_factor = cls.get_vae_sample_size(
+            pipe, rbln_config.vae, return_vae_scale_factor=True
+        )
         return rbln_config
 
     @classmethod
-    def _get_rbln_config(
+    def _update_rbln_config(
         cls,
         preprocessors: Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"],
+        model: "PreTrainedModel",
         model_config: "PretrainedConfig",
-        rbln_kwargs: Dict[str, Any] = {},
-    ) -> RBLNConfig:
-        rbln_batch_size = rbln_kwargs.get("batch_size")
-        sample_size = rbln_kwargs.get("sample_size")
-        is_img2img = rbln_kwargs.get("img2img_pipeline")
-        is_inpaint = rbln_kwargs.get("inpaint_pipeline")
-
-        if rbln_batch_size is None:
-            rbln_batch_size = 1
-
-        if sample_size is None:
-            sample_size = model_config.sample_size
+        rbln_config: RBLNAutoencoderKLConfig,
+    ) -> RBLNAutoencoderKLConfig:
+        if rbln_config.sample_size is None:
+            rbln_config.sample_size = model_config.sample_size
 
-        if isinstance(sample_size, int):
-            sample_size = (sample_size, sample_size)
+        if isinstance(rbln_config.sample_size, int):
+            rbln_config.sample_size = (rbln_config.sample_size, rbln_config.sample_size)
 
-        rbln_kwargs["sample_size"] = sample_size
+        if rbln_config.in_channels is None:
+            rbln_config.in_channels = model_config.in_channels
 
-        if hasattr(model_config, "block_out_channels"):
-            vae_scale_factor = 2 ** (len(model_config.block_out_channels) - 1)
-        else:
-            # vae image processor default value 8 (int)
-            vae_scale_factor = 8
+        if rbln_config.latent_channels is None:
+            rbln_config.latent_channels = model_config.latent_channels
 
-        dec_shape = (sample_size[0] // vae_scale_factor, sample_size[1] // vae_scale_factor)
-        enc_shape = (sample_size[0], sample_size[1])
+        if rbln_config.vae_scale_factor is None:
+            if hasattr(model_config, "block_out_channels"):
+                rbln_config.vae_scale_factor = 2 ** (len(model_config.block_out_channels) - 1)
+            else:
+                # vae image processor default value 8 (int)
+                rbln_config.vae_scale_factor = 8
 
-        if is_img2img or is_inpaint:
+        compile_cfgs = []
+        if rbln_config.uses_encoder:
             vae_enc_input_info = [
                 (
                     "x",
-                    [rbln_batch_size, model_config.in_channels, enc_shape[0], enc_shape[1]],
+                    [
+                        rbln_config.batch_size,
+                        rbln_config.in_channels,
+                        rbln_config.sample_size[0],
+                        rbln_config.sample_size[1],
+                    ],
                     "float32",
                 )
             ]
-            vae_dec_input_info = [
-                (
-                    "z",
-                    [rbln_batch_size, model_config.latent_channels, dec_shape[0], dec_shape[1]],
-                    "float32",
-                )
-            ]
-
-            enc_rbln_compile_config = RBLNCompileConfig(compiled_model_name="encoder", input_info=vae_enc_input_info)
-            dec_rbln_compile_config = RBLNCompileConfig(compiled_model_name="decoder", input_info=vae_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,
+            compile_cfgs.append(RBLNCompileConfig(compiled_model_name="encoder", input_info=vae_enc_input_info))
+
+        vae_dec_input_info = [
+            (
+                "z",
+                [
+                    rbln_config.batch_size,
+                    rbln_config.latent_channels,
+                    rbln_config.latent_sample_size[0],
+                    rbln_config.latent_sample_size[1],
+                ],
+                "float32",
             )
-            return rbln_config
+        ]
+        compile_cfgs.append(RBLNCompileConfig(compiled_model_name="decoder", input_info=vae_dec_input_info))
 
-        vae_config = RBLNCompileConfig(
-            input_info=[
-                (
-                    "z",
-                    [rbln_batch_size, model_config.latent_channels, dec_shape[0], dec_shape[1]],
-                    "float32",
-                )
-            ]
-        )
-        rbln_config = RBLNConfig(
-            rbln_cls=cls.__name__,
-            compile_cfgs=[vae_config],
-            rbln_kwargs=rbln_kwargs,
-        )
+        rbln_config.set_compile_cfgs(compile_cfgs)
         return rbln_config
 
     @classmethod
     def _create_runtimes(
         cls,
         compiled_models: List[rebel.RBLNCompiledModel],
-        rbln_device_map: Dict[str, int],
-        activate_profiler: Optional[bool] = None,
+        rbln_config: RBLNAutoencoderKLConfig,
     ) -> List[rebel.Runtime]:
         if len(compiled_models) == 1:
-            if DEFAULT_COMPILED_MODEL_NAME not in rbln_device_map:
-                cls._raise_missing_compiled_file_error([DEFAULT_COMPILED_MODEL_NAME])
-
-            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
-                )
-            ]
+            # decoder
+            expected_models = ["decoder"]
+        else:
+            # encoder, decoder
+            expected_models = ["encoder", "decoder"]
 
-        if any(model_name not in rbln_device_map for model_name in ["encoder", "decoder"]):
-            cls._raise_missing_compiled_file_error(["encoder", "decoder"])
+        if any(model_name not in rbln_config.device_map for model_name in expected_models):
+            cls._raise_missing_compiled_file_error(expected_models)
 
-        device_vals = [rbln_device_map["encoder"], rbln_device_map["decoder"]]
+        device_vals = [rbln_config.device_map[model_name] for model_name in expected_models]
         return [
-            compiled_model.create_runtime(tensor_type="pt", device=device_val, activate_profiler=activate_profiler)
+            rebel.Runtime(
+                compiled_model,
+                tensor_type="pt",
+                device=device_val,
+                activate_profiler=rbln_config.activate_profiler,
+            )
             for compiled_model, device_val in zip(compiled_models, device_vals)
         ]
 
-    def encode(self, x: torch.FloatTensor, **kwargs) -> torch.FloatTensor:
+    def encode(self, x: torch.FloatTensor, return_dict: bool = True, **kwargs) -> torch.FloatTensor:
         posterior = self.encoder.encode(x)
+        if not return_dict:
+            return (posterior,)
         return AutoencoderKLOutput(latent_dist=posterior)
 
-    def decode(self, z: torch.FloatTensor, **kwargs) -> torch.FloatTensor:
-        return self.decoder.decode(z)
+    def decode(self, z: torch.FloatTensor, return_dict: bool = True, **kwargs) -> torch.FloatTensor:
+        dec = self.decoder.decode(z)
+        if not return_dict:
+            return (dec,)
+        return DecoderOutput(sample=dec)

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

@@ -14,11 +14,9 @@
 
 from typing import TYPE_CHECKING, List
 
-import torch  # noqa: I001
+import torch
 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
 from ....utils.runtime_utils import RBLNPytorchRuntime
@@ -34,12 +32,12 @@ class RBLNRuntimeVAEEncoder(RBLNPytorchRuntime):
     def encode(self, x: torch.FloatTensor, **kwargs) -> torch.FloatTensor:
         moments = self.forward(x.contiguous())
         posterior = DiagonalGaussianDistribution(moments)
-        return AutoencoderKLOutput(latent_dist=posterior)
+        return posterior
 
 
 class RBLNRuntimeVAEDecoder(RBLNPytorchRuntime):
     def decode(self, z: torch.FloatTensor, **kwargs) -> torch.FloatTensor:
-        return (self.forward(z),)
+        return self.forward(z)
 
 
 class _VAEDecoder(torch.nn.Module):
@@ -78,7 +76,7 @@ class _VAEEncoder(torch.nn.Module):
 class RBLNRuntimeVQEncoder(RBLNPytorchRuntime):
     def encode(self, x: torch.FloatTensor, **kwargs) -> torch.FloatTensor:
         h = self.forward(x.contiguous())
-        return VQEncoderOutput(latents=h)
+        return h
 
 
 class RBLNRuntimeVQDecoder(RBLNPytorchRuntime):

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

@@ -12,24 +12,24 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+from typing import TYPE_CHECKING, List, 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 ....configuration_utils import RBLNCompileConfig, RBLNModelConfig
 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 ...configurations.models.configuration_vq_model import RBLNVQModelConfig
+from ...modeling_diffusers import RBLNDiffusionMixin, RBLNDiffusionMixinConfig
 from .vae import RBLNRuntimeVQDecoder, RBLNRuntimeVQEncoder, _VQDecoder, _VQEncoder
 
 
 if TYPE_CHECKING:
-    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer
+    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, PretrainedConfig, PreTrainedModel
 
 logger = get_logger(__name__)
 
@@ -42,126 +42,125 @@ class RBLNVQModel(RBLNModel):
     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")
+        if self.rbln_config.uses_encoder:
+            self.encoder = RBLNRuntimeVQEncoder(runtime=self.model[0], main_input_name="x")
+        else:
+            self.encoder = None
+
+        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]
+        self.image_size = self.rbln_config.image_size
 
     @classmethod
-    def get_compiled_model(cls, model, rbln_config: RBLNConfig):
-        encoder_model = _VQEncoder(model)
-        decoder_model = _VQDecoder(model)
-        encoder_model.eval()
-        decoder_model.eval()
+    def get_compiled_model(cls, model, rbln_config: RBLNModelConfig):
+        if rbln_config.uses_encoder:
+            expected_models = ["encoder", "decoder"]
+        else:
+            expected_models = ["decoder"]
 
-        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])
+        compiled_models = {}
+        for i, model_name in enumerate(expected_models):
+            if model_name == "encoder":
+                wrapped_model = _VQEncoder(model)
+            else:
+                wrapped_model = _VQDecoder(model)
 
-        return {"encoder": enc_compiled_model, "decoder": dec_compiled_model}
+            wrapped_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 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,
-            }
-        )
+            compiled_models[model_name] = cls.compile(wrapped_model, rbln_compile_config=rbln_config.compile_cfgs[i])
+
+        return compiled_models
 
+    @classmethod
+    def update_rbln_config_using_pipe(
+        cls, pipe: RBLNDiffusionMixin, rbln_config: "RBLNDiffusionMixinConfig", submodule_name: str
+    ) -> "RBLNDiffusionMixinConfig":
         return rbln_config
 
     @classmethod
-    def _get_rbln_config(
+    def _update_rbln_config(
         cls,
         preprocessors: Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"],
+        model: "PreTrainedModel",
         model_config: "PretrainedConfig",
-        rbln_kwargs: Dict[str, Any] = {},
-    ) -> RBLNConfig:
-        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
-
+        rbln_config: RBLNVQModelConfig,
+    ) -> RBLNVQModelConfig:
         if hasattr(model_config, "block_out_channels"):
-            scale_factor = 2 ** (len(model_config.block_out_channels) - 1)
+            rbln_config.vqmodel_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)
+            rbln_config.vqmodel_scale_factor = 8
+
+        compile_cfgs = []
+        if rbln_config.uses_encoder:
+            enc_input_info = [
+                (
+                    "x",
+                    [
+                        rbln_config.batch_size,
+                        model_config.in_channels,
+                        rbln_config.sample_size[0],
+                        rbln_config.sample_size[1],
+                    ],
+                    "float32",
+                )
+            ]
+            enc_rbln_compile_config = RBLNCompileConfig(compiled_model_name="encoder", input_info=enc_input_info)
+            compile_cfgs.append(enc_rbln_compile_config)
 
-        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]],
+                [
+                    rbln_config.batch_size,
+                    model_config.latent_channels,
+                    rbln_config.latent_sample_size[0],
+                    rbln_config.latent_sample_size[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.append(dec_rbln_compile_config)
 
-        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,
-        )
+        rbln_config.set_compile_cfgs(compile_cfgs)
         return rbln_config
 
     @classmethod
     def _create_runtimes(
         cls,
         compiled_models: List[rebel.RBLNCompiledModel],
-        rbln_device_map: Dict[str, int],
-        activate_profiler: Optional[bool] = None,
+        rbln_config: RBLNVQModelConfig,
     ) -> 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
-                )
-            ]
+            # decoder
+            expected_models = ["decoder"]
+        else:
+            # encoder, decoder
+            expected_models = ["encoder", "decoder"]
 
-        device_vals = [rbln_device_map["encoder"], rbln_device_map["decoder"]]
+        if any(model_name not in rbln_config.device_map for model_name in expected_models):
+            cls._raise_missing_compiled_file_error(expected_models)
+
+        device_vals = [rbln_config.device_map[model_name] for model_name in expected_models]
         return [
-            compiled_model.create_runtime(tensor_type="pt", device=device_val, activate_profiler=activate_profiler)
+            rebel.Runtime(
+                compiled_model,
+                tensor_type="pt",
+                device=device_val,
+                activate_profiler=rbln_config.activate_profiler,
+            )
             for compiled_model, device_val in zip(compiled_models, device_vals)
         ]
 
-    def encode(self, x: torch.FloatTensor, **kwargs) -> torch.FloatTensor:
+    def encode(self, x: torch.FloatTensor, return_dict: bool = True, **kwargs) -> torch.FloatTensor:
         posterior = self.encoder.encode(x)
+        if not return_dict:
+            return (posterior,)
         return VQEncoderOutput(latents=posterior)
 
-    def decode(self, h: torch.FloatTensor, **kwargs) -> torch.FloatTensor:
+    def decode(self, h: torch.FloatTensor, return_dict: bool = True, **kwargs) -> torch.FloatTensor:
         dec, commit_loss = self.decoder.decode(h, **kwargs)
+        if not return_dict:
+            return (dec, commit_loss)
         return DecoderOutput(sample=dec, commit_loss=commit_loss)