optimum-rbln 0.8.2a0__py3-none-any.whl → 0.9.3__py3-none-any.whl

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

@@ -0,0 +1,201 @@
+# 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 dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Union
+
+import torch
+from diffusers.models.unets.unet_spatio_temporal_condition import (
+    UNetSpatioTemporalConditionModel,
+    UNetSpatioTemporalConditionOutput,
+)
+from transformers import PretrainedConfig
+
+from ....configuration_utils import RBLNCompileConfig
+from ....modeling import RBLNModel
+from ....utils.logging import get_logger
+from ...configurations import RBLNUNetSpatioTemporalConditionModelConfig
+from ...modeling_diffusers import RBLNDiffusionMixin, RBLNDiffusionMixinConfig
+
+
+if TYPE_CHECKING:
+    from transformers import AutoFeatureExtractor, AutoProcessor, PreTrainedModel
+
+logger = get_logger(__name__)
+
+
+class _UNet_STCM(torch.nn.Module):
+    def __init__(self, unet: "UNetSpatioTemporalConditionModel"):
+        super().__init__()
+        self.unet = unet
+
+    def forward(
+        self,
+        sample: torch.Tensor,
+        timestep: Union[torch.Tensor, float, int],
+        encoder_hidden_states: torch.Tensor,
+        added_time_ids: torch.Tensor,
+    ) -> torch.Tensor:
+        unet_out = self.unet(
+            sample=sample,
+            timestep=timestep,
+            encoder_hidden_states=encoder_hidden_states,
+            added_time_ids=added_time_ids,
+            return_dict=False,
+        )
+        return unet_out
+
+
+class RBLNUNetSpatioTemporalConditionModel(RBLNModel):
+    hf_library_name = "diffusers"
+    auto_model_class = UNetSpatioTemporalConditionModel
+    _rbln_config_class = RBLNUNetSpatioTemporalConditionModelConfig
+    output_class = UNetSpatioTemporalConditionOutput
+    output_key = "sample"
+
+    def __post_init__(self, **kwargs):
+        super().__post_init__(**kwargs)
+        self.in_features = self.rbln_config.in_features
+        if self.in_features is not None:
+
+            @dataclass
+            class LINEAR1:
+                in_features: int
+
+            @dataclass
+            class ADDEMBEDDING:
+                linear_1: LINEAR1
+
+            self.add_embedding = ADDEMBEDDING(LINEAR1(self.in_features))
+
+    @classmethod
+    def _wrap_model_if_needed(
+        cls, model: torch.nn.Module, rbln_config: RBLNUNetSpatioTemporalConditionModelConfig
+    ) -> torch.nn.Module:
+        return _UNet_STCM(model).eval()
+
+    @classmethod
+    def get_unet_sample_size(
+        cls,
+        pipe: RBLNDiffusionMixin,
+        rbln_config: RBLNUNetSpatioTemporalConditionModelConfig,
+        image_size: Optional[Tuple[int, int]] = None,
+    ) -> Union[int, Tuple[int, int]]:
+        scale_factor = pipe.vae_scale_factor
+
+        if image_size is None:
+            vae_sample_size = pipe.vae.config.sample_size
+            if isinstance(vae_sample_size, int):
+                vae_sample_size = (vae_sample_size, vae_sample_size)
+
+            sample_size = (
+                vae_sample_size[0] // scale_factor,
+                vae_sample_size[1] // scale_factor,
+            )
+        else:
+            sample_size = (image_size[0] // scale_factor, image_size[1] // scale_factor)
+        return sample_size
+
+    @classmethod
+    def update_rbln_config_using_pipe(
+        cls, pipe: RBLNDiffusionMixin, rbln_config: "RBLNDiffusionMixinConfig", submodule_name: str
+    ) -> Dict[str, Any]:
+        rbln_config.unet.sample_size = cls.get_unet_sample_size(
+            pipe, rbln_config.unet, image_size=rbln_config.image_size
+        )
+        return rbln_config
+
+    @classmethod
+    def _update_rbln_config(
+        cls,
+        preprocessors: Union["AutoFeatureExtractor", "AutoProcessor"],
+        model: "PreTrainedModel",
+        model_config: "PretrainedConfig",
+        rbln_config: RBLNUNetSpatioTemporalConditionModelConfig,
+    ) -> RBLNUNetSpatioTemporalConditionModelConfig:
+        if rbln_config.num_frames is None:
+            rbln_config.num_frames = model_config.num_frames
+
+        if rbln_config.sample_size is None:
+            rbln_config.sample_size = model_config.sample_size
+
+        input_info = [
+            (
+                "sample",
+                [
+                    rbln_config.batch_size,
+                    rbln_config.num_frames,
+                    model_config.in_channels,
+                    rbln_config.sample_size[0],
+                    rbln_config.sample_size[1],
+                ],
+                "float32",
+            ),
+            ("timestep", [], "float32"),
+            ("encoder_hidden_states", [rbln_config.batch_size, 1, model_config.cross_attention_dim], "float32"),
+            ("added_time_ids", [rbln_config.batch_size, 3], "float32"),
+        ]
+
+        if hasattr(model_config, "addition_time_embed_dim"):
+            rbln_config.in_features = model_config.projection_class_embeddings_input_dim
+
+        rbln_compile_config = RBLNCompileConfig(input_info=input_info)
+        rbln_config.set_compile_cfgs([rbln_compile_config])
+
+        return rbln_config
+
+    @property
+    def compiled_batch_size(self):
+        return self.rbln_config.compile_cfgs[0].input_info[0][1][0]
+
+    def forward(
+        self,
+        sample: torch.Tensor,
+        timestep: Union[torch.Tensor, float, int],
+        encoder_hidden_states: torch.Tensor,
+        added_time_ids: torch.Tensor,
+        return_dict: bool = True,
+        **kwargs,
+    ) -> Union[UNetSpatioTemporalConditionOutput, Tuple]:
+        """
+        Forward pass for the RBLN-optimized UNetSpatioTemporalConditionModel.
+
+        Args:
+            sample (torch.Tensor): The noisy input tensor with the following shape `(batch, channel, height, width)`.
+            timestep (Union[torch.Tensor, float, int]): The number of timesteps to denoise an input.
+            encoder_hidden_states (torch.Tensor): The encoder hidden states.
+            added_time_ids (torch.Tensor): A tensor containing additional sinusoidal embeddings and added to the time embeddings.
+            return_dict (bool): Whether or not to return a [`~diffusers.models.unets.unet_spatio_temporal_condition.UNetSpatioTemporalConditionOutput`] instead of a plain tuple.
+
+        Returns:
+            (Union[`~diffusers.models.unets.unet_spatio_temporal_condition.UNetSpatioTemporalConditionOutput`], Tuple)
+        """
+        sample_batch_size = sample.size()[0]
+        compiled_batch_size = self.compiled_batch_size
+        if sample_batch_size != compiled_batch_size and (
+            sample_batch_size * 2 == compiled_batch_size or sample_batch_size == compiled_batch_size * 2
+        ):
+            raise ValueError(
+                f"Mismatch between UNet's runtime batch size ({sample_batch_size}) and compiled batch size ({compiled_batch_size}). "
+                "This may be caused by the 'guidance scale' parameter, which doubles the runtime batch size in Stable Diffusion. "
+                "Adjust the batch size during compilation or modify the 'guidance scale' to match the compiled batch size.\n\n"
+                "For details, see: https://docs.rbln.ai/software/optimum/model_api.html#stable-diffusion"
+            )
+        return super().forward(
+            sample.contiguous(),
+            timestep.float(),
+            encoder_hidden_states,
+            added_time_ids,
+            return_dict=return_dict,
+        )

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",
@@ -54,8 +59,16 @@ _import_structure = {
         "RBLNStableDiffusion3Img2ImgPipeline",
         "RBLNStableDiffusion3InpaintPipeline",
     ],
+    "stable_video_diffusion": [
+        "RBLNStableVideoDiffusionPipeline",
+    ],
 }
 if TYPE_CHECKING:
+    from .auto_pipeline import (
+        RBLNAutoPipelineForImage2Image,
+        RBLNAutoPipelineForInpainting,
+        RBLNAutoPipelineForText2Image,
+    )
     from .controlnet import (
         RBLNMultiControlNetModel,
         RBLNStableDiffusionControlNetImg2ImgPipeline,
@@ -63,11 +76,7 @@ if TYPE_CHECKING:
         RBLNStableDiffusionXLControlNetImg2ImgPipeline,
         RBLNStableDiffusionXLControlNetPipeline,
     )
-    from .cosmos import (
-        RBLNCosmosSafetyChecker,
-        RBLNCosmosTextToWorldPipeline,
-        RBLNCosmosVideoToWorldPipeline,
-    )
+    from .cosmos import RBLNCosmosSafetyChecker, RBLNCosmosTextToWorldPipeline, RBLNCosmosVideoToWorldPipeline
     from .kandinsky2_2 import (
         RBLNKandinskyV22CombinedPipeline,
         RBLNKandinskyV22Img2ImgCombinedPipeline,
@@ -92,6 +101,7 @@ if TYPE_CHECKING:
         RBLNStableDiffusionXLInpaintPipeline,
         RBLNStableDiffusionXLPipeline,
     )
+    from .stable_video_diffusion import RBLNStableVideoDiffusionPipeline
 else:
     import sys
 

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/controlnet/multicontrolnet.py

@@ -96,6 +96,26 @@ class RBLNMultiControlNetModel(RBLNModel):
         guess_mode: bool = False,
         return_dict: bool = True,
     ):
+        """
+        Forward pass for the RBLN-optimized MultiControlNetModel.
+
+        This method processes multiple ControlNet models in sequence, applying each one to the input sample
+        with its corresponding conditioning image and scale factor. The outputs from all ControlNets are
+        merged by addition to produce the final control signals.
+
+        Args:
+            sample (torch.FloatTensor): The noisy input tensor.
+            timestep (Union[torch.Tensor, float, int]): The number of timesteps to denoise an input.
+            encoder_hidden_states (torch.Tensor): The encoder hidden states from the text encoder.
+            controlnet_cond (List[torch.Tensor]): A list of conditional input tensors, one for each ControlNet model.
+            conditioning_scale (List[float]): A list of scale factors for each ControlNet output. Each scale
+                controls the strength of the corresponding ControlNet's influence on the generation.
+            return_dict (bool): Whether or not to return a dictionary instead of a plain tuple. Currently,
+                this method always returns a tuple regardless of this parameter.
+
+        Returns:
+            (Tuple[List[torch.Tensor], torch.Tensor])
+        """
         for i, (image, scale, controlnet) in enumerate(zip(controlnet_cond, conditioning_scale, self.nets)):
             down_samples, mid_sample = controlnet(
                 sample=sample.contiguous(),

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

@@ -12,7 +12,7 @@
 # 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 RBLNSiglipVisionModelConfig
@@ -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,29 +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)
 
-        self.aegis = self.init_submodule_config(RBLNModelConfig, aegis)
-        self.siglip_encoder = self.init_submodule_config(
-            RBLNSiglipVisionModelConfig,
+        if max_seq_len is None:
+            max_seq_len = 512
+
+        tensor_parallel_size = kwargs.get("tensor_parallel_size")
+
+        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.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,
         )