optimum-rbln 0.9.3__py3-none-any.whl → 0.9.3rc0__py3-none-any.whl

optimum/rbln/__init__.py

@@ -186,16 +186,12 @@ _import_structure = {
     "diffusers": [
         "RBLNAutoencoderKL",
         "RBLNAutoencoderKLConfig",
-        "RBLNAutoencoderKLTemporalDecoder",
-        "RBLNAutoencoderKLTemporalDecoderConfig",
         "RBLNAutoencoderKLCosmos",
         "RBLNAutoencoderKLCosmosConfig",
         "RBLNAutoPipelineForImage2Image",
         "RBLNAutoPipelineForInpainting",
         "RBLNAutoPipelineForText2Image",
         "RBLNControlNetModel",
-        "RBLNUNetSpatioTemporalConditionModel",
-        "RBLNStableVideoDiffusionPipeline",
         "RBLNControlNetModelConfig",
         "RBLNCosmosTextToWorldPipeline",
         "RBLNCosmosVideoToWorldPipeline",
@@ -254,8 +250,6 @@ _import_structure = {
         "RBLNUNet2DConditionModelConfig",
         "RBLNVQModel",
         "RBLNVQModelConfig",
-        "RBLNUNetSpatioTemporalConditionModelConfig",
-        "RBLNStableVideoDiffusionPipelineConfig",
     ],
 }
 
@@ -266,8 +260,6 @@ if TYPE_CHECKING:
         RBLNAutoencoderKLConfig,
         RBLNAutoencoderKLCosmos,
         RBLNAutoencoderKLCosmosConfig,
-        RBLNAutoencoderKLTemporalDecoder,
-        RBLNAutoencoderKLTemporalDecoderConfig,
         RBLNAutoPipelineForImage2Image,
         RBLNAutoPipelineForInpainting,
         RBLNAutoPipelineForText2Image,
@@ -326,12 +318,8 @@ if TYPE_CHECKING:
         RBLNStableDiffusionXLInpaintPipelineConfig,
         RBLNStableDiffusionXLPipeline,
         RBLNStableDiffusionXLPipelineConfig,
-        RBLNStableVideoDiffusionPipeline,
-        RBLNStableVideoDiffusionPipelineConfig,
         RBLNUNet2DConditionModel,
         RBLNUNet2DConditionModelConfig,
-        RBLNUNetSpatioTemporalConditionModel,
-        RBLNUNetSpatioTemporalConditionModelConfig,
         RBLNVQModel,
         RBLNVQModelConfig,
     )

optimum/rbln/__version__.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.9.3'
-__version_tuple__ = version_tuple = (0, 9, 3)
+__version__ = version = '0.9.3rc0'
+__version_tuple__ = version_tuple = (0, 9, 3, 'rc0')
 
 __commit_id__ = commit_id = None

optimum/rbln/configuration_utils.py

@@ -24,7 +24,7 @@ import torch
 from packaging.version import Version
 
 from .__version__ import __version__
-from .utils.deprecation import warn_deprecated_npu
+from .utils.depreacate_utils import warn_deprecated_npu
 from .utils.logging import get_logger
 from .utils.runtime_utils import ContextRblnConfig
 
@@ -528,7 +528,6 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
     ]
     submodules: List[str] = []
     subclass_non_save_attributes = []
-    _allow_no_compile_cfgs = False
 
     def initialize_submodule_config(
         self,
@@ -809,8 +808,7 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
             or len(self._compile_cfgs) == 0
             or not all(isinstance(cfg, RBLNCompileConfig) for cfg in self._compile_cfgs)
         ):
-            if not self._allow_no_compile_cfgs:
-                raise RuntimeError("`compile_cfgs` must contain at least one `RBLNCompileConfig` before freezing.")
+            raise RuntimeError("`compile_cfgs` must be set before freezing.")
 
         for submodule_name in self.submodules:
             submodule_config = getattr(self, submodule_name, None)

optimum/rbln/diffusers/__init__.py

@@ -57,9 +57,6 @@ _import_structure = {
         "RBLNSD3Transformer2DModelConfig",
         "RBLNUNet2DConditionModelConfig",
         "RBLNVQModelConfig",
-        "RBLNUNetSpatioTemporalConditionModelConfig",
-        "RBLNStableVideoDiffusionPipelineConfig",
-        "RBLNAutoencoderKLTemporalDecoderConfig",
     ],
     "pipelines": [
         "RBLNAutoPipelineForImage2Image",
@@ -89,17 +86,14 @@ _import_structure = {
         "RBLNStableDiffusion3Pipeline",
         "RBLNStableDiffusion3Img2ImgPipeline",
         "RBLNStableDiffusion3InpaintPipeline",
-        "RBLNStableVideoDiffusionPipeline",
     ],
     "models": [
         "RBLNAutoencoderKL",
         "RBLNAutoencoderKLCosmos",
         "RBLNUNet2DConditionModel",
-        "RBLNUNetSpatioTemporalConditionModel",
         "RBLNControlNetModel",
         "RBLNCosmosTransformer3DModel",
         "RBLNSD3Transformer2DModel",
-        "RBLNAutoencoderKLTemporalDecoder",
         "RBLNPriorTransformer",
         "RBLNVQModel",
     ],
@@ -112,7 +106,6 @@ if TYPE_CHECKING:
     from .configurations import (
         RBLNAutoencoderKLConfig,
         RBLNAutoencoderKLCosmosConfig,
-        RBLNAutoencoderKLTemporalDecoderConfig,
         RBLNControlNetModelConfig,
         RBLNCosmosTextToWorldPipelineConfig,
         RBLNCosmosTransformer3DModelConfig,
@@ -139,22 +132,18 @@ if TYPE_CHECKING:
         RBLNStableDiffusionXLImg2ImgPipelineConfig,
         RBLNStableDiffusionXLInpaintPipelineConfig,
         RBLNStableDiffusionXLPipelineConfig,
-        RBLNStableVideoDiffusionPipelineConfig,
         RBLNUNet2DConditionModelConfig,
-        RBLNUNetSpatioTemporalConditionModelConfig,
         RBLNVQModelConfig,
     )
     from .modeling_diffusers import RBLNDiffusionMixin
     from .models import (
         RBLNAutoencoderKL,
         RBLNAutoencoderKLCosmos,
-        RBLNAutoencoderKLTemporalDecoder,
         RBLNControlNetModel,
         RBLNCosmosTransformer3DModel,
         RBLNPriorTransformer,
         RBLNSD3Transformer2DModel,
         RBLNUNet2DConditionModel,
-        RBLNUNetSpatioTemporalConditionModel,
         RBLNVQModel,
     )
     from .pipelines import (
@@ -185,7 +174,6 @@ if TYPE_CHECKING:
         RBLNStableDiffusionXLImg2ImgPipeline,
         RBLNStableDiffusionXLInpaintPipeline,
         RBLNStableDiffusionXLPipeline,
-        RBLNStableVideoDiffusionPipeline,
     )
 else:
     import sys

optimum/rbln/diffusers/configurations/__init__.py

@@ -1,13 +1,11 @@
 from .models import (
     RBLNAutoencoderKLConfig,
     RBLNAutoencoderKLCosmosConfig,
-    RBLNAutoencoderKLTemporalDecoderConfig,
     RBLNControlNetModelConfig,
     RBLNCosmosTransformer3DModelConfig,
     RBLNPriorTransformerConfig,
     RBLNSD3Transformer2DModelConfig,
     RBLNUNet2DConditionModelConfig,
-    RBLNUNetSpatioTemporalConditionModelConfig,
     RBLNVQModelConfig,
 )
 from .pipelines import (
@@ -33,5 +31,4 @@ from .pipelines import (
     RBLNStableDiffusionXLImg2ImgPipelineConfig,
     RBLNStableDiffusionXLInpaintPipelineConfig,
     RBLNStableDiffusionXLPipelineConfig,
-    RBLNStableVideoDiffusionPipelineConfig,
 )

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

@@ -1,10 +1,8 @@
 from .configuration_autoencoder_kl import RBLNAutoencoderKLConfig
 from .configuration_autoencoder_kl_cosmos import RBLNAutoencoderKLCosmosConfig
-from .configuration_autoencoder_kl_temporal_decoder import RBLNAutoencoderKLTemporalDecoderConfig
 from .configuration_controlnet import RBLNControlNetModelConfig
 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_unet_spatio_temporal_condition import RBLNUNetSpatioTemporalConditionModelConfig
 from .configuration_vq_model import RBLNVQModelConfig

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

@@ -29,6 +29,3 @@ from .configuration_stable_diffusion_xl import (
     RBLNStableDiffusionXLInpaintPipelineConfig,
     RBLNStableDiffusionXLPipelineConfig,
 )
-from .configuration_stable_video_diffusion import (
-    RBLNStableVideoDiffusionPipelineConfig,
-)

optimum/rbln/diffusers/models/__init__.py

@@ -22,11 +22,9 @@ _import_structure = {
         "RBLNAutoencoderKL",
         "RBLNAutoencoderKLCosmos",
         "RBLNVQModel",
-        "RBLNAutoencoderKLTemporalDecoder",
     ],
     "unets": [
         "RBLNUNet2DConditionModel",
-        "RBLNUNetSpatioTemporalConditionModel",
     ],
     "controlnet": ["RBLNControlNetModel"],
     "transformers": [
@@ -37,22 +35,10 @@ _import_structure = {
 }
 
 if TYPE_CHECKING:
-    from .autoencoders import (
-        RBLNAutoencoderKL,
-        RBLNAutoencoderKLCosmos,
-        RBLNAutoencoderKLTemporalDecoder,
-        RBLNVQModel,
-    )
+    from .autoencoders import RBLNAutoencoderKL, RBLNAutoencoderKLCosmos, RBLNVQModel
     from .controlnet import RBLNControlNetModel
-    from .transformers import (
-        RBLNCosmosTransformer3DModel,
-        RBLNPriorTransformer,
-        RBLNSD3Transformer2DModel,
-    )
-    from .unets import (
-        RBLNUNet2DConditionModel,
-        RBLNUNetSpatioTemporalConditionModel,
-    )
+    from .transformers import RBLNCosmosTransformer3DModel, RBLNPriorTransformer, RBLNSD3Transformer2DModel
+    from .unets import RBLNUNet2DConditionModel
 else:
     import sys
 

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

@@ -14,5 +14,4 @@
 
 from .autoencoder_kl import RBLNAutoencoderKL
 from .autoencoder_kl_cosmos import RBLNAutoencoderKLCosmos
-from .autoencoder_kl_temporal_decoder import RBLNAutoencoderKLTemporalDecoder
 from .vq_model import RBLNVQModel

optimum/rbln/diffusers/models/autoencoders/autoencoder_kl_cosmos.py

@@ -68,7 +68,7 @@ class RBLNAutoencoderKLCosmos(RBLNModel):
         self.image_size = self.rbln_config.image_size
 
     @classmethod
-    def _wrap_model_if_needed(
+    def wrap_model_if_needed(
         cls, model: torch.nn.Module, rbln_config: RBLNAutoencoderKLCosmosConfig
     ) -> torch.nn.Module:
         decoder_model = _VAECosmosDecoder(model)
@@ -98,7 +98,7 @@ class RBLNAutoencoderKLCosmos(RBLNModel):
 
             compiled_models = {}
             if rbln_config.uses_encoder:
-                encoder_model, decoder_model = cls._wrap_model_if_needed(model, rbln_config)
+                encoder_model, decoder_model = cls.wrap_model_if_needed(model, rbln_config)
                 enc_compiled_model = cls.compile(
                     encoder_model,
                     rbln_compile_config=rbln_config.compile_cfgs[0],
@@ -107,7 +107,7 @@ class RBLNAutoencoderKLCosmos(RBLNModel):
                 )
                 compiled_models["encoder"] = enc_compiled_model
             else:
-                decoder_model = cls._wrap_model_if_needed(model, rbln_config)
+                decoder_model = cls.wrap_model_if_needed(model, rbln_config)
             dec_compiled_model = cls.compile(
                 decoder_model,
                 rbln_compile_config=rbln_config.compile_cfgs[-1],

optimum/rbln/diffusers/models/autoencoders/vae.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, List
 
 import torch
 from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution, IdentityDistribution
@@ -21,7 +21,7 @@ from ....utils.runtime_utils import RBLNPytorchRuntime
 
 
 if TYPE_CHECKING:
-    from diffusers import AutoencoderKL, AutoencoderKLCosmos, AutoencoderKLTemporalDecoder, VQModel
+    from diffusers import AutoencoderKL, AutoencoderKLCosmos, VQModel
 
 
 class RBLNRuntimeVAEEncoder(RBLNPytorchRuntime):
@@ -67,37 +67,18 @@ class _VAEDecoder(torch.nn.Module):
         return vae_out
 
 
-class _VAETemporalDecoder(torch.nn.Module):
-    def __init__(self, vae: "AutoencoderKLTemporalDecoder"):
-        super().__init__()
-        self.vae = vae
-        self.num_frames = None
-
-    def forward(self, z):
-        vae_out = self.vae.decode(z, num_frames=self.num_frames, return_dict=False)
-        return vae_out
-
-
 class _VAEEncoder(torch.nn.Module):
-    def __init__(self, vae: Union["AutoencoderKL", "AutoencoderKLTemporalDecoder"]):
+    def __init__(self, vae: "AutoencoderKL"):
         super().__init__()
         self.vae = vae
 
     def encode(self, x: torch.FloatTensor, return_dict: bool = True):
-        if hasattr(self, "use_tiling") and hasattr(self, "use_slicing"):
-            if self.use_tiling and (
-                x.shape[-1] > self.tile_sample_min_size or x.shape[-2] > self.tile_sample_min_size
-            ):
-                return self.tiled_encode(x, return_dict=return_dict)
-
-            if self.use_slicing and x.shape[0] > 1:
-                encoded_slices = [self.encoder(x_slice) for x_slice in x.split(1)]
-                h = torch.cat(encoded_slices)
-            else:
-                h = self.encoder(x)
-                if self.quant_conv is not None:
-                    h = self.quant_conv(h)
+        if self.use_tiling and (x.shape[-1] > self.tile_sample_min_size or x.shape[-2] > self.tile_sample_min_size):
+            return self.tiled_encode(x, return_dict=return_dict)
 
+        if self.use_slicing and x.shape[0] > 1:
+            encoded_slices = [self.encoder(x_slice) for x_slice in x.split(1)]
+            h = torch.cat(encoded_slices)
         else:
             h = self.encoder(x)
             if self.quant_conv is not None:

optimum/rbln/diffusers/models/controlnet.py

@@ -118,7 +118,7 @@ class RBLNControlNetModel(RBLNModel):
         )
 
     @classmethod
-    def _wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNModelConfig) -> torch.nn.Module:
+    def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNModelConfig) -> torch.nn.Module:
         use_encoder_hidden_states = False
         for down_block in model.down_blocks:
             if use_encoder_hidden_states := getattr(down_block, "has_cross_attention", False):
@@ -219,21 +219,6 @@ class RBLNControlNetModel(RBLNModel):
         return_dict: bool = True,
         **kwargs,
     ):
-        """
-        Forward pass for the RBLN-optimized ControlNetModel.
-
-        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.
-            controlnet_cond (torch.FloatTensor): The conditional input tensor of shape `(batch_size, max_seq_len, hidden_size)`.
-            conditioning_scale (torch.Tensor): The scale factor for ControlNet outputs.
-            added_cond_kwargs (Dict[str, torch.Tensor]): Additional conditions for the Stable Diffusion XL UNet.
-            return_dict (bool): Whether or not to return a [`~diffusers.models.controlnets.controlnet.ControlNetOutput`] instead of a plain tuple
-
-        Returns:
-            (Union[`~diffusers.models.controlnets.controlnet.ControlNetOutput`], Tuple)
-        """
         sample_batch_size = sample.size()[0]
         compiled_batch_size = self.compiled_batch_size
         if sample_batch_size != compiled_batch_size and (

optimum/rbln/diffusers/models/transformers/prior_transformer.py

@@ -77,7 +77,7 @@ class RBLNPriorTransformer(RBLNModel):
         self.clip_std = artifacts["clip_std"]
 
     @classmethod
-    def _wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNModelConfig) -> torch.nn.Module:
+    def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNModelConfig) -> torch.nn.Module:
         return _PriorTransformer(model).eval()
 
     @classmethod
@@ -128,27 +128,13 @@ class RBLNPriorTransformer(RBLNModel):
 
     def forward(
         self,
-        hidden_states: torch.Tensor,
+        hidden_states,
         timestep: Union[torch.Tensor, float, int],
         proj_embedding: torch.Tensor,
         encoder_hidden_states: Optional[torch.Tensor] = None,
         attention_mask: Optional[torch.Tensor] = None,
         return_dict: bool = True,
     ):
-        """
-        Forward pass for the RBLN-optimized PriorTransformer.
-
-        Args:
-            hidden_states (torch.Tensor): The currently predicted image embeddings.
-            timestep (Union[torch.Tensor, float, int]): Current denoising step.
-            proj_embedding (torch.Tensor): Projected embedding vector the denoising process is conditioned on.
-            encoder_hidden_states (Optional[torch.Tensor]): Hidden states of the text embeddings the denoising process is conditioned on.
-            attention_mask (Optional[torch.Tensor]): Text mask for the text embeddings.
-            return_dict (bool): Whether or not to return a [`~diffusers.models.transformers.prior_transformer.PriorTransformerOutput`] instead of a plain tuple.
-
-        Returns:
-            (Union[`~diffusers.models.transformers.prior_transformer.PriorTransformerOutput`, Tuple])
-        """
         # Convert timestep(long) and attention_mask(bool) to float
         return super().forward(
             hidden_states,

optimum/rbln/diffusers/models/transformers/transformer_cosmos.py

@@ -185,7 +185,7 @@ class RBLNCosmosTransformer3DModel(RBLNModel):
         )
 
     @classmethod
-    def _wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNModelConfig) -> torch.nn.Module:
+    def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNModelConfig) -> torch.nn.Module:
         num_latent_frames = rbln_config.num_latent_frames
         latent_height = rbln_config.latent_height
         latent_width = rbln_config.latent_width
@@ -303,21 +303,6 @@ class RBLNCosmosTransformer3DModel(RBLNModel):
         padding_mask: Optional[torch.Tensor] = None,
         return_dict: bool = True,
     ):
-        """
-        Forward pass for the RBLN-optimized CosmosTransformer3DModel.
-
-        Args:
-            hidden_states (torch.Tensor): The currently predicted image embeddings.
-            timestep (torch.Tensor): Current denoising step.
-            encoder_hidden_states (torch.Tensor): Conditional embeddings (embeddings computed from the input conditions such as prompts) to use.
-            fps: (Optional[int]): Frames per second for the video being generated.
-            condition_mask (Optional[torch.Tensor]): Tensor of condition mask.
-            padding_mask (Optional[torch.Tensor]): Tensor of padding mask.
-            return_dict (bool): Whether or not to return a [`~diffusers.models.modeling_output.Transformer2DModelOutput`] instead of a plain tuple.
-
-        Returns:
-            (Union[`~diffusers.models.modeling_output.Transformer2DModelOutput`, Tuple])
-        """
         (
             hidden_states,
             temb,

optimum/rbln/diffusers/models/transformers/transformer_sd3.py

@@ -77,7 +77,7 @@ class RBLNSD3Transformer2DModel(RBLNModel):
         super().__post_init__(**kwargs)
 
     @classmethod
-    def _wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNModelConfig) -> torch.nn.Module:
+    def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNModelConfig) -> torch.nn.Module:
         return SD3Transformer2DModelWrapper(model).eval()
 
     @classmethod
@@ -161,19 +161,6 @@ class RBLNSD3Transformer2DModel(RBLNModel):
         return_dict: bool = True,
         **kwargs,
     ):
-        """
-        Forward pass for the RBLN-optimized SD3Transformer2DModel.
-
-        Args:
-            hidden_states (torch.FloatTensor): The currently predicted image embeddings.
-            encoder_hidden_states (torch.FloatTensor): Conditional embeddings (embeddings computed from the input conditions such as prompts) to use.
-            pooled_projections (torch.FloatTensor): Embeddings projected from the embeddings of input conditions.
-            timestep (torch.LongTensor): Current denoising step.
-            return_dict (bool): Whether or not to return a [`~diffusers.models.modeling_output.Transformer2DModelOutput`] instead of a plain tuple.
-
-        Returns:
-            (Union[`~diffusers.models.modeling_output.Transformer2DModelOutput`, Tuple])
-        """
         sample_batch_size = hidden_states.size()[0]
         compiled_batch_size = self.compiled_batch_size
         if sample_batch_size != compiled_batch_size and (

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

@@ -13,4 +13,3 @@
 # limitations under the License.
 
 from .unet_2d_condition import RBLNUNet2DConditionModel
-from .unet_spatio_temporal_condition import RBLNUNetSpatioTemporalConditionModel

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

@@ -171,7 +171,7 @@ class RBLNUNet2DConditionModel(RBLNModel):
             self.add_embedding = ADDEMBEDDING(LINEAR1(self.in_features))
 
     @classmethod
-    def _wrap_model_if_needed(
+    def wrap_model_if_needed(
         cls, model: torch.nn.Module, rbln_config: RBLNUNet2DConditionModelConfig
     ) -> torch.nn.Module:
         if model.config.addition_embed_type == "text_time":
@@ -349,22 +349,6 @@ class RBLNUNet2DConditionModel(RBLNModel):
         return_dict: bool = True,
         **kwargs,
     ) -> Union[UNet2DConditionOutput, Tuple]:
-        """
-        Forward pass for the RBLN-optimized UNet2DConditionModel.
-
-        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_cond_kwargs (Dict[str, torch.Tensor]): A kwargs dictionary containing additional embeddings that
-                if specified are added to the embeddings that are passed along to the UNet blocks.
-            down_block_additional_residuals (Optional[Tuple[torch.Tensor]]): A tuple of tensors that if specified are added to the residuals of down unet blocks.
-            mid_block_additional_residual (Optional[torch.Tensor]): A tensor that if specified is added to the residual of the middle unet block.
-            return_dict (bool): Whether or not to return a [`~diffusers.models.unets.unet_2d_condition.UNet2DConditionOutput`] instead of a plain tuple.
-
-        Returns:
-            (Union[`~diffusers.models.unets.unet_2d_condition.UNet2DConditionOutput`], Tuple)
-        """
         sample_batch_size = sample.size()[0]
         compiled_batch_size = self.compiled_batch_size
         if sample_batch_size != compiled_batch_size and (

optimum/rbln/diffusers/pipelines/__init__.py

@@ -59,9 +59,6 @@ _import_structure = {
         "RBLNStableDiffusion3Img2ImgPipeline",
         "RBLNStableDiffusion3InpaintPipeline",
     ],
-    "stable_video_diffusion": [
-        "RBLNStableVideoDiffusionPipeline",
-    ],
 }
 if TYPE_CHECKING:
     from .auto_pipeline import (
@@ -101,7 +98,6 @@ if TYPE_CHECKING:
         RBLNStableDiffusionXLInpaintPipeline,
         RBLNStableDiffusionXLPipeline,
     )
-    from .stable_video_diffusion import RBLNStableVideoDiffusionPipeline
 else:
     import sys
 

optimum/rbln/diffusers/pipelines/controlnet/multicontrolnet.py

@@ -96,26 +96,6 @@ 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/modeling.py

@@ -34,6 +34,49 @@ if TYPE_CHECKING:
 logger = get_logger(__name__)
 
 
+def _get_dtype(
+    cls,
+    dtype: Optional[Union[str, torch.dtype, dict]],
+    config: PretrainedConfig,
+) -> tuple[PretrainedConfig, Optional[torch.dtype], Optional[torch.dtype]]:
+    dtype_orig = None
+
+    if dtype is not None:
+        if isinstance(dtype, str):
+            if dtype == "auto":
+                if hasattr(config, "dtype") and config.dtype is not None:
+                    dtype = config.dtype
+                else:
+                    dtype = torch.get_default_dtype()
+            elif hasattr(torch, dtype):
+                dtype = getattr(torch, dtype)
+                config.dtype = dtype
+        elif isinstance(dtype, torch.dtype):
+            config.dtype = dtype
+        elif isinstance(dtype, dict):
+            for key, curr_dtype in dtype.items():
+                if hasattr(config, key):
+                    value = getattr(config, key)
+                    curr_dtype = curr_dtype if not isinstance(curr_dtype, str) else getattr(torch, curr_dtype)
+                    value.dtype = curr_dtype
+            # main torch dtype for modules that aren't part of any sub-config
+            dtype = dtype.get("")
+            dtype = dtype if not isinstance(dtype, str) else getattr(torch, dtype)
+            config.dtype = dtype
+            if dtype is None:
+                dtype = torch.float32
+        else:
+            raise ValueError(f"Invalid dtype: {dtype}")
+
+        dtype_orig = cls._set_default_dtype(dtype)
+    else:
+        # Use default dtype
+        default_dtype = torch.get_default_dtype()
+        config.dtype = default_dtype
+
+    return config, dtype, dtype_orig
+
+
 class RBLNModel(RBLNBaseModel):
     @classmethod
     def update_kwargs(cls, kwargs):
@@ -54,16 +97,13 @@ class RBLNModel(RBLNBaseModel):
         pass
 
     @classmethod
-    def _wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNModelConfig) -> torch.nn.Module:
+    def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNModelConfig) -> torch.nn.Module:
         # Wrap the model if needed.
         return model
 
     @classmethod
     def get_compiled_model(cls, model: "PreTrainedModel", rbln_config: RBLNModelConfig):
