optimum-rbln 0.9.3rc0__py3-none-any.whl → 0.9.4a2__py3-none-any.whl

optimum/rbln/__init__.py

@@ -186,12 +186,16 @@ _import_structure = {
     "diffusers": [
         "RBLNAutoencoderKL",
         "RBLNAutoencoderKLConfig",
+        "RBLNAutoencoderKLTemporalDecoder",
+        "RBLNAutoencoderKLTemporalDecoderConfig",
         "RBLNAutoencoderKLCosmos",
         "RBLNAutoencoderKLCosmosConfig",
         "RBLNAutoPipelineForImage2Image",
         "RBLNAutoPipelineForInpainting",
         "RBLNAutoPipelineForText2Image",
         "RBLNControlNetModel",
+        "RBLNUNetSpatioTemporalConditionModel",
+        "RBLNStableVideoDiffusionPipeline",
         "RBLNControlNetModelConfig",
         "RBLNCosmosTextToWorldPipeline",
         "RBLNCosmosVideoToWorldPipeline",
@@ -250,6 +254,8 @@ _import_structure = {
         "RBLNUNet2DConditionModelConfig",
         "RBLNVQModel",
         "RBLNVQModelConfig",
+        "RBLNUNetSpatioTemporalConditionModelConfig",
+        "RBLNStableVideoDiffusionPipelineConfig",
     ],
 }
 
@@ -260,6 +266,8 @@ if TYPE_CHECKING:
         RBLNAutoencoderKLConfig,
         RBLNAutoencoderKLCosmos,
         RBLNAutoencoderKLCosmosConfig,
+        RBLNAutoencoderKLTemporalDecoder,
+        RBLNAutoencoderKLTemporalDecoderConfig,
         RBLNAutoPipelineForImage2Image,
         RBLNAutoPipelineForInpainting,
         RBLNAutoPipelineForText2Image,
@@ -318,8 +326,12 @@ if TYPE_CHECKING:
         RBLNStableDiffusionXLInpaintPipelineConfig,
         RBLNStableDiffusionXLPipeline,
         RBLNStableDiffusionXLPipelineConfig,
+        RBLNStableVideoDiffusionPipeline,
+        RBLNStableVideoDiffusionPipelineConfig,
         RBLNUNet2DConditionModel,
         RBLNUNet2DConditionModelConfig,
+        RBLNUNetSpatioTemporalConditionModel,
+        RBLNUNetSpatioTemporalConditionModelConfig,
         RBLNVQModel,
         RBLNVQModelConfig,
     )

optimum/rbln/__version__.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.9.3rc0'
-__version_tuple__ = version_tuple = (0, 9, 3, 'rc0')
+__version__ = version = '0.9.4a2'
+__version_tuple__ = version_tuple = (0, 9, 4, 'a2')
 
 __commit_id__ = commit_id = None

optimum/rbln/configuration_utils.py

@@ -24,7 +24,7 @@ import torch
 from packaging.version import Version
 
 from .__version__ import __version__
-from .utils.depreacate_utils import warn_deprecated_npu
+from .utils.deprecation import warn_deprecated_npu
 from .utils.logging import get_logger
 from .utils.runtime_utils import ContextRblnConfig
 
@@ -117,9 +117,14 @@ class RBLNCompileConfig:
         return self
 
     def get_dummy_inputs(
-        self, fill=0, static_tensors: Dict[str, torch.Tensor] = {}, meta_tensor_names: List[str] = []
+        self,
+        fill=0,
+        static_tensors: Optional[Dict[str, torch.Tensor]] = None,
+        meta_tensor_names: Optional[List[str]] = None,
     ):
         dummy = []
+        static_tensors = static_tensors if static_tensors is not None else {}
+        meta_tensor_names = meta_tensor_names if meta_tensor_names is not None else []
         for name, shape, dtype in self.input_info:
             if name in static_tensors:
                 tensor = static_tensors[name]
