optimum-rbln 0.8.0.post2__py3-none-any.whl → 0.8.1__py3-none-any.whl

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"]

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

@@ -19,6 +19,13 @@ from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNStableDiffusionImg2ImgPipeline(RBLNDiffusionMixin, StableDiffusionImg2ImgPipeline):
+    """
+    RBLN-accelerated implementation of Stable Diffusion pipeline for image-to-image generation.
+
+    This pipeline compiles Stable Diffusion models to run efficiently on RBLN NPUs, enabling high-performance
+    inference for transforming input images based on text prompts with controlled strength and guidance.
+    """
+
     original_class = StableDiffusionImg2ImgPipeline
     _rbln_config_class = RBLNStableDiffusionImg2ImgPipelineConfig
     _submodules = ["text_encoder", "unet", "vae"]

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

@@ -19,6 +19,13 @@ from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNStableDiffusionInpaintPipeline(RBLNDiffusionMixin, StableDiffusionInpaintPipeline):
+    """
+    RBLN-accelerated implementation of Stable Diffusion pipeline for image inpainting.
+
+    This pipeline compiles Stable Diffusion models to run efficiently on RBLN NPUs, enabling high-performance
+    inference for filling masked regions of images based on text prompts with seamless integration.
+    """
+
     original_class = StableDiffusionInpaintPipeline
     _rbln_config_class = RBLNStableDiffusionInpaintPipelineConfig
     _submodules = ["text_encoder", "unet", "vae"]

optimum/rbln/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py

@@ -19,6 +19,13 @@ from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNStableDiffusion3Pipeline(RBLNDiffusionMixin, StableDiffusion3Pipeline):
+    """
+    RBLN-accelerated implementation of Stable Diffusion 3 pipeline for advanced text-to-image generation.
+
+    This pipeline compiles Stable Diffusion 3 models to run efficiently on RBLN NPUs, enabling high-performance
+    inference with improved text understanding, enhanced image quality, and superior prompt adherence.
+    """
+
     original_class = StableDiffusion3Pipeline
     _rbln_config_class = RBLNStableDiffusion3PipelineConfig
     _submodules = ["transformer", "text_encoder_3", "text_encoder", "text_encoder_2", "vae"]

optimum/rbln/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_img2img.py

@@ -19,6 +19,13 @@ from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNStableDiffusion3Img2ImgPipeline(RBLNDiffusionMixin, StableDiffusion3Img2ImgPipeline):
+    """
+    RBLN-accelerated implementation of Stable Diffusion 3 pipeline for advanced image-to-image generation.
+
+    This pipeline compiles Stable Diffusion 3 models to run efficiently on RBLN NPUs, enabling high-performance
+    inference for transforming input images with superior text understanding and enhanced visual quality.
+    """
+
     original_class = StableDiffusion3Img2ImgPipeline
     _rbln_config_class = RBLNStableDiffusion3Img2ImgPipelineConfig
     _submodules = ["transformer", "text_encoder_3", "text_encoder", "text_encoder_2", "vae"]

optimum/rbln/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_inpaint.py

@@ -19,6 +19,13 @@ from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNStableDiffusion3InpaintPipeline(RBLNDiffusionMixin, StableDiffusion3InpaintPipeline):
+    """
+    RBLN-accelerated implementation of Stable Diffusion 3 pipeline for advanced image inpainting.
+
+    This pipeline compiles Stable Diffusion 3 models to run efficiently on RBLN NPUs, enabling high-performance
+    inference for filling masked regions with superior text understanding and seamless content generation.
+    """
+
     original_class = StableDiffusion3InpaintPipeline
     _rbln_config_class = RBLNStableDiffusion3InpaintPipelineConfig
     _submodules = ["transformer", "text_encoder_3", "text_encoder", "text_encoder_2", "vae"]

optimum/rbln/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py

@@ -19,6 +19,13 @@ from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNStableDiffusionXLPipeline(RBLNDiffusionMixin, StableDiffusionXLPipeline):
+    """
+    RBLN-accelerated implementation of Stable Diffusion XL pipeline for high-resolution text-to-image generation.
+
+    This pipeline compiles Stable Diffusion XL models to run efficiently on RBLN NPUs, enabling high-performance
+    inference for generating high-quality images with enhanced detail and improved prompt adherence.
+    """
+
     original_class = StableDiffusionXLPipeline
     _rbln_config_class = RBLNStableDiffusionXLPipelineConfig
     _submodules = ["text_encoder", "text_encoder_2", "unet", "vae"]

optimum/rbln/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py

@@ -19,6 +19,13 @@ from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNStableDiffusionXLImg2ImgPipeline(RBLNDiffusionMixin, StableDiffusionXLImg2ImgPipeline):
+    """
+    RBLN-accelerated implementation of Stable Diffusion XL pipeline for high-resolution image-to-image generation.
+
+    This pipeline compiles Stable Diffusion XL models to run efficiently on RBLN NPUs, enabling high-performance
+    inference for transforming input images with enhanced quality and detail preservation.
+    """
+
     original_class = StableDiffusionXLImg2ImgPipeline
     _rbln_config_class = RBLNStableDiffusionXLImg2ImgPipelineConfig
     _submodules = ["text_encoder", "text_encoder_2", "unet", "vae"]

optimum/rbln/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py

@@ -19,6 +19,13 @@ from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNStableDiffusionXLInpaintPipeline(RBLNDiffusionMixin, StableDiffusionXLInpaintPipeline):
+    """
+    RBLN-accelerated implementation of Stable Diffusion XL pipeline for high-resolution image inpainting.
+
+    This pipeline compiles Stable Diffusion XL models to run efficiently on RBLN NPUs, enabling high-performance
+    inference for filling masked regions with enhanced quality and seamless blending capabilities.
+    """
+
     original_class = StableDiffusionXLInpaintPipeline
     _rbln_config_class = RBLNStableDiffusionXLInpaintPipelineConfig
     _submodules = ["text_encoder", "text_encoder_2", "unet", "vae"]

optimum/rbln/modeling.py

@@ -14,7 +14,7 @@
 
 from pathlib import Path
 from tempfile import TemporaryDirectory
-from typing import TYPE_CHECKING, Dict, List, Optional, Union, get_args, get_origin, get_type_hints
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union, get_args, get_origin, get_type_hints
 
 import rebel
 import torch
@@ -35,27 +35,12 @@ logger = get_logger(__name__)
 
 
 class RBLNModel(RBLNBaseModel):
-    """
-    A class that inherits from RBLNBaseModel for models consisting of a single `torch.nn.Module`.
-
-    This class supports all the functionality of RBLNBaseModel, including loading and saving models using
-    the `from_pretrained` and `save_pretrained` methods, compiling PyTorch models for execution on RBLN NPU
-    devices.
-
-    Example:
-        ```python
-        model = RBLNModel.from_pretrained("model_id", export=True, rbln_npu="npu_name")
-        outputs = model(**inputs)
-        ```
-    """
-
     _output_class = None
 
     @classmethod
     def update_kwargs(cls, kwargs):
-        """
-        Update user-given kwargs to get proper pytorch model.
-        """
+        # Update user-given kwargs to get proper pytorch model.
+
         return kwargs
 
     @classmethod
@@ -66,10 +51,9 @@ class RBLNModel(RBLNBaseModel):
         subfolder: str,
         rbln_config: RBLNModelConfig,
     ):
-        """
-        If you are unavoidably running on a CPU rather than an RBLN device,
-        store the torch tensor, weight, etc. in this function.
-        """
+        # If you are unavoidably running on a CPU rather than an RBLN device,
+        # store the torch tensor, weight, etc. in this function.
+        pass
 
     @classmethod
     def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNModelConfig) -> torch.nn.Module:
@@ -80,7 +64,12 @@ class RBLNModel(RBLNBaseModel):
     def get_compiled_model(cls, model: "PreTrainedModel", rbln_config: RBLNModelConfig):
         model = cls.wrap_model_if_needed(model, rbln_config)
         rbln_compile_config = rbln_config.compile_cfgs[0]
-        compiled_model = cls.compile(model, rbln_compile_config=rbln_compile_config)
+        compiled_model = cls.compile(
+            model,
+            rbln_compile_config=rbln_compile_config,
+            create_runtimes=rbln_config.create_runtimes,
+            device=rbln_config.device,
+        )
         return compiled_model
 
     @classmethod
@@ -88,11 +77,32 @@ class RBLNModel(RBLNBaseModel):
         cls,
         model: "PreTrainedModel",
         config: Optional[PretrainedConfig] = None,
-        rbln_config: Optional[RBLNModelConfig] = None,
+        rbln_config: Optional[Union[RBLNModelConfig, Dict]] = None,
         model_save_dir: Optional[Union[str, Path, TemporaryDirectory]] = None,
         subfolder: str = "",
-        **kwargs,
-    ):
+        **kwargs: Dict[str, Any],
+    ) -> "RBLNModel":
+        """
+        Converts and compiles a pre-trained HuggingFace library model into a RBLN model.
+        This method performs the actual model conversion and compilation process.
+
+        Args:
+            model: The PyTorch model to be compiled. The object must be an instance of the HuggingFace transformers PreTrainedModel class.
+            rbln_config: Configuration for RBLN model compilation and runtime. This can be provided as a dictionary or an instance of the model's configuration class (e.g., `RBLNLlamaForCausalLMConfig` for Llama models).
+                For detailed configuration options, see the specific model's configuration class documentation.
+
+            kwargs: Additional keyword arguments. Arguments with the prefix 'rbln_' are passed to rbln_config, while the remaining arguments are passed to the HuggingFace library.
+
+        The method performs the following steps:
+
+        1. Compiles the PyTorch model into an optimized RBLN graph
+        2. Configures the model for the specified NPU device
+        3. Creates the necessary runtime objects if requested
+        4. Saves the compiled model and configurations
+
+        Returns:
+            A RBLN model instance ready for inference on RBLN NPU devices.
+        """
         preprocessors = kwargs.pop("preprocessors", [])
         rbln_config, kwargs = cls.prepare_rbln_config(rbln_config=rbln_config, **kwargs)
 
@@ -232,7 +242,38 @@ class RBLNModel(RBLNBaseModel):
             for compiled_model in compiled_models
         ]
 
-    def forward(self, *args, return_dict: Optional[bool] = None, **kwargs):
+    def forward(self, *args: Any, return_dict: Optional[bool] = None, **kwargs: Dict[str, Any]) -> Any:
+        """
+        Defines the forward pass of the RBLN model, providing a drop-in replacement for HuggingFace PreTrainedModel.
+
+        This method executes the compiled RBLN model on RBLN NPU devices while maintaining full compatibility
+        with HuggingFace transformers and diffusers APIs. The RBLNModel can be used as a direct substitute
+        for any HuggingFace nn.Module/PreTrainedModel, enabling seamless integration into existing workflows.
+
+        Args:
+            *args: Variable length argument list containing model inputs. The format matches the original
+                HuggingFace model's forward method signature (e.g., input_ids, attention_mask for
+                transformers models, or sample, timestep for diffusers models).
+            return_dict:
+                Whether to return outputs as a dictionary-like object or as a tuple. When `None`:
+                - For transformers models: Uses `self.config.use_return_dict` (typically `True`)
+                - For diffusers models: Defaults to `True`
+            **kwargs: Arbitrary keyword arguments containing additional model inputs and parameters,
+                matching the original HuggingFace model's interface.
+
+        Returns:
+            Model outputs in the same format as the original HuggingFace model.
+
+            - If `return_dict=True`: Returns a dictionary-like object (e.g., BaseModelOutput,
+                CausalLMOutput) with named fields such as `logits`, `hidden_states`, etc.
+            - If `return_dict=False`: Returns a tuple containing the raw model outputs.
+
+        Note:
+            - This method maintains the exact same interface as the original HuggingFace model's forward method
+            - The compiled model runs on RBLN NPU hardware for accelerated inference
+            - All HuggingFace model features (generation, attention patterns, etc.) are preserved
+            - Can be used directly in HuggingFace pipelines, transformers.Trainer, and other workflows
+        """
         if self.hf_library_name == "transformers":
             return_dict = return_dict if return_dict is not None else self.config.use_return_dict
         else:
@@ -246,12 +287,7 @@ class RBLNModel(RBLNBaseModel):
 
     @classmethod
     def get_hf_output_class(cls):
-        """
-        Dynamically gets the output class from the corresponding HuggingFace model class.
-
-        Returns:
-            type: The appropriate output class from transformers or diffusers
-        """
+        # Dynamically gets the output class from the corresponding HuggingFace model class.
         if cls._output_class:
             return cls._output_class
 
@@ -278,10 +314,8 @@ class RBLNModel(RBLNBaseModel):
         return BaseModelOutput
 
     def _prepare_output(self, output, return_dict):
-        """
-        Prepare model output based on return_dict flag.
-        This method can be overridden by subclasses to provide task-specific output handling.
-        """
+        # Prepare model output based on return_dict flag.
+        # This method can be overridden by subclasses to provide task-specific output handling.
         tuple_output = (output,) if not isinstance(output, (tuple, list)) else tuple(output)
         if not return_dict:
             return tuple_output