optimum-rbln 0.8.1a5__py3-none-any.whl → 0.8.1a7__py3-none-any.whl

optimum/rbln/__init__.py

@@ -138,8 +138,17 @@ _import_structure = {
     "diffusers": [
         "RBLNAutoencoderKL",
         "RBLNAutoencoderKLConfig",
+        "RBLNAutoencoderKLCosmos",
+        "RBLNAutoencoderKLCosmosConfig",
         "RBLNControlNetModel",
         "RBLNControlNetModelConfig",
+        "RBLNCosmosTextToWorldPipeline",
+        "RBLNCosmosVideoToWorldPipeline",
+        "RBLNCosmosTextToWorldPipelineConfig",
+        "RBLNCosmosVideoToWorldPipelineConfig",
+        "RBLNCosmosSafetyChecker",
+        "RBLNCosmosTransformer3DModel",
+        "RBLNCosmosTransformer3DModelConfig",
         "RBLNDiffusionMixin",
         "RBLNKandinskyV22CombinedPipeline",
         "RBLNKandinskyV22CombinedPipelineConfig",
@@ -202,8 +211,17 @@ if TYPE_CHECKING:
     from .diffusers import (
         RBLNAutoencoderKL,
         RBLNAutoencoderKLConfig,
+        RBLNAutoencoderKLCosmos,
+        RBLNAutoencoderKLCosmosConfig,
         RBLNControlNetModel,
         RBLNControlNetModelConfig,
+        RBLNCosmosSafetyChecker,
+        RBLNCosmosTextToWorldPipeline,
+        RBLNCosmosTextToWorldPipelineConfig,
+        RBLNCosmosTransformer3DModel,
+        RBLNCosmosTransformer3DModelConfig,
+        RBLNCosmosVideoToWorldPipeline,
+        RBLNCosmosVideoToWorldPipelineConfig,
         RBLNDiffusionMixin,
         RBLNKandinskyV22CombinedPipeline,
         RBLNKandinskyV22CombinedPipelineConfig,

optimum/rbln/__version__.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.8.1a5'
-__version_tuple__ = version_tuple = (0, 8, 1, 'a5')
+__version__ = version = '0.8.1a7'
+__version_tuple__ = version_tuple = (0, 8, 1, 'a7')

optimum/rbln/diffusers/__init__.py

@@ -18,14 +18,21 @@ from diffusers.pipelines.pipeline_utils import ALL_IMPORTABLE_CLASSES, LOADABLE_
 from transformers.utils import _LazyModule
 
 
-LOADABLE_CLASSES["optimum.rbln"] = {"RBLNBaseModel": ["save_pretrained", "from_pretrained"]}
+LOADABLE_CLASSES["optimum.rbln"] = {
+    "RBLNBaseModel": ["save_pretrained", "from_pretrained"],
+    "RBLNCosmosSafetyChecker": ["save_pretrained", "from_pretrained"],
+}
 ALL_IMPORTABLE_CLASSES.update(LOADABLE_CLASSES["optimum.rbln"])
 
 
 _import_structure = {
     "configurations": [
         "RBLNAutoencoderKLConfig",
+        "RBLNAutoencoderKLCosmosConfig",
         "RBLNControlNetModelConfig",
+        "RBLNCosmosTextToWorldPipelineConfig",
+        "RBLNCosmosVideoToWorldPipelineConfig",
+        "RBLNCosmosTransformer3DModelConfig",
         "RBLNKandinskyV22CombinedPipelineConfig",
         "RBLNKandinskyV22Img2ImgCombinedPipelineConfig",
         "RBLNKandinskyV22Img2ImgPipelineConfig",
@@ -52,6 +59,9 @@ _import_structure = {
         "RBLNVQModelConfig",
     ],
     "pipelines": [
+        "RBLNCosmosTextToWorldPipeline",
+        "RBLNCosmosVideoToWorldPipeline",
+        "RBLNCosmosSafetyChecker",
         "RBLNKandinskyV22CombinedPipeline",
         "RBLNKandinskyV22Img2ImgCombinedPipeline",
         "RBLNKandinskyV22InpaintCombinedPipeline",
@@ -76,8 +86,10 @@ _import_structure = {
     ],
     "models": [
         "RBLNAutoencoderKL",
+        "RBLNAutoencoderKLCosmos",
         "RBLNUNet2DConditionModel",
         "RBLNControlNetModel",
+        "RBLNCosmosTransformer3DModel",
         "RBLNSD3Transformer2DModel",
         "RBLNPriorTransformer",
         "RBLNVQModel",
@@ -90,7 +102,11 @@ _import_structure = {
 if TYPE_CHECKING:
     from .configurations import (
         RBLNAutoencoderKLConfig,
+        RBLNAutoencoderKLCosmosConfig,
         RBLNControlNetModelConfig,
+        RBLNCosmosTextToWorldPipelineConfig,
+        RBLNCosmosTransformer3DModelConfig,
+        RBLNCosmosVideoToWorldPipelineConfig,
         RBLNKandinskyV22CombinedPipelineConfig,
         RBLNKandinskyV22Img2ImgCombinedPipelineConfig,
         RBLNKandinskyV22Img2ImgPipelineConfig,
@@ -120,12 +136,16 @@ if TYPE_CHECKING:
     from .models import (
         RBLNAutoencoderKL,
         RBLNControlNetModel,
+        RBLNCosmosTransformer3DModel,
         RBLNPriorTransformer,
         RBLNSD3Transformer2DModel,
         RBLNUNet2DConditionModel,
         RBLNVQModel,
     )
     from .pipelines import (
+        RBLNCosmosSafetyChecker,
+        RBLNCosmosTextToWorldPipeline,
+        RBLNCosmosVideoToWorldPipeline,
         RBLNKandinskyV22CombinedPipeline,
         RBLNKandinskyV22Img2ImgCombinedPipeline,
         RBLNKandinskyV22Img2ImgPipeline,

optimum/rbln/diffusers/configurations/__init__.py

@@ -1,12 +1,16 @@
 from .models import (
     RBLNAutoencoderKLConfig,
+    RBLNAutoencoderKLCosmosConfig,
     RBLNControlNetModelConfig,
+    RBLNCosmosTransformer3DModelConfig,
     RBLNPriorTransformerConfig,
     RBLNSD3Transformer2DModelConfig,
     RBLNUNet2DConditionModelConfig,
     RBLNVQModelConfig,
 )
 from .pipelines import (
+    RBLNCosmosTextToWorldPipelineConfig,
+    RBLNCosmosVideoToWorldPipelineConfig,
     RBLNKandinskyV22CombinedPipelineConfig,
     RBLNKandinskyV22Img2ImgCombinedPipelineConfig,
     RBLNKandinskyV22Img2ImgPipelineConfig,

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

@@ -1,5 +1,7 @@
 from .configuration_autoencoder_kl import RBLNAutoencoderKLConfig
+from .configuration_autoencoder_kl_cosmos import RBLNAutoencoderKLCosmosConfig
 from .configuration_controlnet import RBLNControlNetModelConfig
+from .configuration_cosmos_transformer import RBLNCosmosTransformer3DModelConfig
 from .configuration_prior_transformer import RBLNPriorTransformerConfig
 from .configuration_transformer_sd3 import RBLNSD3Transformer2DModelConfig
 from .configuration_unet_2d_condition import RBLNUNet2DConditionModelConfig

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

@@ -0,0 +1,82 @@
+# 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 Optional
+
+from ....configuration_utils import RBLNModelConfig
+from ....utils.logging import get_logger
+
+
+logger = get_logger(__name__)
+
+
+class RBLNAutoencoderKLCosmosConfig(RBLNModelConfig):
+    def __init__(
+        self,
+        batch_size: Optional[int] = None,
+        uses_encoder: Optional[bool] = None,
+        num_frames: Optional[int] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_channels_latents: Optional[int] = None,
+        vae_scale_factor_temporal: Optional[int] = None,
+        vae_scale_factor_spatial: Optional[int] = None,
+        use_slicing: Optional[bool] = None,
+        **kwargs,
+    ):
+        """
+        Args:
+            batch_size (Optional[int]): The batch size for inference. Defaults to 1.
+            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-video conversion).
+            num_frames (Optional[int]): The number of frames in the generated video. Defaults to 121.
+            height (Optional[int]): The height in pixels of the generated video. Defaults to 704.
+            width (Optional[int]): The width in pixels of the generated video. Defaults to 1280.
+            num_channels_latents (Optional[int]): The number of channels in latent space.
+            vae_scale_factor_temporal (Optional[int]): The scaling factor between time space and latent space.
+                Determines how much shorter the latent representations are compared to the original videos.
+            vae_scale_factor_spatial (Optional[int]): The scaling factor between pixel space and latent space.
+                Determines how much smaller the latent representations are compared to the original videos.
+            use_slicing (Optional[Bool]): Enable sliced VAE encoding and decoding.
+                If True, the VAE will split the input tensor in slices to compute encoding or decoding in several steps.
+            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+
+        Raises:
+            ValueError: If batch_size is not a positive integer.
+        """
+        super().__init__(**kwargs)
+        # Since the Cosmos VAE Decoder already requires approximately 7.9 GiB of memory,
+        # Optimum-rbln cannot execute this model on RBLN-CA12 when the batch size > 1.
+        # However, the Cosmos VAE Decoder propose batch slicing when the batch size is greater than 1,
+        # Optimum-rbln utilize this method by compiling with batch_size=1 to enable batch slicing.
+        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}")
+        elif self.batch_size > 1:
+            logger.warning("The batch size of Cosmos VAE Decoder will be explicitly 1 for memory efficiency.")
+            self.batch_size = 1
+
+        self.uses_encoder = uses_encoder
+        self.num_frames = num_frames or 121
+        self.height = height or 704
+        self.width = width or 1280
+
+        self.num_channels_latents = num_channels_latents
+        self.vae_scale_factor_temporal = vae_scale_factor_temporal
+        self.vae_scale_factor_spatial = vae_scale_factor_spatial
+        self.use_slicing = use_slicing or False
+
+    @property
+    def image_size(self):
+        return (self.height, self.width)

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

@@ -0,0 +1,68 @@
+# 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 Optional
+
+from ....configuration_utils import RBLNModelConfig
+
+
+class RBLNCosmosTransformer3DModelConfig(RBLNModelConfig):
+    def __init__(
+        self,
+        batch_size: Optional[int] = None,
+        num_frames: Optional[int] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        fps: Optional[int] = None,
+        max_seq_len: Optional[int] = None,
+        embedding_dim: Optional[int] = None,
+        num_channels_latents: Optional[int] = None,
+        num_latent_frames: Optional[int] = None,
+        latent_height: Optional[int] = None,
+        latent_width: Optional[int] = None,
+        **kwargs,
+    ):
+        """
+        Args:
+            batch_size (Optional[int]): The batch size for inference. Defaults to 1.
+            num_frames (Optional[int]): The number of frames in the generated video. Defaults to 121.
+            height (Optional[int]): The height in pixels of the generated video. Defaults to 704.
+            width (Optional[int]): The width in pixels of the generated video. Defaults to 1280.
+            fps (Optional[int]): The frames per second of the generated video.  Defaults to 30.
+            max_seq_len (Optional[int]): Maximum sequence length of prompt embeds.
+            embedding_dim (Optional[int]): Embedding vector dimension of prompt embeds.
+            num_channels_latents (Optional[int]): The number of channels in latent space.
+            latent_height (Optional[int]): The height in pixels in latent space.
+            latent_width (Optional[int]): The width in pixels in latent space.
+            **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
+        self.num_frames = num_frames or 121
+        self.height = height or 704
+        self.width = width or 1280
+        self.fps = fps or 30
+
+        self.max_seq_len = max_seq_len
+        self.num_channels_latents = num_channels_latents
+        self.num_latent_frames = num_latent_frames
+        self.latent_height = latent_height
+        self.latent_width = latent_width
+        self.embedding_dim = embedding_dim
+
+        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}")

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

@@ -4,6 +4,7 @@ from .configuration_controlnet import (
     RBLNStableDiffusionXLControlNetImg2ImgPipelineConfig,
     RBLNStableDiffusionXLControlNetPipelineConfig,
 )
+from .configuration_cosmos import RBLNCosmosTextToWorldPipelineConfig, RBLNCosmosVideoToWorldPipelineConfig
 from .configuration_kandinsky2_2 import (
     RBLNKandinskyV22CombinedPipelineConfig,
     RBLNKandinskyV22Img2ImgCombinedPipelineConfig,

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

@@ -20,10 +20,6 @@ from ..models import RBLNAutoencoderKLConfig, RBLNControlNetModelConfig, RBLNUNe
 
 
 class RBLNStableDiffusionControlNetPipelineBaseConfig(RBLNModelConfig):
-    """
-    Base configuration for Stable Diffusion ControlNet pipelines.
-    """
-
     submodules = ["text_encoder", "unet", "vae", "controlnet"]
     _vae_uses_encoder = False
 

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

@@ -0,0 +1,110 @@
+# 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 Optional
+
+from ....configuration_utils import RBLNModelConfig
+from ....transformers import RBLNT5EncoderModelConfig
+from ....utils.logging import get_logger
+from ...pipelines.cosmos.cosmos_guardrail import RBLNCosmosSafetyCheckerConfig
+from ..models import RBLNAutoencoderKLCosmosConfig, RBLNCosmosTransformer3DModelConfig
+
+
+logger = get_logger(__name__)
+
+
+class _RBLNCosmosPipelineBaseConfig(RBLNModelConfig):
+    submodules = ["text_encoder", "transformer", "vae", "safety_checker"]
+    _vae_uses_encoder = False
+
+    def __init__(
+        self,
+        text_encoder: Optional[RBLNT5EncoderModelConfig] = None,
+        transformer: Optional[RBLNCosmosTransformer3DModelConfig] = None,
+        vae: Optional[RBLNAutoencoderKLCosmosConfig] = None,
+        safety_checker: Optional[RBLNCosmosSafetyCheckerConfig] = None,
+        *,
+        batch_size: Optional[int] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_frames: Optional[int] = None,
+        fps: Optional[int] = None,
+        max_seq_len: Optional[int] = None,
+        **kwargs,
+    ):
+        """
+        Args:
+            text_encoder (Optional[RBLNT5EncoderModelConfig]): Configuration for the text encoder component.
+                Initialized as RBLNT5EncoderModelConfig if not provided.
+            transformer (Optional[RBLNCosmosTransformer3DModelConfig]): Configuration for the UNet model component.
+                Initialized as RBLNCosmosTransformer3DModelConfig if not provided.
+            vae (Optional[RBLNAutoencoderKLCosmosConfig]): Configuration for the VAE model component.
+                Initialized as RBLNAutoencoderKLCosmosConfig if not provided.
+            safety_checker (Optional[RBLNCosmosSafetyCheckerConfig]): Configuration for the safety checker component.
+                Initialized as RBLNCosmosSafetyCheckerConfig if not provided.
+            batch_size (Optional[int]): Batch size for inference, applied to all submodules.
+            height (Optional[int]): Height of the generated videos.
+            width (Optional[int]): Width of the generated videos.
+            num_frames (Optional[int]): The number of frames in the generated video.
+            fps (Optional[int]): The frames per second of the generated video.
+            max_seq_len (Optional[int]): Maximum sequence length supported by the model.
+            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+        """
+        super().__init__(**kwargs)
+
+        self.text_encoder = self.init_submodule_config(
+            RBLNT5EncoderModelConfig, text_encoder, batch_size=batch_size, max_seq_len=max_seq_len
+        )
+        self.transformer = self.init_submodule_config(
+            RBLNCosmosTransformer3DModelConfig,
+            transformer,
+            batch_size=batch_size,
+            max_seq_len=max_seq_len,
+            height=height,
+            width=width,
+            num_frames=num_frames,
+            fps=fps,
+        )
+        self.vae = self.init_submodule_config(
+            RBLNAutoencoderKLCosmosConfig,
+            vae,
+            batch_size=batch_size,
+            uses_encoder=self.__class__._vae_uses_encoder,
+            height=height,
+            width=width,
+            num_frames=num_frames,
+        )
+        self.safety_checker = self.init_submodule_config(
+            RBLNCosmosSafetyCheckerConfig,
+            safety_checker,
+            batch_size=batch_size,
+            height=height,
+            width=width,
+        )
+
+    @property
+    def batch_size(self):
+        return self.vae.batch_size
+
+    @property
+    def max_seq_len(self):
+        return self.text_encoder.max_seq_len
+
+
+class RBLNCosmosTextToWorldPipelineConfig(_RBLNCosmosPipelineBaseConfig):
+    _vae_uses_encoder = False
+
+
+class RBLNCosmosVideoToWorldPipelineConfig(_RBLNCosmosPipelineBaseConfig):
+    _vae_uses_encoder = True

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

@@ -21,8 +21,6 @@ from ..models.configuration_prior_transformer import RBLNPriorTransformerConfig
 
 
 class RBLNKandinskyV22PipelineBaseConfig(RBLNModelConfig):
-    """Base configuration class for Kandinsky V2.2 decoder pipelines."""
-
     submodules = ["unet", "movq"]
     _movq_uses_encoder = False
 

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

@@ -20,10 +20,6 @@ from ..models import RBLNAutoencoderKLConfig, RBLNUNet2DConditionModelConfig
 
 
 class RBLNStableDiffusionPipelineBaseConfig(RBLNModelConfig):
-    """
-    Base configuration for Stable Diffusion pipelines.
-    """
-
     submodules = ["text_encoder", "unet", "vae"]
     _vae_uses_encoder = False
 

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

@@ -20,10 +20,6 @@ from ..models import RBLNAutoencoderKLConfig, RBLNSD3Transformer2DModelConfig
 
 
 class RBLNStableDiffusion3PipelineBaseConfig(RBLNModelConfig):
-    """
-    Base configuration for Stable Diffusion 3 pipelines.
-    """
-
     submodules = ["transformer", "text_encoder", "text_encoder_2", "text_encoder_3", "vae"]
     _vae_uses_encoder = False
 
@@ -115,6 +111,7 @@ class RBLNStableDiffusion3PipelineBaseConfig(RBLNModelConfig):
             text_encoder_3,
             batch_size=batch_size,
             max_seq_len=max_seq_len,
+            model_input_names=["input_ids"],
         )
         self.transformer = self.init_submodule_config(
             RBLNSD3Transformer2DModelConfig,

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

@@ -20,10 +20,6 @@ from ..models import RBLNAutoencoderKLConfig, RBLNUNet2DConditionModelConfig
 
 
 class RBLNStableDiffusionXLPipelineBaseConfig(RBLNModelConfig):
-    """
-    Base configuration for Stable Diffusion XL pipelines.
-    """
-
     submodules = ["text_encoder", "text_encoder_2", "unet", "vae"]
     _vae_uses_encoder = False
 

optimum/rbln/diffusers/modeling_diffusers.py

@@ -45,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
@@ -55,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
@@ -67,6 +68,7 @@ class RBLNDiffusionMixin:
 
     _connected_classes = {}
     _submodules = []
+    _optional_submodules = []
     _prefix = {}
     _rbln_config_class = None
     _hf_class = None
@@ -184,31 +186,42 @@ class RBLNDiffusionMixin:
         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 = 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
-                )
                 kwargs[submodule_name] = submodule
 
         with ContextRblnConfig(
@@ -352,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, "")
@@ -386,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,
     )