optimum-rbln 0.8.2a4__py3-none-any.whl → 0.9.3rc0__py3-none-any.whl

optimum/rbln/configuration_utils.py

@@ -21,8 +21,10 @@ from typing import Any, Dict, List, Optional, Protocol, Tuple, Type, Union, runt
 
 import numpy as np
 import torch
+from packaging.version import Version
 
 from .__version__ import __version__
+from .utils.depreacate_utils import warn_deprecated_npu
 from .utils.logging import get_logger
 from .utils.runtime_utils import ContextRblnConfig
 
@@ -31,7 +33,6 @@ logger = get_logger(__name__)
 
 
 DEFAULT_COMPILED_MODEL_NAME = "compiled_model"
-DEFAULT_MOD_NAME = "default"
 TypeInputInfo = List[Tuple[str, Tuple[int], str]]
 
 
@@ -39,6 +40,9 @@ TypeInputInfo = List[Tuple[str, Tuple[int], str]]
 class RBLNSerializableConfigProtocol(Protocol):
     def _prepare_for_serialization(self) -> Dict[str, Any]: ...
 
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({self._prepare_for_serialization()})"
+
 
 @dataclass
 class RBLNCompileConfig:
@@ -47,17 +51,13 @@ class RBLNCompileConfig:
 
     Attributes:
         compiled_model_name (str): Name of the compiled model.
-        mod_name (str): Name of the RBLN module.
         input_info (Union[List[TypeInputInfo], TypeInputInfo]): Information about input tensors.
-        fusion (Optional[bool]): Whether to use fusion optimization.
         npu (Optional[str]): NPU configuration.
         tensor_parallel_size (Optional[int]): Size for tensor parallelism.
     """
 
     compiled_model_name: str = DEFAULT_COMPILED_MODEL_NAME
-    mod_name: str = DEFAULT_MOD_NAME
     input_info: Union[List[TypeInputInfo], TypeInputInfo] = None
-    fusion: Optional[bool] = None
     npu: Optional[str] = None
     tensor_parallel_size: Optional[int] = None
 
@@ -111,9 +111,7 @@ class RBLNCompileConfig:
 
     def update(self, kwargs: Dict[str, Any]):
         self.compiled_model_name = kwargs.get("compiled_model_name", self.compiled_model_name)
-        self.mod_name = kwargs.get("mod_name", self.mod_name)
         self.input_info = kwargs.get("input_info", self.input_info)
-        self.fusion = kwargs.get("fusion", self.fusion)
         self.npu = kwargs.get("npu", self.npu)
         self.tensor_parallel_size = kwargs.get("tensor_parallel_size", self.tensor_parallel_size)
         return self
@@ -147,7 +145,7 @@ class RBLNCompileConfig:
         return asdict(self)
 
 
-RUNTIME_KEYWORDS = ["create_runtimes", "optimize_host_memory", "device", "device_map", "activate_profiler", "timeout"]
+RUNTIME_KEYWORDS = ["create_runtimes", "device", "device_map", "activate_profiler", "timeout"]
 CONFIG_MAPPING: Dict[str, Type["RBLNModelConfig"]] = {}
 
 
@@ -183,6 +181,15 @@ def load_config(path: str) -> Tuple[Type["RBLNModelConfig"], Dict[str, Any]]:
 
 
 class RBLNAutoConfig:
+    """
+    Resolver and factory for RBLN model configurations.
+
+    This class selects the concrete `RBLNModelConfig` subclass, validates the
+    provided data, and returns a frozen configuration object that serves as the
+    single source of truth during export and load. It does not define the schema
+    or control model behavior.
+    """
+
     def __new__(cls, **kwargs):
         cls_name = kwargs.get("cls_name")
         if cls_name is None:
@@ -192,6 +199,33 @@ class RBLNAutoConfig:
 
     @staticmethod
     def load_from_dict(config_dict: Dict[str, Any]) -> "RBLNModelConfig":
+        """
+        Build a `RBLNModelConfig` from a plain dictionary.
+
+        The dictionary must contain `cls_name`, which identifies the concrete
+        configuration class to instantiate. All other keys are forwarded to the
+        target class initializer. This method does not mutate `config_dict`.
+
+        Args:
+            config_dict: Mapping typically created by `json.load` or `yaml.safe_load`.
+                For example, the parsed contents of `rbln_config.json`.
+
+        Returns:
+            RBLNModelConfig: A configuration instance. The specific subclass is
+            selected by `config_dict["cls_name"]`.
+
+        Raises:
+            ValueError: If `cls_name` is missing.
+            Exception: Any error raised by the target config class during init.
+
+        Examples:
+            >>> data = {
+            ...     "cls_name": "RBLNLlamaForCausalLMConfig",
+            ...     "create_runtimes": False,
+            ...     "tensor_parallel_size": 4
+            ... }
+            >>> cfg = RBLNAutoConfig.load_from_dict(data)
+        """
         cls_name = config_dict.get("cls_name")
         if cls_name is None:
             raise ValueError("`cls_name` is required.")
@@ -204,7 +238,8 @@ class RBLNAutoConfig:
         Register a new configuration for this class.
 
         Args:
-            config ([`RBLNModelConfig`]): The config to register.
+            config (RBLNModelConfig): The config to register.
+            exist_ok (bool): Whether to allow registering an already registered model.
         """
         if not issubclass(config, RBLNModelConfig):
             raise ValueError("`config` must be a subclass of RBLNModelConfig.")
@@ -246,9 +281,6 @@ class RBLNAutoConfig:
             if key[5:] not in RUNTIME_KEYWORDS and key[5:] not in cls.submodules
         }
 
-        if len(rbln_kwargs) > 0:
-            raise ValueError(f"Cannot set the following arguments: {list(rbln_kwargs.keys())}")
-
         # Process submodule's rbln_config
         for submodule in cls.submodules:
             if submodule not in config_file:
@@ -263,6 +295,16 @@ class RBLNAutoConfig:
 
         config_file.update(rbln_runtime_kwargs)
 
+        rbln_config = cls(**config_file)
+
+        if len(rbln_kwargs) > 0:
+            for key, value in rbln_kwargs.items():
+                if getattr(rbln_config, key) != value:
+                    raise ValueError(
+                        f"Cannot set the following arguments: {list(rbln_kwargs.keys())} "
+                        f"Since the value is already set to {getattr(rbln_config, key)}"
+                    )
+
         if return_unused_kwargs:
             return cls(**config_file), kwargs
         else:
@@ -273,6 +315,7 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
     """Base configuration class for RBLN models that handles compilation settings, runtime options, and submodules.
 
     This class provides functionality for:
+
     1. Managing compilation configurations for RBLN devices
     2. Configuring runtime behavior such as device placement
     3. Handling nested configuration objects for complex model architectures
@@ -474,10 +517,10 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
     non_save_attributes = [
         "_frozen",
         "_runtime_options",
+        "torch_dtype",
         "npu",
         "tensor_parallel_size",
         "create_runtimes",
-        "optimize_host_memory",
         "device",
         "device_map",
         "activate_profiler",
@@ -486,18 +529,18 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
     submodules: List[str] = []
     subclass_non_save_attributes = []
 
-    def init_submodule_config(
+    def initialize_submodule_config(
         self,
-        submodule_config_cls: Type["RBLNModelConfig"],
         submodule_config: Optional[Union[Dict[str, Any], "RBLNModelConfig"]] = None,
-        **kwargs: Dict[str, Any],
+        force_kwargs: bool = False,
+        **kwargs: Any,
     ) -> "RBLNModelConfig":
-        # Initialize a submodule config from a dict or a RBLNModelConfig.
-        # kwargs is specified from the predecessor config.
-
         if submodule_config is None:
             submodule_config = {}
 
+        if isinstance(submodule_config, RBLNModelConfig):
+            return submodule_config
+
         if isinstance(submodule_config, dict):
             from_predecessor = self._runtime_options.copy()
             from_predecessor.update(
@@ -511,13 +554,60 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
 
             init_kwargs = from_predecessor
             init_kwargs.update(submodule_config)
-            submodule_config = submodule_config_cls(**init_kwargs)
 
-        if not isinstance(submodule_config, submodule_config_cls):
+            if force_kwargs:
+                for key, value in kwargs.items():
+                    if key in init_kwargs:
+                        if init_kwargs[key] != value:
+                            raise ValueError(
+                                f"Parameter conflict for '{key}': submodule_config has {init_kwargs[key]}, "
+                                f"but kwargs has {value}. Using kwargs value: {value}"
+                            )
+                        init_kwargs[key] = value
+
+            if "cls_name" in init_kwargs:
+                config_cls = get_rbln_config_class(init_kwargs["cls_name"])
+            else:
+                return init_kwargs
+
+            submodule_config = config_cls(**init_kwargs)
+
+        if not isinstance(submodule_config, RBLNModelConfig):
             raise TypeError(f"Invalid submodule config type: {type(submodule_config)}")
 
         return submodule_config
 
+    def filter_parameters(self, config_cls: Type["RBLNModelConfig"], parameters: Dict[str, Any]) -> Dict[str, Any]:
+        import importlib
+
+        model_cls_name = config_cls.__name__.replace("Config", "")
+        modeling_module_name = config_cls.__module__.replace("configuration_", "modeling_")
+
+        model_cls = None
+        try:
+            modeling_module = importlib.import_module(modeling_module_name)
+            if hasattr(modeling_module, model_cls_name):
+                model_cls = getattr(modeling_module, model_cls_name)
+        except ImportError:
+            logger.debug(f"Could not import modeling module: {modeling_module_name}")
+
+        filtered_out_params = set()
+
+        if model_cls is not None:
+            if not getattr(model_cls, "_tp_support", False):
+                filtered_out_params.add("tensor_parallel_size")
+
+        filtered_params = {}
+        for key, value in parameters.items():
+            if key in filtered_out_params:
+                logger.debug(
+                    f"Parameter '{key}' filtered out for {config_cls.__name__} (not supported by model flags)."
+                )
+            else:
+                filtered_params[key] = value
+
+        return filtered_params
+
     def __setattr__(self, key, value):
         if (
             key != "_attributes_map"
@@ -556,7 +646,6 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
         self,
         cls_name: Optional[str] = None,
         create_runtimes: Optional[bool] = None,
-        optimize_host_memory: Optional[bool] = None,
         device: Optional[Union[int, List[int]]] = None,
         device_map: Optional[Dict[str, Union[int, List[int]]]] = None,
         activate_profiler: Optional[bool] = None,
@@ -564,8 +653,11 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
         tensor_parallel_size: Optional[int] = None,
         timeout: Optional[int] = None,
         optimum_rbln_version: Optional[str] = None,
+        _torch_dtype: Optional[str] = None,
         _compile_cfgs: List[RBLNCompileConfig] = [],
-        **kwargs: Dict[str, Any],
+        *,
+        optimize_host_memory: Optional[bool] = None,
+        **kwargs: Any,
     ):
         """
         Initialize a RBLN model configuration with runtime options and compile configurations.
@@ -573,7 +665,6 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
         Args:
             cls_name (Optional[str]): The class name of the configuration. Defaults to the current class name.
             create_runtimes (Optional[bool]): Whether to create RBLN runtimes. Defaults to True.
-            optimize_host_memory (Optional[bool]): Whether to optimize host memory usage. Defaults to True.
             device (Optional[Union[int, List[int]]]): The device(s) to load the model onto. Can be a single device ID or a list.
             device_map (Optional[Dict[str, Union[int, List[int]]]]): Mapping from compiled model names to device IDs.
             activate_profiler (Optional[bool]): Whether to activate the profiler for performance analysis.
@@ -581,8 +672,9 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
             tensor_parallel_size (Optional[int]): Size for tensor parallelism to distribute the model across devices.
             timeout (Optional[int]): The timeout for the runtime in seconds. If it isn't provided, it will be set to 60 by default.
             optimum_rbln_version (Optional[str]): The optimum-rbln version used for this configuration.
+            _torch_dtype (Optional[str]): The data type to use for the model.
             _compile_cfgs (List[RBLNCompileConfig]): List of compilation configurations for the model.
-            **kwargs: Additional keyword arguments.
+            kwargs: Additional keyword arguments.
 
         Raises:
             ValueError: If unexpected keyword arguments are provided.
@@ -598,16 +690,19 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
 
         self._runtime_options = {}
         self._runtime_options["create_runtimes"] = create_runtimes
-        self._runtime_options["optimize_host_memory"] = optimize_host_memory
         self._runtime_options["device"] = device
         self._runtime_options["device_map"] = device_map
         self._runtime_options["activate_profiler"] = activate_profiler
         self._runtime_options["timeout"] = timeout
 
+        if optimize_host_memory is not None:
+            logger.warning("`optimize_host_memory` is deprecated and will be removed in future versions.")
+
         # Automatically pass npu, tensor_parallel_size to compile_cfgs
         self.npu = npu
         self.tensor_parallel_size = tensor_parallel_size
 
+        self._torch_dtype = _torch_dtype or "float32"
         self.optimum_rbln_version = optimum_rbln_version
         if self.optimum_rbln_version is None:
             self.optimum_rbln_version = __version__
@@ -620,8 +715,34 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
             self.set_compile_cfgs([RBLNCompileConfig(**cfg) for cfg in self._compile_cfgs])
 
         if len(kwargs) > 0:
+            if optimum_rbln_version is not None:  # loaded from file
+                if Version(__version__) < Version(optimum_rbln_version):
+                    diff = "newer"
+                elif Version(__version__) > Version(optimum_rbln_version):
+                    diff = "older"
+                else:
+                    diff = None
+                if diff is not None:
+                    raise ValueError(
+                        f"Unexpected arguments: {kwargs.keys()}\n"
+                        f"Maybe you are trying to load a model compiled with {diff} version of optimum-rbln. "
+                        "It is recommended to use the same version to compile and load the model.\n"
+                        f"Current version: {__version__}, Loaded version: {optimum_rbln_version}"
+                    )
+
             raise ValueError(f"Unexpected arguments: {kwargs.keys()}")
 
+    @property
+    def torch_dtype(self):
+        return getattr(torch, self._torch_dtype)
+
+    @torch_dtype.setter
+    def torch_dtype(self, torch_dtype: Union[str, torch.dtype]):
+        if isinstance(torch_dtype, torch.dtype):
+            torch_dtype = RBLNCompileConfig.normalize_dtype(torch_dtype)
+
+        self._torch_dtype = torch_dtype
+
     @property
     def rbln_model_cls_name(self) -> str:
         return self.__class__.__name__[:-6]
@@ -675,6 +796,9 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
             compile_cfg.npu = self.npu
             compile_cfg.tensor_parallel_size = self.tensor_parallel_size
 
+        target_npu = self.npu or next((cfg.npu for cfg in self._compile_cfgs if cfg.npu is not None), None)
+        warn_deprecated_npu(target_npu)
+
     def freeze(self):
         if self._frozen:
             raise RuntimeError(f"`{self.__class__.__name__}` is already frozen.")
@@ -713,13 +837,13 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
             json.dump(serializable_data, jsonf, indent=2)
 
     @classmethod
-    def load(cls, path: str, **kwargs: Dict[str, Any]) -> "RBLNModelConfig":
+    def load(cls, path: str, **kwargs: Any) -> "RBLNModelConfig":
         """
         Load a RBLNModelConfig from a path.
 
         Args:
             path (str): Path to the RBLNModelConfig file or directory containing the config file.
-            **kwargs: Additional keyword arguments to override configuration values.
+            kwargs: Additional keyword arguments to override configuration values.
                       Keys starting with 'rbln_' will have the prefix removed and be used
                       to update the configuration.
 
@@ -746,7 +870,7 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
     def initialize_from_kwargs(
         cls: Type["RBLNModelConfig"],
         rbln_config: Optional[Union[Dict[str, Any], "RBLNModelConfig"]] = None,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ) -> Tuple["RBLNModelConfig", Dict[str, Any]]:
         # Initialize RBLNModelConfig from kwargs.
         kwargs_keys = list(kwargs.keys())
@@ -791,19 +915,6 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
     def create_runtimes(self, create_runtimes: bool):
         self._runtime_options["create_runtimes"] = create_runtimes
 
-    @property
-    def optimize_host_memory(self):
-        context = ContextRblnConfig.get_current_context()["optimize_host_memory"]
-        if context is not None:
-            return context
-        elif self._runtime_options["optimize_host_memory"] is None:
-            return True
-        return self._runtime_options["optimize_host_memory"]
-
-    @optimize_host_memory.setter
-    def optimize_host_memory(self, optimize_host_memory: bool):
-        self._runtime_options["optimize_host_memory"] = optimize_host_memory
-
     @property
     def device(self):
         context = ContextRblnConfig.get_current_context()["device"]

optimum/rbln/diffusers/__init__.py

@@ -59,6 +59,9 @@ _import_structure = {
         "RBLNVQModelConfig",
     ],
     "pipelines": [
+        "RBLNAutoPipelineForImage2Image",
+        "RBLNAutoPipelineForInpainting",
+        "RBLNAutoPipelineForText2Image",
         "RBLNCosmosTextToWorldPipeline",
         "RBLNCosmosVideoToWorldPipeline",
         "RBLNCosmosSafetyChecker",
@@ -135,6 +138,7 @@ if TYPE_CHECKING:
     from .modeling_diffusers import RBLNDiffusionMixin
     from .models import (
         RBLNAutoencoderKL,
+        RBLNAutoencoderKLCosmos,
         RBLNControlNetModel,
         RBLNCosmosTransformer3DModel,
         RBLNPriorTransformer,
@@ -143,6 +147,9 @@ if TYPE_CHECKING:
         RBLNVQModel,
     )
     from .pipelines import (
+        RBLNAutoPipelineForImage2Image,
+        RBLNAutoPipelineForInpainting,
+        RBLNAutoPipelineForText2Image,
         RBLNCosmosSafetyChecker,
         RBLNCosmosTextToWorldPipeline,
         RBLNCosmosVideoToWorldPipeline,

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

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Dict, Optional, Tuple
+from typing import Any, Optional, Tuple
 
 from ....configuration_utils import RBLNModelConfig
 
@@ -33,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: Dict[str, Any],
+        **kwargs: Any,
     ):
         """
         Args:
@@ -46,7 +46,7 @@ class RBLNAutoencoderKLConfig(RBLNModelConfig):
                 Determines how much smaller the latent representations are compared to the original images.
             in_channels (Optional[int]): Number of input channels for the model.
             latent_channels (Optional[int]): Number of channels in the latent space.
-            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+            kwargs: Additional arguments passed to the parent RBLNModelConfig.
 
         Raises:
             ValueError: If batch_size is not a positive integer.

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

@@ -52,7 +52,7 @@ class RBLNAutoencoderKLCosmosConfig(RBLNModelConfig):
                 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.
+            kwargs: Additional arguments passed to the parent RBLNModelConfig.
 
         Raises:
             ValueError: If batch_size is not a positive integer.

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

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Dict, Optional, Tuple
+from typing import Any, Optional, Tuple
 
 from ....configuration_utils import RBLNModelConfig
 
@@ -29,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: Dict[str, Any],
+        **kwargs: Any,
     ):
         """
         Args:
@@ -42,7 +42,7 @@ class RBLNControlNetModelConfig(RBLNModelConfig):
                 of the VAE input/output images.
             text_model_hidden_size (Optional[int]): Hidden size of the text encoder model used
                 for conditioning.
-            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+            kwargs: Additional arguments passed to the parent RBLNModelConfig.
 
         Raises:
             ValueError: If batch_size is not a positive integer.

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

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Dict, Optional
+from typing import Any, Optional
 
 from ....configuration_utils import RBLNModelConfig
 
@@ -22,7 +22,7 @@ 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.
+    for Transformer models used in diffusion models like Kandinsky V2.2.
     """
 
     subclass_non_save_attributes = ["_batch_size_is_specified"]
@@ -32,14 +32,14 @@ class RBLNPriorTransformerConfig(RBLNModelConfig):
         batch_size: Optional[int] = None,
         embedding_dim: Optional[int] = None,
         num_embeddings: Optional[int] = None,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ):
         """
         Args:
             batch_size (Optional[int]): The batch size for inference. Defaults to 1.
             embedding_dim (Optional[int]): Dimension of the embedding vectors in the model.
             num_embeddings (Optional[int]): Number of discrete embeddings in the codebook.
-            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+            kwargs: Additional arguments passed to the parent RBLNModelConfig.
 
         Raises:
             ValueError: If batch_size is not a positive integer.

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

@@ -12,13 +12,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Dict, Optional
+from typing import Any, Optional
 
 from ....configuration_utils import RBLNModelConfig
 
 
 class RBLNCosmosTransformer3DModelConfig(RBLNModelConfig):
-    """Configuration class for RBLN Cosmos Transformer models."""
+    """
+    Configuration class for RBLN Cosmos Transformer models.
+
+    This class inherits from RBLNModelConfig and provides specific configuration options
+    for Transformer models used in diffusion models like Cosmos.
+    """
 
     def __init__(
         self,
@@ -33,7 +38,7 @@ class RBLNCosmosTransformer3DModelConfig(RBLNModelConfig):
         num_latent_frames: Optional[int] = None,
         latent_height: Optional[int] = None,
         latent_width: Optional[int] = None,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ):
         """
         Args:
@@ -47,7 +52,7 @@ class RBLNCosmosTransformer3DModelConfig(RBLNModelConfig):
             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.
+            kwargs: Additional arguments passed to the parent RBLNModelConfig.
 
         Raises:
             ValueError: If batch_size is not a positive integer.

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

@@ -12,13 +12,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Dict, Optional, Tuple, Union
+from typing import Any, Optional, Tuple, Union
 
 from ....configuration_utils import RBLNModelConfig
 
 
 class RBLNSD3Transformer2DModelConfig(RBLNModelConfig):
-    """Configuration class for RBLN Stable Diffusion 3 Transformer models."""
+    """
+    Configuration class for RBLN Stable Diffusion 3 Transformer models.
+
+    This class inherits from RBLNModelConfig and provides specific configuration options
+    for Transformer models used in diffusion models like Stable Diffusion 3.
+    """
 
     subclass_non_save_attributes = ["_batch_size_is_specified"]
 
@@ -27,7 +32,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: Dict[str, Any],
+        **kwargs: Any,
     ):
         """
         Args:
@@ -36,7 +41,7 @@ class RBLNSD3Transformer2DModelConfig(RBLNModelConfig):
                 of the generated samples. If an integer is provided, it's used for both height and width.
             prompt_embed_length (Optional[int]): The length of the embedded prompt vectors that
                 will be used to condition the transformer model.
-            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+            kwargs: Additional arguments passed to the parent RBLNModelConfig.
 
         Raises:
             ValueError: If batch_size is not a positive integer.

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

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Dict, Optional, Tuple
+from typing import Any, Optional, Tuple
 
 from ....configuration_utils import RBLNModelConfig
 
@@ -38,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: Dict[str, Any],
+        **kwargs: Any,
     ):
         """
         Args:
@@ -52,7 +52,7 @@ class RBLNUNet2DConditionModelConfig(RBLNModelConfig):
             in_features (Optional[int]): Number of input features for the model.
             text_model_hidden_size (Optional[int]): Hidden size of the text encoder model.
             image_model_hidden_size (Optional[int]): Hidden size of the image encoder model.
-            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+            kwargs: Additional arguments passed to the parent RBLNModelConfig.
 
         Raises:
             ValueError: If batch_size is not a positive integer.

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

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Dict, Optional, Tuple
+from typing import Any, Optional, Tuple
 
 from ....configuration_utils import RBLNModelConfig
 
@@ -33,7 +33,7 @@ class RBLNVQModelConfig(RBLNModelConfig):
         vqmodel_scale_factor: Optional[float] = None,  # TODO: rename to scaling_factor
         in_channels: Optional[int] = None,
         latent_channels: Optional[int] = None,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ):
         """
         Args:
@@ -46,7 +46,7 @@ class RBLNVQModelConfig(RBLNModelConfig):
                 Determines the downsampling ratio between original images and latent representations.
             in_channels (Optional[int]): Number of input channels for the model.
             latent_channels (Optional[int]): Number of channels in the latent space.
-            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+            kwargs: Additional arguments passed to the parent RBLNModelConfig.
 
         Raises:
             ValueError: If batch_size is not a positive integer.