optimum-rbln 0.8.0.post2__py3-none-any.whl → 0.8.1__py3-none-any.whl

optimum/rbln/diffusers/modeling_diffusers.py

@@ -19,10 +19,11 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type, Union
 
 import torch
 
-from ..configuration_utils import ContextRblnConfig, RBLNModelConfig
+from ..configuration_utils import ContextRblnConfig, RBLNModelConfig, get_rbln_config_class
 from ..modeling import RBLNModel
 from ..utils.decorator_utils import remove_compile_time_kwargs
 from ..utils.logging import get_logger
+from ..utils.model_utils import get_rbln_model_cls
 
 
 logger = get_logger(__name__)
@@ -44,7 +45,7 @@ class RBLNDiffusionMixin:
     To use this mixin:
 
     1. Create a new pipeline class that inherits from both this mixin and the original StableDiffusionPipeline.
-    2. Define the required _submodules class variable listing the components to be compiled.
+    2. Define the required _submodules and _optional_submodules class variable listing the components to be compiled.
 
     Example:
         ```python
@@ -54,6 +55,7 @@ class RBLNDiffusionMixin:
 
     Class Variables:
         _submodules: List of submodule names that should be compiled (typically ["text_encoder", "unet", "vae"])
+        _optional_submodules: List of submodule names compiled without inheriting RBLNModel (typically ["safety_checker"])
 
     Methods:
         from_pretrained: Creates and optionally compiles a model from a pretrained checkpoint
@@ -66,6 +68,7 @@ class RBLNDiffusionMixin:
 
     _connected_classes = {}
     _submodules = []
+    _optional_submodules = []
     _prefix = {}
     _rbln_config_class = None
     _hf_class = None
@@ -110,18 +113,10 @@ class RBLNDiffusionMixin:
 
     @classmethod
     def get_rbln_config_class(cls) -> Type[RBLNModelConfig]:
-        """
-        Lazily loads and caches the corresponding RBLN model config class.
-        """
+        # Lazily loads and caches the corresponding RBLN model config class.
         if cls._rbln_config_class is None:
             rbln_config_class_name = cls.__name__ + "Config"
-            library = importlib.import_module("optimum.rbln")
-            cls._rbln_config_class = getattr(library, rbln_config_class_name, None)
-            if cls._rbln_config_class is None:
-                raise ValueError(
-                    f"RBLN config class {rbln_config_class_name} not found. This is an internal error. "
-                    "Please report it to the developers."
-                )
+            cls._rbln_config_class = get_rbln_config_class(rbln_config_class_name)
         return cls._rbln_config_class
 
     @classmethod
@@ -143,7 +138,7 @@ class RBLNDiffusionMixin:
         lora_ids: Optional[Union[str, List[str]]] = None,
         lora_weights_names: Optional[Union[str, List[str]]] = None,
         lora_scales: Optional[Union[float, List[float]]] = None,
-        **kwargs,
+        **kwargs: Dict[str, Any],
     ) -> "RBLNDiffusionMixin":
         """
         Load a pretrained diffusion pipeline from a model checkpoint, with optional compilation for RBLN NPUs.
@@ -157,24 +152,25 @@ class RBLNDiffusionMixin:
         Args:
             model_id (`str`):
                 The model ID or path to the pretrained model to load. Can be either:
+
                 - A model ID from the HuggingFace Hub
                 - A local path to a saved model directory
-            export (`bool`, *optional*, defaults to `False`):
+            export:
                 If True, takes a PyTorch model from `model_id` and compiles it for RBLN NPU execution.
                 If False, loads an already compiled RBLN model from `model_id` without recompilation.
-            model_save_dir (`os.PathLike`, *optional*):
+            model_save_dir:
                 Directory to save the compiled model artifacts. Only used when `export=True`.
                 If not provided and `export=True`, a temporary directory is used.
-            rbln_config (`Dict[str, Any]`, *optional*, defaults to `{}`):
+            rbln_config:
                 Configuration options for RBLN compilation. Can include settings for specific submodules
                 such as `text_encoder`, `unet`, and `vae`. Configuration can be tailored to the specific
                 pipeline being compiled.
-            lora_ids (`str` or `List[str]`, *optional*):
+            lora_ids:
                 LoRA adapter ID(s) to load and apply before compilation. LoRA weights are fused
                 into the model weights during compilation. Only used when `export=True`.
-            lora_weights_names (`str` or `List[str]`, *optional*):
+            lora_weights_names:
                 Names of specific LoRA weight files to load, corresponding to lora_ids. Only used when `export=True`.
-            lora_scales (`float` or `List[float]`, *optional*):
+            lora_scales:
                 Scaling factor(s) to apply to the LoRA adapter(s). Only used when `export=True`.
             **kwargs:
                 Additional arguments to pass to the underlying diffusion pipeline constructor or the
@@ -182,39 +178,50 @@ class RBLNDiffusionMixin:
                 or the particular diffusion pipeline being used.
 
         Returns:
-            `RBLNDiffusionMixin`: A compiled or loaded diffusion pipeline that can be used for inference on RBLN NPU.
-            The returned object is an instance of the class that called this method, inheriting from RBLNDiffusionMixin.
+            A compiled or loaded diffusion pipeline that can be used for inference on RBLN NPU.
+                The returned object is an instance of the class that called this method, inheriting from RBLNDiffusionMixin.
         """
         rbln_config, kwargs = cls.get_rbln_config_class().initialize_from_kwargs(rbln_config, **kwargs)
 
         if export:
             # keep submodules if user passed any of them.
             passed_submodules = {
-                name: kwargs.pop(name) for name in cls._submodules if isinstance(kwargs.get(name), RBLNModel)
+                name: kwargs.pop(name)
+                for name in cls._submodules + cls._optional_submodules
+                if isinstance(kwargs.get(name), RBLNModel)
             }
 
         else:
             # raise error if any of submodules are torch module.
             model_index_config = cls.load_config(pretrained_model_name_or_path=model_id)
-            for submodule_name in cls._submodules:
-                if isinstance(kwargs.get(submodule_name), torch.nn.Module):
-                    raise AssertionError(
-                        f"{submodule_name} is not compiled torch module. If you want to compile, set `export=True`."
+            for submodule_name in cls._submodules + cls._optional_submodules:
+                passed_submodule = kwargs.get(submodule_name, None)
+
+                if passed_submodule is None:
+                    module_name, class_name = model_index_config[submodule_name]
+                    if module_name != "optimum.rbln":
+                        raise ValueError(
+                            f"Invalid module_name '{module_name}' found in model_index.json for "
+                            f"submodule '{submodule_name}'. "
+                            "Expected 'optimum.rbln'. Please check the model_index.json configuration."
+                            "If you want to compile, set `export=True`."
+                        )
+
+                    submodule_cls = get_rbln_model_cls(class_name)
+                    submodule_config = getattr(rbln_config, submodule_name)
+                    submodule = submodule_cls.from_pretrained(
+                        model_id, export=False, subfolder=submodule_name, rbln_config=submodule_config
                     )
 
-                module_name, class_name = model_index_config[submodule_name]
-                if module_name != "optimum.rbln":
-                    raise ValueError(
-                        f"Invalid module_name '{module_name}' found in model_index.json for "
-                        f"submodule '{submodule_name}'. "
-                        "Expected 'optimum.rbln'. Please check the model_index.json configuration."
-                    )
+                else:
+                    if passed_submodule.__class__.__name__.startswith("RBLN"):
+                        submodule = passed_submodule
+
+                    elif isinstance(passed_submodule, torch.nn.Module):
+                        raise AssertionError(
+                            f"{submodule_name} is not compiled torch module. If you want to compile, set `export=True`."
+                        )
 
-                submodule_cls: Type[RBLNModel] = getattr(importlib.import_module("optimum.rbln"), class_name)
-                submodule_config = getattr(rbln_config, submodule_name)
-                submodule = submodule_cls.from_pretrained(
-                    model_id, export=False, subfolder=submodule_name, rbln_config=submodule_config
-                )
                 kwargs[submodule_name] = submodule
 
         with ContextRblnConfig(
@@ -293,7 +300,6 @@ class RBLNDiffusionMixin:
             elif isinstance(submodule, RBLNModel):
                 pass
             elif submodule_name == "controlnet" and hasattr(submodule, "nets"):
-                # In case of multicontrolnet
                 submodule = cls._compile_multicontrolnet(
                     controlnets=submodule,
                     model_save_dir=model_save_dir,
@@ -301,11 +307,8 @@ class RBLNDiffusionMixin:
                     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(
+                submodule = submodule_rbln_cls.from_model(
                     model=submodule,
                     subfolder=subfolder,
                     model_save_dir=model_save_dir,
@@ -362,10 +365,16 @@ class RBLNDiffusionMixin:
             # Causing warning messeages.
 
         update_dict = {}
-        for submodule_name in cls._submodules:
+        for submodule_name in cls._submodules + cls._optional_submodules:
             # replace submodule
-            setattr(model, submodule_name, submodules[submodule_name])
-            update_dict[submodule_name] = ("optimum.rbln", submodules[submodule_name].__class__.__name__)
+            if submodule_name in submodules:
+                setattr(model, submodule_name, submodules[submodule_name])
+                update_dict[submodule_name] = ("optimum.rbln", submodules[submodule_name].__class__.__name__)
+            else:
+                # It assumes that the modules in _optional_components is compiled
+                # and already registered as an attribute of the model.
+                update_dict[submodule_name] = ("optimum.rbln", getattr(model, submodule_name).__class__.__name__)
+
         if cls._load_connected_pipes:
             for connected_pipe_name, connected_pipe_cls in cls._connected_classes.items():
                 prefix = cls._prefix.get(connected_pipe_name, "")
@@ -396,31 +405,29 @@ class RBLNDiffusionMixin:
         return model
 
     def get_compiled_image_size(self):
-        if hasattr(self, "vae"):
+        if hasattr(self, "vae") and hasattr(self.vae, "image_size"):
             compiled_image_size = self.vae.image_size
         else:
             compiled_image_size = None
         return compiled_image_size
 
     def handle_additional_kwargs(self, **kwargs):
-        """
-        Function to handle additional compile-time parameters during inference.
-
-        If the additional variable is determined by another module, this method should be overrided.
-
-        Example:
-            ```python
-            if hasattr(self, "movq"):
-                compiled_image_size = self.movq.image_size
-                kwargs["height"] = compiled_image_size[0]
-                kwargs["width"] = compiled_image_size[1]
-
-            compiled_num_frames = self.unet.rbln_config.num_frames
-            if compiled_num_frames is not None:
-                kwargs["num_frames"] = compiled_num_frames
-            return kwargs
-            ```
-        """
+        # Function to handle additional compile-time parameters during inference.
+
+        # If the additional variable is determined by another module, this method should be overrided.
+
+        # Example:
+        #     ```python
+        #     if hasattr(self, "movq"):
+        #         compiled_image_size = self.movq.image_size
+        #         kwargs["height"] = compiled_image_size[0]
+        #         kwargs["width"] = compiled_image_size[1]
+
+        #     compiled_num_frames = self.unet.rbln_config.num_frames
+        #     if compiled_num_frames is not None:
+        #         kwargs["num_frames"] = compiled_num_frames
+        #     return kwargs
+        #     ```
         return kwargs
 
     @remove_compile_time_kwargs

optimum/rbln/diffusers/models/__init__.py

@@ -20,6 +20,7 @@ from transformers.utils import _LazyModule
 _import_structure = {
     "autoencoders": [
         "RBLNAutoencoderKL",
+        "RBLNAutoencoderKLCosmos",
         "RBLNVQModel",
     ],
     "unets": [
@@ -28,6 +29,7 @@ _import_structure = {
     "controlnet": ["RBLNControlNetModel"],
     "transformers": [
         "RBLNPriorTransformer",
+        "RBLNCosmosTransformer3DModel",
         "RBLNSD3Transformer2DModel",
     ],
 }
@@ -35,10 +37,12 @@ _import_structure = {
 if TYPE_CHECKING:
     from .autoencoders import (
         RBLNAutoencoderKL,
+        RBLNAutoencoderKLCosmos,
         RBLNVQModel,
     )
     from .controlnet import RBLNControlNetModel
     from .transformers import (
+        RBLNCosmosTransformer3DModel,
         RBLNPriorTransformer,
         RBLNSD3Transformer2DModel,
     )

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

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

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

@@ -38,6 +38,17 @@ logger = get_logger(__name__)
 
 
 class RBLNAutoencoderKL(RBLNModel):
+    """
+    RBLN implementation of AutoencoderKL (VAE) for diffusion models.
+
+    This model is used to accelerate AutoencoderKL (VAE) models from diffusers library on RBLN NPUs.
+    It can be configured to include both encoder and decoder, or just the decoder part for latent-to-image
+    conversion.
+
+    This class inherits from [`RBLNModel`]. Check the superclass documentation for the generic methods
+    the library implements for all its models.
+    """
+
     auto_model_class = AutoencoderKL
     hf_library_name = "diffusers"
     _rbln_config_class = RBLNAutoencoderKLConfig
@@ -69,7 +80,12 @@ class RBLNAutoencoderKL(RBLNModel):
 
             wrapped_model.eval()
 
-            compiled_models[model_name] = cls.compile(wrapped_model, rbln_compile_config=rbln_config.compile_cfgs[i])
+            compiled_models[model_name] = cls.compile(
+                wrapped_model,
+                rbln_compile_config=rbln_config.compile_cfgs[i],
+                create_runtimes=rbln_config.create_runtimes,
+                device=rbln_config.device_map[model_name],
+            )
 
         return compiled_models
 

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

@@ -0,0 +1,219 @@
+# Copyright 2025 Rebellions Inc. All rights reserved.
+
+# 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, Dict, List, Union
+
+import rebel
+import torch
+from diffusers.models.autoencoders.autoencoder_kl_cosmos import AutoencoderKLCosmos, CosmosCausalConv3d
+from diffusers.models.autoencoders.vae import DecoderOutput
+from diffusers.models.modeling_outputs import AutoencoderKLOutput
+from torch.nn import functional as F
+from transformers import PretrainedConfig
+
+from ....configuration_utils import RBLNCompileConfig
+from ....modeling import RBLNModel
+from ....utils.logging import get_logger
+from ...configurations import RBLNAutoencoderKLCosmosConfig
+from .vae import RBLNRuntimeCosmosVAEDecoder, RBLNRuntimeCosmosVAEEncoder, _VAECosmosDecoder, _VAECosmosEncoder
+
+
+if TYPE_CHECKING:
+    import torch
+    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, PreTrainedModel
+
+    from ...modeling_diffusers import RBLNDiffusionMixin, RBLNDiffusionMixinConfig
+
+logger = get_logger(__name__)
+
+
+class RBLNAutoencoderKLCosmos(RBLNModel):
+    """
+    RBLN implementation of AutoencoderKLCosmos for diffusion models.
+
+    This model is used to accelerate AutoencoderKLCosmos models from diffusers library on RBLN NPUs.
+    It can be configured to include both encoder and decoder, or just the decoder part for latent-to-video
+    conversion.
+
+    This class inherits from [`RBLNModel`]. Check the superclass documentation for the generic methods
+    the library implements for all its models.
+    """
+
+    auto_model_class = AutoencoderKLCosmos
+    hf_library_name = "diffusers"
+    _rbln_config_class = RBLNAutoencoderKLCosmosConfig
+
+    def __post_init__(self, **kwargs):
+        super().__post_init__(**kwargs)
+
+        if self.rbln_config.uses_encoder:
+            self.encoder = RBLNRuntimeCosmosVAEEncoder(
+                runtime=self.model[0], main_input_name="x", use_slicing=self.rbln_config.use_slicing
+            )
+
+        self.decoder = RBLNRuntimeCosmosVAEDecoder(
+            runtime=self.model[-1], main_input_name="z", use_slicing=self.rbln_config.use_slicing
+        )
+        self.image_size = self.rbln_config.image_size
+
+    @classmethod
+    def wrap_model_if_needed(
+        cls, model: torch.nn.Module, rbln_config: RBLNAutoencoderKLCosmosConfig
+    ) -> torch.nn.Module:
+        decoder_model = _VAECosmosDecoder(model)
+        decoder_model.eval()
+
+        if rbln_config.uses_encoder:
+            encoder_model = _VAECosmosEncoder(model)
+            encoder_model.eval()
+            return encoder_model, decoder_model
+        else:
+            return decoder_model
+
+    @classmethod
+    def get_compiled_model(
+        cls, model, rbln_config: RBLNAutoencoderKLCosmosConfig
+    ) -> Dict[str, rebel.RBLNCompiledModel]:
+        def replaced_forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+            if self.temporal_pad != 0:
+                hidden_states_prev = hidden_states[:, :, :1, ...].repeat(1, 1, self.temporal_pad, 1, 1)
+                hidden_states = torch.cat([hidden_states_prev, hidden_states], dim=2)
+            hidden_states = F.pad(hidden_states, (*self.spatial_pad, 0, 0), mode=self.pad_mode, value=0.0)
+            return super(CosmosCausalConv3d, self).forward(hidden_states)
+
+        try:
+            original_forward = CosmosCausalConv3d.forward
+            CosmosCausalConv3d.forward = replaced_forward
+
+            compiled_models = {}
+            if rbln_config.uses_encoder:
+                encoder_model, decoder_model = cls.wrap_model_if_needed(model, rbln_config)
+                enc_compiled_model = cls.compile(
+                    encoder_model,
+                    rbln_compile_config=rbln_config.compile_cfgs[0],
+                    create_runtimes=rbln_config.create_runtimes,
+                    device=rbln_config.device_map["encoder"],
+                )
+                compiled_models["encoder"] = enc_compiled_model
+            else:
+                decoder_model = cls.wrap_model_if_needed(model, rbln_config)
+            dec_compiled_model = cls.compile(
+                decoder_model,
+                rbln_compile_config=rbln_config.compile_cfgs[-1],
+                create_runtimes=rbln_config.create_runtimes,
+                device=rbln_config.device_map["decoder"],
+            )
+            compiled_models["decoder"] = dec_compiled_model
+
+        finally:
+            CosmosCausalConv3d.forward = original_forward
+
+        return compiled_models
+
+    @classmethod
+    def update_rbln_config_using_pipe(
+        cls, pipe: "RBLNDiffusionMixin", rbln_config: "RBLNDiffusionMixinConfig", submodule_name: str
+    ) -> "RBLNDiffusionMixinConfig":
+        rbln_config.vae.num_channels_latents = pipe.transformer.config.out_channels
+        rbln_config.vae.vae_scale_factor_temporal = pipe.vae_scale_factor_temporal
+        rbln_config.vae.vae_scale_factor_spatial = pipe.vae_scale_factor_spatial
+        return rbln_config
+
+    @classmethod
+    def _update_rbln_config(
+        cls,
+        preprocessors: Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"],
+        model: "PreTrainedModel",
+        model_config: "PretrainedConfig",
+        rbln_config: RBLNAutoencoderKLCosmosConfig,
+    ) -> RBLNAutoencoderKLCosmosConfig:
+        batch_size = 1 if rbln_config.use_slicing else rbln_config.batch_size
+        compile_cfgs = []
+        if rbln_config.uses_encoder:
+            vae_enc_input_info = [
+                (
+                    "x",
+                    [
+                        batch_size,
+                        model_config.in_channels,
+                        rbln_config.num_frames,
+                        rbln_config.height,
+                        rbln_config.width,
+                    ],
+                    "float32",
+                ),
+            ]
+            compile_cfgs.append(RBLNCompileConfig(compiled_model_name="encoder", input_info=vae_enc_input_info))
+
+        num_latent_frames = (rbln_config.num_frames - 1) // rbln_config.vae_scale_factor_temporal + 1
+        latent_height = rbln_config.height // rbln_config.vae_scale_factor_spatial
+        latent_width = rbln_config.width // rbln_config.vae_scale_factor_spatial
+
+        vae_dec_input_info = [
+            (
+                "z",
+                [
+                    batch_size,
+                    rbln_config.num_channels_latents,
+                    num_latent_frames,
+                    latent_height,
+                    latent_width,
+                ],
+                "float32",
+            ),
+        ]
+        compile_cfgs.append(RBLNCompileConfig(compiled_model_name="decoder", input_info=vae_dec_input_info))
+
+        rbln_config.set_compile_cfgs(compile_cfgs)
+        return rbln_config
+
+    @classmethod
+    def _create_runtimes(
+        cls,
+        compiled_models: List[rebel.RBLNCompiledModel],
+        rbln_config: RBLNAutoencoderKLCosmosConfig,
+    ) -> List[rebel.Runtime]:
+        if len(compiled_models) == 1:
+            # decoder
+            expected_models = ["decoder"]
+        else:
+            expected_models = ["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_config.device_map[model_name] for model_name in expected_models]
+        return [
+            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, 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, return_dict: bool = True) -> torch.FloatTensor:
+        decoded = self.decoder.decode(z)
+
+        if not return_dict:
+            return (decoded,)
+
+        return DecoderOutput(sample=decoded)

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

@@ -15,17 +15,13 @@
 from typing import TYPE_CHECKING, List
 
 import torch
-from diffusers import AutoencoderKL, VQModel
-from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution
+from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution, IdentityDistribution
 
-from ....utils.logging import get_logger
 from ....utils.runtime_utils import RBLNPytorchRuntime
 
 
 if TYPE_CHECKING:
-    import torch
-
-logger = get_logger(__name__)
+    from diffusers import AutoencoderKL, AutoencoderKLCosmos, VQModel
 
 
 class RBLNRuntimeVAEEncoder(RBLNPytorchRuntime):
@@ -40,6 +36,27 @@ class RBLNRuntimeVAEDecoder(RBLNPytorchRuntime):
         return self.forward(z)
 
 
+class RBLNRuntimeCosmosVAEEncoder(RBLNPytorchRuntime):
+    def encode(self, x: torch.FloatTensor, **kwargs) -> torch.FloatTensor:
+        if self.use_slicing and x.shape[0] > 1:
+            encoded_slices = [self.forward(x_slice) for x_slice in x.split(1)]
+            h = torch.cat(encoded_slices)
+        else:
+            h = self.forward(x)
+        posterior = IdentityDistribution(h)
+        return posterior
+
+
+class RBLNRuntimeCosmosVAEDecoder(RBLNPytorchRuntime):
+    def decode(self, z: torch.FloatTensor, **kwargs) -> torch.FloatTensor:
+        if self.use_slicing and z.shape[0] > 1:
+            decoded_slices = [self.forward(z_slice) for z_slice in z.split(1)]
+            decoded = torch.cat(decoded_slices)
+        else:
+            decoded = self.forward(z)
+        return decoded
+
+
 class _VAEDecoder(torch.nn.Module):
     def __init__(self, vae: "AutoencoderKL"):
         super().__init__()
@@ -73,6 +90,26 @@ class _VAEEncoder(torch.nn.Module):
         return vae_out
 
 
+class _VAECosmosEncoder(torch.nn.Module):
+    def __init__(self, vae: "AutoencoderKLCosmos"):
+        super().__init__()
+        self.vae = vae
+
+    def forward(self, x):
+        vae_out = self.vae._encode(x)
+        return vae_out
+
+
+class _VAECosmosDecoder(torch.nn.Module):
+    def __init__(self, vae: "AutoencoderKLCosmos"):
+        super().__init__()
+        self.vae = vae
+
+    def forward(self, z):
+        vae_out = self.vae._decode(z, return_dict=False)
+        return vae_out
+
+
 class RBLNRuntimeVQEncoder(RBLNPytorchRuntime):
     def encode(self, x: torch.FloatTensor, **kwargs) -> torch.FloatTensor:
         h = self.forward(x.contiguous())
@@ -91,7 +128,7 @@ class RBLNRuntimeVQDecoder(RBLNPytorchRuntime):
 
 
 class _VQEncoder(torch.nn.Module):
-    def __init__(self, vq_model: VQModel):
+    def __init__(self, vq_model: "VQModel"):
         super().__init__()
         self.vq_model = vq_model
 
@@ -106,7 +143,7 @@ class _VQEncoder(torch.nn.Module):
 
 
 class _VQDecoder(torch.nn.Module):
-    def __init__(self, vq_model: VQModel):
+    def __init__(self, vq_model: "VQModel"):
         super().__init__()
         self.vq_model = vq_model
 

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

@@ -35,6 +35,17 @@ logger = get_logger(__name__)
 
 
 class RBLNVQModel(RBLNModel):
+    """
+    RBLN implementation of VQModel for diffusion models.
+
+    This model is used to accelerate VQModel models from diffusers library on RBLN NPUs.
+    It can be configured to include both encoder and decoder, or just the decoder part for latent-to-image
+    conversion.
+
+    This class inherits from [`RBLNModel`]. Check the superclass documentation for the generic methods
+    the library implements for all its models.
+    """
+
     auto_model_class = VQModel
     config_name = "config.json"
     hf_library_name = "diffusers"
@@ -67,7 +78,12 @@ class RBLNVQModel(RBLNModel):
 
             wrapped_model.eval()
 
-            compiled_models[model_name] = cls.compile(wrapped_model, rbln_compile_config=rbln_config.compile_cfgs[i])
+            compiled_models[model_name] = cls.compile(
+                wrapped_model,
+                rbln_compile_config=rbln_config.compile_cfgs[i],
+                create_runtimes=rbln_config.create_runtimes,
+                device=rbln_config.device_map[model_name],
+            )
 
         return compiled_models