other: update docstrings and minor updates (#270)

Co-authored-by: rebel-junamsong <[email protected]>
Co-authored-by: rebel-seinpark <[email protected]>
Co-authored-by: rebel-thkim <[email protected]>
Co-authored-by: thkim <[email protected]>
Co-authored-by: rebel-dkhong <[email protected]>
Co-authored-by: rebel-kblee <[email protected]>

Author: 6 people

src/optimum/rbln/configuration_utils.py

@@ -185,6 +185,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:
@@ -194,6 +203,33 @@ def __new__(cls, **kwargs):
 
     @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.")
@@ -206,7 +242,8 @@ def register(config: Type["RBLNModelConfig"], exist_ok=False):
         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.")
@@ -282,6 +319,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
@@ -594,7 +632,7 @@ def __init__(
             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.
@@ -761,7 +799,7 @@ def load(cls, path: str, **kwargs: Any) -> "RBLNModelConfig":
 
         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.
 

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

@@ -46,7 +46,7 @@ def __init__(
                 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.

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

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Optional
+from typing import Any, Dict, Optional
 
 from ....configuration_utils import RBLNModelConfig
 from ....utils.logging import get_logger
@@ -35,7 +35,7 @@ def __init__(
         vae_scale_factor_temporal: Optional[int] = None,
         vae_scale_factor_spatial: Optional[int] = None,
         use_slicing: Optional[bool] = None,
-        **kwargs: Any,
+        **kwargs: Dict[str, Any],
     ):
         """
         Args:
@@ -52,7 +52,7 @@ def __init__(
                 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.

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

@@ -42,7 +42,7 @@ def __init__(
                 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.

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

@@ -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"]
@@ -39,7 +39,7 @@ def __init__(
             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.

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

@@ -18,7 +18,12 @@
 
 
 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,
@@ -47,7 +52,7 @@ def __init__(
             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.

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

@@ -18,7 +18,12 @@
 
 
 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"]
 
@@ -36,7 +41,7 @@ def __init__(
                 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.

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

@@ -52,7 +52,7 @@ def __init__(
             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.

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

@@ -46,7 +46,7 @@ def __init__(
                 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.

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

@@ -59,7 +59,7 @@ def __init__(
             image_size (Optional[Tuple[int, int]]): Alternative way to specify image dimensions.
                 Cannot be used together with img_height/img_width.
             guidance_scale (Optional[float]): Scale for classifier-free guidance.
-            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+            kwargs: Additional arguments passed to the parent RBLNModelConfig.
 
         Raises:
             ValueError: If both image_size and img_height/img_width are provided.
@@ -201,7 +201,7 @@ def __init__(
             image_size (Optional[Tuple[int, int]]): Alternative way to specify image dimensions.
                 Cannot be used together with img_height/img_width.
             guidance_scale (Optional[float]): Scale for classifier-free guidance.
-            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+            kwargs: Additional arguments passed to the parent RBLNModelConfig.
 
         Raises:
             ValueError: If both image_size and img_height/img_width are provided.