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

optimum/rbln/diffusers/configurations/pipelines/configuration_stable_diffusion_3.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 RBLNCLIPTextModelWithProjectionConfig, RBLNT5EncoderModelConfig
 from ..models import RBLNAutoencoderKLConfig, RBLNSD3Transformer2DModelConfig
 
 
-class _RBLNStableDiffusion3PipelineBaseConfig(RBLNModelConfig):
+class RBLNStableDiffusion3PipelineBaseConfig(RBLNModelConfig):
+    """
+    Base configuration for Stable Diffusion 3 pipelines.
+    """
+
     submodules = ["transformer", "text_encoder", "text_encoder_2", "text_encoder_3", "vae"]
     _vae_uses_encoder = False
 
@@ -40,7 +44,7 @@ class _RBLNStableDiffusion3PipelineBaseConfig(RBLNModelConfig):
         height: Optional[int] = None,
         width: Optional[int] = None,
         guidance_scale: Optional[float] = None,
-        **kwargs,
+        **kwargs: Dict[str, Any],
     ):
         """
         Args:
@@ -153,13 +157,19 @@ class _RBLNStableDiffusion3PipelineBaseConfig(RBLNModelConfig):
         return self.vae.sample_size
 
 
-class RBLNStableDiffusion3PipelineConfig(_RBLNStableDiffusion3PipelineBaseConfig):
+class RBLNStableDiffusion3PipelineConfig(RBLNStableDiffusion3PipelineBaseConfig):
+    """Config for SD3 Text2Img Pipeline"""
+
     _vae_uses_encoder = False
 
 
-class RBLNStableDiffusion3Img2ImgPipelineConfig(_RBLNStableDiffusion3PipelineBaseConfig):
+class RBLNStableDiffusion3Img2ImgPipelineConfig(RBLNStableDiffusion3PipelineBaseConfig):
+    """Config for SD3 Img2Img Pipeline"""
+
     _vae_uses_encoder = True
 
 
-class RBLNStableDiffusion3InpaintPipelineConfig(_RBLNStableDiffusion3PipelineBaseConfig):
+class RBLNStableDiffusion3InpaintPipelineConfig(RBLNStableDiffusion3PipelineBaseConfig):
+    """Config for SD3 Inpainting Pipeline"""
+
     _vae_uses_encoder = True

optimum/rbln/diffusers/configurations/pipelines/configuration_stable_diffusion_xl.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, RBLNUNet2DConditionModelConfig
 
 
-class _RBLNStableDiffusionXLPipelineBaseConfig(RBLNModelConfig):
+class RBLNStableDiffusionXLPipelineBaseConfig(RBLNModelConfig):
+    """
+    Base configuration for Stable Diffusion XL pipelines.
+    """
+
     submodules = ["text_encoder", "text_encoder_2", "unet", "vae"]
     _vae_uses_encoder = False
 
@@ -38,7 +42,7 @@ class _RBLNStableDiffusionXLPipelineBaseConfig(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:
@@ -134,13 +138,19 @@ class _RBLNStableDiffusionXLPipelineBaseConfig(RBLNModelConfig):
         return self.vae.sample_size
 
 
-class RBLNStableDiffusionXLPipelineConfig(_RBLNStableDiffusionXLPipelineBaseConfig):
+class RBLNStableDiffusionXLPipelineConfig(RBLNStableDiffusionXLPipelineBaseConfig):
+    """Config for SDXL Text2Img Pipeline"""
+
     _vae_uses_encoder = False
 
 
-class RBLNStableDiffusionXLImg2ImgPipelineConfig(_RBLNStableDiffusionXLPipelineBaseConfig):
+class RBLNStableDiffusionXLImg2ImgPipelineConfig(RBLNStableDiffusionXLPipelineBaseConfig):
+    """Config for SDXL Img2Img Pipeline"""
+
     _vae_uses_encoder = True
 
 
-class RBLNStableDiffusionXLInpaintPipelineConfig(_RBLNStableDiffusionXLPipelineBaseConfig):
+class RBLNStableDiffusionXLInpaintPipelineConfig(RBLNStableDiffusionXLPipelineBaseConfig):
+    """Config for SDXL Inpainting Pipeline"""
+
     _vae_uses_encoder = True

optimum/rbln/diffusers/modeling_diffusers.py

@@ -19,10 +19,11 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type, Union
 
 import torch
 
-from ..configuration_utils import ContextRblnConfig, RBLNModelConfig
+from ..configuration_utils import ContextRblnConfig, RBLNModelConfig, get_rbln_config_class
 from ..modeling import RBLNModel
 from ..utils.decorator_utils import remove_compile_time_kwargs
 from ..utils.logging import get_logger
+from ..utils.model_utils import get_rbln_model_cls
 
 
 logger = get_logger(__name__)
@@ -110,18 +111,10 @@ class RBLNDiffusionMixin:
 
     @classmethod
     def get_rbln_config_class(cls) -> Type[RBLNModelConfig]:
-        """
-        Lazily loads and caches the corresponding RBLN model config class.
-        """
+        # Lazily loads and caches the corresponding RBLN model config class.
         if cls._rbln_config_class is None:
             rbln_config_class_name = cls.__name__ + "Config"
-            library = importlib.import_module("optimum.rbln")
-            cls._rbln_config_class = getattr(library, rbln_config_class_name, None)
-            if cls._rbln_config_class is None:
-                raise ValueError(
-                    f"RBLN config class {rbln_config_class_name} not found. This is an internal error. "
-                    "Please report it to the developers."
-                )
+            cls._rbln_config_class = get_rbln_config_class(rbln_config_class_name)
         return cls._rbln_config_class
 
     @classmethod
@@ -143,7 +136,7 @@ class RBLNDiffusionMixin:
         lora_ids: Optional[Union[str, List[str]]] = None,
         lora_weights_names: Optional[Union[str, List[str]]] = None,
         lora_scales: Optional[Union[float, List[float]]] = None,
-        **kwargs,
+        **kwargs: Dict[str, Any],
     ) -> "RBLNDiffusionMixin":
         """
         Load a pretrained diffusion pipeline from a model checkpoint, with optional compilation for RBLN NPUs.
@@ -157,24 +150,25 @@ class RBLNDiffusionMixin:
         Args:
             model_id (`str`):
                 The model ID or path to the pretrained model to load. Can be either:
+
                 - A model ID from the HuggingFace Hub
                 - A local path to a saved model directory
-            export (`bool`, *optional*, defaults to `False`):
+            export:
                 If True, takes a PyTorch model from `model_id` and compiles it for RBLN NPU execution.
                 If False, loads an already compiled RBLN model from `model_id` without recompilation.
-            model_save_dir (`os.PathLike`, *optional*):
+            model_save_dir:
                 Directory to save the compiled model artifacts. Only used when `export=True`.
                 If not provided and `export=True`, a temporary directory is used.
-            rbln_config (`Dict[str, Any]`, *optional*, defaults to `{}`):
+            rbln_config:
                 Configuration options for RBLN compilation. Can include settings for specific submodules
                 such as `text_encoder`, `unet`, and `vae`. Configuration can be tailored to the specific
                 pipeline being compiled.
-            lora_ids (`str` or `List[str]`, *optional*):
+            lora_ids:
                 LoRA adapter ID(s) to load and apply before compilation. LoRA weights are fused
                 into the model weights during compilation. Only used when `export=True`.
-            lora_weights_names (`str` or `List[str]`, *optional*):
+            lora_weights_names:
                 Names of specific LoRA weight files to load, corresponding to lora_ids. Only used when `export=True`.
-            lora_scales (`float` or `List[float]`, *optional*):
+            lora_scales:
                 Scaling factor(s) to apply to the LoRA adapter(s). Only used when `export=True`.
             **kwargs:
                 Additional arguments to pass to the underlying diffusion pipeline constructor or the
@@ -182,8 +176,8 @@ class RBLNDiffusionMixin:
                 or the particular diffusion pipeline being used.
 
         Returns:
