optimum-rbln 0.8.4a6__tar.gz → 0.8.4a8__tar.gz

{optimum_rbln-0.8.4a6 → optimum_rbln-0.8.4a8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: optimum-rbln
-Version: 0.8.4a6
+Version: 0.8.4a8
 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

{optimum_rbln-0.8.4a6 → optimum_rbln-0.8.4a8}/src/optimum/rbln/__version__.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.8.4a6'
-__version_tuple__ = version_tuple = (0, 8, 4, 'a6')
+__version__ = version = '0.8.4a8'
+__version_tuple__ = version_tuple = (0, 8, 4, 'a8')
 
 __commit_id__ = commit_id = None

{optimum_rbln-0.8.4a6 → optimum_rbln-0.8.4a8}/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 @@ 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.")
@@ -206,7 +242,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.")
@@ -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 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
             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 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
 
         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.
 

{optimum_rbln-0.8.4a6 → optimum_rbln-0.8.4a8}/src/optimum/rbln/diffusers/configurations/models/configuration_autoencoder_kl.py

@@ -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-0.8.4a6 → optimum_rbln-0.8.4a8}/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 @@ class RBLNAutoencoderKLCosmosConfig(RBLNModelConfig):
         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 @@ 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-0.8.4a6 → optimum_rbln-0.8.4a8}/src/optimum/rbln/diffusers/configurations/models/configuration_controlnet.py

@@ -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-0.8.4a6 → optimum_rbln-0.8.4a8}/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 @@ class RBLNPriorTransformerConfig(RBLNModelConfig):
             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-0.8.4a6 → optimum_rbln-0.8.4a8}/src/optimum/rbln/diffusers/configurations/models/configuration_transformer_cosmos.py

@@ -18,7 +18,12 @@ 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,
@@ -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-0.8.4a6 → optimum_rbln-0.8.4a8}/src/optimum/rbln/diffusers/configurations/models/configuration_transformer_sd3.py

@@ -18,7 +18,12 @@ 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"]
 
@@ -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-0.8.4a6 → optimum_rbln-0.8.4a8}/src/optimum/rbln/diffusers/configurations/models/configuration_unet_2d_condition.py

@@ -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-0.8.4a6 → optimum_rbln-0.8.4a8}/src/optimum/rbln/diffusers/configurations/models/configuration_vq_model.py

@@ -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.

{optimum_rbln-0.8.4a6 → optimum_rbln-0.8.4a8}/src/optimum/rbln/diffusers/configurations/pipelines/configuration_controlnet.py

@@ -59,7 +59,7 @@ class RBLNStableDiffusionControlNetPipelineBaseConfig(RBLNModelConfig):
             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 @@ class RBLNStableDiffusionXLControlNetPipelineBaseConfig(RBLNModelConfig):
             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.

{optimum_rbln-0.8.4a6 → optimum_rbln-0.8.4a8}/src/optimum/rbln/diffusers/configurations/pipelines/configuration_cosmos.py

@@ -59,7 +59,7 @@ class RBLNCosmosPipelineBaseConfig(RBLNModelConfig):
             num_frames (Optional[int]): The number of frames in the generated video.
             fps (Optional[int]): The frames per second of the generated video.
             max_seq_len (Optional[int]): Maximum sequence length supported by the model.
-            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+            kwargs: Additional arguments passed to the parent RBLNModelConfig.
         """
         super().__init__(**kwargs)
 

{optimum_rbln-0.8.4a6 → optimum_rbln-0.8.4a8}/src/optimum/rbln/diffusers/configurations/pipelines/configuration_kandinsky2_2.py

@@ -54,7 +54,7 @@ class RBLNKandinskyV22PipelineBaseConfig(RBLNModelConfig):
             img_width (Optional[int]): Width of the generated images.
             height (Optional[int]): Height of the generated images.
             width (Optional[int]): Width of the generated images.
-            **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.
@@ -166,7 +166,7 @@ class RBLNKandinskyV22PriorPipelineConfig(RBLNModelConfig):
                 Initialized as RBLNPriorTransformerConfig if not provided.
             batch_size (Optional[int]): Batch size for inference, applied to all submodules.
             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.
 
         Note:
             When guidance_scale > 1.0, the prior batch size is automatically doubled to
@@ -259,7 +259,7 @@ class RBLNKandinskyV22CombinedPipelineBaseConfig(RBLNModelConfig):
                 Used if decoder_pipe is not provided.
             movq (Optional[RBLNVQModelConfig]): Direct configuration for the MoVQ (VQ-GAN) model.
                 Used if decoder_pipe is not provided.
-            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+            kwargs: Additional arguments passed to the parent RBLNModelConfig.
         """
         super().__init__(**kwargs)
 

