optimum-rbln 0.8.1rc0__py3-none-any.whl → 0.8.2a0__py3-none-any.whl

optimum/rbln/__version__.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.8.1rc0'
-__version_tuple__ = version_tuple = (0, 8, 1, 'rc0')
+__version__ = version = '0.8.2a0'
+__version_tuple__ = version_tuple = (0, 8, 2, 'a0')

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

@@ -1,8 +1,8 @@
 from .configuration_autoencoder_kl import RBLNAutoencoderKLConfig
 from .configuration_autoencoder_kl_cosmos import RBLNAutoencoderKLCosmosConfig
 from .configuration_controlnet import RBLNControlNetModelConfig
-from .configuration_cosmos_transformer import RBLNCosmosTransformer3DModelConfig
 from .configuration_prior_transformer import RBLNPriorTransformerConfig
+from .configuration_transformer_cosmos import RBLNCosmosTransformer3DModelConfig
 from .configuration_transformer_sd3 import RBLNSD3Transformer2DModelConfig
 from .configuration_unet_2d_condition import RBLNUNet2DConditionModelConfig
 from .configuration_vq_model import RBLNVQModelConfig

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 Optional
+from typing import Any, Dict, Optional
 
 from ....configuration_utils import RBLNModelConfig
 from ....utils.logging import get_logger
@@ -22,6 +22,8 @@ logger = get_logger(__name__)
 
 
 class RBLNAutoencoderKLCosmosConfig(RBLNModelConfig):
+    """Configuration class for RBLN Cosmos Variational Autoencoder (VAE) models."""
+
     def __init__(
         self,
         batch_size: Optional[int] = None,
@@ -33,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,
+        **kwargs: Dict[str, Any],
     ):
         """
         Args:
@@ -48,7 +50,7 @@ class RBLNAutoencoderKLCosmosConfig(RBLNModelConfig):
                 Determines how much shorter the latent representations are compared to the original videos.
             vae_scale_factor_spatial (Optional[int]): The scaling factor between pixel space and latent space.
                 Determines how much smaller the latent representations are compared to the original videos.
-            use_slicing (Optional[Bool]): Enable sliced VAE encoding and decoding.
+            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.
 

optimum/rbln/diffusers/configurations/models/{configuration_cosmos_transformer.py → configuration_transformer_cosmos.py}

@@ -12,12 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional
+from typing import Any, Dict, Optional
 
 from ....configuration_utils import RBLNModelConfig
 
 
 class RBLNCosmosTransformer3DModelConfig(RBLNModelConfig):
+    """Configuration class for RBLN Cosmos Transformer models."""
+
     def __init__(
         self,
         batch_size: Optional[int] = None,
@@ -31,7 +33,7 @@ class RBLNCosmosTransformer3DModelConfig(RBLNModelConfig):
         num_latent_frames: Optional[int] = None,
         latent_height: Optional[int] = None,
         latent_width: Optional[int] = None,
-        **kwargs,
+        **kwargs: Dict[str, Any],
     ):
         """
         Args:

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

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional
+from typing import Any, Dict, Optional
 
 from ....configuration_utils import RBLNModelConfig
 from ....transformers import RBLNT5EncoderModelConfig
@@ -24,7 +24,7 @@ from ..models import RBLNAutoencoderKLCosmosConfig, RBLNCosmosTransformer3DModel
 logger = get_logger(__name__)
 
 
-class _RBLNCosmosPipelineBaseConfig(RBLNModelConfig):
+class RBLNCosmosPipelineBaseConfig(RBLNModelConfig):
     submodules = ["text_encoder", "transformer", "vae", "safety_checker"]
     _vae_uses_encoder = False
 
@@ -41,13 +41,13 @@ class _RBLNCosmosPipelineBaseConfig(RBLNModelConfig):
         num_frames: Optional[int] = None,
         fps: Optional[int] = None,
         max_seq_len: Optional[int] = None,
-        **kwargs,
+        **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 UNet model component.
+            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.
@@ -102,9 +102,13 @@ class _RBLNCosmosPipelineBaseConfig(RBLNModelConfig):
         return self.text_encoder.max_seq_len
 
 
-class RBLNCosmosTextToWorldPipelineConfig(_RBLNCosmosPipelineBaseConfig):
+class RBLNCosmosTextToWorldPipelineConfig(RBLNCosmosPipelineBaseConfig):
+    """Config for Cosmos Text2World Pipeline"""
+
     _vae_uses_encoder = False
 
 
-class RBLNCosmosVideoToWorldPipelineConfig(_RBLNCosmosPipelineBaseConfig):
+class RBLNCosmosVideoToWorldPipelineConfig(RBLNCosmosPipelineBaseConfig):
+    """Config for Cosmos Video2World Pipeline"""
+
     _vae_uses_encoder = True