-            `RBLNDiffusionMixin`: A compiled or loaded diffusion pipeline that can be used for inference on RBLN NPU.
-            The returned object is an instance of the class that called this method, inheriting from RBLNDiffusionMixin.
+            A compiled or loaded diffusion pipeline that can be used for inference on RBLN NPU.
+                The returned object is an instance of the class that called this method, inheriting from RBLNDiffusionMixin.
         """
         rbln_config, kwargs = cls.get_rbln_config_class().initialize_from_kwargs(rbln_config, **kwargs)
 
@@ -210,7 +204,7 @@ class RBLNDiffusionMixin:
                         "Expected 'optimum.rbln'. Please check the model_index.json configuration."
                     )
 
-                submodule_cls: Type[RBLNModel] = getattr(importlib.import_module("optimum.rbln"), class_name)
+                submodule_cls = get_rbln_model_cls(class_name)
                 submodule_config = getattr(rbln_config, submodule_name)
                 submodule = submodule_cls.from_pretrained(
                     model_id, export=False, subfolder=submodule_name, rbln_config=submodule_config
@@ -293,7 +287,6 @@ class RBLNDiffusionMixin:
             elif isinstance(submodule, RBLNModel):
                 pass
             elif submodule_name == "controlnet" and hasattr(submodule, "nets"):
-                # In case of multicontrolnet
                 submodule = cls._compile_multicontrolnet(
                     controlnets=submodule,
                     model_save_dir=model_save_dir,
@@ -301,11 +294,8 @@ class RBLNDiffusionMixin:
                     prefix=prefix,
                 )
             elif isinstance(submodule, torch.nn.Module):
-                submodule_cls: RBLNModel = getattr(
-                    importlib.import_module("optimum.rbln"), f"RBLN{submodule.__class__.__name__}"
-                )
                 subfolder = prefix + submodule_name
-                submodule = submodule_cls.from_model(
+                submodule = submodule_rbln_cls.from_model(
                     model=submodule,
                     subfolder=subfolder,
                     model_save_dir=model_save_dir,

optimum/rbln/diffusers/models/autoencoders/autoencoder_kl.py

@@ -38,6 +38,17 @@ logger = get_logger(__name__)
 
 
 class RBLNAutoencoderKL(RBLNModel):
+    """
+    RBLN implementation of AutoencoderKL (VAE) for diffusion models.
+
+    This model is used to accelerate AutoencoderKL (VAE) models from diffusers library on RBLN NPUs.
+    It can be configured to include both encoder and decoder, or just the decoder part for latent-to-image
+    conversion.
+
+    This class inherits from [`RBLNModel`]. Check the superclass documentation for the generic methods
+    the library implements for all its models.
+    """
+
     auto_model_class = AutoencoderKL
     hf_library_name = "diffusers"
     _rbln_config_class = RBLNAutoencoderKLConfig

optimum/rbln/diffusers/models/autoencoders/vae.py

@@ -12,22 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import TYPE_CHECKING, List
+from typing import List
 
 import torch
 from diffusers import AutoencoderKL, VQModel
 from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution
 
-from ....utils.logging import get_logger
 from ....utils.runtime_utils import RBLNPytorchRuntime
 
 
-if TYPE_CHECKING:
-    import torch
-
-logger = get_logger(__name__)
-
-
 class RBLNRuntimeVAEEncoder(RBLNPytorchRuntime):
     def encode(self, x: torch.FloatTensor, **kwargs) -> torch.FloatTensor:
         moments = self.forward(x.contiguous())

optimum/rbln/diffusers/models/autoencoders/vq_model.py

@@ -35,6 +35,17 @@ logger = get_logger(__name__)
 
 
 class RBLNVQModel(RBLNModel):
+    """
+    RBLN implementation of VQModel for diffusion models.
+
+    This model is used to accelerate VQModel models from diffusers library on RBLN NPUs.
+    It can be configured to include both encoder and decoder, or just the decoder part for latent-to-image
+    conversion.
+
+    This class inherits from [`RBLNModel`]. Check the superclass documentation for the generic methods
+    the library implements for all its models.
+    """
+
     auto_model_class = VQModel
     config_name = "config.json"
     hf_library_name = "diffusers"

optimum/rbln/diffusers/models/controlnet.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import importlib
 from typing import TYPE_CHECKING, Dict, Optional, Union
 
 import torch
@@ -23,6 +22,7 @@ from transformers import PretrainedConfig
 from ...configuration_utils import RBLNCompileConfig, RBLNModelConfig
 from ...modeling import RBLNModel
 from ...utils.logging import get_logger
+from ...utils.model_utils import get_rbln_model_cls
 from ..configurations import RBLNControlNetModelConfig
 from ..modeling_diffusers import RBLNDiffusionMixin, RBLNDiffusionMixinConfig
 
@@ -98,6 +98,15 @@ class _ControlNetModel_Cross_Attention(torch.nn.Module):
 
 
 class RBLNControlNetModel(RBLNModel):
+    """
+    RBLN implementation of ControlNetModel for diffusion models.
+
+    This model is used to accelerate ControlNetModel models from diffusers library on RBLN NPUs.
+
+    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 = ControlNetModel
     output_class = ControlNetOutput
@@ -122,13 +131,10 @@ class RBLNControlNetModel(RBLNModel):
 
     @classmethod
     def update_rbln_config_using_pipe(
-        cls,
-        pipe: RBLNDiffusionMixin,
-        rbln_config: "RBLNDiffusionMixinConfig",
-        submodule_name: str,
+        cls, pipe: RBLNDiffusionMixin, rbln_config: "RBLNDiffusionMixinConfig", submodule_name: str
     ) -> "RBLNDiffusionMixinConfig":
-        rbln_vae_cls = getattr(importlib.import_module("optimum.rbln"), f"RBLN{pipe.vae.__class__.__name__}")
-        rbln_unet_cls = getattr(importlib.import_module("optimum.rbln"), f"RBLN{pipe.unet.__class__.__name__}")
+        rbln_vae_cls = get_rbln_model_cls(f"RBLN{pipe.vae.__class__.__name__}")
+        rbln_unet_cls = get_rbln_model_cls(f"RBLN{pipe.unet.__class__.__name__}")
 
         rbln_config.controlnet.max_seq_len = pipe.text_encoder.config.max_position_embeddings
         text_model_hidden_size = pipe.text_encoder_2.config.hidden_size if hasattr(pipe, "text_encoder_2") else None

optimum/rbln/diffusers/models/transformers/prior_transformer.py

@@ -56,6 +56,16 @@ class _PriorTransformer(torch.nn.Module):
 
 
 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
+    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 = PriorTransformer
     _output_class = PriorTransformerOutput

optimum/rbln/diffusers/models/transformers/transformer_sd3.py

@@ -59,6 +59,8 @@ class SD3Transformer2DModelWrapper(torch.nn.Module):
 
 
 class RBLNSD3Transformer2DModel(RBLNModel):
+    """RBLN wrapper for the Stable Diffusion 3 MMDiT Transformer model."""
+
     hf_library_name = "diffusers"
     auto_model_class = SD3Transformer2DModel
     _output_class = Transformer2DModelOutput

optimum/rbln/diffusers/models/unets/unet_2d_condition.py

@@ -140,6 +140,13 @@ class _UNet_Kandinsky(torch.nn.Module):
 
 
 class RBLNUNet2DConditionModel(RBLNModel):
+    """
+    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.
+    """
+
     hf_library_name = "diffusers"
     auto_model_class = UNet2DConditionModel
     _rbln_config_class = RBLNUNet2DConditionModelConfig

optimum/rbln/diffusers/pipelines/controlnet/multicontrolnet.py

@@ -14,7 +14,7 @@
 
 import os
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, Union
 
 import torch
 from diffusers.pipelines.controlnet.multicontrolnet import MultiControlNetModel
@@ -24,9 +24,6 @@ from ....utils.logging import get_logger
 from ...models.controlnet import RBLNControlNetModel
 
 
-if TYPE_CHECKING:
-    pass
-
 logger = get_logger(__name__)
 
 

optimum/rbln/diffusers/pipelines/controlnet/pipeline_controlnet.py

@@ -49,6 +49,13 @@ logger = get_logger(__name__)
 
 
 class RBLNStableDiffusionControlNetPipeline(RBLNDiffusionMixin, StableDiffusionControlNetPipeline):
+    """
+    RBLN-accelerated implementation of Stable Diffusion pipeline with ControlNet for guided text-to-image generation.
+
+    This pipeline compiles Stable Diffusion and ControlNet models to run efficiently on RBLN NPUs, enabling high-performance
+    inference for generating images with precise structural control using conditioning inputs like edges, depth, or poses.
+    """
+
     original_class = StableDiffusionControlNetPipeline
     _rbln_config_class = RBLNStableDiffusionControlNetPipelineConfig
     _submodules = ["text_encoder", "unet", "vae", "controlnet"]

optimum/rbln/diffusers/pipelines/controlnet/pipeline_controlnet_img2img.py

@@ -47,6 +47,13 @@ logger = logging.get_logger(__name__)
 
 
 class RBLNStableDiffusionControlNetImg2ImgPipeline(RBLNDiffusionMixin, StableDiffusionControlNetImg2ImgPipeline):
+    """
+    RBLN-accelerated implementation of Stable Diffusion pipeline with ControlNet for guided image-to-image generation.
+
+    This pipeline compiles Stable Diffusion and ControlNet models to run efficiently on RBLN NPUs, enabling high-performance
+    inference for transforming input images with precise structural control and conditioning guidance.
+    """
+
     original_class = StableDiffusionControlNetImg2ImgPipeline
     _submodules = ["text_encoder", "unet", "vae", "controlnet"]
     _rbln_config_class = RBLNStableDiffusionControlNetImg2ImgPipelineConfig

optimum/rbln/diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl.py

@@ -47,6 +47,13 @@ logger = logging.get_logger(__name__)
 
 
 class RBLNStableDiffusionXLControlNetPipeline(RBLNDiffusionMixin, StableDiffusionXLControlNetPipeline):
+    """
+    RBLN-accelerated implementation of Stable Diffusion XL pipeline with ControlNet for high-resolution guided text-to-image generation.
+
+    This pipeline compiles Stable Diffusion XL and ControlNet models to run efficiently on RBLN NPUs, enabling high-performance
+    inference for generating high-quality images with precise structural control and enhanced detail preservation.
+    """
+
     original_class = StableDiffusionXLControlNetPipeline
     _rbln_config_class = RBLNStableDiffusionXLControlNetPipelineConfig
     _submodules = ["text_encoder", "text_encoder_2", "unet", "vae", "controlnet"]

optimum/rbln/diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl_img2img.py

@@ -47,6 +47,13 @@ logger = logging.get_logger(__name__)
 
 
 class RBLNStableDiffusionXLControlNetImg2ImgPipeline(RBLNDiffusionMixin, StableDiffusionXLControlNetImg2ImgPipeline):
+    """
+    RBLN-accelerated implementation of Stable Diffusion XL pipeline with ControlNet for high-resolution guided image-to-image generation.
+
+    This pipeline compiles Stable Diffusion XL and ControlNet models to run efficiently on RBLN NPUs, enabling high-performance
+    inference for transforming input images with precise structural control and enhanced quality preservation.
+    """
+
     original_class = StableDiffusionXLControlNetImg2ImgPipeline
     _rbln_config_class = RBLNStableDiffusionXLControlNetImg2ImgPipelineConfig
     _submodules = ["text_encoder", "text_encoder_2", "unet", "vae", "controlnet"]

optimum/rbln/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2.py

@@ -19,6 +19,13 @@ from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNKandinskyV22Pipeline(RBLNDiffusionMixin, KandinskyV22Pipeline):
+    """
+    RBLN-accelerated implementation of Kandinsky 2.2 pipeline for text-to-image generation.
+
+    This pipeline compiles Kandinsky 2.2 models to run efficiently on RBLN NPUs, enabling high-performance
+    inference for generating images with distinctive artistic style and enhanced visual quality.
+    """
+
     original_class = KandinskyV22Pipeline
     _rbln_config_class = RBLNKandinskyV22PipelineConfig
     _submodules = ["unet", "movq"]

optimum/rbln/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_combined.py

@@ -38,6 +38,13 @@ from .pipeline_kandinsky2_2_prior import RBLNKandinskyV22PriorPipeline
 
 
 class RBLNKandinskyV22CombinedPipeline(RBLNDiffusionMixin, KandinskyV22CombinedPipeline):
+    """
+    RBLN-accelerated implementation of Kandinsky 2.2 combined pipeline for end-to-end text-to-image generation.
+
+    This pipeline compiles both prior and decoder Kandinsky 2.2 models to run efficiently on RBLN NPUs, enabling
+    high-performance inference for complete text-to-image generation with distinctive artistic style.
+    """
+
     original_class = KandinskyV22CombinedPipeline
     _rbln_config_class = RBLNKandinskyV22CombinedPipelineConfig
     _connected_classes = {"prior_pipe": RBLNKandinskyV22PriorPipeline, "decoder_pipe": RBLNKandinskyV22Pipeline}
@@ -46,15 +53,15 @@ class RBLNKandinskyV22CombinedPipeline(RBLNDiffusionMixin, KandinskyV22CombinedP
 
     def __init__(
         self,
-        unet: "UNet2DConditionModel",
-        scheduler: "DDPMScheduler",
-        movq: "VQModel",
-        prior_prior: "PriorTransformer",
-        prior_image_encoder: "CLIPVisionModelWithProjection",
-        prior_text_encoder: "CLIPTextModelWithProjection",
-        prior_tokenizer: "CLIPTokenizer",
-        prior_scheduler: "UnCLIPScheduler",
-        prior_image_processor: "CLIPImageProcessor",
+        unet: UNet2DConditionModel,
+        scheduler: DDPMScheduler,
+        movq: VQModel,
+        prior_prior: PriorTransformer,
+        prior_image_encoder: CLIPVisionModelWithProjection,
+        prior_text_encoder: CLIPTextModelWithProjection,
+        prior_tokenizer: CLIPTokenizer,
+        prior_scheduler: UnCLIPScheduler,
+        prior_image_processor: CLIPImageProcessor,
     ):
         RBLNDiffusionMixin.__init__(self)
         super(KandinskyV22CombinedPipeline, self).__init__()
@@ -90,6 +97,13 @@ class RBLNKandinskyV22CombinedPipeline(RBLNDiffusionMixin, KandinskyV22CombinedP
 
 
 class RBLNKandinskyV22Img2ImgCombinedPipeline(RBLNDiffusionMixin, KandinskyV22Img2ImgCombinedPipeline):
+    """
+    RBLN-accelerated implementation of Kandinsky 2.2 combined pipeline for end-to-end image-to-image generation.
+
+    This pipeline compiles both prior and decoder Kandinsky 2.2 models to run efficiently on RBLN NPUs, enabling
+    high-performance inference for complete image-to-image transformation with distinctive artistic style.
+    """
+
     original_class = KandinskyV22Img2ImgCombinedPipeline
     _connected_classes = {"prior_pipe": RBLNKandinskyV22PriorPipeline, "decoder_pipe": RBLNKandinskyV22Img2ImgPipeline}
     _submodules = ["prior_image_encoder", "prior_text_encoder", "prior_prior", "unet", "movq"]
@@ -97,15 +111,15 @@ class RBLNKandinskyV22Img2ImgCombinedPipeline(RBLNDiffusionMixin, KandinskyV22Im
 
     def __init__(
         self,
-        unet: "UNet2DConditionModel",
-        scheduler: "DDPMScheduler",
-        movq: "VQModel",
-        prior_prior: "PriorTransformer",
-        prior_image_encoder: "CLIPVisionModelWithProjection",
-        prior_text_encoder: "CLIPTextModelWithProjection",
-        prior_tokenizer: "CLIPTokenizer",
-        prior_scheduler: "UnCLIPScheduler",
-        prior_image_processor: "CLIPImageProcessor",
+        unet: UNet2DConditionModel,
+        scheduler: DDPMScheduler,
+        movq: VQModel,
+        prior_prior: PriorTransformer,
+        prior_image_encoder: CLIPVisionModelWithProjection,
+        prior_text_encoder: CLIPTextModelWithProjection,
+        prior_tokenizer: CLIPTokenizer,
+        prior_scheduler: UnCLIPScheduler,
+        prior_image_processor: CLIPImageProcessor,
     ):
         RBLNDiffusionMixin.__init__(self)
         super(KandinskyV22Img2ImgCombinedPipeline, self).__init__()
@@ -141,6 +155,13 @@ class RBLNKandinskyV22Img2ImgCombinedPipeline(RBLNDiffusionMixin, KandinskyV22Im
 
 
 class RBLNKandinskyV22InpaintCombinedPipeline(RBLNDiffusionMixin, KandinskyV22InpaintCombinedPipeline):
+    """
+    RBLN-accelerated implementation of Kandinsky 2.2 combined pipeline for end-to-end image inpainting.
+
+    This pipeline compiles both prior and decoder Kandinsky 2.2 models to run efficiently on RBLN NPUs, enabling
+    high-performance inference for complete image inpainting with distinctive artistic style and seamless integration.
+    """
+
     original_class = KandinskyV22InpaintCombinedPipeline
     _connected_classes = {"prior_pipe": RBLNKandinskyV22PriorPipeline, "decoder_pipe": RBLNKandinskyV22InpaintPipeline}
     _submodules = ["prior_image_encoder", "prior_text_encoder", "prior_prior", "unet", "movq"]
@@ -148,15 +169,15 @@ class RBLNKandinskyV22InpaintCombinedPipeline(RBLNDiffusionMixin, KandinskyV22In
 
     def __init__(
         self,
-        unet: "UNet2DConditionModel",
-        scheduler: "DDPMScheduler",
-        movq: "VQModel",
-        prior_prior: "PriorTransformer",
-        prior_image_encoder: "CLIPVisionModelWithProjection",
-        prior_text_encoder: "CLIPTextModelWithProjection",
-        prior_tokenizer: "CLIPTokenizer",
-        prior_scheduler: "UnCLIPScheduler",
-        prior_image_processor: "CLIPImageProcessor",
+        unet: UNet2DConditionModel,
+        scheduler: DDPMScheduler,
+        movq: VQModel,
+        prior_prior: PriorTransformer,
+        prior_image_encoder: CLIPVisionModelWithProjection,
+        prior_text_encoder: CLIPTextModelWithProjection,
+        prior_tokenizer: CLIPTokenizer,
+        prior_scheduler: UnCLIPScheduler,
+        prior_image_processor: CLIPImageProcessor,
     ):
         RBLNDiffusionMixin.__init__(self)
         super(KandinskyV22InpaintCombinedPipeline, self).__init__()

optimum/rbln/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_img2img.py

@@ -19,6 +19,13 @@ from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNKandinskyV22Img2ImgPipeline(RBLNDiffusionMixin, KandinskyV22Img2ImgPipeline):
+    """
+    RBLN-accelerated implementation of Kandinsky 2.2 pipeline for image-to-image generation.
+
+    This pipeline compiles Kandinsky 2.2 models to run efficiently on RBLN NPUs, enabling high-performance
+    inference for transforming input images with distinctive artistic style and enhanced visual fidelity.
+    """
+
     original_class = KandinskyV22Img2ImgPipeline
     _rbln_config_class = RBLNKandinskyV22Img2ImgPipelineConfig
     _submodules = ["unet", "movq"]

optimum/rbln/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_inpaint.py

@@ -19,6 +19,13 @@ from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNKandinskyV22InpaintPipeline(RBLNDiffusionMixin, KandinskyV22InpaintPipeline):
+    """
+    RBLN-accelerated implementation of Kandinsky 2.2 pipeline for image inpainting.
+
+    This pipeline compiles Kandinsky 2.2 models to run efficiently on RBLN NPUs, enabling high-performance
+    inference for filling masked regions with distinctive artistic style and seamless content integration.
+    """
+
     original_class = KandinskyV22InpaintPipeline
     _rbln_config_class = RBLNKandinskyV22InpaintPipelineConfig
     _submodules = ["unet", "movq"]

optimum/rbln/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior.py

@@ -19,6 +19,13 @@ from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNKandinskyV22PriorPipeline(RBLNDiffusionMixin, KandinskyV22PriorPipeline):
+    """
+    RBLN-accelerated implementation of Kandinsky 2.2 prior pipeline for text and image embedding generation.
+
+    This pipeline compiles Kandinsky 2.2 prior models to run efficiently on RBLN NPUs, enabling high-performance
+    inference for generating image embeddings from text prompts and image inputs for downstream generation tasks.
+    """
+
     original_class = KandinskyV22PriorPipeline
     _rbln_config_class = RBLNKandinskyV22PriorPipelineConfig
     _submodules = ["text_encoder", "image_encoder", "prior"]

optimum/rbln/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py

@@ -20,6 +20,13 @@ from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNStableDiffusionPipeline(RBLNDiffusionMixin, StableDiffusionPipeline):
+    """
+    RBLN-accelerated implementation of Stable Diffusion pipeline for text-to-image generation.
+
+    This pipeline compiles Stable Diffusion models to run efficiently on RBLN NPUs, enabling high-performance
+    inference for generating images from text prompts with optimized memory usage and throughput.
+    """
+
     original_class = StableDiffusionPipeline
     _rbln_config_class = RBLNStableDiffusionPipelineConfig
     _submodules = ["vae", "text_encoder", "unet"]