optimum-rbln 0.8.1a0__py3-none-any.whl → 0.8.1a2__py3-none-any.whl

optimum/rbln/__init__.py

@@ -26,6 +26,7 @@ _import_structure = {
         "RBLNModel",
     ],
     "configuration_utils": [
+        "RBLNAutoConfig",
         "RBLNCompileConfig",
         "RBLNModelConfig",
     ],
@@ -192,6 +193,7 @@ _import_structure = {
 
 if TYPE_CHECKING:
     from .configuration_utils import (
+        RBLNAutoConfig,
         RBLNCompileConfig,
         RBLNModelConfig,
     )

optimum/rbln/__version__.py

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

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,19 +489,25 @@ 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 = {}
 
         if isinstance(submodule_config, dict):
             from_predecessor = self._runtime_options.copy()
+            from_predecessor.update(
+                {
+                    "npu": self.npu,
+                    "tensor_parallel_size": self.tensor_parallel_size,
+                    "optimum_rbln_version": self.optimum_rbln_version,
+                }
+            )
             from_predecessor.update(kwargs)
+
             init_kwargs = from_predecessor
             init_kwargs.update(submodule_config)
             submodule_config = submodule_config_cls(**init_kwargs)
@@ -530,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.
@@ -600,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):
@@ -678,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.
 
@@ -711,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_")}
 
@@ -733,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/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_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_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:

optimum/rbln/diffusers/configurations/models/configuration_vq_model.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 RBLNVQModelConfig(RBLNModelConfig):
+    """
+    Configuration class for RBLN VQModel models, used in Kandinsky.
+
+    This class inherits from RBLNModelConfig and provides specific configuration options
+    for VQModel, which acts similarly to a VAE but uses vector quantization.
+    """
+
     def __init__(
         self,
         batch_size: Optional[int] = None,
@@ -26,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,
+        **kwargs: Dict[str, Any],
     ):
         """
         Args:

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

@@ -12,14 +12,18 @@
 # 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
 from ....transformers import RBLNCLIPTextModelConfig, RBLNCLIPTextModelWithProjectionConfig
 from ..models import RBLNAutoencoderKLConfig, RBLNControlNetModelConfig, RBLNUNet2DConditionModelConfig
 
 
-class _RBLNStableDiffusionControlNetPipelineBaseConfig(RBLNModelConfig):
+class RBLNStableDiffusionControlNetPipelineBaseConfig(RBLNModelConfig):
+    """
+    Base configuration for Stable Diffusion ControlNet pipelines.
+    """
+
     submodules = ["text_encoder", "unet", "vae", "controlnet"]
     _vae_uses_encoder = False
 
@@ -38,7 +42,7 @@ class _RBLNStableDiffusionControlNetPipelineBaseConfig(RBLNModelConfig):
         sample_size: Optional[Tuple[int, int]] = None,
         image_size: Optional[Tuple[int, int]] = None,
         guidance_scale: Optional[float] = None,
-        **kwargs,
+        **kwargs: Dict[str, Any],
     ):
         """
         Args:
@@ -138,15 +142,27 @@ class _RBLNStableDiffusionControlNetPipelineBaseConfig(RBLNModelConfig):
         return self.vae.sample_size
 
 
-class RBLNStableDiffusionControlNetPipelineConfig(_RBLNStableDiffusionControlNetPipelineBaseConfig):
+class RBLNStableDiffusionControlNetPipelineConfig(RBLNStableDiffusionControlNetPipelineBaseConfig):
+    """
+    Configuration for Stable Diffusion ControlNet pipeline.
+    """
+
     _vae_uses_encoder = False
 
 
-class RBLNStableDiffusionControlNetImg2ImgPipelineConfig(_RBLNStableDiffusionControlNetPipelineBaseConfig):
+class RBLNStableDiffusionControlNetImg2ImgPipelineConfig(RBLNStableDiffusionControlNetPipelineBaseConfig):
+    """
+    Configuration for Stable Diffusion ControlNet image-to-image pipeline.
+    """
+
     _vae_uses_encoder = True
 
 
-class _RBLNStableDiffusionXLControlNetPipelineBaseConfig(RBLNModelConfig):
+class RBLNStableDiffusionXLControlNetPipelineBaseConfig(RBLNModelConfig):
+    """
+    Base configuration for Stable Diffusion XL ControlNet pipelines.
+    """
+
     submodules = ["text_encoder", "text_encoder_2", "unet", "vae", "controlnet"]
     _vae_uses_encoder = False
 
@@ -166,7 +182,7 @@ class _RBLNStableDiffusionXLControlNetPipelineBaseConfig(RBLNModelConfig):
         sample_size: Optional[Tuple[int, int]] = None,
         image_size: Optional[Tuple[int, int]] = None,
         guidance_scale: Optional[float] = None,