-        if rbln_config._allow_no_compile_cfgs:
-            return {}
-
-        model = cls._wrap_model_if_needed(model, rbln_config)
+        model = cls.wrap_model_if_needed(model, rbln_config)
         rbln_compile_config = rbln_config.compile_cfgs[0]
         compiled_model = cls.compile(
             model,
@@ -73,18 +113,6 @@ class RBLNModel(RBLNBaseModel):
         )
         return compiled_model
 
-    @classmethod
-    def _update_rbln_config(
-        cls,
-        preprocessors: Optional[Any],
-        model: Optional["PreTrainedModel"] = None,
-        model_config: Optional["PretrainedConfig"] = None,
-        rbln_config: Optional[RBLNModelConfig] = None,
-    ) -> RBLNModelConfig:
-        # Default implementation: return config as-is
-        # Subclasses should override to set compile_cfgs if needed
-        return rbln_config
-
     @classmethod
     def _reconstruct_model_if_needed(cls, model: "PreTrainedModel"):
         return model
@@ -249,9 +277,6 @@ class RBLNModel(RBLNBaseModel):
         compiled_models: List[rebel.RBLNCompiledModel],
         rbln_config: RBLNModelConfig,
     ) -> List[rebel.Runtime]:
-        if len(rbln_config.compile_cfgs) == 0:
-            return []
-
         if DEFAULT_COMPILED_MODEL_NAME not in rbln_config.device_map:
             cls._raise_missing_compiled_file_error([DEFAULT_COMPILED_MODEL_NAME])
 

optimum/rbln/modeling_base.py

@@ -71,6 +71,7 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
         self.rbln_config = rbln_config
         if not rbln_config.is_frozen():
             raise RuntimeError("`rbln_config` must be frozen. Please call `rbln_config.freeze()` first.")
+
         self.compiled_models = rbln_compiled_models
 
         # Registers the RBLN classes into the transformers AutoModel classes to avoid warnings when creating