optimum-rbln 0.8.2a4__py3-none-any.whl → 0.9.3rc0__py3-none-any.whl

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

optimum/rbln/diffusers/pipelines/__init__.py

@@ -18,6 +18,11 @@ from transformers.utils import _LazyModule
 
 
 _import_structure = {
+    "auto_pipeline": [
+        "RBLNAutoPipelineForImage2Image",
+        "RBLNAutoPipelineForInpainting",
+        "RBLNAutoPipelineForText2Image",
+    ],
     "controlnet": [
         "RBLNMultiControlNetModel",
         "RBLNStableDiffusionControlNetImg2ImgPipeline",
@@ -56,6 +61,11 @@ _import_structure = {
     ],
 }
 if TYPE_CHECKING:
+    from .auto_pipeline import (
+        RBLNAutoPipelineForImage2Image,
+        RBLNAutoPipelineForInpainting,
+        RBLNAutoPipelineForText2Image,
+    )
     from .controlnet import (
         RBLNMultiControlNetModel,
         RBLNStableDiffusionControlNetImg2ImgPipeline,
@@ -63,11 +73,7 @@ if TYPE_CHECKING:
         RBLNStableDiffusionXLControlNetImg2ImgPipeline,
         RBLNStableDiffusionXLControlNetPipeline,
     )
-    from .cosmos import (
-        RBLNCosmosSafetyChecker,
-        RBLNCosmosTextToWorldPipeline,
-        RBLNCosmosVideoToWorldPipeline,
-    )
+    from .cosmos import RBLNCosmosSafetyChecker, RBLNCosmosTextToWorldPipeline, RBLNCosmosVideoToWorldPipeline
     from .kandinsky2_2 import (
         RBLNKandinskyV22CombinedPipeline,
         RBLNKandinskyV22Img2ImgCombinedPipeline,

optimum/rbln/diffusers/pipelines/auto_pipeline.py

@@ -0,0 +1,307 @@
+# 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.
+
+
+import importlib
+from pathlib import Path
+from typing import Any, Dict, Type, Union
+
+from diffusers.models.controlnets import ControlNetUnionModel
+from diffusers.pipelines.auto_pipeline import (
+    AUTO_IMAGE2IMAGE_PIPELINES_MAPPING,
+    AUTO_INPAINT_PIPELINES_MAPPING,
+    AUTO_TEXT2IMAGE_PIPELINES_MAPPING,
+    AutoPipelineForImage2Image,
+    AutoPipelineForInpainting,
+    AutoPipelineForText2Image,
+    _get_task_class,
+)
+from huggingface_hub.utils import validate_hf_hub_args
+
+from optimum.rbln.configuration_utils import RBLNModelConfig
+from optimum.rbln.modeling_base import RBLNBaseModel
+from optimum.rbln.utils.model_utils import (
+    MODEL_MAPPING,
+    convert_hf_to_rbln_model_name,
+    convert_rbln_to_hf_model_name,
+    get_rbln_model_cls,
+)
+
+
+class RBLNAutoPipelineBase:
+    _model_mapping = None
+    _model_mapping_names = None
+
+    @classmethod
+    def get_rbln_cls(cls, pretrained_model_name_or_path: Union[str, Path], export: bool = None, **kwargs):
+        if isinstance(pretrained_model_name_or_path, Path):
+            pretrained_model_name_or_path = pretrained_model_name_or_path.as_posix()
+
+        if export is None:
+            export = not cls._is_compiled_pipeline(pretrained_model_name_or_path, **kwargs)
+
+        if export:
+            hf_model_class = cls.infer_hf_model_class(pretrained_model_name_or_path, **kwargs)
+            rbln_class_name = convert_hf_to_rbln_model_name(hf_model_class.__name__)
+        else:
+            rbln_class_name = cls.get_rbln_model_cls_name(pretrained_model_name_or_path, **kwargs)
+            if convert_rbln_to_hf_model_name(rbln_class_name) not in cls._model_mapping_names.values():
+                raise ValueError(
+                    f"The architecture '{rbln_class_name}' is not supported by the `{cls.__name__}.from_pretrained()` method. "
+                    "Please use the `from_pretrained()` method of the appropriate class to load this model, "
+                    f"or directly use '{rbln_class_name}.from_pretrained()`."
+                )
+
+        try:
+            rbln_cls = get_rbln_model_cls(rbln_class_name)
+        except AttributeError as e:
+            raise AttributeError(
+                f"Class '{rbln_class_name}' not found in 'optimum.rbln' module for model ID '{pretrained_model_name_or_path}'. "
+                "Ensure that the class name is correctly mapped and available in the 'optimum.rbln' module."
+            ) from e
+
+        return rbln_cls
+
+    @classmethod
+    def get_rbln_model_cls_name(cls, pretrained_model_name_or_path: Union[str, Path], **kwargs):
+        """
+        Retrieve the path to the compiled model directory for a given RBLN model.
+
+        Args:
+            pretrained_model_name_or_path (str): Identifier of the model.
+
+        Returns:
+            str: Path to the compiled model directory.
+        """
+        model_index_config = cls.load_config(pretrained_model_name_or_path)
+
+        if "_class_name" not in model_index_config:
+            raise ValueError(
+                "The `_class_name` field is missing from model_index_config. This is unexpected and should be reported as an issue. "
+                "Please use the `from_pretrained()` method of the appropriate class to load this model."
+            )
+
+        return model_index_config["_class_name"]
+
+    @classmethod
+    def _is_compiled_pipeline(
+        cls,
+        pretrained_model_name_or_path: Union[str, Path],
+        cache_dir=None,
+        force_download=False,
+        proxies=None,
+        token=None,
+        local_files_only=False,
+        revision=None,
+        **kwargs,
+    ):
+        config: dict = cls.load_config(
+            pretrained_model_name_or_path,
+            cache_dir=cache_dir,
+            force_download=force_download,
+            proxies=proxies,
+            token=token,
+            local_files_only=local_files_only,
+            revision=revision,
+        )
+        for value in config.values():
+            if isinstance(value, list) and len(value) > 0 and value[0] == "optimum.rbln":
+                return True
+        return False
+
+    @classmethod
+    def infer_hf_model_class(
+        cls,
+        pretrained_model_or_path: Union[str, Path],
+        cache_dir=None,
+        force_download=False,
+        proxies=None,
+        token=None,
+        local_files_only=False,
+        revision=None,
+        **kwargs,
+    ):
+        config = cls.load_config(
+            pretrained_model_or_path,
+            cache_dir=cache_dir,
+            force_download=force_download,
+            proxies=proxies,
+            token=token,
+            local_files_only=local_files_only,
+            revision=revision,
+        )
+        pipeline_key_name = cls.get_pipeline_key_name(config, **kwargs)
+
+        pipeline_cls = _get_task_class(cls._model_mapping, pipeline_key_name)
+
+        return pipeline_cls
+
+    @classmethod
+    def get_pipeline_key_name(cls, config, **kwargs):
+        orig_class_name = config["_class_name"]
+        if "ControlPipeline" in orig_class_name:
+            to_replace = "ControlPipeline"
+        else:
+            to_replace = "Pipeline"
+
+        if "controlnet" in kwargs:
+            if isinstance(kwargs["controlnet"], ControlNetUnionModel):
+                orig_class_name = config["_class_name"].replace(to_replace, "ControlNetUnionPipeline")
+            else:
+                orig_class_name = config["_class_name"].replace(to_replace, "ControlNetPipeline")
+        if "enable_pag" in kwargs:
+            enable_pag = kwargs.pop("enable_pag")
+            if enable_pag:
+                orig_class_name = orig_class_name.replace(to_replace, "PAGPipeline")
+
+        return orig_class_name
+
+    @classmethod
+    @validate_hf_hub_args
+    def from_pretrained(
+        cls,
+        model_id: Union[str, Path],
+        *,
+        export: bool = None,
+        rbln_config: Union[Dict[str, Any], RBLNModelConfig] = {},
+        **kwargs: Any,
+    ):
+        """
+        Load an RBLN-accelerated Diffusers pipeline from a pretrained checkpoint or a compiled RBLN artifact.
+
+        This method determines the concrete `RBLN*` model class that corresponds to the
+        underlying Diffusers pipeline architecture and dispatches to that class's
+        `from_pretrained()` implementation. If a compiled RBLN folder is detected at `model_id`
+        (or `export=False` is explicitly passed), it loads the compiled artifacts; otherwise it
+        compiles from the original Diffusers checkpoint.
+
+        Args:
+            model_id:
+                HF repo id or local path. For compiled models, this should point to a directory
+                (optionally under `subfolder`) that contains `*.rbln` files and `rbln_config.json`.
+            export:
+                Force compilation from a Diffusers checkpoint. When `None`, this is inferred by
+                checking whether compiled artifacts exist at `model_id`.
+            rbln_config:
+                RBLN compilation/runtime configuration. May be provided as a dictionary or as an
+                instance of the specific model's config class (e.g., `RBLNFluxPipelineConfig`).
+            kwargs: Additional keyword arguments.
+                - Arguments prefixed with `rbln_` are forwarded to the RBLN config.
+                - Remaining arguments are forwarded to the Diffusers loader.
+
+        Returns:
+            RBLNBaseModel: An instantiated RBLN model wrapping the Diffusers pipeline, ready for
+            inference on RBLN NPUs.
+
+        """
+        rbln_cls = cls.get_rbln_cls(model_id, export=export, **kwargs)
+        return rbln_cls.from_pretrained(model_id, export=export, rbln_config=rbln_config, **kwargs)
+
+    @staticmethod
+    def register(rbln_cls: Type[RBLNBaseModel], exist_ok=False):
+        """
+        Register a new RBLN model class.
+
+        Args:
+            rbln_cls (Type[RBLNBaseModel]): The RBLN model class to register.
+            exist_ok (bool): Whether to allow registering an already registered model.
+        """
+        if not issubclass(rbln_cls, RBLNBaseModel):
+            raise ValueError("`rbln_cls` must be a subclass of RBLNBaseModel.")
+
+        native_cls = getattr(importlib.import_module("optimum.rbln"), rbln_cls.__name__, None)
+        if rbln_cls.__name__ in MODEL_MAPPING or native_cls is not None:
+            if not exist_ok:
+                raise ValueError(f"Model for {rbln_cls.__name__} already registered.")
+
+        MODEL_MAPPING[rbln_cls.__name__] = rbln_cls
+
+
+class RBLNAutoPipelineForText2Image(RBLNAutoPipelineBase, AutoPipelineForText2Image):
+    """Text2Image AutoPipeline for RBLN NPUs."""
+
+    _model_mapping = AUTO_TEXT2IMAGE_PIPELINES_MAPPING
+    _model_mapping_names = {x[0]: x[1].__name__ for x in AUTO_TEXT2IMAGE_PIPELINES_MAPPING.items()}
+
+
+class RBLNAutoPipelineForImage2Image(RBLNAutoPipelineBase, AutoPipelineForImage2Image):
+    """Image2Image AutoPipeline for RBLN NPUs."""
+
+    _model_mapping = AUTO_IMAGE2IMAGE_PIPELINES_MAPPING
+    _model_mapping_names = {x[0]: x[1].__name__ for x in AUTO_IMAGE2IMAGE_PIPELINES_MAPPING.items()}
+
+    @classmethod
+    def get_pipeline_key_name(cls, config, **kwargs):
+        orig_class_name = config["_class_name"]
+        # the `orig_class_name` can be:
+        # `- *Pipeline` (for regular text-to-image checkpoint)
+        #  - `*ControlPipeline` (for Flux tools specific checkpoint)
+        # `- *Img2ImgPipeline` (for refiner checkpoint)
+        if "Img2Img" in orig_class_name:
+            to_replace = "Img2ImgPipeline"
+        elif "ControlPipeline" in orig_class_name:
+            to_replace = "ControlPipeline"
+        else:
+            to_replace = "Pipeline"
+
+        if "controlnet" in kwargs:
+            if isinstance(kwargs["controlnet"], ControlNetUnionModel):
+                orig_class_name = orig_class_name.replace(to_replace, "ControlNetUnion" + to_replace)
+            else:
+                orig_class_name = orig_class_name.replace(to_replace, "ControlNet" + to_replace)
+        if "enable_pag" in kwargs:
+            enable_pag = kwargs.pop("enable_pag")
+            if enable_pag:
+                orig_class_name = orig_class_name.replace(to_replace, "PAG" + to_replace)
+
+        if to_replace == "ControlPipeline":
+            orig_class_name = orig_class_name.replace(to_replace, "ControlImg2ImgPipeline")
+
+        return orig_class_name
+
+
+class RBLNAutoPipelineForInpainting(RBLNAutoPipelineBase, AutoPipelineForInpainting):
+    """Inpainting AutoPipeline for RBLN NPUs."""
+
+    _model_mapping = AUTO_INPAINT_PIPELINES_MAPPING
+    _model_mapping_names = {x[0]: x[1].__name__ for x in AUTO_INPAINT_PIPELINES_MAPPING.items()}
+
+    @classmethod
+    def get_pipeline_key_name(cls, config, **kwargs):
+        orig_class_name = config["_class_name"]
+
+        # The `orig_class_name`` can be:
+        # `- *InpaintPipeline` (for inpaint-specific checkpoint)
+        #  - `*ControlPipeline` (for Flux tools specific checkpoint)
+        #  - or *Pipeline (for regular text-to-image checkpoint)
+        if "Inpaint" in orig_class_name:
+            to_replace = "InpaintPipeline"
+        elif "ControlPipeline" in orig_class_name:
+            to_replace = "ControlPipeline"
+        else:
+            to_replace = "Pipeline"
+
+        if "controlnet" in kwargs:
+            if isinstance(kwargs["controlnet"], ControlNetUnionModel):
+                orig_class_name = orig_class_name.replace(to_replace, "ControlNetUnion" + to_replace)
+            else:
+                orig_class_name = orig_class_name.replace(to_replace, "ControlNet" + to_replace)
+        if "enable_pag" in kwargs:
+            enable_pag = kwargs.pop("enable_pag")
+            if enable_pag:
+                orig_class_name = orig_class_name.replace(to_replace, "PAG" + to_replace)
+        if to_replace == "ControlPipeline":
+            orig_class_name = orig_class_name.replace(to_replace, "ControlInpaintPipeline")
+
+        return orig_class_name

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

@@ -12,10 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Dict, Optional, Tuple
+from typing import Any, Optional, Tuple
 
 from ....configuration_utils import RBLNAutoConfig, RBLNModelConfig
-from ....transformers import RBLNLlamaForCausalLMConfig, RBLNSiglipVisionModelConfig
+from ....transformers import RBLNSiglipVisionModelConfig
 
 
 class RBLNVideoSafetyModelConfig(RBLNModelConfig):
@@ -56,11 +56,11 @@ class RBLNCosmosSafetyCheckerConfig(RBLNModelConfig):
     Configuration class for RBLN Cosmos Safety Checker.
     """
 
-    submodules = ["aegis", "video_safety_model", "face_blur_filter", "siglip_encoder"]
+    submodules = ["llamaguard3", "video_safety_model", "face_blur_filter", "siglip_encoder"]
 
     def __init__(
         self,
-        aegis: Optional[RBLNModelConfig] = None,
+        llamaguard3: Optional[RBLNModelConfig] = None,
         video_safety_model: Optional[RBLNModelConfig] = None,
         face_blur_filter: Optional[RBLNModelConfig] = None,
         siglip_encoder: Optional[RBLNSiglipVisionModelConfig] = None,
@@ -69,37 +69,40 @@ class RBLNCosmosSafetyCheckerConfig(RBLNModelConfig):
         image_size: Optional[Tuple[int, int]] = None,
         height: Optional[int] = None,
         width: Optional[int] = None,
-        **kwargs: Dict[str, Any],
+        max_seq_len: Optional[int] = None,
+        **kwargs: Any,
     ):
         super().__init__(**kwargs)
         if height is not None and width is not None:
             image_size = (height, width)
 
+        if max_seq_len is None:
+            max_seq_len = 512
+
         tensor_parallel_size = kwargs.get("tensor_parallel_size")
 
-        self.aegis = self.init_submodule_config(
-            RBLNLlamaForCausalLMConfig,
-            aegis,
+        self.llamaguard3 = self.initialize_submodule_config(
+            llamaguard3,
+            cls_name="RBLNLlamaForCausalLMConfig",
             batch_size=batch_size,
             tensor_parallel_size=tensor_parallel_size,
+            max_seq_len=max_seq_len,
         )
-
-        self.siglip_encoder = self.init_submodule_config(
-            RBLNSiglipVisionModelConfig,
+        self.siglip_encoder = self.initialize_submodule_config(
             siglip_encoder,
+            cls_name="RBLNSiglipVisionModelConfig",
             batch_size=batch_size,
             image_size=(384, 384),
         )
-
-        self.video_safety_model = self.init_submodule_config(
-            RBLNVideoSafetyModelConfig,
+        self.video_safety_model = self.initialize_submodule_config(
             video_safety_model,
+            cls_name="RBLNVideoSafetyModelConfig",
             batch_size=batch_size,
             input_size=1152,
         )
-        self.face_blur_filter = self.init_submodule_config(
-            RBLNRetinaFaceFilterConfig,
+        self.face_blur_filter = self.initialize_submodule_config(
             face_blur_filter,
+            cls_name="RBLNRetinaFaceFilterConfig",
             batch_size=batch_size,
             image_size=image_size,
         )