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

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

@@ -12,12 +12,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional, Tuple
+from typing import Any, Dict, Optional, Tuple
 
 from ....configuration_utils import RBLNModelConfig
 
 
 class RBLNVQModelConfig(RBLNModelConfig):
+    """
+    Configuration class for RBLN VQModel models, used in Kandinsky.
+
+    This class inherits from RBLNModelConfig and provides specific configuration options
+    for VQModel, which acts similarly to a VAE but uses vector quantization.
+    """
+
     def __init__(
         self,
         batch_size: Optional[int] = None,
@@ -26,7 +33,7 @@ class RBLNVQModelConfig(RBLNModelConfig):
         vqmodel_scale_factor: Optional[float] = None,  # TODO: rename to scaling_factor
         in_channels: Optional[int] = None,
         latent_channels: Optional[int] = None,
-        **kwargs,
+        **kwargs: Dict[str, Any],
     ):
         """
         Args:

optimum/rbln/diffusers/configurations/pipelines/__init__.py

@@ -4,6 +4,7 @@ from .configuration_controlnet import (
     RBLNStableDiffusionXLControlNetImg2ImgPipelineConfig,
     RBLNStableDiffusionXLControlNetPipelineConfig,
 )
+from .configuration_cosmos import RBLNCosmosTextToWorldPipelineConfig, RBLNCosmosVideoToWorldPipelineConfig
 from .configuration_kandinsky2_2 import (
     RBLNKandinskyV22CombinedPipelineConfig,
     RBLNKandinskyV22Img2ImgCombinedPipelineConfig,

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

@@ -12,14 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional, Tuple
+from typing import Any, Dict, Optional, Tuple
 
 from ....configuration_utils import RBLNModelConfig
 from ....transformers import RBLNCLIPTextModelConfig, RBLNCLIPTextModelWithProjectionConfig
 from ..models import RBLNAutoencoderKLConfig, RBLNControlNetModelConfig, RBLNUNet2DConditionModelConfig
 
 
-class _RBLNStableDiffusionControlNetPipelineBaseConfig(RBLNModelConfig):
+class RBLNStableDiffusionControlNetPipelineBaseConfig(RBLNModelConfig):
     submodules = ["text_encoder", "unet", "vae", "controlnet"]
     _vae_uses_encoder = False
 
@@ -38,7 +38,7 @@ class _RBLNStableDiffusionControlNetPipelineBaseConfig(RBLNModelConfig):
         sample_size: Optional[Tuple[int, int]] = None,
         image_size: Optional[Tuple[int, int]] = None,
         guidance_scale: Optional[float] = None,
-        **kwargs,
+        **kwargs: Dict[str, Any],
     ):
         """
         Args:
@@ -138,15 +138,27 @@ class _RBLNStableDiffusionControlNetPipelineBaseConfig(RBLNModelConfig):
         return self.vae.sample_size
 
 
-class RBLNStableDiffusionControlNetPipelineConfig(_RBLNStableDiffusionControlNetPipelineBaseConfig):
+class RBLNStableDiffusionControlNetPipelineConfig(RBLNStableDiffusionControlNetPipelineBaseConfig):
+    """
+    Configuration for Stable Diffusion ControlNet pipeline.
+    """
+
     _vae_uses_encoder = False
 
 
-class RBLNStableDiffusionControlNetImg2ImgPipelineConfig(_RBLNStableDiffusionControlNetPipelineBaseConfig):
+class RBLNStableDiffusionControlNetImg2ImgPipelineConfig(RBLNStableDiffusionControlNetPipelineBaseConfig):
+    """
+    Configuration for Stable Diffusion ControlNet image-to-image pipeline.
+    """
+
     _vae_uses_encoder = True
 
 
-class _RBLNStableDiffusionXLControlNetPipelineBaseConfig(RBLNModelConfig):
+class RBLNStableDiffusionXLControlNetPipelineBaseConfig(RBLNModelConfig):
+    """
+    Base configuration for Stable Diffusion XL ControlNet pipelines.
+    """
+
     submodules = ["text_encoder", "text_encoder_2", "unet", "vae", "controlnet"]
     _vae_uses_encoder = False
 
@@ -166,7 +178,7 @@ class _RBLNStableDiffusionXLControlNetPipelineBaseConfig(RBLNModelConfig):
         sample_size: Optional[Tuple[int, int]] = None,
         image_size: Optional[Tuple[int, int]] = None,
         guidance_scale: Optional[float] = None,
-        **kwargs,
+        **kwargs: Dict[str, Any],
     ):
         """
         Args:
@@ -272,9 +284,17 @@ class _RBLNStableDiffusionXLControlNetPipelineBaseConfig(RBLNModelConfig):
         return self.vae.sample_size
 
 
-class RBLNStableDiffusionXLControlNetPipelineConfig(_RBLNStableDiffusionXLControlNetPipelineBaseConfig):
+class RBLNStableDiffusionXLControlNetPipelineConfig(RBLNStableDiffusionXLControlNetPipelineBaseConfig):
+    """
+    Configuration for Stable Diffusion XL ControlNet pipeline.
+    """
+
     _vae_uses_encoder = False
 
 
-class RBLNStableDiffusionXLControlNetImg2ImgPipelineConfig(_RBLNStableDiffusionXLControlNetPipelineBaseConfig):
+class RBLNStableDiffusionXLControlNetImg2ImgPipelineConfig(RBLNStableDiffusionXLControlNetPipelineBaseConfig):
+    """
+    Configuration for Stable Diffusion XL ControlNet image-to-image pipeline.
+    """
+
     _vae_uses_encoder = True

optimum/rbln/diffusers/configurations/pipelines/configuration_cosmos.py

@@ -0,0 +1,114 @@
+# Copyright 2025 Rebellions Inc. All rights reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any, Dict, Optional
+
+from ....configuration_utils import RBLNModelConfig
+from ....transformers import RBLNT5EncoderModelConfig
+from ....utils.logging import get_logger
+from ...pipelines.cosmos.cosmos_guardrail import RBLNCosmosSafetyCheckerConfig
+from ..models import RBLNAutoencoderKLCosmosConfig, RBLNCosmosTransformer3DModelConfig
+
+
+logger = get_logger(__name__)
+
+
+class RBLNCosmosPipelineBaseConfig(RBLNModelConfig):
+    submodules = ["text_encoder", "transformer", "vae", "safety_checker"]
+    _vae_uses_encoder = False
+
+    def __init__(
+        self,
+        text_encoder: Optional[RBLNT5EncoderModelConfig] = None,
+        transformer: Optional[RBLNCosmosTransformer3DModelConfig] = None,
+        vae: Optional[RBLNAutoencoderKLCosmosConfig] = None,
+        safety_checker: Optional[RBLNCosmosSafetyCheckerConfig] = None,
+        *,
+        batch_size: Optional[int] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_frames: Optional[int] = None,
+        fps: Optional[int] = None,
+        max_seq_len: Optional[int] = None,
+        **kwargs: Dict[str, Any],
+    ):
+        """
+        Args:
+            text_encoder (Optional[RBLNT5EncoderModelConfig]): Configuration for the text encoder component.
+                Initialized as RBLNT5EncoderModelConfig if not provided.
+            transformer (Optional[RBLNCosmosTransformer3DModelConfig]): Configuration for the Transformer model component.
+                Initialized as RBLNCosmosTransformer3DModelConfig if not provided.
+            vae (Optional[RBLNAutoencoderKLCosmosConfig]): Configuration for the VAE model component.
+                Initialized as RBLNAutoencoderKLCosmosConfig if not provided.
+            safety_checker (Optional[RBLNCosmosSafetyCheckerConfig]): Configuration for the safety checker component.
+                Initialized as RBLNCosmosSafetyCheckerConfig if not provided.
+            batch_size (Optional[int]): Batch size for inference, applied to all submodules.
+            height (Optional[int]): Height of the generated videos.
+            width (Optional[int]): Width of the generated videos.
+            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.
+        """
+        super().__init__(**kwargs)
+
+        self.text_encoder = self.init_submodule_config(
+            RBLNT5EncoderModelConfig, text_encoder, batch_size=batch_size, max_seq_len=max_seq_len
+        )
+        self.transformer = self.init_submodule_config(
+            RBLNCosmosTransformer3DModelConfig,
+            transformer,
+            batch_size=batch_size,
+            max_seq_len=max_seq_len,
+            height=height,
+            width=width,
+            num_frames=num_frames,
+            fps=fps,
+        )
+        self.vae = self.init_submodule_config(
+            RBLNAutoencoderKLCosmosConfig,
+            vae,
+            batch_size=batch_size,
+            uses_encoder=self.__class__._vae_uses_encoder,
+            height=height,
+            width=width,
+            num_frames=num_frames,
+        )
+        self.safety_checker = self.init_submodule_config(
+            RBLNCosmosSafetyCheckerConfig,
+            safety_checker,
+            batch_size=batch_size,
+            height=height,
+            width=width,
+        )
+
+    @property
+    def batch_size(self):
+        return self.vae.batch_size
+
+    @property
+    def max_seq_len(self):
+        return self.text_encoder.max_seq_len
+
+
+class RBLNCosmosTextToWorldPipelineConfig(RBLNCosmosPipelineBaseConfig):
+    """Config for Cosmos Text2World Pipeline"""
+
+    _vae_uses_encoder = False
+
+
+class RBLNCosmosVideoToWorldPipelineConfig(RBLNCosmosPipelineBaseConfig):
+    """Config for Cosmos Video2World Pipeline"""
+
+    _vae_uses_encoder = True

optimum/rbln/diffusers/configurations/pipelines/configuration_kandinsky2_2.py

@@ -12,7 +12,7 @@
 # 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, RBLNCLIPVisionModelWithProjectionConfig
@@ -20,7 +20,7 @@ from ..models import RBLNUNet2DConditionModelConfig, RBLNVQModelConfig
 from ..models.configuration_prior_transformer import RBLNPriorTransformerConfig
 
 
-class _RBLNKandinskyV22PipelineBaseConfig(RBLNModelConfig):
+class RBLNKandinskyV22PipelineBaseConfig(RBLNModelConfig):
     submodules = ["unet", "movq"]
     _movq_uses_encoder = False
 
@@ -37,7 +37,7 @@ class _RBLNKandinskyV22PipelineBaseConfig(RBLNModelConfig):
         img_width: Optional[int] = None,
         height: Optional[int] = None,
         width: Optional[int] = None,
-        **kwargs,
+        **kwargs: Dict[str, Any],
     ):
         """
         Args:
@@ -117,19 +117,27 @@ class _RBLNKandinskyV22PipelineBaseConfig(RBLNModelConfig):
         return self.movq.sample_size
 
 
-class RBLNKandinskyV22PipelineConfig(_RBLNKandinskyV22PipelineBaseConfig):
+class RBLNKandinskyV22PipelineConfig(RBLNKandinskyV22PipelineBaseConfig):
+    """Configuration class for the Kandinsky V2.2 text-to-image decoder pipeline."""
+
     _movq_uses_encoder = False
 
 
-class RBLNKandinskyV22Img2ImgPipelineConfig(_RBLNKandinskyV22PipelineBaseConfig):
+class RBLNKandinskyV22Img2ImgPipelineConfig(RBLNKandinskyV22PipelineBaseConfig):
+    """Configuration class for the Kandinsky V2.2 image-to-image decoder pipeline."""
+
     _movq_uses_encoder = True
 
 
-class RBLNKandinskyV22InpaintPipelineConfig(_RBLNKandinskyV22PipelineBaseConfig):
+class RBLNKandinskyV22InpaintPipelineConfig(RBLNKandinskyV22PipelineBaseConfig):
+    """Configuration class for the Kandinsky V2.2 inpainting decoder pipeline."""
+
     _movq_uses_encoder = True
 
 
 class RBLNKandinskyV22PriorPipelineConfig(RBLNModelConfig):