@@ -255,7 +260,7 @@ class RBLNAutoConfig:
     def load(
         path: str,
         passed_rbln_config: Optional["RBLNModelConfig"] = None,
-        kwargs: Optional[Dict[str, Any]] = {},
+        kwargs: Optional[Dict[str, Any]] = None,
         return_unused_kwargs: bool = False,
     ) -> Union["RBLNModelConfig", Tuple["RBLNModelConfig", Dict[str, Any]]]:
         """
@@ -269,6 +274,8 @@ class RBLNAutoConfig:
         Returns:
             RBLNModelConfig: The loaded RBLNModelConfig.
         """
+        if kwargs is None:
+            kwargs = {}
         cls, config_file = load_config(path)
 
         rbln_keys = [key for key in kwargs.keys() if key.startswith("rbln_")]
@@ -528,6 +535,7 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
     ]
     submodules: List[str] = []
     subclass_non_save_attributes = []
+    _allow_no_compile_cfgs = False
 
     def initialize_submodule_config(
         self,
@@ -654,7 +662,7 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
         timeout: Optional[int] = None,
         optimum_rbln_version: Optional[str] = None,
         _torch_dtype: Optional[str] = None,
-        _compile_cfgs: List[RBLNCompileConfig] = [],
+        _compile_cfgs: Optional[List[RBLNCompileConfig]] = None,
         *,
         optimize_host_memory: Optional[bool] = None,
         **kwargs: Any,
@@ -707,7 +715,8 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
         if self.optimum_rbln_version is None:
             self.optimum_rbln_version = __version__
 
-        self._compile_cfgs: List[RBLNCompileConfig] = _compile_cfgs
+        compile_cfgs = _compile_cfgs if _compile_cfgs is not None else []
+        self._compile_cfgs: List[RBLNCompileConfig] = compile_cfgs
 
         if not isinstance(self._compile_cfgs, list):
             raise ValueError("`compile_cfgs` must be a list of `RBLNCompileConfig`.")
@@ -808,7 +817,8 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
             or len(self._compile_cfgs) == 0
             or not all(isinstance(cfg, RBLNCompileConfig) for cfg in self._compile_cfgs)
         ):
-            raise RuntimeError("`compile_cfgs` must be set before freezing.")
+            if not self._allow_no_compile_cfgs:
+                raise RuntimeError("`compile_cfgs` must contain at least one `RBLNCompileConfig` before freezing.")
 
         for submodule_name in self.submodules:
             submodule_config = getattr(self, submodule_name, None)

optimum/rbln/diffusers/__init__.py

@@ -57,6 +57,9 @@ _import_structure = {
         "RBLNSD3Transformer2DModelConfig",
         "RBLNUNet2DConditionModelConfig",
         "RBLNVQModelConfig",
+        "RBLNUNetSpatioTemporalConditionModelConfig",
+        "RBLNStableVideoDiffusionPipelineConfig",
+        "RBLNAutoencoderKLTemporalDecoderConfig",
     ],
     "pipelines": [
         "RBLNAutoPipelineForImage2Image",
@@ -86,14 +89,17 @@ _import_structure = {
         "RBLNStableDiffusion3Pipeline",
         "RBLNStableDiffusion3Img2ImgPipeline",
         "RBLNStableDiffusion3InpaintPipeline",
+        "RBLNStableVideoDiffusionPipeline",
     ],
     "models": [
         "RBLNAutoencoderKL",
         "RBLNAutoencoderKLCosmos",
         "RBLNUNet2DConditionModel",
+        "RBLNUNetSpatioTemporalConditionModel",
         "RBLNControlNetModel",
         "RBLNCosmosTransformer3DModel",
         "RBLNSD3Transformer2DModel",
+        "RBLNAutoencoderKLTemporalDecoder",
         "RBLNPriorTransformer",
         "RBLNVQModel",
     ],
@@ -106,6 +112,7 @@ if TYPE_CHECKING:
     from .configurations import (
         RBLNAutoencoderKLConfig,
         RBLNAutoencoderKLCosmosConfig,
+        RBLNAutoencoderKLTemporalDecoderConfig,
         RBLNControlNetModelConfig,
         RBLNCosmosTextToWorldPipelineConfig,
         RBLNCosmosTransformer3DModelConfig,
@@ -132,18 +139,22 @@ if TYPE_CHECKING:
         RBLNStableDiffusionXLImg2ImgPipelineConfig,
         RBLNStableDiffusionXLInpaintPipelineConfig,
         RBLNStableDiffusionXLPipelineConfig,
+        RBLNStableVideoDiffusionPipelineConfig,
         RBLNUNet2DConditionModelConfig,
+        RBLNUNetSpatioTemporalConditionModelConfig,
         RBLNVQModelConfig,
     )
     from .modeling_diffusers import RBLNDiffusionMixin
     from .models import (
         RBLNAutoencoderKL,
         RBLNAutoencoderKLCosmos,
+        RBLNAutoencoderKLTemporalDecoder,
         RBLNControlNetModel,
         RBLNCosmosTransformer3DModel,
         RBLNPriorTransformer,
         RBLNSD3Transformer2DModel,
         RBLNUNet2DConditionModel,
+        RBLNUNetSpatioTemporalConditionModel,
         RBLNVQModel,
     )
     from .pipelines import (
@@ -174,6 +185,7 @@ if TYPE_CHECKING:
         RBLNStableDiffusionXLImg2ImgPipeline,
         RBLNStableDiffusionXLInpaintPipeline,
         RBLNStableDiffusionXLPipeline,
+        RBLNStableVideoDiffusionPipeline,
     )
 else:
     import sys

optimum/rbln/diffusers/configurations/__init__.py

@@ -1,11 +1,13 @@
 from .models import (
     RBLNAutoencoderKLConfig,
     RBLNAutoencoderKLCosmosConfig,
+    RBLNAutoencoderKLTemporalDecoderConfig,
     RBLNControlNetModelConfig,
     RBLNCosmosTransformer3DModelConfig,
     RBLNPriorTransformerConfig,
     RBLNSD3Transformer2DModelConfig,
     RBLNUNet2DConditionModelConfig,
+    RBLNUNetSpatioTemporalConditionModelConfig,
     RBLNVQModelConfig,
 )
 from .pipelines import (
@@ -31,4 +33,5 @@ from .pipelines import (
     RBLNStableDiffusionXLImg2ImgPipelineConfig,
     RBLNStableDiffusionXLInpaintPipelineConfig,
     RBLNStableDiffusionXLPipelineConfig,
+    RBLNStableVideoDiffusionPipelineConfig,
 )

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

@@ -1,8 +1,10 @@
 from .configuration_autoencoder_kl import RBLNAutoencoderKLConfig
 from .configuration_autoencoder_kl_cosmos import RBLNAutoencoderKLCosmosConfig
+from .configuration_autoencoder_kl_temporal_decoder import RBLNAutoencoderKLTemporalDecoderConfig
 from .configuration_controlnet import RBLNControlNetModelConfig
 from .configuration_prior_transformer import RBLNPriorTransformerConfig
 from .configuration_transformer_cosmos import RBLNCosmosTransformer3DModelConfig
 from .configuration_transformer_sd3 import RBLNSD3Transformer2DModelConfig
 from .configuration_unet_2d_condition import RBLNUNet2DConditionModelConfig
+from .configuration_unet_spatio_temporal_condition import RBLNUNetSpatioTemporalConditionModelConfig
 from .configuration_vq_model import RBLNVQModelConfig

optimum/rbln/diffusers/configurations/models/configuration_autoencoder_kl_temporal_decoder.py

@@ -0,0 +1,67 @@
+# 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 Any, Optional, Tuple
+
+from ....configuration_utils import RBLNModelConfig
+
+
+class RBLNAutoencoderKLTemporalDecoderConfig(RBLNModelConfig):
+    def __init__(
+        self,
+        batch_size: Optional[int] = None,
+        sample_size: Optional[Tuple[int, int]] = None,
+        uses_encoder: Optional[bool] = None,
+        num_frames: Optional[int] = None,
+        decode_chunk_size: Optional[int] = None,
+        vae_scale_factor: Optional[float] = None,
+        **kwargs: Any,
+    ):
+        """
+        Args:
+            batch_size (Optional[int]): The batch size for inference. Defaults to 1.
+            sample_size (Optional[Tuple[int, int]]): The spatial dimensions (height, width) of the input/output images.
+                If an integer is provided, it's used for both height and width.
+            uses_encoder (Optional[bool]): Whether to include the encoder part of the VAE in the model.
+                When False, only the decoder is used (for latent-to-image conversion).
+            num_frames (Optional[int]): The number of frames in the generated video.
+            decode_chunk_size (Optional[int]): The number of frames to decode at once during VAE decoding.
+                Useful for managing memory usage during video generation.
+            vae_scale_factor (Optional[float]): The scaling factor between pixel space and latent space.
+                Determines how much smaller the latent representations are compared to the original images.
+            kwargs: Additional arguments passed to the parent RBLNModelConfig.
+
+        Raises:
+            ValueError: If batch_size is not a positive integer.
+        """
+        super().__init__(**kwargs)
+        self.batch_size = batch_size or 1
+        if not isinstance(self.batch_size, int) or self.batch_size < 0:
+            raise ValueError(f"batch_size must be a positive integer, got {self.batch_size}")
+
+        self.uses_encoder = uses_encoder
+        self.num_frames = num_frames
+        self.decode_chunk_size = decode_chunk_size
+        self.vae_scale_factor = vae_scale_factor
+        self.sample_size = sample_size
+        if isinstance(sample_size, int):
+            self.sample_size = (sample_size, sample_size)
+
+    @property
+    def image_size(self):
+        return self.sample_size
+
+    @property
+    def latent_sample_size(self):
+        return (self.image_size[0] // self.vae_scale_factor, self.image_size[1] // self.vae_scale_factor)

optimum/rbln/diffusers/configurations/models/configuration_unet_spatio_temporal_condition.py

@@ -0,0 +1,59 @@
+# 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 Any, Optional, Tuple
+
+from ....configuration_utils import RBLNModelConfig
+
+
+class RBLNUNetSpatioTemporalConditionModelConfig(RBLNModelConfig):
+    subclass_non_save_attributes = ["_batch_size_is_specified"]
+
+    def __init__(
+        self,
+        batch_size: Optional[int] = None,
+        sample_size: Optional[Tuple[int, int]] = None,
+        in_features: Optional[int] = None,
+        num_frames: Optional[int] = None,
+        **kwargs: Any,
+    ):
+        """
+        Args:
+            batch_size (Optional[int]): The batch size for inference. Defaults to 1.
+            sample_size (Optional[Tuple[int, int]]): The spatial dimensions (height, width) of the generated samples.
+                If an integer is provided, it's used for both height and width.
+            in_features (Optional[int]): Number of input features for the model.
+            num_frames (Optional[int]): The number of frames in the generated video.
+            kwargs: Additional arguments passed to the parent RBLNModelConfig.
+
+        Raises:
+            ValueError: If batch_size is not a positive integer.
+        """
+        super().__init__(**kwargs)
+        self._batch_size_is_specified = batch_size is not None
+
+        self.batch_size = batch_size or 1
+        if not isinstance(self.batch_size, int) or self.batch_size < 0:
+            raise ValueError(f"batch_size must be a positive integer, got {self.batch_size}")
+
+        self.in_features = in_features
+        self.num_frames = num_frames
+
+        self.sample_size = sample_size
+        if isinstance(sample_size, int):
+            self.sample_size = (sample_size, sample_size)
+
+    @property
+    def batch_size_is_specified(self):
+        return self._batch_size_is_specified

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

@@ -29,3 +29,6 @@ from .configuration_stable_diffusion_xl import (
     RBLNStableDiffusionXLInpaintPipelineConfig,
     RBLNStableDiffusionXLPipelineConfig,
 )
+from .configuration_stable_video_diffusion import (
+    RBLNStableVideoDiffusionPipelineConfig,
+)

optimum/rbln/diffusers/configurations/pipelines/configuration_stable_video_diffusion.py

@@ -0,0 +1,114 @@
+# 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 Any, Optional
+
+from ....configuration_utils import RBLNModelConfig
+from ....transformers import RBLNCLIPVisionModelWithProjectionConfig
+from ..models import RBLNAutoencoderKLTemporalDecoderConfig, RBLNUNetSpatioTemporalConditionModelConfig
+
+
+class RBLNStableVideoDiffusionPipelineConfig(RBLNModelConfig):
+    submodules = ["image_encoder", "unet", "vae"]
+    _vae_uses_encoder = True
+
+    def __init__(
+        self,
+        image_encoder: Optional[RBLNCLIPVisionModelWithProjectionConfig] = None,
+        unet: Optional[RBLNUNetSpatioTemporalConditionModelConfig] = None,
+        vae: Optional[RBLNAutoencoderKLTemporalDecoderConfig] = None,
+        *,
+        batch_size: Optional[int] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_frames: Optional[int] = None,
+        decode_chunk_size: Optional[int] = None,
+        guidance_scale: Optional[float] = None,
+        **kwargs: Any,
+    ):
+        """
+        Args:
+            image_encoder (Optional[RBLNCLIPVisionModelWithProjectionConfig]): Configuration for the image encoder component.
+                Initialized as RBLNCLIPVisionModelWithProjectionConfig if not provided.
+            unet (Optional[RBLNUNetSpatioTemporalConditionModelConfig]): Configuration for the UNet model component.
+                Initialized as RBLNUNetSpatioTemporalConditionModelConfig if not provided.
+            vae (Optional[RBLNAutoencoderKLTemporalDecoderConfig]): Configuration for the VAE model component.
+                Initialized as RBLNAutoencoderKLTemporalDecoderConfig if not provided.
+            batch_size (Optional[int]): Batch size for inference, applied to all submodules.
+            height (Optional[int]): Height of the generated images.
+            width (Optional[int]): Width of the generated images.
+            num_frames (Optional[int]): The number of frames in the generated video.
+            decode_chunk_size (Optional[int]): The number of frames to decode at once during VAE decoding.
+                Useful for managing memory usage during video generation.
+            guidance_scale (Optional[float]): Scale for classifier-free guidance.
+            kwargs: Additional arguments passed to the parent RBLNModelConfig.
+
+        Raises:
+            ValueError: If both image_size and height/width are provided.
+
+        Note:
+            When guidance_scale > 1.0, the UNet batch size is automatically doubled to
+            accommodate classifier-free guidance.
+        """
+        super().__init__(**kwargs)
+        if height is not None and width is not None:
+            image_size = (height, width)
+        else:
+            # Get default image size from original class to set UNet, VAE image size
+            height = self.get_default_values_for_original_cls("__call__", ["height"])["height"]
+            width = self.get_default_values_for_original_cls("__call__", ["width"])["width"]
+            image_size = (height, width)
+
+        self.image_encoder = self.initialize_submodule_config(
+            image_encoder, cls_name="RBLNCLIPVisionModelWithProjectionConfig", batch_size=batch_size
+        )
+        self.unet = self.initialize_submodule_config(
+            unet,
+            cls_name="RBLNUNetSpatioTemporalConditionModelConfig",
+            num_frames=num_frames,
+        )
+        self.vae = self.initialize_submodule_config(
+            vae,
+            cls_name="RBLNAutoencoderKLTemporalDecoderConfig",
+            batch_size=batch_size,
+            num_frames=num_frames,
+            decode_chunk_size=decode_chunk_size,
+            uses_encoder=self.__class__._vae_uses_encoder,
+            sample_size=image_size,  # image size is equal to sample size in vae
+        )
+
+        # Get default guidance scale from original class to set UNet batch size
+        if guidance_scale is None:
+            guidance_scale = self.get_default_values_for_original_cls("__call__", ["max_guidance_scale"])[
+                "max_guidance_scale"
+            ]
+
+        if not self.unet.batch_size_is_specified:
+            do_classifier_free_guidance = guidance_scale > 1.0
+            if do_classifier_free_guidance:
+                self.unet.batch_size = self.image_encoder.batch_size * 2
+            else:
+                self.unet.batch_size = self.image_encoder.batch_size
+
+    @property
+    def batch_size(self):
+        return self.vae.batch_size
+
+    @property
+    def sample_size(self):
+        return self.unet.sample_size
+
+    @property
+    def image_size(self):
+        return self.vae.sample_size

optimum/rbln/diffusers/modeling_diffusers.py

@@ -136,7 +136,7 @@ class RBLNDiffusionMixin:
         *,
         export: bool = None,
         model_save_dir: Optional[PathLike] = None,
-        rbln_config: Dict[str, Any] = {},
+        rbln_config: Optional[Dict[str, Any]] = None,
         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,

optimum/rbln/diffusers/models/__init__.py

@@ -22,9 +22,11 @@ _import_structure = {
         "RBLNAutoencoderKL",
         "RBLNAutoencoderKLCosmos",
         "RBLNVQModel",
+        "RBLNAutoencoderKLTemporalDecoder",
     ],
     "unets": [
         "RBLNUNet2DConditionModel",
+        "RBLNUNetSpatioTemporalConditionModel",
     ],
     "controlnet": ["RBLNControlNetModel"],
     "transformers": [
@@ -35,10 +37,22 @@ _import_structure = {
 }
 
 if TYPE_CHECKING:
-    from .autoencoders import RBLNAutoencoderKL, RBLNAutoencoderKLCosmos, RBLNVQModel
+    from .autoencoders import (
+        RBLNAutoencoderKL,
+        RBLNAutoencoderKLCosmos,
+        RBLNAutoencoderKLTemporalDecoder,
+        RBLNVQModel,
+    )
     from .controlnet import RBLNControlNetModel
-    from .transformers import RBLNCosmosTransformer3DModel, RBLNPriorTransformer, RBLNSD3Transformer2DModel
-    from .unets import RBLNUNet2DConditionModel
+    from .transformers import (
+        RBLNCosmosTransformer3DModel,
+        RBLNPriorTransformer,
+        RBLNSD3Transformer2DModel,
+    )
+    from .unets import (
+        RBLNUNet2DConditionModel,
+        RBLNUNetSpatioTemporalConditionModel,
+    )
 else:
     import sys
 

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

@@ -14,4 +14,5 @@
 
 from .autoencoder_kl import RBLNAutoencoderKL
 from .autoencoder_kl_cosmos import RBLNAutoencoderKLCosmos
+from .autoencoder_kl_temporal_decoder import RBLNAutoencoderKLTemporalDecoder
 from .vq_model import RBLNVQModel

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

@@ -68,7 +68,7 @@ class RBLNAutoencoderKLCosmos(RBLNModel):
         self.image_size = self.rbln_config.image_size
 
     @classmethod
-    def wrap_model_if_needed(
+    def _wrap_model_if_needed(
         cls, model: torch.nn.Module, rbln_config: RBLNAutoencoderKLCosmosConfig
     ) -> torch.nn.Module:
         decoder_model = _VAECosmosDecoder(model)
@@ -98,7 +98,7 @@ class RBLNAutoencoderKLCosmos(RBLNModel):
 
             compiled_models = {}
             if rbln_config.uses_encoder:
-                encoder_model, decoder_model = cls.wrap_model_if_needed(model, rbln_config)
+                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],
@@ -107,7 +107,7 @@ class RBLNAutoencoderKLCosmos(RBLNModel):
                 )
                 compiled_models["encoder"] = enc_compiled_model
             else:
-                decoder_model = cls.wrap_model_if_needed(model, rbln_config)
+                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],