{optimum_rbln-0.8.4a6 → optimum_rbln-0.8.4a8}/src/optimum/rbln/diffusers/configurations/pipelines/configuration_stable_diffusion.py

@@ -56,7 +56,7 @@ class RBLNStableDiffusionPipelineBaseConfig(RBLNModelConfig):
             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.

{optimum_rbln-0.8.4a6 → optimum_rbln-0.8.4a8}/src/optimum/rbln/diffusers/configurations/pipelines/configuration_stable_diffusion_3.py

@@ -64,7 +64,7 @@ class RBLNStableDiffusion3PipelineBaseConfig(RBLNModelConfig):
             height (Optional[int]): Height of the generated images.
             width (Optional[int]): Width of the generated images.
             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.

{optimum_rbln-0.8.4a6 → optimum_rbln-0.8.4a8}/src/optimum/rbln/diffusers/configurations/pipelines/configuration_stable_diffusion_xl.py

@@ -59,7 +59,7 @@ class RBLNStableDiffusionXLPipelineBaseConfig(RBLNModelConfig):
             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.

{optimum_rbln-0.8.4a6 → optimum_rbln-0.8.4a8}/src/optimum/rbln/diffusers/modeling_diffusers.py

@@ -33,6 +33,10 @@ if TYPE_CHECKING:
 
 
 class RBLNDiffusionMixinConfig(RBLNModelConfig):
+    """
+    Configuration class for RBLN diffusion pipelines.
+    """
+
     pass
 
 
@@ -54,8 +58,8 @@ class RBLNDiffusionMixin:
         ```
 
     Class Variables:
-        _submodules: List of submodule names that should be compiled (typically ["text_encoder", "unet", "vae"])
-        _optional_submodules: List of submodule names compiled without inheriting RBLNModel (typically ["safety_checker"])
+        - `_submodules`: List of submodule names that should be compiled (typically ["text_encoder", "unet", "vae"])
+        - `_optional_submodules`: List of submodule names compiled without inheriting RBLNModel (typically ["safety_checker"])
 
     Methods:
         from_pretrained: Creates and optionally compiles a model from a pretrained checkpoint
@@ -170,7 +174,7 @@ class RBLNDiffusionMixin:
                 Names of specific LoRA weight files to load, corresponding to lora_ids. Only used when `export=True`.
             lora_scales:
                 Scaling factor(s) to apply to the LoRA adapter(s). Only used when `export=True`.
-            **kwargs:
+            kwargs:
                 Additional arguments to pass to the underlying diffusion pipeline constructor or the
                 RBLN compilation process. These may include parameters specific to individual submodules
                 or the particular diffusion pipeline being used.

{optimum_rbln-0.8.4a6 → optimum_rbln-0.8.4a8}/src/optimum/rbln/diffusers/models/autoencoders/autoencoder_kl.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import TYPE_CHECKING, Dict, List, Tuple, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Tuple, Union
 
 import rebel
 import torch
@@ -214,13 +214,41 @@ class RBLNAutoencoderKL(RBLNModel):
             for compiled_model, device_val in zip(compiled_models, device_vals)
         ]
 
-    def encode(self, x: torch.FloatTensor, return_dict: bool = True, **kwargs) -> torch.FloatTensor:
+    def encode(
+        self, x: torch.FloatTensor, return_dict: bool = True, **kwargs: Dict[str, Any]
+    ) -> Union[torch.FloatTensor, AutoencoderKLOutput]:
+        """
+        Encode an input image into a latent representation.
+
+        Args:
+            x: The input image to encode.
+            return_dict:
+                Whether to return output as a dictionary. Defaults to True.
+            kwargs: Additional arguments to pass to the encoder.
+
+        Returns:
+            The latent representation or AutoencoderKLOutput if return_dict=True
+        """
         posterior = self.encoder.encode(x)
         if not return_dict:
             return (posterior,)
         return AutoencoderKLOutput(latent_dist=posterior)
 
-    def decode(self, z: torch.FloatTensor, return_dict: bool = True, **kwargs) -> torch.FloatTensor:
+    def decode(
+        self, z: torch.FloatTensor, return_dict: bool = True, **kwargs: Dict[str, Any]
+    ) -> Union[torch.FloatTensor, DecoderOutput]:
+        """
+        Decode a latent representation into an image.
+
+        Args:
+            z: The latent representation to decode.
+            return_dict:
+                Whether to return output as a dictionary. Defaults to True.
+            kwargs: Additional arguments to pass to the decoder.
+
+        Returns:
+            The decoded image or DecoderOutput if return_dict=True
+        """
         dec = self.decoder.decode(z)
         if not return_dict:
             return (dec,)

{optimum_rbln-0.8.4a6 → optimum_rbln-0.8.4a8}/src/optimum/rbln/diffusers/models/autoencoders/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 TYPE_CHECKING, Dict, List, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Union
 
 import rebel
 import torch
@@ -205,13 +205,38 @@ class RBLNAutoencoderKLCosmos(RBLNModel):
             for compiled_model, device_val in zip(compiled_models, device_vals)
         ]
 
-    def encode(self, x: torch.FloatTensor, return_dict: bool = True, **kwargs) -> torch.FloatTensor:
+    def encode(
+        self, x: torch.FloatTensor, return_dict: bool = True, **kwargs: Dict[str, Any]
+    ) -> Union[torch.FloatTensor, AutoencoderKLOutput]:
+        """
+        Encode an input video into a latent representation.
+
+        Args:
+            x: The input video to encode.
+            return_dict:
+                Whether to return output as a dictionary. Defaults to True.
+            kwargs: Additional arguments to pass to the encoder.
+
+        Returns:
+            The latent representation or AutoencoderKLOutput if return_dict=True
+        """
         posterior = self.encoder.encode(x)
         if not return_dict:
             return (posterior,)
         return AutoencoderKLOutput(latent_dist=posterior)
 
-    def decode(self, z: torch.FloatTensor, return_dict: bool = True) -> torch.FloatTensor:
+    def decode(self, z: torch.FloatTensor, return_dict: bool = True) -> Union[torch.FloatTensor, DecoderOutput]:
+        """
+        Decode a latent representation into a video.
+
+        Args:
+            z: The latent representation to decode.
+            return_dict:
+                Whether to return output as a dictionary. Defaults to True.
+
+        Returns:
+            The decoded video or DecoderOutput if return_dict=True
+        """
         decoded = self.decoder.decode(z)
 
         if not return_dict:

{optimum_rbln-0.8.4a6 → optimum_rbln-0.8.4a8}/src/optimum/rbln/diffusers/models/autoencoders/vq_model.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import TYPE_CHECKING, List, Union
+from typing import TYPE_CHECKING, Any, List, Union
 
 import rebel
 import torch
@@ -170,13 +170,41 @@ class RBLNVQModel(RBLNModel):
             for compiled_model, device_val in zip(compiled_models, device_vals)
         ]
 
-    def encode(self, x: torch.FloatTensor, return_dict: bool = True, **kwargs) -> torch.FloatTensor:
+    def encode(
+        self, x: torch.FloatTensor, return_dict: bool = True, **kwargs: Any
+    ) -> Union[torch.FloatTensor, VQEncoderOutput]:
+        """
+        Encode an input image into a quantized latent representation.
+
+        Args:
+            x: The input image to encode.
+            return_dict:
+                Whether to return output as a dictionary. Defaults to True.
+            kwargs: Additional arguments to pass to the encoder/quantizer.
+
+        Returns:
+            The quantized latent representation or a specific output object.
+        """
         posterior = self.encoder.encode(x)
         if not return_dict:
             return (posterior,)
         return VQEncoderOutput(latents=posterior)
 
-    def decode(self, h: torch.FloatTensor, return_dict: bool = True, **kwargs) -> torch.FloatTensor:
+    def decode(
+        self, h: torch.FloatTensor, return_dict: bool = True, **kwargs: Any
+    ) -> Union[torch.FloatTensor, DecoderOutput]:
+        """
+        Decode a quantized latent representation back into an image.
+
+        Args:
+            h: The quantized latent representation to decode.
+            return_dict:
+                Whether to return output as a dictionary. Defaults to True.
+            kwargs: Additional arguments to pass to the decoder.
+
+        Returns:
+            The decoded image or a DecoderOutput object.
+        """
         dec, commit_loss = self.decoder.decode(h, **kwargs)
         if not return_dict:
             return (dec, commit_loss)

{optimum_rbln-0.8.4a6 → optimum_rbln-0.8.4a8}/src/optimum/rbln/diffusers/models/transformers/prior_transformer.py

@@ -59,7 +59,7 @@ class RBLNPriorTransformer(RBLNModel):
     """
     RBLN implementation of PriorTransformer for diffusion models like Kandinsky V2.2.
 
-    The Prior Transformer takes text and/or image embeddings from encoders (like CLIP) and
+    The PriorTransformer takes text and/or image embeddings from encoders (like CLIP) and
     maps them to a shared latent space that guides the diffusion process to generate the desired image.
 
     This class inherits from [`RBLNModel`]. Check the superclass documentation for the generic methods

{optimum_rbln-0.8.4a6 → optimum_rbln-0.8.4a8}/src/optimum/rbln/diffusers/models/transformers/transformer_cosmos.py

@@ -94,7 +94,15 @@ class CosmosTransformer3DModelWrapper(torch.nn.Module):
 
 
 class RBLNCosmosTransformer3DModel(RBLNModel):
-    """RBLN wrapper for the Cosmos Transformer model."""
+    """
+    RBLN implementation of CosmosTransformer3DModel for diffusion models like Cosmos.
+
+    The CosmosTransformer3DModel takes text and/or image embeddings from encoders (like CLIP) and
+    maps them to a shared latent space that guides the diffusion process to generate the desired image.
+
+    This class inherits from [`RBLNModel`]. Check the superclass documentation for the generic methods
+    the library implements for all its models.
+    """
 
     hf_library_name = "diffusers"
     auto_model_class = CosmosTransformer3DModel

{optimum_rbln-0.8.4a6 → optimum_rbln-0.8.4a8}/src/optimum/rbln/diffusers/models/transformers/transformer_sd3.py

@@ -59,7 +59,15 @@ class SD3Transformer2DModelWrapper(torch.nn.Module):
 
 
 class RBLNSD3Transformer2DModel(RBLNModel):
-    """RBLN wrapper for the Stable Diffusion 3 MMDiT Transformer model."""
+    """
+    RBLN implementation of SD3Transformer2DModel for diffusion models like Stable Diffusion 3.
+
+    The SD3Transformer2DModel takes text and/or image embeddings from encoders (like CLIP) and
+    maps them to a shared latent space that guides the diffusion process to generate the desired image.
+
+    This class inherits from [`RBLNModel`]. Check the superclass documentation for the generic methods
+    the library implements for all its models.
+    """
 
     hf_library_name = "diffusers"
     auto_model_class = SD3Transformer2DModel

{optimum_rbln-0.8.4a6 → optimum_rbln-0.8.4a8}/src/optimum/rbln/diffusers/models/unets/unet_2d_condition.py

@@ -141,10 +141,13 @@ class _UNet_Kandinsky(torch.nn.Module):
 
 class RBLNUNet2DConditionModel(RBLNModel):
     """
-    Configuration class for RBLN UNet2DCondition models.
+    RBLN implementation of UNet2DConditionModel for diffusion models.
 
-    This class inherits from RBLNModelConfig and provides specific configuration options
-    for UNet2DCondition models used in diffusion-based image generation.
+    This model is used to accelerate UNet2DCondition models from diffusers library on RBLN NPUs.
+    It is a key component in diffusion-based image generation models like Stable Diffusion.
+
+    This class inherits from [`RBLNModel`]. Check the superclass documentation for the generic methods
+    the library implements for all its models.
     """
 
     hf_library_name = "diffusers"