+    """Configuration class for the Kandinsky V2.2 Prior pipeline."""
+
     submodules = ["text_encoder", "image_encoder", "prior"]
 
     def __init__(
@@ -140,7 +148,7 @@ class RBLNKandinskyV22PriorPipelineConfig(RBLNModelConfig):
         *,
         batch_size: Optional[int] = None,
         guidance_scale: Optional[float] = None,
-        **kwargs,
+        **kwargs: Dict[str, Any],
     ):
         """
         Initialize a configuration for Kandinsky 2.2 prior pipeline optimized for RBLN NPU.
@@ -194,7 +202,9 @@ class RBLNKandinskyV22PriorPipelineConfig(RBLNModelConfig):
         return self.image_encoder.image_size
 
 
-class _RBLNKandinskyV22CombinedPipelineBaseConfig(RBLNModelConfig):
+class RBLNKandinskyV22CombinedPipelineBaseConfig(RBLNModelConfig):
+    """Base configuration class for Kandinsky V2.2 combined pipelines."""
+
     submodules = ["prior_pipe", "decoder_pipe"]
     _decoder_pipe_cls = RBLNKandinskyV22PipelineConfig
 
@@ -216,7 +226,7 @@ class _RBLNKandinskyV22CombinedPipelineBaseConfig(RBLNModelConfig):
         prior_text_encoder: Optional[RBLNCLIPTextModelWithProjectionConfig] = None,
         unet: Optional[RBLNUNet2DConditionModelConfig] = None,
         movq: Optional[RBLNVQModelConfig] = None,
-        **kwargs,
+        **kwargs: Dict[str, Any],
     ):
         """
         Initialize a configuration for combined Kandinsky 2.2 pipelines optimized for RBLN NPU.
@@ -325,13 +335,19 @@ class _RBLNKandinskyV22CombinedPipelineBaseConfig(RBLNModelConfig):
         return self.decoder_pipe.movq
 
 
-class RBLNKandinskyV22CombinedPipelineConfig(_RBLNKandinskyV22CombinedPipelineBaseConfig):
+class RBLNKandinskyV22CombinedPipelineConfig(RBLNKandinskyV22CombinedPipelineBaseConfig):
+    """Configuration class for the Kandinsky V2.2 combined text-to-image pipeline."""
+
     _decoder_pipe_cls = RBLNKandinskyV22PipelineConfig
 
 
-class RBLNKandinskyV22InpaintCombinedPipelineConfig(_RBLNKandinskyV22CombinedPipelineBaseConfig):
+class RBLNKandinskyV22InpaintCombinedPipelineConfig(RBLNKandinskyV22CombinedPipelineBaseConfig):
+    """Configuration class for the Kandinsky V2.2 combined inpainting pipeline."""
+
     _decoder_pipe_cls = RBLNKandinskyV22InpaintPipelineConfig
 
 
-class RBLNKandinskyV22Img2ImgCombinedPipelineConfig(_RBLNKandinskyV22CombinedPipelineBaseConfig):
+class RBLNKandinskyV22Img2ImgCombinedPipelineConfig(RBLNKandinskyV22CombinedPipelineBaseConfig):
+    """Configuration class for the Kandinsky V2.2 combined image-to-image pipeline."""
+
     _decoder_pipe_cls = RBLNKandinskyV22Img2ImgPipelineConfig

optimum/rbln/diffusers/configurations/pipelines/configuration_stable_diffusion.py

@@ -12,14 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional, Tuple
+from typing import Any, Dict, Optional, Tuple
 
 from ....configuration_utils import RBLNModelConfig
 from ....transformers import RBLNCLIPTextModelConfig
 from ..models import RBLNAutoencoderKLConfig, RBLNUNet2DConditionModelConfig
 
 
-class _RBLNStableDiffusionPipelineBaseConfig(RBLNModelConfig):
+class RBLNStableDiffusionPipelineBaseConfig(RBLNModelConfig):
     submodules = ["text_encoder", "unet", "vae"]
     _vae_uses_encoder = False
 
@@ -37,7 +37,7 @@ class _RBLNStableDiffusionPipelineBaseConfig(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:
@@ -128,13 +128,25 @@ class _RBLNStableDiffusionPipelineBaseConfig(RBLNModelConfig):
         return self.vae.sample_size
 
 
-class RBLNStableDiffusionPipelineConfig(_RBLNStableDiffusionPipelineBaseConfig):
+class RBLNStableDiffusionPipelineConfig(RBLNStableDiffusionPipelineBaseConfig):
+    """
+    Configuration for Stable Diffusion pipeline.
+    """
+
     _vae_uses_encoder = False
 
 
-class RBLNStableDiffusionImg2ImgPipelineConfig(_RBLNStableDiffusionPipelineBaseConfig):
+class RBLNStableDiffusionImg2ImgPipelineConfig(RBLNStableDiffusionPipelineBaseConfig):
+    """
+    Configuration for Stable Diffusion image-to-image pipeline.
+    """
+
     _vae_uses_encoder = True
 
 
-class RBLNStableDiffusionInpaintPipelineConfig(_RBLNStableDiffusionPipelineBaseConfig):
+class RBLNStableDiffusionInpaintPipelineConfig(RBLNStableDiffusionPipelineBaseConfig):
+    """
+    Configuration for Stable Diffusion inpainting pipeline.
+    """
+
     _vae_uses_encoder = True

optimum/rbln/diffusers/configurations/pipelines/configuration_stable_diffusion_3.py

@@ -12,14 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional, Tuple
+from typing import Any, Dict, Optional, Tuple
 
 from ....configuration_utils import RBLNModelConfig
 from ....transformers import RBLNCLIPTextModelWithProjectionConfig, RBLNT5EncoderModelConfig
 from ..models import RBLNAutoencoderKLConfig, RBLNSD3Transformer2DModelConfig
 
 
-class _RBLNStableDiffusion3PipelineBaseConfig(RBLNModelConfig):
+class RBLNStableDiffusion3PipelineBaseConfig(RBLNModelConfig):
     submodules = ["transformer", "text_encoder", "text_encoder_2", "text_encoder_3", "vae"]
     _vae_uses_encoder = False
 
@@ -40,7 +40,7 @@ class _RBLNStableDiffusion3PipelineBaseConfig(RBLNModelConfig):
         height: Optional[int] = None,
         width: Optional[int] = None,
         guidance_scale: Optional[float] = None,
-        **kwargs,
+        **kwargs: Dict[str, Any],
     ):
         """
         Args:
@@ -111,6 +111,7 @@ class _RBLNStableDiffusion3PipelineBaseConfig(RBLNModelConfig):
             text_encoder_3,
             batch_size=batch_size,
             max_seq_len=max_seq_len,
+            model_input_names=["input_ids"],
         )
         self.transformer = self.init_submodule_config(
             RBLNSD3Transformer2DModelConfig,
@@ -153,13 +154,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,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional, Tuple
+from typing import Any, Dict, Optional, Tuple
 
 from ....configuration_utils import RBLNModelConfig
 from ....transformers import RBLNCLIPTextModelConfig, RBLNCLIPTextModelWithProjectionConfig
 from ..models import RBLNAutoencoderKLConfig, RBLNUNet2DConditionModelConfig
 
 
-class _RBLNStableDiffusionXLPipelineBaseConfig(RBLNModelConfig):
+class RBLNStableDiffusionXLPipelineBaseConfig(RBLNModelConfig):
     submodules = ["text_encoder", "text_encoder_2", "unet", "vae"]
     _vae_uses_encoder = False
 
@@ -38,7 +38,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 +134,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