optimum/rbln/diffusers/pipelines/cosmos/cosmos_guardrail.py

@@ -130,7 +130,8 @@ class RBLNSigLIPEncoder(SigLIPEncoder):
                 self.checkpoint_dir,
                 rbln_device=rbln_config.siglip_encoder.device,
                 rbln_create_runtimes=rbln_config.siglip_encoder.create_runtimes,
-                rbln_activate_profiler=rbln_config.aegis.activate_profiler,
+                rbln_activate_profiler=rbln_config.siglip_encoder.activate_profiler,
+                rbln_optimize_host_memory=rbln_config.siglip_encoder.optimize_host_memory,
             )
         else:
             super().__init__(model_name, checkpoint_id)
@@ -143,6 +144,7 @@ class RBLNSigLIPEncoder(SigLIPEncoder):
                 rbln_npu=rbln_config.siglip_encoder.npu,
                 rbln_create_runtimes=rbln_config.siglip_encoder.create_runtimes,
                 rbln_activate_profiler=rbln_config.siglip_encoder.activate_profiler,
+                rbln_optimize_host_memory=rbln_config.siglip_encoder.optimize_host_memory,
             )
         self.rbln_config = rbln_config
 
@@ -339,6 +341,7 @@ class RBLNAegis(Aegis):
                 rbln_device=rbln_config.aegis.device,
                 rbln_create_runtimes=rbln_config.aegis.create_runtimes,
                 rbln_activate_profiler=rbln_config.aegis.activate_profiler,
+                rbln_optimize_host_memory=rbln_config.aegis.optimize_host_memory,
             )
 
         else:
@@ -353,6 +356,7 @@ class RBLNAegis(Aegis):
                 rbln_create_runtimes=rbln_config.aegis.create_runtimes,
                 rbln_npu=rbln_config.aegis.npu,
                 rbln_activate_profiler=rbln_config.aegis.activate_profiler,
+                rbln_optimize_host_memory=rbln_config.aegis.optimize_host_memory,
             )
 
         self.rbln_config = rbln_config

optimum/rbln/diffusers/pipelines/cosmos/pipeline_cosmos_text2world.py

@@ -35,7 +35,7 @@ class RBLNCosmosTextToWorldPipeline(RBLNDiffusionMixin, CosmosTextToWorldPipelin
     RBLN-accelerated implementation of Cosmos Text to World pipeline for text-to-video generation.
 
     This pipeline compiles Cosmos Text to World models to run efficiently on RBLN NPUs, enabling high-performance
-    inference for generating images with distinctive artistic style and enhanced visual quality.
+    inference for generating videos with distinctive artistic style and enhanced visual quality.
     """
 
     original_class = CosmosTextToWorldPipeline

optimum/rbln/diffusers/pipelines/cosmos/pipeline_cosmos_video2world.py

@@ -35,7 +35,7 @@ class RBLNCosmosVideoToWorldPipeline(RBLNDiffusionMixin, CosmosVideoToWorldPipel
     RBLN-accelerated implementation of Cosmos Video to World pipeline for video-to-video generation.
 
     This pipeline compiles Cosmos Video to World models to run efficiently on RBLN NPUs, enabling high-performance
-    inference for generating images with distinctive artistic style and enhanced visual quality.
+    inference for generating videos with distinctive artistic style and enhanced visual quality.
     """
 
     original_class = CosmosVideoToWorldPipeline

optimum/rbln/transformers/models/blip_2/configuration_blip_2.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional
+from typing import Any, Dict, Optional
 
 from ....configuration_utils import RBLNModelConfig
 
@@ -62,7 +62,7 @@ class RBLNBlip2ForConditionalGenerationConfig(RBLNModelConfig):
         vision_model: Optional[RBLNModelConfig] = None,
         qformer: Optional[RBLNModelConfig] = None,
         language_model: Optional[RBLNModelConfig] = None,
-        **kwargs,
+        **kwargs: Dict[str, Any],
     ):
         """
         Args:

optimum/rbln/transformers/models/colpali/configuration_colpali.py

@@ -11,7 +11,7 @@
 # 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 List, Optional, Union
+from typing import Any, Dict, List, Optional, Union
 
 from ....configuration_utils import RBLNModelConfig
 
@@ -50,7 +50,7 @@ class RBLNColPaliForRetrievalConfig(RBLNModelConfig):
         max_seq_lens: Union[int, List[int]] = None,
         output_hidden_states: Optional[bool] = None,
         vision_tower: Optional[RBLNModelConfig] = None,
-        **kwargs,
+        **kwargs: Dict[str, Any],
     ):
         """
         Args: