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

optimum/rbln/__init__.py

@@ -26,6 +26,7 @@ _import_structure = {
         "RBLNModel",
     ],
     "configuration_utils": [
+        "RBLNAutoConfig",
         "RBLNCompileConfig",
         "RBLNModelConfig",
     ],
@@ -69,6 +70,8 @@ _import_structure = {
         "RBLNCLIPVisionModelConfig",
         "RBLNCLIPVisionModelWithProjection",
         "RBLNCLIPVisionModelWithProjectionConfig",
+        "RBLNColPaliForRetrieval",
+        "RBLNColPaliForRetrievalConfig",
         "RBLNDecoderOnlyModelForCausalLM",
         "RBLNDecoderOnlyModelForCausalLMConfig",
         "RBLNDistilBertForQuestionAnswering",
@@ -135,8 +138,17 @@ _import_structure = {
     "diffusers": [
         "RBLNAutoencoderKL",
         "RBLNAutoencoderKLConfig",
+        "RBLNAutoencoderKLCosmos",
+        "RBLNAutoencoderKLCosmosConfig",
         "RBLNControlNetModel",
         "RBLNControlNetModelConfig",
+        "RBLNCosmosTextToWorldPipeline",
+        "RBLNCosmosVideoToWorldPipeline",
+        "RBLNCosmosTextToWorldPipelineConfig",
+        "RBLNCosmosVideoToWorldPipelineConfig",
+        "RBLNCosmosSafetyChecker",
+        "RBLNCosmosTransformer3DModel",
+        "RBLNCosmosTransformer3DModelConfig",
         "RBLNDiffusionMixin",
         "RBLNKandinskyV22CombinedPipeline",
         "RBLNKandinskyV22CombinedPipelineConfig",
@@ -192,14 +204,24 @@ _import_structure = {
 
 if TYPE_CHECKING:
     from .configuration_utils import (
+        RBLNAutoConfig,
         RBLNCompileConfig,
         RBLNModelConfig,
     )
     from .diffusers import (
         RBLNAutoencoderKL,
         RBLNAutoencoderKLConfig,
+        RBLNAutoencoderKLCosmos,
+        RBLNAutoencoderKLCosmosConfig,
         RBLNControlNetModel,
         RBLNControlNetModelConfig,
+        RBLNCosmosSafetyChecker,
+        RBLNCosmosTextToWorldPipeline,
+        RBLNCosmosTextToWorldPipelineConfig,
+        RBLNCosmosTransformer3DModel,
+        RBLNCosmosTransformer3DModelConfig,
+        RBLNCosmosVideoToWorldPipeline,
+        RBLNCosmosVideoToWorldPipelineConfig,
         RBLNDiffusionMixin,
         RBLNKandinskyV22CombinedPipeline,
         RBLNKandinskyV22CombinedPipelineConfig,
@@ -295,6 +317,8 @@ if TYPE_CHECKING:
         RBLNCLIPVisionModelConfig,
         RBLNCLIPVisionModelWithProjection,
         RBLNCLIPVisionModelWithProjectionConfig,
+        RBLNColPaliForRetrieval,
+        RBLNColPaliForRetrievalConfig,
         RBLNDecoderOnlyModelForCausalLM,
         RBLNDecoderOnlyModelForCausalLMConfig,
         RBLNDistilBertForQuestionAnswering,

optimum/rbln/__version__.py

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

optimum/rbln/configuration_utils.py

@@ -19,6 +19,7 @@ from dataclasses import asdict, dataclass
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Protocol, Tuple, Type, Union, runtime_checkable
 
+import numpy as np
 import torch
 
 from .__version__ import __version__
@@ -61,7 +62,7 @@ class RBLNCompileConfig:
     tensor_parallel_size: Optional[int] = None
 
     @staticmethod
-    def normalize_dtype(dtype):
+    def normalize_dtype(dtype: Union[str, torch.dtype, np.dtype]) -> str:
         """
         Convert framework-specific dtype to string representation.
         i.e. torch.float32 -> "float32"
@@ -70,7 +71,7 @@ class RBLNCompileConfig:
             dtype: The input dtype (can be string, torch dtype, or numpy dtype).
 
         Returns:
-            str: The normalized string representation of the dtype.
+            The normalized string representation of the dtype.
         """
         if isinstance(dtype, str):
             return dtype
@@ -147,6 +148,17 @@ class RBLNCompileConfig:
 
 
 RUNTIME_KEYWORDS = ["create_runtimes", "optimize_host_memory", "device", "device_map", "activate_profiler"]
+CONFIG_MAPPING: Dict[str, Type["RBLNModelConfig"]] = {}
+
+
+def get_rbln_config_class(rbln_config_class_name: str) -> Type["RBLNModelConfig"]:
+    cls = getattr(importlib.import_module("optimum.rbln"), rbln_config_class_name, None)
+    if cls is None:
+        if rbln_config_class_name in CONFIG_MAPPING:
+            cls = CONFIG_MAPPING[rbln_config_class_name]
+        else:
+            raise ValueError(f"Configuration for {rbln_config_class_name} not found.")
+    return cls
 
 
 def load_config(path: str) -> Tuple[Type["RBLNModelConfig"], Dict[str, Any]]:
@@ -166,7 +178,7 @@ def load_config(path: str) -> Tuple[Type["RBLNModelConfig"], Dict[str, Any]]:
             )
 
     cls_name = config_file["cls_name"]
-    cls = getattr(importlib.import_module("optimum.rbln"), cls_name)
+    cls = get_rbln_config_class(cls_name)
     return cls, config_file
 
 
@@ -175,7 +187,7 @@ class RBLNAutoConfig:
         cls_name = kwargs.get("cls_name")
         if cls_name is None:
             raise ValueError("`cls_name` is required.")
-        cls = getattr(importlib.import_module("optimum.rbln"), cls_name)
+        cls = get_rbln_config_class(cls_name)
         return cls(**kwargs)
 
     @staticmethod
@@ -183,9 +195,27 @@ class RBLNAutoConfig:
         cls_name = config_dict.get("cls_name")
         if cls_name is None:
             raise ValueError("`cls_name` is required.")
-        cls = getattr(importlib.import_module("optimum.rbln"), cls_name)
+        cls = get_rbln_config_class(cls_name)
         return cls(**config_dict)
 
+    @staticmethod
+    def register(config: Type["RBLNModelConfig"], exist_ok=False):
+        """
+        Register a new configuration for this class.
+
+        Args:
+            config ([`RBLNModelConfig`]): The config to register.
+        """
+        if not issubclass(config, RBLNModelConfig):
+            raise ValueError("`config` must be a subclass of RBLNModelConfig.")
+
+        native_cls = getattr(importlib.import_module("optimum.rbln"), config.__name__, None)
+        if config.__name__ in CONFIG_MAPPING or native_cls is not None:
+            if not exist_ok:
+                raise ValueError(f"Configuration for {config.__name__} already registered.")
+
+        CONFIG_MAPPING[config.__name__] = config
+
     @staticmethod
     def load(
         path: str,
@@ -307,9 +337,6 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
         # Save to disk
         config.save("/path/to/model")
 
-        # Load configuration from disk
-        loaded_config = RBLNModelConfig.load("/path/to/model")
-
         # Using AutoConfig
         loaded_config = RBLNAutoConfig.load("/path/to/model")
         ```
@@ -462,13 +489,11 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
         self,
         submodule_config_cls: Type["RBLNModelConfig"],
         submodule_config: Optional[Union[Dict[str, Any], "RBLNModelConfig"]] = None,
-        **kwargs,
+        **kwargs: Dict[str, Any],
     ) -> "RBLNModelConfig":
-        """
-        Initialize a submodule config from a dict or a RBLNModelConfig.
+        # Initialize a submodule config from a dict or a RBLNModelConfig.
+        # kwargs is specified from the predecessor config.
 
-        kwargs is specified from the predecessor config.
-        """
         if submodule_config is None:
             submodule_config = {}
 
@@ -538,7 +563,7 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
         tensor_parallel_size: Optional[int] = None,
         optimum_rbln_version: Optional[str] = None,
         _compile_cfgs: List[RBLNCompileConfig] = [],
-        **kwargs,
+        **kwargs: Dict[str, Any],
     ):
         """
         Initialize a RBLN model configuration with runtime options and compile configurations.
@@ -608,10 +633,8 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
         return rbln_model_cls
 
     def _prepare_for_serialization(self) -> Dict[str, Any]:
-        """
-        Prepare the attributes map for serialization by converting nested RBLNModelConfig
-        objects to their serializable form.
-        """
+        # Prepare the attributes map for serialization by converting nested RBLNModelConfig
+        # objects to their serializable form.
         serializable_map = {}
         for key, value in self._attributes_map.items():
             if isinstance(value, RBLNSerializableConfigProtocol):
@@ -686,7 +709,7 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
             json.dump(serializable_data, jsonf, indent=2)
 
     @classmethod
-    def load(cls, path: str, **kwargs) -> "RBLNModelConfig":
+    def load(cls, path: str, **kwargs: Dict[str, Any]) -> "RBLNModelConfig":
         """
         Load a RBLNModelConfig from a path.
 
@@ -719,11 +742,9 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
     def initialize_from_kwargs(
         cls: Type["RBLNModelConfig"],
         rbln_config: Optional[Union[Dict[str, Any], "RBLNModelConfig"]] = None,
-        **kwargs,
+        **kwargs: Dict[str, Any],
     ) -> Tuple["RBLNModelConfig", Dict[str, Any]]:
-        """
-        Initialize RBLNModelConfig from kwargs.
-        """
+        # Initialize RBLNModelConfig from kwargs.
         kwargs_keys = list(kwargs.keys())
         rbln_kwargs = {key[5:]: kwargs.pop(key) for key in kwargs_keys if key.startswith("rbln_")}
 
@@ -741,16 +762,7 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
         return rbln_config, kwargs
 
     def get_default_values_for_original_cls(self, func_name: str, keys: List[str]) -> Dict[str, Any]:
-        """
-        Get default values for original class attributes from RBLNModelConfig.
-
-        Args:
-            func_name (str): The name of the function to get the default values for.
-            keys (List[str]): The keys of the attributes to get.
-
-        Returns:
-            Dict[str, Any]: The default values for the attributes.
-        """
+        # Get default values for original class attributes from RBLNModelConfig.
         model_cls = self.rbln_model_cls.get_hf_class()
         func = getattr(model_cls, func_name)
         func_signature = inspect.signature(func)

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,6 +1,8 @@
 from .configuration_autoencoder_kl import RBLNAutoencoderKLConfig
+from .configuration_autoencoder_kl_cosmos import RBLNAutoencoderKLCosmosConfig
 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_vq_model import RBLNVQModelConfig

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

@@ -12,12 +12,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional, Tuple
+from typing import Any, Dict, Optional, Tuple
 
 from ....configuration_utils import RBLNModelConfig
 
 
 class RBLNAutoencoderKLConfig(RBLNModelConfig):
+    """
+    Configuration class for RBLN Variational Autoencoder (VAE) models.
+
+    This class inherits from RBLNModelConfig and provides specific configuration options
+    for VAE models used in diffusion-based image generation.
+    """
+
     def __init__(
         self,
         batch_size: Optional[int] = None,
@@ -26,7 +33,7 @@ class RBLNAutoencoderKLConfig(RBLNModelConfig):
         vae_scale_factor: Optional[float] = None,  # TODO: rename to scaling_factor
         in_channels: Optional[int] = None,
         latent_channels: Optional[int] = None,
-        **kwargs,
+        **kwargs: Dict[str, Any],
     ):
         """
         Args:

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

@@ -0,0 +1,84 @@
+# 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, Dict, Optional
+
+from ....configuration_utils import RBLNModelConfig
+from ....utils.logging import get_logger
+
+
+logger = get_logger(__name__)
+
+
+class RBLNAutoencoderKLCosmosConfig(RBLNModelConfig):
+    """Configuration class for RBLN Cosmos Variational Autoencoder (VAE) models."""
+
+    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: Dict[str, Any],
+    ):
+        """
+        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_controlnet.py

@@ -12,12 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional, Tuple
+from typing import Any, Dict, Optional, Tuple
 
 from ....configuration_utils import RBLNModelConfig
 
 
 class RBLNControlNetModelConfig(RBLNModelConfig):
+    """Configuration class for RBLN ControlNet models."""
+
     subclass_non_save_attributes = ["_batch_size_is_specified"]
 
     def __init__(
@@ -27,7 +29,7 @@ class RBLNControlNetModelConfig(RBLNModelConfig):
         unet_sample_size: Optional[Tuple[int, int]] = None,
         vae_sample_size: Optional[Tuple[int, int]] = None,
         text_model_hidden_size: Optional[int] = None,
-        **kwargs,
+        **kwargs: Dict[str, Any],
     ):
         """
         Args:

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

@@ -12,12 +12,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional
+from typing import Any, Dict, Optional
 
 from ....configuration_utils import RBLNModelConfig
 
 
 class RBLNPriorTransformerConfig(RBLNModelConfig):
+    """
+    Configuration class for RBLN Prior Transformer models.
+
+    This class inherits from RBLNModelConfig and provides specific configuration options
+    for Prior Transformer models used in diffusion models like Kandinsky V2.2.
+    """
+
     subclass_non_save_attributes = ["_batch_size_is_specified"]
 
     def __init__(
@@ -25,7 +32,7 @@ class RBLNPriorTransformerConfig(RBLNModelConfig):
         batch_size: Optional[int] = None,
         embedding_dim: Optional[int] = None,
         num_embeddings: Optional[int] = None,
-        **kwargs,
+        **kwargs: Dict[str, Any],
     ):
         """
         Args:

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

@@ -0,0 +1,70 @@
+# 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, Dict, Optional
+
+from ....configuration_utils import RBLNModelConfig
+
+
+class RBLNCosmosTransformer3DModelConfig(RBLNModelConfig):
+    """Configuration class for RBLN Cosmos Transformer models."""
+
+    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: Dict[str, Any],
+    ):
+        """
+        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/models/configuration_transformer_sd3.py

@@ -12,12 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional, Tuple, Union
+from typing import Any, Dict, Optional, Tuple, Union
 
 from ....configuration_utils import RBLNModelConfig
 
 
 class RBLNSD3Transformer2DModelConfig(RBLNModelConfig):
+    """Configuration class for RBLN Stable Diffusion 3 Transformer models."""
+
     subclass_non_save_attributes = ["_batch_size_is_specified"]
 
     def __init__(
@@ -25,7 +27,7 @@ class RBLNSD3Transformer2DModelConfig(RBLNModelConfig):
         batch_size: Optional[int] = None,
         sample_size: Optional[Union[int, Tuple[int, int]]] = None,
         prompt_embed_length: Optional[int] = None,
-        **kwargs,
+        **kwargs: Dict[str, Any],
     ):
         """
         Args:

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

@@ -12,12 +12,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional, Tuple
+from typing import Any, Dict, Optional, Tuple
 
 from ....configuration_utils import RBLNModelConfig
 
 
 class RBLNUNet2DConditionModelConfig(RBLNModelConfig):
+    """
+    Configuration class for RBLN UNet2DCondition models.
+
+    This class inherits from RBLNModelConfig and provides specific configuration options
+    for UNet2DCondition models used in diffusion-based image generation.
+    """
+
     subclass_non_save_attributes = ["_batch_size_is_specified"]
 
     def __init__(
@@ -31,7 +38,7 @@ class RBLNUNet2DConditionModelConfig(RBLNModelConfig):
         in_features: Optional[int] = None,
         text_model_hidden_size: Optional[int] = None,
         image_model_hidden_size: Optional[int] = None,
-        **kwargs,
+        **kwargs: Dict[str, Any],
     ):
         """
         Args: