optimum-rbln 0.8.1a5__py3-none-any.whl → 0.8.1a7__py3-none-any.whl

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

@@ -13,4 +13,5 @@
 # limitations under the License.
 
 from .autoencoder_kl import RBLNAutoencoderKL
+from .autoencoder_kl_cosmos import RBLNAutoencoderKLCosmos
 from .vq_model import RBLNVQModel

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

@@ -80,7 +80,12 @@ class RBLNAutoencoderKL(RBLNModel):
 
             wrapped_model.eval()
 
-            compiled_models[model_name] = cls.compile(wrapped_model, rbln_compile_config=rbln_config.compile_cfgs[i])
+            compiled_models[model_name] = cls.compile(
+                wrapped_model,
+                rbln_compile_config=rbln_config.compile_cfgs[i],
+                create_runtimes=rbln_config.create_runtimes,
+                device=rbln_config.device_map[model_name],
+            )
 
         return compiled_models
 

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

@@ -0,0 +1,219 @@
+# Copyright 2025 Rebellions Inc. All rights reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import TYPE_CHECKING, Dict, List, Union
+
+import rebel
+import torch
+from diffusers.models.autoencoders.autoencoder_kl_cosmos import AutoencoderKLCosmos, CosmosCausalConv3d
+from diffusers.models.autoencoders.vae import DecoderOutput
+from diffusers.models.modeling_outputs import AutoencoderKLOutput
+from torch.nn import functional as F
+from transformers import PretrainedConfig
+
+from ....configuration_utils import RBLNCompileConfig
+from ....modeling import RBLNModel
+from ....utils.logging import get_logger
+from ...configurations import RBLNAutoencoderKLCosmosConfig
+from .vae import RBLNRuntimeCosmosVAEDecoder, RBLNRuntimeCosmosVAEEncoder, _VAECosmosDecoder, _VAECosmosEncoder
+
+
+if TYPE_CHECKING:
+    import torch
+    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, PreTrainedModel
+
+    from ...modeling_diffusers import RBLNDiffusionMixin, RBLNDiffusionMixinConfig
+
+logger = get_logger(__name__)
+
+
+class RBLNAutoencoderKLCosmos(RBLNModel):
+    """
+    RBLN implementation of AutoencoderKLCosmos for diffusion models.
+
+    This model is used to accelerate AutoencoderKLCosmos models from diffusers library on RBLN NPUs.
+    It can be configured to include both encoder and decoder, or just the decoder part for latent-to-video
+    conversion.
+
+    This class inherits from [`RBLNModel`]. Check the superclass documentation for the generic methods
+    the library implements for all its models.
+    """
+
+    auto_model_class = AutoencoderKLCosmos
+    hf_library_name = "diffusers"
+    _rbln_config_class = RBLNAutoencoderKLCosmosConfig
+
+    def __post_init__(self, **kwargs):
+        super().__post_init__(**kwargs)
+
+        if self.rbln_config.uses_encoder:
+            self.encoder = RBLNRuntimeCosmosVAEEncoder(
+                runtime=self.model[0], main_input_name="x", use_slicing=self.rbln_config.use_slicing
+            )
+
+        self.decoder = RBLNRuntimeCosmosVAEDecoder(
+            runtime=self.model[-1], main_input_name="z", use_slicing=self.rbln_config.use_slicing
+        )
+        self.image_size = self.rbln_config.image_size
+
+    @classmethod
+    def wrap_model_if_needed(
+        cls, model: torch.nn.Module, rbln_config: RBLNAutoencoderKLCosmosConfig
+    ) -> torch.nn.Module:
+        decoder_model = _VAECosmosDecoder(model)
+        decoder_model.eval()
+
+        if rbln_config.uses_encoder:
+            encoder_model = _VAECosmosEncoder(model)
+            encoder_model.eval()
+            return encoder_model, decoder_model
+        else:
+            return decoder_model
+
+    @classmethod
+    def get_compiled_model(
+        cls, model, rbln_config: RBLNAutoencoderKLCosmosConfig
+    ) -> Dict[str, rebel.RBLNCompiledModel]:
+        def replaced_forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+            if self.temporal_pad != 0:
+                hidden_states_prev = hidden_states[:, :, :1, ...].repeat(1, 1, self.temporal_pad, 1, 1)
+                hidden_states = torch.cat([hidden_states_prev, hidden_states], dim=2)
+            hidden_states = F.pad(hidden_states, (*self.spatial_pad, 0, 0), mode=self.pad_mode, value=0.0)
+            return super(CosmosCausalConv3d, self).forward(hidden_states)
+
+        try:
+            original_forward = CosmosCausalConv3d.forward
+            CosmosCausalConv3d.forward = replaced_forward
+
+            compiled_models = {}
+            if rbln_config.uses_encoder:
+                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],
+                    create_runtimes=rbln_config.create_runtimes,
+                    device=rbln_config.device_map["encoder"],
+                )
+                compiled_models["encoder"] = enc_compiled_model
+            else:
+                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],
+                create_runtimes=rbln_config.create_runtimes,
+                device=rbln_config.device_map["decoder"],
+            )
+            compiled_models["decoder"] = dec_compiled_model
+
+        finally:
+            CosmosCausalConv3d.forward = original_forward
+
+        return compiled_models
+
+    @classmethod
+    def update_rbln_config_using_pipe(
+        cls, pipe: "RBLNDiffusionMixin", rbln_config: "RBLNDiffusionMixinConfig", submodule_name: str
+    ) -> "RBLNDiffusionMixinConfig":
+        rbln_config.vae.num_channels_latents = pipe.transformer.config.out_channels
+        rbln_config.vae.vae_scale_factor_temporal = pipe.vae_scale_factor_temporal
+        rbln_config.vae.vae_scale_factor_spatial = pipe.vae_scale_factor_spatial
+        return rbln_config
+
+    @classmethod
+    def _update_rbln_config(
+        cls,
+        preprocessors: Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"],
+        model: "PreTrainedModel",
+        model_config: "PretrainedConfig",
+        rbln_config: RBLNAutoencoderKLCosmosConfig,
+    ) -> RBLNAutoencoderKLCosmosConfig:
+        batch_size = 1 if rbln_config.use_slicing else rbln_config.batch_size
+        compile_cfgs = []
+        if rbln_config.uses_encoder:
+            vae_enc_input_info = [
+                (
+                    "x",
+                    [
+                        batch_size,
+                        model_config.in_channels,
+                        rbln_config.num_frames,
+                        rbln_config.height,
+                        rbln_config.width,
+                    ],
+                    "float32",
+                ),
+            ]
+            compile_cfgs.append(RBLNCompileConfig(compiled_model_name="encoder", input_info=vae_enc_input_info))
+
+        num_latent_frames = (rbln_config.num_frames - 1) // rbln_config.vae_scale_factor_temporal + 1
+        latent_height = rbln_config.height // rbln_config.vae_scale_factor_spatial
+        latent_width = rbln_config.width // rbln_config.vae_scale_factor_spatial
+
+        vae_dec_input_info = [
+            (
+                "z",
+                [
+                    batch_size,
+                    rbln_config.num_channels_latents,
+                    num_latent_frames,
+                    latent_height,
+                    latent_width,
+                ],
+                "float32",
+            ),
+        ]
+        compile_cfgs.append(RBLNCompileConfig(compiled_model_name="decoder", input_info=vae_dec_input_info))
+
+        rbln_config.set_compile_cfgs(compile_cfgs)
+        return rbln_config
+
+    @classmethod
+    def _create_runtimes(
+        cls,
+        compiled_models: List[rebel.RBLNCompiledModel],
+        rbln_config: RBLNAutoencoderKLCosmosConfig,
+    ) -> List[rebel.Runtime]:
+        if len(compiled_models) == 1:
+            # decoder
+            expected_models = ["decoder"]
+        else:
+            expected_models = ["encoder", "decoder"]
+
+        if any(model_name not in rbln_config.device_map for model_name in expected_models):
+            cls._raise_missing_compiled_file_error(expected_models)
+
+        device_vals = [rbln_config.device_map[model_name] for model_name in expected_models]
+        return [
+            rebel.Runtime(
+                compiled_model,
+                tensor_type="pt",
+                device=device_val,
+                activate_profiler=rbln_config.activate_profiler,
+            )
+            for compiled_model, device_val in zip(compiled_models, device_vals)
+        ]
+
+    def encode(self, x: torch.FloatTensor, return_dict: bool = True, **kwargs) -> torch.FloatTensor:
+        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:
+        decoded = self.decoder.decode(z)
+
+        if not return_dict:
+            return (decoded,)
+
+        return DecoderOutput(sample=decoded)

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

@@ -12,15 +12,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import List
+from typing import TYPE_CHECKING, List
 
 import torch
-from diffusers import AutoencoderKL, VQModel
-from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution
+from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution, IdentityDistribution
 
 from ....utils.runtime_utils import RBLNPytorchRuntime
 
 
+if TYPE_CHECKING:
+    from diffusers import AutoencoderKL, AutoencoderKLCosmos, VQModel
+
+
 class RBLNRuntimeVAEEncoder(RBLNPytorchRuntime):
     def encode(self, x: torch.FloatTensor, **kwargs) -> torch.FloatTensor:
         moments = self.forward(x.contiguous())
@@ -33,6 +36,27 @@ class RBLNRuntimeVAEDecoder(RBLNPytorchRuntime):
         return self.forward(z)
 
 
+class RBLNRuntimeCosmosVAEEncoder(RBLNPytorchRuntime):
+    def encode(self, x: torch.FloatTensor, **kwargs) -> torch.FloatTensor:
+        if self.use_slicing and x.shape[0] > 1:
+            encoded_slices = [self.forward(x_slice) for x_slice in x.split(1)]
+            h = torch.cat(encoded_slices)
+        else:
+            h = self.forward(x)
+        posterior = IdentityDistribution(h)
+        return posterior
+
+
+class RBLNRuntimeCosmosVAEDecoder(RBLNPytorchRuntime):
+    def decode(self, z: torch.FloatTensor, **kwargs) -> torch.FloatTensor:
+        if self.use_slicing and z.shape[0] > 1:
+            decoded_slices = [self.forward(z_slice) for z_slice in z.split(1)]
+            decoded = torch.cat(decoded_slices)
+        else:
+            decoded = self.forward(z)
+        return decoded
+
+
 class _VAEDecoder(torch.nn.Module):
     def __init__(self, vae: "AutoencoderKL"):
         super().__init__()
@@ -66,6 +90,26 @@ class _VAEEncoder(torch.nn.Module):
         return vae_out
 
 
+class _VAECosmosEncoder(torch.nn.Module):
+    def __init__(self, vae: "AutoencoderKLCosmos"):
+        super().__init__()
+        self.vae = vae
+
+    def forward(self, x):
+        vae_out = self.vae._encode(x)
+        return vae_out
+
+
+class _VAECosmosDecoder(torch.nn.Module):
+    def __init__(self, vae: "AutoencoderKLCosmos"):
+        super().__init__()
+        self.vae = vae
+
+    def forward(self, z):
+        vae_out = self.vae._decode(z, return_dict=False)
+        return vae_out
+
+
 class RBLNRuntimeVQEncoder(RBLNPytorchRuntime):
     def encode(self, x: torch.FloatTensor, **kwargs) -> torch.FloatTensor:
         h = self.forward(x.contiguous())
@@ -84,7 +128,7 @@ class RBLNRuntimeVQDecoder(RBLNPytorchRuntime):
 
 
 class _VQEncoder(torch.nn.Module):
-    def __init__(self, vq_model: VQModel):
+    def __init__(self, vq_model: "VQModel"):
         super().__init__()
         self.vq_model = vq_model
 
@@ -99,7 +143,7 @@ class _VQEncoder(torch.nn.Module):
 
 
 class _VQDecoder(torch.nn.Module):
-    def __init__(self, vq_model: VQModel):
+    def __init__(self, vq_model: "VQModel"):
         super().__init__()
         self.vq_model = vq_model
 

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

@@ -78,7 +78,12 @@ class RBLNVQModel(RBLNModel):
 
             wrapped_model.eval()
 
-            compiled_models[model_name] = cls.compile(wrapped_model, rbln_compile_config=rbln_config.compile_cfgs[i])
+            compiled_models[model_name] = cls.compile(
+                wrapped_model,
+                rbln_compile_config=rbln_config.compile_cfgs[i],
+                create_runtimes=rbln_config.create_runtimes,
+                device=rbln_config.device_map[model_name],
+            )
 
         return compiled_models
 

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

@@ -13,4 +13,5 @@
 # limitations under the License.
 
 from .prior_transformer import RBLNPriorTransformer
+from .transformer_cosmos import RBLNCosmosTransformer3DModel
 from .transformer_sd3 import RBLNSD3Transformer2DModel

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

@@ -0,0 +1,321 @@
+# 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 pathlib import Path
+from typing import TYPE_CHECKING, List, Optional, Union
+
+import rebel
+import torch
+from diffusers import CosmosTransformer3DModel
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.models.transformers.transformer_cosmos import (
+    CosmosEmbedding,
+    CosmosLearnablePositionalEmbed,
+    CosmosPatchEmbed,
+    CosmosRotaryPosEmbed,
+)
+from torchvision import transforms
+
+from ....configuration_utils import DEFAULT_COMPILED_MODEL_NAME, RBLNCompileConfig, RBLNModelConfig
+from ....modeling import RBLNModel
+from ....utils.logging import get_logger
+from ...configurations import RBLNCosmosTransformer3DModelConfig
+
+
+if TYPE_CHECKING:
+    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, PretrainedConfig, PreTrainedModel
+
+    from ...modeling_diffusers import RBLNCosmosTransformer3DModelConfig, RBLNDiffusionMixin, RBLNDiffusionMixinConfig
+
+
+logger = get_logger(__name__)
+
+
+class CosmosTransformer3DModelWrapper(torch.nn.Module):
+    def __init__(
+        self,
+        model: CosmosTransformer3DModel,
+        num_latent_frames: int = 16,
+        latent_height: int = 88,
+        latent_width: int = 160,
+    ) -> None:
+        super().__init__()
+        self.model = model
+        self.num_latent_frames = num_latent_frames
+        self.latent_height = latent_height
+        self.latent_width = latent_width
+        self.p_t, self.p_h, self.p_w = model.config.patch_size
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        embedded_timestep: torch.Tensor,
+        temb: torch.Tensor,
+        image_rotary_emb_0: torch.Tensor,
+        image_rotary_emb_1: torch.Tensor,
+        extra_pos_emb: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        return_dict: bool = False,
+    ):
+        image_rotary_emb = [image_rotary_emb_0, image_rotary_emb_1]
+        for block in self.model.transformer_blocks:
+            hidden_states = block(
+                hidden_states=hidden_states,
+                encoder_hidden_states=encoder_hidden_states,
+                embedded_timestep=embedded_timestep,
+                temb=temb,
+                image_rotary_emb=image_rotary_emb,
+                extra_pos_emb=extra_pos_emb,
+                attention_mask=attention_mask,
+            )
+        post_patch_num_frames = self.num_latent_frames // self.p_t
+        post_patch_height = self.latent_height // self.p_h
+        post_patch_width = self.latent_width // self.p_w
+        hidden_states = self.model.norm_out(hidden_states, embedded_timestep, temb)
+        hidden_states = self.model.proj_out(hidden_states)
+        hidden_states = hidden_states.unflatten(2, (self.p_h, self.p_w, self.p_t, -1))
+        hidden_states = hidden_states.unflatten(1, (post_patch_num_frames, post_patch_height, post_patch_width))
+        hidden_states = hidden_states.permute(0, 7, 1, 6, 2, 4, 3, 5)
+        hidden_states = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)
+
+        return (hidden_states,)
+
+
+class RBLNCosmosTransformer3DModel(RBLNModel):
+    """RBLN wrapper for the Cosmos Transformer model."""
+
+    hf_library_name = "diffusers"
+    auto_model_class = CosmosTransformer3DModel
+
+    def __post_init__(self, **kwargs):
+        super().__post_init__(**kwargs)
+        artifacts = torch.load(self.model_save_dir / self.subfolder / "torch_artifacts.pth", weights_only=False)
+
+        hidden_size = self.config.num_attention_heads * self.config.attention_head_dim
+        patch_embed_in_channels = (
+            self.config.in_channels + 1 if self.config.concat_padding_mask else self.config.in_channels
+        )
+        self.rope = CosmosRotaryPosEmbed(
+            hidden_size=self.config.attention_head_dim,
+            max_size=self.config.max_size,
+            patch_size=self.config.patch_size,
+            rope_scale=self.config.rope_scale,
+        )
+        self.rope.load_state_dict(artifacts["rope"])
+        if artifacts["learnable_pos_embed"] is None:
+            self.learnable_pos_embed = None
+        else:
+            self.learnable_pos_embed = CosmosLearnablePositionalEmbed(
+                hidden_size=hidden_size,
+                max_size=self.config.max_size,
+                patch_size=self.config.patch_size,
+            )
+            self.learnable_pos_embed.load_state_dict(artifacts["learnable_pos_embed"])
+        self.patch_embed = CosmosPatchEmbed(patch_embed_in_channels, hidden_size, self.config.patch_size, bias=False)
+        self.patch_embed.load_state_dict(artifacts["patch_embed"])
+        self.time_embed = CosmosEmbedding(hidden_size, hidden_size)
+        self.time_embed.load_state_dict(artifacts["time_embed"])
+
+    def compute_embedding(
+        self,
+        hidden_states: torch.Tensor,
+        timestep: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        fps: Optional[int] = None,
+        condition_mask: Optional[torch.Tensor] = None,
+        padding_mask: Optional[torch.Tensor] = None,
+    ):
+        batch_size, num_channels, num_frames, height, width = hidden_states.shape
+
+        # 1. Concatenate padding mask if needed & prepare attention mask
+        if condition_mask is not None:
+            hidden_states = torch.cat([hidden_states, condition_mask], dim=1)
+
+        if self.config.concat_padding_mask:
+            padding_mask = transforms.functional.resize(
+                padding_mask, list(hidden_states.shape[-2:]), interpolation=transforms.InterpolationMode.NEAREST
+            )
+            hidden_states = torch.cat(
+                [hidden_states, padding_mask.unsqueeze(2).repeat(batch_size, 1, num_frames, 1, 1)], dim=1
+            )
+
+        if attention_mask is not None:
+            attention_mask = attention_mask.unsqueeze(1).unsqueeze(1)  # [B, 1, 1, S]
+
+        # 2. Generate positional embeddings
+        image_rotary_emb = self.rope(hidden_states, fps=fps)
+        extra_pos_emb = self.learnable_pos_embed(hidden_states) if self.config.extra_pos_embed_type else None
+
+        # 3. Patchify input
+        p_t, p_h, p_w = self.config.patch_size
+        hidden_states = self.patch_embed(hidden_states)
+        hidden_states = hidden_states.flatten(1, 3)  # [B, T, H, W, C] -> [B, THW, C]
+
+        # 4. Timestep embeddings
+        temb, embedded_timestep = self.time_embed(hidden_states, timestep)
+
+        return (
+            hidden_states,
+            temb,
+            embedded_timestep,
+            image_rotary_emb[0],
+            image_rotary_emb[1],
+            extra_pos_emb,
+            attention_mask,
+        )
+
+    @classmethod
+    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
+        return CosmosTransformer3DModelWrapper(
+            model=model,
+            num_latent_frames=num_latent_frames,
+            latent_height=latent_height,
+            latent_width=latent_width,
+        ).eval()
+
+    @classmethod
+    def update_rbln_config_using_pipe(
+        cls, pipe: "RBLNDiffusionMixin", rbln_config: "RBLNDiffusionMixinConfig", submodule_name: str
+    ) -> RBLNCosmosTransformer3DModelConfig:
+        rbln_config.transformer.num_latent_frames = (
+            rbln_config.transformer.num_frames - 1
+        ) // pipe.vae_scale_factor_temporal + 1
+        rbln_config.transformer.latent_height = rbln_config.transformer.height // pipe.vae_scale_factor_spatial
+        rbln_config.transformer.latent_width = rbln_config.transformer.width // pipe.vae_scale_factor_spatial
+        rbln_config.transformer.max_seq_len = pipe.text_encoder.config.n_positions
+        rbln_config.transformer.embedding_dim = pipe.text_encoder.encoder.embed_tokens.embedding_dim
+
+        return rbln_config
+
+    @classmethod
+    def save_torch_artifacts(
+        cls,
+        model: "PreTrainedModel",
+        save_dir_path: Path,
+        subfolder: str,
+        rbln_config: RBLNModelConfig,
+    ):
+        save_dict = {}
+        save_dict["rope"] = model.rope.state_dict()
+        if model.learnable_pos_embed is not None:
+            save_dict["learnable_pos_embed"] = model.learnable_pos_embed.state_dict()
+        save_dict["patch_embed"] = model.patch_embed.state_dict()
+        save_dict["time_embed"] = model.time_embed.state_dict()
+        torch.save(save_dict, save_dir_path / subfolder / "torch_artifacts.pth")
+
+    @classmethod
+    def _update_rbln_config(
+        cls,
+        preprocessors: Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"],
+        model: "PreTrainedModel",
+        model_config: "PretrainedConfig",
+        rbln_config: "RBLNCosmosTransformer3DModelConfig",
+    ) -> RBLNCosmosTransformer3DModelConfig:
+        p_t, p_h, p_w = model_config.patch_size
+        hidden_dim = (
+            (rbln_config.num_latent_frames // p_t)
+            * (rbln_config.latent_height // p_h)
+            * (rbln_config.latent_width // p_w)
+        )
+        attention_head_dim = model_config.attention_head_dim
+        hidden_size = model.config.num_attention_heads * model.config.attention_head_dim
+        input_info = [
+            (
+                "hidden_states",
+                [
+                    rbln_config.batch_size,
+                    hidden_dim,
+                    hidden_size,
+                ],
+                "float32",
+            ),
+            (
+                "encoder_hidden_states",
+                [
+                    rbln_config.batch_size,
+                    rbln_config.max_seq_len,
+                    rbln_config.embedding_dim,
+                ],
+                "float32",
+            ),
+            ("embedded_timestep", [rbln_config.batch_size, hidden_size], "float32"),
+            ("temb", [1, hidden_size * 3], "float32"),
+            ("image_rotary_emb_0", [hidden_dim, attention_head_dim], "float32"),
+            ("image_rotary_emb_1", [hidden_dim, attention_head_dim], "float32"),
+            ("extra_pos_emb", [rbln_config.batch_size, hidden_dim, hidden_size], "float32"),
+        ]
+
+        compile_config = RBLNCompileConfig(input_info=input_info)
+        rbln_config.set_compile_cfgs([compile_config])
+        return rbln_config
+
+    @classmethod
+    def _create_runtimes(
+        cls,
+        compiled_models: List[rebel.RBLNCompiledModel],
+        rbln_config: RBLNModelConfig,
+    ) -> List[rebel.Runtime]:
+        if DEFAULT_COMPILED_MODEL_NAME not in rbln_config.device_map:
+            cls._raise_missing_compiled_file_error([DEFAULT_COMPILED_MODEL_NAME])
+
+        return [
+            rebel.Runtime(
+                compiled_model,
+                tensor_type="pt",
+                device=rbln_config.device_map[DEFAULT_COMPILED_MODEL_NAME],
+                activate_profiler=rbln_config.activate_profiler,
+                timeout=120,
+            )
+            for compiled_model in compiled_models
+        ]
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        timestep: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        fps: Optional[int] = None,
+        condition_mask: Optional[torch.Tensor] = None,
+        padding_mask: Optional[torch.Tensor] = None,
+        return_dict: bool = True,
+    ):
+        (
+            hidden_states,
+            temb,
+            embedded_timestep,
+            image_rotary_emb_0,
+            image_rotary_emb_1,
+            extra_pos_emb,
+            attention_mask,
+        ) = self.compute_embedding(hidden_states, timestep, attention_mask, fps, condition_mask, padding_mask)
+
+        hidden_states = self.model[0].forward(
+            hidden_states,
+            encoder_hidden_states,
+            embedded_timestep,
+            temb,
+            image_rotary_emb_0,
+            image_rotary_emb_1,
+            extra_pos_emb,
+        )
+
+        if not return_dict:
+            return (hidden_states,)
+        else:
+            return Transformer2DModelOutput(sample=hidden_states)