-        **kwargs,
+        **kwargs: Dict[str, Any],
     ):
         """
         Args:
@@ -272,9 +288,17 @@ class _RBLNStableDiffusionXLControlNetPipelineBaseConfig(RBLNModelConfig):
         return self.vae.sample_size
 
 
-class RBLNStableDiffusionXLControlNetPipelineConfig(_RBLNStableDiffusionXLControlNetPipelineBaseConfig):
+class RBLNStableDiffusionXLControlNetPipelineConfig(RBLNStableDiffusionXLControlNetPipelineBaseConfig):
+    """
+    Configuration for Stable Diffusion XL ControlNet pipeline.
+    """
+
     _vae_uses_encoder = False
 
 
-class RBLNStableDiffusionXLControlNetImg2ImgPipelineConfig(_RBLNStableDiffusionXLControlNetPipelineBaseConfig):
+class RBLNStableDiffusionXLControlNetImg2ImgPipelineConfig(RBLNStableDiffusionXLControlNetPipelineBaseConfig):
+    """
+    Configuration for Stable Diffusion XL ControlNet image-to-image pipeline.
+    """
+
     _vae_uses_encoder = True

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

@@ -12,7 +12,7 @@
 # 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
 from ....transformers import RBLNCLIPTextModelWithProjectionConfig, RBLNCLIPVisionModelWithProjectionConfig
@@ -20,7 +20,9 @@ from ..models import RBLNUNet2DConditionModelConfig, RBLNVQModelConfig
 from ..models.configuration_prior_transformer import RBLNPriorTransformerConfig
 
 
-class _RBLNKandinskyV22PipelineBaseConfig(RBLNModelConfig):
+class RBLNKandinskyV22PipelineBaseConfig(RBLNModelConfig):
+    """Base configuration class for Kandinsky V2.2 decoder pipelines."""
+
     submodules = ["unet", "movq"]
     _movq_uses_encoder = False
 
@@ -37,7 +39,7 @@ class _RBLNKandinskyV22PipelineBaseConfig(RBLNModelConfig):
         img_width: Optional[int] = None,
         height: Optional[int] = None,
         width: Optional[int] = None,
-        **kwargs,
+        **kwargs: Dict[str, Any],
     ):
         """
         Args:
@@ -117,19 +119,27 @@ class _RBLNKandinskyV22PipelineBaseConfig(RBLNModelConfig):
         return self.movq.sample_size
 
 
-class RBLNKandinskyV22PipelineConfig(_RBLNKandinskyV22PipelineBaseConfig):
+class RBLNKandinskyV22PipelineConfig(RBLNKandinskyV22PipelineBaseConfig):
+    """Configuration class for the Kandinsky V2.2 text-to-image decoder pipeline."""
+
     _movq_uses_encoder = False
 
 
-class RBLNKandinskyV22Img2ImgPipelineConfig(_RBLNKandinskyV22PipelineBaseConfig):
+class RBLNKandinskyV22Img2ImgPipelineConfig(RBLNKandinskyV22PipelineBaseConfig):
+    """Configuration class for the Kandinsky V2.2 image-to-image decoder pipeline."""
+
     _movq_uses_encoder = True
 
 
-class RBLNKandinskyV22InpaintPipelineConfig(_RBLNKandinskyV22PipelineBaseConfig):
+class RBLNKandinskyV22InpaintPipelineConfig(RBLNKandinskyV22PipelineBaseConfig):
+    """Configuration class for the Kandinsky V2.2 inpainting decoder pipeline."""
+
     _movq_uses_encoder = True
 
 
 class RBLNKandinskyV22PriorPipelineConfig(RBLNModelConfig):
+    """Configuration class for the Kandinsky V2.2 Prior pipeline."""
+
     submodules = ["text_encoder", "image_encoder", "prior"]
 
     def __init__(
@@ -140,7 +150,7 @@ class RBLNKandinskyV22PriorPipelineConfig(RBLNModelConfig):
         *,
         batch_size: Optional[int] = None,
         guidance_scale: Optional[float] = None,
-        **kwargs,
+        **kwargs: Dict[str, Any],
     ):
         """
         Initialize a configuration for Kandinsky 2.2 prior pipeline optimized for RBLN NPU.
@@ -194,7 +204,9 @@ class RBLNKandinskyV22PriorPipelineConfig(RBLNModelConfig):
         return self.image_encoder.image_size
 
 
-class _RBLNKandinskyV22CombinedPipelineBaseConfig(RBLNModelConfig):
+class RBLNKandinskyV22CombinedPipelineBaseConfig(RBLNModelConfig):
+    """Base configuration class for Kandinsky V2.2 combined pipelines."""
+
     submodules = ["prior_pipe", "decoder_pipe"]
     _decoder_pipe_cls = RBLNKandinskyV22PipelineConfig
 
@@ -216,7 +228,7 @@ class _RBLNKandinskyV22CombinedPipelineBaseConfig(RBLNModelConfig):
         prior_text_encoder: Optional[RBLNCLIPTextModelWithProjectionConfig] = None,
         unet: Optional[RBLNUNet2DConditionModelConfig] = None,
         movq: Optional[RBLNVQModelConfig] = None,
-        **kwargs,
+        **kwargs: Dict[str, Any],
     ):
         """
         Initialize a configuration for combined Kandinsky 2.2 pipelines optimized for RBLN NPU.
