optimum-rbln 0.8.1a0__tar.gz → 0.8.1a2__tar.gz

optimum_rbln-0.8.1a2/.github/version.yaml

@@ -0,0 +1 @@
+rebel_compiler_version: 0.8.1.dev142+gab6ad3c7

{optimum_rbln-0.8.1a0 → optimum_rbln-0.8.1a2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: optimum-rbln
-Version: 0.8.1a0
+Version: 0.8.1a2
 Summary: Optimum RBLN is the interface between the HuggingFace Transformers and Diffusers libraries and RBLN accelerators. It provides a set of tools enabling easy model loading and inference on single and multiple rbln device settings for different downstream tasks.
 Project-URL: Homepage, https://rebellions.ai
 Project-URL: Documentation, https://docs.rbln.ai
@@ -28,7 +28,7 @@ Requires-Dist: packaging>=24.1
 Requires-Dist: torch==2.6.0
 Requires-Dist: torchaudio<=2.6.0
 Requires-Dist: torchvision<=0.21.0
-Requires-Dist: transformers==4.50.3
+Requires-Dist: transformers==4.51.3
 Description-Content-Type: text/markdown
 
 

optimum_rbln-0.8.1a2/examples/advanced/custom_class.py

@@ -0,0 +1,159 @@
+"""
+# Advanced Tutorial: Creating Custom RBLN Classes
+
+This tutorial demonstrates how to create custom RBLN model classes for use with the optimum-rbln framework.
+
+## Requirements for Custom Class Creation:
+
+1. You must create two classes:
+   - A custom model class that inherits from RBLNModel
+   - A custom config class that inherits from RBLNModelConfig
+
+2. The naming convention is critical:
+   - Model class: RBLN<OriginalModelName> (e.g., RBLNResNetModel)
+   - Config class: RBLN<OriginalModelName>Config (e.g., RBLNResNetModelConfig)
+
+3. Your custom class MUST implement the _update_rbln_config method
+   - This method configures input shapes and compilation settings
+
+4. You must register both classes with the auto registration system
+   - This enables the framework to discover and use your custom classes
+
+Important: Custom RBLN classes can only be created for models that exist in
+the transformers or diffusers libraries. You cannot create custom RBLN classes
+for completely custom architectures without a corresponding HuggingFace implementation.
+
+This example demonstrates creating a custom RBLN class for the ResNet model from transformers.
+"""
+
+from typing import TYPE_CHECKING, Optional, Tuple, Union
+
+import torch
+from transformers import ResNetModel  # noqa: F401
+from transformers.models.resnet.modeling_resnet import BaseModelOutputWithPoolingAndNoAttention
+
+from optimum.rbln import RBLNAutoConfig, RBLNAutoModel, RBLNCompileConfig, RBLNModel, RBLNModelConfig
+
+
+if TYPE_CHECKING:
+    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, PretrainedConfig, PreTrainedModel
+
+
+# --------------------------------------------------------
+# STEP 1: Create a custom model class that extends RBLNModel
+# --------------------------------------------------------
+# A custom RBLN model class must:
+# - Follow the naming pattern: RBLN<OriginalModelName>
+# - Implement the _update_rbln_config method (required)
+# - Define a proper forward method matching the original model's inputs/outputs
+#
+# The _update_rbln_config method is critical - it sets up the input tensor
+# specifications and other compilation parameters needed for the RBLN compiler.
+#
+# For our ResNet example, we need to:
+# 1. Specify the input tensor shape for images (batch_size, channels, height, width)
+# 2. Set the proper configuration options
+# 3. Return the updated RBLN configuration
+class RBLNResNetModel(RBLNModel):
+    @classmethod
+    def _update_rbln_config(
+        cls,
+        preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]] = None,
+        model: Optional["PreTrainedModel"] = None,
+        model_config: Optional["PretrainedConfig"] = None,
+        rbln_config: Optional["RBLNResNetModelConfig"] = None,
+    ) -> "RBLNResNetModelConfig":
+        # Set image_size if not provided
+        if rbln_config.image_size is None:
+            if rbln_config.image_size is None:
+                rbln_config.image_size = model_config.image_size
+
+            if rbln_config.image_size is None:
+                raise ValueError("`image_size` should be specified!")
+
+        # Define input tensor specification for the compiler
+        # Format: (tensor_name, tensor_shape, dtype)
+        input_info = [
+            (
+                "pixel_values",
+                [rbln_config.batch_size, 3, rbln_config.image_size[0], rbln_config.image_size[1]],
+                "float32",
+            )
+        ]
+
+        # Configure compilation settings
+        rbln_config.set_compile_cfgs([RBLNCompileConfig(input_info=input_info)])
+        return rbln_config
+
+    def forward(self, pixel_values, return_dict: Optional[bool] = None, **kwargs):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # self.model is a list of rebel.Runtime objects
+        # See https://docs.rbln.ai/software/api/python/python_api.html#rebel.rebel_runtime.Runtime for more details
+        output = self.model[0](pixel_values)
+
+        if not return_dict:
+            return output
+        else:
+            return BaseModelOutputWithPoolingAndNoAttention(last_hidden_state=output[0], pooler_output=output[1])
+
+
+# ----------------------------------------------------------------
+# STEP 2: Create a custom configuration class that extends RBLNModelConfig
+# ----------------------------------------------------------------
+# The configuration class must:
+# - Follow the naming pattern: RBLN<OriginalModelName>Config
+# - Define model-specific parameters needed for compilation
+#
+# For ResNet, we need:
+# - batch_size: Batch size for inference
+# - image_size: Input image dimensions (height, width)
+class RBLNResNetModelConfig(RBLNModelConfig):
+    def __init__(self, batch_size: int = None, image_size: Optional[Tuple[int, int]] = None, **kwargs):
+        super().__init__(**kwargs)
+        self.batch_size = batch_size or 1
+        self.image_size = image_size or (224, 224)
+
+
+# ----------------------------------------------------------------
+# STEP 3: Register both classes with the auto registration system
+# ----------------------------------------------------------------
+# Registration is essential for the RBLN framework to:
+# 1. Discover your custom classes when needed
+# 2. Enable using RBLNAutoModel.from_pretrained() with your model type
+# 3. Connect your configuration class to the model class in the framework
+RBLNAutoModel.register(RBLNResNetModel)
+RBLNAutoConfig.register(RBLNResNetModelConfig)
+
+
+# ----------------------------------------------------------------
+# STEP 4: Usage Example - Creating and using the custom RBLN model
+# ----------------------------------------------------------------
+# Initialize the model from a pretrained HuggingFace model
+# The 'export=True' parameter triggers the model compilation process
+# rbln_image_size and rbln_batch_size are passed to the custom configuration
+my_model = RBLNResNetModel.from_pretrained(
+    "microsoft/resnet-50", export=True, rbln_image_size=(224, 224), rbln_batch_size=1
+)
+
+# Save the compiled model for later use
+my_model.save_pretrained("my_resnet_model_saved")
+
+# Create a random input tensor for demonstration
+# The shape must match what we defined in _update_rbln_config
+random_image_input = torch.randn(1, 3, 224, 224)
+
+# Run inference with the compiled model
+output = my_model(random_image_input)
+
+# Print each key-value pair in the output dictionary
+for key, value in output.items():
+    print(key, value)
+
+
+# ----------------------------------------------------------------
+# STEP 5: Loading a saved RBLN model
+# ----------------------------------------------------------------
+# Load the model we just saved - no need to recompile
+my_model_reloaded = RBLNResNetModel.from_pretrained("my_resnet_model_saved")
+print(my_model_reloaded)

{optimum_rbln-0.8.1a0 → optimum_rbln-0.8.1a2}/pyproject.toml

@@ -32,7 +32,7 @@ dependencies = [
     "torchaudio<=2.6.0",
     "torchvision<=0.21.0",
     "accelerate>=1.0.1",
-    "transformers==4.50.3",
+    "transformers==4.51.3",
     "diffusers<=0.31.0",
     "packaging>=24.1",
 ]

{optimum_rbln-0.8.1a0 → optimum_rbln-0.8.1a2}/src/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-0.8.1a0 → optimum_rbln-0.8.1a2}/src/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-0.8.1a0 → optimum_rbln-0.8.1a2}/src/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-0.8.1a0 → optimum_rbln-0.8.1a2}/src/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-0.8.1a0 → optimum_rbln-0.8.1a2}/src/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-0.8.1a0 → optimum_rbln-0.8.1a2}/src/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-0.8.1a0 → optimum_rbln-0.8.1a2}/src/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-0.8.1a0 → optimum_rbln-0.8.1a2}/src/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-0.8.1a0 → optimum_rbln-0.8.1a2}/src/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-0.8.1a0 → optimum_rbln-0.8.1a2}/src/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