@@ -325,13 +337,19 @@ class _RBLNKandinskyV22CombinedPipelineBaseConfig(RBLNModelConfig):
         return self.decoder_pipe.movq
 
 
-class RBLNKandinskyV22CombinedPipelineConfig(_RBLNKandinskyV22CombinedPipelineBaseConfig):
+class RBLNKandinskyV22CombinedPipelineConfig(RBLNKandinskyV22CombinedPipelineBaseConfig):
+    """Configuration class for the Kandinsky V2.2 combined text-to-image pipeline."""
+
     _decoder_pipe_cls = RBLNKandinskyV22PipelineConfig
 
 
-class RBLNKandinskyV22InpaintCombinedPipelineConfig(_RBLNKandinskyV22CombinedPipelineBaseConfig):
+class RBLNKandinskyV22InpaintCombinedPipelineConfig(RBLNKandinskyV22CombinedPipelineBaseConfig):
+    """Configuration class for the Kandinsky V2.2 combined inpainting pipeline."""
+
     _decoder_pipe_cls = RBLNKandinskyV22InpaintPipelineConfig
 
 
-class RBLNKandinskyV22Img2ImgCombinedPipelineConfig(_RBLNKandinskyV22CombinedPipelineBaseConfig):
+class RBLNKandinskyV22Img2ImgCombinedPipelineConfig(RBLNKandinskyV22CombinedPipelineBaseConfig):
+    """Configuration class for the Kandinsky V2.2 combined image-to-image pipeline."""
+
     _decoder_pipe_cls = RBLNKandinskyV22Img2ImgPipelineConfig

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

@@ -12,14 +12,18 @@
 # 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
 from ....transformers import RBLNCLIPTextModelConfig
 from ..models import RBLNAutoencoderKLConfig, RBLNUNet2DConditionModelConfig
 
 
-class _RBLNStableDiffusionPipelineBaseConfig(RBLNModelConfig):
+class RBLNStableDiffusionPipelineBaseConfig(RBLNModelConfig):
+    """
+    Base configuration for Stable Diffusion pipelines.
+    """
+
     submodules = ["text_encoder", "unet", "vae"]
     _vae_uses_encoder = False
 
@@ -37,7 +41,7 @@ class _RBLNStableDiffusionPipelineBaseConfig(RBLNModelConfig):
         sample_size: Optional[Tuple[int, int]] = None,
         image_size: Optional[Tuple[int, int]] = None,
         guidance_scale: Optional[float] = None,
-        **kwargs,
+        **kwargs: Dict[str, Any],
     ):
         """
         Args:
@@ -128,13 +132,25 @@ class _RBLNStableDiffusionPipelineBaseConfig(RBLNModelConfig):
         return self.vae.sample_size
 
 
-class RBLNStableDiffusionPipelineConfig(_RBLNStableDiffusionPipelineBaseConfig):
+class RBLNStableDiffusionPipelineConfig(RBLNStableDiffusionPipelineBaseConfig):
+    """
+    Configuration for Stable Diffusion pipeline.
+    """
+
     _vae_uses_encoder = False
 
 
-class RBLNStableDiffusionImg2ImgPipelineConfig(_RBLNStableDiffusionPipelineBaseConfig):
+class RBLNStableDiffusionImg2ImgPipelineConfig(RBLNStableDiffusionPipelineBaseConfig):
+    """
+    Configuration for Stable Diffusion image-to-image pipeline.
+    """
+
     _vae_uses_encoder = True
 
 
-class RBLNStableDiffusionInpaintPipelineConfig(_RBLNStableDiffusionPipelineBaseConfig):
+class RBLNStableDiffusionInpaintPipelineConfig(RBLNStableDiffusionPipelineBaseConfig):
+    """
+    Configuration for Stable Diffusion inpainting pipeline.
+    """
+
     _vae_uses_encoder = True